// Knows only protobuf_sgx objects, protobuf header.
#include <stdlib.h>
#include <string.h>
#include <errno.h>
#include<unistd.h>
#include <stdio.h>
#include "ProtobufLAMessages.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"

// For socket programming
#include <arpa/inet.h>
#include <sys/socket.h>
#include <netinet/in.h>


// 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");
    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);
  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 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;
  uint8_t encrypted_hash_and_tag[150];// uint8_t encrypted_tag[16]; 
  uint32_t total_length;
size_t post_la_bytes_written;
  // For socket to listen to the Apache enclave.

  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;
  }
  memset(encrypted_hash_and_tag, 0, 150); 
  protobuf_sgx_ret=generate_encrypted_rsa_keypair_hash(encrypted_hash_and_tag, &total_length); 
  if(protobuf_sgx_ret==0)
  {
	printf("Done encryption of hash.\n"); fflush(stdout); 
  }
  else 
	{
		printf("Error in enc/dec of hash: 0x%x", protobuf_sgx_ret); fflush(stdout); return protobuf_sgx_ret; 
	}
  

   post_la_bytes_written = write(decryptor_fd, encrypted_hash_and_tag, total_length);
   printf("Wrote the hash and the tag to the decryptor socket.\n Wrote this many bytes: %d\n", post_la_bytes_written); fflush(stdout); 

   if(close(decryptor_fd)!= 0)
   {
	printf("Error in closing the socket connection.\n"); fflush(stdout); return 0xfd; 
   } 
  

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