
#include "Common.h"
#include "DHKeyGenTest.h"
#include "DSATest.h"
#include "RSATest.h"
#include "RC2Test.h"
#include "RC4Test.h"
#include "DesTest.h"
#include "Des3Test.h"
#include "SSF33Test.h"
#include <unistd.h>
using namespace std;

void * m_hSo;
CK_FUNCTION_LIST_PTR m_gToken;
CK_SLOT_ID_PTR pSlotList;
CK_SESSION_HANDLE hSession;
typedef CK_RV (* C_GETFUNCTIONLISTPROC)(CK_FUNCTION_LIST_PTR_PTR);

CK_RV BaseAllEnd()
{
	cout <<"\n[ENTER:CBaseAll::BaseAllEnd:Close session and free the PKCS#11's dll]"<<endl;
	CK_RV rv = CKR_OK;
	cout << "Release the session handle..." << endl;
	if(m_gToken)
	{
		rv = m_gToken->C_CloseSession(hSession);
		if(CKR_OK != rv)
		{
			cout << "....C_CloseSession failed!" << endl;
			return rv;
		}
		rv = m_gToken->C_Finalize(0);
		if(CKR_OK != rv)
		{
			cout << "...C_Finalize failed!" << endl;
			return rv;
		}
		m_gToken = NULL_PTR;
	}
	usleep(500);
	if(m_hSo != NULL_PTR)
	{
		dlclose(m_hSo);
		m_hSo = NULL;
	}
	cout << "[EXIT:CBaseAll::BaseAllEnd()]" <<endl;
	return CKR_OK;
}

CK_RV BaseAllStart()
{
	cout <<"\n[ENTER: CBaseAll::BaseAllStart:LoadLibrary PKCS#11's dll and initialize it!]" <<endl;
	CK_RV rv = CKR_OK;
	m_hSo = dlopen("libepsng_p11.so",RTLD_NOW);
	if(m_hSo == NULL)
	{
		cout << "......LoadLibrary [" << "libepsng_p11.so"<< "] failed!" << endl;
		return CKR_HOST_MEMORY;
	}

	//Get the list of the PKCS#11 function entries.
	C_GETFUNCTIONLISTPROC pC_GetFunctionList = (C_GETFUNCTIONLISTPROC)dlsym(m_hSo,"C_GetFunctionList");
	if(pC_GetFunctionList == NULL)
	{
		cout << "...Err: (pC_GetFunctionList == NULL)!" << endl;
		return CKR_HOST_MEMORY;
	}

	//Get PKCS#11 function list.
	if(CKR_OK != (rv = pC_GetFunctionList(&m_gToken)))
	{
		cout << "...Get PKCS#11 function list failed!" << endl;
		return rv;
	}
	cout << "Initialize the UsbToken PKCS#11 library..." << endl;
	rv = m_gToken->C_Initialize(NULL_PTR);
	if((CKR_OK != rv) && (CKR_CRYPTOKI_ALREADY_INITIALIZED != rv))
	{
		cout << "...C_Initialize failed! Return Error Code=" << rv << endl;
		return rv;
	}

	// Get the number of slots with a token attached.
	cout << "Get the slots information..." << endl;
	unsigned long ulCount;
	rv = m_gToken->C_GetSlotList(TRUE, NULL_PTR, &ulCount);
	if(CKR_OK != rv)
	{
		cout << "...C_GetSlotList failed! Return Error Code=" << rv << endl;
		return rv;
	}
	if(ulCount <= 0)
	{
		cout << "...No slot with token attached." << endl;
		return CKR_TOKEN_NOT_PRESENT;
	}

	// Allocate memory buffer for the slot list.
	cout << "Allocate memory for slots..." << endl;
	pSlotList = (CK_SLOT_ID_PTR) malloc(ulCount * sizeof(CK_SLOT_ID));
	if(!pSlotList)
	{
		cout << "...No slots!" << endl;
		return CKR_HOST_MEMORY;
	}
	else
	{
		cout << "Get " << ulCount << " slot attached to USB port!" << endl;
	}
	// Okay, we have the buffer allocated. Let's get the slot list.
	cout << "Get the slots information." << endl;
	rv = m_gToken->C_GetSlotList(TRUE, pSlotList, &ulCount);
	if(CKR_OK != rv)
	{
		cout << "........C_GetSLotList failed! Return Error Code=" << rv
			 << endl;
		CK_RV rv1 = m_gToken->C_Finalize(0);
		if(CKR_OK != rv1)
		{
			cout << "..........C_Finalize failed!" << endl;
			return rv1;
		}
		m_gToken = NULL_PTR;
		usleep(500);
		if(m_hSo)
		{
			dlclose(m_hSo);
			m_hSo = NULL;
		}
		return rv;
	}
	// Open a session to communicate with the token.
	cout << "Open a session to communicate with the UsbToken..." << endl;
	rv = m_gToken->C_OpenSession(pSlotList[0], CKF_RW_SESSION | CKF_SERIAL_SESSION, NULL_PTR, NULL_PTR, &hSession);
	if(CKR_OK != rv)
	{
		cout << "....C_OpenSession failed! Return Error Code=" << rv << endl;
		BaseAllEnd();
		return rv;
	}

	cout << "Input user pin to login first:" << endl;
	string m_ucPin;
	cin >> m_ucPin;
	rv = m_gToken->C_Login(hSession, CKU_USER, (CK_BYTE_PTR)&m_ucPin[0], m_ucPin.length()); 
	if(rv != CKR_OK)
	{
		cout << "...C_Login failed!" << rv << endl;
		BaseAllEnd();
		return rv;
	} 
	cout << "...C_Login OK!" << endl;
	return rv;
}

