import hashlib import hmac import secrets import threading import time from collections import defaultdict, deque from datetime import datetime, timezone from argon2 import PasswordHasher from argon2.exceptions import InvalidHashError, VerificationError from fastapi import Depends, HTTPException, Request, status from psycopg import Connection from .config import settings from .db import get_connection password_hasher = PasswordHasher() def hash_password(password: str) -> str: return password_hasher.hash(password) def verify_password(password: str, password_hash: str) -> bool: try: return password_hasher.verify(password_hash, password) except (InvalidHashError, VerificationError): return False def token_digest(token: str) -> str: return hashlib.sha256(token.encode("utf-8")).hexdigest() def new_token() -> str: return secrets.token_urlsafe(32) def new_csrf_token() -> str: return secrets.token_urlsafe(32) def code_digest(code: str) -> str: return hmac.new(settings.csrf_secret.encode(), code.encode(), hashlib.sha256).hexdigest() def utc_now() -> datetime: return datetime.now(timezone.utc) def check_origin(request: Request) -> None: origin = request.headers.get("origin") if origin and origin.rstrip("/") not in settings.allowed_origins: raise HTTPException(status_code=403, detail={"code": "ORIGIN_NOT_ALLOWED", "message": "请求来源不被允许"}) def require_csrf(request: Request) -> None: check_origin(request) header = request.headers.get("x-csrf-token") cookie = request.cookies.get("kaotings_csrf") if not header or not cookie or not hmac.compare_digest(header, cookie): raise HTTPException(status_code=403, detail={"code": "CSRF_FAILED", "message": "请求校验失败"}) def set_session_cookie(response, token: str) -> None: response.set_cookie( settings.session_cookie_name, token, max_age=settings.session_ttl_seconds, httponly=True, secure=settings.session_cookie_secure, samesite="lax", path="/", ) def clear_session_cookie(response) -> None: response.delete_cookie(settings.session_cookie_name, path="/") def set_csrf_cookie(response, token: str) -> None: response.set_cookie("kaotings_csrf", token, max_age=settings.session_ttl_seconds, httponly=False, secure=settings.session_cookie_secure, samesite="lax", path="/") def current_user(request: Request, connection: Connection = Depends(get_connection)) -> dict: token = request.cookies.get(settings.session_cookie_name) if not token: raise HTTPException(status_code=401, detail={"code": "AUTH_REQUIRED", "message": "请先登录"}) user = connection.execute( """ SELECT u.*, s.id AS session_id FROM sessions s JOIN users u ON u.id = s.user_id WHERE s.token_hash = %s AND s.revoked_at IS NULL AND s.expires_at > now() """, (token_digest(token),), ).fetchone() if not user: raise HTTPException(status_code=401, detail={"code": "AUTH_REQUIRED", "message": "会话已失效"}) if user["status"] != "active": raise HTTPException(status_code=403, detail={"code": "ACCOUNT_DISABLED", "message": "账户不可用"}) return user def require_admin(user: dict = Depends(current_user)) -> dict: if user["role"] != "admin": raise HTTPException(status_code=403, detail={"code": "FORBIDDEN", "message": "无权执行此操作"}) return user class RateLimiter: def __init__(self) -> None: self._events: dict[str, deque[float]] = defaultdict(deque) self._lock = threading.Lock() def allow(self, key: str, limit: int, window_seconds: int) -> bool: now = time.monotonic() with self._lock: events = self._events[key] while events and events[0] <= now - window_seconds: events.popleft() if len(events) >= limit: return False events.append(now) return True rate_limiter = RateLimiter() _captcha_lock = threading.Lock() _captchas: dict[str, dict[str, object]] = {} _captcha_alphabet = "23456789ABCDEFGHJKLMNPQRSTUVWXYZ" def new_captcha(client_key: str) -> tuple[str, str]: captcha_id = secrets.token_urlsafe(18) answer = "".join(secrets.choice(_captcha_alphabet) for _ in range(5)) now = time.monotonic() with _captcha_lock: for key, item in list(_captchas.items()): if float(item["expires_at"]) <= now: _captchas.pop(key, None) _captchas[captcha_id] = {"answer": answer, "client_key": client_key, "expires_at": now + 300, "attempts": 0} return captcha_id, answer def consume_captcha(captcha_id: str, answer: str, client_key: str) -> bool: with _captcha_lock: item = _captchas.get(captcha_id) if not item or float(item["expires_at"]) <= time.monotonic() or item["client_key"] != client_key: return False item["attempts"] = int(item["attempts"]) + 1 valid = hmac.compare_digest(str(item["answer"]), answer.strip().upper()) if valid or int(item["attempts"]) >= 5: _captchas.pop(captcha_id, None) return valid