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

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

CK_RV rsa_test(CK_FUNCTION_LIST_PTR pToken, CK_SESSION_HANDLE hSession, 
			   CK_OBJECT_HANDLE hPubKey, CK_OBJECT_HANDLE hPriKey);

CK_RV rsa_test2(CK_FUNCTION_LIST_PTR pToken, CK_SESSION_HANDLE hSession, 
				CK_OBJECT_HANDLE &hPubKey, CK_OBJECT_HANDLE &hPriKey);

/*.....RSA 1024 Test Vector.......*/

CK_OBJECT_CLASS class1 = CKO_PRIVATE_KEY;
CK_OBJECT_CLASS class2 = CKO_PUBLIC_KEY;

#define labelPu		"My public key"		//Label of public key
#define labelPr		"My private key"	//Label of private key


CK_BYTE modulus[128] = {
	0xcc
	, 0x01, 0x01, 0x48, 0xa8, 0xf9, 0x14, 0x20, 0xd4, 0xbe, 0x59, 0x9a, 0xb1, 0x42, 0xc3, 0xb7, 0x53
	, 0x9d, 0x4f, 0xfb, 0x48, 0xcc, 0x84, 0xac, 0x80, 0x16, 0x0b, 0x7f, 0xa0, 0xf3, 0xc2, 0x66, 0x74
	, 0xe8, 0xed, 0x83, 0xc5, 0xf9, 0x21, 0x23, 0x33, 0xc1, 0xa9, 0x17, 0x83, 0x60, 0x43, 0xd5, 0x16
	, 0x5a, 0x32, 0x71, 0x25, 0x01, 0xe5, 0xab, 0xac, 0xf1, 0xd2, 0x34, 0x94, 0xb2, 0x49, 0x47, 0x6f
	, 0x98, 0x9a, 0x77, 0xdd, 0x0c, 0x70, 0xd8, 0xd3, 0x82, 0xdc, 0x98, 0x93, 0xce, 0x5b, 0x79, 0xd3
	, 0xc5, 0xb2, 0x62, 0xc9, 0x53, 0xb7, 0xbb, 0xd3, 0x54, 0xd5, 0x4e, 0xcd, 0x08, 0x6a, 0x65, 0xf2
	, 0xb7, 0xa6, 0x96, 0x57, 0xca, 0xb2, 0x36, 0xde, 0xa8, 0xd5, 0xd9, 0xfe, 0x3b, 0x79, 0x30, 0xef
	, 0x73, 0xd7, 0x99, 0x2d, 0x49, 0xb9, 0x2f, 0xba, 0x82, 0x6b, 0x21, 0xac, 0xeb, 0xcd, 0x2b
};


CK_BYTE publicExponent[3] = {
	0x01, 0x00, 0x01
};

CK_BYTE privateExponent[128] = {
	0x9a
	, 0xac, 0x14, 0x66, 0x1b, 0xbf, 0x25, 0xd4, 0xb2, 0x98, 0x53, 0xe5, 0xa6, 0x8a, 0x41, 0x4b, 0xc8
	, 0xf8, 0x98, 0x00, 0xa4, 0xb4, 0x73, 0xcf, 0x3a, 0x19, 0xfc, 0xc3, 0xda, 0xab, 0x96, 0x74, 0x0a
	, 0x25, 0xe4, 0x97, 0x29, 0x4d, 0xf7, 0x22, 0xf2, 0x85, 0xdd, 0x87, 0x77, 0x53, 0x50, 0x7a, 0xca
	, 0x44, 0xbe, 0xc3, 0xec, 0x54, 0x20, 0xe7, 0x65, 0xeb, 0xbe, 0xfd, 0xa0, 0xd4, 0x00, 0x1b, 0xef
	, 0x83, 0x9d, 0x8c, 0x66, 0xb9, 0x06, 0x2e, 0x59, 0x8b, 0x40, 0x2b, 0xc4, 0x45, 0x6a, 0x76, 0x54
	, 0x61, 0x20, 0x6c, 0xb0, 0x7f, 0xe4, 0x82, 0xb3, 0x58, 0x8b, 0x67, 0x95, 0xd4, 0x12, 0xa7, 0x0a
	, 0x0d, 0x9f, 0xbf, 0xb2, 0xca, 0xb8, 0x2b, 0xf3, 0xdd, 0xee, 0x10, 0x03, 0x28, 0x71, 0x3c, 0xcc
	, 0xa5, 0xa3, 0xab, 0x8f, 0x49, 0x9f, 0xa7, 0x8e, 0x68, 0xa2, 0x95, 0xcb, 0xdf, 0x2c, 0xd1
};

