#include "test_common.h"

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

void usage(void);
CK_RV test_crypto();
int test_gost28147_pkcs8_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<10; 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;
   CK_SESSION_HANDLE hSession;
	CK_MECHANISM_INFO minfo;

	printf("CKM_GOST28147_PKCS8_KEY_WRAP and CKM_GOSTR3410_PUBLIC_KEY_DERIVE test\n");

	rc = funcs->C_GetMechanismInfo(SlotId, CKM_GOSTR3410_KEY_PAIR_GEN, &minfo);
	if (rc != CKR_OK) {
		fprintf(stderr, "Mechanism CKM_GOSTR3410_KEY_PAIR_GEN not supported\n");
		return rc;
	}
	rc = funcs->C_GetMechanismInfo(SlotId, CKM_GOST28147_PKCS8_KEY_WRAP, &minfo);
	if (rc != CKR_OK) {
		fprintf(stderr, "Mechanism CKM_GOST28147_PKCS8_KEY_WRAP not supported\n");
		return rc;
	}
	rc = funcs->C_GetMechanismInfo(SlotId, CKM_GOSTR3410_PUBLIC_KEY_DERIVE, &minfo);
	if (rc != CKR_OK) {
		fprintf(stderr, "Mechanism CKM_GOSTR3410_PUBLIC_KEY_DERIVE not supported\n");
		return rc;
	}

   // 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;
   }
/*
   // 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;
   }
*/   
   rc = test_gost28147_pkcs8_key_wrap(hSession);

   if (rc != CKR_OK) {
      fprintf(stderr, "CKM_GOST28147_PKCS8_KEY_WRAP and CKM_GOSTR3410_PUBLIC_KEY_DERIVE test failed, rc = 0x%x\n", rc);
   } else {
      printf("CKM_GOST28147_PKCS8_KEY_WRAP and CKM_GOSTR3410_PUBLIC_KEY_DERIVE test SUCCESS\n");
   }
 
//out_close:
   funcs->C_CloseSession(hSession);
out:
   return rc;
}

int test_gost28147_pkcs8_key_wrap(CK_SESSION_HANDLE sess)
{
	int rc = 0;
    char str[4096];
	CK_BYTE value[1024];
	CK_ULONG len;
	CK_OBJECT_HANDLE pub_key = CK_INVALID_HANDLE;
	CK_OBJECT_HANDLE priv_key = CK_INVALID_HANDLE;
	CK_OBJECT_HANDLE unw_priv_key = CK_INVALID_HANDLE;
    CK_OBJECT_HANDLE der_pub_key = CK_INVALID_HANDLE;
	CK_OBJECT_HANDLE cipher_key = CK_INVALID_HANDLE;
   	CK_MECHANISM mechanism_desc = 
		{CKM_GOST28147_PKCS8_KEY_WRAP, NULL, 0};
   	CK_MECHANISM_PTR mechanism = &mechanism_desc;
   	CK_MECHANISM mechanism_der_desc = 
		{CKM_GOSTR3410_PUBLIC_KEY_DERIVE, NULL, 0};
   	CK_MECHANISM_PTR mechanism_der = &mechanism_der_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 const CK_OBJECT_CLASS oclass_pub = CKO_PUBLIC_KEY;
	static const CK_OBJECT_CLASS oclass_priv = CKO_PRIVATE_KEY;
	static CK_KEY_TYPE key_type = CKK_GOSTR3410;
	// PAR ECC A OID
	static CK_BYTE gostR3410params[] = { 
		0x06, 0x07, 0x2a, 0x85, 0x03, 0x02, 0x02, 0x23, 0x01 
	};
	// 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 
	};
