#include "test_common.h"

CK_RV test_crypto();
CK_RV test_kuznyechik_cbc(CK_SESSION_HANDLE sess);
CK_RV test_magma_cbc(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%lx\n", rc);
        return rc;
    }
    rc = test_crypto();
    if (rc != CKR_OK)
    {
        fprintf(stderr, "test_crypto failed: 0x%lx\n", rc);
        return rc;
    }
    return CKR_OK;
}

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

    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%lx\n", rc);
        return rc;
    }
    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 = 0x%lx\n", rc);
                goto out_close;
        }
        fprintf(stderr, "C_Login success\n");
 */
    rc = funcs->C_GetMechanismInfo(SlotId, CKM_GOST_CIPHER_CBC, &minfo);
    if (rc != CKR_OK)
    {
        fprintf(stderr,
                "\n======== Mechanism CKM_GOST_CIPHER_CBC not supported ===========\n");
		goto done;
    }
    else
    {
        fprintf(stderr,
                "\n==================== CKM_GOST_CIPHER_CBC test ==================\n");
        rc = test_kuznyechik_cbc(hSession);
        if (rc != CKR_OK)
        {
            fprintf(stderr,
                    "ERROR test_kuznyechik_cbc failed, rc = 0x%lx\n", rc);
			goto done;
        }
        else
        {
            fprintf(stderr, "test kuznyechik_cbc passed.\n");
        }
        rc = test_magma_cbc(hSession);
        if (rc != CKR_OK)
        {
            fprintf(stderr,
                    "ERROR test_magma_cbc failed, rc = 0x%lx\n", rc);
			goto done;
        }
        else
        {
            fprintf(stderr, "test_magma_cbc passed.\n");
        }
    }

 done:
	funcs->C_CloseSession(hSession);
    return rc;
}

