#include "test_common.h"

CK_RV test_kuznyechik_key_wrap(CK_SESSION_HANDLE hSession);
CK_RV test_magma_key_wrap(CK_SESSION_HANDLE hSession);

CK_FUNCTION_LIST_PTR Pkcs11FuncList     = NULL;

int main(int argc, char* argv[])
{
    CK_RV                rvResult;
#ifdef WIN32
    HMODULE              hPkcsLib  = NULL;
#else
    void                 *hPkcsLib  = NULL;
#endif
    CK_C_GetFunctionList pcGetFunctionList  = 0;
    CK_SLOT_ID_PTR       pSlotList          = NULL;
    CK_SLOT_ID           SlotId = 0;
    CK_ULONG             ulSlotCount;
    CK_ULONG             i;
    CK_SLOT_INFO         SlotInfo;
    CK_SESSION_HANDLE    hSession;


    char							*api_path  = STRINGIFY(PKCS11_API_PATH);


    printf("Starting CKM_GOST_CIPHER_KEY_WRAP test\n");


    for (i = 1; i < (CK_ULONG) argc; i++)
    {
        if (strcmp("-api", argv[i]) == 0)
        {
            ++i;
            api_path = argv[i];
        }
        else if (strcmp("-user_pin", argv[i]) == 0)
        {
            ++i;
            user_pin = (CK_UTF8CHAR *)argv[i];
        }
    }

#ifdef WIN32
    hPkcsLib = LoadLibrary(api_path);
#else
    hPkcsLib = dlopen(api_path, RTLD_LOCAL|RTLD_LAZY|RTLD_DEEPBIND);
#endif
    if (hPkcsLib == NULL)
    {
        printf(
            "Can't load PKCS#11 API library. "
            "Check API library path.\n");
#ifndef WIN32
        printf("dlerror: %s\n", dlerror());
#endif
        return -1;
    }
#ifdef WIN32
    pcGetFunctionList =
        (CK_C_GetFunctionList) GetProcAddress(
            hPkcsLib, "C_GetFunctionList");
#else
    pcGetFunctionList =
        (CK_C_GetFunctionList) dlsym(
            hPkcsLib, "C_GetFunctionList");
#endif

    // get PKCS #11 function list
    rvResult = pcGetFunctionList(&Pkcs11FuncList);
    printf("Load PKCS #11 function list result: 0x%lx\n", rvResult);
    if (rvResult != CKR_OK)
        return rvResult;

    // initialize Cryptoki
    rvResult = Pkcs11FuncList->C_Initialize(NULL);
    printf("Initialize Cryptoki result: 0x%lx\n", rvResult);
    if (rvResult != CKR_OK && rvResult != CKR_CRYPTOKI_ALREADY_INITIALIZED)
        return rvResult;

    // get slot list
    rvResult = Pkcs11FuncList->C_GetSlotList(
        CK_FALSE, NULL, &ulSlotCount);
    printf("Get slot list result: 0x%lx\n", rvResult);
    if (rvResult != CKR_OK)
        return rvResult;

    if (ulSlotCount > 0)
    {       // allocate memory for slot list
        pSlotList = (CK_SLOT_ID_PTR) malloc(
            ulSlotCount * sizeof(CK_SLOT_ID));

        rvResult =
            Pkcs11FuncList->C_GetSlotList(CK_FALSE,
                                          pSlotList, &ulSlotCount);

        if (rvResult != CKR_OK)
            return rvResult;
        printf("Slot count: %ld\n", ulSlotCount);
    }
    else
        return -3;

    // get information about sender and recipient slots.
    for (i = 0; i < ulSlotCount; ++i)
    {
        rvResult = Pkcs11FuncList->C_GetSlotInfo(pSlotList[i], &SlotInfo);
        if (rvResult == CKR_OK)
        {          // if a token is present in this slot
            if ((SlotInfo.flags & CKF_TOKEN_PRESENT) == CKF_TOKEN_PRESENT)
            {
                SlotId = pSlotList[i];
                break;
            }
        }
    }
    if (i >= ulSlotCount)
    {
        printf("No slots with token present\n");
        return -3;
    }
    printf("Slot ID: sender %ld\n", SlotId);
    // open session for slot with ID = SlotId[0]
    rvResult = Pkcs11FuncList->C_OpenSession(SlotId,
                                             (CKF_SERIAL_SESSION | CKF_RW_SESSION),
                                             NULL,
                                             0,
                                             &hSession);
    printf("Open session result: 0x%lx\n", rvResult);
    if (rvResult != CKR_OK)
        return rvResult;
	
	rvResult = test_kuznyechik_key_wrap(hSession);
    if (rvResult != CKR_OK)
        return rvResult;

    printf("test_kuznyechik_key_wrap SUCCESS\n");

	rvResult = test_magma_key_wrap(hSession);
    if (rvResult != CKR_OK)
        return rvResult;

    printf("test_magma_key_wrap SUCCESS\n");

    Pkcs11FuncList->C_CloseSession(hSession);

	Pkcs11FuncList->C_Finalize(NULL);

	if (pSlotList) free(pSlotList);

    printf("SUCCESS\n");
    return rvResult;
}

