yuxuanhui
56 lines
2.4 KiB
Python
56 lines
2.4 KiB
Python
"""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", "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))
|