#include <regex.h>
#include <string.h>
#include <dirent.h>
#include <err.h>
#include <stdio.h>
#include <inttypes.h>
#include <fcntl.h>
#include <unistd.h>
#include <sys/socket.h>
#include <sys/stat.h>
#include <limits.h>
#include <netinet/in.h>
#include <memory>

#include <mbedtls/pk.h>
#include <mbedtls/entropy.h>
#include "mbedtls/ctr_drbg.h"
#include "mbedtls/gcm.h"

#include <string>
#include <map>
#include <set>
#include <iostream>

const unsigned short PORT = 58132;

static const unsigned char privkey[] =
"-----BEGIN RSA PRIVATE KEY-----\n"
"MIIEpAIBAAKCAQEAwWlNmLHOOzZpdrfp+EAANwqab0FhQwCyZ/u+ySBW5XxPf6mo\n"
"bySvtJrLdsWzdwnup/UfwZiEhJk/4wpD4Qf/2+syuJi3Rf7L+Jfh//Qf9uXAS80+\n"
"LYad7dW0c1r5nt+F9Can5fBn7futnd8n672T+y8QpHRwX9GtaILvYQe5GQac8cfq\n"
"2rGUd5iYj5KSdcaIZnJ4YgnjLHg2PMbtEJwqcV+2oAkcOPzTAJoNE7XjLZTwXmLl\n"
"FgL/2cN4uJZBDnwv3RZSAhpdYF4KOJmE2GFs2jdvRUrYO7WSl8fM16vRH4vz5MNN\n"
"caprg2MlXheVTPQa+WMdcz7dyQx8s9kRVPPUSwIDAQABAoIBAH1KD8A4flYRO2Ry\n"
"YxgzrW/6aGxlt/HFg8ykYcS8NE5Yps8WQkwtQb0HAYKhM06LmpQm0DmC6WVUOPSE\n"
"c9BUdEQsKiE2nJK1KcCR8w7xP7uavWTdQcgQCkJFS63mYwmt1oKAgAcOIuUhQiig\n"
"pKWrmy7+IBPIcftAQssO9q6uaBNy+ONu6KU/FYd4UAoEt07MzuIAl5rybROOWCrA\n"
"cjuOKK830Q2Mi2ESwwlO0+3Vyz9VhSha5wW2WwFv9zx8YQblaVXxfXC6O5XcXb0j\n"
"O2rTpgHMmOVik3Zrg3XoRXsNZCvFQqbIevwGhRNEFQTnzakh/5g0VDa54RAK/APC\n"
"t3ABEAECgYEA+rnqDqmXQc/xMCLOhJduxRFvZbT6EbUVJG5+l6XtATTa4LWzqx1B\n"
"QDXS/Dixxc9FA1vSH2rAW//6wbr8KXihSeI7YxgIfWyrSIxbS3Dd1qwddvniYIpx\n"
"ms9vYpQKAewv2p310nf7fyuURES8YikhpSuff3DhzXBEi7s3Y7gIIEsCgYEAxXrD\n"
"6xjgLSgbbdyqxKeNk8ADMifbj8ZoiNIHkJ6sShFaTEbyHweOJ4RY7OmTBUDhgzUY\n"
"1oynztilADFaq8KhsiMqzI3DgZ3/2OElGYReEAbXljvgudL3tmiBWc5j9S7XNLBe\n"
"u9f3WYAYDu7/BVqmQWa/QtD9JquoZ1xgV2D6nAECgYEA29w9t9/VSJvM9xX+nNyi\n"
"AON6GOjrRK3TPWA7WEXjH+S2bshHJi0ANAs+2Xfpw/kunnRdPLmCtuowfMO4LbGf\n"
"VcexpgLEJyAszvBteikeDwpcyCD19wxP9J4kIYCJiggQKpfLoWUfP/P6DydrPnSt\n"
"EUbAlaNqDpl9Mj7YonQVhCMCgYEAv0c9M5+hrDuX7d76zYaZvI4UymUO54E/yZ7e\n"
"UvdOXGPYed+SL/oKeD5aQAeyLzl79bHdgBs3g0QW9kvXzly0cC5eC0oZH5hhs7nI\n"
"TKII1i86bLtM3dD5vQYWnF0sNtWK/+8Bo6L5ZAiNxRE7lP0L4ndaNKbnPaixcoRo\n"
"kNpPhAECgYAg7jmNlw+7VVurzR36LKKE+d6veraF5ltpJiboDb3j38RGe3982LLq\n"
"WaBKm1gkHfXgBjkNzS4r2kyRijw0GQ9JgmWooR7BXFH30HkPNl4gFTSsrOG5zGLi\n"
"0aexkDpXQuKsgBzqU0Wn94GZMM+RhuOYec/56JFT+8U5Tcntb26wwA==\n"
"-----END RSA PRIVATE KEY-----\n";

