#include "p11-tests.h"

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

#include <openssl/evp.h>
#include <openssl/sha.h>
#include <openssl/md5.h>
#include <openssl/objects.h>
#include <openssl/x509.h>

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

	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 < (int)n_data);

	switch(mech_type)
	{
	case CKM_RSA_PKCS:
		P11T_CHECK_NOTE("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:
		P11T_CHECK_NOTE("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 CONTINUE;
	};

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

	P11T_SECTION("C_DecryptInit");

	rv = (p11t_module_funcs->C_DecryptInit)(session, &mech, key);
	P11T_CHECK_RV("Normal call", rv, CKR_OK);

	P11T_SECTION("C_Decrypt");

	n_decrypted = sizeof(decrypted);
	rv = (p11t_module_funcs->C_Decrypt)(session, encrypted, n_encrypted, decrypted, &n_decrypted);

	/* Requires authentication */
	if (rv == CKR_USER_NOT_LOGGED_IN) {
		rv = p11t_key_login_context_specific (session, key);
		P11T_CHECK_RV("Always authenticate", rv, CKR_OK);

		/* Try it again */
		rv = (p11t_module_funcs->C_Decrypt)(session, encrypted, n_encrypted, decrypted, &n_decrypted);
	}

	P11T_CHECK_RV("Normal call", rv, CKR_OK);

	if(n_decrypted != n_data)
		P11T_CHECK_FAIL("RSA decrypt failed, wrong length");
	if(memcmp(data, decrypted, n_data) != 0)
		P11T_CHECK_FAIL("RSA decrypt failed, mangled data");

	return CONTINUE;
}

static void
hash_for_rsa_ssl3_sign(const CK_BYTE* data, CK_ULONG n_data,
                       CK_BYTE_PTR output, CK_ULONG_PTR n_output)
{
	assert(*n_output >= MD5_DIGEST_LENGTH + SHA_DIGEST_LENGTH);
	MD5(data, n_data, output);
	SHA1(data, n_data, output + MD5_DIGEST_LENGTH);
	*n_output = MD5_DIGEST_LENGTH + SHA_DIGEST_LENGTH;
}

static void
hash_for_rsa_pkcs_sign(int algo, int wrap, const CK_BYTE* data,
                       CK_ULONG n_data, CK_BYTE_PTR output, CK_ULONG_PTR n_output)
{
	CK_BYTE hash[EVP_MAX_MD_SIZE];
	X509_SIG sig;
	X509_ALGOR algor;
	ASN1_TYPE parameter;
	ASN1_OCTET_STRING digest;
	EVP_MD_CTX ctx;
	const EVP_MD* evp;
	size_t n_hash;
	int val;

	if(algo == NID_md5_sha1)
	{
		hash_for_rsa_ssl3_sign(data, n_data, output, n_output);
		return;
	}

	evp = EVP_get_digestbynid(algo);
	assert(evp);

	val = EVP_DigestInit(&ctx, evp);
	assert(val >= 0);
	EVP_DigestUpdate(&ctx, data, n_data);
	EVP_DigestFinal(&ctx, hash, &n_hash);

	/* Not wrapping? */
	if(!wrap)
	{
		assert(*n_output >= n_hash);
		*n_output = (CK_ULONG)n_hash;
		memcpy(output, hash, n_hash);
		return;
	}

	sig.algor = &algor;
	sig.algor->algorithm = OBJ_nid2obj(algo);
	assert(sig.algor->algorithm);
	assert(sig.algor->algorithm->length);
	parameter.type = V_ASN1_NULL;
	parameter.value.ptr = NULL;
	sig.algor->parameter = &parameter;
	sig.digest = &digest;
	sig.digest->data = hash;
	sig.digest->length = (int)n_hash;

	val = i2d_X509_SIG(&sig, &output);
	assert(*n_output >= (CK_ULONG)val);
	*n_output = val;
}

static int
test_rsa_pkcs_sign_hash(CK_SESSION_HANDLE session, CK_OBJECT_HANDLE key,
                        RSA* rsa, int algo)
{
	CK_BYTE hash[P11T_BLOCK];
	CK_BYTE sig[P11T_BLOCK];
	CK_MECHANISM mech;
	CK_ULONG n_hash, n_sig;
	CK_RV rv;
	int res;

	/* Hash the data with appropriate wrappers */
	n_hash = sizeof(hash);
	hash_for_rsa_pkcs_sign(algo, 1, p11t_test_data, p11t_test_data_size, hash, &n_hash);

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

	P11T_SECTION("C_SignInit");

	rv = (p11t_module_funcs->C_SignInit)(session, &mech, key);
	P11T_CHECK_RV("Normal call", 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, key);
		P11T_CHECK_RV("Always authenticate", rv, CKR_OK);

		/* Try it again */
		rv = (p11t_module_funcs->C_Sign)(session, hash, n_hash, sig, &n_sig);
	}

	P11T_CHECK_RV("Normal call", rv, CKR_OK);

	/* Hash the data again this time without wrapping */
	n_hash = sizeof(hash);
	hash_for_rsa_pkcs_sign(algo, 0, p11t_test_data, p11t_test_data_size, hash, &n_hash);

	/* Verify the signature */
	res = RSA_verify(algo, hash, n_hash, sig, n_sig, rsa);
	if(res != 1)
		P11T_CHECK_FAIL("RSA PKCS#1.5 or SSLv3 signature did not verify");

	return CONTINUE;
}

