import base64
import binascii
from Crypto.Cipher import PKCS1_OAEP
from Crypto.PublicKey import RSA
from Crypto.Signature import pkcs1_15
from Crypto.Hash import SHA256


def generate_keys(priKeySavePath, pubKeySavePath):
    # 生成一个RSA密钥对象
    key = RSA.generate(2048)

    # 导出私钥
    private_key = key.export_key()
    with open(priKeySavePath, "wb") as f:
        f.write(private_key)

    # 导出公钥
    public_key = key.publickey().export_key()
    with open(pubKeySavePath, "wb") as f:
        f.write(public_key)


def encrypt_message(message, public_key_base64):
    # 加载公钥
    public_key_bin = base64.b64decode(public_key_base64)
    public_key = RSA.import_key(public_key_bin)
    # 使用公钥加密消息
    cipher_rsa = PKCS1_OAEP.new(public_key)
    encrypted_message = cipher_rsa.encrypt(message.encode('utf-8'))

    return binascii.hexlify(encrypted_message).decode('utf-8')


def decrypt_message(encrypted_message, private_key_base64):
    # 加载私钥
    private_key_bin = base64.b64decode(private_key_base64)
    private_key = RSA.import_key(private_key_bin)
    # 使用私钥解密消息
    cipher_rsa = PKCS1_OAEP.new(private_key)
    decrypted_message = cipher_rsa.decrypt(binascii.unhexlify(encrypted_message))

    return decrypted_message.decode('utf-8')

def sign_message(message, private_key_base64):
    # 加载私钥
    private_key_bin = base64.b64decode(private_key_base64)
    private_key = RSA.import_key(private_key_bin)
    # 计算消息的哈希值
    hash_obj = SHA256.new(message.encode('utf-8'))
    # 使用私钥对哈希值进行签名
    signature = pkcs1_15.new(private_key).sign(hash_obj)

    return binascii.hexlify(signature).decode('utf-8')


def verify_signature(message, signature, public_key_base64):
    # 加载公钥
    public_key_bin = base64.b64decode(public_key_base64)
    public_key = RSA.import_key(public_key_bin)
    # 计算消息的哈希值
    hash_obj = SHA256.new(message.encode('utf-8'))
    # 使用公钥验证签名
    try:
        pkcs1_15.new(public_key).verify(hash_obj, binascii.unhexlify(signature))
        return True
    except (ValueError, TypeError):
        return False