M2-PT-DRP/source/utils/crypto/universal.py

57 lines
1.7 KiB
Python

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.")