// Контрольный пример ТК 26 для Кузнечика	
//	Экспортируемый ключ:
unsigned char TC26_K[] = {
	0x88, 0x99, 0xAA, 0xBB, 0xCC, 0xDD, 0xEE, 0xFF, 
	0x00, 0x11, 0x22, 0x33, 0x44, 0x55, 0x66, 0x77, 
	0xFE, 0xDC, 0xBA, 0x98, 0x76, 0x54, 0x32, 0x10, 
	0x01, 0x23, 0x45, 0x67, 0x89, 0xAB, 0xCD, 0xEF,
};
//	Ключ имитовставки KUZNYECHIK MAC:
unsigned char TC26_K_MAC[] = {
	0x08, 0x09, 0x0A, 0x0B, 0x0C, 0x0D, 0x0E, 0x0F, 
	0x00, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 
	0x10, 0x11, 0x12, 0x13, 0x14, 0x15, 0x16, 0x17, 
	0x18, 0x19, 0x1A, 0x1B, 0x1C, 0x1D, 0x1E, 0x1F,
};
//	Ключ шифрования KUZNYECHIK CTR:
unsigned char TC26_K_ENC[] = {
	0x20, 0x21, 0x22, 0x23, 0x24, 0x25, 0x26, 0x27, 
	0x28, 0x29, 0x2A, 0x2B, 0x2C, 0x2D, 0x2E, 0x2F, 
	0x38, 0x39, 0x3A, 0x3B, 0x3C, 0x3D, 0x3E, 0x3F, 
	0x30, 0x31, 0x32, 0x33, 0x34, 0x35, 0x36, 0x37,
};
//	Вектор инициализации IV:
unsigned char TC26_K_IV[] = {
	0x09, 0x09, 0x47, 0x2D, 0xD9, 0xF2, 0x6B, 0xE8,
};
// Имитовставка от IV|K
unsigned char TC26_K_KEYMAC[] = { 
	0x10, 0x02, 0x2A, 0xDE, 0x94, 0xEE, 0x55, 0xB4, 
	0x34, 0xD2, 0x07, 0x7F, 0x5A, 0x13, 0xAF, 0xF4,
};
// Зашифрованный K|KEYMAC
unsigned char TC26_K_KEYENC[] = { 
	0xE3, 0x61, 0x84, 0xE8, 0x4E, 0x8D, 0x73, 0x6F, 
	0xF3, 0x6C, 0xC2, 0xE5, 0xAE, 0x06, 0x5D, 0xC6, 
	0x56, 0xB2, 0x3C, 0x20, 0xF5, 0x49, 0xB0, 0x2F, 
	0xDF, 0xF8, 0x8E, 0x1F, 0x3F, 0x30, 0xD8, 0xC2, 
	0x9A, 0x53, 0xF3, 0xCA, 0x55, 0x4D, 0xBA, 0xD8, 
	0x0D, 0xE1, 0x52, 0xB9, 0xA4, 0x62, 0x5B, 0x32, 
};
// Результат экспорта
unsigned char TC26_K_KEYEXP[]= {
	0x09, 0x09, 0x47, 0x2D, 0xD9, 0xF2, 0x6B, 0xE8, 
	0xE3, 0x61, 0x84, 0xE8, 0x4E, 0x8D, 0x73, 0x6F, 
	0xF3, 0x6C, 0xC2, 0xE5, 0xAE, 0x06, 0x5D, 0xC6, 
	0x56, 0xB2, 0x3C, 0x20, 0xF5, 0x49, 0xB0, 0x2F, 
	0xDF, 0xF8, 0x8E, 0x1F, 0x3F, 0x30, 0xD8, 0xC2, 
	0x9A, 0x53, 0xF3, 0xCA, 0x55, 0x4D, 0xBA, 0xD8, 
	0x0D, 0xE1, 0x52, 0xB9, 0xA4, 0x62, 0x5B, 0x32,
};

