//
// 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; const char *field_string;
    uint32_t* input_field_sizes, *output_field_sizes;
    std::string client_public_key_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;
    uint32_t counter;

    // Initialize the client public key
    printf("Initialize the client public key\n"); fflush(stdout);
    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 - first element is the length of the public key
    // and rest are sizes of each of the incoming ciphertext fields (an input arg. to the ecall).
    // Initialize an array to store the sizes of the *outgoing* fields too.
    printf("Initialize an array of sizes \n"); fflush(stdout);
    no_of_fields = incomingMsg.fields_size() + 1; // extra one for public key.
    input_field_sizes = (uint32_t*) malloc(no_of_fields);
    output_field_sizes = (uint32_t*) malloc(no_of_fields);
    total_input_length = client_public_key_length;

    printf("Setting input_field_sizes\n"); fflush(stdout);
    input_field_sizes[0] = client_public_key_length;
    for(field_counter=1; field_counter<no_of_fields; field_counter++)
    {
        input_field_sizes[field_counter] = incomingMsg.fields(field_counter-1).field().length();
        total_input_length+=input_field_sizes[field_counter];
        printf("%d ", input_field_sizes[field_counter]);
    }

    printf("Setting incoming client data\n"); fflush(stdout);
    // 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);
    // Did not tighten the upper bound for the length of the output data as the edl file cant handle functions of input arguments as the sizes.
    outgoing_client_data = (uint8_t*) malloc(total_input_length);
    printf("Setting public key of length %d\n", client_public_key_length);
    for(byte_counter=0; byte_counter<client_public_key_length; byte_counter++)
    {
        incoming_client_data[byte_counter] = client_public_key[byte_counter];
        printf("%02x ", incoming_client_data[byte_counter]);
    }
    field_size_accumulator=client_public_key_length;
    printf("\nSetting the other fields.\n");
    for(field_counter=1; field_counter<no_of_fields; field_counter++)
    {
        field_size = input_field_sizes[field_counter];
        printf("%d bytes\n", field_size);
        field_string = incomingMsg.fields(field_counter-1).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];
            printf("%02x ",  incoming_client_data[field_size_accumulator + byte_counter]);
        }
        printf("\n");
        field_size_accumulator += field_size;
    }

    for(counter=0; counter<total_input_length; counter++)
        printf("%02x ", incoming_client_data[counter]);
    for(counter=0; counter<no_of_fields; counter++)
        printf("%d ", input_field_sizes[counter]);
    printf("\nAbout to ecall.\n");
    sgx_ret_status=0;
    Decryptor_decrypt_client_data_wrapper(enclave_id, &sgx_ret_status, incoming_client_data, input_field_sizes,
            no_of_fields, total_input_length,  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: %u\n", sgx_ret_status);
        free(outgoing_client_data);
        free(output_field_sizes);
        return sgx_ret_status;
    }
    printf("returned successfully\n");

    // 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;
    printf("No of fields: %d\n", no_of_fields - 1);
    for(field_counter=0; field_counter<no_of_fields - 1; field_counter++)
    {
        field_size = output_field_sizes[field_counter];
        outgoingMsg.add_fields()->set_field(outgoing_client_data + field_size_accumulator,field_size);
        printf("%u bytes\n", field_size);
        for(byte_counter=field_size_accumulator; byte_counter<field_size_accumulator + field_size; byte_counter++)
            printf("%02x ", outgoing_client_data[byte_counter]);
        printf("\n");
        field_size_accumulator += field_size;
    }
    printf("\n");

    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;

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

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

    printf("Writing the msg back to the target.\n"); fflush(stdout);
    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;
    }
    uint32_t counter;
    for(counter=0;counter<value_length;counter++)
        printf("%02x ", header_value[counter]);
    printf("\n"); fflush(stdout);

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