buff3.0/backend-client/app/middleware/security.py

"""
安全中间件 - 支持AES+RSA混合加密和签名验证
"""
import os
import hmac
import hashlib
import base64
import time
import secrets
import json
import inspect
from typing import Callable, Optional, Dict, Any, Tuple
from fastapi import Request, HTTPException, Depends
from fastapi.responses import JSONResponse
from sqlalchemy.orm import Session
from datetime import datetime, timedelta

from app.config import settings
from app.core.encrypt import encryptor
from app.database import SessionLocal, get_db
from app.models.api_nonce import APINonce
from app.models.api_key import APIKey
from app.models.user import User
from app.models.session import Session as SessionModel
from app.models.software import Software
from app.models.user_software_subscription import UserSoftwareSubscription
from app.utils.timezone import beijing_now_naive
from app.utils.api_logger import api_logger
from app.utils.api_context import get_request_context, set_request_context
from app.utils.machine_code import raw_machine_info_to_machine_code
from app.core.logger import logger


def _normalize_machine_code(decrypted_params: dict) -> str:
    """与登录接口一致：对 raw_machine_info 规范化后做 MD5，再与 session.machine_code 比较"""
    raw = (decrypted_params.get("raw_machine_info") or decrypted_params.get("machine_code") or "").strip()
    return raw_machine_info_to_machine_code(raw)


# 服务端签名密钥：从环境变量获取，未设置时默认 buff-*-secret-{当前年}（每次调用取当前年，跨年无需重启）
def get_server_secret() -> bytes:
    """获取服务端签名密钥（响应签名用）"""
    default = f"buff-server-secret-{datetime.now().year}"
    return os.getenv("BUFF_API_SERVER_SECRET", default).encode()


def get_client_secret() -> bytes:
    """获取客户端签名密钥（请求验签用）"""
    default = f"buff-client-secret-{datetime.now().year}"
    return os.getenv("BUFF_API_CLIENT_SECRET", default).encode()


# 服务端内置AES密钥 (用于公钥获取接口)
SERVER_AES_KEY = os.getenv("SERVER_AES_KEY", "server-aes-key-32bytes!").encode()


def generate_signature(data: str, secret_key: bytes = None) -> str:
    """生成HMAC-SHA256签名"""
    if secret_key is None:
        secret_key = get_server_secret()
    signature = hmac.new(
        secret_key,
        data.encode(),
        hashlib.sha256
    ).digest()
    return base64.b64encode(signature).decode()


def verify_signature(data: str, signature: str, secret_key: bytes = None) -> bool:
    """验证HMAC-SHA256签名"""
    if secret_key is None:
        secret_key = get_client_secret()
    try:
        expected = generate_signature(data, secret_key)
        return hmac.compare_digest(signature, expected)
    except Exception:
        return False


def verify_timestamp(timestamp: int, tolerance: int = 300) -> bool:
    """验证时间戳 (容忍±5分钟)"""
    current_time = int(time.time())
    return abs(current_time - timestamp) <= tolerance


def generate_nonce() -> str:
    """生成随机nonce"""
    return secrets.token_hex(16)


def verify_nonce(nonce: str, endpoint: str, db: Session) -> bool:
    """验证nonce (防重放)"""
    if not nonce:
        return False
    
    # 检查nonce是否已存在
    existing = db.query(APINonce).filter(
        APINonce.nonce == nonce,
        APINonce.endpoint == endpoint,
        APINonce.expires_at > beijing_now_naive()
    ).first()
    
    if existing:
        return False
    
    # 存储新nonce
    new_nonce = APINonce(
        nonce=nonce,
        endpoint=endpoint,
        expires_at=beijing_now_naive() + timedelta(minutes=5)
    )
    db.add(new_nonce)
    db.commit()
    
    return True