static int
test_rsa_x509_sign(CK_SESSION_HANDLE session, CK_OBJECT_HANDLE key, RSA* rsa)
{
	const CK_BYTE* data;
	CK_BYTE sig[P11T_BLOCK];
	CK_BYTE check[P11T_BLOCK];
	CK_MECHANISM mech;
	CK_ULONG n_data, n_sig;
	CK_RV rv;
	int res, size;

	size = RSA_size(rsa);
	data = p11t_test_data;
	n_data = size / 2;
	assert(n_data <= p11t_test_data_size);

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

	P11T_SECTION("C_SignInit");

	rv = (p11t_module_funcs->C_SignInit)(session, &mech, key);
	P11T_CHECK_RV("RSA X509 Call", rv, CKR_OK);

	P11T_SECTION("C_Sign");

	n_sig = sizeof(sig);
	rv = (p11t_module_funcs->C_Sign)(session, (CK_BYTE*)data, n_data, sig, &n_sig);

	/* Requires authentication */
	if (rv == CKR_USER_NOT_LOGGED_IN) {
		rv = p11t_key_login_context_specific (session, key);
		P11T_CHECK_RV("Always authenticate", rv, CKR_OK);

		/* Try it again */
		rv = (p11t_module_funcs->C_Sign)(session, (CK_BYTE*)data, n_data, sig, &n_sig);
	}

	P11T_CHECK_RV("RSA X509 Call", rv, CKR_OK);
	P11T_CHECK_ULONG("C_Sign: rsa x509 result length", n_sig, size);

	res = RSA_public_decrypt(n_sig, sig, check, rsa, RSA_NO_PADDING);
	if(res < 0)
		P11T_CHECK_FAIL("RSA x509 signature was invalid");

	assert(res > (int)n_data);
	if(memcmp(check + (res - n_data), data, n_data) != 0)
		P11T_CHECK_FAIL("RSA x509 signature did not verify");

	return CONTINUE;
}

