import chipwhisperer as cw
import time

# Path to the Speck .hex file for reflashing
PATH="/home/msc/documents/obsidian_notes/master-aits/subjects/implementation_attacks_and_countermeasures/praktikum/speck_cpa_cw/cw_firmware/simple-speck-CWLITEARM.hex"
#PATH="/home/juan/documents/master-aits/subjects/implementation_attacks_and_countermeasures/praktikum/speck_cpa_cw/cw_firmware_masked/simple-speck-CWLITEARM.hex"
PATH="/home/msc/documents/obsidian_notes/master-aits/subjects/implementation_attacks_and_countermeasures/praktikum/speck_cpa_cw/cw_firmware_hiding/simple-speck-CWLITEARM.hex"

scope.dis()

def flash(scope, prog):
    cw.program_target(scope, prog, PATH)

def reset_target(scope):
    scope.io.nrst = 'low'
    time.sleep(0.05)
    scope.io.nrst = 'high_z'
    time.sleep(0.05)

try:
    if not scope.connectStatus:
        scope.con()
except NameError:
    scope = cw.scope()

try:
    target = cw.target(scope)
except IOError:
    print("INFO: Caught exception on reconnecting to target - attempting to reconnect to scope first.")
    print("INFO: This is a work-around when USB has died without Python knowing. Ignore errors above this line.")
    scope = cw.scope()
    target = cw.target(scope)

print("INFO: Found ChipWhisperer😍")

prog = cw.programmers.STM32FProgrammer
time.sleep(0.05)
scope.default_setup()

reset_target(scope)

flash(scope, prog)

# 32 bytes of encryption key
#encryption_key = b"\x00\x01\x02\x03\x04\x05\x06\x07\x08\x09\x0a\x0b\x0c\x0d\x0e\x0f\x10\x11\x12\x13\x14\x15\x16\x17\x18\x19\x1a\x1b\x1c\x1d\x1e\x1f"

# 8 byte encryption key
encryption_key = b"\x11\x22\x33\x44\x55\x66\x77\x88"
if len(encryption_key) == 8:
    target.simpleserial_write("s", encryption_key)

target.simpleserial_write("k", b'\x00'*4)
print(target.simpleserial_read("o", 8))

#pt = b"\x70\x6f\x6f\x6e\x65\x72\x2e\x20\x49\x6e\x20\x74\x68\x6f\x73\x65"
pt = b"\x4c\x69\x74\x65"
target.simpleserial_write("e", pt)
print(target.simpleserial_read("c", 4))

from tqdm.notebook import trange
import random
import numpy as np

ktp = cw.ktp.Basic()
trace_array = []
textin_array = []

pt = b"\x4c\x69\x74\x65" 
random.seed(0x5222223) 


N = 800
for i in trange(N, desc='Capturing traces'):
    pt = bytes([random.randint(0, 255) for i in range(4)])
    scope.arm()
    
    
    target.simpleserial_write('e', pt)
    
    ret = scope.capture()
    if ret:
        print("Target timed out!")
        continue
    
    response = target.simpleserial_read('c', 4)
    
    trace_array.append(scope.get_last_trace())
    textin_array.append(pt)

    
trace_array = np.array(trace_array)

np.save("sample_traces/2200_encryption_traces.npy", trace_array)
np.save("sample_traces/2200_plaintext_traces.npy", textin_array)

import numpy as np
trace_array = np.load("sample_traces/2200_encryption_traces.npy")
textin_array = np.load("sample_traces/2200_plaintext_traces.npy")

%matplotlib notebook
import matplotlib.pylab as plt
plt.figure()
plt.plot(trace_array[2], 'orange')
#plt.plot(trace_array[4], 'g')
plt.show()

#Hamming Weight
HW = [bin(n).count("1") for n in range(0, 256)]
def popcount(x):
    x -= (x >> 1) & 0x5555555555555555
    x = (x & 0x3333333333333333) + ((x >> 2) & 0x3333333333333333)
    x = (x + (x >> 4)) & 0x0f0f0f0f0f0f0f0f
    return ((x * 0x0101010101010101) & 0xffffffffffffffff ) >> 56

def mean(X):
    return np.sum(X, axis=0)/len(X)

def std_dev(X, X_bar):
    return np.sqrt(np.sum((X-X_bar)**2, axis=0))

