import hashlib import json import logging import secrets from datetime import datetime, timezone from sqlalchemy import select from sqlalchemy.ext.asyncio import AsyncSession from app.models.api.api_key import ApiKey from app.utils.id_gen import generate_id logger = logging.getLogger("videogen") API_KEY_PREFIX = "vk_" async def create_api_key( db: AsyncSession, company_name: str, callable_models: list[dict] | None = None, quota_limit: float | None = None, quota_cycle: str | None = None, valid_from: datetime | None = None, valid_until: datetime | None = None, max_concurrent_video_tasks: int | None = None, description: str | None = None, ) -> tuple[ApiKey, str]: """创建新的 API Key。 Returns: (ApiKey 对象, 明文 API Key) — 明文仅返回这一次。 """ # 生成密钥: vk_ + 32字节随机hex raw_key = API_KEY_PREFIX + secrets.token_hex(24) # vk_ + 48位hex = 51字符 key_hash = hashlib.sha256(raw_key.encode()).hexdigest() key_prefix = raw_key[:8] # 前8位用于展示: vk_xxxxx api_key = ApiKey( id=generate_id(), company_name=company_name, api_key_hash=key_hash, api_key_prefix=key_prefix, description=description, callable_models=json.dumps(callable_models or [], ensure_ascii=False), quota_limit=quota_limit, quota_cycle=quota_cycle, quota_used=0.0, valid_from=valid_from, valid_until=valid_until, max_concurrent_video_tasks=max_concurrent_video_tasks, is_active=True, ) api_key.set_plaintext_key(raw_key) # 加密存储完整 Key db.add(api_key) await db.flush() logger.info("API Key created: id=%s company=%s prefix=%s", api_key.id, company_name, key_prefix) return api_key, raw_key async def list_api_keys( db: AsyncSession, skip: int = 0, limit: int = 50, company_name: str | None = None, is_active: bool | None = None, ) -> tuple[int, list[ApiKey]]: """列出 API Key(分页+筛选)。""" from sqlalchemy import func query = select(ApiKey).where(ApiKey.deleted_at.is_(None)) count_query = select(func.count(ApiKey.id)).where(ApiKey.deleted_at.is_(None)) if company_name: query = query.where(ApiKey.company_name.ilike(f"%{company_name}%")) count_query = count_query.where(ApiKey.company_name.ilike(f"%{company_name}%")) if is_active is not None: query = query.where(ApiKey.is_active == is_active) count_query = count_query.where(ApiKey.is_active == is_active) total_result = await db.execute(count_query) total = total_result.scalar_one() query = query.order_by(ApiKey.created_at.desc()).offset(skip).limit(limit) result = await db.execute(query) keys = list(result.scalars().all()) return total, keys async def get_api_key(db: AsyncSession, key_id: str) -> ApiKey | None: """获取单个 API Key 详情。""" result = await db.execute( select(ApiKey).where(ApiKey.id == key_id, ApiKey.deleted_at.is_(None)).limit(1) ) return result.scalar_one_or_none() async def update_api_key(db: AsyncSession, key: ApiKey, **kwargs) -> ApiKey: """更新 API Key 配置。""" updatable_fields = { "company_name", "description", "callable_models", "quota_limit", "quota_cycle", "valid_from", "valid_until", "max_concurrent_video_tasks", "is_active", } for field, value in kwargs.items(): if field in updatable_fields and value is not None: if field == "callable_models" and isinstance(value, list): value = json.dumps(value, ensure_ascii=False) setattr(key, field, value) await db.flush() return key async def adjust_quota( db: AsyncSession, key: ApiKey, action: str, quota_limit_delta: float | None = None, quota_limit: float | None = None, quota_cycle: str | None = None, ) -> tuple[ApiKey, dict]: """调整 API Key 配额。 返回 (更新后的 key, 变更详情 dict)。 action: - adjust: 增加总额,quota_limit_delta 累加到当前 quota_limit - reset_usage: 重置 quota_used 为 0 - set_limit: 直接设置 quota_limit - change_cycle: 修改 quota_cycle """ old_limit = key.quota_limit old_used = key.quota_used old_cycle = key.quota_cycle if action == "adjust": delta = quota_limit_delta or 0 key.quota_limit = round((key.quota_limit or 0) + delta, 2) elif action == "reset_usage": key.quota_used = 0.0 elif action == "set_limit": key.quota_limit = quota_limit # 允许设为 None(无限) elif action == "change_cycle": key.quota_cycle = quota_cycle # 允许设为 None(无限) else: raise ValueError(f"未知的调整操作: {action}") await db.flush() changes = { "old_limit": old_limit, "new_limit": key.quota_limit, "old_used": old_used, "new_used": key.quota_used, "old_cycle": old_cycle, "new_cycle": key.quota_cycle, } return key, changes async def delete_api_key(db: AsyncSession, key: ApiKey) -> None: """软删除 API Key。""" key.deleted_at = datetime.now(timezone.utc) key.is_active = False await db.flush() async def reset_quota_if_needed(db: AsyncSession, key: ApiKey) -> ApiKey: """检查并重置过期周期的配额。 - daily: 如果上次重置不是今天,重置 quota_used=0 - monthly: 如果上次重置不是本月,重置 quota_used=0 """ if key.quota_limit is None or key.quota_cycle is None: return key now = datetime.now(timezone.utc) # 使用 quota_used 的 updated_at 作为周期判断依据 last_reset = key.updated_at or key.created_at if last_reset is None: return key should_reset = False if key.quota_cycle == "daily": should_reset = last_reset.date() < now.date() elif key.quota_cycle == "monthly": should_reset = (last_reset.year, last_reset.month) < (now.year, now.month) if should_reset and key.quota_used > 0: key.quota_used = 0.0 await db.flush() logger.info("Quota reset for API Key %s (cycle=%s)", key.id, key.quota_cycle) return key