'''
Created on May 23, 2020

@author: cvkennedy
'''

from itertools import combinations
import matplotlib.pyplot as plt
import numpy as np

def find_outliers(corpus, num_samps, i, j):
    idxs = corpus[i][j].argsort()[:num_samps]
    return idxs
    
def guess_bytes(corpus, known_keybytes, num_samps, avg):
    candidates = []
    for base in range(4):
        family = [base, base + 4, base + 8, base + 12]
        for combo in combinations(family, 2):
            i,j = combo
            guesses = find_outliers(corpus, num_samps, i, j)
            guesses2 = []
            for guess in guesses:
                cnt = corpus[i][j][guess]
                if cnt-avg < -10:
                    guesses2.append((i, j, guess, cnt-avg))
                print(i, j, guess, cnt - avg)
            candidates.append(tuple(guesses2))
    print(candidates)
            
if __name__ == '__main__':
    known_keybytes = bytes.fromhex("64c7072487f2")
    secret_data = ("c1a5fe7beb2c70bfab98926627dcff8b9671edc5" +
                   "2441f89fa47797aa023f15f67907ee837b93cd9b" +
                   "194922ebb7c3ca3bd1cbfbc888efe147e8055404" +
                   "7d82872fcee564c1bfd2e0a809568acb5cc08f48" +
                   "36a5f91f43b576a4ee1c6f097c15e1cd4056917f" +
                   "c51c1e5d8157409b11f1600d")
    
    data = set()
    with open("test.txt", "r") as fp:
        for line in fp:
            pt, timing = line.strip().split(',')
            pt = bytes.fromhex(pt)
            timing = int(timing)
            data.add((pt, timing))
            
    tavg = sum((d[1] for d in data)) / len(data)
    print("tavg: %d" % tavg)
    
    known_tly = np.zeros((16, 16, 256))
    
    for base in range(4):
        print("Building corpus for family %d" % base)
        family = [base, base + 4, base + 8, base + 12]
        for combo in combinations(family, 2):
            times = np.zeros(256)
            counts = np.zeros(256)
            i,j = combo
            print("Working on %d, %d" % (i, j))
            for d in data:
                n = d[0][i] ^ d[0][j]
                c = d[1]
                times[n] += c
                counts[n] += 1
            for c in range(256):
                cnorm = times[c] / counts[c]
                known_tly[i][j][c] = cnorm
                known_tly[j][i][c] = cnorm
                
    guess_bytes(known_tly, known_keybytes, 4, tavg)