CK_RV test_kuznyechik_key_wrap(CK_SESSION_HANDLE hSession)
{
    /**********************************************************/
    static CK_BBOOL blTrue              = CK_TRUE,
             blFalse             = CK_FALSE;
    static CK_ULONG ulKeyType = CKK_KUZNYECHIK,
             ulClass_SecKey      = CKO_SECRET_KEY;
    // template for secret key
    CK_ATTRIBUTE caSecretKeyTemplate[] =
    {
        { CKA_CLASS,            &ulClass_SecKey,      sizeof(CK_ULONG)          },
        { CKA_KEY_TYPE,         &ulKeyType,           sizeof(CK_ULONG)          },
        { CKA_PRIVATE,          &blFalse,             sizeof(CK_BBOOL)          },
        { CKA_ENCRYPT,          &blTrue,              sizeof(CK_BBOOL)          },
        { CKA_DECRYPT,          &blTrue,              sizeof(CK_BBOOL)          },
        { CKA_WRAP,             &blTrue,              sizeof(CK_BBOOL)          },
        { CKA_EXTRACTABLE,      &blTrue,              sizeof(CK_BBOOL)          },
    };
    CK_ULONG ulSecKeyCount =
        sizeof(caSecretKeyTemplate) / sizeof(CK_ATTRIBUTE);
    CK_ATTRIBUTE caTC26_SecretKeyTemplate[] =
    {
		{ CKA_VALUE,			NULL,				  0},
        { CKA_CLASS,            &ulClass_SecKey,      sizeof(CK_ULONG)          },
        { CKA_KEY_TYPE,         &ulKeyType,           sizeof(CK_ULONG)          },
        { CKA_PRIVATE,          &blFalse,             sizeof(CK_BBOOL)          },
        { CKA_ENCRYPT,          &blTrue,              sizeof(CK_BBOOL)          },
        { CKA_DECRYPT,          &blTrue,              sizeof(CK_BBOOL)          },
        { CKA_SIGN,             &blTrue,              sizeof(CK_BBOOL)          },
        { CKA_VERIFY,           &blTrue,              sizeof(CK_BBOOL)          },
        { CKA_WRAP,             &blTrue,              sizeof(CK_BBOOL)          },
        { CKA_EXTRACTABLE,      &blTrue,              sizeof(CK_BBOOL)          },
    };
    CK_ULONG ulTC26_SecKeyCount =
        sizeof(caTC26_SecretKeyTemplate) / sizeof(CK_ATTRIBUTE);

    CK_OBJECT_HANDLE hKeyForWrap         = 0, // key which will be wrapped
                     hEncKey       = 0, // encryption key
                     hMacKey       = 0, // MAC key
                     hUnwrappedKey = 0; // unwrapped key
	CK_OBJECT_HANDLE hTC26_K = CK_INVALID_HANDLE,
		hTC26_K_MAC = CK_INVALID_HANDLE,
		hTC26_K_ENC = CK_INVALID_HANDLE;
	CK_TC26_KEY_WRAP_PARAMS params = {
		TC26_K_IV,
		sizeof(TC26_K_IV),
		CK_INVALID_HANDLE
	};

    CK_MECHANISM cmKeyGenMechanism,         // mechanism for key pair genaration
                 cmWrapMechanism;           // mechanism for key wrap/unwrap

    CK_BYTE_PTR pbWrappedKey;           // wrapped key

    CK_ULONG    ulWrappedKeyLen;            // length of wrapped key
	CK_BYTE     wrap_key_value[32];
	CK_RV rvResult = CKR_OK;

	printf("CKK_KUZNYECHIK\n");
	caTC26_SecretKeyTemplate[0].pValue = TC26_K;
	caTC26_SecretKeyTemplate[0].ulValueLen = sizeof(TC26_K);
	rvResult = Pkcs11FuncList->C_CreateObject(hSession, 
		caTC26_SecretKeyTemplate,
		ulTC26_SecKeyCount,
		&hTC26_K);
    printf("   Sender generate key hTC26_K "
           "for wrapping result: 0x%lx\n",
           rvResult);
    if (rvResult != CKR_OK)
        return rvResult;
    printf("   key handle: %ld\n", (unsigned long) hTC26_K);

	caTC26_SecretKeyTemplate[0].pValue = TC26_K_MAC;
	caTC26_SecretKeyTemplate[0].ulValueLen = sizeof(TC26_K_MAC);
	rvResult = Pkcs11FuncList->C_CreateObject(hSession, 
		caTC26_SecretKeyTemplate,
		ulTC26_SecKeyCount,
		&hTC26_K_MAC);
    printf("   Sender generate key hTC26_K_MAC "
           "for wrapping result: 0x%lx\n",
           rvResult);
    if (rvResult != CKR_OK)
        return rvResult;
    printf("   key handle: %ld\n", (unsigned long) hTC26_K_MAC);

	caTC26_SecretKeyTemplate[0].pValue = TC26_K_ENC;
	caTC26_SecretKeyTemplate[0].ulValueLen = sizeof(TC26_K_ENC);
	rvResult = Pkcs11FuncList->C_CreateObject(hSession, 
		caTC26_SecretKeyTemplate,
		ulTC26_SecKeyCount,
		&hTC26_K_ENC);
    printf("   Sender generate key hTC26_K_ENC "
           "for wrapping result: 0x%lx\n",
           rvResult);
    if (rvResult != CKR_OK)
        return rvResult;
    printf("   key handle: %ld\n", (unsigned long) hTC26_K_ENC);

	params.pIV = TC26_K_IV;
	params.ulIVLen = sizeof(TC26_K_IV);
	params.hMacKey = hTC26_K_MAC;
    cmWrapMechanism.mechanism      = CKM_GOST_CIPHER_KEY_WRAP;
    cmWrapMechanism.pParameter     = &params;
    cmWrapMechanism.ulParameterLen = sizeof(params); 

    rvResult = Pkcs11FuncList->C_WrapKey(hSession,
                                         &cmWrapMechanism,
                                         hTC26_K_ENC,
                                         hTC26_K,
                                         NULL,
                                         &ulWrappedKeyLen);
    if (rvResult == CKR_OK)
    {
		printf("Get wrapped key output buffer size = %ld\n", 
			ulWrappedKeyLen);
        pbWrappedKey = (CK_BYTE_PTR) malloc(ulWrappedKeyLen);
		printf("Wrap key with wrapping keys\n");
        rvResult         = Pkcs11FuncList->C_WrapKey(hSession,
                                                     &cmWrapMechanism,
                                                     hTC26_K_ENC,
                                                     hTC26_K,
                                                     pbWrappedKey,
                                                     &ulWrappedKeyLen);
    }

    printf("   wrap result: 0x%lx\n", rvResult);
    if (rvResult != CKR_OK)
        return rvResult;

    printf("   wrapped key:\n");
    print_bytes(pbWrappedKey, ulWrappedKeyLen);

	if (ulWrappedKeyLen != sizeof(TC26_K_KEYEXP)){
		fprintf(stderr, "Invalid wrapped key length\n");
		rvResult = CKR_FUNCTION_FAILED;
		return rvResult;
	}
	if (memcmp(pbWrappedKey, TC26_K_KEYEXP, sizeof(TC26_K_KEYEXP)) != 0){
		fprintf(stderr, "Invalid wrapped key\n");
		rvResult = CKR_FUNCTION_FAILED;
		return rvResult;
	}
	printf("Wrapped key OK\n");

    // unwrap wrapped key with wrapping key
    printf("Unwrap key hUnwrappedKey "
           "from pbWrappedKey with wrapping keys\n");

    rvResult = Pkcs11FuncList->C_UnwrapKey(hSession,
                                            &cmWrapMechanism,
                                            hTC26_K_ENC,
                                            pbWrappedKey,
                                            ulWrappedKeyLen,
                                            caSecretKeyTemplate,
                                            ulSecKeyCount,
                                            &hUnwrappedKey);
    printf("   unwrap key: result: 0x%lx\n", rvResult);
    if (rvResult != CKR_OK)
        return rvResult;

    printf("   key handle (hUnwrappedKey): %ld\n",
           (unsigned long) hUnwrappedKey);
    {
        CK_BYTE      value[32];
        CK_ATTRIBUTE attr_get[] = {
            { CKA_VALUE, NULL, 0 },
        };
        attr_get[0].pValue     = value;
        attr_get[0].ulValueLen = sizeof(value);
        rvResult               = Pkcs11FuncList->C_GetAttributeValue(
            hSession, hUnwrappedKey,
            attr_get, sizeof(attr_get) / sizeof(CK_ATTRIBUTE));
        printf("   recipient unwrapped key:\n");
        print_bytes(value, attr_get[0].ulValueLen);
		if (attr_get[0].ulValueLen != sizeof(TC26_K)){
			fprintf(stderr, "Invalid unwrapped key length\n");
			rvResult = CKR_FUNCTION_FAILED;
			return rvResult;
		}
		if (memcmp(value, TC26_K, sizeof(TC26_K)) != 0){
			fprintf(stderr, "Invalid unwrapped key\n");
			rvResult = CKR_FUNCTION_FAILED;
			return rvResult;
		}
		printf("Unwrapped key OK\n");
    }
    printf("CKM_GOST_CIPHER_KEY_WRAP test SUCCESS\n");

    rvResult =
        Pkcs11FuncList->C_DestroyObject(hSession, hTC26_K);
    printf("   C_DestroyObject result: 0x%lx\n", rvResult);
    rvResult =
        Pkcs11FuncList->C_DestroyObject(hSession, hTC26_K_ENC);
    printf("   C_DestroyObject result: 0x%lx\n", rvResult);
    rvResult =
        Pkcs11FuncList->C_DestroyObject(hSession, hTC26_K_MAC);
    printf("   C_DestroyObject result: 0x%lx\n", rvResult);
    rvResult =
        Pkcs11FuncList->C_DestroyObject(hSession, hUnwrappedKey);
    printf("   C_DestroyObject result: 0x%lx\n", rvResult);

	// Далее проверяется вариант с генерацией случайных ключей и IV,
	// когда в структуре параметра pIV == NULL.

	printf("Testing with random input data...\n");

    cmKeyGenMechanism.mechanism      = CKM_GOST_CIPHER_KEY_GEN;
    cmKeyGenMechanism.pParameter     = NULL;
    cmKeyGenMechanism.ulParameterLen = 0;

    // generate secret key which will be wrapped and unwrapped
    printf(
        "Sender generate key for wrapping (hKeyForWrap)\n");
    printf("   mechanism type: CKM_GOST_CIPHER_KEY_GEN\n");

    rvResult = Pkcs11FuncList->C_GenerateKey(hSession,
                                             &cmKeyGenMechanism,
                                             caSecretKeyTemplate,
                                             ulSecKeyCount,
                                             &hKeyForWrap);
    printf("   Sender generate key hKeyForWrap "
           "for wrapping result: 0x%lx\n",
           rvResult);
    if (rvResult != CKR_OK)
        return rvResult;

    printf("   key handle: %ld\n", (unsigned long) hKeyForWrap);
    {
        CK_BYTE      value[32];
        CK_ATTRIBUTE attr_get[] = {
            { CKA_VALUE, NULL, 0 },
        };
        attr_get[0].pValue     = value;
        attr_get[0].ulValueLen = sizeof(value);
        rvResult               = Pkcs11FuncList->C_GetAttributeValue(
            hSession, hKeyForWrap,
            attr_get, sizeof(attr_get) / sizeof(CK_ATTRIBUTE));
        printf("   key to wrap:\n");
        print_bytes(value, attr_get[0].ulValueLen);
		memcpy(wrap_key_value, value, attr_get[0].ulValueLen);
    }

    // generate encryption key for wrapping and unwrapping
    printf(
        "Sender encryption key for wrapping and unwrapping (hEncKey)\n");
    printf("   mechanism type: CKM_GOST_CIPHER_KEY_GEN\n");

    rvResult = Pkcs11FuncList->C_GenerateKey(hSession,
                                             &cmKeyGenMechanism,
                                             caSecretKeyTemplate,
                                             ulSecKeyCount,
                                             &hEncKey);
    printf("   Sender generate key hEncKey result: 0x%lx\n",
           rvResult);
    if (rvResult != CKR_OK)
        return rvResult;

    printf("   key handle: %ld\n", (unsigned long) hEncKey);

    // generate MAC key for wrapping and unwrapping
    printf(
        "Sender MAC key for wrapping and unwrapping (hMacKey)\n");
    printf("   mechanism type: CKM_GOST_CIPHER_KEY_GEN\n");

    rvResult = Pkcs11FuncList->C_GenerateKey(hSession,
                                             &cmKeyGenMechanism,
                                             caSecretKeyTemplate,
                                             ulSecKeyCount,
                                             &hMacKey);
    printf("   Sender generate key hMacKey result: 0x%lx\n",
           rvResult);
    if (rvResult != CKR_OK)
        return rvResult;

    printf("   key handle: %ld\n", (unsigned long) hMacKey);

	params.pIV = NULL;
	params.ulIVLen = 0;
	params.hMacKey = hMacKey;
    cmWrapMechanism.mechanism      = CKM_KUZNYECHIK_KEY_WRAP;
    cmWrapMechanism.pParameter     = &params;
    cmWrapMechanism.ulParameterLen = sizeof(params); 

    rvResult = Pkcs11FuncList->C_WrapKey(hSession,
                                         &cmWrapMechanism,
                                         hEncKey,
                                         hKeyForWrap,
                                         NULL,
                                         &ulWrappedKeyLen);
    if (rvResult == CKR_OK)
    {
		printf("Get wrapped key output buffer size = %ld\n", 
			ulWrappedKeyLen);
        pbWrappedKey = (CK_BYTE_PTR) malloc(ulWrappedKeyLen);
		printf("Wrap key with wrapping keys\n");
        rvResult         = Pkcs11FuncList->C_WrapKey(hSession,
                                                     &cmWrapMechanism,
                                                     hEncKey,
                                                     hKeyForWrap,
                                                     pbWrappedKey,
                                                     &ulWrappedKeyLen);
    }

    printf("   wrap result: 0x%lx\n", rvResult);
    if (rvResult != CKR_OK)
        return rvResult;

    printf("   wrapped key:\n");
    print_bytes(pbWrappedKey, ulWrappedKeyLen);

    // unwrap wrapped key with wrapping keys
    printf("Unwrap key hUnwrappedKey "
           "from pbWrappedKey with wrapping keys\n");

    rvResult = Pkcs11FuncList->C_UnwrapKey(hSession,
                                            &cmWrapMechanism,
                                            hEncKey,
                                            pbWrappedKey,
                                            ulWrappedKeyLen,
                                            caSecretKeyTemplate,
                                            ulSecKeyCount,
                                            &hUnwrappedKey);
    printf("   unwrap key: result: 0x%lx\n", rvResult);
    if (rvResult != CKR_OK)
        return rvResult;

    printf("   key handle (hUnwrappedKey): %ld\n",
           (unsigned long) hUnwrappedKey);
    {
        CK_BYTE      value[32];
        CK_ATTRIBUTE attr_get[] = {
            { CKA_VALUE, NULL, 0 },
        };
        attr_get[0].pValue     = value;
        attr_get[0].ulValueLen = sizeof(value);
        rvResult               = Pkcs11FuncList->C_GetAttributeValue(
            hSession, hUnwrappedKey,
            attr_get, sizeof(attr_get) / sizeof(CK_ATTRIBUTE));
        printf("   recipient unwrapped key:\n");
        print_bytes(value, attr_get[0].ulValueLen);
		if (attr_get[0].ulValueLen != sizeof(wrap_key_value)){
			fprintf(stderr, "Invalid unwrapped key length\n");
			rvResult = CKR_FUNCTION_FAILED;
			return rvResult;
		}
		if (memcmp(value, wrap_key_value, sizeof(wrap_key_value)) != 0){
			fprintf(stderr, "Invalid unwrapped key\n");
			rvResult = CKR_FUNCTION_FAILED;
			return rvResult;
		}
		printf("Unwrapped key OK\n");
    }

    rvResult =
        Pkcs11FuncList->C_DestroyObject(hSession, hKeyForWrap);
    printf("   C_DestroyObject result: 0x%lx\n", rvResult);
    rvResult =
        Pkcs11FuncList->C_DestroyObject(hSession, hEncKey);
    printf("   C_DestroyObject result: 0x%lx\n", rvResult);
    rvResult =
        Pkcs11FuncList->C_DestroyObject(hSession, hMacKey);
    printf("   C_DestroyObject result: 0x%lx\n", rvResult);
    rvResult =
        Pkcs11FuncList->C_DestroyObject(hSession, hUnwrappedKey);
    printf("   C_DestroyObject result: 0x%lx\n", rvResult);
	if (pbWrappedKey) free(pbWrappedKey);

    printf("SUCCESS\n");
	return CKR_OK;
}

