#include "p11-tests-lib.h"

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

CK_INFO slot_global;

int p11t_slot_virtual = 0;
CK_ULONG p11t_slot_count = 0;
static CK_SLOT_ID_PTR slot_ids = NULL;
static CK_SLOT_INFO_PTR slot_info = NULL;
static CK_TOKEN_INFO_PTR slot_token_info = NULL;

static CK_ULONG *slot_mech_count;
static CK_MECHANISM_TYPE_PTR *slot_mech_type;
static CK_MECHANISM_INFO_PTR *slot_mech_info;

/* TODO: Get this from a header */
#define CKF_VIRTUAL_SLOTS 0x40000000

/* ----------------------------------------------------------------------------------
 * TESTS
 */

static int
test_slot_global(void)
{
	CK_RV rv;

	P11T_SECTION("C_GetInfo");

	assert(p11t_module_funcs);

	if(p11t_test_unexpected)
	{
		rv = (p11t_module_funcs->C_GetInfo)(NULL);
		P11T_CHECK_RV("Null argument", rv, CKR_ARGUMENTS_BAD);
	}

	/* Obvious crap fill */
	memset(&slot_global, 0xFF, sizeof(CK_INFO));

	rv = (p11t_module_funcs->C_GetInfo)(&slot_global);
	P11T_CHECK_RV("Normal call", rv, CKR_OK);

	P11T_SECTION("CK_INFO");

	P11T_CHECK_PADDED("manufacturerID", slot_global.manufacturerID);
	P11T_CHECK_PADDED("libraryDescription", slot_global.libraryDescription);

	P11T_CHECK_ULONG("flags", slot_global.flags, CKF_VIRTUAL_SLOTS);
	if (slot_global.flags & CKF_VIRTUAL_SLOTS)
		p11t_slot_virtual = 1;

	return CONTINUE;
}

static int
compar_slot_id(const void *one, const void *two)
{
	CK_SLOT_ID a = *((CK_SLOT_ID*)one);
	CK_SLOT_ID b = *((CK_SLOT_ID*)two);
	if(a == b)
		return 0;
	if(a > b)
		return 1;
	return -1;
}