def decrypt_request_data(
    encrypted_params: str,
    key_version: int,
    db: Session,
    encrypted_aes_key: Optional[str] = None,
) -> Dict[str, Any]:
    """
    解密客户端请求数据。
    若同时提供 encrypted_params 与 encrypted_aes_key，按客户端分离格式解密（IV+AES-CTR + RSA 加密的 AES 密钥）；
    否则按单 blob 格式解密（兼容旧逻辑）。
    """
    try:
        if not encrypted_params:
            return {}
        if not settings.ENCRYPT_ENABLED:
            return json.loads(encrypted_params)
        if encrypted_aes_key:
            decrypted = encryptor.decrypt_from_client_separate(
                encrypted_params, encrypted_aes_key, key_version, db
            )
        else:
            decrypted = encryptor.decrypt_from_client(encrypted_params, key_version, db)
        return decrypted
    except Exception as e:
        logger.exception("请求数据解密失败")
        raise HTTPException(status_code=400, detail="请求数据解密失败")


def encrypt_response_data(data: Dict[str, Any], db: Session) -> Tuple[str, int]:
    """加密响应数据（用于非安全端点或无客户端密钥时）"""
    try:
        encrypted_data, key_version = encryptor.encrypt_for_client(data, db)
        return encrypted_data, key_version
    except Exception as e:
        logger.exception("响应数据加密失败")
        return json.dumps(data, ensure_ascii=False), 0


def encrypt_response_data_with_client_key(
    data: Dict[str, Any],
    encrypted_aes_key_b64: str,
    key_version: int,
    db: Session,
) -> Tuple[str, int]:
    """使用请求中的客户端 AES 密钥加密响应（IV + AES-CTR），与 client-tauri 约定一致"""
    try:
        return encryptor.encrypt_response_with_client_key(
            data, encrypted_aes_key_b64, key_version, db
        )
    except Exception as e:
        logger.exception("使用客户端密钥加密响应失败")
        raise