// Контрольный пример ТК 26 для Магмы	
//	Экспортируемый ключ:
unsigned char TC26_M[] = {
	0x88, 0x99, 0xAA, 0xBB, 0xCC, 0xDD, 0xEE, 0xFF, 
	0x00, 0x11, 0x22, 0x33, 0x44, 0x55, 0x66, 0x77, 
	0xFE, 0xDC, 0xBA, 0x98, 0x76, 0x54, 0x32, 0x10, 
	0x01, 0x23, 0x45, 0x67, 0x89, 0xAB, 0xCD, 0xEF,
};
//	Ключ имитовставки MAGMA MAC:
unsigned char TC26_M_MAC[] = {
	0x08, 0x09, 0x0A, 0x0B, 0x0C, 0x0D, 0x0E, 0x0F, 
	0x00, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 
	0x10, 0x11, 0x12, 0x13, 0x14, 0x15, 0x16, 0x17, 
	0x18, 0x19, 0x1A, 0x1B, 0x1C, 0x1D, 0x1E, 0x1F,
};
//	Ключ шифрования MAGMA CTR:
unsigned char TC26_M_ENC[] = {
	0x20, 0x21, 0x22, 0x23, 0x24, 0x25, 0x26, 0x27, 
	0x28, 0x29, 0x2A, 0x2B, 0x2C, 0x2D, 0x2E, 0x2F, 
	0x38, 0x39, 0x3A, 0x3B, 0x3C, 0x3D, 0x3E, 0x3F, 
	0x30, 0x31, 0x32, 0x33, 0x34, 0x35, 0x36, 0x37,
};
//	Вектор инициализации IV:
unsigned char TC26_M_IV[] = {
	0x67, 0xBE, 0xD6, 0x54,
};
// Имитовставка от IV|K
unsigned char TC26_M_KEYMAC[] = { 
	0x75, 0xA7, 0x66, 0x18, 0xE9, 0x0F, 0x49, 0x73,
};
// Зашифрованный K|KEYMAC
unsigned char TC26_M_KEYENC[] = { 
	0xCF, 0xD5, 0xA1, 0x2D, 0x5B, 0x81, 0xB6, 0xE1, 
	0xE9, 0x9C, 0x91, 0x6D, 0x07, 0x90, 0x0C, 0x6A, 
	0xC1, 0x27, 0x03, 0xFB, 0x3A, 0xBD, 0xED, 0x55, 
	0x56, 0x7B, 0xF3, 0x74, 0x2C, 0x89, 0x9C, 0x75, 
	0x5D, 0xAF, 0xE7, 0xB4, 0x2E, 0x3A, 0x8B, 0xD9, 
};
// Результат экспорта
unsigned char TC26_M_KEYEXP[]= {
	0x67, 0xBE, 0xD6, 0x54, 0xCF, 0xD5, 0xA1, 0x2D, 
	0x5B, 0x81, 0xB6, 0xE1, 0xE9, 0x9C, 0x91, 0x6D, 
	0x07, 0x90, 0x0C, 0x6A, 0xC1, 0x27, 0x03, 0xFB, 
	0x3A, 0xBD, 0xED, 0x55, 0x56, 0x7B, 0xF3, 0x74, 
	0x2C, 0x89, 0x9C, 0x75, 0x5D, 0xAF, 0xE7, 0xB4, 
	0x2E, 0x3A, 0x8B, 0xD9,
};

