/*--
	algorithmTest.cpp

	Test algorithm.

	Include:
		GENERATE DES
		DERIVE DES
		RSA
--*/

#include "global.h"

#define BUF_SIZE	1024
#define DATALEN		100

/*	Generate algorithm test	*/

int GenerateAlgTest(ULONG ulALG)
{
	HCRYPTPROV hCryptProv = NULL;
	char cKeyContainer[32]= {0};
	unsigned long ulName;
	
	srand((unsigned)time(NULL));
	ulName = rand() * rand();
	ltoa(ulName, cKeyContainer, 16);

	printf("\nCSP name:%s\n", CSP_NAME);
	printf("\n\nNOTICE:Get the handle from UsbToken...\n");
	printf("CryptAcquireContext\n");

	if(!CryptAcquireContext(&hCryptProv,CONTAINER,CSP_NAME,PROV_RSA_FULL,0))
	{
		if(!CryptAcquireContext(&hCryptProv,CONTAINER,CSP_NAME,PROV_RSA_FULL,CRYPT_NEWKEYSET))
		{
			printf("\tErr: CryptAcquireContext() ERR CODE=0x%0x\n", GetLastError());
			return CSP_FAILED;
		}
		else
		{
			printf("...OK!\n");
		}
	}
	else
	{
		printf("...OK!\n");
	}
	
	//Enum algs those CSP support:

	printf("CryptGetProvParam:");
	BYTE *ptr = NULL;
	DWORD i;
	ALG_ID aiAlgid;
	DWORD dwBits;
	DWORD dwNameLen;
	CHAR szName[100] = {0}; // Often allocated dynamically
	BYTE lpbData[1000] = {0}; // Often allocated dynamically
	DWORD cbData;
	DWORD dwFlags;
	CHAR *lpszAlgType = NULL;
	
	// Enumerate the supported algorithms.
	for(i=0 ; ; i++) 
	{
		// Set the CRYPT_FIRST flag the first time through the loop.
		if(i == 0) 
			dwFlags = CRYPT_FIRST;
		else 
			dwFlags = 0;
		
		// Retrieve information about an algorithm.
		cbData = 1000;
		if(!CryptGetProvParam(hCryptProv, PP_ENUMALGS, lpbData, &cbData, dwFlags)) 
		{
			if(GetLastError() == ERROR_NO_MORE_ITEMS)
			{
				break; // Exit the loop.
			}
			else 
			{
				printf("\tErr: CryptGetProvParam() ERR CODE=0x%0x\n", GetLastError());
				printf("CryptReleaseContext:");

				if(!CryptReleaseContext(hCryptProv, 0))
					printf("Err: CryptReleaseContext() ERR CODE=0x%0x\n",GetLastError());
				else
					printf("...OK!\n");
				return CSP_FAILED;
			}
		}
		else
		{
			printf("...OK!\n");
		}

		// Extract algorithm information from the 'pbData' buffer.
		ptr = lpbData;
		aiAlgid = *(ALG_ID *)ptr;
		ptr += sizeof(ALG_ID);
		dwBits = *(DWORD *)ptr;
		ptr += sizeof(DWORD);
		dwNameLen = *(DWORD *)ptr;
		ptr += sizeof(DWORD);
		strncpy(szName, (char *)ptr, dwNameLen);
		
		// Determine the algorithm type.
		switch(GET_ALG_CLASS(aiAlgid)) 
		{
		case ALG_CLASS_DATA_ENCRYPT: lpszAlgType = "Encrypt ";
			break;
		case ALG_CLASS_HASH: lpszAlgType = "Hash ";
			break;
		case ALG_CLASS_KEY_EXCHANGE: lpszAlgType = "Exchange ";
			break;
		case ALG_CLASS_SIGNATURE: lpszAlgType = "Signature";
			break;
		default: lpszAlgType = "Unknown ";
		}
		
		// Print information about the algorithm.  
		printf("Algid:%8.8xh, Bits:%-4d, Type:%s, NameLen:%-2d, Name:%s\n",
			aiAlgid, dwBits, lpszAlgType, dwNameLen, szName);
	}

	HCRYPTKEY hKey = 0;
	printf("\nCryptGenKey:"); 
	if (!CryptGenKey(hCryptProv, ulALG, /*CRYPT_EXPORTABLE*/0,  &hKey))
	{
		printf("\tERR: CryptGenKey() ERR CODE=0x%0x\n", GetLastError());
		printf("CryptReleaseContext:");

		if(!CryptReleaseContext(hCryptProv, 0))
			printf("\tErr:CryptReleaseContext() ERR CODE=0x%0x\n",GetLastError());
		else
			printf("...OK!\n");
		return CSP_FAILED;
	}
	else
	{
		printf("...OK!\n");
		printf("hKey=0x%0x\n", hKey); 
	}

	//Encrypt Data with a key:
	BYTE lpData[BUF_SIZE] = {0};
	BYTE lpOrinData[BUF_SIZE] = {0};
	DWORD dwLen = 45;
	DWORD dwRetLen = dwLen;
	for(BYTE j = 0; j < 255; ++j)
		lpData[j] = j;
	memcpy(lpOrinData, lpData, dwLen);


	//encrypt:
	printf("Data to be encrypted:\n"); 
	ShowData(lpData, dwLen);
	printf("\nCryptEncrypt:");

	//Get size:
	CryptEncrypt(hKey, 0, TRUE, 0, NULL, &dwRetLen, BUF_SIZE);
	
	//Get encrypted value:
	if(!CryptEncrypt(hKey, 0, TRUE, 0, lpData, &dwLen, BUF_SIZE))
	{
		printf("\tERR: CryptEncrypt() ERR CODE=0x%0x\n", GetLastError());
		printf("CryptReleaseContext:");

		if(!CryptReleaseContext(hCryptProv, 0))
			printf("\tErr: CryptReleaseContext() ERR CODE=0x%0x\n",GetLastError());
		else
			printf("...OK!\n");
		return CSP_FAILED;
	}
	else
	{
		printf("...OK!\n");
		printf("Data after encrypted:\n");
		ShowData(lpData, dwLen);
	}

	//Decrypt Data with a key:
	printf("\nCryptDecrypt:");
	if(!CryptDecrypt(hKey, 0, TRUE, 0, lpData, &dwLen))
	{
		printf("\tERR: CryptDecrypt() ERR CODE=0x%0x\n", GetLastError());
		printf("CryptReleaseContext:");

		if(!CryptReleaseContext(hCryptProv, 0))
			printf("\tErr: CryptReleaseContext() ERR CODE=0x%0x\n",GetLastError());
		else
			printf("...OK!\n");
		return CSP_FAILED;
	}
	else
	{
		printf("...OK!\n");
		ShowData(lpData, dwLen);
	}

	if(0 != memcmp(lpData, lpOrinData, dwLen))
	{
		printf("\n..... Decrypted data != Orginal Data.  \n");
	}
	else
	{
		printf("\nEncrypt and Decrypt are OK!\n");
	}

	printf("CryptReleaseContext:");
	if(!CryptReleaseContext(hCryptProv, 0))
		printf("\tErr: CryptReleaseContext() ERR CODE=0x%0x\n",GetLastError());
	else
		printf("...OK!\n");
	return CSP_OK;
}


