#include "api.h"
#include <string.h> 
#include <stdio.h>
#include <stdlib.h>
#include "elephant_160.h"

// #include "elephant_176.h"
// #include "elephant_200.h"


//spongent.c 
#if		defined(SPONGENT160)
#define nBits		160
#define nSBox		20 // 8 bit
#define nRounds		80
#define lfsrIV	    0x75
#elif   defined(SPONGENT176)
#define nBits		176
#define nSBox		22
#define nRounds		90
#define lfsrIV	    0x45
#else
#define nBits		0
#define nSBox		0
#define nRounds		0
#define lfsrIV	    0
#endif

#define GET_BIT(x,y) (x >> y) & 0x1

/* Spongent eight bit S-box */
int  sBoxLayer[256] = {
	0xee, 0xed, 0xeb, 0xe0, 0xe2, 0xe1, 0xe4, 0xef, 0xe7, 0xea, 0xe8, 0xe5, 0xe9, 0xec, 0xe3, 0xe6, 
	0xde, 0xdd, 0xdb, 0xd0, 0xd2, 0xd1, 0xd4, 0xdf, 0xd7, 0xda, 0xd8, 0xd5, 0xd9, 0xdc, 0xd3, 0xd6, 
	0xbe, 0xbd, 0xbb, 0xb0, 0xb2, 0xb1, 0xb4, 0xbf, 0xb7, 0xba, 0xb8, 0xb5, 0xb9, 0xbc, 0xb3, 0xb6, 
	0x0e, 0x0d, 0x0b, 0x00, 0x02, 0x01, 0x04, 0x0f, 0x07, 0x0a, 0x08, 0x05, 0x09, 0x0c, 0x03, 0x06, 
	0x2e, 0x2d, 0x2b, 0x20, 0x22, 0x21, 0x24, 0x2f, 0x27, 0x2a, 0x28, 0x25, 0x29, 0x2c, 0x23, 0x26, 
	0x1e, 0x1d, 0x1b, 0x10, 0x12, 0x11, 0x14, 0x1f, 0x17, 0x1a, 0x18, 0x15, 0x19, 0x1c, 0x13, 0x16, 
	0x4e, 0x4d, 0x4b, 0x40, 0x42, 0x41, 0x44, 0x4f, 0x47, 0x4a, 0x48, 0x45, 0x49, 0x4c, 0x43, 0x46, 
	0xfe, 0xfd, 0xfb, 0xf0, 0xf2, 0xf1, 0xf4, 0xff, 0xf7, 0xfa, 0xf8, 0xf5, 0xf9, 0xfc, 0xf3, 0xf6, 
	0x7e, 0x7d, 0x7b, 0x70, 0x72, 0x71, 0x74, 0x7f, 0x77, 0x7a, 0x78, 0x75, 0x79, 0x7c, 0x73, 0x76, 
	0xae, 0xad, 0xab, 0xa0, 0xa2, 0xa1, 0xa4, 0xaf, 0xa7, 0xaa, 0xa8, 0xa5, 0xa9, 0xac, 0xa3, 0xa6, 
	0x8e, 0x8d, 0x8b, 0x80, 0x82, 0x81, 0x84, 0x8f, 0x87, 0x8a, 0x88, 0x85, 0x89, 0x8c, 0x83, 0x86, 
	0x5e, 0x5d, 0x5b, 0x50, 0x52, 0x51, 0x54, 0x5f, 0x57, 0x5a, 0x58, 0x55, 0x59, 0x5c, 0x53, 0x56, 
	0x9e, 0x9d, 0x9b, 0x90, 0x92, 0x91, 0x94, 0x9f, 0x97, 0x9a, 0x98, 0x95, 0x99, 0x9c, 0x93, 0x96, 
	0xce, 0xcd, 0xcb, 0xc0, 0xc2, 0xc1, 0xc4, 0xcf, 0xc7, 0xca, 0xc8, 0xc5, 0xc9, 0xcc, 0xc3, 0xc6, 
	0x3e, 0x3d, 0x3b, 0x30, 0x32, 0x31, 0x34, 0x3f, 0x37, 0x3a, 0x38, 0x35, 0x39, 0x3c, 0x33, 0x36, 
	0x6e, 0x6d, 0x6b, 0x60, 0x62, 0x61, 0x64, 0x6f, 0x67, 0x6a, 0x68, 0x65, 0x69, 0x6c, 0x63, 0x66 
};

