#include "test_common.h"
#ifdef WIN32
#include "getopt.h"
#else 
extern char *optarg;
#endif

void usage(void);
CK_RV test_crypto();
int test_gost3410_key_wrap(CK_SESSION_HANDLE sess);

CK_SLOT_ID SlotId;

CK_SLOT_ID_PTR SlotList = NULL;
CK_ULONG SlotCount = 0;

int main(int argc, char *argv[]) {
   CK_RV rc = 1;
   CK_FLAGS flags = 0;
   unsigned long slot_num = 0;  
   int i;

   // parse the command line parameters
	for (i=1; i<argc; i++) {
		if (strcmp("-api", argv[i]) == 0) {
			i++;
			api_path = argv[i];
		} else if (strcmp("-user_pin", argv[i]) == 0) {
			i++;
			user_pin = argv[i];
		}
	}

   // load the PKCS11 library 
	rc = init(api_path, &funcs, NULL);
   if (rc != CKR_OK) {
      fprintf(stderr, "ERROR calling init, rc = %p.\n", (void *)rc);
      return rc;
   }
	rc = get_slot_list(funcs, 
                    &SlotList, 
                    &SlotCount);
	if (rc != CKR_OK) {
		printf("get_slot_list failed, rc = 0x%x\n", rc );
		return rc;
	}
	rc = find_slot_with_token(funcs, 
                           SlotList, 
                           SlotCount, 
                           &SlotId);
	if (rc != CKR_OK) {
		printf("find_slot_with_token failed, rc = 0x%x\n", rc );
		return rc;
	}
//	for (i=0; i<5; i++) {
		// test the crypto functions 
		rc = test_crypto();
		if (rc != CKR_OK) {
			fprintf(stderr, "ERROR call to lissi failed.\n");
			return rc;
		}
//	}
	return CKR_OK;
}


void usage(void)
{
   fprintf(stderr, "Usage: lissi [-c <slotId>]\n"
      "Default slotID is 0\n"
      "To get a list of slotIDs, call pkcs11_cfg -s\n");
   exit(-1);
}

CK_RV test_crypto() 
{
   CK_RV rc, ret;
   CK_SESSION_HANDLE hSession;
	CK_MECHANISM_INFO minfo;

   // open a R/W cryptoki session, CKR_SERIAL_SESSION is a legacy bit we have to set
   rc = funcs->C_OpenSession(SlotId, CKF_RW_SESSION | CKF_SERIAL_SESSION, NULL_PTR,
         NULL_PTR, &hSession);
   if (rc != CKR_OK) {
      fprintf(stderr, "ERROR call to C_OpenSession failed, rc = %p\n", (void *)rc);
      goto out;
   }
  fprintf(stderr, "C_OpenSession success\n");

   // log in as normal user 
   rc = funcs->C_Login(hSession, CKU_USER, user_pin, strlen(user_pin));
   if (rc != CKR_OK) {
      fprintf(stderr, "ERROR call to C_Login failed, rc = %p\n", (void *)rc);
      goto out_close;
   }
  fprintf(stderr, "C_Login success\n");


	rc = funcs->C_GetMechanismInfo(SlotId, CKM_GOSTR3410_KEY_WRAP, &minfo);
	if (rc != CKR_OK) {
		fprintf(stderr, "\n======== Mechanism CKM_GOSTR3410_KEY_WRAP not supported ===========\n");
	} else {
		fprintf(stderr, "\n==================== CKM_GOSTR3410_KEY_WRAP test ==================\n");
		rc = test_gost3410_key_wrap(hSession);
		if (rc != CKR_OK) {
			fprintf(stderr, "ERROR CKM_GOSTR3410_KEY_WRAP failed, rc = %p\n", (void *)rc);
		} else {
			fprintf(stderr, "CKM_GOSTR3410_KEY_WRAP test passed.\n");
		}
	}

out_close:
   if( (ret = funcs->C_CloseSession(hSession)) != CKR_OK ) {
      fprintf(stderr, "Error: C_CloseSession failed with %p\n", (void *)ret);
	  rc = ret;
   }
   else {
      fprintf(stderr, "C_CloseSession success\n");
   }

out:
   return rc;
}

