#include "p11-tests-lib.h"

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

#include <openssl/dh.h>

int
test_bad_wrap (CK_SESSION_HANDLE session, CK_MECHANISM_PTR mech,
               CK_OBJECT_HANDLE wrapper, CK_OBJECT_HANDLE wrapped)
{
	CK_MECHANISM other;
	CK_BYTE buffer[16384];
	CK_ULONG n_buffer;
	CK_RV rv;

	P11T_SECTION("C_WrapKey");

	if (!p11t_test_unexpected)
		return CONTINUE;

	rv = (p11t_module_funcs->C_WrapKey) ((CK_ULONG)-88, mech, wrapper, wrapped, buffer, &n_buffer);
	P11T_CHECK_RV ("Without valid session handle", rv, CKR_SESSION_HANDLE_INVALID);

	rv = (p11t_module_funcs->C_WrapKey) (session, NULL, wrapper, wrapped, buffer, &n_buffer);
	P11T_CHECK_RV ("With null mechanism", rv, CKR_ARGUMENTS_BAD);

	/* A mechanism not on the key */
	other.mechanism = CKM_DH_PKCS_DERIVE;
	other.pParameter = NULL;
	other.ulParameterLen = 0;

	rv = (p11t_module_funcs->C_WrapKey) (session, &other, wrapper, wrapped, buffer, &n_buffer);
	P11T_CHECK_RV ("With bad mechanism", rv, CKR_KEY_TYPE_INCONSISTENT);

	other.mechanism = mech->mechanism;
	rv = (p11t_module_funcs->C_WrapKey) (session, &other, wrapper, wrapped, buffer, &n_buffer);
	P11T_CHECK_RV ("With bad mechanism parameter", rv, CKR_MECHANISM_PARAM_INVALID);

	rv = (p11t_module_funcs->C_WrapKey) (session, mech, (CK_ULONG)-82, wrapped, buffer, &n_buffer);
	P11T_CHECK_RV ("With bad wrapping key handle", rv, CKR_WRAPPING_KEY_HANDLE_INVALID);

	rv = (p11t_module_funcs->C_WrapKey) (session, mech, wrapper, (CK_ULONG)-82, buffer, &n_buffer);
	P11T_CHECK_RV ("With bad wrapped key handle", rv, CKR_KEY_HANDLE_INVALID);

	rv = (p11t_module_funcs->C_WrapKey) (session, mech, wrapper, wrapped, buffer, NULL);
	P11T_CHECK_RV ("With null wrapped byte buffer count", rv, CKR_ARGUMENTS_BAD);

	return CONTINUE;
}

int
test_aes_wrap (CK_SESSION_HANDLE session, CK_OBJECT_HANDLE wrapper,
               CK_OBJECT_HANDLE wrapped, CK_BYTE_PTR buffer, CK_ULONG_PTR n_buffer)
{
	CK_BYTE iv[16] = { '0', '1', '2', '3', '4', '5', '6', '7', '8', '9', 'A', 'B', 'C', 'D', 'E', 'F' };
	CK_MECHANISM mech = { CKM_AES_CBC_PAD, iv, 16 };
	CK_BYTE check[16384];
	CK_BYTE value[512];
	CK_ATTRIBUTE attr;
	EVP_CIPHER_CTX ctx;
	int outlen, tmplen;
	CK_ULONG length;
	CK_RV rv;

	P11T_SECTION("C_WrapKey");

	test_bad_wrap (session, &mech, wrapper, wrapped);

	/* Just get the length */
	rv = (p11t_module_funcs->C_WrapKey) (session, &mech, wrapper, wrapped, NULL, &length);
	P11T_CHECK_RV ("CKM_AES_CBC_PAD", rv, CKR_OK);

	rv = (p11t_module_funcs->C_WrapKey) (session, &mech, wrapper, wrapped, buffer, n_buffer);
	P11T_CHECK_RV ("CKM_AES_CBC_PAD", rv, CKR_OK);

	if (*n_buffer != length)
		P11T_CHECK_FAIL_MSG ("CKM_AES_CBC_PAD", "Buffer length retrieved in advance should match output length");

	/* Compare to our generated key */
	attr.type = CKA_VALUE;
	attr.pValue = value;
	attr.ulValueLen = sizeof (value);
	if (p11t_object_get (session, wrapper, &attr, 1) && attr.ulValueLen != (CK_ULONG)-1) {

		switch (attr.ulValueLen) {
		case 16:
			EVP_EncryptInit (&ctx, EVP_aes_128_cbc (), value, iv);
			break;
		case 24:
			EVP_EncryptInit (&ctx, EVP_aes_192_cbc (), value, iv);
			break;
		case 32:
			EVP_EncryptInit (&ctx, EVP_aes_256_cbc (), value, iv);
			break;
		default:
			p11t_check_warn ("Couldn't check CKM_AES_CBC_PAD wrapping: %s",
			                 "The AES key length was not 16, 24 or 32");
			return CONTINUE;
		};

		if (!EVP_EncryptUpdate (&ctx, check, &outlen, value, attr.ulValueLen))
			assert (0);
		assert (outlen == attr.ulValueLen);

		if (!EVP_EncryptFinal (&ctx, check + outlen, &tmplen))
			assert (0);
		outlen += tmplen;

		EVP_CIPHER_CTX_cleanup (&ctx);

		if (outlen != *n_buffer)
			P11T_CHECK_FAIL_MSG ("CKM_AES_CBC_PAD", "Output data must be of correct length");

		if (memcmp (check, buffer, outlen) != 0)
			P11T_CHECK_FAIL_MSG ("CKM_AES_CBC_PAD", "Output data must be CBC and PKCS#7 padded");
	}

	return CONTINUE;
}

