#include "contiki.h"
#include "ecc.h"
#include "ecc_dh.h"
#include "ecc_dsa.h"
#include "test_ecc_utils.h"
#include "test_uti.h"
#include "myencrypt.h"
#include "sha256.h"
#include "sys/energest.h"
#include "base64.h"

#include "random.h"
#include <assert.h>
#include <stdio.h> 
#include <stdlib.h>
#include <string.h>
#include <stdint.h>
#include <stddef.h>
#include <stdbool.h>
#include <clock.h>


PROCESS(Experiments, "Experimentation Section");
AUTOSTART_PROCESSES(&Experiments);

static unsigned long
to_seconds(uint64_t time)
{
	return (unsigned long)(time/ENERGEST_SECOND);
}

struct GenKeyPair
{
	uint8_t privkey[NUM_ECC_BYTES];
	uint8_t pubkey[2 * NUM_ECC_BYTES];
};


//Function to generate the ECC keypair. Call it when initializing either client.
struct GenKeyPair ECC_GenKeyPair()
{
	static struct GenKeyPair genpair1;

	const struct uECC_Curve_t * curve = uECC_secp256r1();

	clock_time_t start_time = clock_time();
	uECC_make_key(genpair1.pubkey, genpair1.privkey, curve);
	clock_time_t end_time = clock_time();

	unsigned long time_taken = end_time - start_time;
	printf("Time Taken to Generate Key Pair: %lu ticks\n", time_taken);
	
	show_str("ECC Private Key:", genpair1.privkey, sizeof(genpair1.privkey));
  	show_str("ECC Public Key:", genpair1.pubkey, sizeof(genpair1.pubkey));
  	printf("\n");


	return genpair1;
}


//Function to generate the shared secret key after exchange of public keys
// input:
//		user1 private key (privkey);
//		user2 public key (pubkey);
//
//output:
//		shared secret key (secret);
//
int ECDH_GenSharedKey(uint8_t *privkey, uint8_t *pubkey, uint8_t *secret, int secretlen)
{
	const struct uECC_Curve_t * curve = uECC_secp256r1();

	clock_time_t start_time = clock_time();
	uECC_shared_secret(pubkey, privkey, secret, curve);
	clock_time_t end_time = clock_time();

	unsigned long time_taken = end_time - start_time;
	printf("Time Taken to Generate Shared Secret: %lu ticks\n", time_taken);

	show_str("ECDH Shared Key: ", secret, secretlen);
	printf("\n\n");

	return 0;
}

//Function to perform AES 128 symmetric encryption 
// input:
//		message to be encrypted (m);
//		the key to be used for by the AES scheme (s);
//
//output:
//		ciphertext (ciphertext);
//

int AES_Encrypt(uint8_t *m, uint8_t * ciphertext, TCAesKeySched_t s, int len)
{
	clock_time_t start_time = clock_time();
	tc_aes_encrypt(ciphertext, m, s);
	
	clock_time_t end_time = clock_time();
	
	unsigned long time_taken = end_time - start_time;
	printf("Time Taken to Encrypt Message: %lu ticks\n", time_taken);

	printf("Start: %lu ticks\n", start_time);
	printf("End: %lu ticks\n", end_time);

	show_str("Encrypted Message: ", ciphertext, len);
	printf("\n");

	return 0;
}

//Function to perform AES 128 symmetric decryption 
// input:
//		ciphertext to be decrypted (ciphertext);
//		the key to be used for by the AES scheme (s);
//
//output:
//		original message before encryption (plaintext);
//

int AES_Decrypt(uint8_t * ciphertext, uint8_t * plaintext, TCAesKeySched_t s, int len)
{
	clock_time_t start_time = clock_time();
	tc_aes_decrypt(plaintext, ciphertext, s);
	
	clock_time_t end_time = clock_time();
	
	unsigned long time_taken = end_time - start_time;
	printf("Time Taken to Decrypt Message: %lu ticks\n", time_taken);

	printf("Start: %lu ticks\n", start_time);
	printf("End: %lu ticks\n", end_time);

	show_str("Decrypted Message: ", plaintext, len);
	printf("\n");

	return 0;
}

//Function to perform ECDSA signing algorithm 
// input:
//		user's private key (privkey);
//		message to be signed (msg)
//
//output:
//		signed message tag (signtext);

