worldquant-alpha-system/backend/app/security.py

"""Server-side authentication and encryption utilities."""

import hashlib
import secrets
from datetime import timedelta

from argon2 import PasswordHasher
from argon2.exceptions import VerificationError
from cryptography.fernet import Fernet
from fastapi import Depends, HTTPException, Request
from fastapi.security import APIKeyCookie
from sqlalchemy import delete

from .models import Account, Admin, LoginSession, now

password_hasher = PasswordHasher()
session_cookie = APIKeyCookie(name="wq_session", auto_error=False)


def token_hash(token: str) -> str:
    return hashlib.sha256(token.encode()).hexdigest()


def cipher(settings) -> Fernet:
    return Fernet(settings.encryption_key.get_secret_value().encode())


async def bootstrap(db, settings):
    """Initialize singletons and apply environment credentials without resetting the admin password."""
    if not await db.get(Admin, 1):
        db.add(
            Admin(
                id=1,
                username=settings.admin_username,
                password_hash=password_hasher.hash(settings.admin_password.get_secret_value()),
            )
        )
    account = await db.get(Account, 1)
    if account is None:
        account = Account(id=1)
        db.add(account)
    if settings.wq_email:
        # The environment must not bypass the single-account data ownership boundary.
        if account.wq_user_id and (account.email or "").casefold() != settings.wq_email.casefold():
            raise ValueError("WQ_EMAIL conflicts with the bound WorldQuant account")
        account.email = settings.wq_email
        account.password_encrypted = cipher(settings).encrypt(
            settings.wq_password.get_secret_value().encode()
        ).decode()
    await db.execute(delete(LoginSession).where(LoginSession.expires_at < now()))
    await db.commit()


def valid_password(encoded: str, value: str) -> bool:
    try:
        return password_hasher.verify(encoded, value)
    except VerificationError:
        return False


async def issue_session(db, hours: int) -> str:
    token = secrets.token_urlsafe(32)
    db.add(LoginSession(token_hash=token_hash(token), expires_at=now() + timedelta(hours=hours)))
    await db.commit()
    return token


async def require_auth(request: Request, token: str | None = Depends(session_cookie)):
    if not token:
        raise HTTPException(401, "请先登录系统")
    async with request.app.state.sessions() as db:
        row = await db.get(LoginSession, token_hash(token))
        # SQLite fixtures lose tzinfo; production PostgreSQL preserves it.
        if row is None or row.expires_at.replace(tzinfo=now().tzinfo) <= now():
            raise HTTPException(401, "系统登录已过期，请重新登录")