buff/backend/app/core/encrypt.py

"""
RSA+AES混合加密工具 - 支持密钥版本管理
"""
import base64
import json
import time
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

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()  # 且未过期
            ).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()  # 且未过期
            ).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]:
        """
        解密客户端发送的数据

        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 _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 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()