/*--
	SignVerifyByDSTU4145.cpp

	2015-10-15

	Test pkcs Generate DSTU4145 Key and sign/verify
--*/

#include "../global.h"

CK_FUNCTION_LIST_PTR fl = 0; 
CK_BYTE message[1024];
CK_BYTE signature[500];    
CK_ULONG signature_length = 500;

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  = CKK_DSTU4145;
	int loop = 0;
	CK_ULONG msgLen = 0;
	/* setup public key attributes, specific name so we can find it 
	 * later, DSTU4145 key type, saved on the UniMate/UniToken and public to all
	 * users of the UniMate/UniToken
	 */
	CK_OBJECT_CLASS class_public_key = CKO_PUBLIC_KEY;
	CK_ULONG vecModulusBits  = 163;
	CK_BBOOL b = TRUE;
	CK_BYTE KID = 0;
	CK_BYTE ID = '0';
	CK_CHAR	label_public[] = labelPu;
	CK_CHAR label_private[] = labelPr;
	CK_ATTRIBUTE public_key_template[] = 
	{
		{ CKA_CLASS,			&class_public_key,  sizeof (class_public_key) },
		{ CKA_ID,				&ID,				sizeof(CK_BYTE)},
		{ CKA_KEY_TYPE,			&key_type,          sizeof (key_type) },
		{ CKA_ENCRYPT,			&b,					sizeof (b) },
		{ CKA_LABEL,			label_public,       sizeof (label_public) },
		{ CKA_TOKEN,			&bTrue,             sizeof (bTrue) },
		{ CKA_PRIVATE,			&bFalse,            sizeof (bFalse) },
		{ CKA_MODULUS_BITS,		&vecModulusBits,	sizeof (vecModulusBits) }
	};
	
	/* setup private key attributes, specific name so we can find it 
	 * later, DSTU4145 key type, saved on the UniMate/UniToken and private to only
	 * the user of the UniMate/UniToken
	 */
	CK_OBJECT_CLASS class_private_key = CKO_PRIVATE_KEY;
	CK_ATTRIBUTE private_key_template[] = 
	{
		{ CKA_CLASS,    &class_private_key, sizeof (class_private_key) },
		{ CKA_ID,		&ID,				sizeof(CK_BYTE)},
		{ CKA_KEY_TYPE, &key_type,          sizeof (key_type) },
		{ CKA_DECRYPT,  &b,					sizeof (b) },
		{ CKA_SIGN,		&b,					sizeof (b) },
		{ CKA_LABEL,    label_private,      sizeof (label_private) },
		{ CKA_TOKEN,    &bTrue,             sizeof (bTrue) },
		{ CKA_PRIVATE,  &bTrue,             sizeof (bTrue) },
		{ CKA_KEY_ID,   &KID,				sizeof (KID) }
	};
	/* setup mechanism to generate DSTU4145 key pair
	 */
	CK_MECHANISM mechanism = {0};
	CK_OBJECT_HANDLE public_key;
	CK_OBJECT_HANDLE private_key;
	int i = 0;

	char szMsg[] = {"the dstu4145 sign and verify test!"};
	memcpy(message, szMsg, strlen(szMsg));
	msgLen = strlen(szMsg);
	
	printf("Sign with Dstu4145 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 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 ("Found UniMate/UniToken.\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 UniMate/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 pair.
	*/
	printf("Please selected the DSTU4145 keypair bit:\n");
	printf("0: DSTU4145-163\n");
	printf("1: DSTU4145-167\n");
	printf("2: DSTU4145-173\n");
	printf("3: DSTU4145-179\n");
	printf("4: DSTU4145-191\n");
	printf("5: DSTU4145-233\n");
	printf("6: DSTU4145-257\n");
	printf("7: DSTU4145-307\n");
	printf("8: DSTU4145-367\n");
	printf("9: DSTU4145-431\n");
	switch (getchar())
	{
	case '0':
		ID = '0';
		vecModulusBits = 163;
		break;
	case '1':
		ID = '1';
		vecModulusBits = 167;
		break;
	case '2':
		ID = '2';
		vecModulusBits = 173;
		break;
	case '3':
		ID = '3';
		vecModulusBits = 179;
		break;
	case '4':
		ID = '4';
		vecModulusBits = 191;
		break;
	case '5':
		ID = '5';
		vecModulusBits = 233;
		break;
	case '6':
		ID = '6';
		vecModulusBits = 257;
		break;
	case '7':
		ID = '7';
		vecModulusBits = 307;
		break;
	case '8':
		ID = '8';
		vecModulusBits = 367;
		break;
	case '9':
		ID = '9';
		vecModulusBits = 431;
		break;
	default:
		printf("Invalid bit!\n");
		goto log_out;
		break;
	}

	mechanism.mechanism = CKM_DSTU4145_KEY_PAIR_GEN;
	private_key = find_object(session, private_key_template, 3);
	public_key = find_object(session,  public_key_template,  3);
	
	if (!(public_key && private_key))
	{
		printf ("Wait For Generate DSTU4145 KEY PAIR Now ... \n");
		
		if((rv = fl->C_GenerateKeyPair(session,&mechanism, public_key_template, 8,
			private_key_template, 9, &public_key,&private_key))!=PKCS_OK)
		{
			fl->C_Logout(session);
			fl->C_CloseSession(session);
			return leave("C_GenerateKeyPair failed.  Please run InitToken first.");
		}
	}
	
	/* setup the DSTU4145 mechanism
	 */
	memset (&mechanism, 0, sizeof (mechanism));
	mechanism.mechanism = CKM_GOST34311_DSTU4145;
	printf("The Mechanism:CKM_GOST34311_DSTU4145\n");
	
	/* setup signature operation with private key
	 */
	rv = fl->C_SignInit(session, &mechanism, private_key);
	if (CKR_OK != rv)
	{
		fl->C_Logout(session);
		
		return leave("C_SignInit failed...\n");
	}                       
	
	/* sign the message and get back the signature
	 */
	printf("Sign Message:\n");
	printf("%s\n", szMsg);
	printf("\n");

	rv = fl->C_Sign(session, message, msgLen, signature, &signature_length);
	if (CKR_OK != rv)
	{
		fl->C_Logout(session);
		
		if (CKR_CANCEL == rv)
		{
			return leave("The signature was canceled!");
		}
		else
		{
			return leave("C_Sign failed...\n");
		}
	}

	printf("Signature:\n");
	for(loop;loop<signature_length;loop++)
	{
		printf ("%x",signature[loop]);
	}
	printf("\n");
	printf("Sign successful!\n");
	/******************************************/


	// Verify
	rv = fl->C_VerifyInit(session, &mechanism, public_key);
	if (CKR_OK != rv)
	{
		fl->C_Logout(session);	
		return leave("C_VerifyInit failed...\n");
	} 

	rv = fl->C_Verify(session, message, msgLen, signature, signature_length);
	if (CKR_OK != rv)
	{
		fl->C_Logout(session);
		return leave("C_Verify failed...\n");
	}
	printf("\n");
	printf("Verify successful!\n");
	printf("\n");
	
	/* destroy both DSTU4145 key objects
	 */
	//if (PKCS_OK != fl->C_DestroyObject (session, public_key))
	//	return leave("C_DestroyObject Public Obj failed.\n");
	
	//if (PKCS_OK != fl->C_DestroyObject (session, private_key))
	//	return leave("C_DestroyObject Private Obj failed.\n");
	
	/* logout of the UniToken
	 */

log_out:
	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 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]);
}