#include "test_common.h"

CK_RV test_crypto();
CK_RV test_gost28147(CK_SESSION_HANDLE sess, CK_SESSION_HANDLE sess2);

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

    rc = test_main(argc, argv);
    PROCESS_ERROR("test_main", rc, end);

    rc = test_crypto();
    PROCESS_ERROR("test_crypto", rc, end);
 end:
    return rc;
}

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

    rc = funcs->C_OpenSession(SlotId,
                              CKF_RW_SESSION | CKF_SERIAL_SESSION,
                              NULL_PTR, NULL_PTR, &hSession);
    PROCESS_ERROR("C_OpenSession 1", rc, end);

    rc = funcs->C_OpenSession(SlotId,
                              CKF_RW_SESSION | CKF_SERIAL_SESSION,
                              NULL_PTR, NULL_PTR, &hSession2);
    PROCESS_ERROR("C_OpenSession 2", rc, end);

    rc = funcs->C_GetMechanismInfo(SlotId, CKM_GOST28147, &minfo);
    PROCESS_ERROR("C_GetMechanismInfo", rc, end);

    fprintf(stderr,
            "\n==================== CKM_GOST28147 test ==================\n");
    rc = test_gost28147(hSession, hSession2);
    PROCESS_ERROR("test_gost28147", rc, end);
    fprintf(stderr,
            "\n===================== CKM_GOST28147 test SUCCESS =========\n");

    rc = funcs->C_CloseSession(hSession2);
    PROCESS_ERROR("C_CloseSession 2", rc, end);

    rc = funcs->C_CloseSession(hSession);
    PROCESS_ERROR("C_CloseSession 1", rc, end);

 end:
    return rc;
}

