from Crypto.Util.number import bytes_to_long
from Crypto.Util.number import long_to_bytes
from sage.all import GF

x = GF(2)["x"].gen()
gf2e = GF(2 ** 128, name="y", modulus=x ** 128 + x ** 7 + x ** 2 + x + 1)


# Converts an integer to a gf2e element, little endian.
def _to_gf2e(n):
    return gf2e([(n >> i) & 1 for i in range(127, -1, -1)])


# Converts a gf2e element to an integer, little endian.
def _from_gf2e(p):
    n = p.integer_representation()
    ans = n & 1
    for _ in range(127):
        ans <<= 1
        n >>= 1
        ans |= n & 1

    return ans


# Calculates the GHASH polynomial.
def _ghash(h, a, c):
    la = len(a)
    lc = len(c)
    p = 0
    for i in range(la // 16):
        p += _to_gf2e(bytes_to_long(a[16 * i:16 * (i + 1)]))
        p *= h

    if la % 16 != 0:
        p += _to_gf2e(bytes_to_long(a[-(la % 16):] + bytes(16 - la % 16)))
        p *= h

    for i in range(lc // 16):
        p += _to_gf2e(bytes_to_long(c[16 * i:16 * (i + 1)]))
        p *= h

    if lc % 16 != 0:
        p += _to_gf2e(bytes_to_long(c[-(lc % 16):] + bytes(16 - lc % 16)))
        p *= h

    p += _to_gf2e(((8 * la) << 64) | (8 * lc))
    p *= h
    return p


def recover_possible_auth_keys(a1, c1, t1, a2, c2, t2):
    """
    Recovers possible authentication keys from two messages encrypted with the same authentication key.
    More information: Joux A., "Authentication Failures in NIST version of GCM"
    :param a1: the associated data of the first message (bytes)
    :param c1: the ciphertext of the first message (bytes)
    :param t1: the authentication tag of the first message (bytes)
    :param a2: the associated data of the second message (bytes)
    :param c2: the ciphertext of the second message (bytes)
    :param t2: the authentication tag of the second message (bytes)
    :return: a generator generating possible authentication keys (gf2e element)
    """
    h = gf2e["h"].gen()
    p1 = _ghash(h, a1, c1) + _to_gf2e(bytes_to_long(t1))
    p2 = _ghash(h, a2, c2) + _to_gf2e(bytes_to_long(t2))
    for h in (p1 + p2).roots():
        yield h[0]


def forge_tag(h, a, c, t, target_a, target_c):
    """
    Forges an authentication tag for a target message given a message with a known tag.
    This method is best used with the authentication keys generated by the recover_possible_auth_keys method.
    More information: Joux A., "Authentication Failures in NIST version of GCM"
    :param h: the authentication key to use (gf2e element)
    :param a: the associated data of the message with the known tag (bytes)
    :param c: the ciphertext of the message with the known tag (bytes)
    :param t: the known authentication tag (bytes)
    :param target_a: the target associated data (bytes)
    :param target_c: the target ciphertext (bytes)
    :return: the forged authentication tag (bytes)
    """
    ghash = _from_gf2e(_ghash(h, a, c))
    target_ghash = _from_gf2e(_ghash(h, target_a, target_c))
    return long_to_bytes(ghash ^ bytes_to_long(t) ^ target_ghash)
