
#include "baseall.h"
#include "es_ssl/pkcs11t.h"

#define MODULUS_BIT_LENGTH 1024

CP11TestBase::CP11TestBase(char *libFilePath) :
	m_hDll(NULL),
	m_pSlotList(NULL_PTR)
{
	memset(m_libFileName, 0, MAX_PATH);
	memcpy(m_libFileName, libFilePath, lstrlen(libFilePath));
	memset(m_ucPin, 0, 32);

	m_hStdErr = GetStdHandle(STD_ERROR_HANDLE);
	m_hStdOut = GetStdHandle(STD_OUTPUT_HANDLE);
}

CP11TestBase::~CP11TestBase()
{
	if(m_pSlotList != NULL_PTR)
	{
		free(m_pSlotList);
		m_pSlotList = NULL_PTR;
	}
	if(m_hDll != NULL_PTR)
	{
		FreeLibrary(m_hDll);
		m_hDll = NULL;
	}
}
CK_RV CP11TestBase::Initialize(void)
{
	CK_RV rv = CKR_OK;

	printf("\nLoad PKCS#11 Library!\n");
	m_hDll = LoadLibrary(m_libFileName);
	if(m_hDll == NULL)
	{
		printf("    ERROR: LoadLibrary(%s) failed. [SysErrCode: 0x%08X]\n",
			m_libFileName, GetLastError());
		return CKR_GENERAL_ERROR;
	}

	//Get the list of the PKCS#11 function entries.
	CK_C_GetFunctionList pC_GetFunctionList = (CK_C_GetFunctionList)
			GetProcAddress(m_hDll, "C_GetFunctionList");
	if(pC_GetFunctionList == NULL)
	{
		printf("    ERROR: Can not locate PKCS#11 function list! [SysErrCode: 0x%08X]\n", GetLastError());
		Finalize();
		return CKR_GENERAL_ERROR;
	}

	//Get PKCS#11 function list.
	rv = pC_GetFunctionList(&m_p11);
	if(!_CheckRV("Get PKCS#11 function list", rv))
	{
		Finalize();
		return rv;
	}

	rv = m_p11->C_Initialize(NULL_PTR);
	if(!_CheckRV("C_Initialize()", rv))
	{
		Finalize();
		return rv;
	}

	// Get the number of slots with a token attached.
	unsigned long ulCount;
	rv = m_p11->C_GetSlotList(TRUE, NULL_PTR, &ulCount);
	if(!_CheckRV("C_GetSlotList(TRUE, NULL)", rv))
	{
		Finalize();
		return rv;
	}

	if(ulCount <= 0)
	{
		printf("    WARNING: No slot with token attached.\n");
		Finalize();
		return CKR_TOKEN_NOT_PRESENT;
	}

	// Allocate memory buffer for the slot list.
	m_pSlotList = (CK_SLOT_ID_PTR) malloc(ulCount * sizeof(CK_SLOT_ID));
	if(!m_pSlotList)
	{
		printf("    ERROR: Not enough memory!\n");
		Finalize();
		return CKR_HOST_MEMORY;
	}
	else
	{
		printf("There are %d slot%s attached to USB port.\n", ulCount, ulCount > 1 ? "s" :"");
	}

	// Okay, we have the buffer allocated. Let's get the slot list.
	//cout << "Get the slots information." << endl;
	rv = m_p11->C_GetSlotList(TRUE, m_pSlotList, &ulCount);
	if(!_CheckRV("C_GetSlotList(TRUE, buffer)", rv))
	{
		Finalize();
		return rv;
	}

	// Open a session to communicate with the token.
	rv = m_p11->C_OpenSession(m_pSlotList[0],
		 CKF_RW_SESSION | CKF_SERIAL_SESSION,
		 NULL_PTR, NULL_PTR, &m_hSession);
	if(!_CheckRV("C_OpenSession()", rv))
	{
		Finalize();
		return rv;
	}
    //Ìí¼Ó
	CK_TOKEN_INFO tokenInfo = {0};
    rv = m_p11->C_GetTokenInfo(m_pSlotList[0], &tokenInfo);
	if (CKR_OK != rv)
	{
		printf("\nCan not get token information!");
		return rv;
	}
	
	if(tokenInfo.flags & CKF_PROTECTED_AUTHENTICATION_PATH)
	{
		printf("\nPlease touch the sensor to verify your fingerprint in 8 seconds: ");
		rv = m_p11->C_Login(m_hSession, CKU_USER, 0, 0);
	}
	else
	{
		printf("\nPlease enter your User-PIN: ");
		scanf("%s", m_ucPin);
      	rv = m_p11->C_Login(m_hSession, CKU_USER, (unsigned char*)m_ucPin, lstrlen(m_ucPin)); 
	}

	if(!_CheckRV("C_Login(CKU_USER)", rv))
	{
		Finalize();
		return rv;
	} 
	return rv;
}