/*
typedef struct CK_GOST28147_PKCS8_KEY_WRAP_PARAMS {
    CK_CHAR_PTR      pPassword;
    CK_ULONG         ulPasswordLen;
	CK_BYTE_PTR      pHashParOID;
	CK_ULONG         ulHashParOIDLen;
	CK_BYTE_PTR      pSalt;
	CK_ULONG         ulSaltLen;
    CK_ULONG         ulIterCount;
	CK_BYTE_PTR      pCipherParOID;
	CK_ULONG         ulCipherParOIDLen;
	CK_BYTE_PTR      pIV;
	CK_ULONG         ulIVLen;
} CK_GOST28147_PKCS8_KEY_WRAP_PARAMS;
*/
	CK_GOST28147_PKCS8_KEY_WRAP_PARAMS params;

	CK_BYTE salt[] = { 
		0x12, 0x34, 0x56, 0x78, 0x78, 0x56, 0x34, 0x26 
	};
	CK_ULONG iter = 2048;
	CK_BYTE iv[] = { 
		0x12, 0x34, 0x56, 0x78, 0x78, 0x56, 0x34, 0xab 
	};
    // Use UTF-8 password
	CK_UTF8CHAR_PTR password = "4444";
	CK_ULONG password_len = 4;
	CK_ATTRIBUTE pub_template[] = {
		{ CKA_GOSTR3410PARAMS, gostR3410params, sizeof(gostR3410params) },
		{ CKA_GOSTR3411PARAMS, gostR3411params, sizeof(gostR3411params) },
	};
	// Template for token private key generation.
	CK_ATTRIBUTE priv_template[] = {
		{ CKA_EXTRACTABLE, &ltrue, sizeof(ltrue) },
		{ CKA_SENSITIVE, &ltrue, sizeof(ltrue) }
	};
	// Additional attributes for token unwrapped private key generation.
	CK_ATTRIBUTE unw_template[] = {
		{ CKA_CLASS, (CK_VOID_PTR)&oclass_priv, sizeof(oclass_priv) },
		{ CKA_KEY_TYPE, &key_type, sizeof(key_type) }
	};
	CK_ATTRIBUTE pub_der_template[] = {
        { CKA_CLASS, (CK_VOID_PTR)&oclass_pub, sizeof(oclass_pub) },
        { CKA_KEY_TYPE, &key_type, sizeof(key_type) },
		{ CKA_GOSTR3410PARAMS, gostR3410params, sizeof(gostR3410params) },
		{ CKA_GOSTR3411PARAMS, gostR3411params, sizeof(gostR3411params) },
	};
	// При упаковке на выходе должна получиться такая структура:
    CK_BYTE pr_key_bag[] = {
// 0
        0x30, 0x81, 0xa7,
// 3
          0x30, 0x5c,
// 5
            0x06, 0x09, 
			  0x2a, 0x86, 0x48, 0x86, 0xf7, 0x0d, 0x01, 0x05, 0x0d,
// 16
            0x30, 0x4f,
// 18
              0x30, 0x2e,
// 20
                0x06, 0x09, 
				  0x2a, 0x86, 0x48, 0x86, 0xf7, 0x0d, 0x01, 0x05, 0x0c,
// 31
                0x30, 0x21,
// 33
                  0x04, 0x08,
// 35 - salt 
                    0x12, 0x34, 0x56, 0x78, 0x78, 0x56, 0x34, 0x26,
// 43
                  0x02, 0x02,
// 45 - iter (2048)
                    0x08, 0x00,
// 47
                  0x30, 0x11,
// 49
                    0x06, 0x06, 
					  0x2a, 0x85, 0x03, 0x02, 0x02, 0x0a,
// 57 - oid_3411_par 
                    0x06, 0x07, 
					  0x2a, 0x85, 0x03, 0x02, 0x02, 0x1e, 0x01,
// 66
                    0x30, 0x1d,
// 68
                      0x06, 0x06, 
						0x2a, 0x85, 0x03, 0x02, 0x02, 0x15,
// 76
                      0x30, 0x13,
// 78
                        0x04, 0x08,
// 80 - iv
                          0x12, 0x34, 0x56, 0x78, 0x78, 0x56, 0x34, 0xab,
// 88 - oid_28147_par
                        0x06, 0x07, 
						  0x2a, 0x85, 0x03, 0x02, 0x02, 0x1f, 0x01,
// 97
          0x04, 0x47, 
// 99 - encr_pr_key
            0x46, 0x5c, 0xc8, 0x21, 0xcd, 0xcc, 0xb8, 0x99, 
            0x53, 0xa7, 0x72, 0xda, 0x20, 0x0f, 0x3d, 0xd5, 
            0xab, 0x59, 0x11, 0x6f, 0x4f, 0x8a, 0x75, 0x9b, 
            0xf4, 0xd1, 0x91, 0x7e, 0x9d, 0x2f, 0x79, 0xab, 
            0x95, 0xb8, 0x54, 0xe2, 0x5b, 0x21, 0x61, 0xa8, 
            0xe7, 0x2b, 0x58, 0x5b, 0xe2, 0x10, 0xd0, 0xd8, 
            0xbb, 0xd4, 0x04, 0xdc, 0x06, 0x4c, 0x60, 0x14, 
            0x60, 0x12, 0xd1, 0xf1, 0xb2, 0xbf, 0x39, 0xf7, 
            0xc8, 0x35, 0x16, 0xc4, 0xea, 0xe9, 0xb6
// 170
    };
	// Заметим, что при распаковке во входной структуре допускается 
	// использование 0x05, 0x00 вместо oid_3411_par. Соответственно,
	// длина такой укороченной структуры будет меньше на 7 байтов.
	// LSMS11 воспринимает оба варианта.

	rc = funcs->C_GenerateKeyPair(sess, mechanism_gen,
			pub_template, sizeof(pub_template)/sizeof(CK_ATTRIBUTE),
			priv_template, sizeof(priv_template)/sizeof(CK_ATTRIBUTE),
			&pub_key, &priv_key);
	if (rc != CKR_OK) {
		fprintf(stderr, "C_GenerateKeyPair failed, rc = 0x%x\n", rc);
		goto end;
	}

	params.pPassword = password;
	params.ulPasswordLen = strlen(password);
	params.pHashParOID = gostR3411params;
	params.ulHashParOIDLen = sizeof(gostR3411params);
	params.pSalt = salt;
    params.ulSaltLen = sizeof(salt);
    params.ulIterCount = iter;
	params.pCipherParOID = gost28147params_A;
	params.ulCipherParOIDLen = sizeof(gost28147params_A);
    params.pIV = iv;
    params.ulIVLen = sizeof(iv);

	mechanism->pParameter = &params;
	mechanism->ulParameterLen = sizeof(params);
	rc = funcs->C_WrapKey(sess, mechanism, CK_INVALID_HANDLE, priv_key, value, &len);
	if (rc != CKR_OK) {
		fprintf(stderr, "C_WrapKey failed, rc = 0x%x\n", rc);
		goto end;
	}

    print_bytes(value, len, str);
    printf("PKCS#8:\n%s\n", str);

    if (len != sizeof(pr_key_bag)) {
        fprintf(stderr, "Invalid length: %ld\n", len);
        rc = -1;
		goto end;
    }
    if (memcmp(value, pr_key_bag, len-71) != 0) {
        fprintf(stderr, "Invalid PKCS#8 value\n");
        rc = -1;
		goto end;
    }
