#include "LocalAttestationUntrusted.h"

namespace LocalAttestationUntrusted {
    namespace {
        uint32_t session_id = 0;
        protobuf_sgx_dh_msg1_t protobuf_msg1;

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

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

            protobuf_sgx_ret = protobufLAInitiator::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;
            }

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

    int decrypt_client_data(uint32_t own_enclave_id, int fd, int time_file_fd) {
        extension_to_decryptor_enclosed_msg protobuf_extension_decryptor_msg;
        unsigned char *double_ciphertext, *ciphertext;
        uint32_t *input_sizes_array, *output_sizes_array;
        uint32_t double_ciphertext_length;
        uint32_t number_of_double_ciphertext_fields;
        uint32_t sgx_ret;

        // Get a message of the type decryptor_to_extension msg
        if (!protobufReadWrite::read_protobuf_msg_from_fd(fd, protobuf_extension_decryptor_msg)) {
            printf("Not all of the extension's message was read\n");
            fflush(stdout);
            return 0xf3;
        }

        PostLAProtobufNativeTransforms::get_lengths_for_protobuf_serialized_array(protobuf_extension_decryptor_msg, &double_ciphertext_length,
                &number_of_double_ciphertext_fields);

        double_ciphertext=(unsigned char*) malloc(double_ciphertext_length);
        ciphertext = (unsigned char*) malloc(double_ciphertext_length);

        input_sizes_array = (uint32_t*) malloc(number_of_double_ciphertext_fields * sizeof(uint32_t));
        output_sizes_array = (uint32_t*) malloc(number_of_double_ciphertext_fields * sizeof(uint32_t));
        PostLAProtobufNativeTransforms::create_array_from_protobuf(protobuf_extension_decryptor_msg, double_ciphertext,
            input_sizes_array, &number_of_double_ciphertext_fields);

        // Call the enclave's decryption function with these arguments and get back another vector of ciphertexts.
        Decryptor_process_apache_message_generate_response_wrapper(own_enclave_id, &sgx_ret,
                double_ciphertext,
                double_ciphertext_length,
                input_sizes_array,
                number_of_double_ciphertext_fields,
                ciphertext,
                output_sizes_array);

        free(double_ciphertext);
        free(input_sizes_array);

        // Error checking
        if(!sgx_ret)
        {
            free(ciphertext);
            free(output_sizes_array);
            return 0x32;
        }

        // Clear the protobuf msg above and reset it with the output arguments of the ecall.
        protobuf_extension_decryptor_msg.clear_ciphertext_client_public_key();
        protobuf_extension_decryptor_msg.clear_ciphertext_fields();
        PostLAProtobufNativeTransforms::create_protobuf_from_array(ciphertext, output_sizes_array,
                number_of_double_ciphertext_fields,
                protobuf_extension_decryptor_msg);
        free(ciphertext);
        free(output_sizes_array);

        // write message to apache extension
        if (!protobufReadWrite::write_protobuf_msg_to_fd(fd, protobuf_extension_decryptor_msg)) {
            printf("Not all of the client's ciphertext data was written to the extension.\n");
            fflush(stdout);
            return 31;
        }

        /*
        gettimeofday(&tv2, NULL);
        new_time=tv2.tv_usec + tv2.tv_sec * 1000000;
        old_time=tv1.tv_usec + tv1.tv_sec * 1000000;
        bytes_written=sprintf(time_buf,  "%lu %lu\n", old_time, new_time);
        FileIO::write_to_fd(time_file_fd, time_buf, bytes_written);
        */

        return 0;
    }

    int prepare_local_attestation_as_responder_msg1(uint32_t own_enclave_id)
    {
        uint32_t protobuf_sgx_ret;
        protobuf_sgx_ret = protobufLAInitiator::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;
        }
        return 0;
    }

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

    // TODO: CHANGED SIGNATURE.
    int 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 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[150];
        size_t bytes_read;
        int count, ret_status;

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

        for (count = 0; count < 60; 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,
                                                    60);
        if (protobuf_sgx_ret != 0) {
            printf("Error in decryption: 0x%x\n", protobuf_sgx_ret);
            fflush(stdout);
            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 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[200]; // 176+12 for IV = 188
        uint32_t op_length, internal_return_status, count, sgx_ret;
        uint8_t public_key[64];
        int time_file_fd;

        memset(encrypted_sign_data_and_sign_and_tag, 0x0, 200);

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

        for (count = 0; count < op_length; 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, op_length);
        if (protobufReadWrite::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);
            return 0xfe;
        }

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

        time_file_fd = open("decryptor_time.txt", O_APPEND | O_WRONLY);

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


        close(accept_fd);
        close(time_file_fd);
        return internal_return_status;
    }

};