// ========================================================================================================
// ========================================================================================================
// ***************************************** token_enrollment.c *******************************************
// ========================================================================================================
// ========================================================================================================

#include "common.h"
#include "token_common.h"

// ========================================================================================================
// ========================================================================================================
// Start the PUF engine with the current vector to get the timing values. Once it finishes, fetch each of 
// the timing values.

int GenGetTimingVals(volatile unsigned int *CtrlRegA, volatile unsigned int *DataRegA, int max_outputs, 
   int max_sams, int num_outputs, unsigned short timing_val_arr[max_outputs][max_sams], 
   unsigned short output_pos[max_outputs][max_sams], int sam_num, int ctrl_mask)
   {
   int output_num, timing_val;
   int num_non_zero_timing_vals; 

//printf("GenGetTimingVals(): BEGIN\n"); fflush(stdout);

// Sanity check. PUF engine MUST be 'ready'
   if ( (*DataRegA & (1 << IN_SM_READY)) == 0 )
      { printf("ERROR: GenGetTimingVals(): PUF Engine is NOT ready!\n"); fflush(stdout); exit(EXIT_FAILURE); }

//printf("GenGetTimingVals(): PUF/hash engine is ready -- starting PUF engine!\n"); fflush(stdout);

// DEBUG: Set number of samples to 2.
//   ctrl_mask = ctrl_mask | (0 << OUT_CP_NUM_SAM1) | (0 << OUT_CP_NUM_SAM0);

// Start the PUF engine to generate the timing values. Vector pair has already been loaded.
   *CtrlRegA = ctrl_mask | (1 << OUT_CP_PUF_START);
   *CtrlRegA = ctrl_mask; 

//printf("GenGetTimingVals(): Waiting PUF engine to finish!\n"); fflush(stdout);

// Wait for the PUF engine to finish generating the timing data.
   while ( (*DataRegA & (1 << IN_SM_READY)) == 0 );

//printf("GenGetTimingVals(): PUF engine DONE!\n"); fflush(stdout);

// Reset the PN value pointers in the VHDL code for transferring PNs from VHDL to C code.
   *CtrlRegA = ctrl_mask | (1 << OUT_CP_DTI_RESTART); 
   *CtrlRegA = ctrl_mask; 

//printf("GenGetTimingVals(): Checking 'data_ready' of DataTransferOut!\n"); fflush(stdout);

// Wait for 'data_ready' to become 1 after the pointer reset (should already be 1).
   while ( (*DataRegA & (1 << IN_SM_DTI_DATA_READY)) == 0 );

//printf("GenGetTimingVals(): 'data_ready' is set!\n"); fflush(stdout);

// ==================
// Get timing values
   num_non_zero_timing_vals = 0;
   for ( output_num = 0; output_num < num_outputs; output_num++ )
      {

// Read a timing value -- current version uses only 10 bits (for timing values between 0 and 1023), even though phase
// shift can increase to 1120. NEED TO KEEP AN EYE ON THIS!!!!!!!!!!!!!!!!!!!!!!!! Save 400 FFs in the implementation
// right now. Once we move to SRAM, increase this to 11 bits!
      timing_val = *DataRegA & 0x7FF;

// If timing value is 0, then the path has NO transition.
      if ( timing_val > 0 )
         {

// Don't really need to save these since we are writing them to a file directly above.
         timing_val_arr[num_non_zero_timing_vals][sam_num] = (unsigned short)timing_val;
         output_pos[num_non_zero_timing_vals][sam_num] = (unsigned short)output_num;

//printf("GenGetTimingVals(): Timing val at index %d and sam %d is %d with output pos %d\n", 
//   num_non_zero_timing_vals, sam_num, timing_val_arr[num_non_zero_timing_vals][sam_num], output_pos[num_non_zero_timing_vals][sam_num]); fflush(stdout);

         num_non_zero_timing_vals++;

// Sanity check
         if ( num_non_zero_timing_vals > max_outputs )
            { 
            printf("ERROR: GenGetTimingVals(): Too many timing values %d -- max %d -- increase in program!\n", num_non_zero_timing_vals, max_outputs); 
            fflush(stdout); exit(EXIT_FAILURE); 
            }
         }

// Four phases here.
// 1) Got timing value above, indicate we are done reading.
      *CtrlRegA = ctrl_mask | (1 << OUT_CP_DTI_DONE_READING); 

// 2) Wait for 'data_ready' to become 0.
      while ( (*DataRegA & (1 << IN_SM_DTI_DATA_READY)) != 0 );

// 3) Reset done_reading to 0
      *CtrlRegA = ctrl_mask; 

// 4) Wait for 'data_ready' to become 1.
      while ( (*DataRegA & (1 << IN_SM_DTI_DATA_READY)) == 0 );
      }

   return num_non_zero_timing_vals;
   }


