// encrypt_test.cpp : Defines the entry point for the console application.
//

#include "stdafx.h"
#include "../pkcs11_wrapper.h"

CK_RV encrypt_test_1(CK_FUNCTION_LIST_PTR pToken, CK_SESSION_HANDLE hSession, CK_OBJECT_HANDLE hKey);
CK_RV encrypt_test_2(CK_FUNCTION_LIST_PTR pToken, CK_SESSION_HANDLE hSession, 	CK_OBJECT_HANDLE hKey);
CK_RV encrypt_test_3(CK_FUNCTION_LIST_PTR pToken, CK_SESSION_HANDLE hSession, 	CK_OBJECT_HANDLE hKey);


int _tmain(int argc, _TCHAR* argv[])
{
	if(argc < 2)
		return 1;

	char * p11_dll = argv[1];

	pkcs11_wrapper p11_wrapper;

	if(0 != p11_wrapper.pkcs11_initialize(p11_dll))
	{
		return 1;
	}

	CK_FUNCTION_LIST_PTR pToken = p11_wrapper.GetFunctionsPtr();
	CK_RV rv = CKR_OK;

	CK_SLOT_ID_PTR pSlotList;
	CK_ULONG ulCount;
	CK_SESSION_HANDLE hSession;


	rv = pToken->C_GetSlotList(CK_TRUE, NULL_PTR, &ulCount);
	ERROR_THROW(rv);

	pSlotList = (CK_SLOT_ID_PTR)new CK_SLOT_ID[ulCount];;
	rv = pToken->C_GetSlotList(CK_TRUE, pSlotList, &ulCount);
	ERROR_THROW(rv);

	//Create a session
	rv = pToken->C_OpenSession(pSlotList[0], CKF_RW_SESSION | CKF_SERIAL_SESSION, NULL_PTR, NULL_PTR, &hSession);
	ERROR_THROW(rv);


	CK_UTF8CHAR uPin[] = "11111111";
	rv = pToken->C_Login(hSession, CKU_USER, uPin, 8);
	ERROR_THROW(rv);

	CK_BYTE cka_value[128] = {1,2,3,4,5,6,7,8,1,2,3,4,5,6,7,8};
	CK_OBJECT_CLASS oClass = CKO_SECRET_KEY;
	CK_KEY_TYPE keyType = CKK_AES; 
	CK_BBOOL bTrue = true;
	CK_BBOOL bFalse = false;
	CK_ULONG ulLen = 16;
	CK_ATTRIBUTE AESItem[] = {
		{CKA_CLASS, &oClass, sizeof(CK_OBJECT_CLASS)},
		{CKA_KEY_TYPE, &keyType, sizeof(CK_KEY_TYPE)},
		//		{CKA_TOKEN, &bTrue, sizeof(CK_BBOOL)},
		{CKA_ENCRYPT, &bTrue, sizeof(CK_BBOOL)},
		{CKA_DECRYPT, &bTrue, sizeof(CK_BBOOL)},
		{CKA_EXTRACTABLE, &bFalse, sizeof(CK_BBOOL)},
		{CKA_VALUE, cka_value, ulLen}, 
	};

	ulCount = sizeof(AESItem) / sizeof(CK_ATTRIBUTE);

	CK_OBJECT_HANDLE hKey;
	rv = pToken->C_CreateObject(hSession, AESItem, ulCount, &hKey);
	ERROR_THROW(rv);

	
	rv=  encrypt_test_1( pToken, hSession, hKey);
	ERROR_THROW(rv);

	keyType = CKK_DES;
	ulLen = 8;
	CK_ATTRIBUTE DESItem[] = {
		{CKA_CLASS, &oClass, sizeof(CK_OBJECT_CLASS)},
		{CKA_KEY_TYPE, &keyType, sizeof(CK_KEY_TYPE)},
		//		{CKA_TOKEN, &bTrue, sizeof(CK_BBOOL)},
		{CKA_ENCRYPT, &bTrue, sizeof(CK_BBOOL)},
		{CKA_DECRYPT, &bTrue, sizeof(CK_BBOOL)},
		{CKA_EXTRACTABLE, &bFalse, sizeof(CK_BBOOL)},
		{CKA_VALUE, cka_value, ulLen}, 
	};

	ulCount = sizeof(DESItem) / sizeof(CK_ATTRIBUTE);

	CK_OBJECT_HANDLE hDESKey;
	rv = pToken->C_CreateObject(hSession, DESItem, ulCount, &hDESKey);
	ERROR_THROW(rv);

	rv = encrypt_test_2(pToken, hSession, hDESKey);
	ERROR_THROW(rv);

	rv = encrypt_test_3(pToken, hSession, hDESKey);
	ERROR_THROW(rv);

END_OF_FUN:


	if(rv != CKR_OK)
	{
		printf("Error = %0x \n", rv);
	}
	else
	{
		printf("Successfully!\n");
	}
	getchar();
	return 0;
}

