""" App-level username/password auth for the burn-in dashboard. Sessions are signed cookies (Starlette SessionMiddleware) that carry {user_id, username}. Every request goes through `get_current_user_optional` via the auth middleware in main.py; routes that need an authenticated user import `get_current_user` instead, which raises 401 (or redirects to /login for HTML requests) when there's no session. Passwords are bcrypt with the library's default 12-round cost. We never store plaintext. Bootstrap: if the users table is empty AND `initial_admin_username` / `initial_admin_password` are set, the lifespan creates that admin once at startup. Otherwise, the login template renders the "first user" form when visited and zero users exist. """ from __future__ import annotations import logging import secrets from dataclasses import dataclass from datetime import datetime, timezone from pathlib import Path import aiosqlite import bcrypt from fastapi import HTTPException, Request, status from starlette.responses import RedirectResponse from app.config import settings log = logging.getLogger(__name__) # --------------------------------------------------------------------------- # Session secret — env var > persisted file > generated # --------------------------------------------------------------------------- _SESSION_SECRET_FILE = "session_secret" def get_session_secret() -> str: """Return the HMAC key for SessionMiddleware. env var beats disk.""" if settings.session_secret: return settings.session_secret path = Path(settings.db_path).parent / _SESSION_SECRET_FILE if not path.exists(): path.parent.mkdir(parents=True, exist_ok=True) path.write_bytes(secrets.token_urlsafe(64).encode()) try: path.chmod(0o600) except OSError: pass log.warning( "Generated and persisted session secret to %s. " "Set SESSION_SECRET in env to override.", path, ) return path.read_text().strip() # --------------------------------------------------------------------------- # User model + storage # --------------------------------------------------------------------------- @dataclass(frozen=True) class User: id: int username: str full_name: str | None is_admin: bool def _now() -> str: return datetime.now(timezone.utc).isoformat() def hash_password(plain: str) -> str: return bcrypt.hashpw(plain.encode("utf-8"), bcrypt.gensalt()).decode("utf-8") def verify_password(plain: str, hashed: str) -> bool: try: return bcrypt.checkpw(plain.encode("utf-8"), hashed.encode("utf-8")) except (ValueError, TypeError): return False async def user_count() -> int: async with aiosqlite.connect(settings.db_path) as db: cur = await db.execute("SELECT COUNT(*) FROM users") return (await cur.fetchone())[0] async def get_user_by_username(username: str) -> tuple[User, str] | None: """Returns (user, password_hash) or None. Hash is the only place callers should ever see the raw bcrypt string — for verify_password.""" async with aiosqlite.connect(settings.db_path) as db: db.row_factory = aiosqlite.Row cur = await db.execute( "SELECT id, username, password_hash, full_name, is_admin " "FROM users WHERE username = ? COLLATE NOCASE", (username,), ) row = await cur.fetchone() if not row: return None user = User( id=row["id"], username=row["username"], full_name=row["full_name"], is_admin=bool(row["is_admin"]), ) return user, row["password_hash"] async def get_user_by_id(user_id: int) -> User | None: async with aiosqlite.connect(settings.db_path) as db: db.row_factory = aiosqlite.Row cur = await db.execute( "SELECT id, username, full_name, is_admin " "FROM users WHERE id = ?", (user_id,), ) row = await cur.fetchone() if not row: return None return User( id=row["id"], username=row["username"], full_name=row["full_name"], is_admin=bool(row["is_admin"]), ) async def create_user(username: str, password: str, full_name: str | None = None, is_admin: bool = False, bootstrap_only: bool = False) -> User: """Insert a new user. Raises ValueError if the username collides. bootstrap_only=True: serializes the insert with a check that the users table is empty inside an IMMEDIATE transaction. Used for the /api/v1/auth/setup first-user flow so two concurrent requests can't both create admin accounts during the bootstrap window. """ username = (username or "").strip() if not username: raise ValueError("Username is required.") if len(password) < 8: raise ValueError("Password must be at least 8 characters.") h = hash_password(password) try: async with aiosqlite.connect(settings.db_path) as db: if bootstrap_only: # IMMEDIATE acquires the write lock up-front so a parallel # setup request waits or fails — no two-step race. await db.execute("BEGIN IMMEDIATE") cur = await db.execute("SELECT COUNT(*) FROM users") if (await cur.fetchone())[0] != 0: await db.execute("ROLLBACK") raise ValueError( "Users already exist — first-user setup is closed." ) cur = await db.execute( """INSERT INTO users (username, password_hash, full_name, is_admin, created_at) VALUES (?, ?, ?, ?, ?) RETURNING id""", (username, h, full_name or None, 1 if is_admin else 0, _now()), ) row = await cur.fetchone() await db.commit() except aiosqlite.IntegrityError: raise ValueError(f"Username {username!r} already exists.") return User( id=row[0], username=username, full_name=full_name, is_admin=is_admin, ) async def touch_last_login(user_id: int) -> None: async with aiosqlite.connect(settings.db_path) as db: await db.execute( "UPDATE users SET last_login_at = ? WHERE id = ?", (_now(), user_id), ) await db.commit() async def change_password(user_id: int, current_password: str, new_password: str) -> None: """Verify current password and rotate. Raises ValueError on any failure.""" if len(new_password) < 8: raise ValueError("New password must be at least 8 characters.") async with aiosqlite.connect(settings.db_path) as db: db.row_factory = aiosqlite.Row cur = await db.execute( "SELECT username, password_hash FROM users WHERE id = ?", (user_id,) ) row = await cur.fetchone() if not row or not verify_password(current_password, row["password_hash"]): raise ValueError("Current password is incorrect.") new_hash = hash_password(new_password) await db.execute( "UPDATE users SET password_hash = ? WHERE id = ?", (new_hash, user_id), ) await db.commit() # --------------------------------------------------------------------------- # Login rate limiting (in-memory, per-username + per-source-IP) # --------------------------------------------------------------------------- import time as _time LOGIN_FAILURE_WINDOW_SECONDS = 600 # 10 min LOGIN_FAILURE_THRESHOLD = 10 # this many failures within the window LOGIN_LOCKOUT_SECONDS = 900 # then block for 15 min # {(key,): [(timestamp, ...), ...]} key = (kind, value), kind in {"user","ip"} _login_failures: dict = {} _login_lockouts: dict = {} # key -> unix expiry def _gc_failures(key) -> None: """Drop failure timestamps older than the window.""" arr = _login_failures.get(key, []) cutoff = _time.time() - LOGIN_FAILURE_WINDOW_SECONDS fresh = [t for t in arr if t >= cutoff] if fresh: _login_failures[key] = fresh elif key in _login_failures: del _login_failures[key] def login_locked_until(username: str, ip: str) -> float | None: """Returns the lockout expiry (unix ts) if either dimension is locked, else None. Lazily reaps expired lockouts.""" now = _time.time() soonest = None for key in (("user", username.lower()), ("ip", ip)): exp = _login_lockouts.get(key) if exp is None: continue if now >= exp: del _login_lockouts[key] continue soonest = exp if soonest is None else min(soonest, exp) return soonest def register_login_attempt(username: str, ip: str) -> str: """Atomic check-then-increment for a login attempt. Returns: "ok" — allowed, counter incremented "locked_out" — already locked from a prior attempt "now_locked_out" — THIS attempt is what tripped the lockout The increment runs synchronously (no awaits) so concurrent requests can't slip past the threshold in CPython's single-threaded asyncio loop. Caller must invoke clear_login_failures() on successful auth to roll back this attempt's contribution. """ now = _time.time() # Check existing lockouts first; if already locked, don't even # increment — the lockout window absorbs everything. for key in (("user", username.lower()), ("ip", ip)): exp = _login_lockouts.get(key) if exp is None: continue if now >= exp: del _login_lockouts[key] continue return "locked_out" # Increment + arm lockout if this push crosses the threshold. tripped = False for key in (("user", username.lower()), ("ip", ip)): _gc_failures(key) _login_failures.setdefault(key, []).append(now) if len(_login_failures[key]) >= LOGIN_FAILURE_THRESHOLD: _login_lockouts[key] = now + LOGIN_LOCKOUT_SECONDS _login_failures[key] = [] tripped = True return "now_locked_out" if tripped else "ok" def clear_login_failures(username: str, ip: str) -> None: """Erase counters AND any lockout for a successful auth — caller proved they have credentials, so the brute-force ladder resets.""" for key in (("user", username.lower()), ("ip", ip)): _login_failures.pop(key, None) _login_lockouts.pop(key, None) # --------------------------------------------------------------------------- # Audit events for auth flows # --------------------------------------------------------------------------- async def audit_auth_event(event_type: str, username: str | None, message: str) -> None: """Write a row to audit_events. event_type is one of: user_login / user_login_failed / user_logout / user_password_changed / user_login_locked_out.""" async with aiosqlite.connect(settings.db_path) as db: await db.execute( """INSERT INTO audit_events (event_type, drive_id, burnin_job_id, operator, message) VALUES (?,?,?,?,?)""", (event_type, None, None, username or "?", message), ) await db.commit() async def bootstrap_admin_if_empty() -> None: """Create the env-supplied admin if the users table is empty.""" if await user_count() > 0: return if not (settings.initial_admin_username and settings.initial_admin_password): return try: await create_user( settings.initial_admin_username, settings.initial_admin_password, full_name=None, is_admin=True, ) log.warning( "Bootstrapped initial admin user %r from env. " "Change the password via the UI and remove the env vars from compose.", settings.initial_admin_username, ) except ValueError as exc: log.error("Failed to bootstrap initial admin: %s", exc) # --------------------------------------------------------------------------- # FastAPI dependencies # --------------------------------------------------------------------------- async def get_current_user_optional(request: Request) -> User | None: """Return the logged-in user, or None. Doesn't raise — for templates.""" sess_user_id = request.session.get("user_id") if hasattr(request, "session") else None if not sess_user_id: return None return await get_user_by_id(int(sess_user_id)) def require_admin(request: Request) -> User: """Strict admin gate — for any settings-mutating endpoint. The AuthGate middleware has already populated request.state.current_user; this just enforces is_admin on top.""" user = getattr(request.state, "current_user", None) if not user: raise HTTPException( status_code=status.HTTP_401_UNAUTHORIZED, detail="Authentication required", ) if not user.is_admin: raise HTTPException( status_code=status.HTTP_403_FORBIDDEN, detail="Admin only", ) return user async def get_current_user(request: Request) -> User: """Strict version — for routes. 401 (or redirect for HTML) if missing.""" user = await get_current_user_optional(request) if user is None: # HTML clients prefer a redirect; API clients need a clean 401. accept = request.headers.get("accept", "") if "text/html" in accept and request.method == "GET": raise _RedirectToLogin(request.url.path) raise HTTPException( status_code=status.HTTP_401_UNAUTHORIZED, detail="Authentication required", ) return user class _RedirectToLogin(Exception): """Raised by get_current_user when an HTML page needs to bounce to /login.""" def __init__(self, next_path: str): self.next_path = next_path def login_redirect(next_path: str = "/") -> RedirectResponse: safe_next = next_path if next_path.startswith("/") else "/" target = f"/login?next={safe_next}" if safe_next != "/" else "/login" return RedirectResponse(url=target, status_code=303)