static int
test_rsa_sign(CK_SESSION_HANDLE session, CK_OBJECT_HANDLE key,
              CK_MECHANISM_TYPE mech_type, RSA* rsa)
{
	P11T_SECTION("C_Sign");

	switch(mech_type)
	{
	case CKM_RSA_PKCS:
		P11T_CHECK_NOTE("CKM_RSA_PKCS (SHA1)");
		test_rsa_pkcs_sign_hash(session, key, rsa, NID_sha1);
		P11T_CHECK_NOTE("CKM_RSA_PKCS (MD5)");
		test_rsa_pkcs_sign_hash(session, key, rsa, NID_md5);
		P11T_CHECK_NOTE("CKM_RSA_PKCS (SHA1/MD5/SSL3)");
		test_rsa_pkcs_sign_hash(session, key, rsa, NID_md5_sha1);
		break;
	case CKM_RSA_X_509:
		P11T_CHECK_NOTE("CKM_RSA_X_509");
		test_rsa_x509_sign(session, key, rsa);
		break;
	};

	return CONTINUE;
}

static int
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 CONTINUE;

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

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

	if(can_sign)
		test_rsa_sign(session, key, mech_type, rsa);

	RSA_free(rsa);

	return CONTINUE;
}

static int
test_rsa_encrypt(CK_SESSION_HANDLE session, CK_OBJECT_HANDLE key,
                 CK_MECHANISM_TYPE mech_type, RSA* rsa)
{
	CK_BYTE encrypted[P11T_BLOCK];
	CK_BYTE check[P11T_BLOCK];
	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;

	P11T_SECTION("C_Encrypt");

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

	switch(mech_type)
	{
	case CKM_RSA_PKCS:
		P11T_CHECK_NOTE("CKM_RSA_PKCS");
		n_data = size - 11;
		break;
	case CKM_RSA_X_509:
		P11T_CHECK_NOTE("CKM_RSA_X_509");
		n_data = size;
		break;
	default:
		return CONTINUE;
	};

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

	P11T_SECTION("C_EncryptInit");

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

	P11T_SECTION("C_Encrypt");

	n_encrypted = sizeof(encrypted);
	rv = (p11t_module_funcs->C_Encrypt)(session, (CK_BYTE*)data, n_data, encrypted, &n_encrypted);
	P11T_CHECK_RV("Normal call", rv, CKR_OK);

	if(size != n_encrypted)
		P11T_CHECK_FAIL("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 CONTINUE;

	P11T_SECTION("C_DecryptInit");

	rv = (p11t_module_funcs->C_DecryptInit)(session, &mech, privkey);
	P11T_CHECK_RV("Module encrypted data", rv, CKR_OK);

	P11T_SECTION("C_Decrypt");

	n_check = sizeof(check);
	rv = (p11t_module_funcs->C_Decrypt)(session, encrypted, n_encrypted, check, &n_check);

	/* 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);

		/* Try it again */
		rv = (p11t_module_funcs->C_Decrypt)(session, encrypted, n_encrypted, check, &n_check);
	}

	P11T_CHECK_RV("Module encrypted data", rv, CKR_OK);

	if(n_check != n_data)
		P11T_CHECK_FAIL("RSA validate failed, wrong length");
	if(memcmp(data, check, n_data) != 0)
		P11T_CHECK_FAIL("RSA validate failed, bad data");

	return CONTINUE;
}

static int
test_rsa_pkcs_verify_hash(CK_SESSION_HANDLE session, CK_OBJECT_HANDLE key,
                          CK_OBJECT_HANDLE pubkey, CK_ULONG size, int algo)
{
	CK_BYTE hash[P11T_BLOCK];
	CK_BYTE sig[P11T_BLOCK];
	CK_MECHANISM mech;
	CK_ULONG n_hash, n_sig;
	CK_RV rv;

	/* Hash the data with appropriate wrappers */
	n_hash = sizeof(hash);
	hash_for_rsa_pkcs_sign(algo, 1, p11t_test_data, p11t_test_data_size, hash, &n_hash);

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

	P11T_SECTION("C_SignInit");

	rv = (p11t_module_funcs->C_SignInit)(session, &mech, key);
	P11T_CHECK_RV("Normal call", 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, key);
		P11T_CHECK_RV("Always authenticate", rv, CKR_OK);
		rv = (p11t_module_funcs->C_Sign)(session, hash, n_hash, sig, &n_sig);
	}
	P11T_CHECK_RV("Normal call", rv, CKR_OK);

	P11T_SECTION("C_VerifyInit");

	rv = (p11t_module_funcs->C_VerifyInit)(session, &mech, pubkey);
	P11T_CHECK_RV("RSA PKCS", rv, CKR_OK);

	P11T_SECTION("C_Verify");

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

	return CONTINUE;
}