CK_RV encrypt_test_1(CK_FUNCTION_LIST_PTR pToken, CK_SESSION_HANDLE hSession, CK_OBJECT_HANDLE hKey)
{
	 CK_ULONG ulCount = 0;
	 CK_RV rv = CKR_OK;

	 CK_MECHANISM ckMechanism;
	 CK_ULONG Mechanism[] = {CKM_AES_CBC, CKM_AES_ECB, CKM_AES_CBC_PAD};
	 CK_BYTE iv[16] = {'*','2','1','0','4','a','e','s',1,2,3,4,5,6,7,8};



	unsigned char input[4096] = {0};
	unsigned char output[8192] = {0};
	CK_ULONG input_len = 4096;
	CK_ULONG output_len = 4096;

	for(int i = 0 ; i < input_len ; i ++)
	{
		input[i] = i % 256;
	}


	for(int i = 0 ; i < 3; i ++)
	{
		ckMechanism.mechanism = Mechanism[i];
		ckMechanism.pParameter = iv;
		ckMechanism.ulParameterLen = sizeof(iv);

		if(ckMechanism.mechanism == CKM_AES_CBC_PAD)
		{
			output_len += 16;
		}

		
		rv =  pToken->C_EncryptInit(hSession, &ckMechanism, hKey); 
		ERROR_THROW(rv);

		//system("pause");
		rv =  pToken->C_Encrypt(hSession, input, input_len, output, &output_len);
		ERROR_THROW(rv);

		rv =  pToken->C_DecryptInit(hSession, &ckMechanism, hKey); 
		ERROR_THROW(rv);

		input_len = output_len;
		rv =  pToken->C_Decrypt(hSession, output, output_len, input, &input_len);
		ERROR_THROW(rv);

		for(int j = 0 ; j < input_len ; j ++)
		{
			if( input[j] != j % 256)
			{
				rv = CKR_FUNCTION_FAILED;
				ERROR_THROW(rv);
			}
		}
	}

	input_len = 4096;
	output_len = 4096;

	CK_ULONG len = 0;
	CK_ULONG total = 0;
	for(int i = 0 ; i < 3; i ++)
	{
		ckMechanism.mechanism = Mechanism[i];
		ckMechanism.pParameter = iv;
		ckMechanism.ulParameterLen = sizeof(iv);

		/*if(ckMechanism.mechanism == CKM_AES_CBC_PAD)
		{
			output_len += 16;
		}*/

		rv =  pToken->C_EncryptInit(hSession, &ckMechanism, hKey); 
		ERROR_THROW(rv);

		len = output_len;
		total = 0;

		rv =  pToken->C_EncryptUpdate(hSession, input, input_len, output, &len);
		ERROR_THROW(rv);

		total += len;

		len = output_len - len + 16;
		rv =  pToken->C_EncryptFinal(hSession, output + total, &len);
		ERROR_THROW(rv);

		total += len;

		rv =  pToken->C_DecryptInit(hSession, &ckMechanism, hKey); 
		ERROR_THROW(rv);

		len = input_len;
		CK_ULONG total2 = 0;

		output_len = total;

		len = output_len;
		rv =  pToken->C_DecryptUpdate(hSession, output, output_len, input, &len);
		ERROR_THROW(rv);

		total2 += len;

		len = input_len - len;
		rv =  pToken->C_DecryptFinal(hSession, input + total2, &len);
		ERROR_THROW(rv);

		total2 += len;

		for(int j = 0 ; j < total2 ; j ++)
		{
			if( input[j] != j % 256)
			{
				rv = CKR_FUNCTION_FAILED;
				ERROR_THROW(rv);
			}
		}
	}

END_OF_FUN:

	return rv;
}

