#include <bearssl.h> /* -lbearssl */
#include <assert.h>
#include <string.h>
#include <stdio.h>

int bearssl_dh_nistp256_keypair(unsigned char *pk, unsigned char *sk) {

    static br_hmac_drbg_context rng;
    unsigned char tmpsk[BR_EC_KBUF_PRIV_MAX_SIZE];
    unsigned char tmppk[BR_EC_KBUF_PUB_MAX_SIZE];
    br_ec_private_key brsk;
    br_ec_public_key brpk;
    br_prng_seeder seeder;
    const br_ec_impl *impl = br_ec_get_default();
    size_t ret;

    /* initialize rng */
    br_hmac_drbg_init(&rng, &br_sha512_vtable, 0, 0);
    seeder = br_prng_seeder_system(0);
    if (!seeder) return -1;
    if (!seeder(&rng.vtable)) return -1;

    /* generate random secret-key */
    ret = br_ec_keygen(&rng.vtable, impl, &brsk, tmpsk, BR_EC_secp256r1);
    if (ret != 32) return -1;

    /* generate public-key */ 
    ret = br_ec_compute_pub(impl, &brpk, tmppk, &brsk);
    if (ret != 65) return -1;

    memcpy(pk, tmppk + 1, 64);
    memcpy(sk, tmpsk, 32);
    return 0;
}

int bearssl_dh_nistp256(unsigned char *k, const unsigned char *pk, const unsigned char *sk) {

    unsigned char tmpsk[32];
    unsigned char tmppk[65];
    uint32_t ret;

    tmppk[0] = 4;
    memcpy(tmppk + 1, pk, 64);
    memcpy(tmpsk, sk, 32);

    ret = br_ec_get_default()->mul(tmppk, sizeof tmppk, tmpsk, sizeof tmpsk, BR_EC_secp256r1);
    if (ret != 1) return -1;

    assert(tmppk[0] == 4);
    memcpy(k, tmppk + 1, 64);
    return 0;
}

void h_print(const char *text, unsigned char *h, long long hsize) {
    long long i;
    printf("%s: ", text);
    for (i = 0; i < hsize; ++i) printf("%02x", 255 & (int) h[i]);
    printf("\n");
}


#define LOOPS 100

int main(void) {
    unsigned char k1[64], pk1[64], sk1[32];
    unsigned char k2[64], pk2[64], sk2[32];
    long long i;
    int ret = 1;

    for (i = 0; i < LOOPS; ++i) {

        /* alice - bearssl keypair */
        assert(!bearssl_dh_nistp256_keypair(pk1, sk1));

        /* bob - bearssl keypair */
        assert(!bearssl_dh_nistp256_keypair(pk2, sk2));

        /* alice - bearssl dh */
        assert(!bearssl_dh_nistp256(k1, pk2, sk1));

        /* bob  - bearssl dh */
        assert(!bearssl_dh_nistp256(k2, pk1, sk2));

        ret = !!memcmp(k1, k2, sizeof k1);
    }
    if (ret) {
        printf("failed:\n");
    }
    else {
        printf("OK:\n");
    }
    h_print("pk1", pk1, sizeof pk1);
    h_print("sk1", sk1, sizeof sk1);
    h_print("pk2", pk2, sizeof pk2);
    h_print("sk2", sk2, sizeof sk2);
    h_print(" k1", k1, sizeof k1);
    h_print(" k2", k2, sizeof k2);
    return 0;
}
