"""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, "系统登录已过期,请重新登录")