CK_RV CP11TestBase::Finalize(void)
{
	printf("\nFinalize the PKCS#11's library.\n");

	CK_RV rv = CKR_OK;
	if(m_hDll)
	{
		if(m_p11)
		{
			if(m_hSession)
			{
				rv = m_p11->C_CloseSession(m_hSession);
				_CheckRV("C_CloseSession()", rv);
			}

			rv = m_p11->C_Finalize(0);
			_CheckRV("C_Finalize()", rv);
			m_p11 = NULL_PTR;
		}

		FreeLibrary(m_hDll);
		m_hDll = NULL;
	}

	return CKR_OK;
}

bool CP11TestBase::_CheckRV(char *pInfo, CK_RV rv)
{
	printf("%s ... ", pInfo);

	if(rv == CKR_OK)
	{
		if(NULL != m_hStdOut)
			SetConsoleTextAttribute(m_hStdOut, FOREGROUND_INTENSITY|FOREGROUND_GREEN);

		printf("OK");

		if(NULL != m_hStdOut)
			SetConsoleTextAttribute(m_hStdOut, FOREGROUND_GREEN|FOREGROUND_BLUE|FOREGROUND_RED);

		printf("\n");
		return true;
	}
	else
	{
		if(NULL != m_hStdErr)
			SetConsoleTextAttribute(m_hStdErr, FOREGROUND_INTENSITY|FOREGROUND_RED);

		printf("FAILED");

		if(NULL != m_hStdErr)
			SetConsoleTextAttribute(m_hStdErr, FOREGROUND_GREEN|FOREGROUND_BLUE|FOREGROUND_RED);

		printf(" [ErrCode: 0x%08X]\n", rv);

		return false;
	}
}


void CP11TestBase::GenerateKey_AES()
{
	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_MECHANISM mechanism = {CKM_AES_KEY_GEN, NULL_PTR, 0};
	CK_ATTRIBUTE Aestem[] = {
		{CKA_CLASS, &oClass, sizeof(CK_OBJECT_CLASS)},
		{CKA_KEY_TYPE, &keyType, sizeof(CK_KEY_TYPE)},
		{CKA_TOKEN, &bTrue, sizeof(CK_BBOOL)},
		{CKA_PRIVATE, &bTrue, sizeof(CK_BBOOL)},
		{CKA_ENCRYPT, &bTrue, sizeof(CK_BBOOL)},
		{CKA_DECRYPT, &bTrue, sizeof(CK_BBOOL)},
		{CKA_VALUE_LEN, &ulLen, sizeof(CK_ULONG)}, 
		{CKA_WRAP,		&bTrue,	sizeof(bTrue)},
		{CKA_UNWRAP,	&bTrue,	sizeof(bTrue)},
	};
	
	CK_RV rv = m_p11->C_GenerateKey(m_hSession, &mechanism, Aestem, sizeof(Aestem)/sizeof(CK_ATTRIBUTE), &m_hKey); 
	m_hKey_AES = m_hKey;
	_CheckRV("Generate AES key", rv);
}

CK_RV CP11TestBase::EncrptPkcs8RsaPrivate(CK_OBJECT_HANDLE aesHandle, unsigned char* privatePlain, CK_ULONG pPlainLen, unsigned char* privateEncrypted, CK_ULONG* pEncryptedLen)
{
	CK_RV rv = CKR_OK;
	//ecnrypt init:
	CK_BYTE iv[16] = {'*','2','1','0','4','z','y','b','*','2','1','0','4','z','y','b'};
	CK_MECHANISM ckMechanism = {CKM_AES_CBC_PAD, iv, 16};

	rv =  m_p11->C_EncryptInit(m_hSession, &ckMechanism, aesHandle); 
	if(!_CheckRV("C_EncryptInit()", rv))
	{
		return rv;
	}

	rv =  m_p11->C_Encrypt(m_hSession, privatePlain, pPlainLen, privateEncrypted, pEncryptedLen);
	if(CKR_OK == rv)
	{
		// failed
	}
	_CheckRV("EncrptPkcs8RsaPrivate", rv);
	return rv;
}