// ========================================================================================================
// ========================================================================================================
// Send the timing data to the verifier through the socket. THIS IS ONLY ALLOWED during enrollment.

void SendTimings(int str_length, int verifier_socket_desc, int vec_num, int num_timing_vals, int *overall_num_timing_vals_ptr, 
   int max_outputs, int max_sams, unsigned short timing_val_arr[max_outputs][max_sams], 
   unsigned short output_pos[max_outputs][max_sams], int num_sams, int change_rise_fall)
   {
   char timing_line[str_length];
   char buffer[str_length];
   int sam_num, nz_num;
   int num_lines;

// Load up the timing samples to a string and send to the verifier.
   num_lines = 0;
   for ( nz_num = 0; nz_num < num_timing_vals; nz_num++ )
      {
      sprintf(timing_line, "V: %d\tO: %d\tC: %d\t", vec_num, output_pos[nz_num][0], *overall_num_timing_vals_ptr);
      for ( sam_num = 0; sam_num < num_sams; sam_num++ )
         {
         sprintf(buffer, " %d", timing_val_arr[nz_num][sam_num]);
         strcat(timing_line, buffer);
               
// Sanity check: It must always be true that the sampled output remains the same across samples.
         if ( output_pos[nz_num][0] != output_pos[nz_num][sam_num] )
            { printf("ERROR: SendTimings(): Output that produced timing value has CHANGED %d vs %d!\n", output_pos[nz_num][0], output_pos[nz_num][sam_num]); fflush(stdout); exit(EXIT_FAILURE); }
         }
         
// Send timing value line with all samples included. Be sure to add '+1' to ensure NULL character is transferred to receiver.
      if ( SockSendB((unsigned char *)timing_line, strlen(timing_line) + 1, verifier_socket_desc) < 0 )
         { printf("ERROR: SendTimings(): Send '%s' failed\n", timing_line); fflush(stdout); exit(EXIT_FAILURE); }

// DEBUG
//printf("SendTimings(): Current timing line[%d]:\n'%s'\n", *overall_num_timing_vals_ptr, timing_line); fflush(stdout);
      num_lines++;
      (*overall_num_timing_vals_ptr)++;
      }

// Send the empty string which will cause a blank line to be inserted into the saved file on the receiver, as an indicator of when we change 
// from rising to falling (or vise versa). Be sure to add '+1' to ensure NULL character is transferred to receiver.
   if ( change_rise_fall == 1 )
      if ( SockSendB((unsigned char *)"", strlen("") + 1, verifier_socket_desc) < 0 )
         { printf("ERROR: SendTimings(): Send '<cr>' failed\n"); fflush(stdout); exit(EXIT_FAILURE); }

// DEBUG
#ifdef DEBUG
printf("SendTimings(): Sent %d timing values for vector[%d]\n", nz_num, vec_num); fflush(stdout);
#endif
   }


// ========================================================================================================
// ========================================================================================================
// DEBUG ROUTINE. Check that the vectors received are precisely the same as the ones stored in the file on 
// the server. Write an ASCII file. 

