100 lines
3.0 KiB
Python
100 lines
3.0 KiB
Python
import hashlib
|
|
import json
|
|
import logging
|
|
from datetime import datetime, timezone
|
|
|
|
from fastapi import Depends, HTTPException, status
|
|
from fastapi.security import HTTPAuthorizationCredentials, HTTPBearer
|
|
from sqlalchemy import select
|
|
from sqlalchemy.ext.asyncio import AsyncSession
|
|
|
|
from app.dependencies import get_db
|
|
from app.models.api.api_key import ApiKey
|
|
|
|
logger = logging.getLogger("videogen")
|
|
|
|
security = HTTPBearer(auto_error=False)
|
|
|
|
|
|
class ApiKeyContext:
|
|
"""API Key 验证上下文,携带解析后的可调用模型列表。"""
|
|
|
|
def __init__(self, api_key: ApiKey, callable_models: list[dict]):
|
|
self.api_key = api_key
|
|
self.api_key_id = api_key.id
|
|
self.callable_models = callable_models
|
|
|
|
|
|
async def get_api_key_dependency(
|
|
credentials: HTTPAuthorizationCredentials | None = Depends(security),
|
|
db: AsyncSession = Depends(get_db),
|
|
) -> ApiKeyContext:
|
|
"""FastAPI Dependency: 验证 API Key 并返回上下文。
|
|
|
|
验证流程:
|
|
1. 提取 Bearer <REDACTED>
|
|
2. SHA-256 哈希后查询数据库
|
|
3. 检查 is_active、deleted_at
|
|
4. 检查有效期 (valid_from, valid_until)
|
|
5. 检查配额 (quota_limit, quota_used)
|
|
6. 重置过期周期的配额
|
|
"""
|
|
if not credentials:
|
|
raise HTTPException(
|
|
status_code=status.HTTP_401_UNAUTHORIZED,
|
|
detail="缺少 Authorization 头,请提供 Bearer <REDACTED>",
|
|
)
|
|
|
|
token_hash = hashlib.sha256(credentials.credentials.encode()).hexdigest()
|
|
|
|
result = await db.execute(
|
|
select(ApiKey).where(
|
|
ApiKey.api_key_hash == token_hash,
|
|
ApiKey.is_active == True,
|
|
ApiKey.deleted_at.is_(None),
|
|
).limit(1)
|
|
)
|
|
key = result.scalar_one_or_none()
|
|
|
|
if not key:
|
|
raise HTTPException(
|
|
status_code=status.HTTP_401_UNAUTHORIZED,
|
|
detail="无效的 API Key",
|
|
)
|
|
|
|
now = datetime.now(timezone.utc)
|
|
|
|
# 检查有效期
|
|
if key.valid_from and now < key.valid_from:
|
|
raise HTTPException(
|
|
status_code=status.HTTP_403_FORBIDDEN,
|
|
detail="API Key 尚未生效",
|
|
)
|
|
if key.valid_until and now >= key.valid_until:
|
|
raise HTTPException(
|
|
status_code=status.HTTP_403_FORBIDDEN,
|
|
detail="API Key 已过期",
|
|
)
|
|
|
|
# 配额周期重置
|
|
from app.services.api_v3.key_service import reset_quota_if_needed
|
|
key = await reset_quota_if_needed(db, key)
|
|
|
|
# 检查配额
|
|
if key.quota_limit is not None and key.quota_used >= key.quota_limit:
|
|
raise HTTPException(
|
|
status_code=status.HTTP_429_TOO_MANY_REQUESTS,
|
|
detail=f"API Key 配额已用尽 (已用 {key.quota_used:.2f} / 限额 {key.quota_limit:.2f})",
|
|
)
|
|
|
|
# 解析可调用模型
|
|
try:
|
|
callable_models = json.loads(key.callable_models) if key.callable_models else []
|
|
except (json.JSONDecodeError, TypeError):
|
|
callable_models = []
|
|
|
|
# 更新最后使用时间
|
|
key.last_used_at = now
|
|
|
|
return ApiKeyContext(api_key=key, callable_models=callable_models)
|