CK_BYTE prime1[64] = {
	0xef
	, 0x8e, 0x2b, 0x20, 0x7c, 0x7d, 0xeb, 0x54, 0xcc, 0xdc, 0x84, 0x6c, 0x73, 0x67, 0xb4, 0xf8, 0x31
	, 0xc1, 0x50, 0x64, 0x6b, 0x57, 0xf1, 0x2b, 0xc0, 0xc4, 0xbb, 0x47, 0x40, 0x4b, 0xfc, 0x21, 0xa0
	, 0x78, 0xf6, 0x47, 0xb9, 0xe7, 0x18, 0x90, 0x55, 0x76, 0xf1, 0x47, 0xae, 0xe7, 0x91, 0xb2, 0x5e
	, 0x91, 0x4a, 0x99, 0xb6, 0xba, 0x1e, 0xd8, 0xbd, 0xc6, 0x3d, 0x7f, 0x77, 0xa2, 0x11, 0xf3
};


CK_BYTE prime2[64] = {
	0xda
	, 0x02, 0x12, 0xa5, 0x47, 0xa4, 0xf6, 0xdf, 0xcf, 0xa3, 0x67, 0x71, 0x64, 0xbf, 0x40, 0x63, 0x94
	, 0xf1, 0xab, 0xeb, 0xd8, 0x8d, 0x85, 0x95, 0x06, 0x3c, 0xe9, 0xb2, 0x8e, 0x24, 0xf2, 0xa3, 0x47
	, 0xb7, 0x40, 0xb3, 0xaa, 0x8b, 0xa5, 0x08, 0x78, 0x90, 0x8f, 0xe3, 0x3c, 0xb0, 0xc1, 0xef, 0x4b
	, 0x16, 0x83, 0x3e, 0x7a, 0x5f, 0x5d, 0x55, 0x5f, 0xd4, 0x1a, 0xb9, 0x30, 0x67, 0x6d, 0xe9
};

CK_BYTE exponent1[64] = {
	0x69
	, 0x5b, 0x16, 0x5e, 0x90, 0x41, 0xf4, 0xf8, 0xf3, 0xac, 0x5c, 0x55, 0x32, 0x20, 0x3f, 0x38, 0x14
	, 0xf1, 0xd7, 0x40, 0x82, 0xf1, 0x47, 0x71, 0x1d, 0xb3, 0xa8, 0x45, 0xfa, 0xe2, 0x30, 0x3e, 0x3d
	, 0xcc, 0x55, 0x4c, 0x50, 0x57, 0x71, 0xb7, 0x30, 0xcf, 0x41, 0xe8, 0x6b, 0xe5, 0x0e, 0x6f, 0xeb
	, 0x25, 0x6d, 0x30, 0xd3, 0x6b, 0x61, 0xdb, 0xa3, 0x03, 0x07, 0xd3, 0x02, 0xb3, 0x75, 0xbf
};

CK_BYTE exponent2[64] = {
	0xb9
	, 0x25, 0x4d, 0x84, 0x84, 0xa5, 0xfa, 0xeb, 0xf4, 0xc1, 0x87, 0x96, 0x6b, 0xc3, 0x69, 0x01, 0xa9
	, 0x7a, 0xdd, 0x51, 0x3c, 0xe6, 0x63, 0xee, 0x3e, 0x0f, 0xf9, 0x41, 0x25, 0x2b, 0x07, 0xcd, 0xd7
	, 0x98, 0x30, 0x20, 0xe5, 0xd4, 0x16, 0xe2, 0x6c, 0xd2, 0x09, 0x01, 0x29, 0x87, 0xcc, 0x0c, 0x57
	, 0xae, 0xd4, 0xe3, 0x78, 0xb6, 0x08, 0xe3, 0xfa, 0xb0, 0x7d, 0xc9, 0x8c, 0xec, 0x2a, 0x11
};

