#include "test_common.h"

CK_RV test_crypto();
CK_RV test_gost28147_pkcs8_key_wrap(CK_SESSION_HANDLE sess);

int main(int argc, char *argv[]) {
	CK_RV rc = CKR_OK;

	rc = test_main(argc, argv);
	if (rc != CKR_OK) {
		fprintf(stderr, "test_main failed: 0x%x\n", rc);
		return rc;
	}
	rc = test_crypto();
	if (rc != CKR_OK) {
		fprintf(stderr, "test_crypto failed: 0x%x\n", rc);
		return rc;
	}
	return CKR_OK;
}

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;
	}

   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 = 0x%x\n", rc);
		goto out;
   }
/*
   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 = 0x%x\n", 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");
   }
 
   funcs->C_CloseSession(hSession);
out:
   return rc;
}

CK_RV 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;

	static CK_BYTE salt[] = { 
		0x12, 0x34, 0x56, 0x78, 0x78, 0x56, 0x34, 0x26 
	};
	static CK_ULONG iter = 2048;
	static CK_BYTE iv[] = { 
		0x12, 0x34, 0x56, 0x78, 0x78, 0x56, 0x34, 0xab 
	};
    // Use UTF-8 password
	static CK_UTF8CHAR_PTR password = "4444";
	static CK_ULONG password_len = 4;
	static CK_ATTRIBUTE pub_template[] = {
		{ CKA_GOSTR3410PARAMS, gostR3410params, sizeof(gostR3410params) },
		{ CKA_GOSTR3411PARAMS, gostR3411params, sizeof(gostR3411params) },
	};
	// Template for token private key generation.
	static CK_ATTRIBUTE priv_template[] = {
		{ CKA_EXTRACTABLE, &ltrue, sizeof(ltrue) },
		{ CKA_SENSITIVE, &ltrue, sizeof(ltrue) }
	};
	// Additional attributes for token unwrapped private key generation.
	static CK_ATTRIBUTE unw_template[] = {
		{ CKA_CLASS, (CK_VOID_PTR)&oclass_priv, sizeof(oclass_priv) },
		{ CKA_KEY_TYPE, &key_type, sizeof(key_type) }
	};
	static 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) },
	};
	// При упаковке на выходе должна получиться такая структура:
    static 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;

        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;
		}
	
        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;
		}
 
        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;
}
