buff/backend/app/api/client/keys.py

"""
密钥管理接口
"""
import hmac
import hashlib
import base64
import time
import secrets
import json
from typing import Dict, Any
from fastapi import APIRouter, Depends, HTTPException
from sqlalchemy.orm import Session
from pydantic import BaseModel
from app.database import get_db
from app.models.api_key import APIKey
from app.core.encrypt import encryptor
from app.utils.timestamp_validator import TimestampValidator
from app.utils.timezone import beijing_now
from app.config import settings
from app.core.logger import logger

router = APIRouter()

class PublicKeyRequest(BaseModel):
    """公钥获取请求"""
    aes_key: str  # Base64编码的AES密钥（32字节）
    timestamp: int  # 请求时间戳
    nonce: str  # 随机字符串，防止重放攻击

@router.post("/public")
async def get_public_key(
    request: PublicKeyRequest,
    db: Session = Depends(get_db)
):
    """
    获取公钥 (AES加密传输版本)

    安全特性:
    - AES加密传输公钥数据
    - 时间戳验证防止重放攻击
    - HMAC签名确保数据完整性
    """
    try:
        # 1. 验证时间戳
        is_valid, error_msg = TimestampValidator.validate_timestamp(request.timestamp, db)
        if not is_valid:
            logger.warning(f"公钥获取请求时间戳验证失败: {error_msg}")
            return {
                "code": 400,
                "msg": f"时间戳验证失败: {error_msg}",
                "timestamp": TimestampValidator.generate_timestamp()
            }

        # 2. 获取当前激活的密钥（主is_active，辅expire_at）
        key = db.query(APIKey).filter(
            APIKey.is_active == 1,
            APIKey.is_deleted == 0,
            APIKey.expire_at > beijing_now()
        ).order_by(APIKey.key_version.desc()).first()
        if not key:
            logger.error("未找到激活的有效密钥")
            return {
                "code": 404,
                "msg": "系统密钥服务暂不可用",
                "timestamp": TimestampValidator.generate_timestamp()
            }

        # 3. 验证并解码AES密钥
        try:
            aes_key_bytes = base64.b64decode(request.aes_key)
            if len(aes_key_bytes) != 32:
                raise ValueError("AES密钥长度无效")
        except Exception as e:
            logger.warning(f"AES密钥解码失败: {e}")
            return {
                "code": 400,
                "msg": "无效的AES密钥格式",
                "timestamp": TimestampValidator.generate_timestamp()
            }

        # 4. 准备要加密的数据
        public_key_data = {
            "public_key": key.public_key,
            "key_version": key.key_version,
            "expire_at": key.expire_at.isoformat() if key.expire_at else None
        }

        # 5. 使用AES加密公钥数据
        # 生成新的AES密钥和IV用于加密响应数据
        response_aes_key = secrets.token_bytes(32)  # 256位AES密钥
        response_iv = secrets.token_bytes(16)  # 128位IV

        # AES加密公钥数据
        data_json = json.dumps(public_key_data, ensure_ascii=False)
        encrypted_data = _aes_encrypt(data_json.encode('utf-8'), response_aes_key, response_iv)

        # 使用请求中的AES密钥加密响应AES密钥和IV
        key_iv_combined = response_aes_key + response_iv
        encrypted_key_iv = _aes_encrypt(key_iv_combined, aes_key_bytes, aes_key_bytes[16:32])

        # 6. 生成签名
        timestamp = TimestampValidator.generate_timestamp()
        signature_data = f"{timestamp}:{request.nonce}:{base64.b64encode(encrypted_data).decode()}"
        signature = _generate_signature(signature_data, settings.SECRET_KEY)

        # 7. 返回响应
        response = {
            "code": 200,
            "msg": "成功",
            "data": {
                "encrypted_key": base64.b64encode(encrypted_key_iv).decode(),
                "encrypted_data": base64.b64encode(encrypted_data).decode(),
                "signature": signature,
                "timestamp": timestamp,
                "nonce": request.nonce
            }
        }

        logger.info(f"公钥获取成功，版本: {key.key_version}")
        return response

    except Exception as e:
        logger.exception("公钥获取异常")
        return {
            "code": 500,
            "msg": "系统异常，请稍后重试；若持续发生请联系客服",
            "timestamp": TimestampValidator.generate_timestamp()
        }

def _aes_encrypt(data: bytes, key: bytes, iv: bytes) -> bytes:
    """AES加密"""
    from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes
    from cryptography.hazmat.primitives.padding import PKCS7
    from cryptography.hazmat.backends import default_backend

    cipher = Cipher(algorithms.AES(key), modes.CBC(iv), backend=default_backend())
    encryptor = cipher.encryptor()

    # PKCS7填充
    padder = PKCS7(128).padder()
    padded_data = padder.update(data) + padder.finalize()

    return encryptor.update(padded_data) + encryptor.finalize()

def _generate_signature(data: str, secret_key: str) -> str:
    """生成HMAC签名"""
    signature = hmac.new(
        secret_key.encode(),
        data.encode(),
        hashlib.sha256
    ).digest()
    return base64.b64encode(signature).decode()

@router.get("/version")
async def get_key_version(db: Session = Depends(get_db)):
    """
    获取当前密钥版本信息
    返回当前激活的密钥版本，用于客户端配置验证
    """
    # 获取当前激活的密钥（主is_active，辅expire_at）
    key = db.query(APIKey).filter(
        APIKey.is_active == 1,
        APIKey.is_deleted == 0,
        APIKey.expire_at > beijing_now()
    ).order_by(APIKey.key_version.desc()).first()

    if not key:
        return {"code": 404, "msg": "未找到激活的有效密钥"}

    # 获取所有可用的密钥版本列表（包括当前激活的和新预热的）
    available_keys = db.query(APIKey).filter(
        APIKey.is_active == 1,    # 必须激活
        APIKey.is_deleted == 0,
        APIKey.expire_at > beijing_now()  # 且未过期
    ).order_by(APIKey.key_version.desc()).all()

    available_versions = [k.key_version for k in available_keys]

    return {
        "code": 200,
        "msg": "成功",
        "data": {
            "current_version": key.key_version,
            "available_versions": available_versions,
            "public_key": key.public_key,  # 仍然提供公钥，用于客户端更新配置
            "expire_at": key.expire_at.isoformat() if key.expire_at else None
        }
    }