CK_BYTE coefficient[64] = {
	0x6c
	, 0x5f, 0x6d, 0xe0, 0x19, 0xdb, 0xa7, 0x2d, 0x38, 0x22, 0x6b, 0x3c, 0x88, 0xfa, 0xd3, 0xfb, 0x35
	, 0x64, 0xc8, 0x93, 0x1c, 0x4a, 0xec, 0xd5, 0x31, 0x45, 0x51, 0x7a, 0xf4, 0x7f, 0xfd, 0x22, 0x39
	, 0xbf, 0xe7, 0xb6, 0xaa, 0xb3, 0xea, 0x1a, 0xeb, 0x1d, 0x4e, 0xa7, 0x1d, 0x34, 0x1e, 0x33, 0xac
	, 0x61, 0x3b, 0x53, 0x31, 0xfd, 0xd8, 0xc0, 0x0f, 0xe9, 0x93, 0xc4, 0xb8, 0x3f, 0x26, 0x88
};

CK_ULONG keyType = CKK_RSA;
CK_BBOOL true1 = CK_TRUE;
CK_BBOOL false1 = CK_FALSE;

CK_UTF8CHAR label[32] = "RSAPRIVATE Sample";
CK_UTF8CHAR label1[32] = "RSAPUBLIC Sample";
CK_UTF8CHAR subject[32] = "RSAPRIVATE Subject";

CK_BYTE id[32] = {1,2,3,4,5,6,7,8}; 


CK_ATTRIBUTE PrivateKeyTemplate[] = {
	{CKA_CLASS, &class1, sizeof(class1)},
	{CKA_KEY_TYPE, &keyType, sizeof(keyType)},
	{CKA_TOKEN, &true1, sizeof(true1)},
	{CKA_LABEL, label, sizeof(label)-1},
	{CKA_PRIVATE, &true1, sizeof(true1)},
	{CKA_EXTRACTABLE, &false1, sizeof(false1)},
	{CKA_SUBJECT, subject, sizeof(subject)},
	{CKA_ID, id, sizeof(id)},
	{CKA_SENSITIVE, &true1, sizeof(true1)},
	{CKA_DECRYPT, &true1, sizeof(true1)},
	{CKA_MODIFIABLE, &true1, sizeof(true1)}, 
	{CKA_SIGN, &true1, sizeof(true1)},
	{CKA_SIGN_RECOVER, &true1, sizeof(true1)},
	{CKA_UNWRAP, &true1, sizeof(true1)},
	{CKA_MODULUS, modulus, sizeof(modulus)},
	{CKA_PUBLIC_EXPONENT, publicExponent, sizeof(publicExponent)},
	{CKA_PRIVATE_EXPONENT, privateExponent, sizeof(privateExponent)},
	{CKA_PRIME_1, prime1, sizeof(prime1)},
	{CKA_PRIME_2, prime2, sizeof(prime2)},
	{CKA_EXPONENT_1, exponent1, sizeof(exponent1)},
	{CKA_EXPONENT_2, exponent2, sizeof(exponent2)},
	{CKA_COEFFICIENT, coefficient, sizeof(coefficient)}
}; 


CK_ATTRIBUTE PublicKeyTemplate[] = {
	{CKA_CLASS, &class2, sizeof(class2)},
	{CKA_KEY_TYPE, &keyType, sizeof(keyType)},
	{CKA_TOKEN, &true1, sizeof(true1)},
	{CKA_LABEL, label1, sizeof(label)-1},
	{CKA_ENCRYPT, &true1, sizeof(true1)},
	{CKA_MODULUS, modulus, sizeof(modulus)},
	{CKA_PUBLIC_EXPONENT, publicExponent, sizeof(publicExponent)},
	{CKA_WRAP, &true1, sizeof(true1)}
};


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 write session
	rv = pToken->C_OpenSession(pSlotList[0], CKF_SERIAL_SESSION | CKF_RW_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_OBJECT_HANDLE pubobj;
	CK_OBJECT_HANDLE priobj;

	rv = pToken->C_CreateObject(hSession, PublicKeyTemplate, countof(PublicKeyTemplate), &pubobj);
	ERROR_THROW(rv);

	rv = pToken->C_CreateObject(hSession, PrivateKeyTemplate, countof(PrivateKeyTemplate), &priobj);
	ERROR_THROW(rv);
	
	rv = rsa_test(pToken, hSession, pubobj, priobj);
	ERROR_THROW(rv);

	rv = rsa_test2(pToken, hSession, pubobj, priobj);
	ERROR_THROW(rv);

	rv = rsa_test(pToken, hSession, pubobj, priobj);
	ERROR_THROW(rv);
	
	 
END_OF_FUN:

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

	getchar();
	return 0;
}