void PrintState(BYTE* state)
{
	for(int i = nSBox-1; i>=0; i--)
		printf("%02X ", state[i]);
	printf("\n");
}

BYTE lCounter(BYTE lfsr)
{
	lfsr = (lfsr << 1) | (((0x40 & lfsr) >> 6) ^ ((0x20 & lfsr) >> 5));
	lfsr &= 0x7f; 
	return lfsr;
}

BYTE retnuoCl(BYTE lfsr)
{
	return ((lfsr & 0x01) <<7) | ((lfsr & 0x02) << 5) | ((lfsr & 0x04) << 3)
		| ((lfsr & 0x08) << 1) | ((lfsr & 0x10) >> 1) | ((lfsr & 0x20) >> 3)
		| ((lfsr & 0x40) >> 5) | ((lfsr & 0x80) >> 7);		
}

int Pi(int i)
{
	if (i != nBits-1)
		return (i*nBits/4)%(nBits-1);
	else
		return nBits-1;
}

void pLayer(BYTE* state)
{
	int	PermutedBitNo;
	BYTE tmp[nSBox], x, y;
	
	for(int i = 0; i < nSBox; i++) tmp[i] = 0;
	
	for(int i = 0; i < nSBox; i++){
		for(int j = 0; j < 8; j++){ 
			x			= GET_BIT(state[i],j);
			PermutedBitNo	= Pi(8*i+j);
			y			= PermutedBitNo/8;
			tmp[y]		^= x << (PermutedBitNo - 8*y);
		}
	}	
	memcpy(state, tmp, nSBox);
}

void permutation(BYTE* state)
{
	BYTE IV = lfsrIV;
	BYTE INV_IV;

	for(int i = 0; i < nRounds; i++){
#ifdef _PrintState_
		printf("%3d\t", i);
		PrintState(state);
#endif
		/* Add counter values */
		state[0] ^= IV;
		INV_IV = retnuoCl(IV);
		state[nSBox-1] ^= INV_IV;
		IV	= lCounter(IV);
		
		/* sBoxLayer layer */
		for(int j = 0; j < nSBox; j++)	
			state[j] =  sBoxLayer[state[j]];
		
		/* pLayer */
		pLayer(state);
	}
#ifdef _PrintState_
	printf("%3d\t", i);
	PrintState(state);
#endif
}

BYTE rotl3(BYTE b)
{
    return (b << 3) | (b >> 5);
}

int constcmp(const BYTE* a, const BYTE* b, SIZE length)
{
    BYTE r = 0;

    for (SIZE i = 0; i < length; ++i)
        r |= a[i] ^ b[i];
    return r; 
}


// State should be BLOCK_SIZE bytes long
// Note: input may be equal to output
void lfsr_step(BYTE* output, BYTE* input)
{
    BYTE temp = rotl3(input[0]) ^ (input[3] << 7) ^ (input[13] >> 7);
    for(SIZE i = 0; i < BLOCK_SIZE - 1; ++i)
        output[i] = input[i + 1];
    output[BLOCK_SIZE - 1] = temp;
}

void xor_block(BYTE* state, const BYTE* block, SIZE size)
{
    for(SIZE i = 0; i < size; ++i)
        state[i] ^= block[i];
}