int test_gost3410_key_wrap(CK_SESSION_HANDLE sess)
{
	CK_RV rc;
	CK_BYTE value[1024];
    static CK_BYTE id[4];
	CK_ULONG len;
	CK_OBJECT_HANDLE send_pub_key = CK_INVALID_HANDLE;
	CK_OBJECT_HANDLE send_priv_key = CK_INVALID_HANDLE;
	CK_OBJECT_HANDLE recp_pub_key = CK_INVALID_HANDLE;
	CK_OBJECT_HANDLE recp_priv_key = CK_INVALID_HANDLE;
	CK_OBJECT_HANDLE cipher_key = CK_INVALID_HANDLE;
	CK_OBJECT_HANDLE unwrapped_cipher_key = CK_INVALID_HANDLE;
	CK_ATTRIBUTE attr;
   	CK_MECHANISM mechanism_desc = {CKM_GOSTR3410_KEY_WRAP, NULL, 0};
   	CK_MECHANISM_PTR mechanism = &mechanism_desc;
   	CK_MECHANISM mechanism_gen_desc = 
		{CKM_GOSTR3410_KEY_PAIR_GEN, NULL, 0};
   	CK_MECHANISM_PTR mechanism_gen = &mechanism_gen_desc;

	static CK_BBOOL ltrue = CK_TRUE;
	static CK_BBOOL lfalse = CK_FALSE;
	static CK_OBJECT_CLASS oclass = CKO_SECRET_KEY;
	static CK_OBJECT_CLASS oclass_pub = CKO_PUBLIC_KEY;
	static CK_OBJECT_CLASS oclass_priv = CKO_PRIVATE_KEY;
	static CK_KEY_TYPE wrapping_key_type = CKK_GOSTR3410;
	static CK_KEY_TYPE wrapped_key_type = CKK_GOST28147;
	// PAR ECC A OID
	static CK_BYTE gostR3410params_A[] = { 
		0x06, 0x07, 0x2a, 0x85, 0x03, 0x02, 0x02, 0x23, 0x01 
	};
	// PAR ECC XA OID
	static CK_BYTE gostR3410params_XA[] = { 
		0x06, 0x07, 0x2a, 0x85, 0x03, 0x02, 0x02, 0x24, 0x00 
	};
	// PAR HASH 1 OID
	static CK_BYTE gostR3411params[] = { 
		0x06, 0x07, 0x2a, 0x85, 0x03, 0x02, 0x02, 0x1e, 0x01 
	};
	// PAR CIPHER A OID
	static CK_BYTE gost28147params_A[] = { 
		0x06, 0x07, 0x2a, 0x85, 0x03, 0x02, 0x02, 0x1f, 0x01 
	};
	// PAR CIPHER B OID
	static CK_BYTE gost28147params_B[] = { 
		0x06, 0x07, 0x2a, 0x85, 0x03, 0x02, 0x02, 0x1f, 0x02 
	};

	// Cipher key value
	static CK_BYTE cipher_key_val[] = {
		0xc3,0x73,0x0c,0x5c,0xbc,0xca,0xcf,0x91,
		0x5a,0xc2,0x92,0x67,0x6f,0x21,0xe8,0xbd,
		0x4e,0xf7,0x53,0x31,0xd9,0x40,0x5e,0x5f,
		0x1a,0x61,0xdc,0x31,0x30,0xa6,0x50,0x11
	};
	
	// Template for cipher key generation
	static CK_ATTRIBUTE cipher_template[] = {
		{ CKA_VALUE, cipher_key_val, sizeof(cipher_key_val) },
		{ CKA_TOKEN, &lfalse, sizeof(lfalse) }, 
		{ CKA_SENSITIVE, &lfalse, sizeof(lfalse) },
		{ CKA_EXTRACTABLE, &ltrue, sizeof(ltrue) },
		{ CKA_ENCRYPT, &ltrue, sizeof(ltrue) },
		{ CKA_DECRYPT, &ltrue, sizeof(ltrue) },
		{ CKA_WRAP, &ltrue, sizeof(ltrue) },
		{ CKA_UNWRAP, &ltrue, sizeof(ltrue) },
		{ CKA_CLASS, &oclass, sizeof(oclass) },
		{ CKA_GOST28147PARAMS, 
			gost28147params_A, sizeof(gost28147params_A) },
		{ CKA_KEY_TYPE, &wrapped_key_type, sizeof(wrapped_key_type) }
	};
    static CK_BYTE ukm[] = {
        0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08
    };
	static CK_GOSTR3410_KEY_WRAP_PARAMS params = {
		gost28147params_A, sizeof(gost28147params_A),
		ukm, 8,	
		CK_INVALID_HANDLE	// Generate and use temporary ephemeral key pair
	};
	
	// Key pair labels
	static CK_BYTE priv_label[] = "Private Key Lissi 5312";
	static CK_BYTE pub_label[] = "Public Key Lissi 5312";
	
	// Templates for key pair generation.
	// Template for token public key generation and search.
	static CK_ATTRIBUTE pub_template[] = {
		{ CKA_LABEL, pub_label, sizeof(pub_label) - 1 }, 
		{ CKA_TOKEN, &ltrue, sizeof(ltrue) }, 
		{ CKA_GOSTR3410PARAMS, gostR3410params_A, sizeof(gostR3410params_A) },
		{ CKA_GOSTR3411PARAMS, gostR3411params, sizeof(gostR3411params) },
		{ CKA_WRAP, &ltrue, sizeof(ltrue) },
		{ CKA_ENCRYPT, &ltrue, sizeof(CK_BBOOL) },
		{ CKA_CLASS, &oclass_pub, sizeof(oclass_pub) },
		{ CKA_KEY_TYPE, &wrapping_key_type, sizeof(wrapping_key_type) },
	};
	// Template for token private key generation and search.
	static CK_ATTRIBUTE priv_template[] = {
		{ CKA_LABEL, priv_label, sizeof(priv_label) - 1 },
		{ CKA_TOKEN, &ltrue, sizeof(ltrue) }, 
		{ CKA_EXTRACTABLE, &ltrue, sizeof(ltrue) },
		{ CKA_SENSITIVE, &lfalse, sizeof(lfalse) },
		{ CKA_DERIVE, &ltrue, sizeof(ltrue) },
		{ CKA_UNWRAP, &ltrue, sizeof(ltrue) },
		{ CKA_DECRYPT, &ltrue, sizeof(CK_BBOOL) },
		{ CKA_CLASS, &oclass_priv, sizeof(oclass_priv) },
		{ CKA_KEY_TYPE, &wrapping_key_type, sizeof(wrapping_key_type) },
	};
	// Buffer for public key value
	static CK_BYTE pub_value[] = {
		0x1F, 0x64, 0x6B, 0x89, 0x3B, 0x28, 0x9F, 0x68, 
		0x9E, 0xD7, 0x51, 0x91, 0xF9, 0xFA, 0xAB, 0xE6, 
		0xA4, 0xDA, 0x78, 0x7A, 0x82, 0xD1, 0xB7, 0xE0, 
		0xCB, 0x0E, 0xBD, 0x74, 0x07, 0x9C, 0x63, 0xC7, 
		0xC5, 0x9D, 0xB2, 0x6A, 0x42, 0xB5, 0x66, 0x20, 
		0x3F, 0xFA, 0x80, 0xF4, 0xE8, 0xBE, 0xC4, 0x08, 
		0x69, 0x30, 0x9E, 0x3F, 0x5A, 0x96, 0x5F, 0x51, 
		0xC4, 0x5C, 0x04, 0x44, 0x7B, 0x10, 0xB0, 0x25
	};

	// Template for public key object creation
	static CK_ATTRIBUTE new_pub_template[] = {
		{ CKA_CLASS, &oclass_pub, sizeof(oclass_pub) },
		{ CKA_KEY_TYPE, &wrapping_key_type, sizeof(wrapping_key_type) },
		{ CKA_GOSTR3410PARAMS, gostR3410params_A, sizeof(gostR3410params_A) },
		{ CKA_GOSTR3411PARAMS, gostR3411params, sizeof(gostR3411params) },
		{ CKA_WRAP, &ltrue, sizeof(CK_BBOOL) },
		{ CKA_ENCRYPT, &ltrue, sizeof(CK_BBOOL) },
		{ CKA_VALUE, pub_value, sizeof(pub_value) }
	};
	// Session public key object handle
	CK_OBJECT_HANDLE new_pub_key = CK_INVALID_HANDLE;
	
	CK_OBJECT_HANDLE hObject = CK_INVALID_HANDLE;
	CK_ULONG ulObjectCount = 0;
	SYSTEMTIME          t1, t2;
	CK_ULONG            diff;


	// Generate sender key pair objects
	rc = funcs->C_GenerateKeyPair(sess, mechanism_gen,
		pub_template, sizeof(pub_template)/sizeof(CK_ATTRIBUTE),
		priv_template, sizeof(priv_template)/sizeof(CK_ATTRIBUTE),
		&send_pub_key, &send_priv_key);
	if (rc != CKR_OK) {
		fprintf(stderr, 
			"%4d: C_GenerateKeyPair failed, rc = 0x%x\n", __LINE__, rc);
		return rc;
	}
	printf("Sender key pair generation OK\n");

	// Generate recipient key pair objects
	rc = funcs->C_GenerateKeyPair(sess, mechanism_gen,
		pub_template, sizeof(pub_template)/sizeof(CK_ATTRIBUTE),
		priv_template, sizeof(priv_template)/sizeof(CK_ATTRIBUTE),
		&recp_pub_key, &recp_priv_key);
	if (rc != CKR_OK) {
		fprintf(stderr, 
			"%4d: C_GenerateKeyPair failed, rc = 0x%x\n", __LINE__, rc);
		return rc;
	}
	printf("Recipient key pair generation OK\n");

	// Get public key value
	attr.type = CKA_VALUE; 
	attr.pValue = pub_value; 
	attr.ulValueLen = sizeof(pub_value);
	rc = funcs->C_GetAttributeValue(sess, recp_pub_key, &attr, 1);
	if (rc != CKR_OK) {
		fprintf(stderr, 
			"%4d: C_GetAttributeValue failed, rc = 0x%x\n", __LINE__, rc);
		return rc;
	}
//    gost_hexdump(stdout, 
//		"Recipient Public Key Value:", pub_value, attr.ulValueLen);

	// Create session public key object from its value.
	// There is no need to do it here but we make session public
	// key for testing only.
	rc = funcs->C_CreateObject(sess, new_pub_template, 
	    sizeof(new_pub_template)/sizeof(CK_ATTRIBUTE), &new_pub_key); 
	if (rc != CKR_OK) {
		fprintf(stderr, 
			"%4d: C_CreateObject failed, rc = 0x%x\n", __LINE__, rc);
		return rc;
	}
	
	// Create cipher key object from its value
	rc = funcs->C_CreateObject(sess, cipher_template, 
	    sizeof(cipher_template)/sizeof(CK_ATTRIBUTE), &cipher_key);
	if (rc != CKR_OK) {
		fprintf(stderr, 
			"%4d: C_CreateObject failed, rc = 0x%x\n", __LINE__, rc);
		return rc;
	}

	// Get cipher key to the value buffer 
	// from the cipher_key object (for testing).
	attr.type = CKA_VALUE; 
	attr.pValue = value; 
	attr.ulValueLen = sizeof(value);
	rc = funcs->C_GetAttributeValue(sess, cipher_key, &attr, 1);
	if (rc != CKR_OK) {
		fprintf(stderr, 
			"%4d: C_GetAttributeValue failed, rc = 0x%x\n", __LINE__, rc);
		return rc;
	}
	len = attr.ulValueLen;
//	gost_hexdump(stdout, "Cipher key value:", value, len);

	mechanism->pParameter = &params; 
	mechanism->ulParameterLen = sizeof(params);

//   avg_time = 0;
//   max_time = 0;
//   min_time = 0xFFFFFFFF;
//   for (i=0; i < 10; i++) {
//      GetSystemTime(&t1);
	params.pWrapOID = gost28147params_A;
	params.ulWrapOIDLen = sizeof(gost28147params_A);
	// Использовать закрытый отправителя 
	// или NULL для генерации эфемерной ключевой пары.
	params.hKey = send_priv_key;	
	// Get wrapping key length only.
	len = sizeof(value);
	rc = funcs->C_WrapKey(sess, mechanism, 
		new_pub_key, cipher_key, NULL, &len);
	if (rc != CKR_OK) {
		fprintf(stderr, 
			"%4d: C_WrapKey failed, rc = 0x%x\n", __LINE__, rc);
		return rc;
	}
//	fprintf(stdout, "Wrapped key length = %d\n", len);

	// Wrap cipher key with public key to the value buffer.
    GetSystemTime(&t1);
	rc = funcs->C_WrapKey(sess, mechanism, 
		new_pub_key, cipher_key, value, &len);
    GetSystemTime(&t2);
    diff = process_time(t1, t2);
	fprintf(stderr, "C_WrapKey time:  %ld msec\n", diff );
	if (rc != CKR_OK) {
		fprintf(stderr, 
			"%4d: C_WrapKey failed, rc = 0x%x\n", __LINE__, rc);
		return rc;
	}
//	gost_hexdump(stdout, "Wrapped cipher key value:", value, len);
	printf("Key wrapped OK\n");


/*
	// Проверяем вариант с отсутствующим OID 3411 (0x05, 0x00)
	memcpy(value+0x4f+2, value+0x4f+9, len-(0x4f+9));
	len -= 7;
*/

	params.pWrapOID = NULL;
	params.ulWrapOIDLen = 0;
	// Использовать открытый ключ отправителя 
	// или NULL при генерации эфемерной ключевой пары.
	if (len <= 65) {
		params.hKey = send_pub_key;
	}
    GetSystemTime(&t1);
    // Unwrap cipher key with recipient private key from the value buffer.
    // Don't use the first CKA_VALUE attribute - 
	// it will be added as a result of C_UnwrapKey for cipher_key.
    rc = funcs->C_UnwrapKey(sess, mechanism, 
		recp_priv_key, value, len, cipher_template + 1, 
	    sizeof(cipher_template)/sizeof(CK_ATTRIBUTE) - 1, &unwrapped_cipher_key);
    GetSystemTime(&t2);
    diff = process_time(t1, t2);
	fprintf(stderr, "C_UnwrapKey time:  %ld msec\n", diff );
	if (rc != CKR_OK) {
		fprintf(stderr, 
			"%4d: C_UnwrapKey failed, rc = 0x%x\n", __LINE__, rc);
		return rc;
	}
	printf("Key unwrapped OK\n");

	// Restore cipher key to the value buffer from the cipher_key object.
	attr.type = CKA_VALUE; 
	attr.pValue = value; 
	attr.ulValueLen = sizeof(value);
	rc = funcs->C_GetAttributeValue(sess, unwrapped_cipher_key, &attr, 1);
	if (rc != CKR_OK) {
		fprintf(stderr, 
			"%4d: C_GetAttributeValue failed, rc = 0x%x\n", __LINE__, rc);
		return rc;
	}
	len = attr.ulValueLen;
	// Check cipher key value
	if (len != sizeof(cipher_key_val)) {
		fprintf(stderr, 
			"%4d: Invalid result length: %d\n", __LINE__, len);
		return -1;
	}
	if (memcmp(value, cipher_key_val, len) != 0) {
		fprintf(stderr, "%4d: Invalid result value\n", __LINE__);
		return -2;
	}
	printf("Unwrapped key is equal to source key\n");

	rc = funcs->C_DestroyObject(sess, unwrapped_cipher_key);
	if (rc != CKR_OK) {
		fprintf(stderr, 
			"%4d: C_DestroyObject failed, rc = 0x%x\n", __LINE__, rc);
		return rc;
	}
//      GetSystemTime(&t2);
//      diff = process_time(t1, t2);
//      avg_time += diff;
//      if (diff < min_time)
//         min_time = diff;
//      if (diff > max_time)
//         max_time = diff;
//   }
//   printf("10 GOST R34.10-2001 Wrap/Unwrap operations:  %ld \n", avg_time );
//   printf("Minimum:                        %ld \n", min_time );
//   printf("Maximum:                        %ld \n", max_time );
//   printf("\n");

	// Destroy session public key
	rc = funcs->C_DestroyObject(sess, new_pub_key);
	if (rc != CKR_OK) {
		fprintf(stderr, 
			"%4d: C_DestroyObject failed, rc = 0x%x\n", __LINE__, rc);
		return rc;
	}
	// Destroy cipher key
	rc = funcs->C_DestroyObject(sess, cipher_key);
	if (rc != CKR_OK) {
		fprintf(stderr, 
			"%4d: C_DestroyObject failed, rc = 0x%x\n", __LINE__, rc);
		return rc;
	}
	rc = funcs->C_DestroyObject(sess, recp_pub_key);
	if (rc != CKR_OK) {
		fprintf(stderr, 
			"%4d: C_DestroyObject failed, rc = 0x%x\n", __LINE__, rc);
		return rc;
	}
	rc = funcs->C_DestroyObject(sess, recp_priv_key);
	if (rc != CKR_OK) {
		fprintf(stderr, 
			"%4d: C_DestroyObject failed, rc = 0x%x\n", __LINE__, rc);
		return rc;
	}
	rc = funcs->C_DestroyObject(sess, send_pub_key);
	if (rc != CKR_OK) {
		fprintf(stderr, 
			"%4d: C_DestroyObject failed, rc = 0x%x\n", __LINE__, rc);
		return rc;
	}
	rc = funcs->C_DestroyObject(sess, send_priv_key);
	if (rc != CKR_OK) {
		fprintf(stderr, 
			"%4d: C_DestroyObject failed, rc = 0x%x\n", __LINE__, rc);
		return rc;
	}
	printf("SUCCESS\n");
	rc = CKR_OK;
	
	return rc;
}