void SaveASCIIVectors(int max_string_len, int max_vecs, int num_vecs, unsigned char *first_vecs_b[max_vecs], 
   unsigned char *second_vecs_b[max_vecs], int vec_len_bits, int has_masks, int mask_len_bits, 
   unsigned char *masks[max_vecs])
   {
   int byte_cnter, bit_cnter;
   unsigned char *vec_ptr;
   int vec_num, vec_pair;
   FILE *OUTFILE;

   if ( (OUTFILE = fopen("ReceivedVectors.txt", "w")) == NULL )
      { printf("ERROR: ReadDatabasePNs(): Could not open ReceivedVectors.txt for writing!\n"); fflush(stdout);  exit(EXIT_FAILURE); }

   for ( vec_num = 0; vec_num < num_vecs; vec_num++ )
      {
      for ( vec_pair = 0; vec_pair < 2; vec_pair++ )
         {

         if ( vec_pair == 0 )
            vec_ptr = first_vecs_b[vec_num];
         else
            vec_ptr = second_vecs_b[vec_num];

// Print ASCII version of bitstring in high order to low order. 
         for ( byte_cnter = vec_len_bits/8 - 1; byte_cnter >= 0; byte_cnter-- )
            {
            for ( bit_cnter = 7; bit_cnter >= 0; bit_cnter-- )
               {

// Check binary bit for 0 or 1
               if ( (vec_ptr[byte_cnter] & (unsigned char)(1 << bit_cnter)) == 0 )
                  fprintf(OUTFILE, "0");
               else
                  fprintf(OUTFILE, "1");
               }
            }
         fprintf(OUTFILE, "\n");
         }

// Extra <cr> between vector pairs.
      fprintf(OUTFILE, "\n");
      }
   fclose(OUTFILE);


// Do the same for the received masks.
   if ( has_masks == 1 )
      {
      if ( (OUTFILE = fopen("ReceivedMasks.txt", "w")) == NULL )
         { printf("ERROR: ReadDatabasePNs(): Could not open ReceivedMasks.txt for writing!\n"); fflush(stdout);  exit(EXIT_FAILURE); }

      for ( vec_num = 0; vec_num < num_vecs; vec_num++ )
         {

// Print ASCII version of bitstring in high order to low order. 
         for ( byte_cnter = mask_len_bits/8 - 1; byte_cnter >= 0; byte_cnter-- )
            {
            for ( bit_cnter = 7; bit_cnter >= 0; bit_cnter-- )
               {

// Check binary bit for 0 or 1
               if ( (masks[vec_num][byte_cnter] & (unsigned char)(1 << bit_cnter)) == 0 )
                  fprintf(OUTFILE, "0");
               else
                  fprintf(OUTFILE, "1");
               }
            }
         fprintf(OUTFILE, "\n");
         }
      fclose(OUTFILE);
      }

   return;
   }


// ========================================================================================================
// ========================================================================================================
// ========================================================================================================
// ========================================================================================================
// Looks like the full blown bitstreams are upto 4,045,564 in size (for Zynq 7020)
#define MAX_PR_BIN_SIZE  5000000

// Only need to be big enough to hold the max number of timing values per vector in this version.
unsigned short timing_val_arr[MAX_OUTPUTS][MAX_SAMS];
unsigned short output_pos[MAX_OUTPUTS][MAX_SAMS];

unsigned char *first_vecs_b[MAX_VECS];
unsigned char *second_vecs_b[MAX_VECS];
unsigned char *masks[MAX_VECS];

// The array that stores the current programming bitstream.
char bstream_buf[MAX_PR_BIN_SIZE];
int bstream_size = 0;


