99 lines
No EOL
3 KiB
Python
99 lines
No EOL
3 KiB
Python
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 |