modelbeast/server/main.py

import asyncio
import json
import secrets
from pathlib import Path

from fastapi import (Depends, FastAPI, File, HTTPException, Request, Response,
                     UploadFile, WebSocket, WebSocketDisconnect)
from fastapi.responses import FileResponse, JSONResponse
from fastapi.staticfiles import StaticFiles

from . import auth, db, registry, settings as settings_mod, store, sysinfo
from .runner import runner

app = FastAPI(title="MODELBEAST")
WEB_DIST = db.ROOT / "web" / "dist"
INBOX = db.DATA / "inbox"
MAX_UPLOAD = 1024 * 1024 * 1024  # 1 GB


@app.on_event("startup")
async def startup():
    runner.operators = registry.load_operators()
    con = db.connect()
    app.state.con = con
    runner.get_settings = lambda: settings_mod.get_all(con)
    auth.ensure_auth_secret(con)
    _bootstrap_owner(con)
    asyncio.create_task(runner.worker())
    asyncio.create_task(watch_inbox())


def _bootstrap_owner(con):
    if con.execute("SELECT 1 FROM users LIMIT 1").fetchone():
        return
    pw = secrets.token_urlsafe(12)
    auth.create_user(con, "monster", pw, role="owner", max_active_jobs=999)
    banner = "=" * 62
    print(f"\n{banner}\n MODELBEAST owner account created (first run):\n"
          f"   username: monster\n   password: {pw}\n"
          f"   change it:  MB_TOKEN=... mb user passwd monster   (or the UI)\n{banner}\n",
          flush=True)


def _con(request: Request):
    return request.app.state.con


def _with_usernames(con, rows: list[dict]) -> list[dict]:
    """Attach a 'username' field to job/asset dicts for display."""
    umap = {u["id"]: u["username"] for u in auth.list_users(con)}
    for r in rows:
        r["username"] = umap.get(r.get("user_id")) or ("—" if r.get("user_id") else "owner")
    return rows


# -- health (public, unauthenticated — for serve.sh + proxy monitoring) --------
@app.get("/api/health")
def health():
    return {"ok": True, "operators": len(runner.operators)}


# -- auth / session ------------------------------------------------------------
@app.post("/api/login")
async def login(payload: dict, request: Request, response: Response):
    # Key the limiter on the DIRECT peer (Caddy) — not the client-spoofable
    # X-Forwarded-For — plus a per-username bucket so the owner account can't be
    # brute-forced regardless of a rotated XFF header.
    ip = request.client.host if request.client else "?"
    username = (payload.get("username") or "").strip().lower()
    if not auth.rate_ok(f"ip:{ip}") or not auth.rate_ok(f"user:{username}"):
        raise HTTPException(429, "too many attempts — wait a minute")
    con = _con(request)
    password = payload.get("password") or ""
    u = auth.check_login(con, username, password)  # constant-time, no user enumeration
    if not u:
        auth.note_fail(f"ip:{ip}"); auth.note_fail(f"user:{username}")
        raise HTTPException(401, "invalid username or password")
    token = auth.make_session(con, u["id"])
    secure = request.headers.get("x-forwarded-proto", "http") == "https"
    response.set_cookie(auth.SESSION_COOKIE, token, max_age=auth.SESSION_MAX_AGE,
                        httponly=True, samesite="lax", secure=secure, path="/")
    return auth.public_user(u)


@app.post("/api/logout")
async def logout(response: Response):
    response.delete_cookie(auth.SESSION_COOKIE, path="/")
    return {"ok": True}


@app.get("/api/me")
def me(user=Depends(auth.current_user)):
    return auth.public_user(user)


# -- users (owner only) --------------------------------------------------------
@app.get("/api/users")
def get_users(request: Request, owner=Depends(auth.require_owner)):
    return auth.list_users(_con(request))


@app.post("/api/users")
def add_user(payload: dict, request: Request, owner=Depends(auth.require_owner)):
    try:
        maxj = int(payload.get("max_active_jobs", 4))
        u = auth.create_user(_con(request), payload.get("username", ""),
                             payload.get("password", ""), role=payload.get("role", "guest"),
                             max_active_jobs=maxj)
    except (ValueError, TypeError) as e:
        raise HTTPException(400, str(e))
    return auth.public_user(u)


