/*
[]=========================================================================[]

	Copyright(C) Feitian Technologies Co., Ltd.
	All rights reserved.

FILE:
	rsatest.cpp

DESC:
	implementation of the RSATest class.
[]=========================================================================[]
*/

#include "RSATest.h"
#include "common.h"
RSATest::RSATest(char* dll_file_path):CBaseAll(dll_file_path)
{
}

RSATest::~RSATest()
{
}

void RSATest::RsaKeyGenerationTest(void)
{
	if(CKR_OK != BaseAllStart())
		return;
	
	SHOW_INFO("\n\nThe next demo will generate a RSA key pair on UsbToken.")
	CK_RV rv = CKR_OK;
	CK_BBOOL bTrue = TRUE;
	CK_ULONG ulModulusBits = MODULUS_BIT_LENGTH;
	CK_BYTE subject[] = "Sample RSA Key Pair";
	CK_ULONG keyType = CKK_RSA;
	
	CK_OBJECT_HANDLE hPubKey = 0;
	CK_OBJECT_CLASS pubClass = CKO_PUBLIC_KEY;
	CK_ATTRIBUTE pubTemplate[] =
	{ 
		{CKA_CLASS,			&pubClass,		sizeof(pubClass)},
		{CKA_KEY_TYPE,		&keyType,		sizeof(keyType)},
		{CKA_SUBJECT,		subject,		sizeof(subject)},
		{CKA_MODULUS_BITS,	&ulModulusBits, sizeof(ulModulusBits)},
		{CKA_ENCRYPT,		&bTrue,			sizeof(bTrue)},
		{CKA_TOKEN,			&bTrue,			sizeof(bTrue)},
		{CKA_WRAP,			&bTrue,			sizeof(bTrue)},
	};	
	CK_OBJECT_HANDLE hPriKey		= 0;
	CK_OBJECT_CLASS priClass	= CKO_PRIVATE_KEY;
	CK_ATTRIBUTE priTemplate[] = {
		{CKA_CLASS,			&priClass,	sizeof(priClass)},
		{CKA_KEY_TYPE,		&keyType,	sizeof(keyType)},
		{CKA_SUBJECT,		subject,	sizeof(subject)},
		{CKA_DECRYPT,		&bTrue,		sizeof(bTrue)},
		{CKA_PRIVATE,		&bTrue,		sizeof(bTrue)},
		{CKA_SENSITIVE,		&bTrue,		sizeof(bTrue)},
		{CKA_TOKEN,			&bTrue,		sizeof(bTrue)},
		{CKA_EXTRACTABLE,	&bTrue,		sizeof(bTrue)},
		{CKA_UNWRAP,		&bTrue,		sizeof(bTrue)},
	};
    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[MODULUS_BIT_LENGTH] = {0};
	CK_ULONG ulSignatureLen = 0;
	CK_BYTE_PTR pbCipherBuffer = NULL;
	CK_ULONG ulCipherLen = 0;
	CK_BYTE_PTR pbRestoredMsg = NULL;
	CK_ULONG ulRestoredMsgLen = 0;
	
	do {
		START_OP("Generating public/private key pair...");
		rv =  m_gToken->C_GenerateKeyPair(hSession, &keyGenMechanism,
								   pubTemplate, countof(pubTemplate),
								   priTemplate, countof(priTemplate),
								   &hPubKey, &hPriKey);
		CHECK_OP(rv)
			
    	// Try Sign and Verify operations with the key pair.
		// Sign a message. 
		SHOW_INFO("\nThe message to be signed is \"")
		ShowData(pbMsg, ulMsgLen);
		
		ulSignatureLen	= sizeof(bSignatureBuffer);
		START_OP("Signing initialize.");
		rv =  m_gToken->C_SignInit(hSession, &ckMechanism, hPriKey);
		CHECK_OP(rv);
		
		START_OP("Sign the message.")
		rv =  m_gToken->C_Sign(hSession, 
							pbMsg,
							ulMsgLen, 
							bSignatureBuffer, &ulSignatureLen);
		CHECK_OP(rv)
		SHOW_INFO("\nSignature is:\n");	
		ShowData(bSignatureBuffer, ulSignatureLen);

		// Verify the previously signed message.
		START_OP("Verifying initialize.")
		rv =  m_gToken->C_VerifyInit(hSession, &ckMechanism, hPubKey);
		CHECK_OP(rv)
			
		START_OP("Verify the message.")
		rv =  m_gToken->C_Verify(hSession, 
									pbMsg, ulMsgLen, 
									bSignatureBuffer, ulSignatureLen);
		CHECK_OP(rv)
			
		// Encrypt a message. 
		SHOW_INFO("\nThe message to be Encrypt is: ")
		ShowData(pbMsg, ulMsgLen);
		
		START_OP("Encrypting initialize.")
			rv =  m_gToken->C_EncryptInit(hSession,
											&ckMechanism,
											hPubKey);
		CHECK_OP(rv)
			
		START_OP("Obtain the size of the encrypted message....")
		rv =  m_gToken->C_Encrypt(hSession, pbMsg, ulMsgLen, NULL_PTR, &ulCipherLen);
		CHECK_OP(rv)
			
		START_OP("Allocate buffer for the encrypted message.")
		pbCipherBuffer = (CK_BYTE_PTR)malloc(ulCipherLen);
		if (! pbCipherBuffer)
			CHECK_OP(CKR_HOST_MEMORY)
		else
			CHECK_OP(CKR_OK);
		
		START_OP("Encrypt the message....");
		memset(pbCipherBuffer, 0, ulCipherLen);
		rv =  m_gToken->C_Encrypt(hSession, pbMsg, ulMsgLen, 
											pbCipherBuffer, &ulCipherLen);
		CHECK_OP(rv);
		SHOW_INFO("Data encrypted: ");
		ShowData(pbCipherBuffer, ulCipherLen);
		
		START_OP("Decrypting initialize....")
		rv =  m_gToken->C_DecryptInit(hSession,
										&ckMechanism,
										hPriKey);
		CHECK_OP(rv)
			
			START_OP("Obtain the size of the decrypted message....")
			rv =  m_gToken->C_Decrypt(hSession, pbCipherBuffer, 
										ulCipherLen, NULL_PTR, &ulRestoredMsgLen);
		CHECK_OP(rv)
		START_OP("Allocate buffer for the decrypted message.")
		pbRestoredMsg = (CK_BYTE_PTR)malloc(ulRestoredMsgLen + 1);
		if (! pbRestoredMsg)
			CHECK_OP(CKR_HOST_MEMORY)
			else
			CHECK_OP(CKR_OK);
		memset(pbRestoredMsg, 0, ulRestoredMsgLen + 1);
		
		START_OP("Decrypt the message.")
		rv =  m_gToken->C_Decrypt(hSession, pbCipherBuffer, ulCipherLen, 
									pbRestoredMsg, &ulRestoredMsgLen);
		CHECK_OP(rv)
			
		// Decrypt a message. 
		SHOW_INFO("\nThe message decrypted is \n")
		ShowData(pbRestoredMsg, ulRestoredMsgLen);
		
		START_OP("Verify the message.");
		if(0 == memcmp(pbMsg, pbRestoredMsg, ulRestoredMsgLen))
			CHECK_OP(CKR_OK)
		else
			SHOW_INFO("....[FAILED]\n")
			
		// Remove the RSA key pair from the UsbToken.
		START_OP("Proceed to remove the RSA key pair from the token.");
		rv =  m_gToken->C_DestroyObject(hSession, hPubKey);
		CHECK_OP(rv);
		
		START_OP("Remove private key obeject.");
		rv =  m_gToken->C_DestroyObject(hSession, hPriKey);
		CHECK_OP(rv);
		
		START_OP("Logout normally...")
			rv = m_gToken->C_Logout(hSession);
		CHECK_OP(rv);
	}while (0);
		
		// Release the memory.
		if(pbCipherBuffer)
		{
			free(pbCipherBuffer);
			pbCipherBuffer = NULL;
		}
		if(pbRestoredMsg)
		{
			free(pbRestoredMsg);
			pbCipherBuffer = NULL;
		}
		BaseAllEnd();
}
