#include <sys/socket.h>
#include <netinet/in.h>
#include <string.h>
#include <err.h>
#include <inttypes.h>
#include <arpa/inet.h>
#include <unistd.h>
#include <netdb.h>

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

static const unsigned char pubkey[] =
"-----BEGIN PUBLIC KEY-----\n"
"MIIBIjANBgkqhkiG9w0BAQEFAAOCAQ8AMIIBCgKCAQEAwWlNmLHOOzZpdrfp+EAA\n"
"Nwqab0FhQwCyZ/u+ySBW5XxPf6mobySvtJrLdsWzdwnup/UfwZiEhJk/4wpD4Qf/\n"
"2+syuJi3Rf7L+Jfh//Qf9uXAS80+LYad7dW0c1r5nt+F9Can5fBn7futnd8n672T\n"
"+y8QpHRwX9GtaILvYQe5GQac8cfq2rGUd5iYj5KSdcaIZnJ4YgnjLHg2PMbtEJwq\n"
"cV+2oAkcOPzTAJoNE7XjLZTwXmLlFgL/2cN4uJZBDnwv3RZSAhpdYF4KOJmE2GFs\n"
"2jdvRUrYO7WSl8fM16vRH4vz5MNNcaprg2MlXheVTPQa+WMdcz7dyQx8s9kRVPPU\n"
"SwIDAQAB \n"
"-----END PUBLIC KEY----- \n";

static char *pubkey_file = 0;

static const unsigned char pp[] = "9bf308b7ae027baa46091d980632e27b";
static const char *logging_host = "endpoint-de9XDJ0fH7.gsmk.de";
static const char *logging_port = "8238";

static struct sockaddr_storage logging_host_address;
static socklen_t logging_host_address_len = 0;
static int logging_socket = -1;
typedef uint64_t SessionId_t;

enum { SESSION_ID_LENGTH = 8, AES_KEY_LENGTH = 16, GCM_IV_LENGTH = 12, GCM_TAG_LENGTH = 16 };
static SessionId_t session_id;
static uint8_t aes_key[16];
static mbedtls_gcm_context ctx;

void send_udp(int sock, mbedtls_ctr_drbg_context *ctr_drbg, const uint8_t *text, size_t len) {
    uint8_t iv[GCM_IV_LENGTH];
    mbedtls_ctr_drbg_random(ctr_drbg, iv, GCM_IV_LENGTH);

    const size_t total_length = 1 + SESSION_ID_LENGTH + GCM_IV_LENGTH + GCM_TAG_LENGTH + len;
    uint8_t *output = alloca(total_length);
    output[0] = 2;
    memcpy(output + 1, (uint8_t*)&session_id, SESSION_ID_LENGTH);
    memcpy(output + 1 + SESSION_ID_LENGTH, iv, GCM_IV_LENGTH);


    if (!mbedtls_gcm_crypt_and_tag(&ctx, MBEDTLS_GCM_ENCRYPT, len, iv, GCM_IV_LENGTH, (uint8_t*)&session_id, SESSION_ID_LENGTH,
                                  text, output + 1 + SESSION_ID_LENGTH + GCM_IV_LENGTH + GCM_TAG_LENGTH, GCM_TAG_LENGTH,
                                  output + 1 + SESSION_ID_LENGTH + GCM_IV_LENGTH)) {
        sendto(sock, output, total_length, 0, (struct sockaddr*)&logging_host_address, logging_host_address_len);
    }
}

void new_session(int sock, mbedtls_ctr_drbg_context *ctr_drbg) {
    mbedtls_ctr_drbg_random(ctr_drbg, (uint8_t*)&session_id, sizeof(session_id));
    mbedtls_ctr_drbg_random(ctr_drbg, aes_key, sizeof(aes_key));

    unsigned char output[512];
    output[0] = 0;
    memcpy(output + 1, (uint8_t*)&session_id, SESSION_ID_LENGTH);

    mbedtls_pk_context pk;
    mbedtls_pk_init(&pk);
    int ret = 0;

    if (pubkey_file) {
        if ((ret = mbedtls_pk_parse_public_keyfile(&pk, pubkey_file) ) != 0 )
            errx(-1, "mbedtls_pk_parse_public_keyfile returned -0x%04x\n", -ret );
    } else {
        if ((ret = mbedtls_pk_parse_public_key(&pk, pubkey, sizeof(pubkey)) ) != 0 )
            errx(-1, "mbedtls_pk_parse_public_key returned -0x%04x\n", -ret );
    }

    size_t olen = 0;
    if ((ret = mbedtls_pk_encrypt(&pk, aes_key, AES_KEY_LENGTH, output + 1 + SESSION_ID_LENGTH, &olen,
                                  sizeof(output) - 1 - SESSION_ID_LENGTH, mbedtls_ctr_drbg_random, ctr_drbg)) != 0)
        errx(-1, "mbedtls_pk_encrypt returned -0x%04x\n", -ret );

    mbedtls_pk_free(&pk);

    sendto(sock, output, olen + 1 + SESSION_ID_LENGTH, 0, (struct sockaddr*)&logging_host_address, logging_host_address_len);

    mbedtls_gcm_init(&ctx);
    mbedtls_gcm_setkey(&ctx, MBEDTLS_CIPHER_ID_AES, aes_key, 8 * AES_KEY_LENGTH);
}

int main(int argc, char **argv) {
    const char * host = logging_host, * port = logging_port;
    char ch;
    while ((ch = getopt(argc, argv, "h:p:c:")) != -1) {
        switch (ch) {
        case 'h':
            host = optarg;
            break;
        case 'p':
            port = optarg;
            break;
        case 'c':
            pubkey_file = optarg;
            break;
        case '?':
        default:
            printf("Usage: %s [-h host] [-p port] [-c cert]", argv[0]);
            exit(0);
        }
    }
    argc -= optind;
    argv += optind;

    struct addrinfo hints, *result, *rp;
    memset (&hints, 0, sizeof (hints));
    hints.ai_socktype = SOCK_DGRAM;

    int sock = -1, res = getaddrinfo(host, port, &hints, &result);
    if (res != 0)
        errx(EXIT_FAILURE, "getaddrinfo: %s\n", gai_strerror(res));

    for (rp = result; rp != NULL; rp = rp->ai_next) {
        sock = socket(rp->ai_family, rp->ai_socktype, rp->ai_protocol);
        if (sock == -1)
            continue;
        memcpy(&logging_host_address, rp->ai_addr, rp->ai_addrlen);
        logging_host_address_len = rp->ai_addrlen;
    }
    if (sock == -1)
        errx(EXIT_FAILURE, "Can't open socket");
    freeaddrinfo(result);

    // Setup
    mbedtls_ctr_drbg_context ctr_drbg;
    mbedtls_entropy_context entropy;

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

    new_session(sock, &ctr_drbg);

    sleep(3);

    // Fire
    const unsigned char *logline = (const unsigned char*)"Juchuuu, es klappt!\n";
    send_udp(sock, &ctr_drbg, logline, strlen((char*)logline));
    send_udp(sock, &ctr_drbg, logline, strlen((char*)logline));
    send_udp(sock, &ctr_drbg, logline, strlen((char*)logline));

    mbedtls_gcm_free(&ctx);
    close(sock);
}