//rsa test
CK_RV rsa_test(CK_FUNCTION_LIST_PTR pToken, CK_SESSION_HANDLE hSession, 
		   CK_OBJECT_HANDLE hPubKey, CK_OBJECT_HANDLE hPriKey)
{
	CK_RV rv = CKR_OK;
	CK_BBOOL bTrue = TRUE;

	CK_OBJECT_CLASS pubClass = CKO_PUBLIC_KEY;

	CK_MECHANISM keyGenMechanism = {CKM_RSA_PKCS_KEY_PAIR_GEN, NULL_PTR, 0};
	CK_MECHANISM ckMechanism = {CKM_RSA_PKCS, NULL_PTR, 0};
	CK_BYTE pbMsg[] = "UsbToken RunRsaKeyGenerationTest...";
	CK_ULONG ulMsgLen = strlen((const char *)pbMsg);
	CK_BYTE bSignatureBuffer[256] = {0};
	CK_ULONG ulSignatureLen = 0;
	CK_BYTE_PTR pbCipherBuffer = NULL;
	CK_ULONG ulCipherLen = 0;
	CK_BYTE_PTR pbRestoredMsg = NULL;
	CK_ULONG ulRestoredMsgLen = 0;

	ulSignatureLen	= sizeof(bSignatureBuffer);
	rv = pToken->C_SignInit(hSession, &ckMechanism, hPriKey);
	ERROR_THROW(rv);

	rv = pToken->C_Sign(hSession, 
		pbMsg,
		ulMsgLen, 
		bSignatureBuffer, &ulSignatureLen);
	ERROR_THROW(rv);

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

	rv = pToken->C_Verify(hSession, 
		pbMsg, ulMsgLen, 
		bSignatureBuffer, ulSignatureLen);
	ERROR_THROW(rv);

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

	rv = pToken->C_Encrypt(hSession, pbMsg, ulMsgLen, NULL_PTR, &ulCipherLen);
	ERROR_THROW(rv);

	pbCipherBuffer = (CK_BYTE_PTR)malloc(ulCipherLen);

	memset(pbCipherBuffer, 0, ulCipherLen);
	rv = pToken->C_Encrypt(hSession, pbMsg, ulMsgLen, 
		pbCipherBuffer, &ulCipherLen);
	ERROR_THROW(rv);

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

	rv = pToken->C_Decrypt(hSession, pbCipherBuffer, 
		ulCipherLen, NULL_PTR, &ulRestoredMsgLen);
	ERROR_THROW(rv);

	pbRestoredMsg = (CK_BYTE_PTR)malloc(ulRestoredMsgLen + 1);
	memset(pbRestoredMsg, 0, ulRestoredMsgLen + 1);

	rv = pToken->C_Decrypt(hSession, pbCipherBuffer, ulCipherLen, 
		pbRestoredMsg, &ulRestoredMsgLen);
	ERROR_THROW(rv);

	if(0 != memcmp(pbMsg, pbRestoredMsg, ulRestoredMsgLen))
		return CKR_GENERAL_ERROR;


	// Release the memory.
	if(pbCipherBuffer)
	{
		free(pbCipherBuffer);
		pbCipherBuffer = NULL;
	}
	if(pbRestoredMsg)
	{
		free(pbRestoredMsg);
		pbCipherBuffer = NULL;
	}

END_OF_FUN:

	return rv;
}

