/* 
 * Generate intermediate test vectors useful to test implementations.
 */

// Number of sample vectors to be generated
#include <stdint.h>
#include <stdio.h>
#include <string.h>

#include "../fips202.h"
#include "../params.h"
#include "../randombytes.h"
#include "../utils.h"
#include "../wots.h"
#include "../xmss_commons.h"
#include "../xmss_core.h"

#include <sys/random.h>
#include <sys/stat.h>

#include <json.h>

#define MESSAGE_LEN 128u

struct param_t {
    uint8_t oid;
    const char *name;
    const char *hash;
    unsigned height;
    unsigned n_samples;
};

void sprint_hex(char *sbuf, unsigned char *buf, int len) {
    for (int i = 0; i < len; i++) { 
        sprintf(&sbuf[2*i], "%02X", buf[i]); 
    }
    sbuf[2*len] = '\0';
}

void vectors_keygen(uint32_t oid, json_object *jreq, json_object *jres, uint32_t n_samples) {
    xmss_params params;
    struct json_object *tc_req, *tcs_req, *tc_res, *tcs_res;
    
    xmss_parse_oid(&params, oid);

    unsigned char seed[params.n * 3];
    unsigned char pk[params.pk_bytes + XMSS_OID_LEN];
    unsigned char sk[params.sk_bytes + XMSS_OID_LEN];
    char *sbuf;
    unsigned i;

    tcs_req = json_object_new_array();
    tcs_res = json_object_new_array();

    for (i = 0; i < XMSS_OID_LEN; i++) {
        pk[XMSS_OID_LEN - i - 1] = (oid >> (8 * i)) & 0xFF;
        sk[XMSS_OID_LEN - i - 1] = (oid >> (8 * i)) & 0xFF;
    }

    for (size_t i=0; i<n_samples; i++) {
        tc_req = json_object_new_object();
        tc_res = json_object_new_object();

        randombytes(seed, 3*params.n);
        xmssmt_core_seed_keypair(&params, &pk[XMSS_OID_LEN], &sk[XMSS_OID_LEN], seed);

        json_object_object_add(tc_req, "tcId", json_object_new_int(i+1));
        sbuf = malloc(2*params.n + 1);
        sprint_hex(sbuf, seed, params.n);
        json_object_object_add(tc_req, "S_XMSS", json_object_new_string(sbuf));
        sprint_hex(sbuf, &seed[params.n], params.n);
        json_object_object_add(tc_req, "SK_PRF", json_object_new_string(sbuf));
        sprint_hex(sbuf, &seed[2*params.n], params.n);
        json_object_object_add(tc_req, "I", json_object_new_string(sbuf));
        free(sbuf);

        json_object_object_add(tc_res, "tcId", json_object_new_int(i+1));
        sbuf = malloc(2*(params.pk_bytes + XMSS_OID_LEN) + 1);
        sprint_hex(sbuf, pk, params.pk_bytes + XMSS_OID_LEN);
        json_object_object_add(tc_res, "publicKey", json_object_new_string(sbuf));
        free(sbuf);

        sbuf = malloc(2*(params.sk_bytes + XMSS_OID_LEN) + 1);
        sprint_hex(sbuf, sk, params.sk_bytes + XMSS_OID_LEN);
        json_object_object_add(tc_res, "secretKey", json_object_new_string(sbuf));
        free(sbuf);

        json_object_array_add(tcs_req,tc_req);
        json_object_array_add(tcs_res,tc_res);
    }
    json_object_object_add(jreq, "tests", tcs_req);
    json_object_object_add(jres, "tests", tcs_res);
}

void json_header_keygen(json_object *jobj) {
    json_object_object_add(jobj, "vsId", json_object_new_int(0));
    json_object_object_add(jobj, "algorithm", json_object_new_string("XMSS"));
    json_object_object_add(jobj, "mode", json_object_new_string("keyGen"));
    json_object_object_add(jobj, "revision", json_object_new_string("1.0"));
    json_object_object_add(jobj, "isSample", json_object_new_boolean(1));
}