/*
	// Проверяем вариант с отсутствующим OID 3411 (0x05, 0x00)
	memcpy(value+35+8+2+2+10+2, value+35+8+2+2+10+9, len-(35+8+2+2+10+9));
	len -= 7;
*/
    mechanism->pParameter = password;
    mechanism->ulParameterLen = strlen(password);
	rc = funcs->C_UnwrapKey(sess, mechanism, CK_INVALID_HANDLE, 
        value, len,
        unw_template, sizeof(unw_template)/sizeof(CK_ATTRIBUTE), 
        &unw_priv_key);
	if (rc != CKR_OK) {
		fprintf(stderr, "C_UnwrapKey failed, rc = 0x%x\n", rc);
		goto end;
	}
	// Получаем открытый ключ по закрытому
	rc = funcs->C_DeriveKey(sess, mechanism_der, unw_priv_key, 
        pub_der_template, sizeof(pub_der_template)/sizeof(CK_ATTRIBUTE), 
        &der_pub_key);
	if (rc != CKR_OK) {
		fprintf(stderr, "C_DeriveKey failed, rc = 0x%x\n", rc);
		goto end;
	}
    // Для проверки распакованного ключа нужно что-нибудь им подписать
    // и проверить подпись открытым ключом.
    {
        CK_MECHANISM mechanism_sign = {CKM_GOSTR3410, NULL, 0};
		// Подписываемый дайджест (32 байта)
        CK_BYTE pData[32] = {
			0x4D, 0x89, 0x9E, 0x48, 0xC5, 0x39, 0x64, 0xD1, 
			0x78, 0xB4, 0x6D, 0x58, 0x40, 0x5F, 0x62, 0x8F, 
			0xA4, 0x46, 0x4C, 0xDC, 0x73, 0x75, 0xB0, 0xE5, 
			0x2E, 0x91, 0xFB, 0x64, 0x9C, 0x06, 0xAB, 0x75		
		};
        CK_ULONG ulDataLen = 32;
        CK_BYTE	pSignatureData[64];
        CK_ULONG signatureDataLen = 64;

        // Sign with unwrapped private key
        rc = funcs->C_SignInit(sess, &mechanism_sign, unw_priv_key);
		if (rc != CKR_OK) {
			fprintf(stderr, "C_SignInit failed, rc = 0x%x\n", rc);
			goto end;
		}
        rc = funcs->C_Sign(sess, 
			pData, ulDataLen, pSignatureData, &signatureDataLen);
		if (rc != CKR_OK) {
			fprintf(stderr, "C_Sign failed, rc = 0x%x\n", rc);
			goto end;
		}
	
        // Now try to verify with source public key
        rc = funcs->C_VerifyInit(sess, &mechanism_sign, pub_key);
		if (rc != CKR_OK) {
			fprintf(stderr, "C_VerifyInit failed, rc = 0x%x\n", rc);
			goto end;
		}
        rc = funcs->C_Verify(sess, 
			pData, ulDataLen, pSignatureData, signatureDataLen);
		if (rc != CKR_OK) {
			fprintf(stderr, "C_Verify failed, rc = 0x%x\n", rc);
			goto end;
		}
 
        // Now try to verify with derived public key
        rc = funcs->C_VerifyInit(sess, &mechanism_sign, der_pub_key);
		if (rc != CKR_OK) {
			fprintf(stderr, "C_VerifyInit failed, rc = 0x%x\n", rc);
			goto end;
		}
        rc = funcs->C_Verify(sess, 
			pData, ulDataLen, pSignatureData, signatureDataLen);
		if (rc != CKR_OK) {
			fprintf(stderr, "C_Verify failed, rc = 0x%x\n", rc);
			goto end;
		}
    }
	printf("SUCCESS\n");
end:
    funcs->C_DestroyObject(sess, der_pub_key);
    funcs->C_DestroyObject(sess, unw_priv_key);
	funcs->C_DestroyObject(sess, pub_key);
	funcs->C_DestroyObject(sess, priv_key);

    return rc;
}

