//
// Created by miti on 2019-12-24.
//
#include "PostLAMessages.pb.h"
#include "PostLAMessaging.h"
#include <unistd.h>
#include <stdio.h>
#include "Decryptor_u.h"
int PostLAMessaging::process_target_data_msg(extension_to_decryptor_msg& incomingMsg, decryptor_to_extension_msg& outgoingMsg)
{
    const char* client_public_key;
    uint8_t* incoming_client_data, *outgoing_client_data;
    uint32_t* input_field_sizes, *output_field_sizes;
    std::string client_public_key_string, field_string;
    uint32_t sgx_ret_status, client_public_key_length, field_counter, byte_counter, field_size, no_of_fields, total_input_length, field_size_accumulator;

    // Initialize the client public key as an input to the decryptor ecall
    client_public_key_string = incomingMsg.ciphertext_client_public_key();
    client_public_key_length = client_public_key_string.length();
    // client_public_key = (uint8_t*) malloc(client_public_key_length);
    client_public_key = client_public_key_string.c_str();

    // Initialize an array of sizes of each of the incoming ciphertext fields (again an input arg. to the ecall).
    // Initialize an array to store the sizes of the *outgoing* fields too.
    no_of_fields = incomingMsg.fields_size();
    input_field_sizes = (uint32_t*) malloc(no_of_fields);
    output_field_sizes = (uint32_t*) malloc(no_of_fields);
    total_input_length = 0;
    for(field_counter=0; field_counter<no_of_fields; field_counter++)
    {
        input_field_sizes[field_counter] = incomingMsg.fields(field_counter).field().length();
        total_input_length+=input_field_sizes[field_counter];
    }

    // Initialize a contiguous array of bytes of the ciphertext fields.
    // The number of bytes in any given ciphertext field will be determined by the above array.
    // Also initialize an array to store the outgoing bytes for all ciphertext fields.
    incoming_client_data = (uint8_t*) malloc(total_input_length);
    outgoing_client_data = (uint8_t*) malloc(total_input_length);
    for(field_counter=0; field_counter<no_of_fields; field_counter++)
    {
        field_size = input_field_sizes[field_counter];
        field_string = incomingMsg.fields(field_counter).field().c_str();
        for(byte_counter=0; byte_counter<field_size; byte_counter++)
        {
            incoming_client_data[field_size_accumulator + byte_counter] = field_string[byte_counter];
        }
        field_size_accumulator += field_size;
    }

    sgx_ret_status=0;
    // TODO: Fix enclave ecall.
    // Ecall into the enclave to obtain the fields to be sent to the target enclave in response.
    // Decryptor_functionX(enclave_id, &sgx_ret_status, client_public_key, incoming_client_data, no_of_fields, input_field_sizes,  outgoing_client_data, output_field_sizes);
    free(input_field_sizes);
    free(incoming_client_data);
    if(sgx_ret_status!=0)
    {
        printf("Could not process client's data. Had error: %ud\n", sgx_ret_status);
        free(outgoing_client_data);
        free(output_field_sizes);
        return sgx_ret_status;
    }

    // Initialize the outgoing protobuf message using the field sizes in the outgoing_field_size array and the bytes in the outgoing_client_data array.
    field_size_accumulator=0;
    for(field_counter=0; field_counter<no_of_fields; field_counter++)
    {
        field_size = output_field_sizes[field_counter];
        outgoingMsg.add_fields()->set_field(outgoing_client_data + field_size_accumulator,field_size);
        field_size_accumulator += field_size;
    }
    free(outgoing_client_data);
    free(output_field_sizes);
    return 0;
}

void PostLAMessaging::set_enclave_id(int given_id)
{
    enclave_id = given_id;
}

int PostLAMessaging::receive_token_from_verifier()
{
    uint8_t* token = (uint8_t*) malloc(400);
    uint32_t token_length  = read(verifier_fd, token, 400);
    uint32_t sgx_ret_status;
    if(token_length < 0)
    {
        printf("Could not read verifier's token.\n"); fflush(stdout);
        return 0x1;
    }
    else if(token_length > 400)
    {
        printf("Could not fit all of verifiers token into the buffer.\n"); fflush(stdout);
        return 0x2;
    }

    Decryptor_process_verifiers_message_wrapper(enclave_id, &sgx_ret_status, token, token_length);
    free(token);
    return sgx_ret_status;
}

int PostLAMessaging::receive_decrypt_client_data_from_target()
{
    extension_to_decryptor_msg incomingMsg;
    decryptor_to_extension_msg outgoingMsg;
    int ret_status;

    if( targetChannelData.read_msg(incomingMsg) != 0)
    {
        printf("Error in protobuf in receiving msg from target.\n"); fflush(stdout);
        return 0x1;
    }

    ret_status = process_target_data_msg(incomingMsg, outgoingMsg);
    if(ret_status != 0)
    {
        printf("Error in processing msg from target.\n"); fflush(stdout);
        return 0x2;
    }

    if( targetChannelData.write_msg(outgoingMsg) != 0)
    {
        printf("Error in protobuf in sending msg to target.\n"); fflush(stdout);
        return 0x3;
    }
    return 0;
}

int PostLAMessaging::read_and_write_header()
{
    uint32_t sgx_ret_status, value_length, protobuf_ret_status;
    uint8_t* header_value = (uint8_t*) malloc(400);

    mitigator_header headerMsg;
    if(targetChannelHeaders.read_msg(headerMsg) != 0)
    {
        printf("Error in reading msg for headers.\n"); fflush(stdout); return 0x1;
    }

    Decryptor_create_and_encrypt_mitigator_header_H_wrapper(enclave_id, &sgx_ret_status, header_value, &value_length);
    if(sgx_ret_status != 0 )
    {
        printf("Could not create mitigator header.\n");
        fflush(stdout);
        return sgx_ret_status;
    }

    headerMsg.set_name("Mitigator-Public-Key:", 21);
    headerMsg.set_value(header_value, value_length);
    free(header_value);

    protobuf_ret_status = targetChannelHeaders.write_msg(headerMsg);
    return protobuf_ret_status;
}

void PostLAMessaging::set_verifier_fd(int fd)
{
    verifier_fd = fd;
}

void PostLAMessaging::set_target_data_fd(int fd)
{
    targetChannelData.set_fd(fd);
}

void PostLAMessaging::set_target_headers_fd(int fd)
{
    targetChannelHeaders.set_fd(fd);
}