// 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>
#include "LocalAttestationUntrusted.h"

  uint32_t LocalAttestationUntrusted::session_id=0;
  protobuf_sgx_dh_msg1_t LocalAttestationUntrusted::protobuf_msg1;



  int LocalAttestationUntrusted::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 LocalAttestationUntrusted::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 LocalAttestationUntrusted::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 LocalAttestationUntrusted::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;
    printf("Writing message 1\n"); fflush(stdout);
    if(write_protobuf_msg_to_fd(accept_fd, protobuf_msg1)!=0)
      return 0x1;

    printf("Reading message 2\n"); fflush(stdout);
    if(read_protobuf_msg_from_fd(accept_fd, protobuf_msg2)!=0)
      return 0x2;

    protobuf_sgx_ret = process_protobuf_dh_msg2_generate_protobuf_dh_msg3(own_enclave_id, protobuf_msg2, protobuf_msg3, &LocalAttestationUntrusted::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;
    }

    printf("Writing message 3\n"); fflush(stdout); 
    if(write_protobuf_msg_to_fd(accept_fd, protobuf_msg3)!=0)
      return 0x3;
    return 0;
  }

  int LocalAttestationUntrusted::decrypt_client_data(uint32_t own_enclave_id, int fd)
  {
    protobuf_post_LA_encrypted_msg_t protobuf_msg;
    unsigned char* protobuf_msg_ptr;
    uint32_t sgx_ret_status=0;
    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(fd, protobuf_msg)!=0)
    {
//	printf("Could not read apache's message post-local attestation\n"); fflush(stdout);
      return 0xfe;
    }
//    printf("Clients data\n"); fflush(stdout); 

    input_ciphertext_plus_tag_length = protobuf_msg.msg().length();
    input_ciphertext_plus_tag = (uint8_t*) malloc(1000); //malloc(input_ciphertext_plus_tag_length);
    output_ciphertext_plus_tag = (uint8_t*) malloc(1000); //malloc(input_ciphertext_plus_tag_length); 
    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);
//	printf("0x%02x ", input_ciphertext_plus_tag[counter]); 
    }
    // Just so that the ciphertext - client data - is returned back to Apache in case this function fails.
    // client data is after public key (64 bytes) + signature (64 bytes) = 128 bytes.
    for(counter=0; counter<input_ciphertext_plus_tag_length-64; counter++)
  		output_ciphertext_plus_tag[counter] = input_ciphertext_plus_tag[counter+64];
	output_ciphertext_plus_tag_length=input_ciphertext_plus_tag_length - 64;
        protobuf_msg.set_msg((void*)  output_ciphertext_plus_tag, output_ciphertext_plus_tag_length); 
    
    // We assume that the output is not changed unless it is successful throughout.
    Decryptor_process_apache_message_generate_response_wrapper(own_enclave_id, &sgx_ret_status, input_ciphertext_plus_tag, input_ciphertext_plus_tag_length, output_ciphertext_plus_tag, &output_ciphertext_plus_tag_length);
    free(input_ciphertext_plus_tag);
      if(sgx_ret_status==0)
      {
      	protobuf_msg.set_msg((void*)  output_ciphertext_plus_tag, output_ciphertext_plus_tag_length);
      }
//	else
//		        printf("\n0x%02x\n", sgx_ret_status); 

/*	printf("Returning this: \n"); 
	for(counter=0;counter<output_ciphertext_plus_tag_length;counter++)
		printf("%02x ",output_ciphertext_plus_tag[counter]); 
	printf("\n"); */
       free(output_ciphertext_plus_tag);
     if(write_protobuf_msg_to_fd(fd, protobuf_msg)!=0)
          return 0xfc;
  	return 0;

  }

 
    int LocalAttestationUntrusted::prepare_local_attestation_as_responder_msg1(uint32_t own_enclave_id) //, int port)
    {
      uint32_t protobuf_sgx_ret;
      int temp_server_fd=0;
      protobuf_sgx_ret = generate_protobuf_dh_msg1(own_enclave_id, protobuf_msg1, &LocalAttestationUntrusted::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;
      }
      return 0;
    }

    int LocalAttestationUntrusted::setup_socket_for_local_attestation_requests(int port)
    {
      struct sockaddr_in own_addr;
      return set_up_socket(port, &own_addr);
    }

    // TODO: CHANGED SIGNATURE.
    int LocalAttestationUntrusted::local_attestation_as_responder_msg2_msg3(uint32_t own_enclave_id, int server_fd, int* accept_fd)
    {
      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, temp_accept_fd);
      return protobuf_sgx_ret;
    }

    int LocalAttestationUntrusted::post_local_attestation_with_verifier(uint32_t own_enclave_id, int accept_fd)
    {
      uint32_t protobuf_sgx_ret;
      uint8_t encrypted_apache_mrsigner_and_tag[48];
      size_t bytes_read;
	int count; 
      printf("Here\n"); fflush(stdout);
      bytes_read=read(accept_fd, encrypted_apache_mrsigner_and_tag, 48);
    	if(bytes_read!=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_process_verifiers_message_wrapper(own_enclave_id, &protobuf_sgx_ret, encrypted_apache_mrsigner_and_tag);
      if(protobuf_sgx_ret!=0)
    	{
        printf("Error in decryption: 0x%x\n", protobuf_sgx_ret); fflush(stdout);
        close(accept_fd);
        return protobuf_sgx_ret;
    	}

      printf("Successful decryption\n"); fflush(stdout);
      close(accept_fd);
	uint8_t output[64];
      Decryptor_get_verifier_mrenclave_apache_mrsigner_wrapper(own_enclave_id, output); 
	uint32_t counter;  
	for(counter=0; counter<32; counter++)
		printf("0x%x ", output[counter]);
	printf("/n"); 
        for(counter=32; counter<64; counter++)
                printf("0x%x ", output[counter]);
        printf("/n"); 

	fflush(stdout);

      return 0;

    }

    int LocalAttestationUntrusted::post_local_attestation_with_apache(uint32_t own_enclave_id, int accept_fd)
    {
      protobuf_post_LA_encrypted_msg_t protobuf_encrypted_msg;
      uint8_t encrypted_sign_data_and_sign_and_tag[176];
  	  memset(encrypted_sign_data_and_sign_and_tag,0x0,176);
	    uint32_t internal_return_status;
	    uint32_t count;
	    uint32_t sgx_ret;

        Decryptor_create_and_encrypt_mitigator_header_H_wrapper(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(accept_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;
        }
	uint8_t public_key[64];
	Decryptor_get_short_term_public_key_wrapper(own_enclave_id, public_key); 
	for(count=0;count<64;count++)
		printf("%02x ",public_key[count]); 
	printf("\n"); fflush(stdout);

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

      return 0; }