modelbeast/server/auth.py

"""Authentication: bcrypt passwords, signed-cookie sessions, bearer API tokens,
and the FastAPI dependency that resolves the current user.

Design premise (owner's hard rule): guests are LOCAL-ONLY. Any operator with a
non-empty requires_env is owner-only and enforced server-side (see main.py) — this
module just answers "who is this request" and "is the password right".
"""
import hashlib
import re
import secrets
import time

import bcrypt
from fastapi import HTTPException, Request
from itsdangerous import BadSignature, SignatureExpired, URLSafeTimedSerializer

from . import db

SESSION_COOKIE = "mb_session"
SESSION_MAX_AGE = 30 * 24 * 3600  # 30 days
USERNAME_RE = re.compile(r"^[a-z0-9_-]{2,24}$")
MIN_PW_LEN = 6

# in-memory login rate limit: ip -> [failure epoch, ...]
_login_fails: dict[str, list[float]] = {}
_RATE_WINDOW = 60.0
_RATE_MAX = 5


# -- passwords ----------------------------------------------------------------
def hash_pw(pw: str) -> str:
    # bcrypt caps at 72 bytes; truncate deterministically (standard practice)
    return bcrypt.hashpw(pw.encode()[:72], bcrypt.gensalt()).decode()


def verify_pw(pw: str, pw_hash: str) -> bool:
    try:
        return bcrypt.checkpw(pw.encode()[:72], pw_hash.encode())
    except (ValueError, TypeError):
        return False


# a real hash to compare against when the user doesn't exist, so login timing is
# the same whether or not the username is valid (defeats username enumeration)
_DUMMY_HASH = bcrypt.hashpw(b"modelbeast-dummy", bcrypt.gensalt()).decode()


def check_login(con, username: str, password: str) -> dict | None:
    """Constant-ish-time credential check. Returns the user row or None."""
    u = get_user_by_name(con, username)
    if u and verify_pw(password, u["pw_hash"]):
        return u
    if not u:
        verify_pw(password, _DUMMY_HASH)  # burn equivalent time
    return None


# -- auth secret (stored hidden in settings; never exposed via the API) -------
def ensure_auth_secret(con) -> str:
    row = con.execute("SELECT value FROM settings WHERE key='auth_secret'").fetchone()
    if row and row["value"]:
        return row["value"]
    secret = secrets.token_hex(32)
    con.execute("INSERT OR REPLACE INTO settings (key, value) VALUES ('auth_secret', ?)", (secret,))
    con.commit()
    return secret


def _serializer(con) -> URLSafeTimedSerializer:
    return URLSafeTimedSerializer(ensure_auth_secret(con), salt="mb-session")


def make_session(con, user_id: str) -> str:
    u = get_user(con, user_id)
    epoch = u.get("session_epoch", 0) if u else 0
    return _serializer(con).dumps({"uid": user_id, "e": epoch})


def read_session(con, token: str) -> str | None:
    """Return the user id iff the signature is valid, unexpired, and the embedded
    session epoch still matches (password change / logout-all bumps the epoch)."""
    try:
        payload = _serializer(con).loads(token, max_age=SESSION_MAX_AGE)
    except (BadSignature, SignatureExpired, Exception):
        return None
    uid = payload.get("uid")
    u = get_user(con, uid) if uid else None
    if not u or payload.get("e", 0) != u.get("session_epoch", 0):
        return None
    return uid


# -- API tokens ----------------------------------------------------------------
def _token_hash(raw: str) -> str:
    return hashlib.sha256(raw.encode()).hexdigest()


def create_token(con, user_id: str, name: str = "") -> str:
    raw = "mbt_" + secrets.token_hex(20)
    con.execute(
        "INSERT INTO api_tokens (token_hash, user_id, name, created_at) VALUES (?, ?, ?, ?)",
        (_token_hash(raw), user_id, name, db.now()))
    con.commit()
    return raw


