#include <iostream>
#include <dlfcn.h>
#include <openssl/evp.h>
#include <vector>
#include <chrono>
#include "openvpn-plugin.h"


#define STRUCT_VERSION 5

#define SMARTCARD_KEY_WRAPPING_PATH     "./plugins/libSmartcardKeyWrappingLibrary.so"
#define SMARTCARD_KEY_DERIVATION_PATH   "./plugins/libSmartcardKeyDerivationLibrary.so"
#define YUBIKEY_PATH                    "./plugins/libYubikeyKeyManagementLib.so"
#define SOFTHSM_PATH                    "./plugins/libSoftHSMKeyWrappingLibrary.so"

#define SOFTHSM_WKC "kTkvfenmeqP7LR3uCgR3dZBt/E2XKVwEd8ZpUlpMqwdY0wx1iND6Yv2021y1VODpM0Z8y6gebW6Syf0Ye/EDjKr/Xv/qSos6l/U+mVeU+iizCr9844qx5iDyQI5VAducK9daf7nTdkZmeQVK80CvhtxbOsGUDcmoPQb6mHsbJXwmyoyLsxM2ubp7EWx1ZBNIYtor9LKq4uyqQFthKkPfp9Ab0//asGNaowHc5q1oCRZLquhDD03uZyMUj0Fb/Aqgol5zroSK2JW+hJzt+ngKbf5twLDwR5ksG6CS2qf6uD4c85CgUVGhf7hMWWs3eX6KVTzXpByeeUXjLJUQ/B9xFKkhPhF6vGsQeO4O8VaQXma3D2pDidqFGpItd+Y7fTDPsbwcGOkVbVRBASs="
#define YUBIKEY_WKC "+JBssE/PrQvQNdUm8Utbe06drdDnpB5+/31VdJLlZE7t5MufueRRHUVdD1xADWQA2QH1oHXUyYT204TxxvS362ZuaoVbYjF3aUfMjqmTQOCUxoE6piJvQdibNOMGPCoeO8NVTnXY92s602yBaIsPoMvtf5VCV3FFd9Fb56Ai0dmF9NErnW58oNsxXCFOJaQxbHmPmRGPqWQPxhirYBsW0xWYykkOQDtwV7Er/I6fPJh6bhDaSU//eVjGkqno+/xIUh5+YWSTQVIuhXCzySCp3CfVGo51RZ8k0NOwaQpSYqomshu7gDYjZiaPlTxdNyNFLLJbArQm9TNifR41hKx+39xjXZlhRa8BrOquih0DUFEaX5KQCfjHCDds6Wj1nSxVLl9KrNXR6tfRASs="
#define SMARTCARD_KEY_WRAPPING_WKC "fUH11DDJbn4VM2IIePOK/Uu1ytVbpTtBoFsTFjKhrM79AFUlx6ySchE0Xck3FSyvwtlBZQ+QRK7Gd/+x0xsNeL/ha7hKkRrW2JK6dPXN9qdAqhf50ZJKIwEboOZMD+pzrzJGeNslTwmRWOxuRXWFBQafjElePDSLZBKSzBJnAT5j6tB0vAXdS34j8IzyuuyLnoiDzrCm6S7aoaRfotLpxsLefPatKGwvoimxLEyfIcU4IaMhNktb7iRlFZAM15TLcEkQloB8Hv2moe7pus5EJHxsHPt/WmJYD7Naz6PXS+ynip++z6bu6LUC8kPIcnhqfHU7ieGnP8L2KD70Y/uQ+OkrcmCRhonkBdPgIFK/+RLBNJQtDwjyb6dXTgvn0LcCmLZx/lVfBy91ASs="
#define SMARTCARD_KEY_DERIVATION_WKC "KW9RjSfhyM6JBgfvs7/SDcFD9eFwO/5FiDOChWE1PZ31pfNhhYl5Iytg7Y6VDGuWNyhYq47koVuXdLSpva5qNSt1uRcqaY1nNbWwBMCgSIjcNJb3+t0FxtbZeYP4w78v18Ad8j3Qf2e3T2NN4VvEebffsihPf0vlvqGaRJvn7W8NXh3Y28/vPb4tXLSalpR5+lZrXHLwn+pl03sH3+nSYIO4Fk+DbSA0aVTKXDulVtNYugMy2P9XqQRtlM1lOuMBqxbyc2FB7oqFvrLTHVZnWqGFimvwasIaktKfkxXT+X1uaoUYzsFNbvs8s/Xul+9NF4YXMZpffMbA5ah5629zG8HW9L1tI9crd5GOKdH5QSThtSZinnSXhmhYCvAzhNdjDQc/oCpIicpeASs="

extern "C"
{
    #include "base64.h"
}