def cov(X, X_bar, Y, Y_bar):
    return np.sum((X-X_bar)*(Y-Y_bar), axis=0)

def to16(byte1, byte2):
    return int((byte1 << 8) + byte2)

import math

NUM_ROUNDS = 22
BLOCK_SIZE = 32
KEY_SIZE = 64
WORD_SIZE = 16


# SHIFTs for SPECK
ALPHA = 7
BETA = 2

mod_mask = (2 ** WORD_SIZE) -1
mod_mask_sub = (2 ** WORD_SIZE)

def ER16(x, y, k):

    rs_x = ((x << (16 - ALPHA)) + (x >> ALPHA)) & mod_mask

    add_sxy = (rs_x + y) & mod_mask

    new_x = k ^ add_sxy

    ls_y = ((y >> (16 - BETA)) + (y << BETA)) & mod_mask

    new_y = new_x ^ ls_y

    return new_x, new_y



def simple_speck(plaintext, key, arg=None):
    Ct_0 = (int(plaintext[1]) << 8) + int(plaintext[0])
    Ct_1 = (int(plaintext[3]) << 8) + int(plaintext[2])
                                                
    
    Ct_1, Ct_0 = ER16(Ct_1, Ct_0, key)   # fixed 16 bit key of 0x55
    return popcount((Ct_1 << 8) + Ct_0)

## According to the Paper "Breaking Speck using CPA"
def PowerRightHalfKey(plaintext, key, knownkey=None):
    y = (int(plaintext[1]) << 8) + int(plaintext[0])
    x = (int(plaintext[3]) << 8) + int(plaintext[2])
    
    x = ((x << (16 - ALPHA)) + (x >> ALPHA)) & mod_mask
    x = (x + y) & mod_mask
    x = x ^ key
    
    return popcount(x)

## According to the Paper "Breaking Speck using CPA" 
def PowerLeftHalfKeyII(plaintext, key, knownkey):
    y = (int(plaintext[1]) << 8) + int(plaintext[0])
    x = (int(plaintext[3]) << 8) + int(plaintext[2])
       
    
    # -------------- for one key -----------------#
    x = ((x << (16 - ALPHA)) + (x >> ALPHA)) & mod_mask           # x = ROR(x, 7)
    x = (x + y) & mod_mask                                        # x = ADD(x, y)
    
    if knownkey == None or len(knownkey) == 0:
        x = x ^ key                                               # x = XOR(x, k)
        return popcount(x)
    else:
        x = x ^ knownkey[0]   
    
    
    # -------------- for second key -----------------#
    
    
    y = ((y >> (16 - BETA)) + (y << BETA)) & mod_mask            # y = ROL(y, 2)
    y = y ^ x                                                    # y = XOR(y, x)
    x = ((x << (16 - ALPHA)) + (x >> ALPHA)) & mod_mask          # x = ROR(x, 7)
    x = (x + y) & mod_mask                                       # x = ADD(x, y)
    
    
    if len(knownkey) == 1:
        x = x ^ key                                             # x = XOR(x, k)
        return popcount(x)
    else:
        x = x ^ knownkey[1]                                     # x = XOR(x, k) 
                                            
    
    # -------------- for third key -----------------#

    y = ((y >> (16 - BETA)) + (y << BETA)) & mod_mask            # y = ROL(y, 2)
    y = y ^ x                                                    # y = XOR(y, x)
    x = ((x << (16 - ALPHA)) + (x >> ALPHA)) & mod_mask          # x = ROR(x, 7)
    x = (x + y) & mod_mask                                       # x = ADD(x, y)
    
    if len(knownkey) == 2:
        x = x ^ key                                              # x = XOR(x, k)
    else:
        x =  x ^ knownkey[2]                                     # x = XOR(x, k)
        
    y = ((y >> (16 - BETA)) + (y << BETA)) & mod_mask            # y = ROL(y, 2)
    y = y ^ x                                                    # y = XOR(y, x)
    
    if len(knownkey) == 2:                                       
        return popcount(y)   
    
    
    # -------------- for fourth key -----------------#
    
    x = ((x << (16 - ALPHA)) + (x >> ALPHA)) & mod_mask          # x = ROR(x, 7)
    x = (x + y) & mod_mask                                       # x = ADD(x, y)
    
    if len(knownkey) == 3:
        x = x ^ key                                              # x = XOR(x, k)
    else:
        x =  x ^ knownkey[3]                                     # x = XOR(x, k)
        
    y = ((y >> (16 - BETA)) + (y << BETA)) & mod_mask            # y = ROL(y, 2)
    y = y ^ x                                                    # y = XOR(y, x)
 
    if len(knownkey) == 3:                                       
        return popcount(y)   

    # -------------- for fith key -----------------#
    x = ((x << (16 - ALPHA)) + (x >> ALPHA)) & mod_mask          # x = ROR(x, 7)
    x = (x + y) & mod_mask                                       # x = ADD(x, y)
    
    x =  x ^ key                                                 # x = XOR(x, k)
    y = ((y >> (16 - BETA)) + (y << BETA)) & mod_mask            # y = ROL(y, 2)
    y = y ^ x                                                    # y = XOR(y, x)
    
    if len(knownkey) == 4:                                       # x = XOR(x, k)
        return popcount(y)   

    return popcount(y)
    
    
    ## According to the Paper "Breaking Speck using CPA" 