/*	Device algorithm test	*/

int DeviceAlgTest(ULONG ulALG)
{
	HCRYPTPROV hCryptProv = NULL;
	HCRYPTKEY hKey = 0;
	HCRYPTHASH hHash = 0;
	DWORD dwLength;
	char* szPassword = "UsbToken Derive key test!";
	char cKeyContainer[32] = {0};
	unsigned long ulName;

	srand((unsigned)time(NULL));
	ulName = rand() * rand();
	ltoa(ulName, cKeyContainer, 16);

	printf("\n\nNOTICE:Get the handle from UsbToken...\n");
	printf("CryptAcquireContext\n");
	printf("\nCSP name:%s\n", CSP_NAME);

	if(!CryptAcquireContext(&hCryptProv,CONTAINER,CSP_NAME,PROV_RSA_FULL,0))
	{
		if(!CryptAcquireContext(&hCryptProv,CONTAINER,CSP_NAME,PROV_RSA_FULL,CRYPT_NEWKEYSET))
		{
			printf("\tERR: CryptAcquireContext() ERR CODE=0x%0x\n", GetLastError());
			return CSP_FAILED;
		}
		else
		{
			printf("...OK!\n");
		}
	}
	else
	{
		printf("...OK!\n");
	}

	// Create a hash object.
	printf("CryptCreateHash\n");
	if(!CryptCreateHash(hCryptProv, CALG_MD5, 0, 0, &hHash)) 
	{
		printf("ERR: Error %x during CryptCreateHash!\n", GetLastError());
		printf("CryptReleaseContext:");

		if(!CryptReleaseContext(hCryptProv, 0))
			printf("Err: CryptReleaseContext() ERR CODE=0x%0x\n",GetLastError());
		else
			printf("...OK!\n");
		return CSP_FAILED;
	}
	else
	{
		printf("...OK!\n");
	}
	
	// Hash the password string.
	dwLength = strlen(szPassword);
	if(dwLength%8 != 0)
		dwLength = (dwLength/8 + 1)*8;

	printf("CryptHashData\n");

	if(!CryptHashData(hHash, (BYTE *)szPassword, dwLength, 0)) 
	{
		printf("ERR: Error %x during CryptHashData!\n", GetLastError());
		printf("CryptReleaseContext:");
		
		if(!CryptReleaseContext(hCryptProv, 0))
			printf("Err: CryptReleaseContext() ERR CODE=0x%0x\n",GetLastError());
		else
			printf("...OK!\n");
		return CSP_FAILED;
	}
	else
	{
		printf("...OK!\n");
	}
	
	// Create a block cipher session key based on the hash of the password.
	printf("CryptDeriveKey\n");
	if(!CryptDeriveKey(hCryptProv, ulALG, hHash, CRYPT_EXPORTABLE, &hKey))
	{
		printf("ERR: Error %x during CryptDeriveKey!\n", GetLastError());
		printf("CryptReleaseContext:");

		if(!CryptReleaseContext(hCryptProv, 0))
			printf("Err: CryptReleaseContext() ERR CODE=0x%0x\n",GetLastError());
		else
			printf("...OK!\n");
		return CSP_FAILED;
	}
	else
	{
		printf("...OK!\n"); 
		printf("\n Derive hKey= 0x%0x\n", hKey);
	}

	//Encrypt Data with key:
	BYTE lpData[1024] = "UsbToken: Encrypt Data with the derive key (des) /!\0";
	DWORD dwLen = 60;

	printf("Data to be encrypted:\n"); 
	printf("\t\t%s\n", lpData);
	printf("Data Length= %d", dwLen);

	//encrypt:
	printf("\nCryptEncrypt:");
	if(!CryptEncrypt(hKey, 0, TRUE, 0, lpData, &dwLen, 1024))
	{
		printf("ERR: CryptEncrypt() failed: ERR CODE=0x%0x\n", GetLastError());
		printf("CryptReleaseContext:");

		if(!CryptReleaseContext(hCryptProv, 0))
			printf("Err: CryptReleaseContext() ERR CODE=0x%0x\n",GetLastError());
		else
			printf("...OK!\n");
		return CSP_FAILED;
	}
	else
	{
		printf("...OK!\n");
		printf("Data after encrypted:\n");
		printf("\t\t%s\n", lpData);
		printf("Data Length= %d\n", dwLen);
	}

	//Decrypt Data with key:
	printf("\nCryptDecrypt:");
	if(!CryptDecrypt(hKey, 0, TRUE, 0, lpData, &dwLen))
	{
		printf("ERR: CryptDecrypt() failed: ERR CODE=0x%0x\n", GetLastError());
		printf("CryptReleaseContext:");

		if(!CryptReleaseContext(hCryptProv, 0))
			printf("Err: CryptReleaseContext() ERR CODE=0x%0x\n",GetLastError());
		else
			printf("...OK!\n");
		return CSP_FAILED;
	}
	else
	{
		printf("...OK!\n");
		printf("Decrypt data:\n\t\t%s\n", lpData);
	}

	//delete pDataBuffer;
	printf("CryptReleaseContext:");
	if(!CryptReleaseContext(hCryptProv, 0))
		printf("Err: CryptReleaseContext() ERR CODE=0x%0x\n",GetLastError());
	else
		printf("...OK!\n");
	return CSP_OK;
}