CK_RV rsa_test2(CK_FUNCTION_LIST_PTR pToken, CK_SESSION_HANDLE hSession, 
			   CK_OBJECT_HANDLE &hPubKey, CK_OBJECT_HANDLE &hPriKey)
{
	CK_MECHANISM mechanism = { 
		CKM_RSA_PKCS_KEY_PAIR_GEN, NULL_PTR, 0 
	}; 

	CK_BBOOL         bFalse = 0;
	CK_BBOOL         bTrue = 1;
	CK_KEY_TYPE      key_type  = CKK_RSA;
		
	/* setup public key attributes, specific name so we can find it 
		* later, RSA key type, saved on the UniToken and public to all
		* users of the UniToken
		*/
	CK_OBJECT_CLASS  class_public_key = CKO_PUBLIC_KEY;
	CK_ULONG vecModulusBits = 1024;//2048
	CK_BYTE	public_exponent[] = { 0x01, 0x00, 0x01 };
	CK_BYTE KID = 0;
	CK_BBOOL bVerify = true;
	CK_CHAR	label_public[] = labelPu;
	CK_CHAR	label_private[] = labelPr;
		
	CK_ATTRIBUTE publicKeyTemplate[] = 
	{
		{ CKA_CLASS,    &class_public_key,  sizeof (class_public_key) },
		{ CKA_LABEL,    label_public,       sizeof (label_public) },
		{ CKA_KEY_TYPE, &key_type,          sizeof (key_type) },
		{ CKA_TOKEN,    &bTrue,             sizeof (bTrue) },
		{ CKA_PRIVATE,  &bFalse,            sizeof (bFalse) },
		{ CKA_MODULUS_BITS, &vecModulusBits,  sizeof (vecModulusBits) },
		{ CKA_PUBLIC_EXPONENT, &public_exponent, sizeof (public_exponent) },
		{ CKA_SIGN,	  &bTrue,			  sizeof (bTrue) },
		{ CKA_VERIFY,   &bVerify,			  sizeof(bVerify)},
		{ CKA_WRAP, &bTrue, sizeof(bTrue)}
	};
		
	/* setup private key attributes, specific name so we can find it 
		* later, RSA key type, saved on the UniToken and private to only
		* the user of the UniToken
		*/
	CK_OBJECT_CLASS  class_private_key = CKO_PRIVATE_KEY;
	CK_ATTRIBUTE privateKeyTemplate[] = 
	{
		{ CKA_CLASS,    &class_private_key, sizeof (class_private_key) },
		{ CKA_LABEL,    label_private,      sizeof (label_private) },
		{ CKA_KEY_TYPE, &key_type,          sizeof (key_type) },
		{ CKA_TOKEN,    &bTrue,             sizeof (bTrue) },
		{ CKA_PRIVATE,  &bTrue,             sizeof (bTrue) },
		{ CKA_SIGN,	  &bTrue,			  sizeof (bTrue) },
		{ CKA_VERIFY,   &bVerify,			  sizeof(bVerify)},
		{ CKA_UNWRAP, &bTrue, sizeof(bTrue)}
	};

	CK_ULONG maxCount = 1024;
	CK_ULONG objCount3 = 0;
	pToken->C_FindObjectsInit(hSession, NULL_PTR, 0);
	pToken->C_FindObjects(hSession, NULL_PTR,maxCount, &objCount3);
	pToken->C_FindObjectsFinal(hSession);

	CK_RV rv = CKR_OK; 

	rv = pToken->C_GenerateKeyPair( 
		hSession, &mechanism, 
		publicKeyTemplate, 9, 
		privateKeyTemplate, 7, 
		&hPubKey, &hPriKey); 

	ERROR_THROW(rv);


	CK_ULONG objCount = 0;
	pToken->C_FindObjectsInit(hSession, NULL_PTR, 0);
	pToken->C_FindObjects(hSession, NULL_PTR,maxCount, &objCount);
	pToken->C_FindObjectsFinal(hSession);

	if(objCount - objCount3 != 2)
	{
		rv = CKR_FUNCTION_FAILED;
		ERROR_THROW(rv);
	}

END_OF_FUN:

	return rv;

}
