import base64
import binascii

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


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)

"""
输入：
需要公私钥加解密的消息：字节类型
公私钥：base64
输出：
加解密后的数据：字节类型
"""
def encrypt_message(message: bytes, 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)

    return encrypted_message


def decrypt_message(encrypted_message: bytes, 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(encrypted_message)

    return decrypted_message


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

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


def verify_signature(message: bytes, 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)
    # 使用公钥验证签名
    try:
        pkcs1_15.new(public_key).verify(hash_obj, binascii.unhexlify(signature))
        return True
    except (ValueError, TypeError):
        return False