// Write the ith assocated data block to "output".
// The nonce is prepended and padding is added as required.
// adlen is the length of the associated data in bytes
void get_ad_block(BYTE* output, const BYTE* ad, SIZE adlen, const BYTE* npub, SIZE i)
{
    SIZE len = 0;
    // First block contains nonce
    // Remark: nonce may not be longer then BLOCK_SIZE
    if(i == 0) {
        memcpy(output, npub, CRYPTO_NPUBBYTES);
        len += CRYPTO_NPUBBYTES;
    }

    const SIZE block_offset = i * BLOCK_SIZE - (i != 0) * CRYPTO_NPUBBYTES;
    // If adlen is divisible by BLOCK_SIZE, add an additional padding block
    if(i != 0 && block_offset == adlen) {
        memset(output, 0x00, BLOCK_SIZE);
        output[0] = 0x01;
        return;
    }
    const SIZE r_outlen = BLOCK_SIZE - len;
    const SIZE r_adlen  = adlen - block_offset;
    // Fill with associated data if available
    if(r_outlen <= r_adlen) { // enough AD
        memcpy(output + len, ad + block_offset, r_outlen);
    } else { // not enough AD, need to pad
        if(r_adlen > 0) // ad might be nullptr
            memcpy(output + len, ad + block_offset, r_adlen);
        memset(output + len + r_adlen, 0x00, r_outlen - r_adlen);
        output[len + r_adlen] = 0x01;
    }
}

// Return the ith ciphertext block.
// clen is the length of the ciphertext in bytes 
void get_c_block(BYTE* output, const BYTE* c, SIZE clen, SIZE i)
{
    const SIZE block_offset = i * BLOCK_SIZE;
    // If clen is divisible by BLOCK_SIZE, add an additional padding block
    if(block_offset == clen) {
        memset(output, 0x00, BLOCK_SIZE);
        output[0] = 0x01;
        return;
    }
    const SIZE r_clen  = clen - block_offset;
    // Fill with ciphertext if available
    if(BLOCK_SIZE <= r_clen) { // enough ciphertext
        memcpy(output, c + block_offset, BLOCK_SIZE);
    } else { // not enough ciphertext, need to pad
        if(r_clen > 0) // c might be nullptr
            memcpy(output, c + block_offset, r_clen);
        memset(output + r_clen, 0x00, BLOCK_SIZE - r_clen);
        output[r_clen] = 0x01;
    }
}

