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

#include "stdafx.h"

#include "../pkcs11_wrapper.h"

long FindPublicKey(CK_FUNCTION_LIST_PTR pFunctionList, CK_SESSION_HANDLE hSession, CK_OBJECT_HANDLE *phPriKey);
long FindPrivateKey(CK_FUNCTION_LIST_PTR pFunctionList, CK_SESSION_HANDLE hSession, CK_OBJECT_HANDLE *phPriKey);

CK_OBJECT_HANDLE FindObject(CK_FUNCTION_LIST_PTR pFunctionList, CK_SESSION_HANDLE hSession, CK_ATTRIBUTE_PTR pAttrTemplates, CK_ULONG ulCount );

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 read only session
	rv = pToken->C_OpenSession(pSlotList[0], CKF_SERIAL_SESSION, NULL_PTR, NULL_PTR, &hSession);
	if (rv != CKR_OK)
	{
		delete[] pSlotList;
		return 2;
	}
	
	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, &bTrue, 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);

	CK_OBJECT_HANDLE hPriKey, hPubKey;
	rv = FindPublicKey(pToken, hSession, &hPubKey);
	ERROR_THROW(rv);

	CK_MECHANISM ckMechanism;
	ckMechanism.mechanism = CKM_RSA_PKCS;
	ckMechanism.pParameter = NULL;
	ckMechanism.ulParameterLen = 0;

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

	CK_ULONG wrapLen = 0;
	rv = pToken->C_WrapKey(hSession, &ckMechanism, hPubKey, hKey, NULL_PTR, &wrapLen);
	ERROR_THROW(rv);

	CK_BYTE_PTR pWrappedData = (CK_BYTE_PTR) malloc(wrapLen);

	rv = pToken->C_WrapKey(hSession, &ckMechanism, hPubKey, hKey, pWrappedData, &wrapLen);
	ERROR_THROW(rv);

	rv = FindPrivateKey(pToken, hSession, &hPriKey);
	ERROR_THROW(rv);
	

	CK_OBJECT_CLASS keyClass = CKO_SECRET_KEY; 
	keyType = CKK_AES; 
	CK_ATTRIBUTE template1[] = { 
		{CKA_CLASS, &keyClass, sizeof(keyClass)}, 
		{CKA_KEY_TYPE, &keyType, sizeof(keyType)}, 
		{CKA_ENCRYPT, &bTrue, sizeof(true)}, 
		{CKA_DECRYPT, &bTrue, sizeof(true)},
		{CKA_VALUE_LEN, &ulLen, sizeof(ulLen)}
	}; 
	ulCount = sizeof(template1) / sizeof(CK_ATTRIBUTE);

	CK_OBJECT_HANDLE hNewKey = 0;
	rv = pToken->C_UnwrapKey(hSession, &ckMechanism, hPriKey, pWrappedData, wrapLen, template1, ulCount, &hNewKey);
	ERROR_THROW(rv);


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

	getchar();
	return 0;
}


long FindPublicKey(CK_FUNCTION_LIST_PTR pFunctionList, CK_SESSION_HANDLE hSession, CK_OBJECT_HANDLE *phPriKey)
{
	CK_OBJECT_CLASS objectClass = CKO_PUBLIC_KEY;
	CK_KEY_TYPE keyType = CKK_RSA;
	CK_BBOOL p11True = TRUE;                  

	CK_ATTRIBUTE kpPriAttr[]={
		{CKA_CLASS, &objectClass, sizeof(objectClass)},
		{CKA_KEY_TYPE, &keyType, sizeof(keyType)},
	};

	CK_ULONG ulCount = sizeof(kpPriAttr) / sizeof(CK_ATTRIBUTE);

	*phPriKey = FindObject(pFunctionList, hSession, kpPriAttr,ulCount);

	if(*phPriKey == NULL) //not found
	{
		return CKR_OBJECT_HANDLE_INVALID;
	}	

	return CKR_OK;
}

long FindPrivateKey(CK_FUNCTION_LIST_PTR pFunctionList, CK_SESSION_HANDLE hSession, CK_OBJECT_HANDLE *phPriKey)
{
	CK_OBJECT_CLASS objectClass = CKO_PRIVATE_KEY;
	CK_KEY_TYPE keyType = CKK_RSA;
	CK_BBOOL p11True = TRUE;                  

	CK_ATTRIBUTE kpPriAttr[]={
		{CKA_CLASS, &objectClass, sizeof(objectClass)},
		{CKA_KEY_TYPE, &keyType, sizeof(keyType)},
	};

	CK_ULONG ulCount = sizeof(kpPriAttr) / sizeof(CK_ATTRIBUTE);

	*phPriKey = FindObject(pFunctionList, hSession, kpPriAttr,ulCount);

	if(*phPriKey == NULL) //not found
	{
		return CKR_OBJECT_HANDLE_INVALID;
	}	

	return CKR_OK;
}

CK_OBJECT_HANDLE FindObject(CK_FUNCTION_LIST_PTR pFunctionList, CK_SESSION_HANDLE hSession, CK_ATTRIBUTE_PTR pAttrTemplates, CK_ULONG ulCount )
{
	CK_RV rv = CKR_OK;

	rv = pFunctionList->C_FindObjectsInit(hSession, pAttrTemplates, ulCount); 
	if (rv != CKR_OK)
	{
		return 0;
	}

	CK_OBJECT_HANDLE foundObj = NULL;
	CK_ULONG objNumber = 1;

	rv = pFunctionList->C_FindObjects(hSession, &foundObj, 1, &objNumber);
	if (rv != 0)
	{
		pFunctionList->C_FindObjectsFinal(hSession);
		return NULL;
	}

	pFunctionList->C_FindObjectsFinal(hSession);

	if(objNumber  ==  0) //not found
		foundObj = NULL;

	return foundObj;
}