def PowerLeftHalfKey(plaintext, key, knownkey):
    pt2 = (int(plaintext[1]) << 8) + int(plaintext[0])
    pt1 = (int(plaintext[3]) << 8) + int(plaintext[2])
    
    temp = ((pt1 << (16 - ALPHA)) + (pt1 >> ALPHA)) & mod_mask
    p1 = (temp + pt2) & mod_mask
    
    r1 = p1 ^ knownkey
    
    temp = ((pt2 >> (16 - BETA)) + (pt2 << BETA)) & mod_mask
    s1 = temp ^ r1
    temp = ((r1 << (16 - ALPHA)) + (r1 >> ALPHA)) & mod_mask
    p2 = (temp + s1) & mod_mask
    
    
    intermediate = (p2) ^ key
    
    
    
    return popcount(intermediate)

def simple_speck_partial(plaintext, key, knownkey):
    Ct_0 = (int(plaintext[1]) << 8) + int(plaintext[0])
    Ct_1 = (int(plaintext[3]) << 8) + int(plaintext[2])
                                                
    Ct_1, Ct_0 = ER16(Ct_1, Ct_0, knownkey)
    

    Ct_1, Ct_0 = ER16(Ct_1, Ct_0, key)   # fixed 16 bit key of 0x55
    return popcount((Ct_1 << 8) + Ct_0)


def speck_keyschedule(plaintext, key, known_keys):
    Ct_0 = (int(plaintext[1]) << 8) + int(plaintext[0])
    Ct_1 = (int(plaintext[3]) << 8) + int(plaintext[2])
    
    for known_key in known_keys:
        Ct_1, Ct_0 = ER16(Ct_1, Ct_0, known_key)    

    Ct_1, Ct_0 = ER16(Ct_1, Ct_0, key)   # fixed 16 bit key of 0x55
    return popcount((Ct_1 << 8) + Ct_0)

   
    

simple_speck(b'\xf5\xf9\xa97', 0x1)

def calculate_correlations_plot(traces, plaintexts, model_callback, leftmost=True, other_keybyte=0x00, argument=None):
    
    maxcpa = [0] * 256   # Correlations
    
    # Calculate mean and standard derivation
    t_bar = mean(traces) 
    o_t = std_dev(traces, t_bar)
    
    
    for key in range(0, 256):
        
        # Run the model
        if leftmost:
            hws = np.array([[model_callback(pt, (key << 8) + other_keybyte, argument) for pt in plaintexts]]).transpose()
        elif not leftmost:
            hws = np.array([[model_callback(pt, (other_keybyte << 8) + key, argument) for pt in plaintexts]]).transpose()
        else:
            raise Exception("[-] Invalid Key Position")
            
        
        hws_bar = mean(hws)
        o_hws = std_dev(hws, hws_bar)
        correlation = cov(traces, t_bar, hws, hws_bar)
        cpaoutput = correlation/(o_t*o_hws)
        maxcpa[key] = max(abs(cpaoutput))
    
    # Return the two best guesses
    best_guess = int(np.argmax(maxcpa))
    second_guess = int(np.argsort(maxcpa, axis=0)[-2])
    
    return (maxcpa, best_guess)