CK_RV encrypt_test_2(CK_FUNCTION_LIST_PTR pToken, CK_SESSION_HANDLE hSession, 	CK_OBJECT_HANDLE hKey)
{
	#define PLAINTEXT_BUF_SZ 200 
	#define CIPHERTEXT_BUF_SZ 256 

	CK_ULONG firstPieceLen, secondPieceLen; 
	CK_BYTE iv[8] = {1,2,3,4,5,6,7,8}; 
	CK_MECHANISM mechanism = { 
		CKM_DES_CBC_PAD, iv, sizeof(iv) 
	}; 

	CK_BYTE data[PLAINTEXT_BUF_SZ]; 
	CK_BYTE encryptedData[CIPHERTEXT_BUF_SZ] ={1,2,3,4,5,6,7,8,9}; 

	CK_ULONG ulEncryptedData1Len; 
	CK_ULONG ulEncryptedData2Len; 
	CK_ULONG ulEncryptedData3Len; 
	CK_RV rv = CKR_OK; 

	firstPieceLen = 90; 
	secondPieceLen = PLAINTEXT_BUF_SZ-firstPieceLen; 

	rv = pToken->C_EncryptInit(hSession, &mechanism, hKey); 
	ERROR_THROW(rv);

	/* Encrypt first piece */ 
	ulEncryptedData1Len = sizeof(encryptedData); 
	rv = pToken->C_EncryptUpdate( 
		hSession,  
		&data[0], firstPieceLen, 
		&encryptedData[0], &ulEncryptedData1Len); 
	ERROR_THROW(rv);

	/* Encrypt second piece */ 
	ulEncryptedData2Len = sizeof(encryptedData)-ulEncryptedData1Len; 

	rv = pToken->C_EncryptUpdate( 
		hSession, 
		&data[firstPieceLen], secondPieceLen, 
		&encryptedData[ulEncryptedData1Len], 
		&ulEncryptedData2Len); 
	ERROR_THROW(rv);

	/* Get last little encrypted bit */ 
	ulEncryptedData3Len = sizeof(encryptedData)-ulEncryptedData1Len-ulEncryptedData2Len; 

	rv = pToken->C_EncryptFinal( 
		hSession, 
		&encryptedData[ulEncryptedData1Len + ulEncryptedData2Len], 
		&ulEncryptedData3Len); 

	ERROR_THROW(rv);

END_OF_FUN:

	return rv;
}

CK_RV encrypt_test_3(CK_FUNCTION_LIST_PTR pToken, CK_SESSION_HANDLE hSession, 	CK_OBJECT_HANDLE hKey)
{

#define CIPHERTEXT_BUF_SZ_1 256 
#define PLAINTEXT_BUF_SZ_1 256 

	CK_ULONG firstEncryptedPieceLen, secondEncryptedPieceLen; 
	CK_BYTE iv[8] = {1,2,3,4,5,6,7,8}; 
	CK_MECHANISM mechanism = { 
		CKM_DES_CBC, iv, sizeof(iv) 
	}; 
	CK_BYTE data[PLAINTEXT_BUF_SZ_1]; 
	CK_BYTE encryptedData[CIPHERTEXT_BUF_SZ_1]; 
	CK_ULONG ulData1Len, ulData2Len, ulData3Len; 
	CK_RV rv = CKR_OK;  

	firstEncryptedPieceLen = 90; 
	secondEncryptedPieceLen = CIPHERTEXT_BUF_SZ_1-firstEncryptedPieceLen; 
	rv = pToken->C_DecryptInit(hSession, &mechanism, hKey); 
	ERROR_THROW(rv);

	/* Decrypt first piece */ 
	ulData1Len = sizeof(data); 
	rv = pToken->C_DecryptUpdate( 
		hSession, 
		&encryptedData[0], firstEncryptedPieceLen, 
		&data[0], &ulData1Len); 
	ERROR_THROW(rv);

	/* Decrypt second piece */ 
	ulData2Len = sizeof(data)-ulData1Len; 
	rv = pToken->C_DecryptUpdate( 
		hSession, 
		&encryptedData[firstEncryptedPieceLen], 
		secondEncryptedPieceLen, 
		&data[ulData1Len], &ulData2Len); 
	ERROR_THROW(rv);

	/* Get last little decrypted bit */ 
	ulData3Len = sizeof(data)-ulData1Len-ulData2Len; 
	rv = pToken->C_DecryptFinal( 
		hSession, 
		&data[ulData1Len+ulData2Len], &ulData3Len); 
	ERROR_THROW(rv);

END_OF_FUN:

	return rv;
}