CK_RV test_gost28147(CK_SESSION_HANDLE sess, CK_SESSION_HANDLE sess2)
{
    CK_RV                  rc = CKR_OK;
    CK_ULONG               len, state_len;
    CK_BYTE                value[1024], value2[1024];
    CK_OBJECT_HANDLE       keyh = CK_INVALID_HANDLE;
    CK_BYTE                *state;
    CK_MECHANISM           mechanism_desc = { CKM_GOST28147, NULL, 0 };
    CK_MECHANISM_PTR       mechanism      = &mechanism_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 param_oclass = CKO_DOMAIN_PARAMETERS;
    static CK_KEY_TYPE     key_type     = CKK_GOST28147;

    static CK_BYTE         key[] = {
        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,
    };
    static CK_BYTE         data[] = {
        0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
        0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
    };
    static CK_BYTE         data2[200];

    //par_cipher_0
    static CK_BYTE oid1[] = {
        0x06, 0x07, 0x2a, 0x85, 0x03, 0x02, 0x02, 0x1f, 0x00,
    };
    static CK_BYTE et1[] = {
        0x08, 0x79, 0x76, 0x86, 0xee, 0xbb, 0x99, 0xe7,
        0x7f, 0xa3, 0x75, 0xff, 0x6b, 0x7c, 0x27, 0x6f,
    };

    //par_cipher_A
    static CK_BYTE oid2[] = {
        0x06, 0x07, 0x2a, 0x85, 0x03, 0x02, 0x02, 0x1f, 0x01,
    };
    static CK_BYTE et2[] = {
        0x49, 0x11, 0x96, 0x11, 0xad, 0x19, 0x12, 0x52,
        0x7c, 0x49, 0x5f, 0x23, 0xb9, 0x23, 0x62, 0xd9,
    };

    //par_cipher_B
    static CK_BYTE oid3[] = {
        0x06, 0x07, 0x2a, 0x85, 0x03, 0x02, 0x02, 0x1f, 0x02,
    };
    static CK_BYTE et3[] = {
        0xab, 0x45, 0x07, 0xb4, 0x96, 0xc5, 0x5b, 0xfd,
        0x61, 0x16, 0x64, 0x9d, 0x43, 0x87, 0x42, 0x53,
    };

    //par_cipher_C
    static CK_BYTE oid4[] = {
        0x06, 0x07, 0x2a, 0x85, 0x03, 0x02, 0x02, 0x1f, 0x03,
    };
    static CK_BYTE et4[] = {
        0x98, 0x7c, 0x0c, 0x18, 0x3c, 0x64, 0xd5, 0x4b,
        0x75, 0x8e, 0xaf, 0x67, 0x52, 0x87, 0x15, 0x7b,
    };

    //par_cipher_D
    static CK_BYTE      oid5[] = {
        0x06, 0x07, 0x2a, 0x85, 0x03, 0x02, 0x02, 0x1f, 0x04,
    };
    static CK_BYTE      et5[] = {
        0xa2, 0x4c, 0x7c, 0x4e, 0x91, 0xed, 0x19, 0x8c,
        0xdc, 0xf5, 0x1f, 0x07, 0x5c, 0x80, 0xf6, 0xa4,
    };
    //par_cipher_Z (TC 26 Z)
    static CK_BYTE      oid6[] = {
        0x06, 0x09, 0x2a, 0x85, 0x03, 0x07, 0x01, 0x02, 0x05, 0x01, 0x01
    };
    static CK_BYTE      et6[] = {
        0xea, 0x2f, 0x3f, 0x5f, 0xdb, 0x6f, 0x19, 0x01,
        0xba, 0x3d, 0x01, 0xcd, 0x2c, 0x64, 0x6d, 0xcb,
    };

    static CK_ATTRIBUTE key_template[] = {
        { CKA_CLASS,           &oclass,   sizeof(oclass)   },
        { CKA_KEY_TYPE,        &key_type, sizeof(key_type) },
        { CKA_GOST28147PARAMS, oid1,      sizeof(oid1)     },
        { CKA_ENCRYPT,         &ltrue,    sizeof(CK_BBOOL) },
        { CKA_DECRYPT,         &ltrue,    sizeof(CK_BBOOL) },
        { CKA_VALUE,           key,       sizeof(key)      }
    };
    static CK_BYTE      ivec[8] = { 1 };

    printf("Zero Data\n");
    printf("Param Set 0\n");
    rc = funcs->C_CreateObject(sess,
                               key_template, sizeof(key_template) / sizeof(CK_ATTRIBUTE), &keyh);
    PROCESS_ERROR("C_CreateObject", rc, end);

	printf("Testing returned error codes\n");
	printf("Try to call C_Encrypt without C_EncryptInit\n");
    len = sizeof(data);
    rc  = funcs->C_Encrypt(sess, data, len, value, &len);
	if (rc != CKR_OPERATION_NOT_INITIALIZED) {
		fprintf(stderr, "Can't call C_Encrypt without C_EncryptInit\n");
		rc = CKR_FUNCTION_FAILED;
		goto end;
	} else {
		printf("Uninitialized C_Encrypt returned CKR_OPERATION_NOT_INITIALIZED - it's OK\n");
	}

    mechanism->pParameter     = ivec;
    mechanism->ulParameterLen = sizeof(ivec);
    rc                        = funcs->C_EncryptInit(sess, mechanism, keyh);
    PROCESS_ERROR("C_EncryptInit", rc, end);

	printf("Try to call C_EncryptInit twice\n");
    rc                        = funcs->C_EncryptInit(sess, mechanism, keyh);
	if (rc != CKR_OPERATION_ACTIVE) {
		fprintf(stderr, "Can't call C_EncryptInit while active\n");
		rc = CKR_FUNCTION_FAILED;
		goto end;
	} else {
		printf("Second C_EncryptInit returned CKR_OPERATION_ACTIVE - it's OK\n");
		printf("The encryption context should not be changed\n");
	}

	printf("Try to set too small output buffer to C_Encrypt\n");
    len = sizeof(data)-1;
    rc  = funcs->C_Encrypt(sess, data, sizeof(data), value, &len);
	if (rc != CKR_BUFFER_TOO_SMALL) {
		fprintf(stderr, "Can't set too small output buffer for C_Encrypt\n");
		rc = CKR_FUNCTION_FAILED;
		goto end;
	} else {
		printf("C_Encrypt returned CKR_BUFFER_TOO_SMALL - it's OK\n");
		printf("The encryption context should not be changed\n");
	}

	printf("Try to get output buffer size only\n");
	printf("The encryption context should not be changed\n");
    len = 0;
    rc  = funcs->C_Encrypt(sess, data, sizeof(data), NULL, &len);
    PROCESS_ERROR("C_EncryptInit", rc, end);

	printf("Use correct output buffer and call C_Encrypt once more\n");
    rc  = funcs->C_Encrypt(sess, data, len, value, &len);
    PROCESS_ERROR("C_Encrypt", rc, end);

    CHECK(value, len, et1);

	printf("Try to call yet another C_Encrypt after C_Encrypt\n");
    rc  = funcs->C_Encrypt(sess, data, len, value, &len);
	if (rc != CKR_OPERATION_NOT_INITIALIZED) {
		fprintf(stderr, "Can't call C_Encrypt just after C_Encrypt\n");
		rc = CKR_FUNCTION_FAILED;
		goto end;
	} else {
		printf("Finalized C_Encrypt returned CKR_OPERATION_NOT_INITIALIZED - it's OK\n");
	}

    rc = funcs->C_DestroyObject(sess, keyh);
    PROCESS_ERROR("C_DestroyObject", rc, end);

    printf("Param Set A\n");
    key_template[2].pValue     = oid2;
    key_template[2].ulValueLen = sizeof(oid2);
    rc                         = funcs->C_CreateObject(sess,
                                                       key_template, sizeof(key_template) / sizeof(CK_ATTRIBUTE), &keyh);
    PROCESS_ERROR("C_DestroyObject", rc, end);

    rc = funcs->C_EncryptInit(sess, mechanism, keyh);
    PROCESS_ERROR("C_EncryptInit", rc, end);

    len = sizeof(value);
    rc  = funcs->C_Encrypt(sess, data, sizeof(data), value, &len);
    PROCESS_ERROR("C_Encrypt", rc, end);

    CHECK(value, len, et2);

    rc = funcs->C_DestroyObject(sess, keyh);
    PROCESS_ERROR("C_DestroyObject", rc, end);

    printf("Param Set B\n");
    key_template[2].pValue     = oid3;
    key_template[2].ulValueLen = sizeof(oid3);
    rc                         = funcs->C_CreateObject(sess,
                                                       key_template,
                                                       sizeof(key_template) / sizeof(CK_ATTRIBUTE), &keyh);
    PROCESS_ERROR("C_CreateObject", rc, end);

    rc = funcs->C_EncryptInit(sess, mechanism, keyh);
    PROCESS_ERROR("C_EncryptInit", rc, end);

    len = sizeof(value);
    rc  = funcs->C_Encrypt(sess, data, sizeof(data), value, &len);
    PROCESS_ERROR("C_Encrypt", rc, end);

    CHECK(value, len, et3);

    rc = funcs->C_DestroyObject(sess, keyh);
    PROCESS_ERROR("C_DestroyObject", rc, end);

    printf("Param Set C\n");
    key_template[2].pValue     = oid4;
    key_template[2].ulValueLen = sizeof(oid4);
    rc                         = funcs->C_CreateObject(sess,
                                                       key_template,
                                                       sizeof(key_template) / sizeof(CK_ATTRIBUTE), &keyh);
    PROCESS_ERROR("C_CreateObject", rc, end);

    rc = funcs->C_EncryptInit(sess, mechanism, keyh);
    PROCESS_ERROR("C_EncryptInit", rc, end);

    len = sizeof(value);
    rc  = funcs->C_Encrypt(sess, data, sizeof(data), value, &len);
    PROCESS_ERROR("C_Encrypt", rc, end);

    CHECK(value, len, et4);

    rc = funcs->C_DestroyObject(sess, keyh);
    PROCESS_ERROR("C_DestroyObject", rc, end);

    printf("Param Set D\n");
    key_template[2].pValue     = oid5;
    key_template[2].ulValueLen = sizeof(oid5);
    rc                         = funcs->C_CreateObject(sess,
                                                       key_template,
                                                       sizeof(key_template) / sizeof(CK_ATTRIBUTE), &keyh);
    PROCESS_ERROR("C_CreateObject", rc, end);

    rc = funcs->C_EncryptInit(sess, mechanism, keyh);
    PROCESS_ERROR("C_EncryptInit", rc, end);

    len = sizeof(value);
    rc  = funcs->C_Encrypt(sess, data, sizeof(data), value, &len);
    PROCESS_ERROR("C_Encrypt", rc, end);

    CHECK(value, len, et5);

    rc = funcs->C_DestroyObject(sess, keyh);
    PROCESS_ERROR("C_DestroyObject", rc, end);

    printf("Param Set Z\n");
    key_template[2].pValue     = oid6;
    key_template[2].ulValueLen = sizeof(oid6);
    rc                         = funcs->C_CreateObject(sess,
                                                       key_template,
                                                       sizeof(key_template) / sizeof(CK_ATTRIBUTE), &keyh);
    PROCESS_ERROR("C_CreateObject", rc, end);

    rc = funcs->C_EncryptInit(sess, mechanism, keyh);
    PROCESS_ERROR("C_EncryptInit", rc, end);

    len = sizeof(value);
    rc  = funcs->C_Encrypt(sess, data, sizeof(data), value, &len);
    PROCESS_ERROR("C_Encrypt", rc, end);

    CHECK(value, len, et6);

    rc = funcs->C_DestroyObject(sess, keyh);
    PROCESS_ERROR("C_DestroyObject", rc, end);

    printf("Two Sessions With Get/SetOperationState\n");

    printf("Random Data\n");
    rc = funcs->C_GenerateRandom(sess, data2, sizeof(data2));
    PROCESS_ERROR("C_GenerateRandom", rc, end);

    printf("Param Set B\n");
    key_template[2].pValue     = oid3;
    key_template[2].ulValueLen = sizeof(oid3);
    rc                         = funcs->C_CreateObject(sess,
                                                       key_template,
                                                       sizeof(key_template) / sizeof(CK_ATTRIBUTE), &keyh);
    PROCESS_ERROR("C_CreateObject", rc, end);

    rc = funcs->C_EncryptInit(sess, mechanism, keyh);
    PROCESS_ERROR("C_EncryptInit", rc, end);

    len = sizeof(value);
    rc  = funcs->C_EncryptUpdate(sess, data2, 192, value, &len);
    PROCESS_ERROR("C_EncryptUpdate", rc, end);

    // Save session operation state
    state_len = 0;
    rc        = funcs->C_GetOperationState(sess, NULL, &state_len);
    PROCESS_ERROR("C_GetOperationState", rc, end);

    state = malloc(state_len);
    rc    = funcs->C_GetOperationState(sess, state, &state_len);
    PROCESS_ERROR("C_GetOperationState", rc, end);

    // Restore saved operation state to other session
    rc = funcs->C_SetOperationState(sess2, state, state_len,
                                    keyh, CK_INVALID_HANDLE);
    PROCESS_ERROR("C_SetOperationState", rc, end);
    memcpy(value2, value, sizeof(value));

    free(state);
    state_len = 0;
    rc        = funcs->C_GetOperationState(sess2, NULL, &state_len);
    PROCESS_ERROR("C_GetOperationState", rc, end);

    state = malloc(state_len);
    rc    = funcs->C_GetOperationState(sess2, state, &state_len);
    PROCESS_ERROR("C_GetOperationState", rc, end);

    free(state);
    state_len = 0;
    rc        = funcs->C_GetOperationState(sess, NULL, &state_len);
    PROCESS_ERROR("C_GetOperationState", rc, end);

    state = malloc(state_len);
    rc    = funcs->C_GetOperationState(sess, state, &state_len);
    PROCESS_ERROR("C_GetOperationState", rc, end);

    rc = funcs->C_SetOperationState(sess, state, state_len,
                                    keyh, CK_INVALID_HANDLE);
    PROCESS_ERROR("C_SetOperationState", rc, end);

    // Continue operations in first session
    len = sizeof(value) - 192;
    rc  = funcs->C_EncryptUpdate(sess,
                                 data2 + 192, sizeof(data2) - 192, value + 192, &len);
    PROCESS_ERROR("C_EncryptUpdate", rc, end);

    // Continue restored operations in second session
    len = sizeof(value2) - 192;
    rc  = funcs->C_EncryptUpdate(sess2,
                                 data2 + 192, len, value2 + 192, &len);
    PROCESS_ERROR("C_EncryptUpdate", rc, end);

    // Finalize operations in first session
	// Get final size only with NULL buffer
	printf("Call C_EncryptFinal with NULL buffer to get output size\n");
	printf("The encryption context should not be finalized so far\n");
    len = 0;
    rc  = funcs->C_EncryptFinal(sess, NULL, &len);
    PROCESS_ERROR("C_EncryptFinal", rc, end);
	// Несмотря на то, что в нашей реализации длина выходного буфера
	// в данной функции всегда равна нулю (поскольку паддинг не применяется),
	// нельзя задавать здесь NULL для указателя на выходной буфер,
	// потому что по стандарту в этом случае будет считаться, 
	// что вызов производится только для определения длины 
	// выходного буфера, и контекст операции останется активным, 
	// не позволяя выполнить в текущей сессии новую инициализацию.
    rc  = funcs->C_EncryptFinal(sess, value + sizeof(data2), &len);
    PROCESS_ERROR("C_EncryptFinal", rc, end);
    memcpy(value2, value, sizeof(value));

    free(state);
    state_len = 0;
    rc        = funcs->C_GetOperationState(sess2, NULL, &state_len);
    PROCESS_ERROR("C_GetOperationState", rc, end);

    state = malloc(state_len);
    rc    = funcs->C_GetOperationState(sess2, state, &state_len);
    PROCESS_ERROR("C_GetOperationState", rc, end);

    // Finalize operations in second session
    len = sizeof(value2) - sizeof(data2);
    rc  = funcs->C_EncryptFinal(sess2, value2 + sizeof(data2), &len);
    PROCESS_ERROR("C_EncryptFinal", rc, end);
    // len == 0

    rc = funcs->C_DecryptInit(sess2, mechanism, keyh);
    PROCESS_ERROR("C_DecryptInit", rc, end);

    len = sizeof(data2);
    rc  = funcs->C_Decrypt(sess2, value2, len, value, &len);
    PROCESS_ERROR("C_Decrypt", rc, end);

    CHECK(value, len, data2);

    rc = funcs->C_DecryptInit(sess2, mechanism, keyh);
    PROCESS_ERROR("C_DecryptInit", rc, end);

    len = 192;
    rc  = funcs->C_DecryptUpdate(sess2, value2, len, value, &len);
    PROCESS_ERROR("C_DecryptUpdate", rc, end);

    len = sizeof(value2) - 192;
    rc  = funcs->C_DecryptUpdate(sess2,
                                 value2 + 192, len, value + 192, &len);
    PROCESS_ERROR("C_DecryptUpdate", rc, end);

    len = sizeof(value2);
    rc  = funcs->C_DecryptFinal(sess2, value + sizeof(value2), &len);
    PROCESS_ERROR("C_DecryptFinal", rc, end);

    // len == 0
    len = sizeof(data2);
    CHECK(value, len, data2);

    rc = funcs->C_DecryptInit(sess, mechanism, keyh);
    PROCESS_ERROR("C_DecryptInit", rc, end);

    free(state);
    state_len = 0;
    rc        = funcs->C_GetOperationState(sess, NULL, &state_len);
    PROCESS_ERROR("C_GetOperationState", rc, end);

    state = malloc(state_len);
    rc    = funcs->C_GetOperationState(sess, state, &state_len);
    PROCESS_ERROR("C_GetOperationState", rc, end);

	len = 192;
    rc = funcs->C_DecryptUpdate(sess, value2, len, value, &len);
    PROCESS_ERROR("C_DecryptUpdate", rc, end);

    free(state);
    state_len = 0;
    rc        = funcs->C_GetOperationState(sess, NULL, &state_len);
    PROCESS_ERROR("C_GetOperationState", rc, end);

    state = malloc(state_len);
    rc    = funcs->C_GetOperationState(sess, state, &state_len);
    PROCESS_ERROR("C_GetOperationState", rc, end);

	len = sizeof(value2) - 192;
    rc = funcs->C_DecryptUpdate(sess,
                                value2 + 192, len, value + 192, &len);
    PROCESS_ERROR("C_DecryptUpdate", rc, end);

    free(state);
    state_len = 0;
    rc        = funcs->C_GetOperationState(sess, NULL, &state_len);
    PROCESS_ERROR("C_GetOperationState", rc, end);

    state = malloc(state_len);
    rc    = funcs->C_GetOperationState(sess, state, &state_len);
    PROCESS_ERROR("C_GetOperationState", rc, end);

    rc = funcs->C_DecryptFinal(sess, value + sizeof(value2), &len);
    PROCESS_ERROR("C_DecryptFinal", rc, end);
    free(state);

    // len == 0
    len = sizeof(data2);
    CHECK(value, len, data2);

    printf("SUCCESS\n");
    rc = CKR_OK;
 end:
    if (keyh != CK_INVALID_HANDLE)
    {
        funcs->C_DestroyObject(sess, keyh);
    }

    return rc;
}

