#include <AES.h>
//#include "./printf.h"



//#include <ESP8266WiFi.h>

const int size_mesg=256;




const int sleepTimeS = 10;


const int h=4;
const int h2=h*h;
byte DK[64];
int rp=1;

const int len=size_mesg/h2;



typedef byte   uchar;

int Pbox[len];
  int PboxRM[h2];
  uchar Sbox1[256];
  uchar Sbox2[256];  
  uchar RM1_cpy[h2];

  








void inverse_tables(uchar *tab, int size_tab,uchar *inv_perm_tabs) {

  for(int i=0;i<size_tab;i++) {
    inv_perm_tabs[tab[i]] = i;
  }

}

void inverse_tables_int(int *tab, int size_tab,int *inv_perm_tabs) {

  for(int i=0;i<size_tab;i++) {
    inv_perm_tabs[tab[i]] = i;
  }

}



void rc4key(uchar *key, uchar *sc, int size_DK) {

  for(int i=0;i<256;i++) {
    sc[i]=i;
  }


  uchar j0 = 0;
  for(int i0=0; i0<256; i0++) {
    j0 = (j0 + sc[i0] + key[i0%size_DK] )&0xFF;
    uchar tmp = sc[i0];
    sc[i0] = sc[j0 ];
    sc[j0] = tmp;
  }
}



void rc4keyperm(uchar *key,int len, int rp,int *sc, int size_DK) {

  //sc=1:len;


  
  for (int i=0;i<len;i++) {
    sc[i]=i;
  }
  for (int it = 0; it < rp; it++) {
    int j0 = 1;
    for(int i0 = 0; i0<len; i0++) {
      j0 = (j0 + sc[i0] + sc[j0] + key[i0%size_DK] )% len;
      int tmp = sc[i0];
      sc[i0] = sc[j0];
      sc[j0] = tmp;
    }

  }
}

void prga(uchar *sc, int ldata, uchar *r) {
  uchar i0=0;
  uchar j0=0;

  for (int it=0; it<ldata; it++) {
    i0 = ((i0+1)&0xFE); //%255);
    j0 = (j0 + sc[i0])&0xFF;
    uchar tmp = sc[i0];
    sc[i0] = sc[j0];
    sc[j0] = tmp;
    r[it]=sc[(sc[i0]+sc[j0])&0xFF];
  }
}



void encrypt_ctr(const uchar* seq_in, uchar *seq_out, int len,uchar* RM1,const int *Pbox, const int *PboxRM, const uchar *Sbox1, const uchar *Sbox2, int enc) {

 //printArray(seq_out,16);
  uchar X[h2];
  uchar fX[h2];
  
  int ind1,ind2;

  
   for(int a=0;a<h2;a++) {
     X[a]=Sbox1[a&0xFF];           //Warning according to the size of h2, we can be outsize of Sbox1[a]
   }

   
  for(int it=0;it<len;it++) {
    if(enc) {
      ind1=it*h2;
      ind2=Pbox[it]*h2;
    }
    else {
      ind2=it*h2;
      ind1=Pbox[it]*h2;
    }
    

   for(int a=0;a<h2;a+=4) {
      X[a]=X[Sbox1[a]];
      X[a+1]=X[Sbox1[a+1]];
      X[a+2]=X[Sbox1[a+2]];
      X[a+3]=X[Sbox1[a+3]];
    }

    for(int a=0;a<h2;a+=4){
      fX[a]=X[a];
      fX[a+1]=X[a+1];
      fX[a+2]=X[a+2];
      fX[a+3]=X[a+3];
    }
   


    //printf("fx %d\n",fX[0] );
    
  /*
    for(int a=0;a<h2;a+=16) {
      *(int*)&fX[a]^=it;
      *(int*)&fX[a+4]^=it;
      *(int*)&fX[a+8]^=it;
      *(int*)&fX[a+12]^=it;
    }  
   */




    for(int a=0;a<h2;a+=4) {
      fX[a]=fX[a]^RM1[a];
      fX[a+1]=fX[a+1]^RM1[a+1];
      fX[a+2]=fX[a+2]^RM1[a+2];
      fX[a+3]=fX[a+3]^RM1[a+3];
    }

 
    for(int a=0;a<h2;a+=4) {
      fX[a]=Sbox2[fX[a]];
      fX[a+1]=Sbox2[fX[a+1]];
      fX[a+2]=Sbox2[fX[a+2]];
      fX[a+3]=Sbox2[fX[a+3]];
    }
    
 

    
     for(int a=0;a<h2;a+=4) {
      fX[a]=fX[a]^seq_in[ind2+a];
      fX[a+1]=fX[a+1]^seq_in[ind2+a+1];
      fX[a+2]=fX[a+2]^seq_in[ind2+a+2];
      fX[a+3]=fX[a+3]^seq_in[ind2+a+3];
    }


 
    for(int a=0;a<h2;a+=4) {
      seq_out[ind1+a]=fX[a];
      seq_out[ind1+a+1]=fX[a+1];
      seq_out[ind1+a+2]=fX[a+2];
      seq_out[ind1+a+3]=fX[a+3];
    }
    
    for(int a=0;a<h2;a+=4) {
      RM1[a]=RM1[PboxRM[a]];
      RM1[a+1]=RM1[PboxRM[a+1]];
      RM1[a+2]=RM1[PboxRM[a+2]];
      RM1[a+3]=RM1[PboxRM[a+3]];
    }


    
  }

//  printf("size mes %d\n",size_mesg);
//  printf("len %d\n",len);


//printArray(seq_out,16);

}