def calculate_correlations(traces, plaintexts, model_callback, leftmost=True, other_keybyte=0x00, argument=None):
    
    maxcpa = [0] * 256   # Correlations
    
    # Calculate mean and standard derivation
    t_bar = mean(traces) 
    o_t = std_dev(traces, t_bar)
    
    
    for key in range(0, 256):
        
        # Run the model
        if leftmost:
            hws = np.array([[model_callback(pt, (key << 8) + other_keybyte, argument) for pt in plaintexts]]).transpose()
        elif not leftmost:
            hws = np.array([[model_callback(pt, (other_keybyte << 8) + key, argument) for pt in plaintexts]]).transpose()
        else:
            raise Exception("[-] Invalid Key Position")
            
        
        hws_bar = mean(hws)
        o_hws = std_dev(hws, hws_bar)
        correlation = cov(traces, t_bar, hws, hws_bar)
        cpaoutput = correlation/(o_t*o_hws)
        maxcpa[key] = max(abs(cpaoutput))
    
    # Return the two best guesses
    best_guess = int(np.argmax(maxcpa))
    second_guess = int(np.argsort(maxcpa, axis=0)[-2])
    
    return ([best_guess, maxcpa[best_guess]], [second_guess, maxcpa[second_guess]])
    

from tqdm import tnrange


def calc_mean_from_trace(traces, plaintexts):

    maxcpa = [0] * 256

    t_bar = mean(traces) 
    o_t = std_dev(traces, t_bar)

    for key in range(0, 256):
        
        hws = np.array([[simple_speck(textin, (key << 8) + 0x00) for textin in textin_array]]).transpose()
        
        # The following line works for a one byte key
        #hws = np.array([[simple_speck(textin, key) for textin in textin_array]]).transpose()
        

        hws_bar = mean(hws)
        o_hws = std_dev(hws, hws_bar)
        correlation = cov(traces, t_bar, hws, hws_bar)
        cpaoutput = correlation/(o_t*o_hws)
        maxcpa[key] = max(abs(cpaoutput))
    
    plt.figure()
    plt.plot(maxcpa, 'orange')
    plt.show()
    guess = np.argmax(maxcpa)
    print(f"Key guess: (xored with) = ", hex(guess))
    return guess

calc_mean_from_trace(trace_array, textin_array)

def calc_mean_from_trace_two_byte_key(traces, plaintexts):

    best, second = calculate_correlations(traces, plaintexts, simple_speck, True, 0x00)
    print(f"[+] Highest Correlation: {hex(best[0])} ({best[1]}) and Second Highest: {hex(second[0])} ({second[1]})")
    sbest, ssecond = calculate_correlations(traces, plaintexts, simple_speck, False, best[0])
    print(f"[+] Highest Correlation: {hex(sbest[0])} ({sbest[1]}) and Second Highest: {hex(ssecond[0])} ({ssecond[1]})")
    
    
    return to16(best[0], sbest[0])

possible_firstkeys = calc_mean_from_trace_two_byte_key(trace_array, textin_array)

def calc_mean_from_trace_two_byte_key_plot(traces, plaintexts):

    plt.figure()

    corr, best = calculate_correlations_plot(traces, plaintexts, simple_speck, True, 0x00)
    scorr, second = calculate_correlations_plot(traces, plaintexts, simple_speck, False, best)
 
    rk = plt.plot(corr, 'orange', label="Left Roundkey")
    lk = plt.plot(scorr, 'blue', label="Right Roundkey")
    plt.ylabel('Correlation')
    plt.xlabel('Key Byte')

    plt.legend(loc="upper left")
    plt.show()
    

calc_mean_from_trace_two_byte_key_plot(trace_array, textin_array)

from tqdm import tnrange


def get_first_keybyte(traces, plaintexts):

    best, second = calculate_correlations(traces, plaintexts, PowerRightHalfKey, False, 0x00)
    print(f"[+] Highest Correlation: {hex(best[0])} ({best[1]}) and Second Highest: {hex(second[0])} ({second[1]})")
    sbest, ssecond = calculate_correlations(traces, plaintexts, PowerRightHalfKey, True, best[0])
    print(f"[+] Highest Correlation: {hex(sbest[0])} ({sbest[1]}) and Second Highest: {hex(ssecond[0])} ({ssecond[1]})")
    
    print(f"Key guess: (xored with) = ", hex(best[0]))
    return to16(sbest[0], best[0])