@app.patch("/api/users/{username}")
def patch_user(username: str, payload: dict, request: Request, owner=Depends(auth.require_owner)):
    con = _con(request)
    u = auth.get_user_by_name(con, username)
    if not u:
        raise HTTPException(404, "no such user")
    if payload.get("password"):
        try:
            auth.set_password(con, u["id"], payload["password"])
        except ValueError as e:
            raise HTTPException(400, str(e))
    if "max_active_jobs" in payload:
        try:
            maxj = int(payload["max_active_jobs"])
        except (ValueError, TypeError):
            raise HTTPException(400, "max_active_jobs must be an integer")
        con.execute("UPDATE users SET max_active_jobs = ? WHERE id = ?", (maxj, u["id"]))
        con.commit()
    return auth.public_user(auth.get_user(con, u["id"]))


@app.delete("/api/users/{username}")
def del_user(username: str, request: Request, owner=Depends(auth.require_owner)):
    con = _con(request)
    u = auth.get_user_by_name(con, username)
    if not u:
        raise HTTPException(404, "no such user")
    if u["role"] == "owner":
        raise HTTPException(400, "cannot delete an owner")
    auth.delete_user(con, u["id"])
    return {"ok": True}


@app.post("/api/tokens")
def make_token(payload: dict, request: Request, owner=Depends(auth.require_owner)):
    # owner mints a token for a user (default: self)
    con = _con(request)
    target = payload.get("username")
    u = auth.get_user_by_name(con, target) if target else owner
    if not u:
        raise HTTPException(404, "no such user")
    raw = auth.create_token(con, u["id"], payload.get("name", ""))
    return {"token": raw, "username": u["username"]}  # shown ONCE


# -- operators -----------------------------------------------------------------
@app.get("/api/operators")
def list_operators(user=Depends(auth.current_user)):
    is_owner = user["role"] == "owner"
    out = []
    for op in runner.operators.values():
        # guests are LOCAL-ONLY: cloud operators are filtered out entirely
        if not is_owner and op.get("requires_env"):
            continue
        out.append({k: v for k, v in op.items() if k != "dir"})
    return out


# -- settings (owner only) -----------------------------------------------------
@app.get("/api/settings")
def get_settings(request: Request, owner=Depends(auth.require_owner)):
    s = settings_mod.get_all(_con(request))
    view = settings_mod.public_view(s)
    view["_secret_keys"] = sorted(settings_mod.SECRET_KEYS)
    view["_env_keys"] = sorted(settings_mod.ENV_MAP.keys())
    view["_env_set"] = sorted(env for key, env in settings_mod.ENV_MAP.items() if s.get(key))
    return view


@app.put("/api/settings")
def put_settings(payload: dict, request: Request, owner=Depends(auth.require_owner)):
    clean = {k: v for k, v in payload.items()
             if not (k in settings_mod.SECRET_KEYS and v and set(str(v)) <= {"•"})}
    settings_mod.set_many(_con(request), clean)
    return get_settings(request, owner)


# -- assets --------------------------------------------------------------------
def _can_see(user, obj) -> bool:
    """owner sees everything; a guest sees only their own jobs/assets."""
    return user["role"] == "owner" or obj.get("user_id") == user["id"]


def _guest_may_run(user, op) -> bool:
    """HARD RULE: guests are local-only. Blocked by TWO signals (defense in depth):
    any declared requires_env (paid keys) AND the net/cloud lane. So a paid operator
    that forgot to declare requires_env is still blocked."""
    if user["role"] == "owner":
        return True
    return not op.get("requires_env") and op.get("resources") != "net"


