from struct import pack, unpack
import ucryptolib

MODE_ECB = 1


class AES_CMAC:
    def gen_subkey(self, K):
        AES_128 = ucryptolib.aes(bytearray(K), MODE_ECB)

        L = AES_128.encrypt(('\x00' * 16).encode())

        LHigh = unpack('>Q', L[:8])[0]
        LLow = unpack('>Q', L[8:])[0]

        K1High = ((LHigh << 1) | (LLow >> 63)) & 0xFFFFFFFFFFFFFFFF
        K1Low = (LLow << 1) & 0xFFFFFFFFFFFFFFFF

        if (LHigh >> 63):
            K1Low ^= 0x87

        K2High = ((K1High << 1) | (K1Low >> 63)) & 0xFFFFFFFFFFFFFFFF
        K2Low = ((K1Low << 1)) & 0xFFFFFFFFFFFFFFFF

        if (K1High >> 63):
            K2Low ^= 0x87

        K1 = pack('>QQ', K1High, K1Low)
        K2 = pack('>QQ', K2High, K2Low)

        return K1, K2

    def xor_128(self, N1, N2):
        J = b''
        for i in range(len(N1)):
            J += bytes([N1[i] ^ N2[i]])
        return J

    def pad(self, N):
        const_Bsize = 16
        padLen = 16 - len(N)
        return N + b'\x80' + b'\x00' * (padLen - 1)

    def encode(self, K, M):
        const_Bsize = 16
        const_Zero = b'\x00' * 16

        AES_128 = ucryptolib.aes(bytearray(K), MODE_ECB)

        K1, K2 = self.gen_subkey(K)
        n = int(len(M) / const_Bsize)

        if n == 0:
            n = 1
            flag = False
        else:
            if (len(M) % const_Bsize) == 0:
                flag = True
            else:
                n += 1
                flag = False

        M_n = M[(n - 1) * const_Bsize:]
        if flag is True:
            M_last = self.xor_128(M_n, K1)
        else:
            M_last = self.xor_128(self.pad(M_n), K2)

        X = const_Zero
        for i in range(n - 1):
            M_i = M[(i) * const_Bsize:][:16]
            Y = self.xor_128(X, M_i)
            X = AES_128.encrypt(Y)
        Y = self.xor_128(M_last, X)
        T = AES_128.encrypt(Y)
        return T