AES_CMAC.py

from Crypto.Cipher import AES
from binascii import unhexlify

from math import ceil

const_Zero = "00000000000000000000000000000000"
const_Rb   = 0x87

const_Bsize = 16

aes_128 = lambda key, data: AES.new(key=key, mode=AES.MODE_ECB).encrypt(data)


padding = lambda x: x + b'\x80' + (b'\x00'*(const_Bsize-1-len(x)))

def MSB(b):
    return b & 0x80
    

def Generate_Subkey(k):
    L = aes_128(k, unhexlify(const_Zero))
    if MSB(L[0]) == 0:
        K1 = (int.from_bytes(L, 'big') << 1 & int('ff'*16, 16)).to_bytes(16, 'big')
    else:
        K1 = (int.from_bytes(L, 'big') << 1 & int('ff'*16, 16) ^ const_Rb).to_bytes(16, 'big')
    if MSB(K1[0]) == 0:
        K2 = (int.from_bytes(K1, 'big') << 1 & int('ff'*16, 16)).to_bytes(16, 'big')
    else:
        K2 = (int.from_bytes(K1, 'big') << 1 & int('ff'*16, 16) ^ const_Rb).to_bytes(16, 'big')
    return K1, K2

def AES_CMAC(k, m, l):
    m = _block(m)

    (K1, K2) = Generate_Subkey(k)   # 1.

    n = ceil(l/const_Bsize)         # 2.

    if n == 0:                      # 3.
        n = 1
        flag = False
    else:
        if l % const_Bsize == 0:
            flag = True
        else:
            flag = False

    if flag is True:                # 4.
        M_last = (int.from_bytes(m[-1], 'big') ^ int.from_bytes(K1, 'big')).to_bytes(16, 'big')
    else:
        M_last = (int.from_bytes(padding(m[-1]), 'big') ^ int.from_bytes(K2, 'big')).to_bytes(16, 'big')

    X = 0                           # 5.

    for i in range(n-1):            # 6.
        Y = (X ^ int.from_bytes(m[i], 'big')).to_bytes(16, 'big')
        X = int.from_bytes(aes_128(k, Y),'big')

    Y = (X ^ int.from_bytes(M_last, 'big')).to_bytes(16, 'big')
    T = aes_128(k, Y)

    return T                        # 7.

tests_subkey = [
    {'K' :'2b7e151628aed2a6abf7158809cf4f3c',
     'K1':'fbeed618357133667c85e08f7236a8de',
     'K2':'f7ddac306ae266ccf90bc11ee46d513b'}
]


def _block(s, bs=const_Bsize):
    # splits a message in blocks
    if s == b'':
        return [b'']
    l = []
    rs = len(s) % bs
    for i in range(int(len(s)/bs)):
        l.append(s[i*bs:((i+1)*bs)])
    if rs:
        return l + [s[-rs:]]
    else:
        return l