#include "jwt.h"
#include <dynstr.h>
#include <sodium.h>
#include <errno.h>
#include <stddef.h>
#include <stdlib.h>
#include <stdio.h>
#include <string.h>

#define VARIANT sodium_base64_VARIANT_ORIGINAL

static char *get_payload(const char *const name)
{
    char *ret = NULL, *s = NULL;
    struct dynstr d;

    dynstr_init(&d);

    if (dynstr_append(&d, "{\"name\": \"%s\"}", name))
    {
        fprintf(stderr, "%s: dynstr_append name failed\n", __func__);
        goto end;
    }

    const size_t n = sodium_base64_encoded_len(d.len, VARIANT);
    const unsigned char *const in = (const unsigned char *)d.str;

    if (!(s = malloc(n)))
    {
        fprintf(stderr, "%s: malloc(3): %s\n", __func__, strerror(errno));
        goto end;
    }
    else if (!sodium_bin2base64(s, n, in, d.len, VARIANT))
    {
        fprintf(stderr, "%s: sodium_bin2base64 failed\n", __func__);
        goto end;
    }

    ret = s;

end:
    if (!ret)
        free(s);

    dynstr_free(&d);
    return ret;
}

static char *get_hmac(const void *const buf, const size_t n,
    const void *const key, const size_t keyn)
{
    unsigned char hmac[crypto_auth_hmacsha256_KEYBYTES];
    char *ret = NULL;

    if (keyn != crypto_auth_hmacsha256_KEYBYTES)
    {
        fprintf(stderr, "%s: invalid key size (%zu), expected %u\n", __func__,
            n, crypto_auth_hmacsha256_KEYBYTES);
        goto failure;
    }
    else if (crypto_auth_hmacsha256(hmac, buf, n, key))
    {
        fprintf(stderr, "%s: HMAC failed\n", __func__);
        goto failure;
    }

    const size_t b64n = sodium_base64_encoded_len(sizeof hmac, VARIANT);

    if (!(ret = malloc(b64n)))
    {
        fprintf(stderr, "%s: malloc(3): %s\n", __func__, strerror(errno));
        goto failure;
    }
    else if (!sodium_bin2base64(ret, b64n, hmac, sizeof hmac, VARIANT))
    {
        fprintf(stderr, "%s: sodium_bin2base64 failed\n", __func__);
        goto failure;
    }

    return ret;

failure:
    free(ret);
    return NULL;
}

char *jwt_encode(const char *const name, const void *const key, const size_t n)
{
    static const char jwt_header[] = "{\"alg\": \"HS256\", \"typ\": \"JWT\"}";
    struct dynstr jwt;
    const size_t hlen = strlen(jwt_header),
        sz = sodium_base64_encoded_len(hlen, VARIANT);
    char *ret = NULL, *header = NULL, *hmac = NULL,
        *const payload = get_payload(name);

    dynstr_init(&jwt);

    if (!payload)
    {
        fprintf(stderr, "%s: get_payload failed\n", __func__);
        goto end;
    }
    else if (!(header = malloc(sz)))
    {
        fprintf(stderr, "%s: malloc(3): %s\n", __func__, strerror(errno));
        goto end;
    }
    else if (!sodium_bin2base64(header, sz, (const unsigned char *)jwt_header,
        hlen, VARIANT))
    {
        fprintf(stderr, "%s: sodium_bin2base64 failed\n", __func__);
        goto end;
    }
    else if (dynstr_append(&jwt, "%s.%s", header, payload))
    {
        fprintf(stderr, "%s: dynstr_append header+payload failed", __func__);
        goto end;
    }
    else if (!(hmac = get_hmac(jwt.str, jwt.len, key, n)))
    {
        fprintf(stderr, "%s: get_hmac failed\n", __func__);
        goto end;
    }
    else if (dynstr_append(&jwt, ".%s", hmac))
    {
        fprintf(stderr, "%s: dynstr_append hmac failed\n", __func__);
        goto end;
    }

    ret = jwt.str;

end:
    free(header);
    free(payload);
    free(hmac);

    if (!ret)
        dynstr_free(&jwt);

    return ret;
}

int jwt_check(const char *const jwt, const void *const key,
    const size_t n)
{
    const char *const p = strrchr(jwt, '.');
    const unsigned char *const in = (const unsigned char *)jwt;
    unsigned char hmac[crypto_auth_hmacsha256_KEYBYTES];

    if (n != crypto_auth_hmacsha256_KEYBYTES)
    {
        fprintf(stderr, "%s: invalid key size (%zu), expected %u\n", __func__,
            n, crypto_auth_hmacsha256_KEYBYTES);
        return -1;
    }
    else if (!p)
    {
        fprintf(stderr, "%s: expected '.'\n", __func__);
        return 1;
    }
    else if (sodium_base642bin(hmac, sizeof hmac, p + 1, strlen(p + 1),
        NULL, NULL, NULL, VARIANT))
    {
        fprintf(stderr, "%s: sodium_base642bin failed\n", __func__);
        return 1;
    }
    else if (crypto_auth_hmacsha256_verify(hmac, in, p - jwt, key))
        return 1;

    return 0;
}