void vlog_func(openvpn_plugin_log_flags_t flags, const char *plugin_name, const char *format, va_list arglist) {
    (void)plugin_name;
    // Ignore notes and debug messages
    vprintf(format, arglist);
    puts("");
}

void plugin_secure_memzero(void *data, size_t len) {
    (void) data;
    (void) len;

    // Not Implemented
}

void plugin_log(openvpn_plugin_log_flags_t flags, const char *plugin_name, const char *format, ...) {
    (void) flags;
    (void) plugin_name;
    (void) format;

    // Not Implemented
}

std::string bytes_to_hex(const unsigned char *data, int size) {
    char const hex_chars[16] = { '0', '1', '2', '3', '4', '5', '6', '7', '8', '9', 'A', 'B', 'C', 'D', 'E', 'F' };
    std::string hex;
    for( int i = 0; i < size; i++) {
        unsigned char byte = data[i];
        hex.push_back(hex_chars[ ( byte & 0xF0 ) >> 4 ]);
        hex.push_back(hex_chars[ ( byte & 0x0F ) >> 0 ]);
    }
    return hex;
}

class Plugin
{
    void *plugin_lib;
    int (*openvpn_plugin_open_v3)(const int v3structver,
                                  struct openvpn_plugin_args_open_in const *args,
                                  struct openvpn_plugin_args_open_return *ret);
    int (*openvpn_plugin_func_v3)(const int version,
                                  struct openvpn_plugin_args_func_in const *arguments,
                                  struct openvpn_plugin_args_func_return *retptr);
    openvpn_plugin_handle_t handle{};
public:
    explicit Plugin(const std::string &plugin_path);
    ~Plugin();
    void open_plugin(int num, ...);
    std::string execute_plugin(int call_type, int num, ...);

    int up_plugin();
};

Plugin::Plugin(const std::string &plugin_path)
{
    this->plugin_lib = dlopen(plugin_path.c_str(), RTLD_LAZY);

    if (!this->plugin_lib) {
        std::cerr << dlerror() << std::endl;
        throw std::invalid_argument("Couldn't load plugin library");
    }

    this->openvpn_plugin_open_v3 = reinterpret_cast<int (*)(const int v3structver, struct openvpn_plugin_args_open_in const *args, struct openvpn_plugin_args_open_return *ret)>
                                     (
                                             dlsym(this->plugin_lib, "openvpn_plugin_open_v3")
                                     );
    this->openvpn_plugin_func_v3 = reinterpret_cast<int (*)(const int version, struct openvpn_plugin_args_func_in const *arguments, struct openvpn_plugin_args_func_return *retptr)>
                                     (
                                             dlsym(this->plugin_lib, "openvpn_plugin_func_v3")
                                     );

    char *error;
    if ((error = dlerror()) != nullptr) {
        std::cerr << error << std::endl;
        throw std::exception();
    }
}

void free_stringlist(openvpn_plugin_string_list **string_list)
{
    auto next = *string_list;
    while (next)
    {
        auto current = next;
        next = current->next;
        free(current);
    }
    free(string_list);
}

void Plugin::open_plugin(int num, ...) {
    va_list valist;
    va_start(valist, num);
    const char **argv_open = static_cast<const char **>(
                             calloc(num, sizeof(char *)));
    for (int i = 0; i < num; i++)
    {
        argv_open[i+1] = va_arg(valist, char *);
    }

    openvpn_plugin_callbacks callbacks = {
            nullptr,
            vlog_func,
            plugin_secure_memzero,
            openvpn_base64_encode,
            openvpn_base64_decode,
    };

    struct openvpn_plugin_args_open_in args_open =
            {
                    0,
                    reinterpret_cast<const char **const>(argv_open),
                    nullptr,
                    &callbacks,
                    SSLAPI_OPENSSL,
                    nullptr,
                    0,
                    0,
                    nullptr
            };

    struct openvpn_plugin_args_open_return ret_open{};
    auto retlist = static_cast<openvpn_plugin_string_list **>(
            calloc(1,sizeof(openvpn_plugin_string_list *)));
    ret_open.return_list = retlist;

    openvpn_plugin_open_v3(STRUCT_VERSION, &args_open, &ret_open);
    this->handle = ret_open.handle;

    free_stringlist(ret_open.return_list);
    ret_open.return_list = nullptr;
}

