from Crypto.Cipher import AES
from Crypto.Util import Counter
from Crypto.Util.number import long_to_bytes, bytes_to_long
from Utils import *
def gf_2_128_mul(x, y):
    assert x < (1 << 128)
    assert y < (1 << 128)
    res = 0
    for i in range(127, -1, -1):
        res ^= x * ((y >> i) & 1)  # branchless
        x = (x >> 1) ^ ((x & 1) * 0xE1000000000000000000000000000000)
    assert res < 1 << 128
    return res


class InvalidInputException(Exception):
    def __init__(self, msg):
        self.msg = msg

    def __str__(self):
        return str(self.msg)


class InvalidTagException(Exception):
    def __str__(self):
        return 'The authenticaiton tag is invalid.'


# Galois/Counter Mode with AES-128 and 96-bit IV

class GhashCry:
    def __init__(self, master_key, init_value):
        print(type(master_key))
        print(type(init_value))
        if master_key >= (1 << 128):
            raise InvalidInputException('Master key should be 128-bit')
        self.init_value = init_value
        self.__master_key = long_to_bytes(master_key, 16)
        self.__aes_ecb = AES.new(self.__master_key, AES.MODE_ECB)
        self.__auth_key = bytes_to_long(self.__aes_ecb.encrypt(b'\x00' * 16))
        self.__last_key = bytes_to_long(self.__aes_ecb.encrypt(
            long_to_bytes((init_value << 32) | 1, 16)))

    def get(self):
        return self.__auth_key, self.__last_key

    def encIV(self, n):
        res = []
        print(long_to_bytes(self.init_value, 16))
        print(long_to_bytes(self.init_value << 32, 16))
        for i in range(n):
            print(long_to_bytes((self.init_value << 32) + i + 2, 16).hex())
            a = long_to_bytes((self.init_value << 32) + i + 2, 16)
            simplePrint("iv", a.hex())
            r = self.__aes_ecb.encrypt(a)
            s = "ENC iv "
            simplePrint(s, r.hex())
            res.insert(i, r)

        print(res)
        return res


class GhashCon:
    def __init__(self, __auth_key, __last_key):
        self.change_key(__auth_key, __last_key)

    def change_key(self, __auth_key, __last_key):
        self.__auth_key = __auth_key
        self.__last_key = __last_key
        # precompute the table for multiplication in finite field
        table = []  # for 8-bit
        for i in range(16):
            row = []
            for j in range(256):
                row.append(gf_2_128_mul(self.__auth_key, j << (8 * i)))
            table.append(tuple(row))
        self.__pre_table = tuple(table)
        # print(f"Ghash_0 {tuple(table)}")

        self.prev_init_value = None  # reset

    def __times_auth_key(self, val):
        res = 0
        for i in range(16):
            res ^= self.__pre_table[i][val & 0xFF]
            val >>= 8
        return res

    def ghash(self, aad, txt):
        len_aad = len(aad)
        len_txt = len(txt)

        # padding
        print(aad)
        if 0 == len_aad % 16:
            data = aad
        else:
            data = aad + b'\x00' * (16 - len_aad % 16)
        print(data)
        if 0 == len_txt % 16:
            data += txt
        else:
            data += txt + b'\x00' * (16 - len_txt % 16)
        # print(data)
        tag = 0
        assert len(data) % 16 == 0
        print(len(data))
        print(len(data) // 16)
        for i in range(len(data) // 16):
            print("tag\t", hex(tag))
            print("data ", i, "\t", data[i * 16: (i + 1) * 16].hex())
            tag ^= bytes_to_long(data[i * 16: (i + 1) * 16])
            print("data XOR tag \t", hex(tag))
            tag = self.__times_auth_key(tag)
            print("__times_auth_key", tag,  "\n\n")
        # print 'X\t', hex(tag)
        tag ^= ((8 * len_aad) << 64) | (8 * len_txt)
        print("XOR ahib \t", ((8 * len_aad) << 64) | (8 * len_txt))
        print("tag \t", tag)

        tag = self.__times_auth_key(tag)
        print("salam man ejra shodam", tag, "\n\n")
        print("salam man ejra shodam", tag.to_bytes(16, "big").hex(), "\n\n")
        return tag ^ self.__last_key