modelbeast/server/db.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

105 lines
3.0 KiB
Python

import json
import os
import sqlite3
import time
import uuid
from pathlib import Path
ROOT = Path(__file__).resolve().parent.parent
# Tests / alternate roots can override the data dir.
DATA = Path(os.environ.get("MODELBEAST_DATA", str(ROOT / "data")))
DB_PATH = DATA / "modelbeast.db"
SCHEMA = """
CREATE TABLE IF NOT EXISTS assets (
id TEXT PRIMARY KEY,
name TEXT NOT NULL,
kind TEXT NOT NULL,
path TEXT NOT NULL,
size INTEGER NOT NULL DEFAULT 0,
meta TEXT NOT NULL DEFAULT '{}',
parent_job TEXT,
created_at REAL NOT NULL
);
CREATE TABLE IF NOT EXISTS jobs (
id TEXT PRIMARY KEY,
operator TEXT NOT NULL,
status TEXT NOT NULL DEFAULT 'queued',
asset_id TEXT,
asset_ids TEXT NOT NULL DEFAULT '[]',
params TEXT NOT NULL DEFAULT '{}',
outdir TEXT,
log TEXT NOT NULL DEFAULT '',
error TEXT,
created_at REAL NOT NULL,
started_at REAL,
finished_at REAL
);
CREATE TABLE IF NOT EXISTS settings (
key TEXT PRIMARY KEY,
value TEXT NOT NULL DEFAULT ''
);
CREATE TABLE IF NOT EXISTS users (
id TEXT PRIMARY KEY,
username TEXT UNIQUE NOT NULL,
pw_hash TEXT NOT NULL,
role TEXT NOT NULL DEFAULT 'guest', -- 'owner' | 'guest'; guests are LOCAL-ONLY
max_active_jobs INTEGER NOT NULL DEFAULT 4,
created_at REAL NOT NULL
);
CREATE TABLE IF NOT EXISTS api_tokens (
token_hash TEXT PRIMARY KEY, -- sha256 hex of the raw token
user_id TEXT NOT NULL,
name TEXT NOT NULL DEFAULT '',
created_at REAL NOT NULL
);
"""
# Columns added after the original Phase 0 schema; applied idempotently so an
# existing data/modelbeast.db upgrades in place.
MIGRATIONS = [
("jobs", "asset_ids", "TEXT NOT NULL DEFAULT '[]'"),
("jobs", "user_id", "TEXT"), # nullable; legacy rows = owner-era
("assets", "user_id", "TEXT"),
("users", "session_epoch", "INTEGER NOT NULL DEFAULT 0"), # bump to revoke sessions
]
def connect() -> sqlite3.Connection:
DATA.mkdir(parents=True, exist_ok=True)
# FastAPI sync endpoints run in a threadpool; python sqlite3 is built in
# serialized threading mode, so sharing one connection across threads is safe.
con = sqlite3.connect(DB_PATH, check_same_thread=False)
con.row_factory = sqlite3.Row
con.execute("PRAGMA journal_mode=WAL")
con.executescript(SCHEMA)
_migrate(con)
return con
def _migrate(con: sqlite3.Connection) -> None:
for table, column, decl in MIGRATIONS:
cols = {r["name"] for r in con.execute(f"PRAGMA table_info({table})")}
if column not in cols:
con.execute(f"ALTER TABLE {table} ADD COLUMN {column} {decl}")
con.commit()
def new_id() -> str:
return uuid.uuid4().hex[:12]
def now() -> float:
return time.time()
def row_to_dict(row: sqlite3.Row) -> dict:
d = dict(row)
for key in ("meta", "params", "asset_ids"):
if key in d and isinstance(d[key], str):
try:
d[key] = json.loads(d[key])
except ValueError:
pass
return d