void printArray(byte *mes, int n) {
  for (byte i = 0; i < n; i++) {
    Serial.print(mes[i]);
    Serial.print(" ");
  }
  Serial.println();
}


void setup ()
{
  Serial.begin (57600) ;
//  printf_begin();
  delay(500);
  printf("\n===testng mode\n") ;


  uchar sc[256];  
  
  uchar RM1[h2];
  uchar RM2[h2];
 
  randomSeed(134);

  for(byte i=0;i<64;i++) {
    DK[i]=random(255);
  }
  rc4key(DK, Sbox1, 8);
  
  
  rc4key(&DK[8], Sbox2, 8);
  
  rc4key(&DK[16], sc, 16);
  
  
  prga(sc, h2, RM1);
  
  
  rc4keyperm(&DK[32], len, rp, Pbox, 16);
  
  
  rc4keyperm(&DK[48], h2, rp, PboxRM, 16);

  for(int i=0;i<h2;i++){
    RM2[i]=RM1[i];
    RM1_cpy[i]=RM1[i];
  }

 //  WiFi.mode( WIFI_OFF );
 
Serial.println("END OF INIT"); 
  

  
//  otfly_test () ;
//  otfly_test256 () ;
}

void loop () 
{
  prekey_test () ;
  delay(1000);

//  Serial.println("ESP8266 in sleep mode");
 // ESP.deepSleep(sleepTimeS * 1000000, WAKE_RF_DISABLED );
  delay(1000);
//  Serial.println("ESP8266 in normal mode");
}

void prekey (int bits)
{
 
 
  byte plain[size_mesg];
  byte cipher [size_mesg] ;
  byte check [size_mesg] ;

  randomSeed(334);
  for(int i=0;i<size_mesg;i++) {
    plain[i]=random(255);
  }
  
  uchar RM1[h2];
  uchar RM2[h2];
 
  for(int i=0;i<h2;i++){
    RM1[i]=RM1_cpy[i];
    RM2[i]=RM1_cpy[i];
  }

 
 

  unsigned long ms1 = micros ();
  encrypt_ctr(plain, cipher,len,RM1,Pbox,PboxRM,Sbox1,Sbox2,1);
  Serial.print("ONE   Encryption took: ");
  Serial.println(micros() - ms1);
  
  ms1 = micros ();
  encrypt_ctr(cipher, check,len,RM2,Pbox,PboxRM,Sbox1,Sbox2,0);
  Serial.print("ONE   Decryption took: ");
  Serial.println(micros() - ms1);

  printf("\n\nPLAIN :");
  printArray(plain,16*2);
  printf("\nCIPHER:");
  printArray(cipher,16*2);
  printf("\nCHECK :");
  printArray(check,16*2);
  
  bool equal=true;
  for(int i=0;i<size_mesg;i++) {
      
      if(check[i]!=plain[i]) {
        equal=false;
      }
    }
  Serial.print("CHECK ");
  Serial.println(equal);
  
}

void prekey_test ()
{
  prekey (128) ;
}