std::string Plugin::execute_plugin(int call_type, int num, ...) {
    va_list valist;
    va_start(valist, num);
    const char **argv_func = static_cast<const char **>(
            calloc(num, sizeof(char *)));
    for (int i = 0; i < num; i++)
    {
        argv_func[i+1] = va_arg(valist, char *);
    }

    struct openvpn_plugin_args_func_in args_func = {
            call_type,
            argv_func,
            nullptr,
            this->handle,
            nullptr,
            0,
            nullptr
    };

    struct openvpn_plugin_args_func_return ret_func{};

    auto retlist = static_cast<openvpn_plugin_string_list **>(
                                          calloc(1,sizeof(openvpn_plugin_string_list *)));
    ret_func.return_list = retlist;

    int ret_code = openvpn_plugin_func_v3(STRUCT_VERSION, &args_func, &ret_func);

    if(ret_code != OPENVPN_PLUGIN_FUNC_SUCCESS)
        return "";
    char * return_base64 = (*ret_func.return_list)->value;
    unsigned char return_bytes[strlen(return_base64)];
    int byte_len = openvpn_base64_decode(return_base64, return_bytes, (int) sizeof(return_bytes));
    std::string return_hex = bytes_to_hex(return_bytes, byte_len);

    free_stringlist(ret_func.return_list);
    ret_func.return_list = nullptr;

    return return_hex;
}

int Plugin::up_plugin() {
    struct openvpn_plugin_args_func_in args_func = {
            OPENVPN_PLUGIN_UP,
            nullptr,
            nullptr,
            this->handle,
            nullptr,
            0,
            nullptr
    };

    struct openvpn_plugin_args_func_return ret_func{};

    auto retlist = static_cast<openvpn_plugin_string_list **>(
            calloc(1,sizeof(openvpn_plugin_string_list *)));
    ret_func.return_list = retlist;

    int ret_code = openvpn_plugin_func_v3(STRUCT_VERSION, &args_func, &ret_func);

    free_stringlist(ret_func.return_list);
    ret_func.return_list = nullptr;

    return ret_code;
}

Plugin::~Plugin() {
    dlclose(this->plugin_lib);
}

enum plugin_types {
    SMARTCARD_KEY_WRAPPING,
    SMARTCARD_KEY_DERIVATION,
    YUBIKEY,
    SOFTHSM
};

Plugin createPlugin(int type) {
    Plugin *plugin = nullptr;

    switch(type)
    {
        case SMARTCARD_KEY_WRAPPING:
            plugin = new Plugin(SMARTCARD_KEY_WRAPPING_PATH);
            plugin->open_plugin(0);
            break;
        case SMARTCARD_KEY_DERIVATION:
            plugin = new Plugin(SMARTCARD_KEY_DERIVATION_PATH);
            plugin->open_plugin(0 ,"/usr/lib/pkcs11/libsofthsm2.so");
            break;
        case YUBIKEY:
            plugin = new Plugin(YUBIKEY_PATH);
            // Slot: 2; Access_Code: 0
            plugin->open_plugin(2 , "2", "0");
            plugin->up_plugin();
            break;
        case SOFTHSM:
            plugin = new Plugin(SOFTHSM_PATH);
            plugin->open_plugin(1 ,"/usr/lib/pkcs11/libsofthsm2.so");
            break;
        default:
            throw std::invalid_argument("Plugin type not implemented!");
    }

    return *plugin;
}

int main() {
    Plugin plugin = createPlugin(SMARTCARD_KEY_DERIVATION);

    int iterations = 10;
    std::string WKc = SMARTCARD_KEY_DERIVATION_WKC;
    auto begin = std::chrono::high_resolution_clock::now();

    for(int i = 0; i < iterations; i++)
    {
        plugin.execute_plugin(OPENVPN_PLUGIN_CLIENT_KEY_WRAPPING, 2, "unwrap", WKc.c_str());

        std::cout << "\r" << i+1 << " / " << iterations << " completed";
        std::cout.flush();
    }

    std::cout << std::endl;

    auto end = std::chrono::high_resolution_clock::now();
    float duration = (static_cast<float>(std::chrono::duration_cast<std::chrono::nanoseconds>(end-begin).count()) / 1000000);
    std::cout << duration << "ms total, average : " << duration / iterations << "ms." << std::endl;

    std::string result = plugin.execute_plugin(OPENVPN_PLUGIN_CLIENT_KEY_WRAPPING, 2, "unwrap", WKc.c_str());
    /*
    if(result != "1298789D57618EC2FEF4936C6DD96CB56BCDF804816E02286F14B4156D13BED63279101A4FFAEE654A3FE6F990428061A1F9DE83ABB1B1347B865BC3A68A4ED57446BC0B6206B45640C50A3CA776DA195FD38126A471EBE296E1F77BB4581EA605DA37EDF4946843B7ABAD840A08EB1F3C53F818D6C4D1664CF7A14412ED79D329ED22B66D7771350397B4EE8B6D9D94E666809D8B8C5E7ACF1E536D78DC07ABA147E73FDA4F0E527FD87B8B28CF425047B128890975F156EC48237D8D4BCEA8EA361573BC187B61984764344AA5C778972E249F8AA08097C0FCE9DB0C69C1A36A01BF71C73E3425B9B66F241990A423DCFF9642CD97E11D4B90A2E7CF0732BF010000000063920AA4")
        throw std::exception();
    */
    return EXIT_SUCCESS;
}