static const unsigned char pp[] = "IJUHZGFDXTZKHJKHGFDHZLUÖDRTFGHHJGHH";

// Close files after 20 minutes
#define ACCESS_TIMEOUT (20*60)

/*
   1 byte type: 0 allocate session, 1 log to session, rest: discard
   8 bytes session id
   type 0:
      256 bytes rsa 2048 data yielding 32 byte AES session key
   type 1:
      16 bytes iv
      16 bytes tag
      rest cipher text

   File name format: 2020-10-13-23-42-05-SSSSSSSSSSSSSSSS-KKKKKKKKKKKKKKKKKKKKKKKKKKKKKKKK.log

*/

// Time helper
static time_t now() {
    struct timespec tp;
    clock_gettime(CLOCK_MONOTONIC, &tp);
    return tp.tv_sec;
}

// Constants
enum { SESSION_ID_LENGTH = 8, AES_KEY_LENGTH = 16, GCM_IV_LENGTH_LEGACY = 16, GCM_IV_LENGTH = 12, GCM_TAG_LENGTH = 16, MIN_PACKET_SIZE = 40 };
enum { DIRNAME_LENGTH = 10, FILENAME_LENGTH = 73, SIDOFFS = 20, KEYOFFS = 37 };

class Session {
public:
    Session(uint64_t session_id, uint8_t key[AES_KEY_LENGTH], const std::string &dirname, const std::string &filename) :
            _session_id(session_id), _dirname(dirname), _filename(filename) {
        memcpy(_key, key, AES_KEY_LENGTH);
        mbedtls_gcm_init(&_ctx);
        mbedtls_gcm_setkey(&_ctx, MBEDTLS_CIPHER_ID_AES, _key, 8 * AES_KEY_LENGTH);
//        std::cerr << "Imported file " << _filename << " in dir " << _dirname << std::endl;
    }

    Session(uint64_t session_id, uint8_t key[AES_KEY_LENGTH]) : _session_id(session_id) {
        memcpy(_key, key, AES_KEY_LENGTH);
        mbedtls_gcm_init(&_ctx);
        mbedtls_gcm_setkey(&_ctx, MBEDTLS_CIPHER_ID_AES, _key, 8 * AES_KEY_LENGTH);

        // Create timestamp
        char tprefix[32];
        time_t t = time(NULL);
        struct tm * jetzt = localtime(&t);
        strftime(tprefix, sizeof(tprefix), "%F-%H-%M-%S", jetzt);

        // Dump key
        char hexkey[2*AES_KEY_LENGTH + 1];
        for (int i=0; i<AES_KEY_LENGTH; ++i)
            sprintf(hexkey + 2 * i, "%02x", _key[i]);

        // Glue together serialisation
        char filename[FILENAME_LENGTH + 1];
        size_t nlen = snprintf(filename, sizeof(filename), "%s-%016" PRIx64 "-%s.log", tprefix, _session_id, hexkey);
        _filename = std::string(filename, filename + nlen);

        _dirname = std::string(filename, DIRNAME_LENGTH);

        // Touch file to save session_id and key
        close(ensure_file());
    }

    ~Session() {
        mbedtls_gcm_free(&_ctx);
        if (_fd != -1)
            ::close(_fd);
        _fd = -1;
    }

    int ensure_file() {
        mkdir(_dirname.c_str(), 0755);
        return open((_dirname + "/" + _filename).c_str(), O_WRONLY | O_APPEND | O_CREAT, 0755);
    }