void keygen_KAT(const struct param_t *h) {
    struct json_object *jres, *jreq, *tgs_req, *tgs_res, *tg_req, *tg_res;
    char buf[256];

    jreq = json_object_new_object();
    jres = json_object_new_object();

    json_header_keygen(jreq);
    json_header_keygen(jres);

    tgs_req = json_object_new_array();
    tgs_res = json_object_new_array();

    tg_req = json_object_new_object();
    tg_res = json_object_new_object();

    // Request file
    json_object_object_add(tg_req, "tgId", json_object_new_int(1));
    json_object_object_add(tg_req, "testType", json_object_new_string("AFT"));
    json_object_object_add(tg_req, "OID", json_object_new_int(h->oid));
    json_object_object_add(tg_req, "param", json_object_new_string(h->name));

    // Response file
    json_object_object_add(tg_res, "tgId", json_object_new_int(1));
    json_object_object_add(tg_req, "OID", json_object_new_int(h->oid));

    vectors_keygen(h->oid, tg_req, tg_res, h->n_samples); 

    json_object_array_add(tgs_req, tg_req);
    json_object_array_add(tgs_res, tg_res);

    json_object_object_add(jreq, "testGroups", tgs_req);
    json_object_object_add(jres, "testGroups", tgs_res);
    
    sprintf(buf, "XMSS-%s-%s-H%u", "keyGen", h->hash, h->height);
    mkdir(buf, S_IRWXU | S_IRWXG | S_IROTH | S_IXOTH);
    sprintf(buf, "XMSS-%s-%s-H%u/%s", "keyGen", h->hash, h->height, "prompt.json");
    json_object_to_file_ext(buf, jreq, JSON_C_TO_STRING_SPACED | JSON_C_TO_STRING_PRETTY);
    sprintf(buf, "XMSS-%s-%s-H%u/%s", "keyGen", h->hash, h->height, "expectedResults.json");
    json_object_to_file_ext(buf, jres, JSON_C_TO_STRING_SPACED | JSON_C_TO_STRING_PRETTY);

    json_object_put(jres);
    json_object_put(jreq);
}

void json_header_siggen(json_object *jobj) {
    json_object_object_add(jobj, "vsId", json_object_new_int(0));
    json_object_object_add(jobj, "algorithm", json_object_new_string("XMSS"));
    json_object_object_add(jobj, "mode", json_object_new_string("sigGen"));
    json_object_object_add(jobj, "revision", json_object_new_string("1.0"));
    json_object_object_add(jobj, "isSample", json_object_new_boolean(1));
}

void vectors_siggen(uint32_t oid, unsigned char *sk, json_object *jreq, json_object *jres, uint32_t n_samples) {
    xmss_params params;
    struct json_object *tc_req, *tcs_req, *tc_res, *tcs_res;
    
    xmss_parse_oid(&params, oid);
    unsigned char sm[params.sig_bytes + MESSAGE_LEN];
    unsigned long long smlen = 0;
    unsigned char msg[MESSAGE_LEN];
    unsigned q;
    unsigned height = 1u << (params.full_height);
    char *sbuf;

    tcs_req = json_object_new_array();
    tcs_res = json_object_new_array();

    for (size_t i=0; i<n_samples; i++) {
        tc_req = json_object_new_object();
        tc_res = json_object_new_object();

        randombytes(msg, MESSAGE_LEN);
        randombytes((void*)&q, sizeof(q));
        q = q % height;
        smlen = 0;

        ull_to_bytes(sk, params.index_bytes, q);
        xmss_core_sign(&params, sk, sm, &smlen, msg, MESSAGE_LEN);

        json_object_object_add(tc_res, "tcId", json_object_new_int(i+1));
        json_object_object_add(tc_req, "tcId", json_object_new_int(i+1));

        sbuf = malloc(2*MESSAGE_LEN + 1);
        sprint_hex(sbuf, msg, MESSAGE_LEN);
        json_object_object_add(tc_req, "message", json_object_new_string(sbuf));
        free(sbuf);

        sbuf = malloc(2*smlen + 1);
        sprint_hex(sbuf, sm, smlen);
        json_object_object_add(tc_res, "signature", json_object_new_string(sbuf));
        free(sbuf);

        json_object_object_add(tc_req, "q", json_object_new_int(q));
        json_object_array_add(tcs_req, tc_req);
        json_object_array_add(tcs_res, tc_res);
    }
    json_object_object_add(jreq, "tests", tcs_req);
    json_object_object_add(jres, "tests", tcs_res);
}

