import re, random, functools, base64, os, itertools
from say_my_name import HeisenbergMatrix, HeisenbergGroup
from __params import N, M, L

from Crypto.Util.number import *
from Crypto.Cipher import AES
from random import randint as ri

pad = lambda s, l: s + bytes([ri(0,31) * bool(i) for i in range(l - len(s))])
unpad = lambda s: s[:s.index(b'\x00')] if b'\x00' in s else s
prod = lambda arr: functools.reduce(lambda a, b: a * b, arr)


def conjugate(w1, w2):
    return w2.inverse() * w1 * w2

def commutator(w1, w2):
    return w1.inverse() * w2.inverse() * w1 * w2

def encode(s):
    return base64.b64encode(s.encode()).decode()

def decode(s):
    return base64.b64decode(s.encode()).decode()

def parsearr(s):
    return [HeisenbergMatrix.fromstr(x) for x in decode(s).split('|')]

def parse(s):
    return HeisenbergMatrix.fromstr(decode(s))

def group(s):
    return encode('|'.join(map(repr, s)))


def load_data():
    global set_A, set_B, pub_A, pub_B, enc
    directory = os.path.join(".", "intercepted_data")
    with open(os.path.join(directory, "set_A.txt"), 'r') as f:
        set_A = parsearr(f.read())
    with open(os.path.join(directory, "set_B.txt"), 'r') as f:
        set_B= parsearr(f.read())
    with open(os.path.join(directory, "pub_A.txt"), 'r') as f:
        pub_A = parsearr(f.read())
    with open(os.path.join(directory, "pub_B.txt"), 'r') as f:
        pub_B = parsearr(f.read())
    with open(os.path.join(".", "encrypted"), 'rb') as f:
        enc = f.read()

def oracle(key):
    key = key.c&(2**128 - 1)
    aes = AES.new(long_to_bytes(key, 16), AES.MODE_ECB)
    dec = unpad(aes.decrypt(enc))
    return b'Dear Bob,' in dec, dec

def test(priv_set_A, k_A):
    L_A = []
    priv_A = prod([word ** k for word, k in zip(priv_set_A, k_A)])
    for word, k in zip(priv_set_A, k_A):
        i = set_A.index(word)
        w = pub_B[i] ** k
        L_A.append(w)
    K_A = priv_A * prod(L_A).inverse()
    return K_A, oracle(K_A)

def bruteforce():
    for idx_A in itertools.permutations(range(M), L):
        for k_A in itertools.product([1, -1], repeat = L):
            print('\rProgress: %s %s' % (idx_A, k_A), end = '')
            priv_set_A = list(map(set_A.__getitem__, idx_A))
            result, (succ, dec) = test(priv_set_A, k_A)
            if succ:
                return result, dec

load_data()

key, dec = bruteforce()
print('\n\n')
if key:
    print("Common Key Found: \033[32m%s\033[0m" % key)
    print()
    print("Message:\n\033[32m%s\033[0m" % dec.decode(encoding = 'utf-8', errors = 'ignore'))
else:
    print("\033[31mFailed\033[0m")