/*
 * GTK VNC Widget, Diffie Hellman tests
 *
 * Copyright (C) 2024  Red Hat, Inc.
 *
 * This library is free software; you can redistribute it and/or
 * modify it under the terms of the GNU Lesser General Public
 * License as published by the Free Software Foundation; either
 * version 2.0 of the License, or (at your option) any later version.
 *
 * This library 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
 * Lesser General Public License for more details.
 *
 * You should have received a copy of the GNU Lesser General Public
 * License along with this library; if not, write to the Free Software
 * Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301 USA
 */

#include <config.h>

#include "dhpriv.h"
#include "testutils.h"

struct DHTestData {
    const gchar *mod;
    const gchar *gen;
    const gchar *itr;
    const gchar *sec;
    const gchar *key;
    size_t keylen;
} test_data[] = {
    {
        .mod = "a655e5905de52b53",
        .gen = "a835819f929dee92",
        .itr = "398a399fee99fa99",
        .sec = "0000000000000000a312ae42bdbb877e",
        .key = "0000000000000000818e6d8128185a7f",
        .keylen = 16,
    }
};

/* Mock this to make test repeatable */
void mpz_urandomb(mpz_ptr mpi, gmp_randstate_ptr rng G_GNUC_UNUSED, mp_bitcnt_t bits)
{
    size_t len = bits / 8;
    guchar *data = g_new0(guchar, len);
    GBytes *bytes;
    mpz_t tmp;
    memset(data, 0x42, len);
    bytes = g_bytes_new_take(data, len);
    vnc_bytes_to_mpi(bytes, tmp);
    mpz_set(mpi, tmp);
    mpz_clear(tmp);
    g_bytes_unref(bytes);
}

static void test_dh(gconstpointer val)
{
    const struct DHTestData *data = val;
    GBytes *mod = bytes_from_hex(data->mod);
    GBytes *gen = bytes_from_hex(data->gen);
    GBytes *inter = bytes_from_hex(data->itr);
    GBytes *sec;
    GBytes *key;
    struct vnc_dh *dh;

    dh = vnc_dh_new(gen, mod);

    sec = vnc_dh_gen_secret(dh, data->keylen);
    key = vnc_dh_gen_key(dh, inter, data->keylen);

    if (getenv("VNC_DH_DEBUG")) {
        gchar *modstr = bytes_to_hex(mod);
        gchar *genstr = bytes_to_hex(gen);
        gchar *interstr = bytes_to_hex(inter);
        gchar *secstr = bytes_to_hex(sec);
        gchar *keystr = bytes_to_hex(key);
        g_printerr(">>>Mod: %s\n", modstr);
        g_printerr(">>>Gen: %s\n", genstr);
        g_printerr(">>>Itr: %s\n", interstr);
        g_printerr(">>>Sec: %s\n", secstr);
        g_printerr(">>>Key: %s\n", keystr);
        g_free(modstr);
        g_free(genstr);
        g_free(interstr);
        g_free(secstr);
        g_free(keystr);
    }

    char *secstr = bytes_to_hex(sec);
    char *keystr = bytes_to_hex(key);
    g_assert_cmpstr(secstr, ==, data->sec);
    g_assert_cmpstr(keystr, ==, data->key);
    g_free(secstr);
    g_free(keystr);

    g_bytes_unref(key);
    g_bytes_unref(sec);
    g_bytes_unref(inter);
    g_bytes_unref(mod);
    g_bytes_unref(gen);
    vnc_dh_free(dh);
}

int main(int argc, char **argv) {
    size_t i;

    g_test_init(&argc, &argv, NULL);

    for (i = 0 ; i < G_N_ELEMENTS(test_data); i++) {
        gchar *name = g_strdup_printf("/crypto/dh%zu", i);
        g_test_add_data_func(name, &test_data[i], test_dh);
        g_free(name);
    }

    return g_test_run();
}
