Spaces:
Paused
Paused
| """鉴权：客户端 API Key 校验 + 管理员 JWT。""" | |
| from __future__ import annotations | |
| import hashlib | |
| import secrets | |
| from datetime import datetime, timedelta, timezone | |
| from fastapi import Depends, HTTPException | |
| from fastapi.security import HTTPAuthorizationCredentials, HTTPBearer | |
| from jose import JWTError, jwt | |
| from passlib.context import CryptContext | |
| from sqlalchemy import select | |
| from sqlalchemy.ext.asyncio import AsyncSession | |
| from . import cache | |
| from .config import get_settings | |
| from .db import SessionLocal, get_db | |
| from .models import AdminUser, ApiKey | |
| settings = get_settings() | |
| pwd = CryptContext(schemes=["bcrypt"], deprecated="auto") | |
| bearer = HTTPBearer(auto_error=False) | |
| def hash_password(p: str) -> str: | |
| return pwd.hash(p) | |
| def verify_password(p: str, h: str) -> bool: | |
| try: | |
| return pwd.verify(p, h) | |
| except Exception: | |
| return False | |
| def generate_api_key() -> tuple[str, str, str]: | |
| """生成新 key，返回 (完整 key, 前缀, 哈希)。""" | |
| raw = "sk-relay-" + secrets.token_urlsafe(32) | |
| prefix = raw[:20] | |
| digest = hashlib.sha256(raw.encode()).hexdigest() | |
| return raw, prefix, digest | |
| def verify_api_key(raw: str, digest: str) -> bool: | |
| return cache.verify_key_hash(raw, digest) | |
| def create_admin_token(username: str, expires_hours: int | None = None) -> str: | |
| exp = timedelta(hours=expires_hours or settings.jwt_expire_hours) | |
| payload = {"sub": username, "exp": datetime.now(timezone.utc) + exp} | |
| return jwt.encode(payload, settings.jwt_secret, algorithm=settings.jwt_algo) | |
| async def get_admin( | |
| credentials: HTTPAuthorizationCredentials | None = Depends(bearer), | |
| db: AsyncSession = Depends(get_db), | |
| ) -> AdminUser: | |
| """后台接口依赖：校验管理员 JWT。""" | |
| if not credentials or not credentials.credentials: | |
| raise HTTPException(status_code=401, detail="未登录") | |
| try: | |
| payload = jwt.decode(credentials.credentials, settings.jwt_secret, algorithms=[settings.jwt_algo]) | |
| except JWTError: | |
| raise HTTPException(status_code=401, detail="凭证无效或已过期") | |
| user = ( | |
| await db.execute(select(AdminUser).where(AdminUser.username == payload.get("sub"))) | |
| ).scalar_one_or_none() | |
| if not user: | |
| raise HTTPException(status_code=401, detail="用户不存在") | |
| return user | |
| async def _load_key_from_db(prefix: str, db: AsyncSession) -> ApiKey | None: | |
| """从数据库按前缀加载 key。""" | |
| return ( | |
| await db.execute(select(ApiKey).where(ApiKey.prefix == prefix)) | |
| ).scalar_one_or_none() | |
| async def get_valid_api_key( | |
| credentials: HTTPAuthorizationCredentials | None = Depends(bearer), | |
| ) -> cache.CachedKey: | |
| """代理接口依赖：校验客户端 relay key，优先走内存缓存。 | |
| 缓存未命中时回库加载，加载失败则抛 401（无效 key）。 | |
| 返回 CachedKey（内存缓存对象），供后续配额/日志使用，避免再查库。 | |
| """ | |
| if not credentials or not credentials.credentials: | |
| raise HTTPException( | |
| status_code=401, | |
| detail="缺少 API Key", | |
| headers={"WWW-Authenticate": "Bearer"}, | |
| ) | |
| raw = credentials.credentials.strip() | |
| # 先查内存缓存 | |
| ck = cache.key_cache.get(raw[:20]) | |
| if ck is not None and verify_api_key(raw, ck.key_hash): | |
| if not ck.is_active: | |
| raise HTTPException(status_code=401, detail="无效的 API Key") | |
| if ck.expires_at and ck.expires_at < datetime.now(timezone.utc).timestamp(): | |
| raise HTTPException(status_code=401, detail="API Key 已过期") | |
| return ck | |
| # 缓存未命中或哈希不匹配，回库加载 | |
| try: | |
| async with SessionLocal() as db: | |
| db_key = await _load_key_from_db(raw[:20], db) | |
| except Exception: | |
| # 数据库挂了：若缓存里还有这个 key（哈希匹配但已过期/停用等分支没进），仍可用 | |
| ck = cache.key_cache.get(raw[:20]) | |
| if ck is not None and verify_api_key(raw, ck.key_hash): | |
| return ck | |
| raise HTTPException(status_code=503, detail="数据库暂时不可用，请稍后再试") | |
| if not db_key or not db_key.is_active or not verify_api_key(raw, db_key.key_hash): | |
| raise HTTPException(status_code=401, detail="无效的 API Key") | |
| if db_key.expires_at and db_key.expires_at < datetime.now(timezone.utc): | |
| raise HTTPException(status_code=401, detail="API Key 已过期") | |
| # 缓存校验信息 | |
| exp = db_key.expires_at.timestamp() if db_key.expires_at else None | |
| ck = cache.CachedKey( | |
| id=db_key.id, name=db_key.name, prefix=db_key.prefix, key_hash=db_key.key_hash, | |
| is_active=db_key.is_active, expires_at=exp, allowed_models=db_key.allowed_models, | |
| rpm_limit=db_key.rpm_limit, tpm_limit=db_key.tpm_limit, | |
| token_quota=db_key.token_quota, tokens_used=db_key.tokens_used or 0, | |
| loaded_at=datetime.now(timezone.utc).timestamp(), | |
| ) | |
| await cache.key_cache.put(ck) | |
| return ck | |