graph_rag_server/utils/asymmetric_encryption.py

from cryptography.hazmat.primitives import serialization, hashes
from cryptography.hazmat.primitives.asymmetric import rsa, padding
from cryptography.hazmat.backends import default_backend
from typing import Tuple, Optional
import base64


class AsymmetricEncryption:
    """非对称加密工具类，使用RSA算法"""

    @staticmethod
    def generate_key_pair(key_size: int = 2048) -> Tuple[bytes, bytes]:
        """
        生成RSA密钥对
        
        Args:
            key_size: 密钥大小，默认为2048位
        
        Returns:
            Tuple[bytes, bytes]: (私钥PEM格式，公钥PEM格式)
        """
        # 生成私钥
        private_key = rsa.generate_private_key(
            public_exponent=65537,
            key_size=key_size,
            backend=default_backend()
        )
        
        # 生成公钥
        public_key = private_key.public_key()
        
        # 将私钥序列化为PEM格式
        private_pem = private_key.private_bytes(
            encoding=serialization.Encoding.PEM,
            format=serialization.PrivateFormat.TraditionalOpenSSL,
            encryption_algorithm=serialization.NoEncryption()
        )
        
        # 将公钥序列化为PEM格式
        public_pem = public_key.public_bytes(
            encoding=serialization.Encoding.PEM,
            format=serialization.PublicFormat.SubjectPublicKeyInfo
        )
        
        return private_pem, public_pem
    
    @staticmethod
    def encrypt(message: str, public_key_pem: bytes) -> str:
        """
        使用公钥加密数据
        
        Args:
            message: 要加密的明文
            public_key_pem: 公钥PEM格式
        
        Returns:
            str: 加密后的base64编码字符串
        """
        # 加载公钥
        public_key = serialization.load_pem_public_key(
            public_key_pem,
            backend=default_backend()
        )
        
        # 加密数据
        encrypted = public_key.encrypt(
            message.encode('utf-8'),
            padding.OAEP(
                mgf=padding.MGF1(algorithm=hashes.SHA256()),
                algorithm=hashes.SHA256(),
                label=None
            )
        )
        
        # 返回base64编码的加密数据
        return base64.b64encode(encrypted).decode('utf-8')
    
    @staticmethod
    def decrypt(encrypted_message: str, private_key_pem: bytes) -> str:
        """
        使用私钥解密数据
        
        Args:
            encrypted_message: 加密后的base64编码字符串
            private_key_pem: 私钥PEM格式
        
        Returns:
            str: 解密后的明文
        """
        # 加载私钥
        private_key = serialization.load_pem_private_key(
            private_key_pem,
            password=None,
            backend=default_backend()
        )
        
        # 解码base64加密数据
        encrypted = base64.b64decode(encrypted_message)
        
        # 解密数据
        decrypted = private_key.decrypt(
            encrypted,
            padding.OAEP(
                mgf=padding.MGF1(algorithm=hashes.SHA256()),
                algorithm=hashes.SHA256(),
                label=None
            )
        )
        
        # 返回解密后的明文
        return decrypted.decode('utf-8')
    
    @staticmethod
    def save_key_to_file(key_pem: bytes, file_path: str) -> None:
        """
        将密钥保存到文件
        
        Args:
            key_pem: 密钥PEM格式
            file_path: 文件路径
        """
        with open(file_path, 'wb') as f:
            f.write(key_pem)
    
    @staticmethod
    def load_key_from_file(file_path: str) -> bytes:
        """
        从文件加载密钥
        
        Args:
            file_path: 文件路径
        
        Returns:
            bytes: 密钥PEM格式
        """
        with open(file_path, 'rb') as f:
            return f.read()
    
    @staticmethod
    def generate_key_pair_and_save(private_key_path: str, public_key_path: str, key_size: int = 2048) -> None:
        """
        生成密钥对并保存到文件
        
        Args:
            private_key_path: 私钥文件路径
            public_key_path: 公钥文件路径
            key_size: 密钥大小，默认为2048位
        """
        private_pem, public_pem = AsymmetricEncryption.generate_key_pair(key_size)
        AsymmetricEncryption.save_key_to_file(private_pem, private_key_path)
        AsymmetricEncryption.save_key_to_file(public_pem, public_key_path)