static int
test_slot_info(void)
{
	CK_SLOT_ID_PTR present, only;
	CK_ULONG n_present, n_only, count;
	CK_ULONG i;
	CK_RV rv;

	P11T_SECTION("C_GetSlotList");

	assert(p11t_module_funcs);

	if(p11t_test_unexpected)
	{
		rv = (p11t_module_funcs->C_GetSlotList)(FALSE, NULL, NULL);
		P11T_CHECK_RV("Null arguments", rv, CKR_ARGUMENTS_BAD);
	}

	rv = (p11t_module_funcs->C_GetSlotList)(FALSE, NULL, &p11t_slot_count);
	P11T_CHECK_RV("Retrieving the count", rv, CKR_OK);

	if(p11t_slot_count == 0)
	{
		p11t_check_info("no slots to test");
		return CONTINUE;
	}

	/* Allocate a bit extra. We're going to try and trip up module */
	slot_ids = calloc(p11t_slot_count + 5, sizeof(CK_SLOT_ID));
	assert(slot_ids);

	slot_info = calloc(p11t_slot_count, sizeof(CK_SLOT_INFO));
	assert(slot_info);

	slot_token_info = calloc(p11t_slot_count, sizeof(CK_TOKEN_INFO));
	assert(slot_token_info);

	if(p11t_test_unexpected)
	{
		count = 0;
		rv = (p11t_module_funcs->C_GetSlotList)(FALSE, slot_ids, &count);
		P11T_CHECK_RV("Passing buffer with zero count", rv, CKR_BUFFER_TOO_SMALL);
		P11T_CHECK_ULONG("Resulting count when buffer with zero count passed", count, p11t_slot_count);

		if(p11t_slot_count > 1)
		{
			count = 1;
			rv = (p11t_module_funcs->C_GetSlotList)(FALSE, slot_ids, &count);
			P11T_CHECK_RV("Passing buffer along with low count", rv, CKR_BUFFER_TOO_SMALL);
			P11T_CHECK_ULONG("Resulting count when buffer with low count passed", count, p11t_slot_count);
		}
	}

	if(p11t_test_unexpected)
		count = p11t_slot_count + 5;
	else
		count = p11t_slot_count;
	rv = (p11t_module_funcs->C_GetSlotList)(FALSE, slot_ids, &count);
	P11T_CHECK_RV("Normal call", rv, CKR_OK);

	if(p11t_test_unexpected)
		P11T_CHECK_ULONG("Count invalid when too much buffer passed", count, p11t_slot_count);

	/* When we ask for just token valid ones, they should only be the ones we noted with flags */
	present = calloc(p11t_slot_count, sizeof(CK_SLOT_ID));
	only = calloc(p11t_slot_count, sizeof(CK_SLOT_ID));
	assert(present && only);
	n_present = 0;

	/* Now go through and load info about each slot */
	for(i = 0; i < p11t_slot_count; ++i)
	{
		P11T_SECTION("C_GetSlotInfo");

		if(p11t_test_unexpected)
		{
			rv = (p11t_module_funcs->C_GetSlotInfo)(slot_ids[i], NULL);
			P11T_CHECK_RV("Null argument", rv, CKR_ARGUMENTS_BAD);
		}

		rv = (p11t_module_funcs->C_GetSlotInfo)(slot_ids[i], &slot_info[i]);
		P11T_CHECK_RV("Normal call", rv, CKR_OK);

		P11T_SECTION("CK_SLOT_INFO");

		P11T_CHECK_PADDED("slotDescription", slot_info[i].slotDescription);
		P11T_CHECK_PADDED("manufacturerID", slot_info[i].manufacturerID);

		P11T_CHECK_MASK("flags", slot_info[i].flags,
		                CKF_TOKEN_PRESENT | CKF_REMOVABLE_DEVICE | CKF_HW_SLOT);

		P11T_CHECK_NOTE("CKF_TOKEN_PRESENT flag is equivalent to C_GetSlotList(TRUE, ...)");
		if((slot_info[i].flags & CKF_TOKEN_PRESENT) == CKF_TOKEN_PRESENT)
		{
			/* Note if token is present for later */
			present[n_present++] = slot_ids[i];

			P11T_SECTION("C_GetTokenInfo");

			if(p11t_test_unexpected)
			{
				rv = (p11t_module_funcs->C_GetTokenInfo)(slot_ids[i], NULL);
				P11T_CHECK_RV("Null arguments", rv, CKR_ARGUMENTS_BAD);
			}

			rv = (p11t_module_funcs->C_GetTokenInfo)(slot_ids[i], &slot_token_info[i]);
			P11T_CHECK_RV("Normal call", rv, CKR_OK);

			P11T_SECTION("CK_TOKEN_INFO");

			P11T_CHECK_PADDED("label", slot_token_info[i].label);
			P11T_CHECK_PADDED("manufacturerID", slot_token_info[i].manufacturerID);
			P11T_CHECK_PADDED("model", slot_token_info[i].model);
			P11T_CHECK_PADDED("serialNumber", slot_token_info[i].serialNumber);

			P11T_CHECK_MASK("flags", slot_token_info[i].flags,
			                CKF_RNG | CKF_WRITE_PROTECTED | CKF_LOGIN_REQUIRED | CKF_USER_PIN_INITIALIZED |
			                CKF_RESTORE_KEY_NOT_NEEDED | CKF_CLOCK_ON_TOKEN | CKF_PROTECTED_AUTHENTICATION_PATH |
			                CKF_DUAL_CRYPTO_OPERATIONS | CKF_TOKEN_INITIALIZED | CKF_SECONDARY_AUTHENTICATION |
			                CKF_USER_PIN_COUNT_LOW | CKF_USER_PIN_FINAL_TRY | CKF_USER_PIN_LOCKED |
			                CKF_USER_PIN_TO_BE_CHANGED | CKF_SO_PIN_COUNT_LOW | CKF_SO_PIN_FINAL_TRY |
			                CKF_SO_PIN_LOCKED | CKF_SO_PIN_TO_BE_CHANGED);

			/* Can't validate all the statistics, any number is valid */

			P11T_CHECK_NOTE("Validate token time when CKF_CLOCK_ON_TOKEN");
			if((slot_token_info[i].flags & CKF_CLOCK_ON_TOKEN) == CKF_CLOCK_ON_TOKEN)
			{
				int year, month, day, hour, minute, second, extra, n;
				char buffer[sizeof(slot_token_info[i].utcTime) + 1];

				memcpy(buffer, (char*)slot_token_info[i].utcTime,
				       sizeof(slot_token_info[i].utcTime));
				buffer[sizeof(slot_token_info[i].utcTime)] = 0;

				n = sscanf(buffer, "%04d%02d%02d%02d%02d%02d%02d",
				           &year, &month, &day, &hour, &minute, &second, &extra);

				if(n != 7 ||
				   year < 0 || year > 9999 ||
				   month < 1 || month > 12 ||
				   day < 1 || day > 31 ||
				   hour < 0 || hour > 23 ||
				   minute < 0 || minute > 59 ||
				   second < 0 || second > 50 ||
				   extra != 0)
				{
					p11t_check_fail("C_GetTokenInfo: invalid time on token: %s", buffer);
				}
			}
		}
		else if(p11t_test_unexpected)
		{
			P11T_SECTION("C_GetTokenInfo");

			rv = (p11t_module_funcs->C_GetTokenInfo)(slot_ids[i], &slot_token_info[i]);
			P11T_CHECK_RV("Calling on slot without token", rv, CKR_TOKEN_NOT_PRESENT);
		}
	}

	P11T_SECTION("C_GetSlotList");

	n_only = p11t_slot_count;
	rv = (p11t_module_funcs->C_GetSlotList)(TRUE, only, &n_only);
	P11T_CHECK_RV("Listing only tokens", rv, CKR_OK);

	P11T_CHECK_ULONG("Number of present tokens doesn't match token info flags. ie: CKF_TOKEN_PRESENT", n_only, n_present);

	qsort(only, n_only, sizeof(CK_SLOT_ID), compar_slot_id);
	qsort(present, n_present, sizeof(CK_SLOT_ID), compar_slot_id);
	if(memcmp(only, present, sizeof(CK_SLOT_ID) * n_only) != 0)
		P11T_CHECK_FAIL("Present tokens don't match those in token info flags. ie: CKF_TOKEN_PRESENT");

	free(only);
	free(present);

	return CONTINUE;
}

