import os

from cryptography.hazmat.backends import default_backend
from cryptography.hazmat.primitives import padding
from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes


def aes_ecb_encrypt(data: bytes, key: bytes) -> bytes:
    """
    Encrypt the message using AES in ECB mode.
    :param data: the data to cipher
    :param key: the key to use for the cipher
    :return: the encrypted data
    """

    # pad the data with PKCS7 for AES to work properly
    padder = padding.PKCS7(128).padder()
    padded_data = padder.update(data) + padder.finalize()

    # create the AES cipher in ECB mode
    cipher = Cipher(algorithms.AES(key), modes.ECB(), backend=default_backend())
    encryptor = cipher.encryptor()

    # encrypt the padded data
    encrypted_data = encryptor.update(padded_data) + encryptor.finalize()

    return encrypted_data

def aes_ecb_decrypt(encrypted_data: bytes, key: bytes) -> bytes:
    """
    Decrypt data encrypted with AES in CBC mode.
    :param encrypted_data: the encrypted data
    :param key: the key used to encrypt it
    :return: the decrypted data
    """

    # create the AES cipher in ECB mode
    cipher = Cipher(algorithms.AES(key), modes.ECB(), backend=default_backend())
    decryptor = cipher.decryptor()

    # decrypt the encrypted data
    decrypted_data = decryptor.update(encrypted_data) + decryptor.finalize()

    # unpad the data
    unpadder = padding.PKCS7(128).unpadder()
    data = unpadder.update(decrypted_data) + unpadder.finalize()

    return data

def aes_cbc_encrypt(data: bytes, key: bytes) -> bytes:
    """
    Encrypt the message using AES in CBC mode.
    :param data: the data to cipher
    :param key: the key to use for the cipher
    :return: the encrypted data
    """

    # pad the data with PKCS7 for AES to work properly
    padder = padding.PKCS7(128).padder()
    padded_data = padder.update(data) + padder.finalize()

    # create an initialisation vector
    iv = os.urandom(16)

    # create the AES cipher in CBC mode
    cipher = Cipher(algorithms.AES(key), modes.CBC(iv), backend=default_backend())
    encryptor = cipher.encryptor()

    # encrypt the padded data
    encrypted_data = encryptor.update(padded_data) + encryptor.finalize()

    # prepend the iv to the encrypted data
    return iv + encrypted_data


def aes_cbc_decrypt(payload: bytes, key: bytes) -> bytes:
    """
    Decrypt data encrypted with AES in CBC mode.
    :param payload: the encrypted data
    :param key: the key used to encrypt it
    :return: the decrypted data
    """

    # split the payload into the iv and the encrypted data
    iv = payload[:16]
    encrypted_data = payload[16:]

    # create the AES cipher in CBC mode
    cipher = Cipher(algorithms.AES(key), modes.CBC(iv), backend=default_backend())
    decryptor = cipher.decryptor()

    # decrypt the encrypted data
    decrypted_data = decryptor.update(encrypted_data) + decryptor.finalize()

    # unpad the data
    unpadder = padding.PKCS7(128).unpadder()
    data = unpadder.update(decrypted_data) + unpadder.finalize()

    return data