int ECDSA_sign(uint8_t * privkey, char * msg, uint8_t * signtext, int signlen)
{
	uint8_t hashtag[32];
	const struct uECC_Curve_t * curve = uECC_secp256r1();
	struct tc_sha256_state_struct s;
	(void)tc_sha256_init(&s);
	tc_sha256_update(&s, (const uint8_t *) msg, strlen(msg));
    (void)tc_sha256_final(hashtag, &s);
	

	clock_time_t start_time = clock_time();
	uECC_sign(privkey, hashtag, sizeof(hashtag), signtext, curve);
	clock_time_t end_time = clock_time();
	
	unsigned long time_taken = end_time - start_time;
	printf("Time Taken to Sign Message: %lu ticks\n", time_taken);

	show_str("Signed Message: ", signtext, signlen);
	printf("\n");

	return 0;
}


//Function to perform ECDSA signature verification algorithm 
// input:
//		message sender's pubkey key (pubkey);
//		message to be signed (msg);
// 		signature tag to be verified (signtext);
//
//output:
//		boolean (1 for true and 0 for false);

int ECDSA_verify(uint8_t * pubkey, char * msg, uint8_t * signtext)
{
	uint8_t hashtag[32];
	const struct uECC_Curve_t * curve = uECC_secp256r1();
	struct tc_sha256_state_struct s;
	(void)tc_sha256_init(&s);
	tc_sha256_update(&s, (const uint8_t *) msg, strlen(msg));
    (void)tc_sha256_final(hashtag, &s);

	clock_time_t start_time = clock_time();
	if(!uECC_verify(pubkey, hashtag, sizeof(hashtag), signtext, curve))
	{
		printf("Verification Unsuccessful! Check Signature!\n");
	}
	else
	{
		
		printf("Successful Verification! Signature is genuine\n");
	}
	clock_time_t end_time = clock_time();	
	unsigned long time_taken = end_time - start_time;
	printf("Time Taken to Verify Signature: %lu ticks\n", time_taken);

	return 0;
}



PROCESS_THREAD(Experiments, ev, data)
{
  	PROCESS_BEGIN();
  	struct GenKeyPair keypair1;
  	struct GenKeyPair keypair2;

  	uint8_t ciphertext[16];
  	uint8_t plaintext[16];
  	uint8_t key[NUM_ECC_BYTES];


  	/* Setup of the Cryptographically Secure PRNG. */
  	srand(8977);
  	uECC_set_rng(&default_CSPRNG);
  	struct tc_aes_key_sched_struct s;

  	char *msg = "Variable message for Testing";

  	
  	// for loop to test functions for a specified number of times.
  	//
  	for (int i = 0; i < 5; ++i)
  	{
  		
  		printf("Round: %d\n", i);
  		//Call the generate key pair function to generate 2 pairs of keys
  		//
  		keypair1  = ECC_GenKeyPair();
  		keypair2  = ECC_GenKeyPair();

  		//Call the generate shared key function to generate a shared key based on 2 public keys
  		ECDH_GenSharedKey(keypair1.privkey, keypair2.pubkey, key, sizeof(key));

  		//Set the AES128 key to be used for Encryption
  		tc_aes128_set_encrypt_key(&s, key);
  		

  		//Perform AES128 encryption on a message using the specified key
  		AES_Encrypt((uint8_t *)msg, ciphertext, &s, sizeof(ciphertext));

  		//Perform AES128 decryption on the ciphertext from the encryption function
  		AES_Decrypt(ciphertext, plaintext, &s, sizeof(plaintext));


  		//Sign the message with one of the generated private keys.
		ECDSA_sign(keypair1.privkey, msg, signtext, sizeof(signtext));

  		//Verify the signature tag generated by the signing algorithm  		
  		ECDSA_verify(keypair2.pubkey, msg, signtext);

  		printf("\n");

  		//Initialize the energest module
  		//Used by Contiki-NG to measure the energy utilization of the board
  		energest_flush();

  		//Print out the energest results on the serial display
	    printf("\nEnergest:\n");
	    printf(" CPU          %4lus LPM      %4lus DEEP LPM %4lus  Total time %lus\n",
	           to_seconds(energest_type_time(ENERGEST_TYPE_CPU)),
	           to_seconds(energest_type_time(ENERGEST_TYPE_LPM)),
	           to_seconds(energest_type_time(ENERGEST_TYPE_DEEP_LPM)),
	           to_seconds(ENERGEST_GET_TOTAL_TIME()));
	    printf(" Radio LISTEN %4lus TRANSMIT %4lus OFF      %4lus\n",
	           to_seconds(energest_type_time(ENERGEST_TYPE_LISTEN)),
	           to_seconds(energest_type_time(ENERGEST_TYPE_TRANSMIT)),
	           to_seconds(ENERGEST_GET_TOTAL_TIME()
	                      - energest_type_time(ENERGEST_TYPE_TRANSMIT)
	                      - energest_type_time(ENERGEST_TYPE_LISTEN)));

	    printf("\n");
  	}

  	PROCESS_END();
}