pigway-frame-flow/pigway_frame_flow/core/accounts.py

"""Browser accounts: salted password hashes and revocable server-side sessions."""

import hashlib
import hmac
import json
import secrets
import threading
import time
from pathlib import Path

from fastapi import Request, Response, Security
from fastapi.security import APIKeyCookie
from pydantic import BaseModel, Field

from pigway_frame_flow.errors import ErrorCode as E
from pigway_frame_flow.errors import FlowError

COOKIE = "frameflow_session"
cookie = APIKeyCookie(name=COOKIE, auto_error=False, scheme_name="AccountSession")


def password_hash(password: str) -> str:
    salt = secrets.token_hex(16)
    digest = hashlib.scrypt(
        password.encode(), salt=bytes.fromhex(salt), n=16384, r=8, p=5
    ).hex()
    return f"scrypt-v2${salt}${digest}"


def verify_password(password: str, encoded: str) -> bool:
    try:
        algorithm, salt, expected = encoded.split("$")
        if algorithm not in {"scrypt", "scrypt-v2"}:
            return False
        actual = hashlib.scrypt(
            password.encode(),
            salt=bytes.fromhex(salt),
            n=16384,
            r=8,
            p=5 if algorithm == "scrypt-v2" else 1,
        ).hex()
        return hmac.compare_digest(actual, expected)
    except (ValueError, TypeError):
        return False


class Login(BaseModel):
    username: str = Field(min_length=1, max_length=80)
    password: str = Field(min_length=1, max_length=256, repr=False)
    remember: bool = True


class AccountView(BaseModel):
    username: str
    csrf_token: str


class PasswordChange(BaseModel):
    current_password: str = Field(min_length=1, max_length=256, repr=False)
    new_password: str = Field(min_length=8, max_length=256, repr=False)


def seed_accounts(db, config_dir: Path):
    db.execute(
        "CREATE TABLE IF NOT EXISTS users (username TEXT PRIMARY KEY, password_hash TEXT NOT NULL, enabled INTEGER NOT NULL DEFAULT 1)"
    )
    db.execute("CREATE UNIQUE INDEX IF NOT EXISTS single_browser_account ON users((1))")
    db.execute(
        "CREATE TABLE IF NOT EXISTS sessions (token_hash TEXT PRIMARY KEY, username TEXT NOT NULL, csrf TEXT NOT NULL, expires_at REAL NOT NULL)"
    )
    db.execute(
        "CREATE TABLE IF NOT EXISTS login_attempts (identity TEXT PRIMARY KEY, failures INTEGER NOT NULL, blocked_until REAL NOT NULL)"
    )
    path = config_dir / "accounts.json"
    if path.exists():
        users = json.loads(path.read_text()).get("users", [])
        if len(users) > 1:
            raise ValueError("Only one browser account is supported")
        for user in users:
            existing = db.one("SELECT username FROM users LIMIT 1")
            if existing and existing["username"] != user["username"]:
                raise ValueError("A browser account already exists")
            db.execute(
                "INSERT OR IGNORE INTO users VALUES(?,?,1)",
                (user["username"], user["password_hash"]),
            )


def session_user(request: Request):
    raw = request.cookies.get(COOKIE, "")
    if not raw:
        return None
    db = request.app.state.service.db
    row = db.one(
        "SELECT s.*,u.enabled FROM sessions s JOIN users u USING(username) WHERE token_hash=?",
        (hashlib.sha256(raw.encode()).hexdigest(),),
    )
    if not row or not row["enabled"] or row["expires_at"] <= time.time():
        return None
    if request.method in {"POST", "PUT", "PATCH", "DELETE"}:
        if not hmac.compare_digest(
            request.headers.get("X-CSRF-Token", ""), row["csrf"]
        ):
            raise FlowError(E.FORBIDDEN, "Session CSRF validation failed", 403)
    return row


