// 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"PostLAMessages.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
// TODO: Rewrite read, write calls for protobuf msgs to and from fds to use message.SerializeToFileDescriptor and ParseFromFileDescriptor
bool read_protobuf_msg_from_fd(int accept_fd, google::protobuf::MessageLite& message)
{
    //message.SerializeToFileDescriptor(accept_fd);
  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 false;
  }
  //  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 false;
  }
  //  printf("Done parsing msg\n"); fflush(stdout);
  coded_input->PopLimit(limit);
  return true;
}

// TODO: private functions
bool write_protobuf_msg_to_fd(int accept_fd, google::protobuf::MessageLite& message)
{
    //return message.ParseFromFileDescriptor(accept_fd);
  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 false;
  }
  // 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 true;
}

void create_vector_from_protobuf(extension_to_decryptor_enclosed_msg &protobuf_ext_to_decryptor,
                                 std::vector<std::string> &binary_ciphertext_client_data)
{
    uint32_t temp_size, counter;

    binary_ciphertext_client_data.push_back(protobuf_ext_to_decryptor.ciphertext_client_public_key());
    temp_size=protobuf_ext_to_decryptor.ciphertext_fields_size();
    for(counter=0; counter<temp_size; counter++)
        binary_ciphertext_client_data.push_back(protobuf_ext_to_decryptor.ciphertext_fields(counter).field());
}

void create_protobuf_from_vector(std::vector<std::string> &double_ciphertext_client_data,
                                 extension_to_decryptor_enclosed_msg &protobuf_ext_decryptor_msg)
{
    uint32_t counter, temp_size;

    protobuf_ext_decryptor_msg.set_ciphertext_client_public_key(double_ciphertext_client_data[0]);
    temp_size=double_ciphertext_client_data.size()-1; // 1 for the public key.
    for(counter=0; counter<temp_size; counter++)
        protobuf_ext_decryptor_msg.mutable_ciphertext_fields(counter)->set_field(double_ciphertext_client_data[counter]);
}


// 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))
    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))
    return -1;

  // Reading msg3 from decryptor enclave
  if(!read_protobuf_msg_from_fd(decryptor_fd, protobuf_msg3))
    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))
  {
    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 encrypt_decrypt_ciphertexts(int encrypt_decrypt, std::vector<std::string> &binary_ciphertext_client_data,
        std::vector<std::string> &double_ciphertext_client_data)
{
    unsigned char *ciphertext; unsigned char *double_ciphertext;
    uint32_t ciphertext_size, double_ciphertext_size, ret;
    double_ciphertext=NULL;

    for (std::string ciphertext_from_client:binary_ciphertext_client_data)
    {
        ciphertext_size = ciphertext_from_client.length();
        ciphertext = (unsigned char*) ciphertext_from_client.c_str();
        double_ciphertext=(unsigned char*) realloc(double_ciphertext, ciphertext_size + 28); // 16 for tag, 12 for IV
        ret = aes_gcm_wrapper(encrypt_decrypt, ciphertext, ciphertext_size, double_ciphertext, &double_ciphertext_size );
        if(ret!=0)
        {
            free(double_ciphertext);
            printf("Failed to encrypt a ciphertext field.\n"); fflush(stdout);
            return 0x2;
        }
        double_ciphertext_client_data.push_back(std::string(reinterpret_cast<const char *> (double_ciphertext), double_ciphertext_size));
    }
    free(double_ciphertext);

    return 0;
}

int decrypt_client_data_through_decryptor( std::vector<std::string> &binary_ciphertext_client_data,
        std::vector<std::string> &plaintext_client_data)
{
    uint32_t ret;
    std::vector<std::string> double_ciphertext_client_data;
    std::vector<std::string> received_ciphertext_client_data;
    extension_to_decryptor_enclosed_msg protobuf_extension_decryptor_msg;

    ret=encrypt_decrypt_ciphertexts(1, binary_ciphertext_client_data, double_ciphertext_client_data);
    if(ret!=0)
        return ret;

    create_protobuf_from_vector(double_ciphertext_client_data,
                                protobuf_extension_decryptor_msg);

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

    protobuf_extension_decryptor_msg.clear_ciphertext_fields();
    protobuf_extension_decryptor_msg.clear_ciphertext_client_public_key(); // not necessary.

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

    create_vector_from_protobuf(protobuf_extension_decryptor_msg, received_ciphertext_client_data);

    ret=encrypt_decrypt_ciphertexts(0, received_ciphertext_client_data,
            plaintext_client_data);
    return ret;
}