class SecureEndpoint:
    """安全端点装饰器类"""

    def __init__(
        self,
        require_login: bool = True,
        require_machine_code: bool = True,
        require_subscription_valid: bool = True,
    ):
        self.require_login = require_login
        self.require_machine_code = require_machine_code
        self.require_subscription_valid = require_subscription_valid
    
    def __call__(self, func: Callable):
        async def wrapper(request: Request, db: Session = Depends(get_db)):
            try:
                # 1. 获取并解密请求参数（仅从 Request 读 body，避免被 FastAPI 当作 body 参数校验导致 422）
                body = await request.body()
                body_str = body.decode() if body else "{}"
                logger.debug(
                    f"[SecureEndpoint] 收到请求 path={request.url.path} body_len={len(body_str)} body_preview={repr(body_str[:200]) if body_str else 'empty'}"
                )
                
                # 尝试解析JSON请求体
                try:
                    body_json = json.loads(body_str) if body_str else {}
                except json.JSONDecodeError as e:
                    logger.warning(f"[SecureEndpoint] JSON 解析失败 path={request.url.path} err={e} body_len={len(body_str)}")
                    body_json = {}
                
                body_keys = list(body_json.keys()) if body_json else []
                logger.debug(f"[SecureEndpoint] body_json keys={body_keys}")
                
                # 获取请求参数（与客户端 JSON 字段一致；数值统一为 int 避免 float 导致签名字符串不一致）
                encrypted_params = body_json.get("encrypted_params") or ""
                aes_key = body_json.get("aes_key") or ""
                _kv = body_json.get("key_version", 1)
                _ts = body_json.get("timestamp", 0)
                key_version = int(_kv) if _kv is not None else 1
                timestamp = int(_ts) if _ts is not None else 0
                nonce = body_json.get("nonce") or ""
                signature = body_json.get("signature") or ""
                
                # 2. 验证时间戳
                if not verify_timestamp(timestamp):
                    logger.warning(f"[SecureEndpoint] 时间戳验证失败 path={request.url.path} ts={timestamp}")
                    return JSONResponse(
                        status_code=400,
                        content={"code": 400, "msg": "请求已过期"}
                    )
                
                # 3. 验证nonce (防重放) - 心跳接口跳过验证，因为高频请求且只读不写
                # 时间戳验证（±5分钟）已足够防止重放攻击
                request_path = request.url.path
                if request_path != "/api/v1/heartbeat" and request_path != "/api/user/heartbeat":
                    if not verify_nonce(nonce, request_path, db):
                        logger.warning(f"[SecureEndpoint] nonce 验证失败 path={request_path} nonce_len={len(nonce)}")
                        return JSONResponse(
                            status_code=400,
                            content={"code": 400, "msg": "重复请求或nonce无效"}
                        )
                
                # 4. 验证客户端签名（确保与客户端构造顺序、类型一致）
                # 客户端格式: timestamp:nonce:encrypted_params:aes_key:key_version，均为字符串拼接
                sign_string = f"{timestamp}:{nonce}:{encrypted_params}:{aes_key}:{key_version}"
                logger.debug(
                    f"[SecureEndpoint] 签名字符串: ts={timestamp} kv={key_version} nonce_len={len(nonce)} "
                    f"enc_len={len(encrypted_params)} aes_key_len={len(aes_key)} sign_str_len={len(sign_string)}"
                )
                
                # 统一使用按年份写死的静态密钥验签（CLIENT_SECRET = buff-client-secret-{年}）
                if not verify_signature(sign_string, signature):
                    try:
                        expected = generate_signature(sign_string, get_client_secret())
                        sign_prefix = sign_string[:80] + "..." if len(sign_string) > 80 else sign_string
                        expect_tail = expected[-12:] if len(expected) >= 12 else expected
                        recv_tail = signature[-12:] if len(signature) >= 12 else signature
                        logger.warning(
                            f"[SecureEndpoint] 签名验证失败: body_len={len(body_str)} sign_str_len={len(sign_string)} "
                            f"ts={timestamp} kv={key_version} sign_prefix={sign_prefix} "
                            f"expect_tail={expect_tail} recv_tail={recv_tail}"
                        )
                    except Exception as e:
                        logger.exception(f"[SecureEndpoint] 签名验证失败且记录日志异常: {e}")
                    return JSONResponse(
                        status_code=400,
                        content={"code": 400, "msg": "签名验证失败"}
                    )
                
                # 5. 解密请求参数（客户端分离格式：encrypted_params + aes_key）
                decrypted_params = decrypt_request_data(
                    encrypted_params, key_version, db, encrypted_aes_key=aes_key
                )
                
                # 6. 登录验证 (如果需要)
                session = None
                user = None
                
                if self.require_login:
                    session_id = (decrypted_params.get("session_id") or "").strip()
                    # 与登录一致：客户端传原始串，后端统一 MD5 后与 session.machine_code 比较
                    machine_code = _normalize_machine_code(decrypted_params)

                    if not session_id:
                        return JSONResponse(
                            status_code=401,
                            content={"code": 401, "msg": "未登录"}
                        )
                    
                    # 验证会话
                    session = db.query(SessionModel).filter(
                        SessionModel.session_id == session_id,
                        SessionModel.is_active == 1,
                        SessionModel.expires_at > beijing_now_naive()
                    ).first()
                    
                    if not session:
                        return JSONResponse(
                            status_code=401,
                            content={"code": 401, "msg": "会话已过期，请重新登录"}
                        )
                    
                    # 验证机器码 (如果需要)：支持 session 存 MD5(32位hex) 或 旧版原始串
                    if self.require_machine_code:
                        raw_from_params = (decrypted_params.get("raw_machine_info") or decrypted_params.get("machine_code") or "").strip()
                        session_mc = (session.machine_code or "").strip()
                        is_session_md5 = len(session_mc) == 32 and all(c in "0123456789abcdef" for c in session_mc.lower())
                        if is_session_md5:
                            machine_code_match = machine_code == session_mc
                        else:
                            machine_code_match = raw_from_params == session_mc
                        if not machine_code_match:
                            logger.warning(
                                "[SecureEndpoint] 机器码不匹配: session_id=%s keys=%s "
                                "session_mc_len=%s req_norm_len=%s session_prefix=%s req_prefix=%s raw_len=%s",
                                session_id[:16] if session_id else "",
                                list(decrypted_params.keys()),
                                len(session_mc),
                                len(machine_code or ""),
                                session_mc[:8],
                                (machine_code or "")[:8],
                                len(raw_from_params),
                            )
                            return JSONResponse(
                                status_code=403,
                                content={"code": 403, "msg": "账号已在其他设备登录，如非本人操作请联系客服"}
                            )
                    
                    # 获取用户信息
                    user = db.query(User).filter(
                        User.id == session.user_id,
                        User.is_deleted == 0
                    ).first()
                    
                    if not user:
                        return JSONResponse(
                            status_code=404,
                            content={"code": 404, "msg": "用户不存在"}
                        )
                    
                    # 检查用户状态
                    if user.status != 1:
                        return JSONResponse(
                            status_code=403,
                            content={"code": 403, "msg": "账号已被禁用"}
                        )
                    
                    if self.require_subscription_valid:
                        # 按软件检查卡密有效期（多软件共用用户表：优先 user_software_subscription）
                        software_code = (decrypted_params.get("software_code") or "").strip() or "buff-xp-steam-sh"
                        sw = db.query(Software).filter(
                            Software.code == software_code,
                            Software.is_deleted == 0
                        ).first()
                        expire_time = None
                        if sw:
                            sub = db.query(UserSoftwareSubscription).filter(
                                UserSoftwareSubscription.user_id == user.id,
                                UserSoftwareSubscription.software_id == sw.id,
                                UserSoftwareSubscription.is_deleted == 0
                            ).first()
                            if sub and sub.expire_time is not None:
                                expire_time = sub.expire_time
                        if expire_time is None:
                            expire_time = user.expire_time
                        if expire_time and expire_time < beijing_now_naive():
                            return JSONResponse(
                                status_code=403,
                                content={"code": 403, "msg": "您的卡密已到期，请前往续费或联系客户支持"}
                            )
                
                # 7. 执行实际业务逻辑
                response_data = await func(decrypted_params, request, db, session, user)
                
                # 8. 使用请求中的客户端 AES 密钥加密响应（IV + AES-CTR），与 client-tauri 一致
                encrypted_response, response_key_version = encrypt_response_data_with_client_key(
                    response_data, aes_key, key_version, db
                )
                
                # 9. 生成响应签名
                response_timestamp = int(time.time())
                response_nonce = generate_nonce()
                response_sign_string = f"{response_timestamp}:{response_nonce}:{encrypted_response}:{response_key_version}"
                response_signature = generate_signature(response_sign_string)
                
                # 10. 记录API日志
                request_user_id = session.user_id if session else (user.id if user else None)
                client_ip = request.headers.get("X-Forwarded-For", "").split(",")[0].strip() or \
                           request.headers.get("X-Real-IP", "") or \
                           "unknown"
                
                # 解密响应数据用于日志记录
                decrypted_response = response_data

                # 更新请求上下文，供 APILoggingASGIMiddleware 写入 api_logs 明文（脱敏）
                try:
                    ctx = get_request_context()
                    if ctx is not None:
                        ctx = dict(ctx)
                        ctx["decrypted_request"] = decrypted_params
                        ctx["decrypted_response"] = decrypted_response
                        set_request_context(ctx)
                except Exception:
                    pass

                api_logger.log_api_call(
                    endpoint=request.url.path,
                    method=request.method,
                    encrypted_request=encrypted_params,
                    decrypted_request=decrypted_params,
                    encrypted_response=encrypted_response,
                    decrypted_response=decrypted_response,
                    processing_time_ms=0,  # 可以添加处理时间计算
                    status_code=200,
                    user_id=request_user_id,
                    client_ip=client_ip,
                )
                
                # 11. 返回加密响应（字段名与客户端验签约定一致）
                return {
                    "code": 200,
                    "msg": "成功",
                    "encrypted_data": encrypted_response,
                    "response_timestamp": response_timestamp,
                    "response_nonce": response_nonce,
                    "response_signature": response_signature,
                    "key_version": response_key_version
                }
                
            except HTTPException:
                raise
            except Exception as e:
                logger.exception(f"安全端点处理异常: {e}")
                # 记录错误日志
                api_logger.log_api_call(
                    endpoint=request.url.path if 'request' in dir() else "unknown",
                    method=request.method if 'request' in dir() else "unknown",
                    error_message=str(e),
                    status_code=500,
                )
                return JSONResponse(
                    status_code=500,
                    content={"code": 500, "msg": "系统异常，请稍后重试"}
                )
            finally:
                pass  # db 由 Depends(get_db) 注入，由 FastAPI 管理生命周期

        # 显式签名：仅 request + db，避免 FastAPI 按被装饰函数的 decrypted_params 解析 body 导致 422
        wrapper.__signature__ = inspect.signature(
            lambda request: None
        ).replace(
            parameters=[
                inspect.Parameter("request", inspect.Parameter.POSITIONAL_OR_KEYWORD, annotation=Request),
                inspect.Parameter("db", inspect.Parameter.POSITIONAL_OR_KEYWORD, annotation=Session, default=Depends(get_db)),
            ]
        )
        return wrapper


