// Knows only protobuf_sgx objects, protobuf header.
// For socket programming
#include <sys/socket.h>
#include <stdlib.h>
#include <netinet/in.h>
#include <string.h>
#include <errno.h>
#include<unistd.h>
#include <stdio.h>
#include "ProtobufLAMessages.pb.h"
#include <google/protobuf/io/coded_stream.h>
#include <google/protobuf/io/zero_copy_stream_impl.h>
using namespace google::protobuf::io;
#include "protobufLAInitiator.h"
#include "../Decryptor/Decryptor_u.h"
#include "sgx_tcrypto.h"
#include <iostream>

// TODO: Make these private functions
int apache_fd;
int decrypt_client_data_wrapper(unsigned char* op_plaintext , uint32_t own_enclave_id);


int read_protobuf_msg_from_fd(int accept_fd, google::protobuf::MessageLite& message)
{
  ZeroCopyInputStream* raw_input;
  CodedInputStream* coded_input;
  uint32_t size;
  CodedInputStream::Limit limit;
  raw_input = new FileInputStream(accept_fd);
  coded_input = new CodedInputStream(raw_input);
  if(!coded_input->ReadVarint32(&size))
  {
    printf("Error in reading size of msg");
    fflush(stdout);
    return -1;
  }
  //printf("size of msg was read to be %" PRIu32 " \n", size);
  fflush(stdout);
  limit = coded_input->PushLimit(size);
  if(!message.ParseFromCodedStream(coded_input))
  {
    printf("Error in parsing msg");
    fflush(stdout);
    return -1;
  }
  coded_input->PopLimit(limit);
  delete raw_input;
  delete coded_input; 
  return 0;
}

// TODO: private functions
int write_protobuf_msg_to_fd(int accept_fd, google::protobuf::MessageLite& message)
{
  ZeroCopyOutputStream* raw_output = new FileOutputStream(accept_fd);
  CodedOutputStream* coded_output  = new CodedOutputStream(raw_output);
  coded_output->WriteVarint32(message.ByteSize());
  if(!message.SerializeToCodedStream(coded_output))
  {
    printf("SerializeToCodedStream failed");
    fflush(stdout);
    return -1;
  }
  // As per this - https://stackoverflow.com/questions/22881876/protocol-buffers-how-to-serialize-and-deserialize-multiple-messages-into-a-file?noredirect=1&lq=1
  // TODO: There may be a better way to do this - 1) this happens with every accept now and 2) make it happen on the stack vs heap - destructor will be called on return from this function (main) and the items will then be written out. (We probably don't want that, actually)
  delete coded_output;
  delete raw_output;
  fflush(stdout);
  return 0;
}

// Sets up a socket to bind and listen to the given port. Returns FD of the socket on success, -1 on failure (and prints a msg to stdout with the errno)
int set_up_socket(int port,   sockaddr_in* address)
{
  int server_fd = 0;

  // Creating socket file descriptor for listening for attestation requests.
  server_fd = socket(AF_INET, SOCK_STREAM | SOCK_CLOEXEC, 0);
  if (server_fd == -1)
  {
     printf("Error in creating a socket - %d", errno);
      return -1;
  }

  // Preparing the address struct for binding
  address->sin_family = AF_INET;
  address->sin_addr.s_addr = INADDR_ANY; // Todo: should this be localhost?
  address->sin_port = htons(port);
  //  memset(address->sin_zero,0,sizeof(address->sin_zero));

  socklen_t addrlen = sizeof(*address);
  // Binding
  if (bind(server_fd, (sockaddr*)address, addrlen)<0)
  {
      printf("Error in binding %d - port was %d - ", errno, port);
      return -1;
  }

  // Listening
  if (listen(server_fd, 128) < 0)
  {
      printf("Error in listening %d", errno);
      return -1;
  }

  return server_fd;
}

