video-gen/video-gen-api/app/services/api_v3/key_service.py

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