from Crypto.Cipher import AES
from Crypto.Util import Padding, strxor



def ctr_enc(key, nonce, message):
    cipher_ctr = AES.new(key, AES.MODE_CTR, nonce=nonce, initial_value=1)
    return cipher_ctr.encrypt(message)


def ctr_dec(key, nonce, ciphertext):
    cipher_ctr = AES.new(key, AES.MODE_CTR, nonce=nonce, initial_value=1)
    return cipher_ctr.decrypt(ciphertext)


def ecb_enc(k, data):
    cipher_ctr = AES.new(k, AES.MODE_ECB)
    return cipher_ctr.encrypt(data)


def byte_xor(ba1, ba2):
    # https://nitratine.net/blog/post/xor-python-byte-strings/
    return bytes([_a ^ _b for _a, _b in zip(ba1, ba2)])


def cbc_mac(k, n, a, m):
    gruen = n + len(m).to_bytes(8, 'big')
    gelb = len(a).to_bytes(8, 'big') + a
    blau = m

    data = gruen + Padding.pad(gelb, 16) + blau
    data = Padding.pad(data, 16)

    cipher = AES.new(k, AES.MODE_CBC, iv=b'\00'*16)
    n0 = ecb_enc(k, n+b'\x00'*8)
    c = cipher.encrypt(data)[-16:]

    return byte_xor(n0, c)


def ccm_enc(k, n, a, m):
    # nonce based CTR ENC
    c = ctr_enc(k, n, m)
    # CBC-MAC with a-data
    t = cbc_mac(k, n, a, m)
    print(f'enc tag: {t}')
    return a, c, t


def ccm_dec(k, n, a, c, t):
    m = ctr_dec(k, n, c)
    # CBC-MAC with a-data
    t = cbc_mac(k, n, a, m)
    print(f'dec tag: {t}')
    return a, m

a = b'das ist gar nicht geheim'
m = b'this is a stupid and totally useless but very secret message'
k = b'einszweidreivier'
n = bytes.fromhex('1122334455667788')
(a1, c, t) = ccm_enc(k, n, a, m)
print("encrypted data:")
print(a1)
print(f' ciphertext: {c}')

(a2,m2) = ccm_dec(k,n,a1,c,t)

print("Checking decryption result: ", a==a2 and m==m2)