/*
 * nanobbs, a tiny forums software.
 * Copyright (C) 2025-2026  Xavier Del Campo Romero
 *
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU Affero General Public License as published by
 * the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU Affero General Public License for more details.
 *
 * You should have received a copy of the GNU Affero General Public License
 * along with this program.  If not, see <https://www.gnu.org/licenses/>.
 */

#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_URLSAFE

static char *b64(const void *const p, const size_t len)
{
    char *ret = NULL;
    const size_t n = sodium_base64_encoded_len(len, VARIANT);

    if (!(ret = malloc(n)))
    {
        fprintf(stderr, "%s: malloc(3): %s\n", __func__, strerror(errno));
        return NULL;
    }
    else if (!sodium_bin2base64(ret, n, p, len, VARIANT))
    {
        fprintf(stderr, "%s: sodium_bin2base64 failed\n", __func__);
        free(ret);
        return NULL;
    }

    return ret;
}

static int b64d(const char *const b64, const size_t len, void **const outbuf,
    size_t *const outsz)
{
    int ret = 1;
    size_t sz = 0, binlen;
    const char *end;
    void *buf = NULL, *p;

again:

    if (!(p = realloc(buf, ++sz)))
    {
        fprintf(stderr, "%s: realloc(3): %s\n", __func__, strerror(errno));
        goto failure;
    }
    else if (sodium_base642bin(buf = p, sz, b64, len, NULL, &binlen, &end,
        VARIANT))
    {
        if (errno == ERANGE)
            goto again;

        fprintf(stderr, "%s: sodium_base642bin failed\n", __func__);
        goto failure;
    }

    *outbuf = buf;
    *outsz = sz;
    return 0;

failure:
    free(buf);
    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;

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

    return ret;
}

char *jwt_encode(const cJSON *const j, const void *const key, const size_t n)
{
    static const char jwt_header[] = "{\"alg\": \"HS256\", \"typ\": \"JWT\"}";
    struct dynstr d;
    char *ret = NULL, *p = NULL, *header = NULL, *hmac = NULL, *encp = NULL;

    dynstr_init(&d);

    if (!(p = cJSON_PrintUnformatted(j)))
    {
        fprintf(stderr, "%s: cJSON_PrintUnformatted failed\n", __func__);
        goto end;
    }
    else if (!(encp = b64(p, strlen(p))))
    {
        fprintf(stderr, "%s: b64 %s failed\n", __func__, p);
        goto end;
    }
    else if (!(header = b64(jwt_header, strlen(jwt_header))))
    {
        fprintf(stderr, "%s: b64 %s failed\n", __func__, jwt_header);
    }
    else if (dynstr_append(&d, "%s.%s", header, encp))
    {
        fprintf(stderr, "%s: dynstr_append header+payload failed", __func__);
        goto end;
    }
    else if (!(hmac = get_hmac(d.str, d.len, key, n)))
    {
        fprintf(stderr, "%s: get_hmac failed\n", __func__);
        goto end;
    }
    else if (dynstr_append(&d, ".%s", hmac))
    {
        fprintf(stderr, "%s: dynstr_append hmac failed\n", __func__);
        goto end;
    }

    ret = d.str;

end:
    free(header);
    free(hmac);
    free(encp);
    cJSON_free(p);

    if (!ret)
        dynstr_free(&d);

    return ret;
}

static size_t cnt(const char *s, const char c)
{
    size_t ret = 0;

    while (*s)
        if (*s++ == c)
            ret++;

    return ret;
}

int jwt_decode(const char *const jwt, const void *const key, const size_t n,
    cJSON **const payload)
{
    int ret = 1;
    cJSON *j = NULL;
    void *p = NULL;
    const char *first = strchr(jwt, '.'), *const last = strrchr(jwt, '.');
    const unsigned char *const in = (const unsigned char *)jwt;
    unsigned char hmac[crypto_auth_hmacsha256_KEYBYTES];
    size_t sz;

    if (n != crypto_auth_hmacsha256_KEYBYTES)
    {
        fprintf(stderr, "%s: invalid key size (%zu), expected %u\n", __func__,
            n, crypto_auth_hmacsha256_KEYBYTES);
        ret = -1;
        goto end;
    }
    else if (cnt(jwt, '.') != 2 || ++first == last)
        goto end;
    else if (sodium_base642bin(hmac, sizeof hmac, last + 1, strlen(last + 1),
        NULL, NULL, NULL, VARIANT))
    {
        fprintf(stderr, "%s: sodium_base642bin failed\n", __func__);
        goto end;
    }
    else if (crypto_auth_hmacsha256_verify(hmac, in, last - jwt, key))
    {
        fprintf(stderr, "%s: crypto_auth_hmacsha256_verify failed\n", __func__);
        goto end;
    }
    else if ((ret = b64d(first, last - first, &p, &sz)))
    {
        if (ret < 0)
            fprintf(stderr, "%s: b64d failed\n", __func__);

        goto end;
    }
    else if (!(j = cJSON_ParseWithLength(p, sz)))
    {
        fprintf(stderr, "%s: cJSON_ParseWithLength failed\n", __func__);
        goto end;
    }
    else if (!cJSON_IsObject(j))
    {
        fprintf(stderr, "%s: cJSON_IsObject failed\n", __func__);
        goto end;
    }

    *payload = j;
    ret = 0;

end:

    if (ret)
        cJSON_Delete(j);

    free(p);
    return ret;
}