// Knows only protobuf_sgx objects, protobuf header.
#include <stdlib.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 "SgxProtobufLAInitiator.h"
#include "crypto.h"
// For socket programming
#include <arpa/inet.h>
#include <sys/socket.h>
#include <netinet/in.h>

int global_decryptor_fd;


// TODO: Make these private functions
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\n");
    fflush(stdout);
    return -1;
  }
  printf("size of msg was read to be %d \n", size);
  fflush(stdout);
  limit = coded_input->PushLimit(size);
  if(!message.ParseFromCodedStream(coded_input))
  {
    printf("Error in parsing msg\n");
    fflush(stdout);
    return -1;
  }
  printf("Done parsing msg\n"); fflush(stdout);
  coded_input->PopLimit(limit);
  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\n");
    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 connected to the port passed as input - returns the socket FD on success and -1 on error.
// Also prints the errno on error.
int set_up_socket_connect(int port)
{
  int sock = 0;
  if ((sock = socket(AF_INET, SOCK_STREAM, 0)) < 0)
  {
    printf("\n Error in socket call - errno is %d \n", errno);
    return -1;
  }

  struct sockaddr_in serv_addr;
  memset(&serv_addr, '0', sizeof(serv_addr));

  serv_addr.sin_family = AF_INET;
  serv_addr.sin_port = htons(port);

  // Convert IPv4 and IPv6 addresses from text to binary form
  if(inet_pton(AF_INET, "127.0.0.1", &serv_addr.sin_addr)<=0)
  {
    printf("\nError in inet_pton - errno is %d\n", errno);
    return -1;
  }

  if (connect(sock, (struct sockaddr *)&serv_addr, sizeof(serv_addr)) < 0)
  {
    printf("\nError in connect - errno is %d \n", errno);
    return -1;
  }

  return sock;
}

int local_attestation_initiator(int port, unsigned char* base64_encoded_token)
{
  // 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;
  uint32_t protobuf_sgx_ret;

//  uint8_t encrypted_hash[32]; uint8_t encrypted_tag[16]; size_t post_la_bytes_written;
//	uint8_t tokenT_and_tag[176]; size_t bytes_read;

  // For socket to listen to the Apache enclave.
//  int server_fd=0; int accept_fd = 0;
//   std::string protobuf_encrypted_msg_string(protobuf_encrypted_msg.msg());
  // memcpy(tokenT_and_tag, protobuf_encrypted_msg_string.c_str(), 176);
//  struct sockaddr_in own_addr;
//  struct sockaddr_storage apache_addr; socklen_t apache_addr_size = sizeof(apache_addr);
  uint32_t session_id;
  //  int counter;

  int decryptor_fd;
  setbuf(stdout,NULL);
  decryptor_fd=set_up_socket_connect(port);
  if(decryptor_fd == -1)
  {
    perror("\nCould not set up the socket: had the following error: ");
    fflush(stderr);
  }
//  printf("");
  if(read_protobuf_msg_from_fd(decryptor_fd, protobuf_msg1)!=0)
    return -1;

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

  if(write_protobuf_msg_to_fd(decryptor_fd, protobuf_msg2)!=0)
    return -1;

  if(read_protobuf_msg_from_fd(decryptor_fd, protobuf_msg3)!=0)
    return -1;

  protobuf_sgx_ret = process_protobuf_dh_msg3(protobuf_msg3, &session_id);
  if(protobuf_sgx_ret != 0)
  {
      printf("Error in process_protobuf_dh_msg3: 0x%x", protobuf_sgx_ret); fflush(stdout); return protobuf_sgx_ret;
  }
          unsigned char tokenT_and_tag[176]; int protobuf_encrypted_msg_size;


          protobuf_post_LA_encrypted_msg_t protobuf_encrypted_msg;
        if(read_protobuf_msg_from_fd(decryptor_fd, protobuf_encrypted_msg)!=0)
   {
	printf("Not all of the post-LA message was read\n"); fflush(stdout); return 0xfe;
   }
   std::string protobuf_encrypted_msg_string(protobuf_encrypted_msg.msg());
   memcpy(tokenT_and_tag, protobuf_encrypted_msg_string.c_str(), 176);

   printf("Read encrypted signature and tag from the decryptor socket.\n"); fflush(stdout);
    uint32_t count;
	printf("Encrypted data:\n");
	for(count=0;count<176;count++)
		printf("0x%02x ", tokenT_and_tag[count]);
	printf("\n"); fflush(stdout);
	uint8_t plaintext_token[160]; int plaintext_length=0;
	uint32_t ret_decrypt=aes_gcm_wrapper(0, tokenT_and_tag, 160, plaintext_token, &plaintext_length, tokenT_and_tag+160);
//int aes_cipher(int enc, unsigned char *key, unsigned char *iv, unsigned char* plaintext, int plaintext_len, unsigned char *ciphertext,  int* op_ciphertext_len, unsigned char* tag);
	if(ret_decrypt!=0)
		{
			printf("Errorin decryption 0x%x", ret_decrypt); fflush(stdout); return 0xed;
		}
        printf("Decryption keypair\n");
        for(count=0;count<64;count++)
                printf("0x%02x ", plaintext_token[count]);
	printf("\n"); fflush(stdout);
	printf("Decryption verifier mrenclave\n");
	for(count=64;count<96;count++)
                printf("0x%02x ", plaintext_token[count]);
	printf("\n"); fflush(stdout);
	count=base64_encoding_wrapper(plaintext_token, base64_encoded_token, 160);
	if(count != 216)
	{
		printf("Somehow not the entire token was encoded in base64:0x%x\n", count); fflush(stdout); return 0x55;
	}
	global_decryptor_fd=decryptor_fd;
	printf("Successfully done Local attestation\n");
	fflush(stdout);
	return 0;
}

