import chipwhisperer as cw
import time, os

KEY=bytes(list(x for x in range(16)))
delay1list=[104,103,118,117]
delay2abs =503
glitchlen=[5,5]

def reset_target(scope, target):
    scope.default_setup()
    scope.io.target_pwr = 0
    time.sleep(0.2)
    scope.io.target_pwr = 1
    time.sleep(0.2)
    # send key and random nonce
    target.simpleserial_write('k', KEY)
    target.simpleserial_wait_ack()
    target.simpleserial_write('n', os.urandom(8)+b'\0'*8)
    target.simpleserial_wait_ack()

def init_cw():
    scope = cw.scope()
    target_type = cw.targets.SimpleSerial
    target = cw.target(scope, target_type)
    print("Connected...")
    time.sleep(0.05)
    scope.default_setup()
    reset_target(scope, target)
    cw.set_all_log_levels(cw.logging.CRITICAL)
    time.sleep(0.1)
    return scope, target

def xor_bytes(a,b):
    out=bytearray()
    for i in range(len(a)):
        out.append(a[i]^b[i])
    return bytes(out)

def or_bytes(a,b):
    out=bytearray()
    for i in range(len(a)):
        out.append(a[i]|b[i])
    return bytes(out)


def do_operation(scope, target, wordnum):
    # setup CW
    scope.adc.basic_mode = "rising_edge"
    scope.trigger.triggers = "tio4"
    scope.glitch.enabled = True
    scope.glitch.clk_src = 'pll'
    scope.clock.pll.update_fpga_vco(900e6)
    scope.glitch.output = 'enable_only'
    scope.glitch.trigger_src = 'ext_single'
    scope.glitch.num_glitches=2
    if not scope.clock.clkgen_locked: time.sleep(0.1)

    # setup parameters
    delay_1=delay1list[wordnum]
    delay_2=delay2abs-delay_1
    scope.glitch.ext_offset=[delay_1,delay_2]
    scope.glitch.repeat = glitchlen

    # arm and send plaintext
    scope.arm()
    target.simpleserial_write('d', b'0'*16)

    # not sure why this is needed, but it doesn't work without it. copied from the internet somewhere
    scope.io.glitch_hp = False
    scope.io.glitch_hp = True
    scope.io.glitch_lp = False
    scope.io.glitch_lp = True

    response1=None
    response2=None
    # receive response
    try: response1=target.simpleserial_read('r',48)
    except Exception as e: 
        print(("serial error",args,e))
    if response1 is None:  
        reset_target(scope, target)
        return b'\0'*16

    # no fault, hopefully we skipped the nonce increment
    target.simpleserial_write('d', b'0'*16)
    try: response2=target.simpleserial_read('r',48)
    except Exception as e: 
        print(("serial error",args,e))
        return b'\0'*16

    if response1 is None or response1[:16]!=response2[:16] or response1[32:]!=response2[32:]:  
        reset_target(scope, target)
        return b'\0'*16
    
    diff=xor_bytes(response1,response2)
    idx=wordnum*4
    word=diff[16+idx:16+idx+4]
    if word==b'\0\0\0\0': 
        return b'\0'*16
    
    #print(idx,word,delay_1,delay_2,bytes(xor_bytes(response1[16:32],response2[16:32])))
    #print(idx,word,bytes(xor_bytes(response1[16:32],response2[16:32])), response1[32:],response2[32:])
    return diff[16:32]
    


scope,target=init_cw()
res=b'\0'*16
for i in range(1000):
    res=or_bytes(do_operation(scope, target, i%4),res)
    print(res.hex())
    if not b'\0'*4 in res: break
if res==KEY: print("Key recovered: "+res.hex())
else: print("Failed to recover key: "+res.hex())

scope.dis()
target.dis()