CK_RV test_magma_key_wrap(CK_SESSION_HANDLE hSession)
{
    CK_RV                rvResult;

    /**********************************************************/
    static CK_BBOOL blTrue              = CK_TRUE,
             blFalse             = CK_FALSE;
//    static CK_ULONG ulKeyType = CKK_GOST28147,
    static CK_ULONG ulKeyType = CKK_MAGMA,
             ulClass_SecKey      = CKO_SECRET_KEY;
    // id-tc26-gost-28147-param-Z (1.2.643.7.1.2.5.1.1)
//    static CK_BYTE         oid_Z[] = {
//        0x06, 0x09, 0x2a, 0x85, 0x03, 0x07, 0x01, 0x02, 0x05, 0x01, 0x01
//    };
    // template for secret GOST key
    CK_ATTRIBUTE caGOST_SecretKeyTemplate[] =
    {
        { CKA_CLASS,            &ulClass_SecKey,      sizeof(CK_ULONG)          },
        { CKA_KEY_TYPE,         &ulKeyType,			  sizeof(CK_ULONG)          },
        { CKA_PRIVATE,          &blFalse,             sizeof(CK_BBOOL)          },
        { CKA_ENCRYPT,          &blTrue,              sizeof(CK_BBOOL)          },
        { CKA_DECRYPT,          &blTrue,              sizeof(CK_BBOOL)          },
        { CKA_WRAP,             &blTrue,              sizeof(CK_BBOOL)          },
        { CKA_EXTRACTABLE,      &blTrue,              sizeof(CK_BBOOL)          },
//        { CKA_GOST28147_PARAMS, oid_Z,    sizeof(oid_Z) },
    };
    CK_ULONG ulSecKeyCount =
        sizeof(caGOST_SecretKeyTemplate) / sizeof(CK_ATTRIBUTE);
    CK_ATTRIBUTE caTC26_SecretKeyTemplate[] =
    {
		{ CKA_VALUE,			NULL,				  0},
        { CKA_CLASS,            &ulClass_SecKey,      sizeof(CK_ULONG)          },
        { CKA_KEY_TYPE,         &ulKeyType,			  sizeof(CK_ULONG)          },
        { CKA_PRIVATE,          &blFalse,             sizeof(CK_BBOOL)          },
        { CKA_ENCRYPT,          &blTrue,              sizeof(CK_BBOOL)          },
        { CKA_DECRYPT,          &blTrue,              sizeof(CK_BBOOL)          },
        { CKA_SIGN,             &blTrue,              sizeof(CK_BBOOL)          },
        { CKA_VERIFY,           &blTrue,              sizeof(CK_BBOOL)          },
        { CKA_WRAP,             &blTrue,              sizeof(CK_BBOOL)          },
        { CKA_EXTRACTABLE,      &blTrue,              sizeof(CK_BBOOL)          },
//        { CKA_GOST28147_PARAMS, oid_Z,    sizeof(oid_Z) },
    };
    CK_ULONG ulTC26_SecKeyCount =
        sizeof(caTC26_SecretKeyTemplate) / sizeof(CK_ATTRIBUTE);

    CK_OBJECT_HANDLE hKeyForWrap         = 0, // key which will be wrapped
                     hEncKey       = 0, // encryption key
                     hMacKey       = 0, // MAC key
                     hUnwrappedKey = 0; // unwrapped key
	CK_OBJECT_HANDLE hTC26_M = CK_INVALID_HANDLE,
		hTC26_M_MAC = CK_INVALID_HANDLE,
		hTC26_M_ENC = CK_INVALID_HANDLE;
	CK_TC26_KEY_WRAP_PARAMS params = {
		TC26_M_IV,
		sizeof(TC26_M_IV),
		CK_INVALID_HANDLE
	};

    CK_MECHANISM cmKeyGenMechanism,         // mechanism for key pair genaration
                 cmWrapMechanism;           // mechanism for key wrap/unwrap

    CK_BYTE_PTR pbWrappedKey;           // wrapped key

    CK_ULONG    ulWrappedKeyLen;            // length of wrapped key
	CK_BYTE     wrap_key_value[32];


	printf("CKK_MAGMA\n");

	caTC26_SecretKeyTemplate[0].pValue = TC26_M;
	caTC26_SecretKeyTemplate[0].ulValueLen = sizeof(TC26_M);
	rvResult = Pkcs11FuncList->C_CreateObject(hSession, 
		caTC26_SecretKeyTemplate,
		ulTC26_SecKeyCount,
		&hTC26_M);
    printf("   Sender generate key hTC26_M "
           "for wrapping result: 0x%lx\n",
           rvResult);
    if (rvResult != CKR_OK)
        return rvResult;
    printf("   key handle: %ld\n", (unsigned long) hTC26_M);

	caTC26_SecretKeyTemplate[0].pValue = TC26_M_MAC;
	caTC26_SecretKeyTemplate[0].ulValueLen = sizeof(TC26_M_MAC);
	rvResult = Pkcs11FuncList->C_CreateObject(hSession, 
		caTC26_SecretKeyTemplate,
		ulTC26_SecKeyCount,
		&hTC26_M_MAC);
    printf("   Sender generate key hTC26_M_MAC "
           "for wrapping result: 0x%lx\n",
           rvResult);
    if (rvResult != CKR_OK)
        return rvResult;
    printf("   key handle: %ld\n", (unsigned long) hTC26_M_MAC);

	caTC26_SecretKeyTemplate[0].pValue = TC26_M_ENC;
	caTC26_SecretKeyTemplate[0].ulValueLen = sizeof(TC26_M_ENC);
	rvResult = Pkcs11FuncList->C_CreateObject(hSession, 
		caTC26_SecretKeyTemplate,
		ulTC26_SecKeyCount,
		&hTC26_M_ENC);
    printf("   Sender generate key hTC26_M_ENC "
           "for wrapping result: 0x%lx\n",
           rvResult);
    if (rvResult != CKR_OK)
        return rvResult;
    printf("   key handle: %ld\n", (unsigned long) hTC26_M_ENC);

	params.pIV = TC26_M_IV;
	params.ulIVLen = sizeof(TC26_M_IV);
	params.hMacKey = hTC26_M_MAC;
    cmWrapMechanism.mechanism      = CKM_GOST_CIPHER_KEY_WRAP;
    cmWrapMechanism.pParameter     = &params;
    cmWrapMechanism.ulParameterLen = sizeof(params); 

    rvResult = Pkcs11FuncList->C_WrapKey(hSession,
                                         &cmWrapMechanism,
                                         hTC26_M_ENC,
                                         hTC26_M,
                                         NULL,
                                         &ulWrappedKeyLen);
    if (rvResult == CKR_OK)
    {
		printf("Get wrapped key output buffer size = %ld\n", 
			ulWrappedKeyLen);
        pbWrappedKey = (CK_BYTE_PTR) malloc(ulWrappedKeyLen);
		printf("Wrap key with wrapping keys\n");
        rvResult         = Pkcs11FuncList->C_WrapKey(hSession,
                                                     &cmWrapMechanism,
                                                     hTC26_M_ENC,
                                                     hTC26_M,
                                                     pbWrappedKey,
                                                     &ulWrappedKeyLen);
    }

    printf("   wrap result: 0x%lx\n", rvResult);
    if (rvResult != CKR_OK)
        return rvResult;

    printf("   wrapped key:\n");
    print_bytes(pbWrappedKey, ulWrappedKeyLen);

	if (ulWrappedKeyLen != sizeof(TC26_M_KEYEXP)){
		fprintf(stderr, "Invalid wrapped key length\n");
		rvResult = CKR_FUNCTION_FAILED;
		return rvResult;
	}
	if (memcmp(pbWrappedKey, TC26_M_KEYEXP, sizeof(TC26_M_KEYEXP)) != 0){
		fprintf(stderr, "Invalid wrapped key\n");
		rvResult = CKR_FUNCTION_FAILED;
		return rvResult;
	}
	printf("Wrapped key OK\n");

    // unwrap wrapped key with wrapping key
    printf("Unwrap key hUnwrappedKey "
           "from pbWrappedKey with wrapping keys\n");

    rvResult = Pkcs11FuncList->C_UnwrapKey(hSession,
                                            &cmWrapMechanism,
                                            hTC26_M_ENC,
                                            pbWrappedKey,
                                            ulWrappedKeyLen,
                                            caGOST_SecretKeyTemplate,
                                            ulSecKeyCount,
                                            &hUnwrappedKey);
    printf("   unwrap key: result: 0x%lx\n", rvResult);
    if (rvResult != CKR_OK)
        return rvResult;

    printf("   key handle (hUnwrappedKey): %ld\n",
           (unsigned long) hUnwrappedKey);
    {
        CK_BYTE      value[32];
        CK_ATTRIBUTE attr_get[] = {
            { CKA_VALUE, NULL, 0 },
        };
        attr_get[0].pValue     = value;
        attr_get[0].ulValueLen = sizeof(value);
        rvResult               = Pkcs11FuncList->C_GetAttributeValue(
            hSession, hUnwrappedKey,
            attr_get, sizeof(attr_get) / sizeof(CK_ATTRIBUTE));
        printf("   recipient unwrapped key:\n");
        print_bytes(value, attr_get[0].ulValueLen);
		if (attr_get[0].ulValueLen != sizeof(TC26_M)){
			fprintf(stderr, "Invalid unwrapped key length\n");
			rvResult = CKR_FUNCTION_FAILED;
			return rvResult;
		}
		if (memcmp(value, TC26_M, sizeof(TC26_M)) != 0){
			fprintf(stderr, "Invalid unwrapped key\n");
			rvResult = CKR_FUNCTION_FAILED;
			return rvResult;
		}
		printf("Unwrapped key OK\n");
    }
    printf("CKM_GOST_CIPHER_KEY_WRAP test SUCCESS\n");

    rvResult =
        Pkcs11FuncList->C_DestroyObject(hSession, hTC26_M);
    printf("   C_DestroyObject result: 0x%lx\n", rvResult);
    rvResult =
        Pkcs11FuncList->C_DestroyObject(hSession, hTC26_M_ENC);
    printf("   C_DestroyObject result: 0x%lx\n", rvResult);
    rvResult =
        Pkcs11FuncList->C_DestroyObject(hSession, hTC26_M_MAC);
    printf("   C_DestroyObject result: 0x%lx\n", rvResult);
    rvResult =
        Pkcs11FuncList->C_DestroyObject(hSession, hUnwrappedKey);
    printf("   C_DestroyObject result: 0x%lx\n", rvResult);

	// Далее проверяется вариант с генерацией случайных ключей и IV,
	// когда в структуре параметра pIV == NULL.

	printf("Testing with random input data...\n");

    cmKeyGenMechanism.mechanism      = CKM_GOST_CIPHER_KEY_GEN;
    cmKeyGenMechanism.pParameter     = NULL;
    cmKeyGenMechanism.ulParameterLen = 0;

    // generate secret key which will be wrapped and unwrapped
    printf(
        "Sender generate key for wrapping (hKeyForWrap)\n");
    printf("   mechanism type: CKM_GOST_CIPHER_KEY_GEN\n");

    rvResult = Pkcs11FuncList->C_GenerateKey(hSession,
                                             &cmKeyGenMechanism,
                                             caGOST_SecretKeyTemplate,
                                             ulSecKeyCount,
                                             &hKeyForWrap);
    printf("   Sender generate key hKeyForWrap "
           "for wrapping result: 0x%lx\n",
           rvResult);
    if (rvResult != CKR_OK)
        return rvResult;

    printf("   key handle: %ld\n", (unsigned long) hKeyForWrap);
    {
        CK_BYTE      value[32];
        CK_ATTRIBUTE attr_get[] = {
            { CKA_VALUE, NULL, 0 },
        };
        attr_get[0].pValue     = value;
        attr_get[0].ulValueLen = sizeof(value);
        rvResult               = Pkcs11FuncList->C_GetAttributeValue(
            hSession, hKeyForWrap,
            attr_get, sizeof(attr_get) / sizeof(CK_ATTRIBUTE));
        printf("   key to wrap:\n");
        print_bytes(value, attr_get[0].ulValueLen);
		memcpy(wrap_key_value, value, attr_get[0].ulValueLen);
    }

    // generate encryption key for wrapping and unwrapping
    printf(
        "Sender encryption key for wrapping and unwrapping (hEncKey)\n");
    printf("   mechanism type: CKM_GOST_CIPHER_KEY_GEN\n");

    rvResult = Pkcs11FuncList->C_GenerateKey(hSession,
                                             &cmKeyGenMechanism,
                                             caGOST_SecretKeyTemplate,
                                             ulSecKeyCount,
                                             &hEncKey);
    printf("   Sender generate key hEncKey result: 0x%lx\n",
           rvResult);
    if (rvResult != CKR_OK)
        return rvResult;

    printf("   key handle: %ld\n", (unsigned long) hEncKey);

    // generate MAC key for wrapping and unwrapping
    printf(
        "Sender MAC key for wrapping and unwrapping (hMacKey)\n");
    printf("   mechanism type: CKM_GOST_CIPHER_KEY_GEN\n");

    rvResult = Pkcs11FuncList->C_GenerateKey(hSession,
                                             &cmKeyGenMechanism,
                                             caGOST_SecretKeyTemplate,
                                             ulSecKeyCount,
                                             &hMacKey);
    printf("   Sender generate key hMacKey result: 0x%lx\n",
           rvResult);
    if (rvResult != CKR_OK)
        return rvResult;

    printf("   key handle: %ld\n", (unsigned long) hMacKey);

	params.pIV = NULL;
	params.ulIVLen = 0;
	params.hMacKey = hMacKey;
    cmWrapMechanism.mechanism      = CKM_GOST_CIPHER_KEY_WRAP;
    cmWrapMechanism.pParameter     = &params;
    cmWrapMechanism.ulParameterLen = sizeof(params); 

    rvResult = Pkcs11FuncList->C_WrapKey(hSession,
                                         &cmWrapMechanism,
                                         hEncKey,
                                         hKeyForWrap,
                                         NULL,
                                         &ulWrappedKeyLen);
    if (rvResult == CKR_OK)
    {
		printf("Get wrapped key output buffer size = %ld\n", 
			ulWrappedKeyLen);
        pbWrappedKey = (CK_BYTE_PTR) malloc(ulWrappedKeyLen);
		printf("Wrap key with wrapping keys\n");
        rvResult         = Pkcs11FuncList->C_WrapKey(hSession,
                                                     &cmWrapMechanism,
                                                     hEncKey,
                                                     hKeyForWrap,
                                                     pbWrappedKey,
                                                     &ulWrappedKeyLen);
    }

    printf("   wrap result: 0x%lx\n", rvResult);
    if (rvResult != CKR_OK)
        return rvResult;

    printf("   wrapped key:\n");
    print_bytes(pbWrappedKey, ulWrappedKeyLen);

    // unwrap wrapped key with wrapping keys
    printf("Unwrap key hUnwrappedKey "
           "from pbWrappedKey with wrapping keys\n");

    rvResult = Pkcs11FuncList->C_UnwrapKey(hSession,
                                            &cmWrapMechanism,
                                            hEncKey,
                                            pbWrappedKey,
                                            ulWrappedKeyLen,
                                            caGOST_SecretKeyTemplate,
                                            ulSecKeyCount,
                                            &hUnwrappedKey);
    printf("   unwrap key: result: 0x%lx\n", rvResult);
    if (rvResult != CKR_OK)
        return rvResult;

    printf("   key handle (hUnwrappedKey): %ld\n",
           (unsigned long) hUnwrappedKey);
    {
        CK_BYTE      value[32];
        CK_ATTRIBUTE attr_get[] = {
            { CKA_VALUE, NULL, 0 },
        };
        attr_get[0].pValue     = value;
        attr_get[0].ulValueLen = sizeof(value);
        rvResult               = Pkcs11FuncList->C_GetAttributeValue(
            hSession, hUnwrappedKey,
            attr_get, sizeof(attr_get) / sizeof(CK_ATTRIBUTE));
        printf("   recipient unwrapped key:\n");
        print_bytes(value, attr_get[0].ulValueLen);
		if (attr_get[0].ulValueLen != sizeof(wrap_key_value)){
			fprintf(stderr, "Invalid unwrapped key length\n");
			rvResult = CKR_FUNCTION_FAILED;
			return rvResult;
		}
		if (memcmp(value, wrap_key_value, sizeof(wrap_key_value)) != 0){
			fprintf(stderr, "Invalid unwrapped key\n");
			rvResult = CKR_FUNCTION_FAILED;
			return rvResult;
		}
		printf("Unwrapped key OK\n");
    }
    printf("CKM_GOST_CIPHER_KEY_WRAP test SUCCESS\n");

    rvResult =
        Pkcs11FuncList->C_DestroyObject(hSession, hKeyForWrap);
    printf("   C_DestroyObject result: 0x%lx\n", rvResult);
    rvResult =
        Pkcs11FuncList->C_DestroyObject(hSession, hEncKey);
    printf("   C_DestroyObject result: 0x%lx\n", rvResult);
    rvResult =
        Pkcs11FuncList->C_DestroyObject(hSession, hMacKey);
    printf("   C_DestroyObject result: 0x%lx\n", rvResult);
    rvResult =
        Pkcs11FuncList->C_DestroyObject(hSession, hUnwrappedKey);
    printf("   C_DestroyObject result: 0x%lx\n", rvResult);

	if (pbWrappedKey) free(pbWrappedKey);

    printf("SUCCESS\n");
    return 0;
}
