#include "crypt.h"

#include <sodium/utils.h>
#include <stdlib.h>

int crypt_key_gen(crypt_key_t *const k) {
    if (crypto_kx_keypair(k->kx_pub, k->kx_sec) < 0) {
        return -1;
    }

    if (crypto_sign_keypair(k->sign_pub, k->sign_sec) < 0) {
        return -1;
    }

    k->hasSecKey = true;

    return 0;
}

int crypt_key_from_hex(crypt_key_t *const k, char phex[CRYPT_PKEY_HEXLEN], char shex[CRYPT_SKEY_HEXLEN]) {
    if (sodium_hex2bin(k->kx_pub, CRYPT_KX_PKEY_LEN, phex, CRYPT_KX_PKEY_LEN * 2, NULL, NULL, NULL) < 0)
        return -1;
    if (sodium_hex2bin(k->kx_sec, CRYPT_KX_SKEY_LEN, shex, CRYPT_KX_SKEY_LEN * 2, NULL, NULL, NULL) < 0)
        return -1;
    if (sodium_hex2bin(k->sign_pub, CRYPT_SIGN_PKEY_LEN, phex+(CRYPT_KX_PKEY_LEN * 2), CRYPT_SIGN_PKEY_LEN * 2, NULL, NULL, NULL) < 0)
        return -1;
    if (sodium_hex2bin(k->sign_sec, CRYPT_SIGN_SKEY_LEN, shex+(CRYPT_KX_SKEY_LEN * 2), CRYPT_SIGN_SKEY_LEN * 2, NULL, NULL, NULL) < 0)
        return -1;

    k->hasSecKey = true;

    return 0;
}

int crypt_key_from_hex_public(crypt_key_t *const k, char phex[CRYPT_PKEY_HEXLEN]) {
    if (sodium_hex2bin(k->kx_pub, CRYPT_KX_PKEY_LEN, phex, CRYPT_KX_PKEY_LEN * 2, NULL, NULL, NULL) < 0)
        return -1;
    if (sodium_hex2bin(k->sign_pub, CRYPT_SIGN_PKEY_LEN, phex+(CRYPT_KX_PKEY_LEN * 2), CRYPT_SIGN_PKEY_LEN * 2, NULL, NULL, NULL) < 0)
        return -1;

    k->hasSecKey = false;

    return 0;
}

void crypt_key_destroy(crypt_key_t *const k) {
    sodium_memzero(k->kx_pub, CRYPT_KX_PKEY_LEN);
    sodium_memzero(k->kx_sec, CRYPT_KX_SKEY_LEN);
    sodium_memzero(k->sign_pub, CRYPT_SIGN_PKEY_LEN);
    sodium_memzero(k->sign_sec, CRYPT_SIGN_SKEY_LEN);
}

int crypt_key_export_public(const crypt_key_t *const k, char hex[CRYPT_PKEY_HEXLEN]) {
    if (sodium_bin2hex(hex, CRYPT_KX_PKEY_LEN * 2 + 1, k->kx_pub, CRYPT_KX_PKEY_LEN) == NULL)
        return -1;
    if (sodium_bin2hex(hex+(CRYPT_KX_PKEY_LEN * 2), CRYPT_SIGN_PKEY_LEN * 2 + 1, k->sign_pub, CRYPT_SIGN_PKEY_LEN) == NULL)
        return -1;

    return 0;
}

int crypt_key_export_secret(const crypt_key_t *const k, char hex[CRYPT_SKEY_HEXLEN]) {
    if (!k->hasSecKey)
        return -1;

    if (sodium_bin2hex(hex, CRYPT_KX_SKEY_LEN * 2 + 1, k->kx_sec, CRYPT_KX_SKEY_LEN) == NULL)
        return -1;
    if (sodium_bin2hex(hex+(CRYPT_KX_SKEY_LEN * 2), CRYPT_SIGN_SKEY_LEN * 2 + 1, k->sign_sec, CRYPT_SIGN_SKEY_LEN) == NULL)
        return -1;

    return 0;
}

int crypt_load_key(crypt_key_t *const k, FILE *const pub, FILE *const sec) {
    if (pub == NULL || sec == NULL)
        return -1;

    char phex[CRYPT_PKEY_HEXLEN], shex[CRYPT_SKEY_HEXLEN];

    if (fread(phex, sizeof(phex[0]), CRYPT_PKEY_HEXLEN, pub) != CRYPT_PKEY_HEXLEN)
        return -1;
    if (fread(shex, sizeof(shex[0]), CRYPT_SKEY_HEXLEN, sec) != CRYPT_SKEY_HEXLEN)
        return -1;

    return crypt_key_from_hex(k, phex, shex);
}

int crypt_store_key(const crypt_key_t *const k, FILE *const pub, FILE *const sec) {
    if (!k->hasSecKey || pub == NULL || sec == NULL)
        return -1;

    char phex[CRYPT_PKEY_HEXLEN], shex[CRYPT_SKEY_HEXLEN];

    if (crypt_key_export_public(k, phex) < 0)
        return -1;
    if (crypt_key_export_secret(k, shex) < 0)
        return -1;

    if (fwrite(phex, sizeof(phex[0]), CRYPT_PKEY_HEXLEN, pub) != CRYPT_PKEY_HEXLEN)
        return -1;
    if (fwrite(shex, sizeof(shex[0]), CRYPT_SKEY_HEXLEN, sec) != CRYPT_SKEY_HEXLEN)
        return -1;

    return 0;
}

unsigned char *crypt_hello(const crypt_key_t *const own, long long unsigned int *const sign_len) {
    char pub_hex[CRYPT_PKEY_HEXLEN];
    crypt_key_export_public(own, pub_hex);
    *sign_len = crypto_sign_BYTES + CRYPT_PKEY_HEXLEN;
    unsigned char *sign = (unsigned char *)malloc(*sign_len * sizeof(unsigned char));
    if (crypto_sign(sign, sign_len, (const unsigned char *)pub_hex, CRYPT_PKEY_HEXLEN, own->sign_sec) == -1)
        return NULL;
    return sign;
}

int crypt_session_init(crypt_session_t *const s, const crypt_key_t *const own, crypt_key_t *const remote, const bool is_client) {
    if (is_client) {
        if (crypto_kx_client_session_keys(s->rx, s->tx, own->kx_pub, own->kx_sec, remote->kx_pub) != 0) {
            fprintf(stderr, "Failed to instantiate a client session.\n");
            return -1;
        }
    } else {
        if (crypto_kx_server_session_keys(s->rx, s->tx, own->kx_pub, own->kx_sec, remote->kx_pub) != 0) {
            fprintf(stderr, "Failed to instantiate a server session.\n");
            return -1;
        }
    }

    s->remote_key = remote;

    return 0;
}

void crypt_session_destroy(crypt_session_t *const s) {
    sodium_memzero(s->rx, CRYPT_SESS_KEY_LEN);
    sodium_memzero(s->tx, CRYPT_SESS_KEY_LEN);
    crypt_key_destroy(s->remote_key);
}