aes.py

#!/usr/bin/env python3

from codecs import decode, encode
from hashlib import md5
from json import dumps, loads

from Crypto.Cipher import AES

N = 16


def pad(s: bytes) -> bytes:
    n = N - len(s) % N
    return s + bytes(n for _ in range(n))


def unpad(s: bytes) -> bytes:
    return s[0:-s[-1]]


def new_aes(secret: str, version: str) -> AES.AESCipher:
    key = md5(secret.encode()).digest()
    iv = md5((secret + version).encode()).hexdigest()[8:-8]
    return AES.new(key, 2, iv.upper().encode())


def encode(secret: str, version: str, raw: object) -> bytes:
    aes = new_aes(secret, version)
    payload = dumps(raw, separators=(",", ":")).encode()
    return encode(aes.encrypt(pad(payload)), 'base64')


def decode(secret: str, version: str, mess: bytes) -> object:
    aes = new_aes(secret, version)
    payload = unpad(aes.decrypt(decode(mess, 'base64')))
    return loads(payload)