get_first_keybyte(trace_array, textin_array)

from tqdm import tnrange


def get_keybytel(traces):

    maxcpa = [0] * 256

    t_bar = mean(traces) 
    o_t = std_dev(traces, t_bar)

    for key in range(0, 256):
        
        hws = np.array([[speck_keyschedule(textin, (0x00 << 8) + key, [0x2211, 0x00dd, 0xa8dc]) for textin in textin_array]]).transpose()
        
        # The following line works for a one byte key
        #hws = np.array([[simple_speck(textin, key) for textin in textin_array]]).transpose()
        

        hws_bar = mean(hws)
        o_hws = std_dev(hws, hws_bar)
        correlation = cov(traces, t_bar, hws, hws_bar)
        cpaoutput = correlation/(o_t*o_hws)
        maxcpa[key] = max(abs(cpaoutput))
    
    plt.figure()
    plt.plot(maxcpa, 'orange')
    plt.show()
    guess = np.argmax(maxcpa)
    second = np.argsort(maxcpa, axis=0)[-2]
    print(f"Key guess: (xored with) = ", hex(guess))
    return maxcpa

cpa2 = get_keybytel(trace_array)

def recv_roundkey(traces, plaintexts, roundkeys):

    best, second = calculate_correlations(traces, plaintexts, speck_keyschedule, False, 0x00, roundkeys)
    print(f"[+] Highest Correlation: {hex(best[0])} ({best[1]}) and Second Highest: {hex(second[0])} ({second[1]})")
    sbest, ssecond = calculate_correlations(traces, plaintexts, speck_keyschedule, True, best[0], roundkeys)
    print(f"[+] Highest Correlation: {hex(sbest[0])} ({sbest[1]}) and Second Highest: {hex(ssecond[0])} ({ssecond[1]})")
    
    return to16(sbest[0], best[0])
    

'''
    This seems to work most reliable (currently, only one byte is wrong....)
'''
def recv_roundkey_two(traces, plaintexts, roundkeys):
    best, second = calculate_correlations(traces, plaintexts, PowerLeftHalfKeyII, False, 0x00, roundkeys)
    print(f"[+] Highest Correlation: {hex(best[0])} ({best[1]}) and Second Highest: {hex(second[0])} ({second[1]})")
    sbest, ssecond = calculate_correlations(traces, plaintexts, PowerLeftHalfKeyII, True, best[0], roundkeys)
    print(f"[+] Highest Correlation: {hex(sbest[0])} ({sbest[1]}) and Second Highest: {hex(ssecond[0])} ({ssecond[1]})")
    
    return to16(sbest[0], best[0])
    
       

a = recv_roundkey_two(trace_array, textin_array, [0x2211, 0x00dd])
print(hex(a))

def get_roundkey(traces, firstkey):
    initial_key = [firstkey]
    
    for i in range(4):
        next_rk = recv_roundkey_two(traces, textin_array, initial_key)    
        initial_key.append(next_rk)
    print([hex(k) for k in initial_key])
    return initial_key

def get_key_from_roundkey(roundkey):
    if len(roundkey) != 5:
        raise Exception("Wrong length of roundkey")
    key = [roundkey[0]]
    for i in range(4):
        for keybyte in range(2**16):
            _, out = ER16(keybyte, known_keys[i], i)
            if out == roundkey[i+1]:
                key.append(keybyte)
                print(f"Found: {hex(keybyte)}")
    return [hex(k) for k in key]

rk = get_roundkey(trace_array, int(0x2211))  # get the roundkeys, based on the first roundkey

get_key_from_roundkey(rk)

# These are the roundkeys that __should__ be there for the key
known_keys = [0x2211, 0x00dd, 0xa8dc, 0x349c, 0xb5de]

get_key_from_roundkey(known_keys)

get_key_from_roundkey([8721, 221, 43228, 13468, 18977])

x = [
        "kb1",
        "kb2",
        "kb3",
        "kb4",
        "kb5",
        "kb6",
        "kb7",
        "kb8",
        "kb9",
        "kb10",
]


y = [
    0.41,
    0.41,
    0.92,
    0.92,
    0.54,
    0.84,
    0.62,
    0.79,
    0.49,
    0.85,
    ]

%matplotlib notebook
import matplotlib.pylab as plt
plt.figure()
plt.bar(x, y, color="orange")
plt.show()