CK_OBJECT_HANDLE CP11TestBase::GetAesKeyHandle()
{
	return m_hKey_AES;
}


void CP11TestBase::CreateObject_rsaPub(unsigned char* pN, int lenN, unsigned char* pE, int lenE)
{
	CK_RV rv = CKR_OK;
	CK_BBOOL bTrue = TRUE;
	CK_BBOOL bFalse = FALSE;
	CK_ULONG ulModulusBits = MODULUS_BIT_LENGTH;
	CK_BYTE subject[] = "Sample RSA Key Pair";
	CK_ULONG keyType = CKK_RSA;

	CK_OBJECT_HANDLE hPubKey = NULL_PTR;
	CK_OBJECT_CLASS pubClass = CKO_PUBLIC_KEY;

	char keyID[MAX_PATH] = {0};
	sprintf(keyID, "%s1", "container name");


	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_LABEL,"nolabel",(CK_ULONG)strlen("nolabel")},
		{CKA_PRIVATE,&bFalse,sizeof(CK_BBOOL)},
		{CKA_LOCAL,&bFalse,sizeof(CK_BBOOL)},
		{CKA_VERIFY, &bTrue, sizeof(CK_BBOOL)},
		{CKA_MODULUS, &pN[0], lenN},
		{CKA_PUBLIC_EXPONENT, &pE[0], lenE},
		{CKA_ID, keyID, sizeof(keyID)},
		{CKA_WRAP,			&bTrue,			sizeof(bTrue)},
		{CKA_UNWRAP,		&bTrue,			sizeof(bTrue)},
	};	

	rv = m_p11->C_CreateObject(m_hSession,pubTemplate,sizeof(pubTemplate)/sizeof(CK_ATTRIBUTE),&hPubKey);
	m_hKey_rsaPub = hPubKey;
	
	if(!_CheckRV("C_CreateObject() rsaPub", rv))
		return;
	
}


CK_OBJECT_HANDLE CP11TestBase::GetRsaPubKeyHandle()
{
	return m_hKey_rsaPub;
}


CK_RV CP11TestBase::RsaWrapAes(CK_OBJECT_HANDLE rsaPubHandle, CK_OBJECT_HANDLE aesHandle, unsigned char* wrappedKey, CK_ULONG* pWrappedKeyLen)
{
	
	CK_OBJECT_HANDLE hWrappingKey = rsaPubHandle;
	CK_OBJECT_HANDLE hKey = aesHandle;

	CK_MECHANISM mechanism = {
		CKM_RSA_PKCS, NULL_PTR, 0
	};


	CK_RV rv;	
	rv = m_p11->C_WrapKey(
		m_hSession, &mechanism,
		hWrappingKey, hKey,
		wrappedKey, pWrappedKeyLen);
	if (rv == CKR_OK)
	{

	}


	return rv;
}

