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
from pydantic import BaseModel, Field

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

COOKIE = "frameflow_session"


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"])
    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"])
    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"])
    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="/")