"""鉴权:客户端 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