import typing

from source.utils.crypto import aes, rsa
from source.utils.crypto.type import CipherType


def encrypt(data: bytes, key: typing.Optional[bytes] = None, cipher_type: CipherType = CipherType.NONE) -> bytes:
    """
    Encrypt data on various cipher type.
    :param data: the data to cipher
    :param key: the key to cipher the data
    :param cipher_type: the type of cipher to use
    :return:
    """

    match cipher_type:
        case CipherType.NONE:
            return data

        case CipherType.AES_ECB:
            if key is None:
                raise ValueError("The key cannot be None.")
            return aes.aes_ecb_encrypt(data, key)

        case CipherType.RSA:
            if key is None:
                raise ValueError("The key cannot be None.")
            return rsa.rsa_encrypt(data, key)

        case _:
            raise KeyError("Unknown cipher mode.")

def decrypt(data: bytes, key: typing.Optional[bytes] = None, cipher_type: CipherType = CipherType.NONE) -> bytes:
    """
    Encrypt data on various cipher type.
    :param data: the data to decipher
    :param key: the key to cipher the data
    :param cipher_type: the type of cipher to use
    :return:
    """

    match cipher_type:
        case CipherType.NONE:
            return data

        case CipherType.AES_ECB:
            if key is None:
                raise ValueError("The key cannot be None.")
            return aes.aes_ecb_decrypt(data, key)

        case CipherType.RSA:
            if key is None:
                raise ValueError("The key cannot be None.")
            return rsa.rsa_decrypt(data, key)

        case _:
            raise KeyError("Unknown cipher mode.")