    void write_log(const uint8_t *packet, size_t len, int version = 2) {
        const size_t iv_len = ( version == 1 ? GCM_IV_LENGTH_LEGACY : GCM_IV_LENGTH);

        // First check if the packet holds enough space for session id, iv and at least one gcm block
        if (len < iv_len + GCM_TAG_LENGTH) {
            std::cerr << "Error: Short packet, size " << len << std::endl;
            return;
        }

        const uint8_t *iv = packet;
        const uint8_t *tag = packet + iv_len;
        const uint8_t *payload = packet + iv_len + GCM_TAG_LENGTH;
        len -= iv_len + GCM_TAG_LENGTH;

        std::string ivs(packet, packet + iv_len);
        if (_used_ivs.find(ivs) != _used_ivs.end()) {
            std::cerr << "Error: Session " << std::hex << _session_id << " reused IV. Dropping packet" << std::endl;
            return;
        }
        _used_ivs.insert(ivs);

        // Create output file if it doesn't exist
        if (_fd < 0)
            _fd = ensure_file();
        if (_fd < 0) {
            std::cerr << "Error: Can't create file " << _filename << " for session " << std::hex << _session_id;
            return;
        }

        // Prepare output scratch space
        uint8_t *output = static_cast<uint8_t*>(alloca(len));

        // This should fail on invalid input sizes
        switch (mbedtls_gcm_auth_decrypt(&_ctx, len, iv, iv_len, (uint8_t*)&_session_id, SESSION_ID_LENGTH, tag, GCM_TAG_LENGTH, payload, output))
        {
            case 0:
                write(_fd, output, len);
                _last_access = now();
                break;
            case MBEDTLS_ERR_GCM_BAD_INPUT:
                std::cerr << "Error: Invalid log data" << std::endl;
                break;
            case MBEDTLS_ERR_GCM_AUTH_FAILED :
                std::cerr << "Error: Can't decrypt" << std::endl;
                break;
            default:
                std::cerr << "Error: Unknown gcm error" << std::endl;
        }
    }

    void check_timeout(time_t when) {
        if (_last_access && _last_access < when && _fd != -1) {
            ::close(_fd);
            _fd = -1;
        }
    }

private:
    uint64_t                _session_id;
    uint8_t                 _key[AES_KEY_LENGTH];
    int                     _fd = -1;
    time_t                  _last_access = 0;
    std::string             _dirname;
    std::string             _filename;
    std::set<std::string>   _used_ivs;
    mbedtls_gcm_context     _ctx;
};

std::map<uint64_t, std::unique_ptr<Session>> g_sessions;

static uint8_t hex2nyble(char c)
{
    return (c>='0'&&c<='9') ? (uint8_t)(c-'0')
        : (c>='a'&&c<='f') ? (uint8_t)(c-'a'+10)
        : (c>='A'&&c<='F') ? (uint8_t)(c-'A'+10)
        : 0;
}

static void import_directory(const char *path, std::string dirname) {
    DIR * dirp = opendir(path);
    if (!dirp)
        errx(-1, "Fatal: Can't open dir %s\n", path);

    regex_t regex;
    if (regcomp(&regex, "^[[:digit:]]{4}-[[:digit:]][[:digit:]]-[[:digit:]][[:digit:]]-[[:digit:]][[:digit:]]-[[:digit:]][[:digit:]]-"
                         "[[:digit:]][[:digit:]]-[[:xdigit:]]{16}-[[:xdigit:]]{32}.log$", REG_EXTENDED))
        errx(-1, "Fatal: Can't compile re");

    struct dirent * entry;
    while ((entry = readdir(dirp)) != NULL) {
        // We expect a very specific format
        if (entry->d_type != DT_REG || strlen(entry->d_name) != FILENAME_LENGTH)
            continue;

        std::string filename(entry->d_name);
        uint64_t session_id;
        uint8_t aeskey[AES_KEY_LENGTH];

        if (regexec(&regex, filename.c_str(), (size_t) 0, NULL, 0) ||
            sscanf(filename.c_str() + SIDOFFS, "%" SCNx64, &session_id) != 1) {
            std::cerr << "Skipping non-re-matching file: " << filename << std::endl;
            continue;
        }

        const char * hexkey = filename.c_str() + KEYOFFS;
        for (int i=0; i<16; ++i)
            aeskey[i] = (hex2nyble(hexkey[2*i]) << 4 ) | hex2nyble(hexkey[2*i+1]);
        g_sessions[session_id] = std::make_unique<Session>(session_id, aeskey, dirname, filename);
    }
    closedir(dirp);
    regfree(&regex);
}

static void import_sessions(const char *root_dir) {
    char dirpath[PATH_MAX];
    DIR * dirp = opendir(root_dir);
    if (!dirp)
        errx(-1, "Fatal: Can't open dir %s\n", root_dir);

    struct dirent * entry;
    while ((entry = readdir(dirp)) != NULL) {
        if (entry->d_type != DT_DIR || strlen(entry->d_name) != DIRNAME_LENGTH)
            continue;
        std::string dirname(entry->d_name);
        snprintf(dirpath, PATH_MAX, "%s/%s", root_dir, entry->d_name);
        import_directory(dirpath, dirname);
    }
    closedir(dirp);
}