int main(int argc, char *argv[])
   {
   volatile unsigned int *CtrlRegA;
   volatile unsigned int *DataRegA;

   char verifier_IP[MAX_STRING_LEN];
   int verifier_socket_desc = 0;
   int port_number;
   
   char chip_name[MAX_STRING_LEN];
   char outfile_prefix[MAX_STRING_LEN];

   int num_vecs, num_sams; 
   int vec_num, sam_num;

//Timing variables.
#if DEBUG
   struct timeval t0, t1;
   long elapsed; 
#endif

   int overall_num_timing_vals, num_rise_vecs;
   int first_num_non_zero_timing_vals, num_non_zero_timing_vals;
   int ctrl_mask;
   int has_masks; 

   char fu_command_string[MAX_STRING_LEN];
   char bsfilename[MAX_STRING_LEN];
   char *char_ptr;

   int vec_len_bits; 
   int num_outputs; 
   int version_num; 

   int BSTREAM_FILE;
   int DEVCFG_FILE; 

// KG_FU
   char version_coords_KG_FU[25][MAX_STRING_LEN] = {
      "X22_Y0", "X22_Y1", "X22_Y2", "X22_Y3", "X22_Y4", "X22_Y5", "X22_Y6", "X22_Y7", "X22_Y8", "X22_Y9", "X22_Y51", "X22_Y52", "X22_Y53", "X22_Y54", 
      "X22_Y55", "X22_Y56", "X22_Y57", "X22_Y58", "X22_Y59"};

// ======================================================================================================================
// COMMAND LINE
   if ( argc != 3 )
      {
      printf("ERROR: token_enrollment.elf(): Chip name (C1)-- verifier_IP\n");
      return(1);
      }

   sscanf(argv[1], "%s", chip_name);
   strcpy(verifier_IP, argv[2]);

// ====================================================== PARAMETERS ====================================================
   port_number = 8888;

// ======================================================================================================================

// Open up the memory mapped device so we can access the GPIO registers.
   int fd = open("/dev/mem", O_RDWR|O_SYNC);

   if (fd < 0) 
      { printf("ERROR: /dev/mem could NOT be opened!\n"); exit(EXIT_FAILURE); }

// Add 2 for the DataReg (for an offset of 8 bytes for 32-bit integer variables)
   DataRegA = mmap(0, getpagesize(), PROT_READ|PROT_WRITE, MAP_SHARED, fd, GPIO_0_BASE_ADDR);
   CtrlRegA = DataRegA + 2;

// Open up a socket connection to the verifier.
   OpenSocketClient(MAX_STRING_LEN, verifier_IP, port_number, &verifier_socket_desc);

// Need this to do DPR. This fails if the file already exists but no big deal.
   system("mknod /dev/xdevcfg c 259 0 > /dev/null");

// ==================
// ==================
// Verifier is controlling labview now. Loop until server sends the 'EXIT' string.
   strcpy(fu_command_string, "");
   while (1)
      {

// Get the FU to see if there is another iteration to carry out.
      if ( SockGetB((unsigned char *)fu_command_string, MAX_STRING_LEN, verifier_socket_desc) < 0 )
         { printf("ERROR: Failed to receive FU/Command string!\n"); fflush(stdout); exit(EXIT_FAILURE); }

// ------------------------------------------
// Check for EXIT command from the verifier.
      if ( strcmp(fu_command_string, "EXIT") == 0 )
         {
         printf("RECEIVED 'EXIT' string!\n"); fflush(stdout); 
         break;
         }

// Parse out the number 1 in, e.g., "KG_SBOX_V1"
      char_ptr = strchr(fu_command_string, 'V');

// Sanity check
      if ( char_ptr == NULL )
         { printf("ERROR: Failed to find 'V' in fu_command_string '%s'!\n", fu_command_string); fflush(stdout); exit(EXIT_FAILURE); }

// Make it a zero-based index into version_coords.
      sscanf(char_ptr+1, "%d", &version_num);
      version_num--;

// Sanity check
      if ( version_num < 0 )
         { printf("ERROR: Version number 'Vx' MUST be > 0 in fu_command_string '%s'!\n", fu_command_string); fflush(stdout); exit(EXIT_FAILURE); }

      printf("fu_command_string received '%s', (extracted version number-1) %d\n", fu_command_string, version_num); fflush(stdout);

// Construct the bitstream name to read from the SD card and set parameters based on the type of functional unit.
      vec_len_bits = 0;
      num_outputs = 0;
      strcpy(bsfilename, "");
      if ( strstr(fu_command_string, "KG_FU") != NULL )
         {
         vec_len_bits = 64;
         num_outputs = 64;
         num_sams = 16;
         sprintf(bsfilename, "/tmp/KG_FU");
         sprintf(bsfilename, "%s_%s.dcp.bit.bin", bsfilename, version_coords_KG_FU[version_num]);
         }
//      else if ( strstr(fu_command_string, "xxx") != NULL )
//         {
//         }
      else
         { printf("ERROR: FU/Command string INVALID '%s'!\n", fu_command_string); fflush(stdout); exit(EXIT_FAILURE); }


      printf("Reading bitstream file '%s'\n", bsfilename); fflush(stdout);
      if ( strlen(bsfilename) == 0 )
         { printf("ERROR: NULL bitstream filename!\n"); fflush(stdout); exit(EXIT_FAILURE); }

// Read the bitstream into a buffer
      if ( (BSTREAM_FILE = open(bsfilename, O_RDONLY)) < 0 )
         { printf("ERROR: Failed to open bstream file '%s'!\n", bsfilename); fflush(stdout); exit(EXIT_FAILURE); }
      if ( (bstream_size = read(BSTREAM_FILE, bstream_buf, MAX_PR_BIN_SIZE)) < 0 )
         { printf("ERROR: Failed to read bstream file '%s'!\n", bsfilename); fflush(stdout); exit(EXIT_FAILURE); }
      close(BSTREAM_FILE);

// Program the PL side with the buffered bitstream
      if ( (DEVCFG_FILE = open("/dev/xdevcfg", O_RDWR)) < 0 )
         { printf("ERROR: Failed to open 'xdevcfg'!\n"); fflush(stdout); exit(EXIT_FAILURE); }

      write(DEVCFG_FILE, bstream_buf, bstream_size);
      close(DEVCFG_FILE);
      usleep (10000);

// Set the control mask to indicate enrollment (also, normally set PUF mode here too but there is only one mode here).
      ctrl_mask = (1 << OUT_CP_MODE1) | (1 << OUT_CP_MODE0); 
      
// ------------------------------------------
// Read all the vectors from the verifier into a set of malloc'ed string arrays. Verifier will send number of rising
// vectors (inspects vectors as it reads them) and indicate whether masks will also be sent. 
      num_vecs = ReceiveVectors(MAX_STRING_LEN, verifier_socket_desc, MAX_VECS, first_vecs_b, second_vecs_b, vec_len_bits, 
         &num_rise_vecs, &has_masks, num_outputs, masks);
      printf("Number of vectors received %d\tNumber of rising vectors %d\tHas masks ? %d\n", num_vecs, num_rise_vecs, has_masks); fflush(stdout);

// DEBUG
// Regenerate the vectors in an ASCII file so that we can verify that they are the same ones that were generated.
#ifdef DEBUG
      SaveASCIIVectors(MAX_STRING_LEN, MAX_VECS, num_vecs, first_vecs_b, second_vecs_b, vec_len_bits, has_masks, num_outputs, masks);
#endif

// Be sure to add '+1' to ensure NULL character is transferred to receiver.
      sprintf(outfile_prefix, "%s_%s", chip_name, fu_command_string);
      if ( SockSendB((unsigned char *)outfile_prefix, strlen(outfile_prefix) + 1, verifier_socket_desc) < 0 )
         { printf("ERROR: Send '%s' failed\n", outfile_prefix); fflush(stdout); exit(EXIT_FAILURE); }

// Run time information 
#ifdef DEBUG
      gettimeofday(&t0, 0);
#endif

// Do a soft RESET
      *CtrlRegA = (1 << OUT_CP_RESET);
      *CtrlRegA = 0;
      usleep (10000);

// Repeat for n vectors. NOTE: Enrollment mode of operation allows the PUF to be run on a per-vector basis, for as many samples as you 
// want (CollectPNs.vhd is NOT RUN -- ONLY LCDT). This iterative version allows very large sets of timing values to be retrieved since 
// the timing values for each vector (ALL SAMPLES) are sent after they are generated. 
      overall_num_timing_vals = 0;
      for ( vec_num = 0; vec_num < num_vecs; vec_num++ )
         {

// Progress indicator
         if ( (vec_num + 1) % 100 == 0 )
            { printf("Processing vector number %d\n", vec_num); fflush(stdout); }

// Transfer a vector pair to the VHDL registers.
         LoadVecPairMask(MAX_STRING_LEN, CtrlRegA, DataRegA, MAX_VECS, vec_num, first_vecs_b, second_vecs_b, ctrl_mask, 
            vec_len_bits, VEC_CHUNK_SIZE, has_masks, num_outputs, masks);

// Get timing values over n samples.
         for ( sam_num = 0; sam_num < num_sams; sam_num++ )
            {

// 'GenGetTimingVals' checks if PUF engine is ready, starts the engine, waits for strobing to complete and then retrieves the timing 
// values for all outputs under this vector (number of timing value lines to be transmitted to the server is given by 
// 'num_non_zero_timing_vals'
            num_non_zero_timing_vals = GenGetTimingVals(CtrlRegA, DataRegA, MAX_OUTPUTS, MAX_SAMS, num_outputs, timing_val_arr, 
               output_pos, sam_num, ctrl_mask);

// Sanity check. Same paths must be timed across ALL n samples.
            if ( sam_num == 0 )
               first_num_non_zero_timing_vals = num_non_zero_timing_vals;
            else if ( first_num_non_zero_timing_vals != num_non_zero_timing_vals )
               { printf("ERROR: Number of paths timed for first sample is NOT equal for other samples!\n"); fflush(stdout); exit(EXIT_FAILURE); }
            }

// Transmit the timing values for this vector (including all samples) to the verifier one line at a time. Last parameter sends the empty string to 
// the verifier after this last rising edge vector is applied, which adds a extra line in the output file between rising and falling data values.
         SendTimings(MAX_STRING_LEN, verifier_socket_desc, vec_num, num_non_zero_timing_vals, &overall_num_timing_vals, 
            MAX_OUTPUTS, MAX_SAMS, timing_val_arr, output_pos, num_sams, (vec_num == num_rise_vecs - 1));

// To allow multiple boards and parallelism, we send a 'CONT' string after the entire dataset (all samples) are collected for a given vector
// EXCEPT FOR THE LAST VECTOR. The 'DONE' string is sent instead.
         if ( vec_num < num_vecs - 1 && SockSendB((unsigned char *)"CONT", strlen("CONT") + 1, verifier_socket_desc) < 0 )
            { printf("ERROR: Send 'CONT' failed\n"); fflush(stdout); exit(EXIT_FAILURE); }
         }

// Send "DONE" string to the data for ALL vectors is complete. Be sure to add '+1' to ensure NULL character is transferred to receiver.
      if ( SockSendB((unsigned char *)"DONE", strlen("DONE") + 1, verifier_socket_desc) < 0 )
         { printf("ERROR: Send 'DONE' failed\n"); fflush(stdout); exit(EXIT_FAILURE); }

// Free up the vector malloced space for next run.
      for ( vec_num = 0; vec_num < num_vecs; vec_num++ )
         {
         if ( first_vecs_b[vec_num] != NULL )
            free(first_vecs_b[vec_num]);
         first_vecs_b[vec_num] = NULL;
         if ( second_vecs_b[vec_num] != NULL )
            free(second_vecs_b[vec_num]);
         second_vecs_b[vec_num] = NULL;
         if ( masks[vec_num] != NULL )
            free(masks[vec_num]);
         masks[vec_num] = NULL;
         }
      }

// Run time information
#ifdef DEBUG
   gettimeofday(&t1, 0); elapsed = (t1.tv_sec-t0.tv_sec)*1000000 + t1.tv_usec-t0.tv_usec; printf("\tDone: Elapsed %ld us\n\n", (long)elapsed);
#endif

   close(verifier_socket_desc);
   return 0;
   }