int local_attestation_initiator(int port, uint32_t own_enclave_id)
{
  // declare msg1, msg2, msg3 protobuf objects
  protobuf_sgx_dh_msg1_t protobuf_msg1;
  protobuf_sgx_dh_msg2_t protobuf_msg2;
  protobuf_sgx_dh_msg3_t protobuf_msg3;
  protobuf_post_LA_encrypted_msg_t protobuf_encrypted_msg; 
  uint32_t protobuf_sgx_ret; uint32_t sgx_ret;

  // For socket to listen to the Apache enclave.
  int server_fd=0; int accept_fd = 0;
  struct sockaddr_in own_addr;
  struct sockaddr_storage apache_addr; socklen_t apache_addr_size = sizeof(apache_addr);
  uint32_t session_id;  uint8_t read_or_write;
  size_t bytes_read_post_la; uint8_t encrypted_apache_mrsigner_and_tag[48];  
  size_t bytes_written_post_la; 
  //  int counter;
  server_fd=set_up_socket(port, &own_addr);
  if(server_fd==-1)
      return -1;
  printf("Successfully set up a socket to communicate with the Apache enclave.\n");
  fflush(stdout);

  protobuf_sgx_ret = generate_protobuf_dh_msg1(own_enclave_id, protobuf_msg1, &session_id);
  if(protobuf_sgx_ret != 0)
  {
    printf("Error in generate_protobuf_dh_msg1: 0x%x", protobuf_sgx_ret); fflush(stdout); return protobuf_sgx_ret;
  }


  accept_fd = accept(server_fd, (struct sockaddr *)&apache_addr,&apache_addr_size);
  if (accept_fd <0)
  {
        printf("Error in accepting %d", errno);
        return -1;
  }
  printf("Accepted fd\n"); fflush(stdout); 

  if(write_protobuf_msg_to_fd(accept_fd, protobuf_msg1)!=0)
    return -1;

  if(read_protobuf_msg_from_fd(accept_fd, protobuf_msg2)!=0)
    return -1;

  protobuf_sgx_ret = process_protobuf_dh_msg2_generate_protobuf_dh_msg3(own_enclave_id, protobuf_msg2, protobuf_msg3, &session_id, &read_or_write);
  if(protobuf_sgx_ret != 0)
  {
      printf("Error in generate_protobuf_dh_msg2: 0x%x", protobuf_sgx_ret); fflush(stdout); return protobuf_sgx_ret;
  }

  if(write_protobuf_msg_to_fd(accept_fd, protobuf_msg3)!=0)
    return -1;
//  read_or_write=0;  

  printf("Here\n"); fflush(stdout); 
  if(read_or_write)
  {  
  	bytes_read_post_la=read(accept_fd, encrypted_apache_mrsigner_and_tag, 32); 
	  if(bytes_read_post_la!=32)
 	 {
       	printf("Not all of the encrypted apache's mrsigner was read from the verifier.\n"); fflush(stdout); return 0xfe;
 	 }
	  bytes_read_post_la=read(accept_fd, encrypted_apache_mrsigner_and_tag+32, 16);
	  if(bytes_read_post_la!=16)
	  {
       printf("Not all of the encrypted apache's mrsigner **tag** was read from the verifier.\n"); fflush(stdout); return 0xfe;
  	}
  	uint32_t count;
  	for(count=0;count<48;count++)
		printf("0x%02x ", encrypted_apache_mrsigner_and_tag[count]);
  	printf("\n");fflush(stdout); 
//  sgx_ret=decrypt_wrapper(own_enclave_id, encrypted_apache_mrsigner_and_tag, 32, encrypted_apache_mrsigner_and_tag+32 , plaintext); 
Decryptor_decrypt_verifiers_message_set_apache_mrsigner(own_enclave_id, &sgx_ret, encrypted_apache_mrsigner_and_tag, encrypted_apache_mrsigner_and_tag+32);
 	 if(sgx_ret!=0)
	  {
	printf("Error in decryption: 0x%x\n", sgx_ret); fflush(stdout); return sgx_ret; 
	  }
	  printf("Successful decryption\n"); fflush(stdout); 
	

  }
  else 
  {
	
	apache_fd=accept_fd;
	uint8_t encrypted_sign_data_and_sign_and_tag[176];
	memset(encrypted_sign_data_and_sign_and_tag,0x0,176);
	uint8_t plaintext_sign_data_and_sign[160];
	uint8_t plaintext_priv_key[32];	
	sgx_ec256_signature_t sig2;
	Decryptor_create_and_encrypt_mitigator_header_value(own_enclave_id, &sgx_ret, plaintext_sign_data_and_sign, encrypted_sign_data_and_sign_and_tag , encrypted_sign_data_and_sign_and_tag+160, plaintext_priv_key,&sig2);
	if(sgx_ret!=0)
	{
		printf("Error in generating encrypted mitigator header:0x%x\n", sgx_ret); fflush(stdout); return 0xf3;
	}
	uint32_t count; 
	for(count=0;count<160;count++)
	{
		printf("0x%02x ", encrypted_sign_data_and_sign_and_tag[count]);
	}
	printf("\n"); fflush(stdout); 
	printf("Plaintext Signature data:\n"); fflush(stdout); 
        for(count=0;count<96;count++)
        {
                printf("0x%02x ", plaintext_sign_data_and_sign[count]);

        }
	printf("\n"); fflush(stdout);
	
	printf("Plaintext signature: \n"); fflush(stdout); 
        for(count=0;count<32;count++)
        {
                printf("%02x", plaintext_sign_data_and_sign[count+96]);
        }
	printf("\n"); fflush(stdout);
        for(count=32;count<64;count++)
        {
                printf("%02x", plaintext_sign_data_and_sign[count+96]);

        }

        printf("\n"); fflush(stdout);
	printf("Heres the private key used to sign this \n"); // TODO: Remove this printf and the private key parts to the ecall
	for(count=0;count<32;count++)
		printf("%02x", plaintext_priv_key[31-count]);
        printf("\n"); fflush(stdout);
	for(count=0;count<8;count++)
	{
		printf("%02x ", sig2.x[count]);
	}
	printf("\n"); fflush(stdout);

        for(count=0;count<8;count++)
        {
                printf("%02x ", sig2.y[count]);
        }
        printf("\n"); fflush(stdout);
        
	protobuf_encrypted_msg.set_msg((void*)encrypted_sign_data_and_sign_and_tag, 176); 
	if(write_protobuf_msg_to_fd(apache_fd, protobuf_encrypted_msg) != 0)
         {
        printf("Not all of the decryptor's signature was written to the Apache.\n"); fflush(stdout); return 0xfe;
         }
	unsigned char op_plaintext[160];
	do {
//	sleep(100);
	decrypt_client_data_wrapper(op_plaintext, own_enclave_id); 
//	sleep(100); 
	} while(true); 


  }

        printf("Successfully done Local attestation\n");
        fflush(stdout);
        return 0;
      
}


