// Knows only protobuf_sgx objects, protobuf header.
// For socket programming
#include <sys/socket.h>
#include <stdlib.h>
#include <netinet/in.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 "protobufLAInitiator.h"
#include "../Decryptor/Decryptor_u.h"
#include <iostream>

class LocalAttestationUntrusted {
  int server_fd;
  int accept_fd;
  uint32_t session_id;
  protobuf_sgx_dh_msg1_t protobuf_msg1;

  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);
    delete raw_input;
    delete coded_input;
    return 0;
  }

  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 to bind and listen to the given port. Returns FD of the socket on success, -1 on failure (and prints a msg to stdout with the errno)
  int set_up_socket(int port,   sockaddr_in* address)
  {
    int server_fd = 0;

    // Creating socket file descriptor for listening for attestation requests.
    server_fd = socket(AF_INET, SOCK_STREAM | SOCK_CLOEXEC, 0);
    if (server_fd == -1)
    {
       printf("Error in creating a socket - %d", errno);
        return -1;
    }

    // Preparing the address struct for binding
    address->sin_family = AF_INET;
    address->sin_addr.s_addr = INADDR_ANY; // Todo: should this be localhost?
    address->sin_port = htons(port);
    //  memset(address->sin_zero,0,sizeof(address->sin_zero));

    socklen_t addrlen = sizeof(*address);
    // Binding
    if (bind(server_fd, (sockaddr*)address, addrlen)<0)
    {
        printf("Error in binding %d - port was %d - ", errno, port);
        return -1;
    }

    // Listening
    if (listen(server_fd, 128) < 0)
    {
        printf("Error in listening %d", errno);
        return -1;
    }

    return server_fd;
  }

  uint32_t local_attestation_msg2_msg3(uint32_t own_enclave_id, int accept_fd)
  {
    uint32_t protobuf_sgx_ret;
    protobuf_sgx_dh_msg2_t protobuf_msg2;
    protobuf_sgx_dh_msg3_t protobuf_msg3;

    if(write_protobuf_msg_to_fd(accept_fd, protobuf_msg1)!=0)
      return 0x1;

    if(read_protobuf_msg_from_fd(accept_fd, protobuf_msg2)!=0)
      return 0x2;

    // TODO: Edit function signature in the definition: last argument read_or_write is used to control the flow of the untrusted program:
    // no point in doing this as it is untrusted. Have an attribute in its class for it..
    protobuf_sgx_ret = process_protobuf_dh_msg2_generate_protobuf_dh_msg3(own_enclave_id, protobuf_msg2, protobuf_msg3, &session_id);
    if(protobuf_sgx_ret != 0)
    {
        printf("Error in generate_protobuf_dh_msg2: 0x%x", protobuf_sgx_ret); fflush(stdout); return protobuf_sgx_ret;
    }

    if(write_protobuf_msg_to_fd(accept_fd, protobuf_msg3)!=0)
      return 0x3;

    return 0;
  }

  int decrypt_client_data_wrapper(uint32_t own_enclave_id, int apache_fd)
  {
    protobuf_post_LA_encrypted_msg_t protobuf_msg;
    unsigned char* protobuf_msg_ptr;
    uint32_t sgx_ret_status;
    // TODO: THIS SHOULD BE WAYYYY GREATER THAN 144 -> CLIENTS PUBLIC KEY (64)+ ENCRYPTION TAG (16)+ SIGNATURE (64) + CLIENT DATA
    // 144 + client data length
    uint8_t* input_ciphertext_plus_tag;
    uint32_t input_ciphertext_plus_tag_length;
  	uint8_t* output_ciphertext_plus_tag;
    uint32_t output_ciphertext_plus_tag_length;
    int counter;

    if(read_protobuf_msg_from_fd(apache_fd, protobuf_msg)!=0)
      return 0xfe;

    input_ciphertext_plus_tag_length = protobuf_msg.msg().length();
    // TODO: MAKE SURE THIS IS NOT 0XFFFFFFFF.
    input_ciphertext_plus_tag = (uint8_t*) malloc(input_ciphertext_plus_tag_length);
    output_ciphertext_plus_tag = (uint8_t*) malloc(input_ciphertext_plus_tag_length - 128 + 10);
    protobuf_msg_ptr = (uint8_t*) protobuf_msg.msg().c_str();

    for(counter=0; counter<input_ciphertext_plus_tag_length; counter++)
      input_ciphertext_plus_tag[counter] = *(protobuf_msg_ptr + counter);

    // Just so that the ciphertext - client data - is returned back to Apache in case decrypt_client_data fails.
    // client data is after public key (64 bytes) + signature (64 bytes) = 128 bytes.
    for(counter=0; counter<input_ciphertext_plus_tag_length; counter++)
  		output_ciphertext_plus_tag[counter] = input_ciphertext_plus_tag[counter+128];

    // We assume that the output is not changed unless it is successful throughout.
    // Return value is not sent back..
    Decryptor_process_apache_message_generate_response(own_enclave_id, &sgx_ret_status, input_ciphertext_plus_tag, input_ciphertext_plus_tag_length, output_ciphertext, &output_ciphertext_plus_tag_length);
    free(input_ciphertext_plus_tag);
    protobuf_msg.set_msg((void*)  output_ciphertext_plus_tag, output_ciphertext_plus_tag_length);
    free(output_ciphertext_plus_tag);
  	if(write_protobuf_msg_to_fd(apache_fd, protobuf_msg)!=0)
  	  return 0xfc;

  	return 0;
  }

  public:
    int prepare_local_attestation_as_responder_fd_and_msg1(uint32_t own_enclave_id, int port)
    {
      struct sockaddr_in own_addr;
      uint32_t protobuf_sgx_ret;
      int temp_server_fd=0;
      protobuf_sgx_ret = generate_protobuf_dh_msg1(own_enclave_id, protobuf_msg1, session_id);
      if(protobuf_sgx_ret != 0)
      {
        printf("Error in generate_protobuf_dh_msg1: 0x%x", protobuf_sgx_ret); fflush(stdout); return -protobuf_sgx_ret;
      }

      temp_server_fd=set_up_socket(port, &own_addr);
      if(temp_server_fd>0)
        server_fd = temp_server_fd;
      return temp_server_fd;
    }

    // TODO: CHANGED SIGNATURE.
    int local_attestation_as_responder_msg2_msg3(uint32_t own_enclave_id, int server_fd)
    {
      //   protobuf_post_LA_encrypted_msg_t protobuf_encrypted_msg;
      uint32_t protobuf_sgx_ret;
      struct sockaddr_storage apache_addr;
      socklen_t apache_addr_size = sizeof(apache_addr);

      int temp_accept_fd;
      temp_accept_fd = accept(server_fd, (struct sockaddr *)&apache_addr,&apache_addr_size);
      if (temp_accept_fd <0)
      {
        printf("Error in accepting %d", errno); fflush(stdout);
        return temp_accept_fd;
      }
      accept_fd=temp_accept_fd;

      protobuf_sgx_ret = local_attestation_msg2_msg3(own_enclave_id, accept_fd);
      return protobuf_sgx_ret;
    }

    int post_local_attestation_with_verifier(uint32_t own_enclave_id)
    {
      uint32_t protobuf_sgx_ret;
      uint8_t encrypted_apache_mrsigner_and_tag[48];
      size_t bytes_read;

      printf("Here\n"); fflush(stdout);
      bytes_read=read(accept_fd, encrypted_apache_mrsigner_and_tag, 48);
    	if(bytes_read_post_la!=48)
     	{
           	printf("Not all of the encrypted apache's mrsigner was read from the verifier.\n"); fflush(stdout);
            close(accept_fd);
            return 0xfe;
     	}

      for(count=0;count<48;count++)
    	  printf("0x%02x ", encrypted_apache_mrsigner_and_tag[count]);
      printf("\n");fflush(stdout);

      Decryptor_decrypt_verifiers_message_set_apache_mrsigner(own_enclave_id, &sgx_ret, encrypted_apache_mrsigner_and_tag);
      if(sgx_ret!=0)
    	{
        printf("Error in decryption: 0x%x\n", sgx_ret); fflush(stdout);
        close(accept_fd);
        return sgx_ret;
    	}

      printf("Successful decryption\n"); fflush(stdout);
      close(accept_fd);
      return 0;
    }

    int post_local_attestation_with_apache(uint32_t own_enclave_id)
    {
      protobuf_post_LA_encrypted_msg_t protobuf_encrypted_msg;
      uint8_t encrypted_sign_data_and_sign_and_tag[176];
      int apache_fd=accept_fd;
  	  memset(encrypted_sign_data_and_sign_and_tag,0x0,176);

      Decryptor_create_and_encrypt_mitigator_header_H(own_enclave_id, &sgx_ret, encrypted_sign_data_and_sign_and_tag);
    	if(sgx_ret!=0)
    	{
    		printf("Error in generating encrypted mitigator header:0x%x\n", sgx_ret); fflush(stdout);
        close(accept_fd);
        return 0xf3;
    	}

    	for(count=0;count<176;count++)
    	{
    		printf("0x%02x ", encrypted_sign_data_and_sign_and_tag[count]);
    	}
    	printf("\n"); fflush(stdout);

    	protobuf_encrypted_msg.set_msg((void*)encrypted_sign_data_and_sign_and_tag, 176);
    	if(write_protobuf_msg_to_fd(apache_fd, protobuf_encrypted_msg) != 0)
      {
        printf("Not all of the mitigator token H was written to the Apache.\n"); fflush(stdout);
        close(accept_fd);
        return 0xfe;
      }

  	  do {
      	internal_return_status = decrypt_client_data_wrapper(own_enclave_id, apache_fd);
      } while(internal_return_status==0);
      close(accept_fd);
      return internal_return_status;
    }
}