// 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;
}

// 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 exchange_ciphertext_fields_with_decryptor(unsigned char* input, uint32_t input_size,
        unsigned char* output, uint32_t* output_size)
{
    // TODO: Initialize the extension_to_decryptor_ciphertext_msg protobuf with ciphertext from above.
    extension_to_decryptor_ciphertext_msg protobuf_extension_to_decryptor_msg;
    bool bool_return;

    protobuf_extension_to_decryptor_msg.set_ciphertext_bytes(input, input_size);

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

    protobuf_extension_to_decryptor_msg.clear_ciphertext_bytes();

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

    // First parse to the extension_to_decryptor_ciphertext_msg protobuf
    bool_return=protobuf_extension_to_decryptor_msg.SerializeToArray(output, protobuf_extension_to_decryptor_msg.ByteSize());
    if(bool_return)
    {
        *output_size=protobuf_extension_to_decryptor_msg.ByteSize();
        return 0;
    }
    else
        return 0x55;
}

int decrypt_client_data_through_decryptor( std::vector<std::string> &binary_ciphertext_client_data,
        std::vector<std::string> &plaintext_client_data)
{
    uint32_t counter, plaintext_size, ciphertext_size, size, double_ciphertext_size, ret;
    unsigned char* plaintext, *ciphertext, *double_ciphertext;

    extension_to_decryptor_enclosed_msg protobuf_client_to_decryptor_msg;

    protobuf_client_to_decryptor_msg.set_ciphertext_client_public_key(binary_ciphertext_client_data[0]);
    counter=0;
    for (std::string ciphertext_from_client:binary_ciphertext_client_data)
    {
        // Put all fields in the order in which they come, into the enclosed protobuf - first one is client pub key, rest are ciphertexts
        protobuf_client_to_decryptor_msg.mutable_ciphertext_fields(counter)->set_field(ciphertext_from_client);
        counter++;
    }

    // Serializing the protobuf with all ciphertext client data.
    ciphertext_size=protobuf_client_to_decryptor_msg.ByteSize();
    ciphertext = (unsigned char*) malloc(ciphertext_size);
    if(!protobuf_client_to_decryptor_msg.SerializeToArray(ciphertext, ciphertext_size))
    {
        free(ciphertext);
        return 0x32;
    }

    // Encrypting the serialized protobuf to the decryptor enclave
    double_ciphertext = (unsigned char*) malloc(ciphertext_size + 20); // +20 for the tag (prolly 16 or w/e)
    ret = aes_gcm_wrapper(1, ciphertext, ciphertext_size, double_ciphertext, &double_ciphertext_size );
    if(ret != 0)
    {
        free(ciphertext); free(double_ciphertext);
        printf("Error in encrypting to decryptor.\n"); fflush(stdout);
        return 34;
    }

    // Communicating with the decryptor enclave by transferring the serialized, encrypted buffer to another protobuf
    // (no crypto in the php extension code for this function)
    ret=exchange_ciphertext_fields_with_decryptor(double_ciphertext, double_ciphertext_size, ciphertext, &ciphertext_size);
    free(double_ciphertext);
    if(ret!=0)
    {
        free(ciphertext);
        return ret;
    }

    // Decrypting from the decryptor enclave
    plaintext=(unsigned char*)malloc(ciphertext_size);
    ret = aes_gcm_wrapper(0, (unsigned char*) ciphertext, ciphertext_size,
                                             plaintext, &plaintext_size);
    free(ciphertext);
    if(ret != 0)
    {
        printf("Error in decrypting content from the decryptor.\n"); fflush(stdout);
        return 36;
    }

    // Getting a deserialized protobuf
    protobuf_client_to_decryptor_msg.clear_ciphertext_fields();
    if(!protobuf_client_to_decryptor_msg.ParseFromArray(plaintext, plaintext_size))
    {
        printf("Error in generating protobuf with plaintexts.\n"); fflush(stdout);
        return 35;
    }
    free(plaintext);

    // Setting the output to elements in that protobuf.
    size=protobuf_client_to_decryptor_msg.ciphertext_fields_size();
    for(counter=0;counter<size;counter++)
    {
        plaintext_client_data.push_back(protobuf_client_to_decryptor_msg.ciphertext_fields(counter).field());
    }
    return 0;
}