import random, struct, hmac
from Crypto.Cipher import AES
import numpy as np

## Computation in GF(2^128)/(x^128 + x^7 + x^2 + x^1 + 1)

class X:
    def __pow__(self, y):
        return 1 << y
x = X()
gcm_modulus = x**128 + x**7 + x**2 + x**1 + 1
gcm_modulus_degree = 128

def gf128_deg(x):
    """Compute the degree of x.
    >>> assert gf128_deg(x**20 + x**5 + 1) == 20
    >>> assert gf128_deg(x**90) == 90
    >>> assert gf128_deg(x**1) == 1
    >>> assert gf128_deg(1) == 0
    >>> assert gf128_deg(0) == -1
    """
    return x.bit_length() - 1

def gf128_add(a, b):
    """Compute a + b.
    >>> assert gf128_add(0, 0) == 0
    >>> assert gf128_add(x**7 + x**2, 0) == x**7 + x**2
    >>> assert gf128_add(x**7 + x**2, x**2 + x**1) == x**7 + x**1
    >>> assert gf128_add(x**7 + x**2, x**7 + x**2) == 0
    """
    return a ^ b

def gf128_mul(a, b):
    """Compute ab.
    >>> assert gf128_mul(x**20 + x**5 + 1, x**13 + x**2) == x**33 + x**22 + x**18 + x**13 + x**7 + x**2
    >>> assert gf128_mul(x**80, x**90) == x**49 + x**44 + x**43 + x**42
    >>> assert gf128_mul(x**23 + x**5, 1) == x**23 + x**5
    """
    if a == 0 or b == 0:
        return 0
    if a > b:
        return gf128_mul(b, a)
    p = 0
    deg_b = gf128_deg(b)
    while a > 0:
        if a & 1:
            p ^= b
        a = a >> 1
        b = b << 1
        deg_b += 1
        if deg_b == gcm_modulus_degree:
            b = b ^ gcm_modulus
            deg_b = gf128_deg(b)
    return p

def gf128_divmod(a, b):
    """Compute q, r such that a = qb + r.
    >>> assert gf128_divmod(x**20 + x**5 + 1, x**13 + x**2) == (x**7, x**9 + x**5 + 1)
    >>> assert gf128_divmod(x**90, x**80) == (x**10, 0)
    """
    q, r = 0, a
    while gf128_deg(r) >= gf128_deg(b):
        d = gf128_deg(r) - gf128_deg(b)
        q = q ^ (1 << d)
        r = r ^ (b << d)
    return q, r

def gf128_mod(a, m=gcm_modulus):
    """Compute b such that a = b; gf128_deg(b) < gf128_deg(m).
    >>> assert gf128_mod(x**220 + x**5 + 1, x**50 + x**4) == x**36 + x**5 + 1
    >>> assert gf128_mod(x**130 + x**64 + x**2) == x**64 + x**9 + x**4 + x**3
    """
    q, r = gf128_divmod(a, m)
    return r

def gf128_egcd(a, b):
    """Compute g, x, y such that g | a, g | b, g = ax + by;
	for any h such that h | a and h | b, gf128_deg(g) >= gf128_deg(h).
    >>> assert gf128_egcd(x**2 + 1, x**1 + 1) == (x**1 + 1, 0, 1)
	>>> assert gf128_egcd(x**12 + x**4, x**5 + 1) == (x**1 + 1, x**3 + x**2 + 1, x**10 + x**9 + x**7 + x**5 + x**4 + x**1 + 1)
    >>> assert gf128_egcd(x**64 + x**2, x**37 + x**12 + 1) == (1, x**35 + x**34 + x**33 + x**31 + x**30 + x**28 + x**27 + x**25 + x**24 + x**21 + x**20 + x**18 + x**17 + x**15 + x**14 + x**12 + x**11 + x**9 + x**7 + x**6 + x**4 + x**3 + x**1 + 1, x**62 + x**61 + x**60 + x**58 + x**57 + x**55 + x**54 + x**52 + x**51 + x**48 + x**47 + x**45 + x**44 + x**42 + x**41 + x**39 + x**38 + x**37 + x**35 + x**34 + x**32 + x**31 + x**29 + x**28 + x**26 + x**25 + x**24 + x**22 + x**21 + x**19 + x**18 + x**16 + x**15 + x**13 + x**12 + x**11 + x**9 + x**8 + x**6 + x**5 + x**3 + x**2 + 1)
    """
    orig_a = a
    orig_b = b
    a_x = 1
    a_y = 0
    b_x = 0
    b_y = 1
    # Invariant = a_x*x + a_y*y = a, b_x*x + b_y*y = b, r_x*r + r_y*y = r
    r_x = 0
    r_y = 1
    if a == 0:
        return b, r_x, r_y
    while True:
        q, r = gf128_divmod(b, a)
        if r == 0:
            # FIX mul so mod isn't necessary; with zeros
            assert a == gf128_mod(gf128_add(gf128_mul(orig_a, r_x), gf128_mul(orig_b, r_y)))
            return a, r_x, r_y
        r_x = gf128_add(b_x, gf128_mul(q, a_x))
        r_y = gf128_add(b_y, gf128_mul(q, a_y))
        a, b = r, a
        b_x, b_y = a_x, a_y
        a_x, a_y = r_x, r_y