static int
test_slot_events(void)
{
	P11T_SECTION("C_WaitForSlotEvent");
	P11T_CHECK_NOTE("Not Tested");
	return CONTINUE;
}

static int
test_slot_mechanisms(void)
{
	CK_MECHANISM_TYPE_PTR mech_list;
	CK_MECHANISM_INFO_PTR mech_info;
	CK_SLOT_ID slot_id;
	CK_ULONG mech_count, value;
	CK_ULONG i, j;
	CK_RV rv;

	assert(p11t_module_funcs);

	P11T_SECTION("C_GetMechanismList");

	if(!p11t_slot_count)
		return CONTINUE;

	if(p11t_test_unexpected && !p11t_slot_virtual)
	{
		mech_count = 0;
		rv = (p11t_module_funcs->C_GetMechanismList)((CK_ULONG)-10, NULL, &mech_count);
		P11T_CHECK_RV("Invalid Slot", rv, CKR_SLOT_ID_INVALID);
	}

	slot_mech_count = calloc(p11t_slot_count, sizeof(CK_ULONG));
	assert(slot_mech_count);

	slot_mech_info = calloc(p11t_slot_count, sizeof(CK_MECHANISM_INFO_PTR));
	assert(slot_mech_info);

	slot_mech_type = calloc(p11t_slot_count, sizeof(CK_MECHANISM_TYPE_PTR));
	assert(slot_mech_type);

	for(i = 0; i < p11t_slot_count; ++i)
	{
		slot_id = slot_ids[i];

		if(p11t_test_unexpected)
		{
			rv = (p11t_module_funcs->C_GetMechanismList)(slot_id, NULL, NULL);
			P11T_CHECK_RV("Null arguments", rv, CKR_ARGUMENTS_BAD);
		}

		rv = (p11t_module_funcs->C_GetMechanismList)(slot_id, NULL, &mech_count);
		P11T_CHECK_RV("Without buffer", rv, CKR_OK);

		mech_list = NULL;
		mech_info = NULL;

		if(mech_count > 0)
		{
			mech_list = calloc(mech_count + 5, sizeof(CK_MECHANISM_TYPE));
			assert(mech_list);

			if(p11t_test_unexpected)
			{
				value = 0;
				rv = (p11t_module_funcs->C_GetMechanismList)(slot_id, mech_list, &value);
				P11T_CHECK_RV("Zero count but buffer present", rv, CKR_BUFFER_TOO_SMALL);
				P11T_CHECK_ULONG("Should return number of mechs", value, mech_count);

				if(mech_count > 1)
				{
					value = 1;
					rv = (p11t_module_funcs->C_GetMechanismList)(slot_id, mech_list, &value);
					P11T_CHECK_RV("Low count but buffer present", rv, CKR_BUFFER_TOO_SMALL);
					P11T_CHECK_ULONG("Should return number of mechs", value, mech_count);
				}
			}

			if(p11t_test_unexpected)
				value = mech_count + 5;
			else
				value = mech_count;
			rv = (p11t_module_funcs->C_GetMechanismList)(slot_id, mech_list, &value);
			P11T_CHECK_RV("Call with too much buffer", rv, CKR_OK);
			P11T_CHECK_ULONG("Should return number of mechs", value, mech_count);


			mech_info = calloc(mech_count, sizeof(CK_MECHANISM_INFO));
			assert(mech_info);

			P11T_SECTION("C_GetMechanismInfo");

			if(p11t_test_unexpected)
			{
				rv = (p11t_module_funcs->C_GetMechanismInfo)(slot_id, (CK_ULONG)-10, mech_info);
				P11T_CHECK_RV("Invalid mechanism", rv, CKR_MECHANISM_INVALID);

				rv = (p11t_module_funcs->C_GetMechanismInfo)(slot_id, mech_list[0], NULL);
				P11T_CHECK_RV("Null arguments", rv, CKR_ARGUMENTS_BAD);

				if (!p11t_slot_virtual)
				{
					rv = (p11t_module_funcs->C_GetMechanismInfo)((CK_ULONG)-11, mech_list[0], mech_info);
					P11T_CHECK_RV("Invalid slot id", rv, CKR_SLOT_ID_INVALID);
				}
			}

			for(j = 0; j < mech_count; ++j)
			{
				rv = (p11t_module_funcs->C_GetMechanismInfo)(slot_id, mech_list[j], &mech_info[j]);
				P11T_CHECK_RV("Normal call", rv, CKR_OK);

				P11T_SECTION("CK_MECHANISM_INFO");

				if(mech_info[j].ulMinKeySize > mech_info[j].ulMaxKeySize)
					P11T_CHECK_FAIL("Mechanism min key size should not be greater than max");

				P11T_CHECK_MASK("flags", mech_info[j].flags,
				                CKF_HW | CKF_ENCRYPT | CKF_DECRYPT | CKF_DIGEST | CKF_SIGN |
				                CKF_SIGN_RECOVER | CKF_VERIFY | CKF_VERIFY_RECOVER | CKF_GENERATE |
				                CKF_GENERATE_KEY_PAIR | CKF_WRAP | CKF_UNWRAP | CKF_DERIVE);

				P11T_CHECK_NFLAG("flags", mech_info[j].flags, CKF_EXTENSION);
			}

		}

		slot_mech_info[i] = mech_info;
		slot_mech_type[i] = mech_list;
		slot_mech_count[i] = mech_count;
	}

	return CONTINUE;
}