int
test_bad_unwrap (CK_SESSION_HANDLE session, CK_MECHANISM_PTR mech, CK_OBJECT_HANDLE wrapper,
                 CK_BYTE_PTR buffer, CK_ULONG n_buffer, CK_ATTRIBUTE_PTR attrs, CK_ULONG n_attrs)
{
	CK_MECHANISM other;
	CK_OBJECT_HANDLE handle;
	CK_RV rv;

	P11T_SECTION("C_UnwrapKey");

	if (!p11t_test_unexpected)
		return CONTINUE;

	rv = (p11t_module_funcs->C_UnwrapKey) ((CK_ULONG)-88, mech, wrapper, buffer, n_buffer, attrs, n_attrs, &handle);
	P11T_CHECK_RV ("Without valid session handle", rv, CKR_SESSION_HANDLE_INVALID);

	rv = (p11t_module_funcs->C_UnwrapKey) (session, NULL, wrapper, buffer, n_buffer, attrs, n_attrs, &handle);
	P11T_CHECK_RV ("With null mechanism", rv, CKR_ARGUMENTS_BAD);

	/* A mechanism not on the key */
	other.mechanism = CKM_DH_PKCS_DERIVE;
	other.pParameter = NULL;
	other.ulParameterLen = 0;

	rv = (p11t_module_funcs->C_UnwrapKey) (session, &other, wrapper, buffer, n_buffer, attrs, n_attrs, &handle);
	P11T_CHECK_RV ("With bad mechanism", rv, CKR_KEY_TYPE_INCONSISTENT);

	other.mechanism = mech->mechanism;
	rv = (p11t_module_funcs->C_UnwrapKey) (session, &other, wrapper, buffer, n_buffer, attrs, n_attrs, &handle);
	P11T_CHECK_RV ("With bad mechanism parameter", rv, CKR_MECHANISM_PARAM_INVALID);

	rv = (p11t_module_funcs->C_UnwrapKey) (session, mech, (CK_ULONG)-82, buffer, n_buffer, attrs, n_attrs, &handle);
	P11T_CHECK_RV ("With bad wrapping key handle", rv, CKR_WRAPPING_KEY_HANDLE_INVALID);

	rv = (p11t_module_funcs->C_UnwrapKey) (session, mech, wrapper, NULL, 0, attrs, n_attrs, &handle);
	P11T_CHECK_RV ("With no wrapped key buffer", rv, CKR_WRAPPED_KEY_LEN_RANGE);

	rv = (p11t_module_funcs->C_UnwrapKey) (session, mech, wrapper, buffer, n_buffer, attrs, n_attrs, NULL);
	P11T_CHECK_RV ("With null result handle", rv, CKR_ARGUMENTS_BAD);

	return CONTINUE;
}

int
test_aes_unwrap (CK_SESSION_HANDLE session, CK_OBJECT_HANDLE wrapper,
                 CK_BYTE_PTR buffer, CK_ULONG n_buffer, CK_OBJECT_HANDLE_PTR unwrapped)
{
	CK_BYTE iv[16] = { '0', '1', '2', '3', '4', '5', '6', '7', '8', '9', 'A', 'B', 'C', 'D', 'E', 'F' };
	CK_MECHANISM mech = { CKM_AES_CBC_PAD, iv, 16 };
	CK_OBJECT_CLASS klass = CKO_SECRET_KEY;
	CK_KEY_TYPE key_type = CKK_AES;
	CK_RV rv;

	CK_ATTRIBUTE attrs[] = {
		{ CKA_CLASS, &klass, sizeof (klass) },
		{ CKA_KEY_TYPE, &key_type, sizeof (key_type) },
	};

	P11T_SECTION("C_UnwrapKey");

	test_bad_unwrap (session, &mech, wrapper, buffer, n_buffer, attrs, 2);

	rv = (p11t_module_funcs->C_UnwrapKey) (session, &mech, wrapper, buffer, n_buffer, attrs, 2, unwrapped);
	P11T_CHECK_RV ("CKM_AES_CBC_PAD", rv, CKR_OK);

	return p11t_key_test (session, *unwrapped, CKO_SECRET_KEY);
}

int
p11t_aes_test_key (CK_SESSION_HANDLE session, CK_OBJECT_HANDLE key)
{
	CK_ULONG mechanisms[128];
	CK_OBJECT_HANDLE unwrapped;
	CK_BBOOL wrap, unwrap;
	CK_BYTE buffer[16384];
	CK_ULONG n_buffer;
	CK_ULONG i;

	CK_ATTRIBUTE attrs[] = {
		{ CKA_ALLOWED_MECHANISMS, mechanisms, sizeof (mechanisms) },
		{ CKA_WRAP, &wrap, sizeof (wrap) },
		{ CKA_UNWRAP, &unwrap, sizeof (unwrap) },
	};

	/* Let's see if it support wrapping */
	if (!p11t_object_get (session, key, attrs, 3))
		return STOP;

	if (!wrap)
		return CONTINUE;

	/* We try to wrap itself, not sure if that's the best idea ... */
	for (i = 0; i < attrs[0].ulValueLen / sizeof (CK_ULONG); ++i) {
		if (mechanisms[i] == CKM_AES_CBC_PAD)
			if (!test_aes_wrap (session, key, key, buffer, &n_buffer))
				return STOP;
	}

	if (!unwrap)
		return CONTINUE;

	if (!test_aes_unwrap (session, key, buffer, n_buffer, &unwrapped))
		return STOP;

	return CONTINUE;
}