def secure_endpoint(
    require_login: bool = True,
    require_machine_code: bool = True,
    require_subscription_valid: bool = True,
):
    """安全端点装饰器工厂函数"""
    def decorator(func: Callable):
        return SecureEndpoint(
            require_login=require_login,
            require_machine_code=require_machine_code,
            require_subscription_valid=require_subscription_valid,
        ).__call__(func)
    return decorator


# ============ 统一安全响应中间件 ============

class SecureResponseMiddleware:
    """
    统一安全响应中间件
    自动为所有响应添加加密和签名
    """
    
    @staticmethod
    def create_secure_response(
        data: Any,
        db: Session,
        code: int = 200,
        msg: str = "成功"
    ) -> Dict[str, Any]:
        """
        创建带签名的安全响应
        符合技术方案文档要求的响应格式
        """
        from app.core.encrypt import (
            encryptor, 
            generate_response_signature,
            get_server_secret
        )
        
        try:
            # 构建响应数据
            response_data = {
                "code": code,
                "msg": msg,
                "data": data
            }
            
            # 加密响应数据
            encrypted_data, key_version = encryptor.encrypt_for_client(response_data, db)
            
            # 生成响应签名
            signature, timestamp, nonce = generate_response_signature(
                encrypted_data, 
                key_version,
                get_server_secret()
            )
            
            # 返回完整的带签名响应
            return {
                "code": 200,
                "msg": "成功",
                "encrypted_data": encrypted_data,
                "timestamp": timestamp,
                "nonce": nonce,
                "signature": signature,
                "key_version": key_version
            }
            
        except Exception as e:
            logger.exception("创建安全响应失败")
            # 如果加密失败，返回错误信息
            return {
                "code": 500,
                "msg": "系统异常，请稍后重试",
                "encrypted_data": None,
                "timestamp": int(time.time()),
                "nonce": secrets.token_hex(16),
                "signature": None,
                "key_version": 0
            }


def create_encrypted_response(
    data: Any,
    db: Session,
    code: int = 200,
    msg: str = "成功"
) -> Dict[str, Any]:
    """
    便捷函数：创建带签名的安全响应
    
    使用方法:
        return create_encrypted_response(user_data, db, 200, "成功")
    """
    return SecureResponseMiddleware.create_secure_response(data, db, code, msg)