static int
test_init_token(void)
{
	P11T_SECTION("C_InitToken");
	P11T_CHECK_NOTE("Not Tested");
	return CONTINUE;
}

void
p11t_slot_tests(void)
{
	test_slot_global();
	test_slot_info();
	test_slot_events();
	test_init_token();
	test_slot_mechanisms();
}

void
p11t_slot_cleanup(void)
{
	CK_ULONG i;

	free(slot_ids);
	slot_ids = NULL;

	for(i = 0; i < p11t_slot_count; ++i)
	{
		free(slot_mech_type[i]);
		free(slot_mech_info[i]);
	}

	free(slot_info);
	slot_info = NULL;

	free(slot_token_info);
	slot_token_info = NULL;

	free(slot_mech_count);
	slot_mech_count = NULL;

	free(slot_mech_type);
	slot_mech_type = NULL;

	free(slot_mech_info);
	slot_mech_info = NULL;

	p11t_slot_count = 0;
}

/* ----------------------------------------------------------------------------------
 * METHODS
 */

static int
get_slot_index(CK_SLOT_ID slot)
{
	CK_ULONG i;

	for(i = 0; i < p11t_slot_count; ++i)
	{
		if(slot == slot_ids[i])
			return i;
	}

	return -1;
}

CK_SLOT_ID
p11t_slot_get_id(int index)
{
	if(index >= (int)p11t_slot_count)
		return CK_INVALID;
	return slot_ids[index];
}

CK_SLOT_INFO_PTR
p11t_slot_get_info(CK_SLOT_ID slot)
{
	int i = get_slot_index(slot);
	if(i < 0)
		return NULL;
	return &slot_info[i];
}

CK_TOKEN_INFO_PTR
p11t_slot_get_token_info(CK_SLOT_ID slot)
{
	int i = get_slot_index(slot);
	if(i < 0)
		return NULL;

	if(slot_info[i].flags & CKF_TOKEN_PRESENT)
		return &slot_token_info[i];

	return NULL;
}

CK_MECHANISM_INFO_PTR
p11t_slot_get_mech_info(CK_SLOT_ID slot, CK_MECHANISM_TYPE mech)
{
	int i;
	CK_ULONG j;

	i = get_slot_index(slot);
	assert(i >= 0);

	for(j = 0; j < slot_mech_count[i]; ++j)
	{
		if(mech == slot_mech_type[i][j])
			return &slot_mech_info[i][j];
	}

	return CK_FALSE;
}

void
p11t_slot_for_each_mech(CK_MECHANISM_TYPE mech_type,
                        P11tSlotMechCallback callback)
{
	CK_ULONG i, j;

	for(i = 0; i < p11t_slot_count; ++i)
	{
		for(j = 0; j < slot_mech_count[i]; ++j)
		{
			if(mech_type == CK_INVALID || slot_mech_type[i][j] == mech_type)
				(callback)(slot_ids[i], slot_mech_type[i][j], &slot_mech_info[i][j]);
		}
	}
}