void ShowTitle(void)
{
	printf("\n");
	printf("[]==================================================[]\n");
	printf(" |             UsbToken PKCS#11 Demo               |\n");				
	printf("[]==================================================[]\n");
}

int main(int argc, char* argv[])
{
	char iTestItem ;
	if(CKR_OK != BaseAllStart())
	{
		exit(0);
	}
	for(;;)
	{
	ShowTitle();
	if(argv[1] == NULL)
	{
		cout<<"1: DES, 2: DES3, 3: RC2, 4:RC4, 5:RSA, 6: DSA, 7: DH, 8: SSF33"<<endl;
		//cout<<"1: DES, 2: DES3, 3: RC2, 4:RC4, 5:RSA, 6: DSA, 7: DH"<<endl;
		cout<< "0: Exit."<<endl;
		cout<< "Please select a testing item: ";
		string s;
		cin >> s;
		if((!cin) || (s.length() > 1))
		{
			cout<<"0 -- 8 only"<<endl;
			cin.clear();
			continue;
		}
		iTestItem = s[0];
	}
	else
	{
		char *temp = NULL;
		char c, *s = NULL ;
		temp = strdup(argv[1]) ;
		argv[1] = NULL;
		s = temp ;
		while(*s)
		{
			c = toupper(*s++) ;
			*(s-1) = c ;
		}
		cout << temp<<endl;
		if(0==strcmp(temp, "-DES"))
			iTestItem = '1';
		else if(0==strcmp(temp, "-DES3"))
			iTestItem = '2';
		else if(0==strcmp(temp, "-RC2"))
			iTestItem = '3';
		else if(0==strcmp(temp, "-RC4"))
			iTestItem = '4';
		else if(0==strcmp(temp, "-RSA"))
			iTestItem = '5';
		else if(0==strcmp(temp, "-DSA"))
			iTestItem = '6';
		else if(0==strcmp(temp, "-DH"))
			iTestItem = '7';
		else if(0==strcmp(temp, "-SSF33"))
			iTestItem = '8'; 
		else 
			iTestItem = '0';
	}
	switch(iTestItem)
	{
	case '1':
		{
			DesTest des;
			des.Test();
		}
		break;
	case '2':
		{
			Des3Test des3;	
			des3.Test();
		}
		break;
	case '3':
		{
			RC2Test rc2;
			rc2.Test();
		}
		break;
	case '4':
		{
			RC4Test rc4;
			rc4.Test();
		}
		break;
	case '5':
		{
			RSATest rsa;
			rsa.RsaKeyGenerationTest();
		}
		break;
	case '6':
		{
			DSATest dsa;
			dsa.TestDSA();
		}
		break;
	case '7':
		{
			DHKeyGenTest dh;
			dh.DHKeyTest();
		}
		break;
	case '8':
		{
			Ssf33Test ssf33;
			ssf33.Test();
		}
		break;
	case '0':
		{
			printf("\n\n Exit!");
		}
		return 0;
	default:
		printf("0-8 only \n");
		break;
	}
	}
	BaseAllEnd();
	return 0;
}