def user_by_token(con, raw: str) -> dict | None:
    row = con.execute(
        "SELECT u.* FROM api_tokens t JOIN users u ON u.id = t.user_id WHERE t.token_hash = ?",
        (_token_hash(raw),)).fetchone()
    return dict(row) if row else None


# -- users ---------------------------------------------------------------------
def get_user(con, user_id: str) -> dict | None:
    row = con.execute("SELECT * FROM users WHERE id = ?", (user_id,)).fetchone()
    return dict(row) if row else None


def get_user_by_name(con, username: str) -> dict | None:
    row = con.execute("SELECT * FROM users WHERE username = ?", (username.lower(),)).fetchone()
    return dict(row) if row else None


def create_user(con, username: str, password: str, role: str = "guest",
                max_active_jobs: int = 4) -> dict:
    username = username.lower()
    if not USERNAME_RE.match(username):
        raise ValueError("username must be 2-24 chars of a-z 0-9 _ -")
    if len(password) < MIN_PW_LEN:
        raise ValueError(f"password must be >= {MIN_PW_LEN} characters")
    if get_user_by_name(con, username):
        raise ValueError("username already exists")
    if role not in ("owner", "guest"):
        raise ValueError("role must be owner or guest")
    uid = db.new_id()
    con.execute(
        "INSERT INTO users (id, username, pw_hash, role, max_active_jobs, created_at) "
        "VALUES (?, ?, ?, ?, ?, ?)",
        (uid, username, hash_pw(password), role, max_active_jobs, db.now()))
    con.commit()
    return get_user(con, uid)


def set_password(con, user_id: str, password: str) -> None:
    if len(password) < MIN_PW_LEN:
        raise ValueError(f"password must be >= {MIN_PW_LEN} characters")
    # bump session_epoch so every existing session/cookie for this user is revoked
    con.execute("UPDATE users SET pw_hash = ?, session_epoch = session_epoch + 1 WHERE id = ?",
                (hash_pw(password), user_id))
    con.commit()


def list_users(con) -> list[dict]:
    rows = con.execute(
        "SELECT id, username, role, max_active_jobs, created_at FROM users ORDER BY created_at"
    ).fetchall()
    return [dict(r) for r in rows]


def delete_user(con, user_id: str) -> None:
    con.execute("DELETE FROM api_tokens WHERE user_id = ?", (user_id,))
    con.execute("DELETE FROM users WHERE id = ?", (user_id,))
    con.commit()


def public_user(u: dict) -> dict:
    return {"username": u["username"], "role": u["role"],
            "is_owner": u["role"] == "owner", "max_active_jobs": u["max_active_jobs"]}


# -- rate limit ----------------------------------------------------------------
def rate_ok(ip: str) -> bool:
    now = time.time()
    fails = [t for t in _login_fails.get(ip, []) if now - t < _RATE_WINDOW]
    _login_fails[ip] = fails
    return len(fails) < _RATE_MAX


def note_fail(ip: str) -> None:
    _login_fails.setdefault(ip, []).append(time.time())


# -- FastAPI dependency --------------------------------------------------------
def resolve_user(request: Request) -> dict | None:
    """Cookie session first, then bearer token. Returns user dict or None."""
    con = request.app.state.con
    tok = request.cookies.get(SESSION_COOKIE)
    if tok:
        uid = read_session(con, tok)
        if uid:
            u = get_user(con, uid)
            if u:
                return u
    authz = request.headers.get("authorization", "")
    if authz.startswith("Bearer "):
        u = user_by_token(con, authz[7:].strip())
        if u:
            return u
    return None


def current_user(request: Request) -> dict:
    u = resolve_user(request)
    if not u:
        raise HTTPException(401, "authentication required")
    return u


def require_owner(request: Request) -> dict:
    u = current_user(request)
    if u["role"] != "owner":
        raise HTTPException(403, "owner only")
    return u