static int
test_rsa_x509_verify(CK_SESSION_HANDLE session, CK_OBJECT_HANDLE key,
                     CK_OBJECT_HANDLE pubkey, CK_ULONG size)
{
	const CK_BYTE* data;
	CK_BYTE sig[P11T_BLOCK];
	CK_MECHANISM mech;
	CK_ULONG n_data, n_sig;
	CK_RV rv;

	data = p11t_test_data;
	n_data = size / 2;
	assert(n_data <= p11t_test_data_size);

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

	P11T_SECTION("C_SignInit");

	rv = (p11t_module_funcs->C_SignInit)(session, &mech, key);
	P11T_CHECK_RV("RSA X509 Call", rv, CKR_OK);

	P11T_SECTION("C_Sign");

	n_sig = sizeof(sig);
	rv = (p11t_module_funcs->C_Sign)(session, (CK_BYTE*)data, n_data, sig, &n_sig);

	/* Requires authentication */
	if (rv == CKR_USER_NOT_LOGGED_IN) {
		rv = p11t_key_login_context_specific (session, key);
		P11T_CHECK_RV("Always authenticate", rv, CKR_OK);
		rv = (p11t_module_funcs->C_Sign)(session, (CK_BYTE*)data, n_data, sig, &n_sig);
	}

	P11T_CHECK_RV("RSA X509 Call", rv, CKR_OK);

	P11T_SECTION("C_VerifyInit");

	rv = (p11t_module_funcs->C_VerifyInit)(session, &mech, pubkey);
	P11T_CHECK_RV("RSA X509 Call", rv, CKR_OK);

	P11T_SECTION("C_Verify");
	rv = (p11t_module_funcs->C_Verify)(session, (CK_BYTE*)data, n_data, sig, n_sig);
	P11T_CHECK_RV("RSA X509 Call", rv, CKR_OK);

	return CONTINUE;
}

static int
test_rsa_verify (CK_SESSION_HANDLE session, CK_OBJECT_HANDLE key,
                 CK_MECHANISM_TYPE mech_type, RSA* rsa)
{
	CK_OBJECT_HANDLE privkey;
	CK_ULONG size;

	P11T_SECTION("C_Verify");

	/* Must find a private key for this public one */
	privkey = p11t_key_get_private(session, key);
	if(privkey == CK_INVALID)
		return CONTINUE;

	size = RSA_size(rsa);
	assert(size);

	switch(mech_type)
	{
	case CKM_RSA_PKCS:
		P11T_CHECK_NOTE("CKM_RSA_PKCS (SHA1)");
		test_rsa_pkcs_verify_hash(session, privkey, key, size, NID_sha1);
		P11T_CHECK_NOTE("CKM_RSA_PKCS (MD5)");
		test_rsa_pkcs_verify_hash(session, privkey, key, size, NID_md5);
		P11T_CHECK_NOTE("CKM_RSA_PKCS (SHA1/MD5/SSL3)");
		test_rsa_pkcs_verify_hash(session, privkey, key, size, NID_md5_sha1);
		break;
	case CKM_RSA_X_509:
		P11T_CHECK_NOTE("CKM_RSA_X_509");
		test_rsa_x509_verify(session, privkey, key, size);
		break;
	};

	return CONTINUE;
}

static int
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 CONTINUE;

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

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

	if(can_verify)
		test_rsa_verify(session, key, mech_type, rsa);

	RSA_free(rsa);

	return CONTINUE;
}

static void
test_rsa(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);
	p11t_slot_for_each_mech(CKM_RSA_X_509, test_rsa);
}