// It is assumed that c is sufficiently long
// Also, tag and c should not overlap
void crypto_aead_impl(
    BYTE* c, BYTE* tag, const BYTE* m, SIZE mlen, const BYTE* ad, SIZE adlen,
    const BYTE* npub, const BYTE* k, int encrypt)
{ 
    // Compute number of blocks
    const SIZE nblocks_c  = 1 + mlen / BLOCK_SIZE;
    const SIZE nblocks_m  = (mlen % BLOCK_SIZE) ? nblocks_c : nblocks_c - 1;
    const SIZE nblocks_ad = 1 + (CRYPTO_NPUBBYTES + adlen) / BLOCK_SIZE;
    const SIZE nb_it = (nblocks_c + 1 > nblocks_ad - 1) ? nblocks_c + 1 : nblocks_ad - 1;

    // Storage for the expanded key L
    BYTE expanded_key[BLOCK_SIZE] = {0};
    memcpy(expanded_key, k, CRYPTO_KEYBYTES);
    permutation(expanded_key);

    // Buffers for storing previous, current and next mask
    BYTE mask_buffer_1[BLOCK_SIZE] = {0};
    BYTE mask_buffer_2[BLOCK_SIZE] = {0};
    BYTE mask_buffer_3[BLOCK_SIZE] = {0};
    memcpy(mask_buffer_2, expanded_key, BLOCK_SIZE);

    BYTE* previous_mask = mask_buffer_1;
    BYTE* current_mask = mask_buffer_2;
    BYTE* next_mask = mask_buffer_3;

    // Buffer to store current ciphertext/AD block
    BYTE buffer[BLOCK_SIZE];
    
    // Tag buffer and initialization of tag to first AD block
    BYTE tag_buffer[BLOCK_SIZE] = {0};
    get_ad_block(tag_buffer, ad, adlen, npub, 0);

    SIZE offset = 0;
    for(SIZE i = 0; i < nb_it; ++i) {
        // Compute mask for the next message
        lfsr_step(next_mask, current_mask);
        
        if(i < nblocks_m) {
            // Compute ciphertext block
            memcpy(buffer, npub, CRYPTO_NPUBBYTES);
            memset(buffer + CRYPTO_NPUBBYTES, 0, BLOCK_SIZE - CRYPTO_NPUBBYTES);
            xor_block(buffer, current_mask, BLOCK_SIZE);
            xor_block(buffer, next_mask, BLOCK_SIZE);
            permutation(buffer);
            xor_block(buffer, current_mask, BLOCK_SIZE);
            xor_block(buffer, next_mask, BLOCK_SIZE);
            const SIZE r_size = (i == nblocks_m - 1) ? mlen - offset : BLOCK_SIZE;
            xor_block(buffer, m + offset, r_size);
            memcpy(c + offset, buffer, r_size);
        }

        if(i > 0 && i <= nblocks_c) {
            // Compute tag for ciphertext block
            get_c_block(buffer, encrypt ? c : m, mlen, i - 1);
            xor_block(buffer, previous_mask, BLOCK_SIZE);
            xor_block(buffer, next_mask, BLOCK_SIZE);
            permutation(buffer);
            xor_block(buffer, previous_mask, BLOCK_SIZE);
            xor_block(buffer, next_mask, BLOCK_SIZE);
            xor_block(tag_buffer, buffer, BLOCK_SIZE);
        }

        // If there is any AD left, compute tag for AD block 
        if(i + 1 < nblocks_ad) {
            get_ad_block(buffer, ad, adlen, npub, i + 1);
            xor_block(buffer, next_mask, BLOCK_SIZE);
            permutation(buffer);
            xor_block(buffer, next_mask, BLOCK_SIZE);
            xor_block(tag_buffer, buffer, BLOCK_SIZE);
        }

        // Cyclically shift the mask buffers 
        // Value of next_mask will be computed in the next iteration
        BYTE* const temp = previous_mask;
        previous_mask = current_mask;
        current_mask = next_mask;
        next_mask = temp;

        offset += BLOCK_SIZE;
    }
    // Compute tag
    xor_block(tag_buffer, expanded_key, BLOCK_SIZE);
    permutation(tag_buffer);
    xor_block(tag_buffer, expanded_key, BLOCK_SIZE);
    memcpy(tag, tag_buffer, CRYPTO_ABYTES);
}

// Remark: c must be at least mlen + CRYPTO_ABYTES long

int crypto_aead_encrypt(
  unsigned char *c, unsigned long long *clen,
  const unsigned char *m, unsigned long long mlen,
  const unsigned char *ad, unsigned long long adlen,
  const unsigned char *nsec,
  const unsigned char *npub,
  const unsigned char *k)
{ 
    (void)nsec;
    *clen = mlen + CRYPTO_ABYTES;
    BYTE tag[CRYPTO_ABYTES];
    crypto_aead_impl(c, tag, m, mlen, ad, adlen, npub, k, 1);
    memcpy(c + mlen, tag, CRYPTO_ABYTES); 
    return 0;
}

int crypto_aead_decrypt(
  unsigned char *m, unsigned long long *mlen,
  unsigned char *nsec,
  const unsigned char *c, unsigned long long clen,
  const unsigned char *ad, unsigned long long adlen,
  const unsigned char *npub,
  const unsigned char *k)
{
    (void)nsec;
    if(clen < CRYPTO_ABYTES)
        return -1;
    *mlen = clen - CRYPTO_ABYTES;
    BYTE tag[CRYPTO_ABYTES];
    crypto_aead_impl(m, tag, c, *mlen, ad, adlen, npub, k, 0);
    return (constcmp(c + *mlen, tag, CRYPTO_ABYTES) == 0) ? 0 : -1;
}

//coding sendiri 
void string2hexString(unsigned char* input, int clen, char* output)
{
    int loop;
    int i; 
    
    i=0;
    loop=0;
    
    for (i=0;i<clen;i+=2){
        sprintf((char*)(output+i),"%02X", input[loop]);
        loop+=1;

    }
    //insert NULL at the end of the output string
    output[i++] = '\0';
}
void *hextobyte(char *hexstring, unsigned char* bytearray ) {

    int i;

    int str_len = strlen(hexstring);

    for (i = 0; i < (str_len / 2); i++) {
        sscanf(hexstring + 2*i, "%02x", &bytearray[i]);
    }

}

