"""
RSA+AES混合加密工具
"""
import base64
import json
from typing import Dict, Any, Optional
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

class RSAEncryptor:
    """RSA+AES混合加密器"""
    
    def __init__(self):
        self.private_key: Optional[rsa.RSAPrivateKey] = None
        self.public_key: Optional[rsa.RSAPublicKey] = None
    
    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):
        """加载私钥"""
        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()
            logger.info("私钥加载成功")
        except Exception as e:
            logger.error(f"私钥加载失败: {e}")
            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:]
                
                # 使用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()
                
                # 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.error(f"数据解密失败: {e}")
            raise

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