int main(int argc, char **argv) {
    mbedtls_ctr_drbg_context ctr_drbg;
    mbedtls_entropy_context entropy;
    mbedtls_pk_context pk;
    int ret = 0, clport = PORT;
    unsigned char result[256];
    unsigned char input[256];
    size_t inputlen = 0;
    char * privkey_file = 0;
    char * me = argv[0];

    char ch;
    while ((ch = getopt(argc, argv, "p:k:")) != -1) {
        switch (ch) {
        case 'k':
            privkey_file = optarg;
            break;
        case 'p':
            clport = std::stoi(optarg);
            break;
        default:
            errx(-1, "Syntax: %s [sessionlogdir] xx", me);
            break;
        }
    }
    argc -= optind;
    argv += optind;

    if (argc != 1)
        errx(-1, "Syntax: %s [sessionlogdir]", me);
    chdir(argv[0]);

    mbedtls_pk_init( &pk );
    mbedtls_entropy_init( &entropy );
    mbedtls_ctr_drbg_init( &ctr_drbg );
    mbedtls_ctr_drbg_seed(&ctr_drbg, mbedtls_entropy_func, &entropy, pp, sizeof(pp));

    if (privkey_file) {
        if ((ret = mbedtls_pk_parse_keyfile(&pk, privkey_file, NULL) ) != 0 )
            errx(-1, "Fatal: mbedtls_pk_parse_keyfile returned -0x%04x\n", -ret );
    } else {
        if ((ret = mbedtls_pk_parse_key(&pk, privkey, sizeof(privkey), NULL, 0) ) != 0 )
            errx(-1, "Fatal: mbedtls_pk_parse_key returned -0x%04x\n", -ret );
    }

    int sock = socket(AF_INET, SOCK_DGRAM, 0);
    if (sock < 0)
        errx(-1, "Fatal: Can not make socket\n");

    struct sockaddr_in servaddr, peer;
    servaddr.sin_family = AF_INET;
    servaddr.sin_addr.s_addr = INADDR_ANY;
    servaddr.sin_port = htons(clport);
    if (bind(sock, (const struct sockaddr *)&servaddr, sizeof(servaddr)) < 0)
        errx(-1, "Fatal: Can't bind to port %d\n", clport);

    import_sessions(".");

    time_t last_close_check = now();
    while(1) {
        uint8_t packet[16*1024];
        uint8_t rsa_plain_text[AES_KEY_LENGTH];
        size_t olen;
        socklen_t peer_len = sizeof(peer);

        ssize_t len = recvfrom(sock, packet, sizeof(packet), MSG_WAITALL, (struct sockaddr *) &peer, &peer_len);
        if (len < 0) {
            errx(-1, "Fatal: recvfrom yields %zd\n", len);
            g_sessions.clear();
        }

        if (len < MIN_PACKET_SIZE) {
            std::cerr << "Skipping short packet" << std::endl;
            continue;
        }

        uint64_t session_id = *(uint64_t*)(packet + 1);
        auto session = g_sessions.find(session_id);

        switch(packet[0]) {
        case 0: // Session setup
            ret = mbedtls_pk_decrypt(&pk, packet + 1 + SESSION_ID_LENGTH, len - 1 - SESSION_ID_LENGTH,
                                     rsa_plain_text, &olen, sizeof(rsa_plain_text), mbedtls_ctr_drbg_random, &ctr_drbg);
            if (ret < 0) {
                std::cerr << "Error: Failed to decrypt key (error -" << -ret << " == -" << std::hex << -ret << ") for session " << session_id << std::dec << std::endl;
                break;
            }

            if (session == g_sessions.end())
                g_sessions[session_id] = std::make_unique<Session>(session_id, rsa_plain_text);
            break;
        case 1:
        case 2:
            if (session != g_sessions.end())
                session->second->write_log(packet + 1 + SESSION_ID_LENGTH, len - 1 - SESSION_ID_LENGTH, static_cast<int>(packet[0]));
            else
                std::cerr << "Error: Can't log to unknown session " << std::hex << session_id << std::dec << std::endl;
            break;
        default:
            break;
        }

        time_t jetzt = now();
        if (jetzt > last_close_check + 60) {
            for (auto &&session: g_sessions)
                session.second->check_timeout(jetzt - ACCESS_TIMEOUT);
            last_close_check = jetzt;
        }
    }
}