/*--
	SymmetricKeyEncryption.cpp

	2015-10-27

	Test PKCS#11 Generate DES/TDES/AES128/AES192/AES256 Key and encryption/decryption
--*/


#include "../global.h"

CK_FUNCTION_LIST_PTR fl = 0;

CK_OBJECT_HANDLE find_object(CK_SESSION_HANDLE hSess, 
							 CK_ATTRIBUTE * t, 
							 CK_ULONG size);

int main(int argc, char* argv[])
{
	CK_RV rv;
	CK_ULONG slot_count = 100;
	CK_SLOT_ID slots[100];
	CK_SESSION_HANDLE session;
	CK_BBOOL         bFalse = FALSE;
	CK_BBOOL         bTrue = TRUE;
	CK_KEY_TYPE      key_type[3]  = {CKK_DES, CKK_DES3, CKK_AES};
	CK_ULONG ulKeyLen[5] = {8, 24, 16, 24, 32};
	CK_ULONG i = 0;
	/* setup key attributes, specific name so we can find it 
	 * later, symmetric key type, saved on the UniMate/UniToken and private to only the
	 * users of the UniMate/UniToken
	 */
	CK_OBJECT_CLASS class_key = CKO_SECRET_KEY;
	CK_BYTE IdGen = '0';
	CK_BYTE IdImport = '1';

	CK_BYTE lpKeyValue_8[] = {
		0xcc, 0x63, 0x98, 0x02, 0xd9, 0x81, 0xac, 0xb4
	};

	CK_BYTE lpKeyValue_24[] = {
		0xCC, 0x63, 0x98, 0x02, 0xD9, 0x81, 0xAC, 0xB4,
		0x34, 0x60, 0x85, 0x0A, 0xCF, 0xF6, 0xCD, 0x54,
		0xCC, 0x63, 0x98, 0x02, 0xD9, 0x81, 0xAC, 0xB4
	};

	CK_BYTE lpKeyValue_16[] = {
		0xCC, 0x63, 0x98, 0x02, 0xD9, 0x81, 0xAC, 0xB4,
		0x34, 0x60, 0x85, 0x0A, 0xCF, 0xF6, 0xCD, 0x54
	};

	CK_BYTE lpKeyValue_32[] = {
		0xCC, 0x63, 0x98, 0x02, 0xD9, 0x81, 0xAC, 0xB4,
		0x34, 0x60, 0x85, 0x0A, 0xCF, 0xF6, 0xCD, 0x54,
		0xCC, 0x63, 0x98, 0x02, 0xD9, 0x81, 0xAC, 0xB4,
		0x34, 0x60, 0x85, 0x0A, 0xCF, 0xF6, 0xCD, 0x54
	};

	CK_BYTE lpData[] = {
		0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08,
		0x09, 0x0a, 0x0b, 0x0c, 0x0d, 0x0e, 0x0f, 0x10
	};
	CK_ULONG dataLen = sizeof(lpData);

	CK_BYTE en_data[1024] = {0};
	CK_ULONG en_len = 1024;
	CK_BYTE de_data[1024] = {0};
	CK_ULONG de_len = 1024;

	CK_BYTE lpIV[16] = {
		0x01, 0x01, 0x01, 0x01, 0x01, 0x01, 0x01, 0x01,
		0x01, 0x01, 0x01, 0x01, 0x01, 0x01, 0x01, 0x01
	};
	CK_ULONG ivLen = sizeof(lpIV);

	CK_ATTRIBUTE Template_A[] =	{
		{CKA_ID,		&IdGen,			sizeof(CK_BYTE)},
		{CKA_CLASS,		&class_key,		sizeof(class_key)},
		{CKA_KEY_TYPE,	&key_type[0],	sizeof(CK_KEY_TYPE)},
		{CKA_PRIVATE,	&bTrue,			sizeof(bTrue)},
		{CKA_TOKEN,		&bTrue,			sizeof(bTrue)},
		{CKA_ENCRYPT,	&bTrue,			sizeof(bTrue)},
		{CKA_DECRYPT,	&bTrue,			sizeof(bTrue)},
		{CKA_SENSITIVE,	&bFalse,		sizeof(bFalse)}
	};	

	CK_ATTRIBUTE Template_AES_A[] =	{
		{CKA_ID,		&IdGen,			sizeof(CK_BYTE)},
		{CKA_CLASS,		&class_key,		sizeof(class_key)},
		{CKA_KEY_TYPE,	&key_type[2],	sizeof(CK_KEY_TYPE)},
		{CKA_VALUE_LEN, &ulKeyLen[2],	sizeof(CK_ULONG)},
		{CKA_PRIVATE,	&bTrue,			sizeof(bTrue)},
		{CKA_TOKEN,		&bTrue,			sizeof(bTrue)},
		{CKA_ENCRYPT,	&bTrue,			sizeof(bTrue)},
		{CKA_DECRYPT,	&bTrue,			sizeof(bTrue)},
		{CKA_SENSITIVE,	&bFalse,		sizeof(bFalse)}
	};

	CK_ATTRIBUTE Template_A_Import[] =	{
		{CKA_ID,		&IdImport,		sizeof(CK_BYTE)},
		{CKA_CLASS,		&class_key,		sizeof(class_key)},
		{CKA_KEY_TYPE,	&key_type[0],	sizeof(CK_KEY_TYPE)},
		{CKA_VALUE_LEN,	&ulKeyLen[0],	sizeof(CK_ULONG)},
		{CKA_PRIVATE,	&bTrue,			sizeof(bTrue)},
		{CKA_TOKEN,		&bTrue,			sizeof(bTrue)},
		{CKA_ENCRYPT,	&bTrue,			sizeof(bTrue)},
		{CKA_DECRYPT,	&bTrue,			sizeof(bTrue)},
		{CKA_SENSITIVE,	&bFalse,		sizeof(bFalse)},
		{CKA_VALUE,		lpKeyValue_8,	sizeof(lpKeyValue_8)}
	};

	/* setup mechanism to generate GOST28147 key 
	 */
	CK_MECHANISM mechanism = {0};
	CK_OBJECT_HANDLE hKeyGen = NULL;
	CK_OBJECT_HANDLE hKeyImp = NULL;
	CK_ATTRIBUTE_PTR pGenKeyTemplate = NULL;
	CK_ULONG ulAttCount = 0;
	CK_ULONG ulFindAttCount = 0;
	CK_BYTE szKeyType[10] = {0};
	char cKeySel;
	
	printf("Generate Symmetric Key Sample:\n");
	
	fl = InitFunctionList();
	if (fl == 0)
		return leave("Can't get function list. \n");
	
	/* Initialize the PKCS #11 library
	 */
	if (PKCS_OK != fl->C_Initialize (0))
		return leave("C_Initialize failed...");                       
	
	/* Get list of slots containing UniMate/UniTokens
	 */
	if (PKCS_OK != fl->C_GetSlotList (TRUE, slots, &slot_count))
		return leave("C_GetSlotList failed...");

	/* One or more UniMate/UniToken must be inserted
	 */
	if (slot_count < 1)
		return leave("No UniMate/UniToken is available.\n");
	
	printf ("UniMate/UniToken Found.\n");
	/* open a read/write session on the first UniMate/UniToken found
	 */
	if (PKCS_OK != fl->C_OpenSession(slots[0], 
		(CKF_SERIAL_SESSION | CKF_RW_SESSION), 0, 0, &session)
		)
		return leave("C_OpenSession failed.\n");
	
	/* login as user so we can use the public and private
	 * keys on the UniToken
	 */
	if (PKCS_OK != fl->C_Login(session, CKU_USER, USER_PIN, 8))
		return leave("C_Login failed.\n");
	
	/* Try to find the two keys using the find_object utility function if exist,
	 * or generate new key.
	*/
	printf("Please selected the key type:\n");
	printf("0: DES\n");
	printf("1: TDES\n");
	printf("2: AES-128\n");
	printf("3: AES-192\n");
	printf("4: AES-256\n");
	cKeySel = getchar();
	switch (cKeySel)
	{
	case '0'://DES
		memcpy(szKeyType, "DES", 3);
		Template_A[2].pValue = &key_type[0];
		
		pGenKeyTemplate = Template_A;
		ulAttCount = 8;

		mechanism.mechanism = CKM_DES_KEY_GEN;
		mechanism.pParameter = 0;
		mechanism.ulParameterLen = 0;

		Template_A_Import[2].pValue = &key_type[0];
		Template_A_Import[3].pValue = &ulKeyLen[0];
		Template_A_Import[9].pValue = lpKeyValue_8;
		Template_A_Import[9].ulValueLen = sizeof(lpKeyValue_8);
		

		// Find DES Key by CKA_ID/CKA_CLASS/CKA_KEY_TYPE
		ulFindAttCount = 3;
		break;
	case '1'://3DES
		memcpy(szKeyType, "TDES", 4);
		Template_A[2].pValue = &key_type[1];

		pGenKeyTemplate = Template_A;
		ulAttCount = 8;
		
		mechanism.mechanism = CKM_DES3_KEY_GEN;
		mechanism.pParameter = 0;
		mechanism.ulParameterLen = 0;
		
		Template_A_Import[2].pValue = &key_type[1];
		Template_A_Import[3].pValue = &ulKeyLen[1];
		Template_A_Import[9].pValue = lpKeyValue_24;
		Template_A_Import[9].ulValueLen = sizeof(lpKeyValue_24);
		
		// Find TDES Key by CKA_ID/CKA_CLASS/CKA_KEY_TYPE
		ulFindAttCount = 3;
		break;
	case '2'://AES-128
		memcpy(szKeyType, "AES-128", 7);
		Template_AES_A[2].pValue = &key_type[2];
		Template_AES_A[3].pValue = &ulKeyLen[2];

		pGenKeyTemplate = Template_AES_A;
		ulAttCount = 9;
		
		mechanism.mechanism = CKM_AES_KEY_GEN;
		mechanism.pParameter = 0;
		mechanism.ulParameterLen = 0;
		
		Template_A_Import[2].pValue = &key_type[2];
		Template_A_Import[3].pValue = &ulKeyLen[2];
		Template_A_Import[9].pValue = lpKeyValue_16;
		Template_A_Import[9].ulValueLen = sizeof(lpKeyValue_16);

		// Find DES Key by CKA_ID/CKA_CLASS/CKA_KEY_TYPE/CKA_VALUE_LEN
		ulFindAttCount = 4;
		break;
	case '3'://AES-192
		memcpy(szKeyType, "AES-192", 7);
		Template_AES_A[2].pValue = &key_type[2];
		Template_AES_A[3].pValue = &ulKeyLen[3];

		pGenKeyTemplate = Template_AES_A;
		ulAttCount = 9;
		
		mechanism.mechanism = CKM_AES_KEY_GEN;
		mechanism.pParameter = 0;
		mechanism.ulParameterLen = 0;
		
		Template_A_Import[2].pValue = &key_type[2];
		Template_A_Import[3].pValue = &ulKeyLen[3];
		Template_A_Import[9].pValue = lpKeyValue_24;
		Template_A_Import[9].ulValueLen = sizeof(lpKeyValue_24);
		
		// Find DES Key by CKA_ID/CKA_CLASS/CKA_KEY_TYPE/CKA_VALUE_LEN
		ulFindAttCount = 4;
		break;
	case '4'://AES-256
		memcpy(szKeyType, "AES-256", 7);
		Template_AES_A[2].pValue = &key_type[2];
		Template_AES_A[3].pValue = &ulKeyLen[4];

		pGenKeyTemplate = Template_AES_A;
		ulAttCount = 9;
		
		mechanism.mechanism = CKM_AES_KEY_GEN;
		mechanism.pParameter = 0;
		mechanism.ulParameterLen = 0;

		Template_A_Import[2].pValue = &key_type[2];
		Template_A_Import[3].pValue = &ulKeyLen[4];
		Template_A_Import[9].pValue = lpKeyValue_32;
		Template_A_Import[9].ulValueLen = sizeof(lpKeyValue_32);

		// Find DES Key by CKA_ID/CKA_CLASS/CKA_KEY_TYPE/CKA_VALUE_LEN
		ulFindAttCount = 4;
		break;
	default:
		printf("Invalid key type!\n");
		goto log_out;
		break;
	}

	// Generate Key
	hKeyGen = find_object(session, pGenKeyTemplate, ulFindAttCount);
	printf("Generate %s Key:\n", szKeyType);
	if (!hKeyGen)
	{
		printf ("Wait For Generating %s Key Now ... \n", szKeyType);
		rv = fl->C_GenerateKey(session, &mechanism, pGenKeyTemplate, ulAttCount, &hKeyGen);
		if (rv != PKCS_OK)
		{
			fl->C_Logout(session);
			fl->C_CloseSession(session);
			return leave("C_GenerateKey failed.  Please initialize UniMate/InitToken first.");
		}

		printf ("Generating %s Key Succeed.\n", szKeyType);
	}
	else
	{
		printf("The %s Key ID %c is already exists.\n", szKeyType, IdGen);	
	}
	printf("\n");

	// Import Key
	hKeyImp = find_object(session, Template_A_Import, ulFindAttCount);
	printf("Import %s Key:\n", szKeyType);
	if (!hKeyImp)
	{
		printf ("Wait For Import %s Key Now ... \n", szKeyType);
		rv = fl->C_CreateObject(session, Template_A_Import, 10, &hKeyImp);
		if (rv != PKCS_OK)
		{
			fl->C_Logout(session);
			fl->C_CloseSession(session);
			return leave("C_CreateObject failed.  Please run InitToken first.");
		}
		
		printf ("Import %s Key Succeed.\n", szKeyType);
	}
	else
	{
		printf("The %s Key ID %c is already exists.\n", szKeyType, IdImport);	
	}
	printf("\n");


	// ECB Encrypt
	printf("%s-ECB Encrypt\n", szKeyType);
	printf("Date to encrypt:\n");
	for (i=0; i<dataLen; i++)
	{
		printf("%02x", lpData[i]);
	}
	printf("\n");

	switch (cKeySel)
	{
	case '0':
		mechanism.mechanism = CKM_DES_ECB;
		mechanism.pParameter = 0;
		mechanism.ulParameterLen = 0;
		break;
	case '1':
		mechanism.mechanism = CKM_DES3_ECB;
		mechanism.pParameter = 0;
		mechanism.ulParameterLen = 0;
		break;
	case '2':
	case '3':
	case '4':
		mechanism.mechanism = CKM_AES_ECB;
		mechanism.pParameter = 0;
		mechanism.ulParameterLen = 0;
		break;
	default:
		printf("Invalid key type!\n");
		goto log_out;
		break;
	}

	rv = fl->C_EncryptInit(session, &mechanism, hKeyGen);
	if(rv != CKR_OK)
	{
		printf("C_EncryptInit failed.\n");
		goto log_out;
	}	

	rv = fl->C_Encrypt(session, lpData, dataLen, NULL_PTR, &en_len);
	if(rv != CKR_OK)
	{
		printf("C_Encrypt failed.\n");
		goto log_out;
	}
	if (en_len > 1024)
	{
		printf("Buffer too small.\n");
		goto log_out;
	}

	rv = fl->C_Encrypt(session, lpData, dataLen, en_data, &en_len);
	if(rv != CKR_OK)
	{
		printf("C_Encrypt failed.\n");
		goto log_out;
	}
	
	printf("\n");
	printf("Encrypt Result:\n");
	for (i=0; i<en_len; i++)
	{
		printf("%02x", en_data[i]);
	}
	printf("\n\n");
	
	// ECB Decrypt
	printf("%s-ECB Decrypt\n", szKeyType);
	rv = fl->C_DecryptInit(session, &mechanism, hKeyGen);
	if(rv != CKR_OK)
	{	
		printf("C_DecryptInit failed.\n");
		goto log_out;
	}
	
	rv = fl->C_Decrypt(session, en_data, en_len, NULL_PTR, &de_len);
	if(rv != CKR_OK)
	{
		printf("C_Decrypt failed.\n");
		goto log_out;
	}
	if (de_len > 1024)
	{
		printf("Buffer too small.\n");
		goto log_out;
	}

	rv = fl->C_Decrypt(session, en_data, en_len, de_data, &de_len);
	if(rv != CKR_OK)
	{
		printf("C_Decrypt failed.\n");
		goto log_out;
	}
	
	printf("Decrypt Result:\n");
	for (i=0; i<de_len; i++)
	{
		printf("%02x", de_data[i]);
	}
	printf("\n\n");

	// CBC Encrypt
	printf("%s-CBC Encrypt\n", szKeyType);
	printf("Date to encrypt:\n");
	for (i=0; i<dataLen; i++)
	{
		printf("%02x", lpData[i]);
	}
	printf("\n");

	switch (cKeySel)
	{
	case '0':
		mechanism.mechanism = CKM_DES_CBC;
		ivLen = 8;
		break;
	case '1':
		mechanism.mechanism = CKM_DES3_CBC;
		ivLen = 8;
		break;
	case '2':
	case '3':
	case '4':
		mechanism.mechanism = CKM_AES_CBC;
		ivLen = 16;
		break;
	default:
		printf("Invalid key type!\n");
		goto log_out;
		break;
	}

	printf("IV:\n");
	for (i=0; i<ivLen; i++)
	{
		printf("%02x", lpIV[i]);
	}
	printf("\n");

	mechanism.pParameter = lpIV;
	mechanism.ulParameterLen = ivLen;
	rv = fl->C_EncryptInit(session, &mechanism, hKeyImp);
	if(rv != CKR_OK)
	{
		printf("C_EncryptInit failed.\n");
		goto log_out;
	}	
	
	memset(en_data, 0, 1024);
	rv = fl->C_Encrypt(session, lpData, dataLen, NULL_PTR, &en_len);
	if(rv != CKR_OK)
	{
		printf("C_Encrypt failed.\n");
		goto log_out;
	}
	if (en_len > 1024)
	{
		printf("Buffer too small.\n");
		goto log_out;
	}
	
	rv = fl->C_Encrypt(session, lpData, dataLen, en_data, &en_len);
	if(rv != CKR_OK)
	{
		printf("C_Encrypt failed.\n");
		goto log_out;
	}
	
	printf("\n");
	printf("Encrypt Result:\n");
	for (i=0; i<en_len; i++)
	{
		printf("%02x", en_data[i]);
	}
	printf("\n\n");
	
	// CBC Decrypt
	printf("%s-CBC Decrypt\n", szKeyType);
	rv = fl->C_DecryptInit(session, &mechanism, hKeyImp);
	if(rv != CKR_OK)
	{	
		printf("C_DecryptInit failed.\n");
		goto log_out;
	}
	
	rv = fl->C_Decrypt(session, en_data, en_len, NULL_PTR, &de_len);
	if(rv != CKR_OK)
	{
		printf("C_Decrypt failed.\n");
		goto log_out;
	}
	if (de_len > 1024)
	{
		printf("Buffer too small.\n");
		goto log_out;
	}
	
	memset(de_data, 0, 1024);
	rv = fl->C_Decrypt(session, en_data, en_len, de_data, &de_len);
	if(rv != CKR_OK)
	{
		printf("C_Decrypt failed.\n");
		goto log_out;
	}
	
	printf("Decrypt Result:\n");
	for (i=0; i<de_len; i++)
	{
		printf("%02x", de_data[i]);
	}
	printf("\n\n");

	// CBC_PAD Encrypt
	printf("%s-CBC_PAD Encrypt\n", szKeyType);
	printf("Date to encrypt:\n");
	for (i=0; i<dataLen; i++)
	{
		printf("%02x", lpData[i]);
	}
	printf("\n");
	
	switch (cKeySel)
	{
	case '0':
		mechanism.mechanism = CKM_DES_CBC_PAD;
		ivLen = 8;
		break;
	case '1':
		mechanism.mechanism = CKM_DES3_CBC_PAD;
		ivLen = 8;
		break;
	case '2':
	case '3':
	case '4':
		mechanism.mechanism = CKM_AES_CBC_PAD;
		ivLen = 16;
		break;
	default:
		printf("Invalid key type!\n");
		goto log_out;
		break;
	}
	
	printf("IV:\n");
	for (i=0; i<ivLen; i++)
	{
		printf("%02x", lpIV[i]);
	}
	printf("\n");
	
	mechanism.pParameter = lpIV;
	mechanism.ulParameterLen = ivLen;
	rv = fl->C_EncryptInit(session, &mechanism, hKeyImp);
	if(rv != CKR_OK)
	{
		printf("C_EncryptInit failed.\n");
		goto log_out;
	}	
	
	memset(en_data, 0, 1024);
	rv = fl->C_Encrypt(session, lpData, dataLen, NULL_PTR, &en_len);
	if(rv != CKR_OK)
	{
		printf("C_Encrypt failed.\n");
		goto log_out;
	}
	
	rv = fl->C_Encrypt(session, lpData, dataLen, en_data, &en_len);
	if(rv != CKR_OK)
	{
		printf("C_Encrypt failed.\n");
		goto log_out;
	}
	
	printf("\n");
	printf("Encrypt Result:\n");
	for (i=0; i<en_len; i++)
	{
		printf("%02x", en_data[i]);
	}
	printf("\n\n");
	
	// CBC_PAD Decrypt
	printf("%s-CBC_PAD Decrypt\n", szKeyType);
	rv = fl->C_DecryptInit(session, &mechanism, hKeyImp);
	if(rv != CKR_OK)
	{	
		printf("C_DecryptInit failed.\n");
		goto log_out;
	}
	
	rv = fl->C_Decrypt(session, en_data, en_len, NULL_PTR, &de_len);
	if(rv != CKR_OK)
	{
		printf("C_Decrypt failed.\n");
		goto log_out;
	}
	
	memset(de_data, 0, 1024);
	rv = fl->C_Decrypt(session, en_data, en_len, de_data, &de_len);
	if(rv != CKR_OK)
	{
		printf("C_Decrypt failed.\n");
		goto log_out;
	}
	
	printf("Decrypt Result:\n");
	for (i=0; i<de_len; i++)
	{
		printf("%02x", de_data[i]);
	}
	printf("\n\n");
	
log_out:
	/* logout of the UniMate/UniToken
	 */
	if (PKCS_OK != fl->C_Logout (session))
		return leave("C_Logout failed.\n");
    
	/* close open session
	 */
	if (PKCS_OK != fl->C_CloseSession (session))
		return leave("C_CloseSession failed.\n");
	
	
	/* close PKCS #11 library
	 */
	if (PKCS_OK != fl->C_Finalize(0))
		return leave("C_Finalize failed.\n");
    Finalize();
	
	printf ("Press Enter to exit...");
	getchar();

	return PKCS_OK;
}


/*
	Forward declaration of utility function
		(To find an object handle, with a certain set of attributes)
*/
CK_OBJECT_HANDLE 
find_object(CK_SESSION_HANDLE hSess, CK_ATTRIBUTE * t, CK_ULONG size)
{
	CK_OBJECT_HANDLE list[1];
	CK_ULONG found = 0;
	
	// initialize the search for any object with attribute 't'
	if (CKR_OK != fl->C_FindObjectsInit (hSess, t, size))
	{
		leave("C_FindObjectsInit failed.\n");
		return 0;
	}
	
	// find the objects which match, we ask for only a single match
	// since we know there is only one on the UniMate/UniToken(InitToken creates a single public/private key pair)
	if (CKR_OK != fl->C_FindObjects(hSess, list, 1, &found))
	{
		leave("C_FindObjects failed.\n");
		return 0;
	}
	
	// cleanup the search
	//
	if (CKR_OK != fl->C_FindObjectsFinal(hSess))
	{
		leave ("C_FindObjectsFinal failed.\n");
		return 0;
	}
	
	// either return NULL if nothing matches or return the object handle
	// of the matching object
	//
	if (found < 1)
		return 0;
	
	return (list[0]);
}

