//
// Created by miti on 2019-12-24.
//
#include "PostLAMessaging.h"
#include "sgx_trts.h" // for sgx_read_rand
#include "crypto.h" // for aes_gcm_128
#include <unistd.h>
#include <stdio.h>
uint32_t PostLAMessaging::aes_gcm_wrapper(int enc, uint8_t* plaintext, uint32_t plaintext_length, uint8_t* ciphertext, uint32_t* ciphertext_length)
{
    uint32_t actual_plaintext_length=plaintext_length;
    uint8_t tag[16];uint32_t counter, return_status;
    uint8_t iv[12];
    if(enc == 0)
    {
        for(counter=0;counter<16;counter++)
            tag[counter]=plaintext[counter+plaintext_length-16];
        for(counter=0;counter<12;counter++)
            iv[counter]=plaintext[counter+plaintext_length-28];
        actual_plaintext_length-=28;
    }
    else
    {
        return_status=sgx_read_rand(iv, 12);
        if(return_status != 0)
            return return_status;
    }
    return_status = aes_gcm_128(enc, key, iv, plaintext, actual_plaintext_length, ciphertext, ciphertext_length, tag);
    if(enc == 1 && return_status == 0)
    {
        for(counter=0;counter<12;counter++)
            ciphertext[counter + *ciphertext_length] = iv[counter];
        for(counter=0;counter<16;counter++)
            ciphertext[counter + 12 + *ciphertext_length] = tag[counter];
        *ciphertext_length=*ciphertext_length + 28;
    }
    return return_status;
}

void PostLAMessaging::set_la_symmetric_key(uint8_t* given_key) {
    uint32_t counter;
    for(counter=0; counter<16; counter++)
    {
        key[counter] = given_key[counter];
    }
}

void PostLAMessaging::set_headers_fd(int given_fd)
{
    headersChannel.set_fd(given_fd);
}

void PostLAMessaging::set_data_fd(int given_fd)
{
    dataChannel.set_fd(given_fd);
}


uint32_t PostLAMessaging::encrypt_decrypt_msgs(int encrypt_decrypt, std::vector<std::vector<unsigned char>>& input_msgs,
                         std::vector<std::vector<unsigned char>>& output_msgs)
{
    unsigned char *input, *output;
    uint32_t input_size, output_size, ret, counter;
    output=NULL;

    printf("In encrypt_decrypt_msgs\n"); fflush(stdout);
    for (auto msg:input_msgs)
    {
        printf("In the loop.\n"); fflush(stdout);
        input_size = msg.size();
        input = (unsigned char*) &msg[0];
        printf("Input length %d\n", input_size);
        for(counter=0; counter<input_size; counter++)
            printf("%02x ", input[counter]);
        printf("\n");
        output = (unsigned char*) realloc(output, input_size + 28); // 16 for tag, 12 for IV
        ret = aes_gcm_wrapper(encrypt_decrypt, input, input_size, output, &output_size );
        if(ret!=0)
        {
            free(output);
            printf("Failed to encrypt an input field.\n"); fflush(stdout);
            return ret;
        }
        printf("Output size %d\n", output_size);
        for(counter=0; counter<output_size; counter++)
            printf("%02x ", output[counter]);
        printf("\n");
        std::vector<unsigned char> output_field(output, output + output_size);
        output_msgs.push_back(output_field);
    }
    printf("Outside da loop.\n"); fflush(stdout);
    if(output != NULL)
        free(output);
    return 0;
}

void PostLAMessaging::create_vector_from_protobuf(decryptor_to_extension_msg &protobuf_msg,
                                 std::vector<std::vector<unsigned char>> &output_vector)
{
    uint32_t counter, no_of_fields;
    std::string msg;

    no_of_fields=protobuf_msg.fields_size();
    for(counter=0; counter<no_of_fields; counter++)
    {
        msg = protobuf_msg.fields(counter).field();
        std::vector<unsigned char> output_field(msg.begin(), msg.end());
        output_vector.push_back(output_field);
    }

}

