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 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=1
    ).hex()
    return f"scrypt${salt}${digest}"


def verify_password(password: str, encoded: str) -> bool:
    try:
        _, salt, expected = encoded.split("$")
        actual = hashlib.scrypt(
            password.encode(), salt=bytes.fromhex(salt), n=16384, r=8, p=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 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():
        for user in json.loads(path.read_text()).get("users", []):
            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", "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):
    @app.post("/api/v1/auth/login", response_model=AccountView, tags=["Accounts"])
    def login(payload: Login, request: Request, response: Response):
        identity = hashlib.sha256(
            (request.client.host if request.client else "unknown").encode()
        ).hexdigest()
        attempt = db.one("SELECT * FROM login_attempts WHERE identity=?", (identity,))
        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$" + "00" * 16 + "$" + "00" * 64
        valid = verify_password(
            payload.password, row["password_hash"] if row else fallback
        )
        if not row or not valid:
            failures = (
                attempt["failures"] if attempt and attempt["blocked_until"] == 0 else 0
            ) + 1
            db.execute(
                "INSERT OR REPLACE INTO login_attempts VALUES(?,?,?)",
                (identity, failures, time.time() + 300 if failures >= 5 else 0),
            )
            raise FlowError(E.UNAUTHORIZED, "Invalid username or password", 401)
        db.execute("DELETE FROM login_attempts WHERE identity=?", (identity,))
        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)
        db.execute(
            "UPDATE users SET password_hash=? WHERE username=?",
            (password_hash(payload.new_password), session["username"]),
        )
        db.execute("DELETE FROM sessions WHERE username=?", (session["username"],))
        response.delete_cookie(COOKIE, path="/")