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


def rsa_encrypt(data: bytes, public_key_data: bytes) -> bytes:
    """
    Encrypt data with RSA using a public key
    :param data: the data to encrypt
    :param public_key_data: the public key to encrypt with
    :return: the encrypted data
    """

    # load the public key
    public_key = serialization.load_der_public_key(public_key_data)
    # verify if the key is loaded
    if not isinstance(public_key, rsa.RSAPublicKey):
        raise ValueError("Could not load the public key.")

    # encrypt the data with the key
    return public_key.encrypt(
        data,
        padding.OAEP(
            mgf=padding.MGF1(algorithm=hashes.SHA256()),
            algorithm=hashes.SHA256(),
            label=None
        )
    )

def rsa_decrypt(encrypted_data: bytes, private_key_data: bytes) -> bytes:
    """
    Decrypt the data with the RSA private key
    :param encrypted_data: the data to decrypt
    :param private_key_data: the private key data
    :return: the decrypted data
    """

    # load the private key
    private_key = serialization.load_der_private_key(private_key_data, None)
    # verify if the key is loaded
    if not isinstance(private_key, rsa.RSAPrivateKey):
        raise ValueError("Could not load the private key.")

    # decrypt the data
    return private_key.decrypt(
        encrypted_data,
        padding.OAEP(
            mgf=padding.MGF1(algorithm=hashes.SHA256()),
            algorithm=hashes.SHA256(),
            label=None
        )
    )