void siggen_KAT(const struct param_t *h) {
    xmss_params params;
    struct json_object *jres, *jreq, *tgs_req, *tgs_res, *tg_req, *tg_res;
    xmss_parse_oid(&params, h->oid);

    char buf[256], *sbuf;
    unsigned char seed[params.n * 3];
    unsigned char pk[params.pk_bytes + XMSS_OID_LEN];
    unsigned char sk[params.sk_bytes + XMSS_OID_LEN];
    size_t i;

    jreq = json_object_new_object();
    jres = json_object_new_object();

    json_header_siggen(jreq);
    json_header_siggen(jres);

    tgs_req = json_object_new_array();
    tgs_res = json_object_new_array();

    tg_req = json_object_new_object();
    tg_res = json_object_new_object();

    // Store key OIDs
    for (i = 0; i < XMSS_OID_LEN; i++) {
        pk[XMSS_OID_LEN - i - 1] = (h->oid >> (8 * i)) & 0xFF;
        sk[XMSS_OID_LEN - i - 1] = (h->oid >> (8 * i)) & 0xFF;
    }

    // Request file
    json_object_object_add(tg_req, "tgId", json_object_new_int(1));
    json_object_object_add(tg_req, "testType", json_object_new_string("AFT"));
    json_object_object_add(tg_req, "OID", json_object_new_int(h->oid));
    json_object_object_add(tg_req, "param", json_object_new_string(h->name));

    // Response file
    json_object_object_add(tg_res, "tgId", json_object_new_int(1));
    json_object_object_add(tg_res, "OID", json_object_new_int(h->oid));

    randombytes(seed, 3*params.n);
    xmssmt_core_seed_keypair(&params, pk + XMSS_OID_LEN, sk + XMSS_OID_LEN, seed);

    sbuf = malloc(2*params.n + 1);
    sprint_hex(sbuf, seed, params.n);
    json_object_object_add(tg_req, "S_XMSS", json_object_new_string(sbuf));
    sprint_hex(sbuf, &seed[params.n], params.n);
    json_object_object_add(tg_req, "SK_PRF", json_object_new_string(sbuf));
    sprint_hex(sbuf, &seed[2*params.n], params.n);
    json_object_object_add(tg_req, "I", json_object_new_string(sbuf));
    free(sbuf);

    sbuf = malloc(2*params.n + 1);
    // Store 'root' from the public key (same as in the secret key)
    sprint_hex(sbuf, pk + XMSS_OID_LEN, params.n);
    json_object_object_add(tg_req, "PK_root", json_object_new_string(sbuf));
    free(sbuf);

    sbuf = malloc(2*(params.pk_bytes + XMSS_OID_LEN) + 1);
    sprint_hex(sbuf, pk, params.pk_bytes  + XMSS_OID_LEN);
    json_object_object_add(tg_res, "publicKey", json_object_new_string(sbuf));
    free(sbuf);

    vectors_siggen(h->oid, sk + XMSS_OID_LEN, tg_req, tg_res, h->n_samples);

    json_object_array_add(tgs_req, tg_req);
    json_object_array_add(tgs_res, tg_res);

    json_object_object_add(jreq, "testGroups", tgs_req);
    json_object_object_add(jres, "testGroups", tgs_res);
    
    sprintf(buf, "XMSS-%s-%s-H%u", "sigGen", h->hash, h->height);
    mkdir(buf, S_IRWXU | S_IRWXG | S_IROTH | S_IXOTH);
    sprintf(buf, "XMSS-%s-%s-H%u/%s", "sigGen", h->hash, h->height, "prompt.json");
    json_object_to_file_ext(buf, jreq, JSON_C_TO_STRING_SPACED | JSON_C_TO_STRING_PRETTY);
    sprintf(buf, "XMSS-%s-%s-H%u/%s", "sigGen", h->hash, h->height, "expectedResults.json");
    json_object_to_file_ext(buf, jres, JSON_C_TO_STRING_SPACED | JSON_C_TO_STRING_PRETTY);

    json_object_put(jreq);
    json_object_put(jres);
}

int main() {
    const struct param_t OIDs[] = {
        {0x01, "XMSS-SHA2_10_256", "SHA256-N32", 10, 10},   // H10
        {0x0D, "XMSS-SHA2_10_192", "SHA256-N24", 10, 10},   // H10
        {0x10, "XMSS-SHAKE256_10_256", "SHAKE256-N32", 10, 10}, // H10
        {0x13, "XMSS-SHAKE256_10_192", "SHAKE256-N24", 10, 10}, // H10

        {0x02, "XMSS-SHA2_16_256", "SHA256-N32", 16, 5},    // H16
        {0x0E, "XMSS-SHA2_16_192", "SHA256-N24", 16, 5},    // H16
        {0x11, "XMSS-SHAKE256_16_256", "SHAKE256-N32", 16, 5},  // H16
        {0x14, "XMSS-SHAKE256_16_192", "SHAKE256-N24", 16, 5},  // H16

        {0x03, "XMSS-SHA2_20_256", "SHA256-N32", 20, 3},    // H20
        {0x0F, "XMSS-SHA2_20_192", "SHA256-N24", 20, 3},    // H20
        {0x12, "XMSS-SHAKE256_20_256", "SHAKE256-N32", 20, 3},  // H20
        {0x15, "XMSS-SHAKE256_20_192", "SHAKE256-N24", 20, 3}   // H20
    };
    
    for (size_t i=0; i<sizeof(OIDs)/sizeof(OIDs[0]); i++) {
        keygen_KAT(&OIDs[i]);
        siggen_KAT(&OIDs[i]);
    }
}