from cryptography.hazmat.primitives.asymmetric import rsa, padding
from cryptography.hazmat.primitives import serialization, hashes
from cryptography.hazmat.backends import default_backend
import os

def encrypt_key():
    # 获取当前脚本所在目录
    current_directory = os.path.dirname(__file__)

    # 读取对称密钥
    symmetric_key_path = os.path.join(current_directory, 'symmetric_key.key')

    if os.path.exists(symmetric_key_path):
        with open(symmetric_key_path, 'rb') as symmetric_key_file:
            symmetric_key = symmetric_key_file.read()

        # 读取B的公钥
        public_key_path = os.path.join(current_directory, 'B_public.txt')

        if os.path.exists(public_key_path):
            with open(public_key_path, 'rb') as public_key_file:
                public_key = serialization.load_pem_public_key(
                    public_key_file.read(),
                    backend=default_backend()
                )

            # 使用RSA公钥加密对称密钥
            encrypted_symmetric_key = public_key.encrypt(
                symmetric_key,
                padding.OAEP(
                    mgf=padding.MGF1(algorithm=hashes.SHA256()),
                    algorithm=hashes.SHA256(),
                    label=None
                )
            )

            # 保存加密后的密文
            encrypted_symmetric_key_path = os.path.join(current_directory, 'encrypted_symmetric_key.bin')

            with open(encrypted_symmetric_key_path, 'wb') as encrypted_symmetric_key_file:
                encrypted_symmetric_key_file.write(encrypted_symmetric_key)

            print(f"对称密钥已读取并加密，加密后的对称密钥已保存到: {encrypted_symmetric_key_path}")
        else:
            print(f"B的公钥文件 {public_key_path} 不存在")
    else:
        print(f"对称密钥文件 {symmetric_key_path} 不存在")

def decrypt_key():
    # 获取当前脚本所在目录
    current_directory = os.path.dirname(__file__)

    # 读取B的私钥
    private_key_path = os.path.join(current_directory, 'B_private.txt')

    if os.path.exists(private_key_path):
        with open(private_key_path, 'rb') as private_key_file:
            private_key = serialization.load_pem_private_key(
                private_key_file.read(),
                password=None,
                backend=default_backend()
            )

        # 读取加密后的对称密钥
        encrypted_symmetric_key_path = os.path.join(current_directory, 'encrypted_symmetric_key.bin')

        if os.path.exists(encrypted_symmetric_key_path):
            with open(encrypted_symmetric_key_path, 'rb') as encrypted_symmetric_key_file:
                encrypted_symmetric_key = encrypted_symmetric_key_file.read()
            
            # 使用B的私钥解密密文获得对称密钥
            symmetric_key = private_key.decrypt(
                encrypted_symmetric_key,
                padding.OAEP(
                    mgf=padding.MGF1(algorithm=hashes.SHA256()),
                    algorithm=hashes.SHA256(),
                    label=None
                )
            )

            # 保存解密后的对称密钥到文件
            decrypt_symmetric_key_path = os.path.join(current_directory, 'decrypt_symmetric_key.key')
            with open(decrypt_symmetric_key_path, 'wb') as decrypt_symmetric_key_file:
                decrypt_symmetric_key_file.write(symmetric_key)

            print(f"对称密钥已解密，得到的对称密钥已保存到: {decrypt_symmetric_key_path}")
        else:
            print(f"加密后的对称密钥文件 {encrypted_symmetric_key_path} 不存在")
    else:
        print(f"B的私钥文件 {private_key_path} 不存在")