modelbeast/server/store.py
MODELBEAST 514ec6cdcc Security hardening from adversarial review (5-agent workflow)
Fixed confirmed findings before public exposure:
HIGH:
- upload filename path traversal → store.safe_name() strips to basename
- login rate-limit XFF bypass → key on request.client.host + per-username bucket;
  auth.check_login() burns bcrypt time on unknown users (no enumeration)
- cross-user read access → per-user isolation: guests see/use/download/delete only
  their own assets & jobs (owner sees all); WS job events scoped per-user
MEDIUM:
- unbounded upload read → bounded chunked streaming to the 1GB cap
- asset member path check → Path.is_relative_to boundary + ownership gate
- WS token-in-query-string leak → session-cookie-only WS auth
LOW:
- retry_job bypassed the per-user job cap → cap now checked on retry
- wholesale API-key injection → env_for_operator injects a paid key only to
  operators that declare it (guest local jobs never receive fal/OpenRouter keys)
- session revocation → users.session_epoch, bumped on password change
- int() 500s → 400; net-lane defense-in-depth (guests blocked by requires_env AND
  resources==net, so a mis-tagged paid op is still blocked)
+ public /api/health for serve.sh & proxy; docs/VPS.md; mb MB_TOKEN bearer auth

tests/smoke.sh: 34 checks passing incl. all new hardening.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
2026-07-13 12:25:13 +10:00

102 lines
3.6 KiB
Python

"""Asset storage: files land in data/assets/<id>/<name>, kinds inferred by extension."""
import json
import shutil
from pathlib import Path
from . import db
ASSETS_DIR = db.DATA / "assets"
KIND_BY_EXT = {
"video": {".mp4", ".mov", ".mkv", ".webm", ".avi", ".m4v"},
"image": {".png", ".jpg", ".jpeg", ".webp", ".heic", ".bmp", ".tif", ".tiff"},
"model": {".glb", ".gltf", ".obj", ".fbx", ".blend", ".usd", ".usdz", ".usda",
".ply", ".stl", ".spz", ".abc", ".dae", ".bvh"},
}
def kind_of(name: str) -> str:
ext = Path(name).suffix.lower()
for kind, exts in KIND_BY_EXT.items():
if ext in exts:
return kind
return "other"
def safe_name(name: str) -> str:
"""Strip any directory components from a user-supplied filename (path-traversal
guard) so the stored file can only ever land inside its own asset dir."""
base = Path(name or "").name
return base if base and base not in (".", "..") else "unnamed"
def register_file(con, src: Path, name: str | None = None, parent_job: str | None = None,
move: bool = False, meta: dict | None = None, user_id: str | None = None) -> dict:
"""Copy/move a file (or directory) into the asset store and record it."""
name = safe_name(name or src.name)
asset_id = db.new_id()
dest_dir = ASSETS_DIR / asset_id
dest_dir.mkdir(parents=True, exist_ok=True)
dest = dest_dir / name
meta = dict(meta or {})
if src.is_dir():
if move:
shutil.move(str(src), dest)
else:
shutil.copytree(src, dest)
size = sum(f.stat().st_size for f in dest.rglob("*") if f.is_file())
kind = "frames" if any(dest.glob("*.jpg")) or any(dest.glob("*.png")) else "folder"
else:
if move:
shutil.move(str(src), dest)
else:
shutil.copy2(src, dest)
size = dest.stat().st_size
kind = kind_of(name)
# operators may tag a semantic kind that isn't inferable from the extension
# (e.g. a splat .ply vs a mesh .ply, or a colmap_dataset folder)
kind = meta.pop("kind", kind)
con.execute(
"INSERT INTO assets (id, name, kind, path, size, meta, parent_job, created_at, user_id) "
"VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)",
(asset_id, name, kind, str(dest), size, json.dumps(meta or {}), parent_job, db.now(), user_id),
)
con.commit()
return get_asset(con, asset_id)
def register_upload(con, name: str, data: bytes, user_id: str | None = None) -> dict:
name = safe_name(name)
asset_id = db.new_id()
dest_dir = ASSETS_DIR / asset_id
dest_dir.mkdir(parents=True, exist_ok=True)
dest = dest_dir / name
dest.write_bytes(data)
con.execute(
"INSERT INTO assets (id, name, kind, path, size, meta, parent_job, created_at, user_id) "
"VALUES (?, ?, ?, ?, ?, '{}', NULL, ?, ?)",
(asset_id, name, kind_of(name), str(dest), len(data), db.now(), user_id),
)
con.commit()
return get_asset(con, asset_id)
def get_asset(con, asset_id: str) -> dict | None:
row = con.execute("SELECT * FROM assets WHERE id = ?", (asset_id,)).fetchone()
return db.row_to_dict(row) if row else None
def list_assets(con) -> list[dict]:
rows = con.execute("SELECT * FROM assets ORDER BY created_at DESC").fetchall()
return [db.row_to_dict(r) for r in rows]
def delete_asset(con, asset_id: str) -> bool:
asset = get_asset(con, asset_id)
if not asset:
return False
shutil.rmtree(ASSETS_DIR / asset_id, ignore_errors=True)
con.execute("DELETE FROM assets WHERE id = ?", (asset_id,))
con.commit()
return True