aes.py

from Crypto.Cipher import AES
from Crypto import Random
import binascii
import base64

SEGMENT_SIZE = 128

# BLOCK_SIZE can be 16, 24 or 32
# key must be same number of bytes
BLOCK_SIZE = 16


def pad_string(value: str):
    """Pads the given string with zero-bytes"""
    length = len(value)

    pad_size = BLOCK_SIZE - (length % BLOCK_SIZE)

    return value.ljust(length + pad_size, "\x00")


def pad_bytes(value: bytes):
    """Pads the given bytes with zero-bytes"""
    length = len(value)

    pad_size = BLOCK_SIZE - (length % BLOCK_SIZE)
    return value.ljust(length + pad_size, binascii.unhexlify("00"))


def unpad_bytes(value: bytes):
    """Removes trailing zero-bytes from the given bytes"""
    while value[-1] == b"\x00":
        value = value[:-1]

    return value


def encrypt(in_string, preimage):
    """
    Takes in a plain string to encrypt, as well as 
    hex representations of the preimage 
    """
    key = binascii.unhexlify(preimage)
    if len(key) not in [16, 24, 32]:
        raise TypeError(f"Preimage has bad length: {len(key)} bytes")

    iv = Random.new().read(16)  # IV is always 16 bytes in CFB mode
    aes = AES.new(key, AES.MODE_CFB, IV=iv, segment_size=SEGMENT_SIZE)
    padded = pad_string(in_string)
    encrypted = aes.encrypt(padded)
    unpadded = unpad_bytes(encrypted)
    return serialize(unpad_bytes(encrypted), iv)


def serialize(ciphertext: bytes, iv: bytes):
    added = iv + ciphertext
    return base64.encodebytes(added)


def deserialize(encrypted: str) -> (bytes, bytes):
    decoded = binascii.a2b_base64(encrypted)
    iv = decoded[:16]
    cipher = decoded[16:]
    return (cipher, iv)


def decrypt(encrypted: bytes, preimage: str):
    """
    Takes in a serialized bytes object to decrypt, 
    as well as hex representation of the preimage 
    """
    key = binascii.unhexlify(preimage)
    if len(key) not in [16, 24, 32]:
        raise TypeError(f"Preimage has bad length: {len(key)} bytes")

    (ciphertext, iv) = deserialize(encrypted)
    aes = AES.new(key, AES.MODE_CFB, IV=iv, segment_size=SEGMENT_SIZE)
    padded = pad_bytes(ciphertext)
    decrypted = aes.decrypt(padded)
    unpadded = decrypted[: -(len(padded) - len(ciphertext))]
    return unpadded


cleartext = "The quick brown fox jumps over the lazy dog"
print(f"cleartext: {cleartext}")

preimage = "e67a00b510bcff7f4a0101ff5f7fb690"

encrypted = encrypt(cleartext, preimage)
print(f"encrypted: {encrypted}")

decrypted = decrypt(encrypted, preimage)
print(f"decrypted: {decrypted.decode('utf-8')}")

print()

encrypted_from_js = b"4XDIApIUztlL1Eqp2D+76C3iEX/Q7PVGNWP45CWcPtQmzWNdI+BwXoKicjuiONICFjD55AI376LLW5URfV2KXxb7Iw=="
(cipher, iv) = deserialize(encrypted_from_js)

print(f"encrypted_from_js: {encrypted_from_js}")
preimage_from_js = "1200f3c48578ac8740782fb03c7b9a08"
decrypted_from_js = decrypt(encrypted_from_js, preimage_from_js)
print(f"decrypted_from_js: {decrypted_from_js.decode('utf-8')}")