def install_accounts(app, db):
    # Serialize authentication work without holding the database lock during scrypt.
    login_lock = threading.Lock()

    @app.post("/api/v1/auth/login", response_model=AccountView, tags=["Accounts"])
    def login(payload: Login, request: Request, response: Response):
        with login_lock:
            # Reject login CSRF before checking credentials. Never trust forwarded hosts.
            origin = request.headers.get("origin")
            if request.headers.get("sec-fetch-site") == "cross-site" or (
                origin and origin != str(request.base_url).rstrip("/")
            ):
                raise FlowError(E.FORBIDDEN, "Cross-origin login is not allowed", 403)
            identity = hashlib.sha256(
                (request.client.host if request.client else "unknown").encode()
            ).hexdigest()
            # Account throttling also applies when an attacker changes source IPs.
            identities = (
                identity,
                "account:" + hashlib.sha256(payload.username.encode()).hexdigest(),
            )
            attempts = [
                db.one("SELECT * FROM login_attempts WHERE identity=?", (key,))
                for key in identities
            ]
            attempt = next(
                (a for a in attempts if a and a["blocked_until"] > time.time()), None
            )
            if attempt and attempt["blocked_until"] > time.time():
                raise FlowError(
                    E.RATE_LIMITED, "Too many login attempts; try again later", 429
                )
            row = db.one(
                "SELECT * FROM users WHERE username=? AND enabled=1",
                (payload.username,),
            )
            # Perform a password calculation even for unknown usernames.
            fallback = "scrypt-v2$" + "00" * 16 + "$" + "00" * 64
            valid = verify_password(
                payload.password, row["password_hash"] if row else fallback
            )
            if not row or not valid:
                for key, previous in zip(identities, attempts, strict=True):
                    failures = (
                        previous["failures"]
                        if previous and previous["blocked_until"] == 0
                        else 0
                    ) + 1
                    db.execute(
                        "INSERT OR REPLACE INTO login_attempts VALUES(?,?,?)",
                        (key, failures, time.time() + 300 if failures >= 5 else 0),
                    )
                raise FlowError(E.UNAUTHORIZED, "Invalid username or password", 401)
            upgraded = (
                password_hash(payload.password)
                if row["password_hash"].startswith("scrypt$")
                else None
            )
            with db.lock:
                # A concurrent password change must not create a session from old credentials.
                current_row = db.one(
                    "SELECT * FROM users WHERE username=? AND enabled=1",
                    (payload.username,),
                )
                if (
                    not current_row
                    or current_row["password_hash"] != row["password_hash"]
                ):
                    raise FlowError(E.UNAUTHORIZED, "Invalid username or password", 401)
                for key in identities:
                    db.execute("DELETE FROM login_attempts WHERE identity=?", (key,))
                if row["password_hash"].startswith("scrypt$"):
                    db.execute(
                        "UPDATE users SET password_hash=? WHERE username=?",
                        (upgraded, row["username"]),
                    )
                raw, csrf = secrets.token_urlsafe(32), secrets.token_urlsafe(32)
                lifetime = 30 * 86400 if payload.remember else 12 * 3600
                db.execute("DELETE FROM sessions WHERE expires_at<=?", (time.time(),))
                old = request.cookies.get(COOKIE)
                if old:
                    db.execute(
                        "DELETE FROM sessions WHERE token_hash=?",
                        (hashlib.sha256(old.encode()).hexdigest(),),
                    )
                db.execute(
                    "INSERT INTO sessions VALUES(?,?,?,?)",
                    (
                        hashlib.sha256(raw.encode()).hexdigest(),
                        row["username"],
                        csrf,
                        time.time() + lifetime,
                    ),
                )
            response.set_cookie(
                COOKIE,
                raw,
                max_age=lifetime if payload.remember else None,
                httponly=True,
                samesite="strict",
                secure=request.url.scheme == "https",
                path="/",
            )
            return AccountView(username=row["username"], csrf_token=csrf)

    def current(request):
        row = session_user(request)
        if not row:
            raise FlowError(E.UNAUTHORIZED, "Account login required", 401)
        return row

    @app.get(
        "/api/v1/auth/me",
        response_model=AccountView,
        tags=["Accounts"],
        dependencies=[Security(cookie)],
    )
    def me(request: Request):
        row = current(request)
        return AccountView(username=row["username"], csrf_token=row["csrf"])

    @app.post(
        "/api/v1/auth/logout",
        status_code=204,
        tags=["Accounts"],
        dependencies=[Security(cookie)],
    )
    def logout(request: Request, response: Response):
        row = current(request)
        db.execute("DELETE FROM sessions WHERE token_hash=?", (row["token_hash"],))
        response.delete_cookie(COOKIE, path="/")

    @app.put(
        "/api/v1/auth/password",
        status_code=204,
        tags=["Accounts"],
        dependencies=[Security(cookie)],
    )
    def change_password(payload: PasswordChange, request: Request, response: Response):
        session = current(request)
        user = db.one("SELECT * FROM users WHERE username=?", (session["username"],))
        if not verify_password(payload.current_password, user["password_hash"]):
            raise FlowError(E.UNAUTHORIZED, "Current password is incorrect", 401)
        with db.transaction() as connection:
            connection.execute(
                "UPDATE users SET password_hash=? WHERE username=?",
                (password_hash(payload.new_password), session["username"]),
            )
            connection.execute(
                "DELETE FROM sessions WHERE username=?", (session["username"],)
            )
        response.delete_cookie(COOKIE, path="/")