import numpy as np
import chipwhisperer as cw
import time
import sys
class Test:
  def __init__(self, FCPU = 0):
    self.log = []
    self.scope = cw.scope()
    self.scope.default_setup()
    self.scope.clock.freq_ctr_src = 'clkgen'
    time.sleep(1)
    print(self.scope.clock.freq_ctr)
    print(self.scope.clock.clkgen_mul, self.scope.clock.clkgen_div)
    cw.program_target(self.scope, cw.programmers.STM32FProgrammer, "./simpleserial-aead-CWLITEARM.hex")
    time.sleep(1)
    self.scope.io.target_pwr = False
    self.scope.dis()
    self.scope = cw.scope()
    self.scope.default_setup()
    if FCPU!=0:
      self.scope.clock.clkgen_freq = FCPU
      print(self.scope.clock.clkgen_mul, self.scope.clock.clkgen_div)
    time.sleep(1)
    self.scope.clock.freq_ctr_src = 'clkgen'
    time.sleep(1)
    print(self.scope.clock.freq_ctr)
    self.scope.io.target_pwr = True
    self.target = cw.target(self.scope)
    self.scope.io.target_pwr = False
    time.sleep(0.5)
    self.scope.io.target_pwr = True
    time.sleep(0.1)
    print("HELLO?", self.target.read().strip('\n'))
    return

  def data_write(self, data, data_len, tag):
    N_input = int(np.ceil(data_len/16))
    for fr in range(0, N_input):
      if fr==(N_input-1):
        frag = tag+(data[(fr*32):].ljust(32, '0'))+"\n"
      else:
        frag = tag+data[(fr*32):(fr*32+32)]+"\n"
      print("Set fragment "+frag)     
      self.target.write(frag)
      time.sleep(0.2)
      discard = self.target.read()
      print(discard)
    return

  def set_data(self, key, nonce, a_data, text):
    A_size = len(a_data)//2
    T_size = len(text)//2
    self.O_size = T_size-16
    print("AD_SIZE =", A_size)
    print("TEXT_SIZE =", T_size)
    size_tag = "x02"+hex(A_size)[2:].zfill(4)+hex(T_size)[2:].zfill(4)+"\n"
    print(size_tag)
    self.target.write(size_tag)
    time.sleep(0.1)
    discard = self.target.read()
    print(discard)
    print("Key:", key)
    self.data_write(key, 16, 'k')
    print("Nonce:", nonce)
    self.data_write(nonce, 16, 'n')
    print("A. Data:", a_data)
    self.data_write(a_data, A_size, 'a')
    print("Text:", text)
    self.data_write(text, T_size, 'i')
    return
  
  def decrypt(self):
    self.target.write("e\n")
    print("Executing (ASCON_AEAD_128_DEC)")
    time.sleep(10.0)
    discard = (self.target.read())
    print(discard)
    return
  
  def get_response(self):
    response = ""
    while True:
      self.target.write("o\n")
      time.sleep(0.1)
      cccc = self.target.read()
      print(cccc)
      size_left = (2*self.O_size)-len(response)
      if size_left>=32:
        response += cccc[1:33]
      else:
        response += cccc[1:(size_left)+1]
      if len(response)==(2*self.O_size):
        break
    print("Plaintext:", response)
    self.log.append(response)
    return response.lower()

  def log_print(self):
    for t in self.log:
      print(t)
    return

  def close(self):
    self.scope.io.target_pwr = False
    time.sleep(2)
    self.scope.io.target_pwr = True
    self.target.dis()
    self.scope.dis()
    return

if __name__=='__main__':
  FC = 5000000
  Temp = Test(FC)
  print("=====================================================================================")
  #Count = 104
  Key   = "000102030405060708090a0b0c0d0e0f"
  Nonce = "000102030405060708090a0b0c0d0e0f"
  Cipher_Tag  = "7763f8a2bbfbb05b3b4f54da1576a863b47409"
  A_DATA     = "00010203"
  Plaintext = "000102"
  Temp.set_data(Key, Nonce, A_DATA, Cipher_Tag)
  Temp.decrypt()
  Results = Temp.get_response()
  print(Plaintext == Results)
  print(Plaintext)
  print(Results)
  print("=====================================================================================")
  #Count = 11
  Key   = "000102030405060708090a0b0c0d0e0f"
  Nonce = "000102030405060708090a0b0c0d0e0f"
  Cipher_Tag  = "4b006a400b6dfb9777bc3446c2b7dc26"
  A_DATA     = "00010203040506070809"
  Plaintext = ""
  Temp.set_data(Key, Nonce, A_DATA, Cipher_Tag)
  Temp.decrypt()
  Results = Temp.get_response()
  print(Plaintext == Results)
  print(Plaintext)
  print(Results)
  #Temp.log_print()
  Temp.close()

