buff2.0/backend-client/app/core/encrypt.py

"""
RSA+AES混合加密工具 - 支持密钥版本管理
支持生成带签名的安全响应
"""
import base64
import json
import time
import hmac
import hashlib
import secrets
from datetime import datetime
from typing import Dict, Any, Optional, Tuple
from sqlalchemy.orm import Session
from cryptography.hazmat.primitives import serialization
from cryptography.hazmat.primitives.asymmetric import rsa, padding
from cryptography.hazmat.primitives import hashes
from cryptography.hazmat.backends import default_backend
from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes
from cryptography.hazmat.primitives.padding import PKCS7
from app.config import settings
from app.core.logger import logger
from app.database import get_db
from app.models.api_key import APIKey
from app.utils.timestamp_validator import TimestampValidator
from app.utils.timezone import beijing_now_naive

class RSAEncryptor:
    """RSA+AES混合加密器 - 支持多版本密钥"""

    def __init__(self):
        self.private_key: Optional[rsa.RSAPrivateKey] = None
        self.public_key: Optional[rsa.RSAPublicKey] = None
        self.current_key_version: int = 0
        self.key_cache: Dict[int, Tuple[str, str]] = {}  # version -> (private_key, public_key)
    
    def generate_key_pair(self) -> tuple[str, str]:
        """生成密钥对"""
        private_key = rsa.generate_private_key(
            public_exponent=65537,
            key_size=2048,
            backend=default_backend()
        )
        public_key = private_key.public_key()
        
        # 序列化为PEM格式
        private_pem = private_key.private_bytes(
            encoding=serialization.Encoding.PEM,
            format=serialization.PrivateFormat.PKCS8,
            encryption_algorithm=serialization.NoEncryption()
        ).decode('utf-8')
        
        public_pem = public_key.public_bytes(
            encoding=serialization.Encoding.PEM,
            format=serialization.PublicFormat.SubjectPublicKeyInfo
        ).decode('utf-8')
        
        return private_pem, public_pem
    
    def load_private_key(self, private_key_pem: str, key_version: int = 0):
        """加载私钥"""
        try:
            self.private_key = serialization.load_pem_private_key(
                private_key_pem.encode(),
                password=None,
                backend=default_backend()
            )
            self.public_key = self.private_key.public_key()
            self.current_key_version = key_version
            logger.info(f"私钥加载成功，版本: {key_version}")
        except Exception as e:
            logger.exception("私钥加载失败")
            raise

    def load_key_by_version(self, db: Session, key_version: int) -> bool:
        """根据版本号加载密钥（主is_active，辅expire_at）"""
        try:
            # 查找指定版本的密钥，必须激活、未删除且未过期
            api_key = db.query(APIKey).filter(
                APIKey.key_version == key_version,
                APIKey.is_active == 1,      # 必须激活
                APIKey.is_deleted == 0,
                APIKey.expire_at > beijing_now_naive()  # 且未过期
            ).first()

            if not api_key:
                logger.warning(f"未找到有效的密钥版本 {key_version}（可能未激活、已过期或已删除）")
                return False

            # 检查缓存是否仍然有效（激活状态可能已改变）
            if key_version in self.key_cache:
                cached_private_pem, cached_public_pem = self.key_cache[key_version]
                # 验证缓存内容是否与数据库完全一致
                if cached_private_pem == api_key.private_key and cached_public_pem == api_key.public_key:
                    # 缓存有效，直接使用
                    self.load_private_key(cached_private_pem, key_version)
                    return True
                else:
                    # 缓存已过期或不一致，清除旧缓存
                    del self.key_cache[key_version]
                    logger.info(f"清除过期缓存: 版本 {key_version}")

            self.load_private_key(api_key.private_key, key_version)
            self.key_cache[key_version] = (api_key.private_key, api_key.public_key)
            return True

        except Exception as e:
            logger.exception(f"加载密钥版本 {key_version} 失败")
            return False

    def get_current_public_key(self, db: Session) -> Tuple[str, int]:
        """获取当前激活的公钥（主is_active，辅expire_at）"""
        try:
            api_key = db.query(APIKey).filter(
                APIKey.is_active == 1,      # 必须是激活状态
                APIKey.is_deleted == 0,
                APIKey.expire_at > beijing_now_naive()  # 且未过期
            ).order_by(APIKey.key_version.desc()).first()

            if api_key:
                return api_key.public_key, api_key.key_version

            logger.warning("未找到激活的有效公钥")
            return "", 0

        except Exception as e:
            logger.exception("获取当前公钥失败")
            return "", 0

    def encrypt_for_client(self, data: Dict[str, Any], db: Session) -> Tuple[str, int]:
        """
        为客户端加密数据

        Args:
            data: 要加密的数据
            db: 数据库会话

        Returns:
            tuple: (加密后的数据, 使用的密钥版本)
        """
        try:
            # 获取当前激活的公钥
            public_key_pem, key_version = self.get_current_public_key(db)
            if not public_key_pem:
                raise ValueError("未找到可用的公钥")

            # 如果当前加载的不是最新的密钥，则重新加载
            if self.current_key_version != key_version:
                self.load_key_by_version(db, key_version)

            # 生成AES密钥
            aes_key = self._generate_aes_key()

            # 使用AES加密数据
            encrypted_payload = self._aes_encrypt(json.dumps(data, ensure_ascii=False).encode('utf-8'), aes_key)

            # 使用RSA加密AES密钥
            encrypted_aes_key = self._rsa_encrypt_key(aes_key, public_key_pem)

            # 组合：encrypted_aes_key(256) + iv(16) + encrypted_payload
            iv = aes_key[16:32]  # 从AES密钥中提取IV
            result = encrypted_aes_key + iv + encrypted_payload

            # Base64编码
            encrypted_data = base64.b64encode(result).decode('utf-8')

            logger.debug(f"客户端数据加密成功，版本: {key_version}")
            return encrypted_data, key_version

        except Exception as e:
            logger.exception("客户端数据加密失败")
            raise

    def decrypt_from_client(self, encrypted_data: str, key_version: int, db: Session) -> Dict[str, Any]:
        """
        解密客户端发送的数据（单 blob：混合加密或纯 RSA）

        Args:
            encrypted_data: 加密的数据
            key_version: 使用的密钥版本
            db: 数据库会话

        Returns:
            dict: 解密后的数据
        """
        try:
            # 确保使用正确的密钥版本
            if self.current_key_version != key_version:
                if not self.load_key_by_version(db, key_version):
                    raise ValueError(f"无法加载密钥版本 {key_version}")

            return self.decrypt(encrypted_data)

        except Exception as e:
            logger.exception("客户端数据解密失败")
            raise

    def decrypt_from_client_separate(
        self,
        encrypted_params_b64: str,
        encrypted_aes_key_b64: str,
        key_version: int,
        db: Session,
    ) -> Dict[str, Any]:
        """
        解密客户端按「分离格式」发送的数据（与 client-tauri 一致）：
        - encrypted_params = base64(IV(16) + AES-256-CTR 密文)
        - aes_key = base64(RSA-OAEP 加密的 32 字节 AES 密钥)
        """
        if not encrypted_params_b64 or not encrypted_aes_key_b64:
            raise ValueError("encrypted_params 与 aes_key 均不能为空")
        if self.current_key_version != key_version:
            if not self.load_key_by_version(db, key_version):
                raise ValueError(f"无法加载密钥版本 {key_version}")
        if not self.private_key:
            raise ValueError("私钥未加载")
        try:
            # 1) RSA-OAEP 解密得到 32 字节 AES 密钥
            enc_aes = base64.b64decode(encrypted_aes_key_b64)
            aes_key = self.private_key.decrypt(
                enc_aes,
                padding.OAEP(
                    mgf=padding.MGF1(algorithm=hashes.SHA256()),
                    algorithm=hashes.SHA256(),
                    label=None,
                ),
            )
            if len(aes_key) != 32:
                raise ValueError(f"AES 密钥长度必须为 32，实际为 {len(aes_key)}")
            # 2) Base64 解码请求体：前 16 字节 IV，其余为 AES-CTR 密文
            params_bytes = base64.b64decode(encrypted_params_b64)
            if len(params_bytes) < 16:
                raise ValueError("encrypted_params 长度不足（需要至少 16 字节 IV）")
            iv = params_bytes[:16]
            ciphertext = params_bytes[16:]
            # 3) AES-256-CTR 解密（与客户端 Ctr64BE::<Aes256> 对应）
            cipher = Cipher(
                algorithms.AES(aes_key),
                modes.CTR(iv),
                backend=default_backend(),
            )
            decryptor = cipher.decryptor()
            decrypted_bytes = decryptor.update(ciphertext) + decryptor.finalize()
            return json.loads(decrypted_bytes.decode("utf-8"))
        except Exception as e:
            logger.exception("客户端数据解密失败(分离格式)")
            raise

    def _generate_aes_key(self) -> bytes:
        """生成AES密钥和IV"""
        import os
        return os.urandom(32)  # 256位AES密钥

    def _aes_encrypt(self, data: bytes, key: bytes) -> bytes:
        """AES加密"""
        iv = key[16:32]  # 使用密钥的后16字节作为IV
        cipher = Cipher(algorithms.AES(key[:16]), 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 _rsa_encrypt_key(self, aes_key: bytes, public_key_pem: str) -> bytes:
        """使用RSA加密AES密钥"""
        public_key = serialization.load_pem_public_key(
            public_key_pem.encode(),
            backend=default_backend()
        )

        return public_key.encrypt(
            aes_key,
            padding.OAEP(
                mgf=padding.MGF1(algorithm=hashes.SHA256()),
                algorithm=hashes.SHA256(),
                label=None
            )
        )

    def encrypt_response_with_client_key(
        self,
        response_data: Dict[str, Any],
        encrypted_aes_key_b64: str,
        key_version: int,
        db: Session,
    ) -> Tuple[str, int]:
        """
        使用请求中的客户端 AES 密钥加密响应（与 client-tauri 约定一致）。
        格式：Base64( IV(16) + AES-256-CTR(JSON) )，客户端用同一密钥解密。
        """
        if not encrypted_aes_key_b64:
            raise ValueError("encrypted_aes_key 为空")
        if self.current_key_version != key_version:
            if not self.load_key_by_version(db, key_version):
                raise ValueError(f"无法加载密钥版本 {key_version}")
        if not self.private_key:
            raise ValueError("私钥未加载")
        try:
            enc_aes = base64.b64decode(encrypted_aes_key_b64)
            aes_key = self.private_key.decrypt(
                enc_aes,
                padding.OAEP(
                    mgf=padding.MGF1(algorithm=hashes.SHA256()),
                    algorithm=hashes.SHA256(),
                    label=None,
                ),
            )
            if len(aes_key) != 32:
                raise ValueError(f"AES 密钥长度必须为 32，实际为 {len(aes_key)}")
            payload = json.dumps(response_data, ensure_ascii=False).encode("utf-8")
            iv = secrets.token_bytes(16)
            cipher = Cipher(
                algorithms.AES(aes_key),
                modes.CTR(iv),
                backend=default_backend(),
            )
            encryptor = cipher.encryptor()
            ciphertext = encryptor.update(payload) + encryptor.finalize()
            result = iv + ciphertext
            return base64.b64encode(result).decode("utf-8"), key_version
        except Exception as e:
            logger.exception("使用客户端密钥加密响应失败")
            raise

    def decrypt(self, encrypted_data: str) -> Dict[str, Any]:
        """解密数据（支持RSA+AES混合加密和纯RSA加密）"""
        if not settings.ENCRYPT_ENABLED:
            return json.loads(encrypted_data)
        
        if not self.private_key:
            raise ValueError("私钥未加载")
        
        try:
            # Base64解码
            encrypted_bytes = base64.b64decode(encrypted_data)
            
            # 判断是混合加密还是纯RSA加密
            # 混合加密：加密的AES密钥(256字节) + IV(16字节) + 加密的数据
            # 纯RSA加密：直接是RSA加密的数据（通常小于256字节）
            
            # 如果数据长度大于256字节，认为是混合加密
            if len(encrypted_bytes) > 256:
                # 混合加密模式
                # 提取加密的AES密钥（前256字节）
                encrypted_aes_key = encrypted_bytes[:256]
                # 提取IV（接下来16字节）
                iv = encrypted_bytes[256:272]
                # 提取加密的数据（剩余部分）
                encrypted_payload = encrypted_bytes[272:]
                
                # 记录调试信息：各部分长度
                try:
                    logger.debug(f"Encrypted total bytes: {len(encrypted_bytes)}, "
                                 f"enc_key_len: {len(encrypted_aes_key)}, iv_len: {len(iv)}, payload_len: {len(encrypted_payload)}")
                except Exception:
                    pass
                
                # 使用RSA解密AES密钥
                aes_key = self.private_key.decrypt(
                    encrypted_aes_key,
                    padding.OAEP(
                        mgf=padding.MGF1(algorithm=hashes.SHA256()),
                        algorithm=hashes.SHA256(),
                        label=None
                    )
                )
                
                # 使用AES解密数据
                cipher = Cipher(algorithms.AES(aes_key), modes.CBC(iv), backend=default_backend())
                decryptor = cipher.decryptor()
                padded_data = decryptor.update(encrypted_payload) + decryptor.finalize()
                
                # 去除PKCS7填充
                unpadder = PKCS7(128).unpadder()
                decrypted_bytes = unpadder.update(padded_data) + unpadder.finalize()

                # 记录解密后字节预览，便于调试（限制长度，避免日志过大）
                try:
                    preview = decrypted_bytes[:300]
                    # 尝试以utf-8安全显示
                    try:
                        preview_text = preview.decode('utf-8', errors='replace')
                    except Exception:
                        preview_text = str(preview)
                    logger.debug(f"Decrypted bytes preview (first 300): {preview_text}")
                    logger.debug(f"Decrypted total length: {len(decrypted_bytes)}")
                except Exception:
                    pass
                
                # JSON解析
                return json.loads(decrypted_bytes.decode('utf-8'))
            else:
                # 纯RSA加密模式（向后兼容）
                decrypted = self.private_key.decrypt(
                    encrypted_bytes,
                    padding.OAEP(
                        mgf=padding.MGF1(algorithm=hashes.SHA256()),
                        algorithm=hashes.SHA256(),
                        label=None
                    )
                )
                
                # JSON解析
                return json.loads(decrypted.decode('utf-8'))
        except Exception as e:
            logger.exception("数据解密失败")
            raise

# 全局加密器实例
encryptor = RSAEncryptor()


# ============ 响应签名相关 ============

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


def generate_response_signature(
    encrypted_data: str, 
    key_version: int, 
    server_secret: bytes = None
) -> Tuple[str, int, str]:
    """
    生成响应签名
    
    Args:
        encrypted_data: 加密后的响应数据 (Base64)
        key_version: 密钥版本
        server_secret: 服务端密钥 (可选，默认从环境变量获取)
        
    Returns:
        (signature, timestamp, nonce)
    """
    if server_secret is None:
        server_secret = get_server_secret()
    
    # 生成时间戳和nonce
    timestamp = int(time.time())
    nonce = secrets.token_hex(16)
    
    # 构建签名字符串: timestamp:nonce:encrypted_data:key_version
    sign_string = f"{timestamp}:{nonce}:{encrypted_data}:{key_version}"
    
    # HMAC-SHA256签名
    signature = hmac.new(
        server_secret,
        sign_string.encode(),
        hashlib.sha256
    ).digest()
    
    return base64.b64encode(signature).decode(), timestamp, nonce


def create_secure_response(
    data: Any, 
    db: Session, 
    code: int = 200, 
    msg: str = "成功"
) -> Dict[str, Any]:
    """
    创建带签名的安全响应 (符合技术方案文档要求)
    
    响应格式:
    {
        "code": 200,
        "msg": "成功",
        "encrypted_data": "base64(AES加密响应数据)",
        "timestamp": 1704105600,
        "nonce": "h8g7f6e5d4c3b2a1...",
        "signature": "base64(HMAC-SHA256签名)",
        "key_version": 1
    }
    
    Args:
        data: 响应数据
        db: 数据库会话
        code: 状态码
        msg: 消息
        
    Returns:
        带签名的安全响应字典
    """
    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)
        
        # 返回完整的带签名响应
        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 verify_response_signature(
    encrypted_data: str,
    signature: str,
    timestamp: int,
    nonce: str,
    key_version: int,
    server_secret: bytes = None
) -> bool:
    """
    验证响应签名 (客户端使用)
    
    Args:
        encrypted_data: 加密的响应数据
        signature: 响应签名
        timestamp: 响应时间戳
        nonce: 响应nonce
        key_version: 密钥版本
        server_secret: 服务端密钥
        
    Returns:
        签名是否有效
    """
    try:
        if server_secret is None:
            server_secret = get_server_secret()
        
        # 构建签名字符串
        sign_string = f"{timestamp}:{nonce}:{encrypted_data}:{key_version}"
        
        # 计算期望的签名
        expected_signature = hmac.new(
            server_secret,
            sign_string.encode(),
            hashlib.sha256
        ).digest()
        
        # 常数时间比较 (防止时序攻击)
        return hmac.compare_digest(
            base64.b64decode(signature),
            expected_signature
        )
        
    except Exception as e:
        logger.exception("验证响应签名失败")
        return False