#include "p11-tests-lib.h"

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

#include <openssl/sha.h>

static int
test_dsa_sign_verify(CK_SESSION_HANDLE session, CK_OBJECT_HANDLE key,
                     CK_OBJECT_HANDLE privkey, DSA* dsa)
{
	CK_BYTE hash[20];
	CK_BYTE sig[P11T_BLOCK];
	CK_BYTE encsig[P11T_BLOCK];
	unsigned char *ptr;
	CK_MECHANISM mech;
	CK_ULONG n_hash, n_sig;
	DSA_SIG *dsa_sig;
	CK_RV rv;
	int res, n_encsig;

	/* Hash the data for DSA */
	SHA1(p11t_test_data, p11t_test_data_size, hash);
	n_hash = 20;

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

	P11T_SECTION("C_SignInit");

	rv = (p11t_module_funcs->C_SignInit)(session, &mech, privkey);
	P11T_CHECK_RV("CKM_DSA", rv, CKR_OK);

	P11T_SECTION("C_Sign");

	n_sig = sizeof(sig);
	rv = (p11t_module_funcs->C_Sign)(session, hash, n_hash, sig, &n_sig);

	/* Requires authentication */
	if (rv == CKR_USER_NOT_LOGGED_IN) {
		rv = p11t_key_login_context_specific (session, privkey);
		P11T_CHECK_RV("Always authenticate", rv, CKR_OK);
		rv = (p11t_module_funcs->C_Sign)(session, hash, n_hash, sig, &n_sig);
	}
	P11T_CHECK_RV("CKM_DSA", rv, CKR_OK);
	if(n_sig != 40)
		P11T_CHECK_FAIL("Length of DSA signature");


	/* Verify the DSA signature via openssl */
	dsa_sig = DSA_SIG_new();
	assert (dsa_sig);
	dsa_sig->r = BN_bin2bn(sig, 20, NULL);
	dsa_sig->s = BN_bin2bn(sig + 20, 20, NULL);
	ptr = encsig;
	n_encsig = i2d_DSA_SIG(dsa_sig, &ptr);
	DSA_SIG_free(dsa_sig);
	res = DSA_verify(0, hash, n_hash, encsig, n_encsig, dsa);
	if(res != 1)
		P11T_CHECK_FAIL_MSG("Check valid DSA signature", p11t_msg_openssl());


	P11T_SECTION("C_VerifyInit");

	rv = (p11t_module_funcs->C_VerifyInit)(session, &mech, key);
	P11T_CHECK_RV("CKM_DSA", rv, CKR_OK);

	P11T_SECTION("C_Verify");

	rv = (p11t_module_funcs->C_Verify)(session, hash, n_hash, sig, n_sig);
	P11T_CHECK_RV("CKM_DSA", rv, CKR_OK);

	return CONTINUE;
}

static int
test_dsa_public_key(CK_SESSION_HANDLE session, CK_OBJECT_HANDLE key,
                    CK_MECHANISM_TYPE mech_type)
{
	CK_OBJECT_HANDLE privkey;
	CK_ATTRIBUTE attrs[1];
	CK_BBOOL can_verify = 0;
	DSA *dsa;

	attrs[0].type = CKA_VERIFY;
	attrs[0].ulValueLen = sizeof(CK_BBOOL);
	attrs[0].pValue = &can_verify;

	if(!p11t_object_get (session, key, attrs, 1))
		return CONTINUE;

	dsa = p11t_key_export_public_dsa (session, key);
	if(!dsa)
		return CONTINUE;

	privkey = p11t_key_get_private (session, key);
	if (privkey == CK_INVALID)
		return CONTINUE;

	if(can_verify)
		test_dsa_sign_verify(session, key, privkey, dsa);

	DSA_free(dsa);

	return CONTINUE;
}

static int
test_dsa_private_key(CK_SESSION_HANDLE session, CK_OBJECT_HANDLE key,
                     CK_MECHANISM_TYPE mech_type)
{
	CK_OBJECT_HANDLE pubkey;
	CK_ATTRIBUTE attrs[2];
	CK_BBOOL can_sign = 0;
	DSA *dsa;

	attrs[0].type = CKA_SIGN;
	attrs[0].ulValueLen = sizeof(CK_BBOOL);
	attrs[0].pValue = &can_sign;

	if(!p11t_object_get(session, key, attrs, 1))
		return CONTINUE;

	pubkey = p11t_key_get_public(session, key);
	if (pubkey == CK_INVALID)
		return CONTINUE;

	dsa = p11t_key_export_public_dsa(session, pubkey);
	if(!dsa)
		return CONTINUE;

	if(can_sign)
		test_dsa_sign_verify(session, pubkey, key, dsa);

	DSA_free(dsa);

	return CONTINUE;
}

static void
test_dsa (CK_SLOT_ID slot, CK_MECHANISM_TYPE mech, CK_MECHANISM_INFO_PTR info)
{
	CK_SESSION_HANDLE session;
	CK_ATTRIBUTE attrs[2];
	CK_ATTRIBUTE attr;
	CK_OBJECT_CLASS klass;
	CK_KEY_TYPE key_type = CKK_DSA;
	CK_BBOOL token = CK_TRUE;
	CK_OBJECT_HANDLE_PTR keys;
	CK_ULONG i, n_keys;

	/* Find all keys that match this mechanism */
	attrs[0].type = CKA_TOKEN;
	attrs[0].ulValueLen = sizeof(token);
	attrs[0].pValue = &token;
	attrs[1].type = CKA_KEY_TYPE;
	attrs[1].ulValueLen = sizeof(key_type);
	attrs[1].pValue = &key_type;

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

	if(!p11t_session_login(session))
		return;

	/* Test signing corner cases */
	p11t_crypto_test_sign (session, mech);
	p11t_crypto_test_verify (session, mech);

	keys = p11t_object_find(session, attrs, 2, &n_keys);
	if(!keys)
		return;

	for(i = 0; i < n_keys; ++i)
	{
		attr.type = CKA_CLASS;
		attr.ulValueLen = sizeof(klass);
		attr.pValue = &klass;

		if(p11t_object_get(session, keys[i], &attr, 1))
		{
			if(klass == CKO_PRIVATE_KEY)
				test_dsa_private_key(session, keys[i], mech);
			else if(klass == CKO_PUBLIC_KEY)
				test_dsa_public_key(session, keys[i], mech);
		}
	}


	free(keys);
}

void
p11t_dsa_tests(void)
{
	p11t_slot_for_each_mech(CKM_DSA, test_dsa);
}

void
p11t_dsa_test_public_key(CK_SESSION_HANDLE session, CK_OBJECT_HANDLE handle)
{
	CK_MECHANISM_TYPE_PTR mechanisms;
	CK_ULONG n_mechanisms, i;

	mechanisms = p11t_key_get_mechanisms (session, handle, &n_mechanisms);
	if (!mechanisms)
		return;

	for (i = 0; i < n_mechanisms; ++i)
		test_dsa_public_key (session, handle, mechanisms[i]);

	free (mechanisms);
}

void
p11t_dsa_test_private_key(CK_SESSION_HANDLE session, CK_OBJECT_HANDLE handle)
{
	CK_MECHANISM_TYPE_PTR mechanisms;
	CK_ULONG n_mechanisms, i;

	mechanisms = p11t_key_get_mechanisms (session, handle, &n_mechanisms);
	if (!mechanisms)
		return;

	for (i = 0; i < n_mechanisms; ++i)
		test_dsa_private_key (session, handle, mechanisms[i]);

	free (mechanisms);
}