#include "config.h"

#include "p11-tests.h"

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

void
test_rsa_decrypt(CK_SESSION_HANDLE session, CK_OBJECT_HANDLE key,
                 CK_MECHANISM_TYPE mech_type, RSA* rsa)
{
	CK_BYTE encrypted[p11t_test_data_size];
	CK_BYTE decrypted[p11t_test_data_size];
	CK_MECHANISM mech;
	const CK_BYTE* data;
	CK_ULONG n_data, n_decrypted;
	int size, n_encrypted;
	CK_RV rv;

	/** C_Decrypt */

	data = p11t_test_data;
	n_data = p11t_test_data_size;

	size = RSA_size(rsa);
	assert(size);
	assert(size < sizeof(encrypted));
	assert(size < sizeof(decrypted));
	assert(size < n_data);

	switch(mech_type)
	{
	case CKM_RSA_PKCS:
		/** - CKM_RSA_PKCS */
		n_data = size - 11;
		n_encrypted = RSA_public_encrypt(n_data, data, encrypted, rsa, RSA_PKCS1_PADDING);
		assert(n_encrypted == size);
		break;
	case CKM_RSA_X_509:
		/** - CKM_RSA_X_509 */
		n_data = size;
		n_encrypted = RSA_public_encrypt(n_data, data, encrypted, rsa, RSA_NO_PADDING);
		assert(n_encrypted == size);
		break;
	default:
		return;
	};

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

	rv = (p11t_module_funcs->C_DecryptInit)(session, &mech, key);
	p11t_check_returns("C_DecryptInit: rsa", rv, CKR_OK);

	n_decrypted = sizeof(decrypted);
	rv = (p11t_module_funcs->C_Decrypt)(session, encrypted, n_encrypted, decrypted, &n_decrypted);
	p11t_check_returns("C_Decrypt: rsa", rv, CKR_OK);

	if(n_decrypted != n_data)
		p11t_check_fail("C_Decrypt: rsa decrypt failed, wrong length");
	if(memcmp(data, decrypted, n_data) != 0)
		p11t_check_fail("C_Decrypt: rsa decrypt failed, bad data");
}

void
test_rsa_encrypt(CK_SESSION_HANDLE session, CK_OBJECT_HANDLE key,
                 CK_MECHANISM_TYPE mech_type, RSA* rsa)
{
	CK_BYTE encrypted[p11t_test_data_size];
	CK_BYTE check[p11t_test_data_size];
	CK_OBJECT_HANDLE privkey;
	CK_MECHANISM mech;
	const CK_BYTE* data;
	CK_ULONG n_data, n_check, n_encrypted;
	int size;
	CK_RV rv;

	data = p11t_test_data;
	n_data = p11t_test_data_size;

	/** C_Encrypt */

	size = RSA_size(rsa);
	assert(size);
	assert(size < sizeof(encrypted));
	assert(size < sizeof(check));
	assert(size < n_data);

	switch(mech_type)
	{
	case CKM_RSA_PKCS:
		/** - CKM_RSA_PKCS */
		n_data = size - 11;
		break;
	case CKM_RSA_X_509:
		/** - CKM_RSA_X_509 */
		n_data = size;
		break;
	default:
		return;
	};

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

	/* Now ask PKCS#11 to decrypt it */
	rv = (p11t_module_funcs->C_EncryptInit)(session, &mech, key);
	p11t_check_returns("C_EncryptInit: rsa", rv, CKR_OK);

	n_encrypted = sizeof(encrypted);
	rv = (p11t_module_funcs->C_Encrypt)(session, (CK_BYTE*)data, n_data, encrypted, &n_encrypted);
	p11t_check_returns("C_Encrypt: rsa", rv, CKR_OK);

	if(size != n_encrypted)
		p11t_check_fail("C_Encrypt: rsa encrypt failed, wrong length");

	/* We need to find a private key in order to validate */
	privkey = p11t_key_get_private(session, key);
	if(privkey == CK_INVALID)
		return;

	rv = (p11t_module_funcs->C_DecryptInit)(session, &mech, privkey);
	p11t_check_returns("C_DecryptInit: rsa validate", rv, CKR_OK);

	n_check = sizeof(check);
	rv = (p11t_module_funcs->C_Decrypt)(session, encrypted, n_encrypted, check, &n_check);
	p11t_check_returns("C_Decrypt: rsa validate", rv, CKR_OK);

	if(n_check != n_data)
		p11t_check_fail("C_Decrypt: rsa validate failed, wrong length");
	if(memcmp(data, check, n_data) != 0)
		p11t_check_fail("C_Decrypt: rsa validate failed, bad data");
}

void
test_rsa_private_key(CK_SESSION_HANDLE session, CK_OBJECT_HANDLE key,
                     CK_MECHANISM_TYPE mech_type)
{
	CK_ATTRIBUTE attrs[2];
	CK_BBOOL can_decrypt = 0;
	CK_BBOOL can_sign = 0;
	RSA *rsa;

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

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

	rsa = p11t_key_export_public_rsa(session, key);
	if(!rsa)
		return;

	if(can_decrypt)
		test_rsa_decrypt(session, key, mech_type, rsa);

	RSA_free(rsa);
}

void
test_rsa_public_key(CK_SESSION_HANDLE session, CK_OBJECT_HANDLE key,
                    CK_MECHANISM_TYPE mech_type)
{
	CK_ATTRIBUTE attrs[2];
	CK_BBOOL can_encrypt = 0;
	CK_BBOOL can_verify = 0;
	RSA *rsa;

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

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

	rsa = p11t_key_export_public_rsa(session, key);
	if(!rsa)
		return;

	if(can_encrypt)
		test_rsa_encrypt(session, key, mech_type, rsa);

	RSA_free(rsa);
}

void
test_rsa_pkcs(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_RSA;
	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;

	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_rsa_private_key(session, keys[i], mech);
			else if(klass == CKO_PUBLIC_KEY)
				test_rsa_public_key(session, keys[i], mech);
		}
	}

	free(keys);
}

void
p11t_rsa_tests(void)
{
	p11t_slot_for_each_mech(CKM_RSA_PKCS, test_rsa_pkcs);
}