CK_RV CP11TestBase::AesUnWrapRsa(CK_OBJECT_HANDLE aesHandle,
								 CK_BYTE_PTR wrappedKey,
								 CK_ULONG wrappedLen,
								 CK_OBJECT_HANDLE_PTR pPrivateHanle,
								 CK_BYTE_PTR dataN, 
								 CK_ULONG lenN, 
								 CK_BYTE_PTR dataE, 
								 CK_ULONG lenE, 
								 CK_BYTE_PTR dataD, 
								 CK_ULONG lenD, 
								 CK_BYTE_PTR dataP, 
								 CK_ULONG lenP, 
								 CK_BYTE_PTR dataQ, 
								 CK_ULONG lenQ, 
								 CK_BYTE_PTR dataDP,
								 CK_ULONG lenDP, 
								 CK_BYTE_PTR dataDQ,
								 CK_ULONG lenDQ, 
								 CK_BYTE_PTR dataQP, 
								 CK_ULONG lenQP)
{
	
	// private key template
	CK_OBJECT_CLASS prvPKCls = CKO_PRIVATE_KEY;
	CK_BBOOL bTrue = CK_TRUE;
	CK_BBOOL bFalse = CK_FALSE;
	CK_KEY_TYPE ktPK = CKK_RSA;

	CK_RV rv = CKR_OK;

	char keyID[MAX_PATH] = {0};
	sprintf(keyID, "%s2", "container name");


	CK_ATTRIBUTE prikeyTemp[] = 
	{
// 		{CKA_CLASS,&prvPKCls,sizeof(CK_OBJECT_CLASS)},
		{CKA_TOKEN,&bTrue,sizeof(CK_BBOOL)},
		{CKA_PRIVATE,&bTrue,sizeof(CK_BBOOL)},
		{CKA_SENSITIVE,&bTrue,sizeof(CK_BBOOL)},
		{CKA_MODIFIABLE,&bFalse, sizeof(CK_BBOOL)},
		{CKA_LABEL,"nolabel",(CK_ULONG)strlen("nolabel")},
// 		{CKA_KEY_TYPE,&ktPK,sizeof(CK_KEY_TYPE)},
// 		{CKA_LOCAL,&bFalse,sizeof(CK_BBOOL)},
		{CKA_SIGN, &bTrue, sizeof(CK_BBOOL)},
		{CKA_DECRYPT, &bTrue, sizeof(CK_BBOOL)},
// 		{CKA_MODULUS, &dataN[0], lenN},
// 		{CKA_PUBLIC_EXPONENT, &dataE[0], lenE},
// 		{CKA_PRIVATE_EXPONENT, &dataD[0], lenD},
// 		{CKA_PRIME_1, &dataP[0], lenP},//p
// 		{CKA_PRIME_2, &dataQ[0], lenQ},//p
// 		{CKA_EXPONENT_1, &dataDP[0],lenDP},//dp
// 		{CKA_EXPONENT_2, &dataDQ[0],lenDQ},//dq
// 		{CKA_COEFFICIENT, &dataQP[0], lenQP},//qiv
		{CKA_ID, keyID, sizeof(keyID)},
// 		{CKA_WRAP,			&bTrue,			sizeof(bTrue)},
		{CKA_UNWRAP,		&bTrue,			sizeof(bTrue)},
		{CKA_EXTRACTABLE, &bTrue,			sizeof(bTrue)}
	};


	
	CK_OBJECT_HANDLE hUnwrappingKey = aesHandle;
	CK_OBJECT_HANDLE hKey;



	//decrypt init:
	CK_BYTE iv[16] = {'*','2','1','0','4','z','y','b','*','2','1','0','4','z','y','b'};
	CK_MECHANISM ckMechanism = {CKM_AES_CBC_PAD, iv, 16};



	rv = m_p11->C_UnwrapKey(
		m_hSession, &ckMechanism, hUnwrappingKey,
		wrappedKey, wrappedLen, prikeyTemp, sizeof(prikeyTemp)/sizeof(CK_ATTRIBUTE), &hKey);
	if (rv == CKR_OK) {
	
		*pPrivateHanle = hKey;
	}
	else
	{
		_CheckRV("UnwrapKey", rv);
		*pPrivateHanle = NULL;
	}

	return rv;
}
CK_RV CP11TestBase::SignTest(CK_OBJECT_HANDLE PrivateKeyHandle)
{
	CK_RV rv = CKR_OK;
	CK_MECHANISM ckMechanism = {CKM_RSA_PKCS, NULL_PTR, 0};
	CK_BYTE pbMsg[] = "UsbToken RunRsaKeyGenerationTest...";
	CK_ULONG ulMsgLen = lstrlen((char *)pbMsg);
	CK_BYTE bSignatureBuffer[MODULUS_BIT_LENGTH] = {0};
	CK_ULONG ulSignatureLen = MODULUS_BIT_LENGTH;

	rv =  m_p11->C_SignInit(m_hSession, &ckMechanism, PrivateKeyHandle);
	if(!_CheckRV("C_SignInit()", rv))
		return rv;

	rv =  m_p11->C_Sign(m_hSession, pbMsg, ulMsgLen, bSignatureBuffer, &ulSignatureLen);
	if(!_CheckRV("C_Sign()", rv))
		return rv;

	return CKR_OK;
}