void PostLAMessaging::create_protobuf_from_vector(std::vector<std::vector<unsigned char>> &input_vector,
                                 extension_to_decryptor_msg &protobuf_msg)
{
    unsigned char* input_ptr;
    std::vector<unsigned char> input_field;
    size_t input_field_length;

    input_field = input_vector.at(0);
    input_ptr = &input_field[0];
    input_field_length = input_field.size();
    protobuf_msg.set_ciphertext_client_public_key(input_ptr, input_field_length);
    input_vector.erase(input_vector.begin());

    for(auto input_field : input_vector) // counter=0; counter<no_of_fields; counter++
    {
        input_field_length = input_field.size();
        input_ptr = &input_field[0];
        protobuf_msg.add_fields()->set_field(input_ptr, input_field_length);
    }

}

uint32_t PostLAMessaging::receive_data_from_decryptor(std::vector<std::vector<unsigned char>> &plaintext_msg_list) {
    std::vector<std::vector<unsigned char>> ciphertext_msg_list;

    decryptor_to_extension_msg msg;

    // read encrypted data
    if(dataChannel.read_msg(msg) != 0)
    {
        printf("Not all of the decryptor's message was read\n"); fflush(stdout);
        return 0xf3;
    }
    printf("Read the msg from the decryptor.\n"); fflush(stdout);
    create_vector_from_protobuf(msg, ciphertext_msg_list);
    printf("created a vector\n"); fflush(stdout);
    return encrypt_decrypt_msgs(0, ciphertext_msg_list, plaintext_msg_list);
}

uint32_t PostLAMessaging::send_data_to_decryptor(std::vector<std::vector<unsigned char>> &plaintext_msg_list)
{
    uint32_t ret;
    extension_to_decryptor_msg msg;
    std::vector<std::vector<unsigned char>> ciphertext_msg_list;

    ret=encrypt_decrypt_msgs(1, plaintext_msg_list, ciphertext_msg_list);
    if(ret!=0)
        return ret;
    printf("About to create a protobuf vector\n"); fflush(stdout);
    create_protobuf_from_vector(ciphertext_msg_list, msg);
    printf("About to write the message to the data channel.\n"); fflush(stdout);
    if(dataChannel.write_msg(msg) != 0)
    {
        printf("Not all of the client's pub key and ciphertext data was written\n"); fflush(stdout);
        return 0xfe;
    }

    return 0;
}

uint32_t PostLAMessaging::receive_header_from_decryptor(std::string &plaintext_header) {
    std::vector<std::vector<unsigned char>> ip_list, binary_list;
    unsigned char* binary_ptr;
    std::vector<unsigned char> binary_header;
    size_t binary_header_length;
    uint8_t* base64_ciphertext_header;
    int base64_ciphertext_header_length;
    uint32_t ret_status;

    printf("About to send-receive header.\n"); fflush(stdout);

    mitigator_header headerMsg;
    headerMsg.set_name("Name:", 5);
    headerMsg.set_value("!!Value", 7);
    if(headersChannel.write_msg(headerMsg) != 0)
    {
        printf("Error in writing msg for headers.\n"); fflush(stdout); return 0x1;
    }
    printf("Wrote a placeholder message to the header.\n"); fflush(stdout);

    if(headersChannel.read_msg(headerMsg) != 0)
    {
        printf("Error in reading msg for headers.\n"); fflush(stdout); return 0x2;
    }
    printf("Read the following header name and value from the decryptor:\nName:-%s-\n", headerMsg.name().c_str()); fflush(stdout);
    plaintext_header = headerMsg.name();

    std::vector<unsigned char> input_field(headerMsg.value().begin(), headerMsg.value().end());
    ip_list.push_back(input_field);

    ret_status = encrypt_decrypt_msgs(0, ip_list, binary_list);
    if(ret_status != 0)
        return ret_status;

    binary_header = binary_list.at(0);
    binary_ptr = &binary_header[0];
    binary_header_length = binary_header.size();

    base64_ciphertext_header = (uint8_t* ) malloc(2*binary_header_length);
    base64_ciphertext_header_length = base64_encoding_wrapper((unsigned char*) binary_ptr, binary_header_length , base64_ciphertext_header);
    if(base64_ciphertext_header_length <= 0)
    {
        printf("Error encoding the ciphertext header into base64 format.\n"); fflush(stdout);  return 0x3;
    }
    printf("Got this header value:-%s-\n", base64_ciphertext_header);
    plaintext_header.append((const char* )base64_ciphertext_header, (size_t) base64_ciphertext_header_length);
    printf("Got this header:%s\n", plaintext_header.c_str()); fflush(stdout);

    return 0;
}