@app.get("/api/assets")
def list_assets(parent_job: str = None, user=Depends(auth.current_user),
                request: Request = None):
    con = _con(request)
    rows = [a for a in store.list_assets(con) if _can_see(user, a)]
    # ?parent_job= used to be accepted and silently ignored, so callers filtering on it
    # received the whole table and typically took rows[0] — correct only by accident of
    # newest-first ordering, and wrong the moment two jobs run at once.
    if parent_job:
        def _pj(a):
            try:
                return a["parent_job"]
            except Exception:
                return None
        rows = [a for a in rows if str(_pj(a)) == str(parent_job)]
    return _with_usernames(con, rows)


@app.post("/api/assets")
async def upload_asset(request: Request, file: UploadFile = File(...),
                       user=Depends(auth.current_user)):
    # stream in bounded chunks so a missing/spoofed Content-Length can't blow memory
    chunks, total = [], 0
    while True:
        chunk = await file.read(1024 * 1024)
        if not chunk:
            break
        total += len(chunk)
        if total > MAX_UPLOAD:
            raise HTTPException(413, "file too large (1 GB max)")
        chunks.append(chunk)
    data = b"".join(chunks)
    if not data:
        raise HTTPException(400, "empty file")
    return store.register_upload(_con(request), file.filename or "unnamed", data, user_id=user["id"])


@app.delete("/api/assets/{asset_id}")
def delete_asset(asset_id: str, user=Depends(auth.current_user), request: Request = None):
    con = _con(request)
    asset = store.get_asset(con, asset_id)
    if not asset:
        raise HTTPException(404, "no such asset")
    if not _can_see(user, asset):
        raise HTTPException(403, "not your asset")
    store.delete_asset(con, asset_id)
    return {"ok": True}


@app.get("/api/assets/{asset_id}/file")
def asset_file(asset_id: str, member: str | None = None,
               user=Depends(auth.current_user), request: Request = None):
    asset = store.get_asset(_con(request), asset_id)
    if not asset:
        raise HTTPException(404, "no such asset")
    if not _can_see(user, asset):
        raise HTTPException(403, "not your asset")
    path = Path(asset["path"])
    if member:
        target = (path / member).resolve()
        if not target.is_relative_to(path.resolve()) or not target.is_file():
            raise HTTPException(404, "no such member")
        return FileResponse(target)
    if path.is_dir():
        files = sorted(p.name for p in path.iterdir() if p.is_file())
        return JSONResponse({"folder": True, "files": files})
    return FileResponse(path, filename=asset["name"])


# -- jobs ----------------------------------------------------------------------
@app.get("/api/jobs")
def list_jobs(user=Depends(auth.current_user), request: Request = None):
    con = _con(request)
    rows = [j for j in runner.list_jobs(con) if _can_see(user, j)]
    return _with_usernames(con, rows)


@app.get("/api/jobs/recent")
def recent_jobs(user=Depends(auth.current_user), request: Request = None):
    con = _con(request)
    uid = None if user["role"] == "owner" else user["id"]
    return sysinfo.recent_jobs(con, runner, user_id=uid)


@app.post("/api/jobs")
async def create_job(payload: dict, user=Depends(auth.current_user), request: Request = None):
    con = _con(request)
    operator = payload.get("operator")
    if operator not in runner.operators:
        raise HTTPException(400, f"unknown operator: {operator}")
    op = runner.operators[operator]
    if not _guest_may_run(user, op):
        raise HTTPException(403, "guests run local models only")
    if user["role"] != "owner" and runner.active_job_count(con, user["id"]) >= user["max_active_jobs"]:
        raise HTTPException(429, "too many active jobs — wait for some to finish")
    asset_ids = payload.get("asset_ids")
    if not asset_ids:
        asset_ids = [payload["asset_id"]] if payload.get("asset_id") else []
    for aid in asset_ids:
        a = store.get_asset(con, aid)
        if not a:
            raise HTTPException(400, f"unknown asset: {aid}")
        if not _can_see(user, a):
            raise HTTPException(403, "not your asset")
    params = payload.get("params") or {}
    job = runner.create_job(con, operator, asset_ids, params, user_id=user["id"])
    await runner.broadcast({"type": "job", "job": job})
    return job


