#include <stdio.h>
#include <stdbool.h>
#include <string.h>

#include "ecdh.h"
#include "log.h"

static bool
_test_ecdh_pass(KEYPAIR_TYPE keypair_type, bool corrupt_key)
{
    const KEYSTORAGE_SLOT alise_slot = KEY_SLOT_1;
    const KEYSTORAGE_SLOT bob_slot = KEY_SLOT_2;

    uint8_t alise_public_key[256];
    size_t alise_public_key_sz = 0;
    KEYPAIR_TYPE alise_keypair_type;

    uint8_t bob_public_key[256];
    size_t bob_public_key_sz = 0;
    KEYPAIR_TYPE bob_keypair_type;

    uint8_t shared_secret_1[128];
    size_t shared_secret_sz_1;

    uint8_t shared_secret_2[128];
    size_t shared_secret_sz_2;

    // Create key pair for Alise
    if (IOTELIC_OK != keypair_create(alise_slot, keypair_type)) {
        LOG("Can't create keypair for Alise (%s)", keypair_name(keypair_type));
        return false;
    }

    if (IOTELIC_OK != keypair_get_pubkey(alise_slot, alise_public_key, sizeof(alise_public_key), &alise_public_key_sz, &alise_keypair_type)) {
        LOG("Can't load public key from slot for Alise (%s)", keypair_name(keypair_type));
        return false;
    }

    if (corrupt_key) {
        ++alise_public_key[1];
    }

    // Create key pair for Bob
    if (IOTELIC_OK != keypair_create(bob_slot, keypair_type)) {
        LOG("Can't create keypair for Bob (%s)", keypair_name(keypair_type));
        return false;
    }

    if (IOTELIC_OK != keypair_get_pubkey(bob_slot, bob_public_key, sizeof(bob_public_key), &bob_public_key_sz, &bob_keypair_type)) {
        LOG("Can't load public key from slot for Bob (%s)", keypair_name(keypair_type));
        return false;
    }

    // ECDH for Alise - Bob
    if (IOTELIC_OK != ecdh(alise_slot,
                           bob_keypair_type, bob_public_key, bob_public_key_sz,
                           shared_secret_1, sizeof(shared_secret_1), &shared_secret_sz_1)) {
        LOG("Can't process ECDH for Alise (%s)", keypair_name(keypair_type));
        return false;
    }

    // ECDH for Bob - Alise
    if (IOTELIC_OK != ecdh(bob_slot,
                           alise_keypair_type, alise_public_key, alise_public_key_sz,
                           shared_secret_2, sizeof(shared_secret_2), &shared_secret_sz_2)) {
        if (!corrupt_key) {
            LOG("Can't process ECDH for Bob (%s)", keypair_name(keypair_type));
        }
        return false;
    }

    // Compare shared secrets
    if (0 == memcmp(shared_secret_1, shared_secret_2, shared_secret_sz_1)) {
        LOG("    DH ok (%s)", keypair_name(keypair_type));
        return true;
    }

    return false;
}

bool
test_ecdh(void)
{
    // Pass
    if (!_test_ecdh_pass(KEYPAIR_EC_SECP256R1, false)) return false;
    if (!_test_ecdh_pass(KEYPAIR_EC_SECP384R1, false)) return false;
    if (!_test_ecdh_pass(KEYPAIR_EC_SECP521R1, false)) return false;
    if (!_test_ecdh_pass(KEYPAIR_EC_ED25519, false)) return false;
    if (!_test_ecdh_pass(KEYPAIR_EC_CURVE25519, false)) return false;

    // Fail
    if (_test_ecdh_pass(KEYPAIR_EC_SECP256R1, true)) return false;
    if (_test_ecdh_pass(KEYPAIR_EC_CURVE25519, true)) return false;

    return true;
}