#include "p11-tests-lib.h"

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

#include <openssl/dh.h>

static int
compare_secret_values (void *one, size_t n_one, void *two, size_t n_two)
{
	if (n_one > n_two)
		return memcmp (one + (n_one - n_two), two, n_two);
	else
		return memcmp (one, two + (n_two - n_one), n_one);
}

static DH*
generate_peer_dh (CK_SESSION_HANDLE session, CK_OBJECT_HANDLE key)
{
	CK_ATTRIBUTE attrs[2];
	CK_BYTE prime[16384];
	CK_BYTE base[1024];
	DH *dh;

	attrs[0].type = CKA_PRIME;
	attrs[0].ulValueLen = sizeof (prime);
	attrs[0].pValue = prime;

	attrs[1].type = CKA_BASE;
	attrs[1].ulValueLen = sizeof (base);
	attrs[1].pValue = base;

	if (!p11t_object_get (session, key, attrs, 2))
		return NULL;

	if(attrs[0].ulValueLen == CK_INVALID ||
	   attrs[1].ulValueLen == CK_INVALID)
		return NULL;

	dh = DH_new ();
	dh->p = BN_bin2bn (prime, attrs[0].ulValueLen, NULL);
	dh->g = BN_bin2bn (base, attrs[1].ulValueLen, NULL);
	assert(dh && dh->p && dh->g);

	if (!DH_generate_key (dh)) {
		DH_free (dh);
		return NULL;
	}

	return dh;
}

int
test_dh_private_derive (CK_SESSION_HANDLE session, CK_OBJECT_HANDLE pub, CK_OBJECT_HANDLE prv)
{
	CK_ATTRIBUTE attrs[3];
	CK_KEY_TYPE key_type = CKK_AES;
	CK_OBJECT_CLASS klass = CKO_SECRET_KEY;
	CK_BYTE buffer[16384];
	CK_BYTE value[16384];
	CK_OBJECT_HANDLE derived;
	CK_MECHANISM mech;
	BIGNUM *our_pub;
	CK_RV rv;
	DH *dh;

	P11T_SECTION("C_DeriveKey");

	/* Setup peer's DH */
	dh = generate_peer_dh (session, pub);
	if (dh == NULL)
		p11t_check_warn ("Couldn't retrieve DH parameters from DH pub key");

	/* Dig out our public key */
	attrs[0].type = CKA_VALUE;
	attrs[0].pValue = value;
	attrs[0].ulValueLen = sizeof (value);
	if (p11t_object_get (session, pub, attrs, 1)) {
		our_pub = BN_bin2bn (value, attrs[0].ulValueLen, NULL);
	} else {
		p11t_check_warn ("Couldn't retrieve DH public value from DH pub key");
		our_pub = NULL;
	}

	mech.mechanism = CKM_DH_PKCS_DERIVE;
	mech.pParameter = buffer;
	mech.ulParameterLen = BN_num_bytes (dh->pub_key);
	BN_bn2bin (dh->pub_key, buffer);

	/* TODO: How can we know this key type is supported? */
	attrs[0].type = CKA_KEY_TYPE;
	attrs[0].pValue = &key_type;
	attrs[0].ulValueLen = sizeof (key_type);

	attrs[1].type = CKA_CLASS;
	attrs[1].pValue = &klass;
	attrs[1].ulValueLen = sizeof (klass);

	rv = (p11t_module_funcs->C_DeriveKey) (session, &mech, prv, attrs, 2, &derived);
	P11T_CHECK_RV ("CKM_DH_PKCS_DERIVE", rv, CKR_OK);

	/* Compare to our generated key */
	attrs[0].type = CKA_VALUE;
	attrs[0].pValue = value;
	attrs[0].ulValueLen = sizeof (value);
	if (p11t_object_get (session, derived, attrs, 1) && attrs[0].ulValueLen != (CK_ULONG)-1 &&
	    dh != NULL && our_pub != NULL) {
		if (!DH_compute_key (buffer, our_pub, dh))
			assert (0);
		if (compare_secret_values (value, attrs[0].ulValueLen, buffer, DH_size (dh)) != 0)
			P11T_CHECK_FAIL_MSG ("CKM_DH_PKCS_DERIVE", "Derived key has wrong secret value");
	}

	BN_free (our_pub);
	DH_free (dh);

	return p11t_key_test (session, derived, CKO_SECRET_KEY);
}

