// 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>
#include"SgxProtobufLAInitiator.h"
#include"crypto.h"
// For socket programming
#include<arpa/inet.h>
#include<sys/socket.h>
#include<netinet/in.h>
#include<queue>
using namespace google::protobuf::io;

int global_decryptor_fd;
#define TOKEN_H_MAX_LENGTH 300
uint8_t ciphertext_to_decryptor[4092]; 

// 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)
{
  // 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;

  uint32_t session_id;
  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);
  }
  global_decryptor_fd=decryptor_fd;

  // Reading msg 1 from decryptor enclave.
  if(read_protobuf_msg_from_fd(decryptor_fd, protobuf_msg1)!=0)
    return -1;

  // Generating message 2 given message 1 as input
  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;
  }

  // Writing msg 2 to decryptor enclave.
  if(write_protobuf_msg_to_fd(decryptor_fd, protobuf_msg2)!=0)
    return -1;

  // Reading msg3 from decryptor enclave
  if(read_protobuf_msg_from_fd(decryptor_fd, protobuf_msg3)!=0)
    return -1;

  // Process message 3
  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;
  }
  return 0;
}


uint32_t post_local_attestation_get_mitigator_header(unsigned char* base64_encoded_token_H, uint32_t* base64_encoded_token_H_length)
{
  unsigned char* encrypted_token_H_and_tag, *plaintext_token_H;
  uint32_t encrypted_token_H_and_tag_length, plaintext_token_H_length;
  protobuf_post_LA_encrypted_msg_t protobuf_encrypted_msg;
  uint32_t count;
  uint32_t internal_return_status;

  // Reading the encrypted token H from the
  if(read_protobuf_msg_from_fd(global_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());
  encrypted_token_H_and_tag_length=protobuf_encrypted_msg_string.length();
  encrypted_token_H_and_tag = (unsigned char*) malloc(encrypted_token_H_and_tag_length);
  memcpy(encrypted_token_H_and_tag, protobuf_encrypted_msg_string.c_str(), protobuf_encrypted_msg_string.length());
  
  printf("Read the following encrypted token T and tag from the decryptor socket.\n"); fflush(stdout);
	printf("Encrypted data:\n");
	for(count=0; count<encrypted_token_H_and_tag_length; count++)
		printf("0x%02x ", encrypted_token_H_and_tag[count]);
	printf("\n"); fflush(stdout);
 plaintext_token_H=(unsigned char*) malloc(encrypted_token_H_and_tag_length); 
	internal_return_status = aes_gcm_wrapper(0, encrypted_token_H_and_tag, encrypted_token_H_and_tag_length, plaintext_token_H, &plaintext_token_H_length);
	if(internal_return_status != 0)
	{
		free(encrypted_token_H_and_tag); 
		free(plaintext_token_H); 
		printf("Error in decryption 0x%x", internal_return_status); fflush(stdout); return internal_return_status;
	}

  printf("Decryptor's short-term public key\n");
  for(count=0;count<64;count++)
    printf("0x%02x ", plaintext_token_H[count]);
	printf("\n"); fflush(stdout);

	printf("Verifier mrenclave\n");
	for(count=64;count<96;count++)
    printf("0x%02x ", plaintext_token_H[count]);
	printf("\n"); fflush(stdout);
	 
	*base64_encoded_token_H_length=base64_encoding_wrapper(base64_encoded_token_H, plaintext_token_H, plaintext_token_H_length);
  
  // TODO: Better return value handling.
  //if(!(base64_encoded_token_H_length <= 4 * Math.ceil(plaintext_token_H_length/3) && base64_encoded_token_H_length > 4 * (Math.ceil(plaintext_token_H_length/3) - 1)))
  //{
  //  printf("Somehow not the entire token was encoded in base64:0x%x\n", base64_encoded_token_H_length); fflush(stdout); return 0x55;
  //}
  
        free(encrypted_token_H_and_tag); 
        free(plaintext_token_H); 
	fflush(stdout);
	return 0;
}


uint32_t base64_decoding_on_all_client_data(unsigned char* ip_base64_client_public_key_ciphertext, 
  uint32_t ip_base64_client_public_key_ciphertext_length,
  unsigned char* op_client_public_key_ciphertext,
  uint32_t* op_client_public_key_ciphertext_length
)
{
/*    uint32_t openssl_ret_status;
    openssl_ret_status=base64_decoding_wrapper(ip_base64_client_public_key_ciphertext, op_client_public_key_ciphertext, ip_base64_client_public_key_ciphertext_length);
    if(openssl_ret_status == -1)
      	return 0xfe;
    *op_client_public_key_ciphertext_length = openssl_ret_status; 
*/
    return 0;

}

int decrypt_client_data_through_decryptor( unsigned char* ip_client_public_key_ciphertext, 
  uint32_t ip_client_public_key_ciphertext_length,
  unsigned char* op_plaintext,
  uint32_t* op_plaintext_length
)
{
	uint32_t ciphertext_to_decryptor_length;
	uint32_t internal_return_status; 
	protobuf_post_LA_encrypted_msg_t protobuf_encrypted_msg; 
	std::string protobuf_encrypted_msg_string;
	uint32_t counter;
	// Encrypting to the decryptor enclave
	internal_return_status = aes_gcm_wrapper(1, ip_client_public_key_ciphertext, ip_client_public_key_ciphertext_length,    ciphertext_to_decryptor, &ciphertext_to_decryptor_length );
	// Error checking on ciphertext that is to be sent to decryptor
	if(internal_return_status != 0)
	{
		printf("Error in encrypting content to the decryptor.\n"); fflush(stdout); 
		return internal_return_status;
	}

	  // write message to decryptor
	protobuf_encrypted_msg.set_msg( ciphertext_to_decryptor , ciphertext_to_decryptor_length );
	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
	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;
	}

	protobuf_encrypted_msg_string=std::string(protobuf_encrypted_msg.msg());
	// Decrypt data from decryptor
	internal_return_status = aes_gcm_wrapper(0, (unsigned char*) protobuf_encrypted_msg.msg().c_str(), protobuf_encrypted_msg_string.length(),  op_plaintext, op_plaintext_length);
	return internal_return_status;
}