CK_RV test_kuznyechik_cbc(CK_SESSION_HANDLE sess)
{
    CK_RV rc = CKR_OK;
    CK_ULONG               len, dummy_len = 0;
    CK_OBJECT_HANDLE       keyh = CK_INVALID_HANDLE;
    CK_MECHANISM           mechanism_desc = { CKM_GOST_CIPHER_CBC, 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_KEY_TYPE     key_type = CKK_KUZNYECHIK;

    /* тестовые значения мастер-ключа */
    static unsigned char kMasterKeyData[] =
    {
        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
    };
    static unsigned char kGost14CbcSV[] =
    {
        0x12, 0x34, 0x56, 0x78, 0x90, 0xab, 0xce, 0xf0,
        0xa1, 0xb2, 0xc3, 0xd4, 0xe5, 0xf0, 0x01, 0x12,
        0x23, 0x34, 0x45, 0x56, 0x67, 0x78, 0x89, 0x90,
        0x12, 0x13, 0x14, 0x15, 0x16, 0x17, 0x18, 0x19
    };

    /* значение открытого текста */
    static unsigned char kGost14PlainText[] =
    {
        0x11, 0x22, 0x33, 0x44, 0x55, 0x66, 0x77, 0x00,
        0xff, 0xee, 0xdd, 0xcc, 0xbb, 0xaa, 0x99, 0x88,
        0x00, 0x11, 0x22, 0x33, 0x44, 0x55, 0x66, 0x77,
        0x88, 0x99, 0xaa, 0xbb, 0xcc, 0xee, 0xff, 0x0a,
        0x11, 0x22, 0x33, 0x44, 0x55, 0x66, 0x77, 0x88,
        0x99, 0xaa, 0xbb, 0xcc, 0xee, 0xff, 0x0a, 0x00,
        0x22, 0x33, 0x44, 0x55, 0x66, 0x77, 0x88, 0x99,
        0xaa, 0xbb, 0xcc, 0xee, 0xff, 0x0a, 0x00, 0x11
    };
    /* значение шифртекста */
    static unsigned char kGost14CbcC[] =
    {
        0x68, 0x99, 0x72, 0xd4, 0xa0, 0x85, 0xfa, 0x4d,
        0x90, 0xe5, 0x2e, 0x3d, 0x6d, 0x7d, 0xcc, 0x27,
        0x28, 0x26, 0xe6, 0x61, 0xb4, 0x78, 0xec, 0xa6,
        0xaf, 0x1e, 0x8e, 0x44, 0x8d, 0x5e, 0xa5, 0xac,
        0xfe, 0x7b, 0xab, 0xf1, 0xe9, 0x19, 0x99, 0xe8,
        0x56, 0x40, 0xe8, 0xb0, 0xf4, 0x9d, 0x90, 0xd0,
        0x16, 0x76, 0x88, 0x06, 0x5a, 0x89, 0x5c, 0x63,
        0x1a, 0x2d, 0x9a, 0x15, 0x60, 0xb6, 0x39, 0x70
    };

    CK_BYTE             value[256];
    CK_BYTE             buf[256];
    static CK_ATTRIBUTE key_template[] = {
        { CKA_CLASS,    &oclass,        sizeof(oclass)         },
        { CKA_KEY_TYPE, &key_type,      sizeof(key_type)       },
        { CKA_VALUE,    kMasterKeyData, sizeof(kMasterKeyData) },
        { CKA_ENCRYPT,  &ltrue,         sizeof(CK_BBOOL)       },
        { CKA_DECRYPT,  &ltrue,         sizeof(CK_BBOOL)       },
    };


    rc = funcs->C_CreateObject(sess,
                               key_template, sizeof(key_template) / sizeof(CK_ATTRIBUTE), &keyh);
    if (rc != CKR_OK)
    {
        fprintf(stderr, "C_CreateObject failed: 0x%lx\n", rc);
        return rc;
    }
    printf("Encrypt\n");
    mechanism->pParameter     = kGost14CbcSV;
    mechanism->ulParameterLen = sizeof(kGost14CbcSV);
    rc                        = funcs->C_EncryptInit(sess, mechanism, keyh);
    if (rc != CKR_OK)
    {
        fprintf(stderr, "C_EncryptInit failed: 0x%lx\n", rc);
        return rc;
    }
    // FIXME: Попытка задать данные, длина которых не кратна размеру блока,
    // приводит к ошибке в последнем неполном блоке.
    len = sizeof(kGost14PlainText);
    rc  = funcs->C_EncryptUpdate(sess, kGost14PlainText, sizeof(kGost14PlainText),
                                 value, &len);
    if (rc != CKR_OK)
    {
        fprintf(stderr, "C_EncryptUpdate failed: 0x%lx\n", rc);
        return rc;
    }
    rc = funcs->C_EncryptFinal(sess, value + sizeof(kGost14PlainText), &dummy_len);
    if (rc != CKR_OK)
    {
        fprintf(stderr, "C_EncryptFinal failed: 0x%lx\n", rc);
        return rc;
    }

    len = sizeof(kGost14PlainText);
    CHECK(value, len, kGost14CbcC);

    printf("Decrypt\n");
    mechanism->pParameter     = kGost14CbcSV;
    mechanism->ulParameterLen = sizeof(kGost14CbcSV);
    rc                        = funcs->C_DecryptInit(sess, mechanism, keyh);
    if (rc != CKR_OK)
    {
        fprintf(stderr, "C_DecryptInit failed: 0x%lx\n", rc);
        return rc;
    }
    // Нельзя указывать для входных и выходных данных
    // один и тот же буфер.
    rc = funcs->C_DecryptUpdate(sess, value, len, buf, &len);
    if (rc != CKR_OK)
    {
        fprintf(stderr, "C_DecryptUpdate failed: 0x%lx\n", rc);
        return rc;
    }
    rc = funcs->C_DecryptFinal(sess,
                               buf + sizeof(kGost14PlainText), &dummy_len);
    if (rc != CKR_OK)
    {
        fprintf(stderr, "C_DecryptFinal failed: 0x%lx\n", rc);
        return rc;
    }
    len = sizeof(kGost14PlainText);

    CHECK(buf, len, kGost14PlainText);

    rc = funcs->C_DestroyObject(sess, keyh);
    if (rc != CKR_OK)
    {
        fprintf(stderr, "C_DestroyObject error, code 0x%lx\n", rc);
        return rc;
    }

    return rc;
}

CK_RV test_magma_cbc(CK_SESSION_HANDLE sess)
{
    CK_RV rc = CKR_OK;
    CK_BYTE                value[256];
    CK_BYTE                value2[256];
    CK_ULONG               len;
    CK_OBJECT_HANDLE       keyh           = CK_INVALID_HANDLE;
    CK_MECHANISM           mechanism_desc = { CKM_GOST_CIPHER_CBC, 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_KEY_TYPE     key_type = CKK_GOST28147;	// Тоже работает
    static CK_KEY_TYPE     key_type = CKK_MAGMA;
    // 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
//    };
    static unsigned char   plaintext_cbc_Z[] = {
        0x61, 0x62, 0x63, 0x64, 0x62, 0x63, 0x64, 0x65,
        0x63, 0x64, 0x65, 0x66, 0x64, 0x65, 0x66, 0x67,
        0x65, 0x66, 0x67, 0x68, 0x66, 0x67, 0x68, 0x69,
        0x67, 0x68, 0x69, 0x6a, 0x68, 0x69, 0x6a, 0x6b,
        0x69, 0x6a, 0x6b, 0x6c, 0x6a, 0x6b, 0x6c, 0x6d,
        0x6b, 0x6c, 0x6d, 0x6e, 0x6c, 0x6d, 0x6e, 0x6f,
        0x6d, 0x6e, 0x6f, 0x70, 0x6e, 0x6f, 0x70, 0x71,
        0x0a, 0x07, 0x07, 0x07, 0x07, 0x07, 0x07, 0x07,
    };
    static unsigned char   iv_cbc_Z[] = {
        0x81, 0x16, 0x29, 0xec, 0x62, 0xb4, 0xd0, 0xb4,
    };
    static unsigned char   key_cbc_Z[] = {
        0x0d, 0xe3, 0xf0, 0x8b, 0x21, 0x18, 0x3c, 0x28,
        0x0c, 0x5c, 0xa8, 0x7f, 0x7b, 0xf2, 0x7a, 0x89,
        0x14, 0x3a, 0xa4, 0x47, 0x68, 0xd5, 0x51, 0x05,
        0x23, 0x0a, 0xbe, 0xdc, 0x86, 0x68, 0x4b, 0x7a,
    };
    static unsigned char   encr_cbc_tc26_Z[] = {
		0xba, 0x1a, 0x50, 0x73, 0xfe, 0xe0, 0x5c, 0x5f,
		0xef, 0x52, 0x44, 0xa4, 0x77, 0xa4, 0xcd, 0x87,
		0x83, 0x11, 0x0f, 0x64, 0x1a, 0x19, 0xa5, 0xab,
		0x4d, 0x5c, 0xdf, 0xcb, 0xcd, 0x5e, 0xda, 0xb9,
		0xc0, 0xc6, 0x4f, 0x70, 0xfe, 0xc7, 0xd6, 0x51,
		0x82, 0x93, 0xb1, 0xeb, 0x55, 0x24, 0xa6, 0x21,
		0x84, 0x04, 0xaa, 0x67, 0x3c, 0xc5, 0x69, 0x53,
		0xa6, 0x83, 0x16, 0xe6, 0x1a, 0x4a, 0x8b, 0x23,
    };

    static CK_ATTRIBUTE    key_template[] = {
        { CKA_CLASS,           &oclass,   sizeof(oclass)    },
        { CKA_KEY_TYPE,        &key_type, sizeof(key_type)  },
//        { CKA_GOST28147PARAMS, oid_Z,      sizeof(oid_Z)      },
        { CKA_VALUE,           key_cbc_Z, sizeof(key_cbc_Z) },
        { CKA_ENCRYPT,         &ltrue,    sizeof(CK_BBOOL)  },
        { CKA_DECRYPT,         &ltrue,    sizeof(CK_BBOOL)  },
    };

    rc = funcs->C_CreateObject(sess,
                               key_template, sizeof(key_template) / sizeof(CK_ATTRIBUTE),
                               &keyh);
    if (rc != CKR_OK)
    {
        fprintf(stderr, "C_CreateObject failed: 0x%lx\n", rc);
        goto end;
    }

    mechanism->pParameter     = iv_cbc_Z;
    mechanism->ulParameterLen = sizeof(iv_cbc_Z);
    rc                        = funcs->C_EncryptInit(sess, mechanism, keyh);
    if (rc != CKR_OK)
    {
        fprintf(stderr, "C_EncryptInit failed: 0x%lx\n", rc);
        goto end;
    }

    fprintf(stderr, "Plain text to encrypt:\n");
    print_bytes(plaintext_cbc_Z, sizeof(plaintext_cbc_Z));
    // Данные для шифрования в блочных режимах должны
    // поступать на вход порциями, длина которых
    // кратна размеру блока, т.е. 8-ми.
    len = sizeof(plaintext_cbc_Z);
    rc  = funcs->C_Encrypt(sess, plaintext_cbc_Z, len, value, &len);
    if (rc != CKR_OK)
    {
        fprintf(stderr, "C_Encrypt failed: 0x%lx\n", rc);
        goto end;
    }
    if (len != sizeof(encr_cbc_tc26_Z))
    {
        fprintf(stderr, "Invalid ciphertext length: %ld\n", len);
        rc = CKR_FUNCTION_FAILED;
        goto end;
    }
    if (memcmp(value, encr_cbc_tc26_Z, len) != 0)
    {
        fprintf(stderr, "Invalid ciphertext\n");
        rc = CKR_FUNCTION_FAILED;
        goto end;
    }
    fprintf(stderr, "CBC encryption result OK\n");
    print_bytes(value, len);
    rc = funcs->C_DecryptInit(sess, mechanism, keyh);
    if (rc != CKR_OK)
    {
        fprintf(stderr, "C_DecryptInit failed: 0x%lx\n", rc);
        goto end;
    }

    rc = funcs->C_Decrypt(sess, value, len, value2, &len);
    if (rc != CKR_OK)
    {
        fprintf(stderr, "C_Decrypt failed: 0x%lx\n", rc);
        goto end;
    }
    fprintf(stderr, "Decrypted text:\n");
    print_bytes(value2, len);

    if (len != sizeof(plaintext_cbc_Z))
    {
        fprintf(stderr, "Invalid decrypted text length: %ld\n", len);
        rc = CKR_FUNCTION_FAILED;
        goto end;
    }
    if (memcmp(value2, plaintext_cbc_Z, len) != 0)
    {
        fprintf(stderr, "Invalid decrypted text\n");
        rc = CKR_FUNCTION_FAILED;
        goto end;
    }
    fprintf(stderr, "CBC decryption result OK\n");

 end:
    if (keyh != CK_INVALID_HANDLE)
    {
        funcs->C_DestroyObject(sess, keyh);
    }

    return rc;
}