int
p11t_dh_test_key_pair (CK_SESSION_HANDLE session, CK_OBJECT_HANDLE pub, CK_OBJECT_HANDLE prv)
{
	CK_ULONG mechanisms[128];
	CK_ATTRIBUTE attr = { CKA_ALLOWED_MECHANISMS, mechanisms, sizeof (mechanisms) };
	CK_ULONG i;

	/* Let's see if it support derivation */
	if (!p11t_object_get (session, prv, &attr, 1))
		return STOP;

	for (i = 0; i < attr.ulValueLen / sizeof (CK_ULONG); ++i) {
		if (mechanisms[i] == CKM_DH_PKCS_DERIVE)
			test_dh_private_derive (session, pub, prv);
	}

	return CONTINUE;
}

static void
test_dh_key_pair_gen (CK_SLOT_ID slot, CK_MECHANISM_TYPE mech, CK_MECHANISM_INFO_PTR info)
{
	CK_SESSION_HANDLE session;

	session = p11t_session_open(slot, 0);
	if(session == CK_INVALID || !p11t_session_login (session))
		return;

	p11t_dh_test_generate_pair (session);
}

void
p11t_dh_tests (void)
{
	p11t_slot_for_each_mech (CKM_DH_PKCS_KEY_PAIR_GEN, test_dh_key_pair_gen);
}

int
p11t_dh_test_generate_pair (CK_SESSION_HANDLE session)
{
	CK_ATTRIBUTE pub_attrs[3];
	CK_ATTRIBUTE priv_attrs[3];
	CK_MECHANISM mech;
	CK_OBJECT_HANDLE pub_key;
	CK_OBJECT_HANDLE priv_key;
	CK_ULONG bits;
	CK_RV rv;
	DH *dh;

	P11T_SECTION ("C_GenerateKeyPair");

	dh = DH_generate_parameters (256, 2, NULL, NULL);
	assert (dh);

	pub_attrs[0].type = CKA_PRIME;
	pub_attrs[0].ulValueLen = BN_num_bytes (dh->p);
	pub_attrs[0].pValue = alloca (pub_attrs[0].ulValueLen);
	BN_bn2bin (dh->p, (unsigned char*)pub_attrs[0].pValue);

	pub_attrs[1].type = CKA_BASE;
	pub_attrs[1].ulValueLen = BN_num_bytes (dh->g);
	pub_attrs[1].pValue = alloca (pub_attrs[1].ulValueLen);
	BN_bn2bin (dh->g, (unsigned char*)pub_attrs[1].pValue);

	priv_attrs[0].type = CKA_PRIME;
	priv_attrs[0].ulValueLen = BN_num_bytes (dh->p);
	priv_attrs[0].pValue = alloca (priv_attrs[0].ulValueLen);
	BN_bn2bin (dh->p, (unsigned char*)priv_attrs[0].pValue);

	priv_attrs[1].type = CKA_BASE;
	priv_attrs[1].ulValueLen = BN_num_bytes (dh->g);
	priv_attrs[1].pValue = alloca (priv_attrs[1].ulValueLen);
	BN_bn2bin (dh->g, (unsigned char*)priv_attrs[1].pValue);

	priv_attrs[2].type = CKA_VALUE_BITS;
	priv_attrs[2].ulValueLen = sizeof (bits);
	priv_attrs[2].pValue = &bits;

	mech.mechanism = CKM_DH_PKCS_KEY_PAIR_GEN;
	mech.pParameter = NULL;
	mech.ulParameterLen = 0;

	if (p11t_test_unexpected) {
		rv = (p11t_module_funcs->C_GenerateKeyPair) (session, &mech, pub_attrs, 0, priv_attrs, 0, &pub_key, &priv_key);
		P11T_CHECK_RV ("DH Key Pair without CKA_PRIME", rv, CKR_TEMPLATE_INCOMPLETE);

		bits = 1024;
		rv = (p11t_module_funcs->C_GenerateKeyPair) (session, &mech, pub_attrs, 2, priv_attrs, 3, &pub_key, &priv_key);
		P11T_CHECK_RV ("DH Key Pair with CKA_VALUE_BITS larger than CKA_PRIME", rv, CKR_TEMPLATE_INCONSISTENT);
	}

	bits = 256;
	rv = (p11t_module_funcs->C_GenerateKeyPair) (session, &mech, pub_attrs, 2, priv_attrs, 3, &pub_key, &priv_key);
	P11T_CHECK_RV ("DH Key Pair", rv, CKR_OK);

	p11t_key_test (session, pub_key, CKO_PUBLIC_KEY);
	p11t_key_test (session, priv_key, CKO_PRIVATE_KEY);

	/* Test corner cases */
	return p11t_object_generate_pair_bad (session, CKM_DH_PKCS_KEY_PAIR_GEN);
}