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")