import chipwhisperer as cw
import time, os

KEY=bytes(list(x for x in range(16)))
NONCE=bytes(list(x for x in range(16)))
CT_EXPECTED=b'\x00\xd5\xc0H\xca\xd3\xa5f\x81v\xae\x1a\xca:\x00m:A\x01\x91\xf1\xe0.\x83\xe9Ii\x84T\x13.\xc1'
resps=[]

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)
    target.simpleserial_write('k', KEY)
    target.simpleserial_wait_ack()
    target.simpleserial_write('n', NONCE)
    target.simpleserial_wait_ack()


def init_target():
    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)
    # GlitchController lets you iterate through parameters without having to write the logic yourself
    # it also does some other things that I'm not using 
    gc = cw.GlitchController(groups=[], parameters=["width", "ext_offset","iter"])
    gc.set_range("width", 5,5)
    gc.set_range("ext_offset", 100, 120)
    gc.set_range("iter", 1, 1)  # number of attempts per delay/width combination
    gc.set_global_step(1)
    gc.set_step("ext_offset", 1)
    gc.set_step("width",1)
    cw.set_all_log_levels(cw.logging.CRITICAL)

    time.sleep(0.2)
    return scope, target, gc


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

def classify_tag(data):
    dat=xor_bytes(data,CT_EXPECTED)[16:]
    if KEY==bytes(dat): return 'fullkey'
    elif KEY[:8]==bytes(dat[:8]): return 'key1'
    elif KEY[8:]==bytes(dat[8:]): return 'key2'
    elif KEY[:4]==bytes(dat[:4]): return 'key01'
    elif KEY[4:8]==bytes(dat[4:8]): return 'key02'
    elif KEY[8:12]==bytes(dat[8:12]): return 'key03'
    elif KEY[12:16]==bytes(dat[12:16]): return 'key04'
    else: return 'corrupt'

def classify_succ(data):
    normal=CT_EXPECTED
    if data[:16]==normal[:16]:
        return classify_tag(data)
    return 'corrupt'

def do_operation(scope, target, args):
    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=1
    if not scope.clock.clkgen_locked: time.sleep(0.1)
    
    scope.glitch.ext_offset=args[1]
    scope.glitch.repeat = args[0]
    scope.arm()
    target.simpleserial_write('p', b'0'*16)
    scope.io.glitch_hp = False
    scope.io.glitch_hp = True
    scope.io.glitch_lp = False
    scope.io.glitch_lp = True
    response=None
    try: response=target.simpleserial_read('r',32)
    except Exception as e: 
        print((args,e))
    if response is not None: 
        if bytes(response)==CT_EXPECTED:
            return 'normal'
        reset_target(scope,target)
        resp_type=classify_succ(response)
        resps.append((resp_type,args[0],args[1],bytes(response)))
        print((resps[-1],bytes(xor_bytes(response,CT_EXPECTED)[16:])))
        return resp_type
    else:
        reset_target(scope,target)
        return 'reset'

scope,target,gc=init_target()
for i in gc.glitch_values():
    res=do_operation(scope, target, i)
    

scope.dis()
target.dis()