note-zero/app/auth.py

from __future__ import annotations

import hashlib
import hmac
import secrets
import time
from collections import defaultdict, deque
from dataclasses import dataclass
from datetime import UTC, datetime, timedelta

from argon2 import PasswordHasher
from argon2.exceptions import InvalidHashError, VerifyMismatchError
from fastapi import HTTPException, Request, Response, status

from .config import Settings
from .db import Database, utc_now

COOKIE_NAME = "xuxiang_session"
SESSION_DAYS = 30
_password_hasher = PasswordHasher()


@dataclass(frozen=True)
class AuthUser:
    id: str
    label: str
    role: str
    debug_sharing: bool

    @property
    def is_admin(self) -> bool:
        return self.role == "admin"

    def public(self) -> dict[str, object]:
        return {
            "id": self.id,
            "label": self.label,
            "role": self.role,
            "debug_sharing": self.debug_sharing,
        }


class LoginLimiter:
    def __init__(self) -> None:
        self.attempts: dict[str, deque[float]] = defaultdict(deque)

    def check(self, key: str) -> None:
        now = time.monotonic()
        attempts = self.attempts[key]
        while attempts and now - attempts[0] > 900:
            attempts.popleft()
        if len(attempts) >= 8:
            raise HTTPException(
                status_code=status.HTTP_429_TOO_MANY_REQUESTS,
                detail="请稍后再试",
            )

    def fail(self, key: str) -> None:
        self.attempts[key].append(time.monotonic())

    def clear(self, key: str) -> None:
        self.attempts.pop(key, None)


limiter = LoginLimiter()


def find_user_for_key(access_key: str, db: Database) -> AuthUser | None:
    for candidate in db.list_auth_users():
        try:
            matches = _password_hasher.verify(
                candidate["access_key_hash"], access_key
            )
        except (VerifyMismatchError, InvalidHashError):
            matches = False
        if matches:
            return AuthUser(
                id=candidate["id"],
                label=candidate["label"],
                role=candidate["role"],
                debug_sharing=bool(candidate["debug_sharing"]),
            )
    return None


def _token_hash(token: str, settings: Settings) -> str:
    secret = (settings.session_secret or "").encode()
    return hmac.new(secret, token.encode(), hashlib.sha256).hexdigest()


def create_session(
    response: Response,
    db: Database,
    settings: Settings,
    user_id: str,
) -> None:
    token = secrets.token_urlsafe(40)
    expires = datetime.now(UTC) + timedelta(days=SESSION_DAYS)
    with db.connect() as connection:
        connection.execute(
            "DELETE FROM sessions WHERE expires_at < ?", (utc_now(),)
        )
        connection.execute(
            """
            INSERT INTO sessions (
                token_hash, user_id, created_at, expires_at
            ) VALUES (?, ?, ?, ?)
            """,
            (
                _token_hash(token, settings),
                user_id,
                utc_now(),
                expires.isoformat(),
            ),
        )
    response.set_cookie(
        COOKIE_NAME,
        token,
        max_age=SESSION_DAYS * 24 * 60 * 60,
        httponly=True,
        secure=settings.cookie_secure,
        samesite="strict",
        path="/",
    )


def delete_session(
    request: Request, response: Response, db: Database, settings: Settings
) -> None:
    token = request.cookies.get(COOKIE_NAME)
    if token:
        with db.connect() as connection:
            connection.execute(
                "DELETE FROM sessions WHERE token_hash = ?",
                (_token_hash(token, settings),),
            )
    response.delete_cookie(
        COOKIE_NAME,
        path="/",
        secure=settings.cookie_secure,
        httponly=True,
        samesite="strict",
    )


def require_auth(
    request: Request, db: Database, settings: Settings
) -> AuthUser:
    if settings.auth_disabled:
        users = db.list_auth_users()
        if users:
            user = users[0]
            return AuthUser(
                id=user["id"],
                label=user["label"],
                role="admin",
                debug_sharing=True,
            )
        return AuthUser(
            id="development",
            label="本地开发",
            role="admin",
            debug_sharing=True,
        )
    if not settings.users or not settings.session_secret:
        raise HTTPException(
            status_code=status.HTTP_503_SERVICE_UNAVAILABLE,
            detail="应用尚未配置访问身份",
        )
    token = request.cookies.get(COOKIE_NAME)
    if not token:
        raise HTTPException(status_code=status.HTTP_401_UNAUTHORIZED)
    with db.connect() as connection:
        row = connection.execute(
            """
            SELECT u.id, u.label, u.role, u.debug_sharing, s.expires_at
            FROM sessions s
            JOIN users u ON u.id = s.user_id
            WHERE s.token_hash = ?
            """,
            (_token_hash(token, settings),),
        ).fetchone()
    if not row or row["expires_at"] <= utc_now():
        raise HTTPException(status_code=status.HTTP_401_UNAUTHORIZED)
    return AuthUser(
        id=row["id"],
        label=row["label"],
        role=row["role"],
        debug_sharing=bool(row["debug_sharing"]),
    )


def require_admin(user: AuthUser) -> None:
    if not user.is_admin:
        raise HTTPException(status_code=status.HTTP_403_FORBIDDEN)


def require_same_origin_intent(request: Request) -> None:
    if request.method in {"GET", "HEAD", "OPTIONS"}:
        return
    if request.url.path == "/api/login":
        return
    if request.headers.get("x-note-client") != "xuxiang-web":
        raise HTTPException(status_code=status.HTTP_403_FORBIDDEN)