// decrypt_client_data function that does a read-then-write loop 
// say msg length amount of bytes are read
////// TODO: first need to decrypt all data received with the aes-gcm apache_iv 
////// Decryptor.cpp function gets the first 64 bytes as first arg and the rest of ciphertext as 2nd arg. 
// then need to call the sgx's compute_shared_key fn for ECDH key - on first 64 bytes. 
// print this key
// call sha256 on it
// TODO: decrypt rest of msglength - 64 bytes using this key. (set IV in crypto function)
////// TODO: then encrypt all user data with apache's key 
////// return here
/// write data.
int decrypt_client_data_wrapper(unsigned char* op_plaintext , uint32_t own_enclave_id)
{
        protobuf_post_LA_encrypted_msg_t protobuf_encrypted_msg; uint32_t sgx_ret_status;

        printf("Reading msg from apache"); fflush(stdout); 
        if(read_protobuf_msg_from_fd(apache_fd, protobuf_encrypted_msg)!=0)
        {
                printf("Not all of the Apache's message was read\n"); fflush(stdout); return 0xfe; 
        }
   
        std::string protobuf_encrypted_msg_string(protobuf_encrypted_msg.msg());
//	printf("%s\n", protobuf_encrypted_msg_string.c_str()); fflush(stdout);
	const char* temp_array = protobuf_encrypted_msg_string.c_str();
	int counter;
	for(counter=0;counter<protobuf_encrypted_msg_string.length();counter++)
		printf("%d ",temp_array[counter]);
	printf("\n"); fflush(stdout); 

        unsigned char client_data_from_apache[160+64]; //protobuf_encrypted_msg_string.size());      // set a large limit for "size" for third argument - client data - of 160 - so 160+64
	memset(client_data_from_apache, 0, 160+64); 
	unsigned char* protobuf_string = (unsigned char*) protobuf_encrypted_msg_string.c_str(); 
//        client_data_from_apache = (unsigned char*) protobuf_encrypted_msg_string.c_str(); //temp_array;//(unsigned char*)&(std::vector<char> (protobuf_encrypted_msg_string.begin(), protobuf_encrypted_msg_string.end()))[0]; // This shit eats up 8 starting bytes - sets them to 0s
	for(counter=0;counter<protobuf_encrypted_msg_string.length();counter++)
		client_data_from_apache[counter]=*(protobuf_string+counter);

	uint32_t ciphertext_length = protobuf_encrypted_msg_string.length()-64; // we hope it's greater than 0 

        printf("Key in big endian form:\n"); fflush(stdout); 
        for(counter=0; counter<64; counter++)
                printf("0x%02x ", *(client_data_from_apache + counter)); 
        printf("\n"); fflush(stdout); 

	unsigned char client_data_to_apache[160+64]; //ciphertext_length); 
	// Just so that the ciphertext is returned back to Apache in case decrypt_client_data fails. 
	for(counter=0;counter<ciphertext_length;counter++)
		client_data_to_apache[counter]=client_data_from_apache[counter+64];
	uint8_t clen;
	Decryptor_decrypt_client_data(own_enclave_id, &sgx_ret_status, client_data_from_apache, ciphertext_length, client_data_from_apache + 64, client_data_to_apache, &clen); 
	if(sgx_ret_status != 0)
	{
		printf("decrypt_client_data returned :0x%x\n", sgx_ret_status); fflush(stdout); //return sgx_ret_status; 
	}

	// TODO: REMOVE THIS -
//	for(counter=0;counter<ciphertext_length;counter++)
//		client_data_to_apache[counter]=0x46;//client_data_from_apache[counter+64];
	// TODO: Print the bytesize of the message. should be ciphertext_length=4 + more bytes for protobuf packing. 
	printf("About to write the following bytes to the apache\n"); fflush(stdout);
	for(counter=0;counter<ciphertext_length;counter++)
		printf("0x%02x ", client_data_to_apache[counter]);
	printf("\n"); fflush(stdout);
	printf("Following shared secret was established.\n"); 
	for(counter=ciphertext_length;counter<32+ciphertext_length;counter++)
                printf("0x%02x ", client_data_to_apache[counter]);
        printf("\n"); fflush(stdout);
        printf("Following key was derived.\n"); 
        for(counter=ciphertext_length+32;counter<64+ciphertext_length;counter++)
                printf("0x%02x ", client_data_to_apache[counter]);
	printf("\n"); fflush(stdout);
	printf("Ciphertext for the string Miti \n"); 
        for(counter=ciphertext_length+64;counter<64+ciphertext_length + 4;counter++)
		printf("%d ", client_data_to_apache[counter]);
        printf("\n"); fflush(stdout); 
        printf("Tag for the string Miti \n"); 
        for(counter=ciphertext_length+64+4;counter<80+ciphertext_length + 4;counter++)
		printf("%d ", client_data_to_apache[counter]);
        printf("\n"); fflush(stdout); 


        protobuf_encrypted_msg.set_msg((void*)  client_data_to_apache, ciphertext_length);// Is this message set tho?
	if(write_protobuf_msg_to_fd(apache_fd, protobuf_encrypted_msg)!=0)
	{
		printf("Not all of the protobuf message was written to Apache.\n"); fflush(stdout); return 0xfe;
	}
	printf("Wrote data to Apache of length %d \n", protobuf_encrypted_msg.ByteSize());fflush(stdout);
	return 0; 
}