@app.get("/api/jobs/{job_id}")
def get_job(job_id: str, user=Depends(auth.current_user), request: Request = None):
    job = runner.get_job(_con(request), job_id)
    if not job:
        raise HTTPException(404, "no such job")
    return job


def _own_job_or_owner(con, job_id, user):
    job = runner.get_job(con, job_id)
    if not job:
        raise HTTPException(404, "no such job")
    if user["role"] != "owner" and job.get("user_id") != user["id"]:
        raise HTTPException(403, "not your job")
    return job


@app.post("/api/jobs/{job_id}/cancel")
async def cancel_job(job_id: str, user=Depends(auth.current_user), request: Request = None):
    con = _con(request)
    _own_job_or_owner(con, job_id, user)
    if not await runner.cancel(con, job_id):
        raise HTTPException(400, "job not cancellable")
    return {"ok": True}


@app.post("/api/jobs/{job_id}/retry")
async def retry_job(job_id: str, user=Depends(auth.current_user), request: Request = None):
    con = _con(request)
    job = _own_job_or_owner(con, job_id, user)
    op = runner.operators.get(job["operator"], {})
    if not _guest_may_run(user, op):
        raise HTTPException(403, "guests run local models only")
    if user["role"] != "owner" and runner.active_job_count(con, user["id"]) >= user["max_active_jobs"]:
        raise HTTPException(429, "too many active jobs — wait for some to finish")
    new = runner.create_job(con, job["operator"], job.get("asset_ids") or [],
                            job["params"], user_id=user["id"])
    await runner.broadcast({"type": "job", "job": new})
    return new


@app.delete("/api/jobs/{job_id}")
def delete_job(job_id: str, user=Depends(auth.current_user), request: Request = None):
    con = _con(request)
    _own_job_or_owner(con, job_id, user)
    if not runner.delete_job(con, job_id):
        raise HTTPException(400, "job running or missing")
    return {"ok": True}


# -- system / dashboard --------------------------------------------------------
@app.get("/api/system")
def system(user=Depends(auth.current_user), request: Request = None):
    return sysinfo.snapshot(_con(request), runner)


@app.get("/api/nodes")
def nodes(user=Depends(auth.current_user), request: Request = None):
    """Per-node capability + live work. /api/system only reports name/remote/busy,
    and its busy flag is never updated — this is the real per-machine view."""
    return sysinfo.nodes_detail(_con(request), runner)


# -- inbox watch folder --------------------------------------------------------
async def watch_inbox():
    INBOX.mkdir(parents=True, exist_ok=True)
    seen: dict[str, int] = {}
    while True:
        await asyncio.sleep(3)
        try:
            for p in INBOX.iterdir():
                if p.name.startswith(".") or not p.exists():
                    continue
                size = p.stat().st_size if p.is_file() else sum(
                    f.stat().st_size for f in p.rglob("*") if f.is_file())
                if seen.get(str(p)) == size and size > 0:
                    asset = store.register_file(app.state.con, p, move=True,
                                                meta={"source": "inbox"})
                    seen.pop(str(p), None)
                    await runner.broadcast({"type": "assets_changed", "assets": [asset["id"]]})
                else:
                    seen[str(p)] = size
        except Exception as e:
            print(f"[inbox] {e}")


# -- websocket (auth: session cookie only — never a token in the query string,
#    which would leak into proxy/access logs) ----------------------------------
@app.websocket("/ws")
async def ws(websocket: WebSocket):
    con = websocket.app.state.con
    user = None
    tok = websocket.cookies.get(auth.SESSION_COOKIE)
    if tok:
        uid = auth.read_session(con, tok)
        if uid:
            user = auth.get_user(con, uid)
    if not user:
        await websocket.close(code=1008)  # policy violation
        return
    await websocket.accept()
    runner.subscribers[websocket] = user
    try:
        while True:
            await websocket.receive_text()
    except WebSocketDisconnect:
        runner.subscribers.pop(websocket, None)


# -- frontend (built) — public so the SPA can load and show the login screen ----
if WEB_DIST.exists():
    app.mount("/", StaticFiles(directory=WEB_DIST, html=True), name="web")