int extract_binary_key_and_client_ciphertext_data_from_base64(unsigned char* base64_encoded_client_pub_key, unsigned char* base64_ciphertext, uint32_t base64_ciphertext_length, unsigned char* op_client_public_key_and_ciphertext, uint32_t* op_ciphertext_length)
{
	int openssl_ret_status;
        openssl_ret_status=base64_decoding_wrapper(base64_encoded_client_pub_key, op_client_public_key_and_ciphertext, 88);
        if(openssl_ret_status == -1)
        	return 0xfe;

	*op_ciphertext_length = base64_decoding_wrapper(base64_ciphertext, op_client_public_key_and_ciphertext + 64, base64_ciphertext_length);
	if(*op_ciphertext_length == -1)
	{
		return 0x33;
	}
	return 0;

}

// should be called with the base64encoded public key and data to be decrypted
// base64 encoded public key should be of length 4*(ceil(64/3)) = 88 (at least)
int decrypt_client_data_through_decryptor(unsigned char* base64_client_public_key, unsigned char* base64_ciphertext_client_data, uint32_t base64_ciphertext_client_data_length, unsigned char* plaintext_client_data, uint32_t* op_plaintext_client_data_length)
{
	int ret_status;
	uint32_t count;
	unsigned char* client_data_to_decryptor;
	unsigned char* client_data_encrypted_to_decryptor;
	unsigned char* client_data_from_decryptor;
	unsigned char tag[16]; 
	const char* msg;
	protobuf_post_LA_encrypted_msg_t protobuf_encrypted_msg;
	int output_length;
	int input_length;
	uint32_t ciphertext_length;
	for(count=0; count<base64_ciphertext_client_data_length; count++)
		plaintext_client_data[count] = base64_ciphertext_client_data[count];
	*op_plaintext_client_data_length = base64_ciphertext_client_data_length;

	client_data_to_decryptor = (unsigned char*) malloc(64 + (base64_ciphertext_client_data_length * 3/4));
	// If ciphertext_length != base64_ciphertext_client_data_length * 3/4 then it will write outside the allocated memory
	ret_status = extract_binary_key_and_client_ciphertext_data_from_base64(base64_client_public_key, base64_ciphertext_client_data, base64_ciphertext_client_data_length, client_data_to_decryptor, &ciphertext_length);
	if(ret_status != 0)
		return ret_status;

	input_length = ciphertext_length + 64; // normal case - no writing outside the allocated memory.
	printf("Public key and ciphertext:\n"); fflush(stdout);
	for(count = 0; count<input_length; count++)
		printf("0x%02x ", client_data_to_decryptor[count]);
	printf("\n"); fflush(stdout);

/*	client_data_encrypted_to_decryptor = (unsigned char*) malloc(input_length + 16);
	// TODO: try to call aes_gcm  *in encryption mode* on ciphertext data of ciphertext_length (where is the iv?)
	// if output_length != ciphertext_length + 64 then tag will be overwritten.
	ret_status = aes_gcm_wrapper(1, client_data_to_decryptor, input_length, client_data_encrypted_to_decryptor, &output_length, tag);
	if(ret_status != 0 || output_length != input_length)
	{
		free(client_data_to_decryptor);
		free(client_data_encrypted_to_decryptor);
		return ret_status;
	}
	for(count=0; count<16; count++)
		client_data_encrypted_to_decryptor[output_length+count]=tag[count]; 
	// TODO: Increment AESGCM IV here.
*/
	protobuf_encrypted_msg.set_msg((void*)client_data_to_decryptor, input_length);
	free(client_data_to_decryptor);
//	free(client_data_encrypted_to_decryptor);
/*
	printf("Writing this to the decryptor: \n"); fflush(stdout); 
	for(count=8;count<output_length+16+8;count++)
		printf("0x%02x ", client_data_to_decryptor[count]);
	printf("\n");	fflush(stdout);
*/

	// write message to decryptor
	if(write_protobuf_msg_to_fd(global_decryptor_fd, protobuf_encrypted_msg)!=0)
	{
		printf("Not all of the client's pub key and ciphertext data was written\n"); fflush(stdout); return 0xfe;
	}

	// clear the message before setting it again
	protobuf_encrypted_msg.clear_msg();

	// read encrypted data
	printf("Reading msg from decryptor"); fflush(stdout);
	if(read_protobuf_msg_from_fd(global_decryptor_fd, protobuf_encrypted_msg)!=0)
	{
		printf("Not all of the decryptor's message was read\n"); fflush(stdout); return 0xf3;
	}
	printf("Done reading msg from decryptor\n"); fflush(stdout);

	// define plaintext array
	input_length = protobuf_encrypted_msg.msg().length();
	msg = protobuf_encrypted_msg.msg().c_str(); 
	for(count=0; count<input_length; count++)
		plaintext_client_data[count] = msg[count]; 
	*op_plaintext_client_data_length = input_length; 
	/* client_data_from_decryptor=(unsigned char*)malloc(input_length);
	memcpy(client_data_from_decryptor, protobuf_encrypted_msg.msg().c_str(), input_length); //TODO: THIS might be a problem - c_str might terminate at null bytes?

	// TODO: decrypt data received in the protobuf_encrypted_msg, set bytes at plaintext_client_data to it.
	for(count=0;count<input_length;count++)
	{
		plaintext_client_data[count]=client_data_from_decryptor[count]; 
		printf("0x%02x ", client_data_from_decryptor[count]);
	}
	printf("\n"); fflush(stdout);
	*op_plaintext_client_data_length=input_length; 

	ret_status = aes_gcm_wrapper(0, client_data_from_decryptor, input_length, plaintext_client_data, &output_length, client_data_from_decryptor + input_length - 16);
	if(ret_status == 0)
		*op_plaintext_client_data_length = output_length;
	free(client_data_from_decryptor);
	*/
	return 0;
}