//codingan dari crypto_aead_encrypt 
// https://github.com/rweather/lwc-finalists/blob/master/src/arduino/LWC_Benchmark/benchmark.cpp

// static unsigned char const nonce[24] = {
//     0x00, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07,
//     0x08, 0x09, 0x0A, 0x0B, 0x0C, 0x0D, 0x0E, 0x0F,
//     0x10, 0x11, 0x12, 0x13, 0x14, 0x15, 0x16, 0x17,
// };

// static unsigned char const key[32] = {
//     0x00, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07,
//     0x08, 0x09, 0x0A, 0x0B, 0x0C, 0x0D, 0x0E, 0x0F,
//     0x10, 0x11, 0x12, 0x13, 0x14, 0x15, 0x16, 0x17,
//     0x18, 0x19, 0x1A, 0x1B, 0x1C, 0x1D, 0x1E, 0x1F
// };

void setup() {
  Serial.begin(9600);

  unsigned long long  plen;
  unsigned long long  clen;
  unsigned long long  nsec;
  unsigned long long  adlen;
  unsigned long long  mlen = 50;
  // unsigned long long  mlen = 128; // keluaran cipher 136 bytes
  // unsigned long long  mlen = 152; //keluaran cipher 160 bytes
  unsigned char       cipher[160];
  unsigned char       plaintext[160];
  unsigned char       chex[160]; 
  String tuliskan; //variable untuk input
  const unsigned char ad[] = "000102030405060708090A0B0C0D0E0F10111213141516171819";
  const unsigned char key[] = "000102030405060708090A0B0C0D0E0F";
  const unsigned char nonce[] = "000102030405060708090A0B"; 
  unsigned long time1;
  unsigned long time2;
  unsigned long waktu;

  Serial.print("Plaintext: ");
                  
    while (Serial.available() == 0) {}  
    tuliskan = Serial.readString();
    Serial.println(tuliskan);

  unsigned char message[sizeof(tuliskan)];
    for(int i=0; i<sizeof(tuliskan); i++){
    message[i] = tuliskan[i];
  }

// crypto_aead_encrypt(cipher, &clen, message, mlen, 0, 0, 0, nonce, key);
//   crypto_aead_encrypt(cipher, &clen, message, mlen, ad, &adlen, nsec, nonce, key);
  crypto_aead_encrypt(cipher, &clen, message, mlen, 0, 0, 0, nonce, key);

  string2hexString(cipher,clen,chex);
  //untuk timer awal 
  Serial.print("Time awal: ");
  time1 = micros();
  Serial.println(time1); //prints time since program started
  Serial.print("Key: ");
  Serial.println((char *)key);
  Serial.print("Nonce: ");
  Serial.print((char *)nonce);
  Serial.println();
  Serial.print("Ad: ");
  Serial.print((char *)ad);
//   for (int i = 0; i < 16; i++) 
{
    // print the current element as a hexadecimal value
 
        // Serial.print(nonce[i], HEX);
  }

//   Serial.print("Plaintext before decrypt: ");
//   Serial.println((char *)plaintext);
// //   unsigned char plaintext[160];
//   unsigned long long plen;
//   //codingan dari https://rweather.github.io/lwc-finalists/elephant-dumbo_8h.html#aba3b4f28013762fa154494b4de96ed9f 
  crypto_aead_decrypt(plaintext, &plen, 0, cipher, clen, 0, 0, nonce, key);
  Serial.println();

//untuk output
  Serial.print("Cipher: ");
  Serial.println((char *)chex);
  Serial.print("Plaintext after decrypt: ");
  Serial.println((char *)plaintext);
  //untuk timer end 
  Serial.print("Time End:");
  time2 = micros();
  Serial.println(time2); //prints time since program started
  
  Serial.print ("Encryption time:");
  waktu = (time2) - (time1) ;
  // waktu = micros;
  Serial.println(waktu);
  

}
void loop() {
}