def gf128_inv(a):
    """Compute b such that ab = 1.
    >>> assert gf128_inv(x**1) == x**127 + x**6 + x**1 + 1
    >>> assert gf128_inv(x**4) == x**127 + x**125 + x**124 + x**6 + x**4 + x**3 + x**1 + 1
    """
    g, xx, yy = gf128_egcd(a, gcm_modulus)
    assert g == 1
    return xx

def gf128_exp(a, n):
    """Compute a^n.
    >>> assert gf128_exp(x**2 + 1, 0) == 1
    >>> assert gf128_exp(x**2 + 1, 1) == x**2 + 1
    >>> assert gf128_exp(x**2 + 1, 2) == x**4 + 1
    >>> assert gf128_exp(x**2 + 1, 10) == x**20 + x**16 + x**4 + 1
    """
    if n == 0:
        return 1
    rec = gf128_exp(a, n//2)
    rec = gf128_mul(rec, rec)
    if n % 2 == 1:
        rec = gf128_mul(rec, a)
    return rec

## Computation in GF(2^128)[X]/(x^128 + x^7 + x^2 + x^1 + 1)

# Get rid of trailing zeros?
def collapse(x):
    if len(x) == 0:
        return x
    if x[-1] == 0:
        return collapse(x[:-1])
# CAN WE REMOVE MOD?
    return [gfmod(y, m) for y in x]

def gf128polyring_deg(x):
    return len(x) - 1
# assert gf128deg([x**1, x**1, x**1]) == 2
def gf128polyring_add(a, b):
    return collapse([gfadd(x, y) for x, y in zip_longest(a, b, fillvalue=0)])
# assert gf128add([x**1, x**1, x**1, x**1], [x**2, x**3, x**4]) == [x**1+x**2, x**1+x**3, x**1+x**4, x**1]
def gf128polyring_sub(a, b):
    return gf128add(a, b)

def gf128polyring_cmul(c, a):
    return collapse([gfmul(x, c) for x in a])

def gf128polyring_mul(a, b):
    a = a.copy()
    b = b.copy()
    
    p = []
    while gf128deg(a) >= 0:
        p = gf128add(p, gf128cmul(a[0], b))
        a = a[1:]
        b = [0] + b
    return collapse(p)

def gf128polyring_divmod(a, b):
    q, r = [], []
    while gf128deg(a) >= gf128deg(b):
        highA = a[-1]
        highB = b[-1]
        highBinv = 1 if highB == 1 else gfmodinv(highB, m)
        #print(highA, highBinv)
        qq = gfmodmul(highA, highBinv, m)
        #qq, rr = gfdivmod(highA, highB)
        #assert rr == 0
        de = gf128deg(a) - gf128deg(b)
        new_term = [0]*(de+1)
        new_term[de] = qq
        q = gf128add(q, new_term)
        fact = gf128mul(b, new_term)
        a = gf128sub(a, fact)
        
    return collapse(q), collapse(a)

def InverseFrobeniusAutomorphismGF128(x):
    return gfmodexp(x, 2**127, m) # https://math.stackexchange.com/questions/943417/square-root-for-galois-fields-gf2m

def gf128polyring_sqrt(p):
    p = p.copy()
    q = []
    for x in p:
        if x == 0 or x == 1:
            q.append(x)
        else:
            inv = InverseFrobeniusAutomorphismGF128(x)
            q.append(inv)
    #print(q)
    p = q
    for i in range(2,len(p),2):
        p[i//2] = gfadd(p[i//2], p[i])
        p[i] = 0
    return collapse(p)

def gf128polyring_monic(p):
    p = p.copy()
    finalCoeff = p[-1]
    if finalCoeff == 1:
        return p
    inv = gfmodinv(finalCoeff, m)
    return [gfmodmul(x, inv, m) for x in p]

def gf128polyring_formal_derivative(p):
    q = []
    for i, x in enumerate(p):
        if i == 0:
            continue
        if i % 2 == 0:
            e = 0
        else:
            e = x
        q.append(e)
    return collapse(q)

def gf128polyring_sff(f):
    R = []
    assert f[-1] == 1, 'monic'
    fprime = gf128formal_derivative(f)
    #print('fprime', fprime)
    c = gf128gcd(f, fprime)
    #print('c', c)
    w, r = gf128divmod(f, c)
    #print('w', w)
    assert r == []
    i = 1
    while w != [1]:
        y = gf128gcd(w, c)
        #print('y', i, y)
        fac, r = gf128divmod(w, y)
        #print('fac', i, fac)
        assert r == []
        R.append((fac, i))
        w = y
        c, r = gf128divmod(c, y)
        #print('c', i, c, r)
        assert r == []
        i = i + 1
    if c != [1]: 
        #print('sqrt of', ds(c))
        c = gf128sqrt(c)
        c = gf128monic(c)
        #print('rec', ds(c))
        Rrec = gf128sff(c)
        R += [(fac, i*2) for (fac, i) in Rrec]
    return R
    # Frobenius automorphism???

def gf128polyring_modmul(a, b, m):
    c = gf128mul(a, b)
    q, r = gf128divmod(c, m)
    return r
def gf128polyring_modexp(a, n, m):
    if n == 0:
        return [1]
    rec = gf128modexp(a, n//2, m)
    rec2 = gf128modmul(rec, rec, m)
    if n % 2 == 1:
        rec2 = gf128modmul(rec2, a, m)
    return rec2
def ddf(f):
    i = 1
    S = set()
    f_ = f.copy()
    
    while gf128deg(f_) >= 2 * i:
        qq = gf128modexp([0, 1], 2**(128*i), f_)
        qq = gf128add(qq, [0, 1])
        g = gf128gcd(f_, qq)
        #print('gcd', ds(f_), '||', qq, '||', ds(g))
        if g != [1]:
            S.add((tuple(g), i))
            q, r = gf128divmod(f_, g)
            assert r == []
            f_ = q
        i += 1
    if f_ != [1]:
        #print('last')
        S.add((tuple(f_), gf128deg(f_)))
    if len(S) == 0:
        return {(tuple(f), 1)}
    else:
        return S

exp = (2**128-1)//3
def gfrand():
    return random.randint(0, 2**128-1)
def gf128polyring_rand(degree):
    qq= [gfrand() for _ in range(degree+1)]
    return qq
def edf(f, d):
    n = gf128deg(f)
    r = n//d
    S = {tuple(f)}
    while len(S) < r:
        h = gf128rand(n-1)
        g = gf128gcd(h, f)
        if g == [1]:
            g = gf128add(gf128modexp(h, exp, f), [1])
        for u in list(S):
            u = list(u)
            if gf128deg(u) == d:
                continue
            gugcd = gf128gcd(g, u)
            if gugcd != [1] and gugcd != u:
                S.remove(tuple(u))
                S.add(tuple(gugcd))
                qq, rr = gf128divmod(u, gugcd)
                assert rr == []
                S.add(tuple(qq))
        #print("Iteration", S)
    return S

def factorize(f):
    #print('Computing factors over GF(2^128)/[y](x^128+x^7+x^2+x+1) of', ds(f))
    factors = set()
    f = gf128monic(f)
    #print('Reduced to monic polynomial', ds(f))
    fs = gf128sff(f)
    for p, _ in fs:
        #print('Computed square-free factor', ds(p))
        qs = ddf(p)
        for q, d in qs:
            #print('Computed distinct-degree factor', ds(q), 'of degree', d)
            rs = edf(list(q), d)
            #for r in rs:
            #   print('Computed equal-degree factor', ds(r))
            factors |= rs
    return factors

def monic_roots(factors):
    roots = []
    for factor in factors:
        if gf128deg(factor) == 1:
            roots.append(factor[0])
    return roots

## AES-GCM

def reverse_mask(b):
    # from https://stackoverflow.com/a/2602885
    b = (b & 0xF0) >> 4 | (b & 0x0F) << 4;
    b = (b & 0xCC) >> 2 | (b & 0x33) << 2;
    b = (b & 0xAA) >> 1 | (b & 0x55) << 1;
    return b;

def bytes_to_gf128(xs, st=None, en=None):
    if st is None and en is None:
        st = 0
        en = 16
    a, b = struct.unpack('<QQ', bytes(reverse_mask(xs[i]) for i in range(st, en)))
    return a + (b << 64)

def gf128_to_bytes(n):
    s = b''
    while n != 0:
        c = n & 0xff
        f = 0
        for i in range(8):
            if c & (1 << i):
                f += (1 << (7-i))
        s += int.to_bytes(f, 1, 'big')
        n >>= 8
    s += b'\x00' * (16 - len(s))
    assert len(s) == 16
    return s

def gmac(h, s, aad, c):
    assert len(n) == 12
    blocks = build_blockseq(ad, c)
    g = 0
    for i in range(len(blocks)//16):
        block = bytes_to_gf128(blocks, i*16, (i + 1)*16)
        g = gf128_add(g, block)
        g = gf128_mul(g, h)
    t = gf128_add(g, s)
    return gf128_to_bytes(t)

def ecb_encrypt(key, xs):
    c = AES.new(key, AES.MODE_ECB)
    return c.encrypt(xs)

def xor(a, b):
    return bytes([(x^y) for (x, y) in zip(a, b)])

def gctr(k, nonce, x):
    y = b''
    ctr = (int.from_bytes(nonce, 'big') << 32) + 2
    i = 0
    while len(y) < len(x):
        c += xor(x[i*16:(i + 1)*16], ecb_encrypt(k, ctr.to_bytes(16, 'big')))
        ctr += 1
        i += 1
    return y

def authentication_key(k):
    return bytes_to_gf128(ecb_encrypt(k, b'\x00'*16))

def blind(k, n):
    return bytes_to_gf128(ecb_encrypt(k, nonce + b'\x00\x00\x00\x01'))

def gcm_encrypt(k, nonce, aad, m, mac_bytes=16):
    c = gctr(k, nonce, m)
    mac = gmac(authentication_key(k), blind(k, nonce), aad, c)
    return c, mac[:mac_bytes]

def gcm_decrypt(k, nonce, aad, c, mac, mac_bytes=16):
    m = gctr(k, nonce, c)
    expected_mac = gmac(authentication_key(k), blind(k, nonce), aad, c)[:mac_bytes]
    assert hmac.compare_digest(mac, exp_mac), f'Expected MAC {exp_mac.hex()}, got {mac.hex()}'
    return m

## Forbidden Attack

def compute_polydiff(ad1, ct1, mac1, ad2, ct2, mac2):
    bs1 = build_blockseq(ad1, ct1)
    bs2 = build_blockseq(ad2, ct2)
    if len(bs1) < len(bs2):
        bs1 = b'\x00'*(len(bs2)-len(bs1)) + bs1
    else:
        bs2 = b'\x00'*(len(bs1)-len(bs2)) + bs2
    assert len(bs1) == len(bs2)
    #print(bs1.hex())
    #print(bs2.hex())
    f = []
    N = len(bs1)//16
    for i in range(N):
        b1 = binToGF128(bs1[i*16:(i+1)*16])
        b2 = binToGF128(bs2[i*16:(i+1)*16])
        c = gfadd(b1, b2)
        f.append(c)
    mac1gf = binToGF128(mac1)
    mac2gf = binToGF128(mac2)
    macSgf = gfadd(mac1gf, mac2gf)
    f.append(macSgf)
    return list(reversed(f))

def forge_gcm_tag(h, origAd, origCt, newAd, newCt, tag, n):
    # First recover s for the nonce
    tag_ = gmac(h, 0, origAd, origCt, n)
    s = gfadd(binToGF128(tag), binToGF128(tag_))
    tag = gmac(h, s, newAd, newCt, n)
    return tag
def gcm_repeatednonce_attack(factors, n, ad1, ct1, mac1, ad2, ct2, mac2):
    #f=compute_polydiff(ad1, ct1, mac1, ad2, ct2, mac2)
    #factors = factorize(f)
    roots = monic_roots(factors)
    h = None
    tags = []
    for root in roots:
        forged_tag = forge_gcm_tag(root, ad1, ct1, ad2, ct2, mac1, n)
        tags.append(forged_tag)
    return tags
# tag_guesses = gcm_repeatednonce_attack(factors, n, ad1, ct1, mac1, ad2, ct2, mac2)


def d(n):
    if n == 0:
        return '0'
    s = []
    k = 0
    while n > 0:
        if n & 1 == 0:
            k += 1
            n >>= 1
            continue
        if k == 0:
            s.append('1')
        elif k == 1:
            s.append('x')
        else:
            s.append('x^' + str(k))
        n >>= 1
        k += 1
    return ' + '.join(reversed(s))