worldquant-alpha-system/backend/app/mcp_api/auth.py

"""Personal access tokens are isolated from browser and upstream credentials."""

import secrets
from dataclasses import dataclass
from datetime import timedelta, timezone
from uuid import uuid4

from fastapi import HTTPException
from sqlalchemy import select

from ..models import Account, Admin, MCPToken, now
from ..security import token_hash

SCOPES = frozenset({"research:read", "research:refresh", "research:write", "backtests:execute", "backtests:control"})


@dataclass(frozen=True)
class Principal:
    token_id: str
    admin_id: int
    account_id: int
    wq_user_id: str
    scopes: frozenset[str]


async def create_token(db, name, scopes=None, days=90):
    """Issue a token for the bound account; caller commits and reveals it once."""
    scopes = set(scopes if scopes is not None else ["research:read"])
    if not name.strip() or len(name) > 100 or not 1 <= days <= 365:
        raise ValueError("名称须为 1–100 字，有效期须为 1–365 天")
    if not scopes <= SCOPES or "research:read" not in scopes:
        raise ValueError("权限无效；所有令牌必须包含 research:read")
    account, admin = await db.get(Account, 1), await db.get(Admin, 1)
    if not account or not account.wq_user_id or not admin:
        raise ValueError("请先初始化系统并确认 WorldQuant 账户身份")
    secret = "wqmcp_" + secrets.token_urlsafe(32)
    row = MCPToken(
        id=str(uuid4()), token_hash=token_hash(secret), name=name.strip(), admin_id=admin.id,
        account_id=account.id, wq_user_id=account.wq_user_id, scopes=sorted(scopes),
        expires_at=now() + timedelta(days=days),
    )
    db.add(row)
    await db.flush()
    return row, secret


async def authenticate(db, secret):
    """Validate every request, including current account binding; return no secrets."""
    row = await db.scalar(select(MCPToken).where(MCPToken.token_hash == token_hash(secret)))
    if not row or row.revoked_at or row.expires_at.replace(tzinfo=row.expires_at.tzinfo or timezone.utc) <= now():
        raise HTTPException(401, "MCP 令牌无效或已过期")
    account, admin = await db.get(Account, row.account_id), await db.get(Admin, row.admin_id)
    if not account or not admin or account.id != 1 or account.wq_user_id != row.wq_user_id:
        raise HTTPException(401, "MCP 令牌账户绑定已失效")
    return Principal(row.id, row.admin_id, row.account_id, row.wq_user_id, frozenset(row.scopes))