#include "test_common.h"

CK_RV test_crypto();
CK_RV test_gost28147_cnt(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, ret;
	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%x\n", rc);
		goto out;
	}
	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%x\n", rc);
		goto out_close;
	}
	fprintf(stderr, "C_Login success\n");
*/
	rc = funcs->C_GetMechanismInfo(SlotId, CKM_GOST28147, &minfo);
	if (rc != CKR_OK) {
		fprintf(stderr, 
"\n======== Mechanism CKM_GOST28147 not supported ===========\n");
	} else {
		fprintf(stderr, 
"\n==================== CKM_GOST28147 CNT test ==================\n");
		rc = test_gost28147_cnt(hSession);
		if (rc != CKR_OK) {
			fprintf(stderr, 
				"ERROR CKM_GOST28147 CNT failed, rc = 0x%x\n", rc);
		} else {
			fprintf(stderr, "CKM_GOST28147 CNT test passed.\n");
		}
	}

	if( (ret = funcs->C_CloseSession(hSession)) != CKR_OK ) {
		fprintf(stderr, 
			"Error: C_CloseSession failed with 0x%x\n", ret);
		rc = ret;
	}
	else {
		fprintf(stderr, "C_CloseSession success\n");
	}

out:
	return rc;
}

// Шифрование в режиме CNT может быть реализовано двумя способами:
// - с помощью нестандартного механизма CKM_GOST28147_CNT;
// - с использованием нестандартных параметров домена (как в тесте CBC).
CK_RV test_gost28147_cnt(CK_SESSION_HANDLE sess)
{
    int rc = 0;
	CK_ULONG len, state_len, dummy_len;
	CK_OBJECT_HANDLE keyh = CK_INVALID_HANDLE;
	CK_BYTE *state;
	CK_BYTE iv[] = {
		0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
	};
   	CK_MECHANISM mechanism_desc = {CKM_GOST28147, iv, sizeof(iv)};
   	CK_MECHANISM_PTR mechanism = &mechanism_desc;
   	CK_MECHANISM mechanism_desc_cnt = {CKM_GOST28147_CNT, iv, sizeof(iv)};
   	CK_MECHANISM_PTR mechanism_cnt = &mechanism_desc_cnt;

	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_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,
    };
	CK_BYTE value[32];
	CK_BYTE buf[32];

	static CK_BYTE oid1[] = { 
		0x06, 0x07, 0x2a, 0x85, 0x03, 0x02, 0x02, 0x1f, 0x00 };
	static CK_BYTE oid_cnt1[] = { 
		0x06, 0x07, 0x2a, 0x85, 0x03, 0x02, 0x02, 0x1f, 0x40 };
	static CK_BYTE et1[] = { 
		0x75, 0xea, 0xe3, 0x45, 0xb5, 0x12, 0xac, 0x47, 
		0xe5, 0xe6, 0x33, 0x04, 0xe0, 0xb9, 0x99, 0xa6 };

    static CK_BYTE oid2[] = { 
		0x06, 0x07, 0x2a, 0x85, 0x03, 0x02, 0x02, 0x1f, 0x01 };
    static CK_BYTE oid_cnt2[] = { 
		0x06, 0x07, 0x2a, 0x85, 0x03, 0x02, 0x02, 0x1f, 0x41 };
	static CK_BYTE et2[] = { 
		0xd4, 0x7d, 0xab, 0xaa, 0x4b, 0x49, 0x3a, 0x8d, 
		0xd6, 0xeb, 0xb6, 0x17, 0xd6, 0xa4, 0xfd, 0x31 };
 
	static CK_BYTE oid3[] = { 
		0x06, 0x07, 0x2a, 0x85, 0x03, 0x02, 0x02, 0x1f, 0x02 };
	static CK_BYTE oid_cnt3[] = { 
		0x06, 0x07, 0x2a, 0x85, 0x03, 0x02, 0x02, 0x1f, 0x42 };
	static CK_BYTE et3[] = { 
		0x22, 0xb3, 0x6e, 0x02, 0x7b, 0x03, 0xa5, 0x6a, 
		0x5f, 0x23, 0xf4, 0xbd, 0x63, 0x6e, 0x03, 0x1f };
 
	static CK_BYTE oid4[] = { 
		0x06, 0x07, 0x2a, 0x85, 0x03, 0x02, 0x02, 0x1f, 0x03 };
	static CK_BYTE oid_cnt4[] = { 
		0x06, 0x07, 0x2a, 0x85, 0x03, 0x02, 0x02, 0x1f, 0x43 };
	static CK_BYTE et4[] = { 
		0xf1, 0x26, 0xfc, 0x4d, 0x20, 0xd5, 0xe2, 0xcd, 
		0xe8, 0x1c, 0x7e, 0x10, 0xed, 0x27, 0x07, 0xe0 };
 
	static CK_BYTE oid5[] = { 
		0x06, 0x07, 0x2a, 0x85, 0x03, 0x02, 0x02, 0x1f, 0x04 };
	static CK_BYTE oid_cnt5[] = { 
		0x06, 0x07, 0x2a, 0x85, 0x03, 0x02, 0x02, 0x1f, 0x44 };
	static CK_BYTE et5[] = { 
		0xfd, 0x65, 0x30, 0xa2, 0xe2, 0x93, 0xfe, 0xc1, 
		0xfe, 0x71, 0xc5, 0x79, 0x3d, 0x20, 0x66, 0x73 };
 
	static CK_ATTRIBUTE key_template[] = {
		{ CKA_CLASS, &oclass, sizeof(oclass) },
		{ CKA_KEY_TYPE, &key_type, sizeof(key_type) },
		{ CKA_GOST28147PARAMS, NULL, 0 },
		{ CKA_VALUE, key, sizeof(key) },
		{ CKA_ENCRYPT, &ltrue, sizeof(CK_BBOOL) },
		{ CKA_DECRYPT, &ltrue, sizeof(CK_BBOOL) },
	};
	static CK_OBJECT_CLASS dp_class = CKO_DOMAIN_PARAMETERS;
	const CK_OBJECT_CLASS oc_dp = CKO_DOMAIN_PARAMETERS;

	CK_ATTRIBUTE t_find_dp[] = {
		{CKA_CLASS,    (CK_VOID_PTR)&oc_dp, sizeof(oc_dp)},
		{CKA_KEY_TYPE, &key_type, sizeof(key_type)},
		{CKA_OBJECT_ID, NULL, 0},
	};	
	
	// Фактически это набор параметров A из RFC 4357,
	// в котором режим шифрования изменен на CNT.
	CK_BYTE dp_cnt[] =
	{
		0x30, 0x53,
			0x04, 0x40,
				// +4 (таблица замен)
				0x93, 0xee, 0xb3, 0x1b, 0x67, 0x47, 0x5a, 0xda,
				0x3e, 0x6a, 0x1d, 0x2f, 0x29, 0x2c, 0x9c, 0x95,
				0x88, 0xbd, 0x81, 0x70, 0xba, 0x31, 0xd2, 0xac,
				0x1f, 0xd3, 0xf0, 0x6e, 0x70, 0x89, 0x0b, 0x08,
				0xa5, 0xc0, 0xe7, 0x86, 0x42, 0xf2, 0x45, 0xc2,
				0xe6, 0x5b, 0x29, 0x43, 0xfc, 0xa4, 0x34, 0x59,
				0xcb, 0x0f, 0xc8, 0xf1, 0x04, 0x78, 0x7f, 0x37,
				0xdd, 0x15, 0xae, 0xbd, 0x51, 0x96, 0x66, 0xe4,	
			0x02, 0x01,
				// +70 (mode: 0-CNT, 1-CFB, 2-CBC)
				0x00,
			0x02, 0x01, 0x40,
			0x30, 0x09,
				0x06, 0x07, 0x2a, 0x85, 0x03, 0x02, 0x02, 0x0e,
					// +84 (key meshing: 0-Null, 1-CryptoPro)
					0x01,
	};
	CK_ULONG uldp_cnt = sizeof(dp_cnt);
	// Этот OID и название не являются предопределенными в LCC,
	// а используются для определения дополнительных параметров домена. 
	CK_BYTE oid_cnt[] = {
		0x06, 0x07, 0x2a, 0x85, 0x03, 0x02, 0x02, 0x1f, 0x42};
	CK_CHAR *dp_label = "id-Gost28147-89-CryptoPro-A-ParamSet-CNT";

	CK_ATTRIBUTE t_dp_cnt[] = {
		{CKA_CLASS, &dp_class, sizeof(dp_class)},
		{CKA_TOKEN, &lfalse, sizeof(lfalse)},
		{CKA_KEY_TYPE, &key_type, sizeof(key_type)},
		{CKA_OBJECT_ID, oid_cnt, sizeof(oid_cnt)},
		{CKA_VALUE, dp_cnt, uldp_cnt},
		{CKA_LABEL, dp_label, strlen(dp_label)+1},
		{CKA_COPYABLE, &ltrue, sizeof(ltrue)}
	};

	CK_BYTE dp_value[4096];
	CK_ULONG uldp_value = sizeof(dp_value);

	CK_ATTRIBUTE t_dp_value = {
		CKA_VALUE, dp_value, sizeof(dp_value)
	};

	CK_ATTRIBUTE t_copy_dp_cnt[] = {
		{CKA_TOKEN, &lfalse, sizeof(lfalse)},
		{CKA_OBJECT_ID, NULL, 0},
		{CKA_VALUE, dp_value, 0},
	};
	CK_OBJECT_HANDLE hDP_CNT;
	CK_ULONG ulCB = 0;

	key_template[2].pValue = oid1;
	key_template[2].ulValueLen = sizeof(oid1);
	rc = funcs->C_CreateObject(sess, 
		key_template, sizeof(key_template)/sizeof(CK_ATTRIBUTE), &keyh);
	if (rc != CKR_OK) {
		fprintf(stderr, "C_CreateObject failed: 0x%x\n", rc);
		return rc;
	}

	rc = funcs->C_EncryptInit(sess, mechanism_cnt, keyh);
	if (rc != CKR_OK) {
		fprintf(stderr, "C_EncryptInit failed: 0x%x\n", rc);
		return rc;
	}

	len = sizeof(value);
	rc = funcs->C_EncryptUpdate(sess, data, 4, value, &len);
	if (rc != CKR_OK) {
		fprintf(stderr, "C_EncryptUpdate failed: 0x%x\n", rc);
		return rc;
	}
	// Save partial encryption result
	memcpy(buf, value, 4);
	state_len = 0;
	rc = funcs->C_GetOperationState(sess, NULL, &state_len);
	if (rc != CKR_OK) {
		fprintf(stderr, "C_GetOperationState failed: 0x%x\n", rc);
		goto end;
	}
	state = malloc(state_len);
	rc = funcs->C_GetOperationState(sess, state, &state_len);
	if (rc != CKR_OK) {
		fprintf(stderr, "C_GetOperationState failed: 0x%x\n", rc);
		goto end;
	}

	rc = funcs->C_SetOperationState(sess, state, state_len, 
		keyh, CK_INVALID_HANDLE);
	if (rc != CKR_OK) {
		fprintf(stderr, "C_SetOperationState failed: 0x%x\n", rc);
		goto end;
	}
	rc = funcs->C_EncryptUpdate(sess, 
		data+4, sizeof(data)-4, value+4, &len);
	if (rc != CKR_OK) {
		fprintf(stderr, "C_EncryptUpdate failed: 0x%x\n", rc);
		return rc;
	}
	rc = funcs->C_EncryptFinal(sess, value+sizeof(data), &dummy_len);
	if (rc != CKR_OK) {
		fprintf(stderr, "C_EncryptFinal failed: 0x%x\n", rc);
		return rc;
	}

	len = sizeof(data);
	CHECK(value, len, et1);

	rc = funcs->C_DecryptInit(sess, mechanism_cnt, keyh);
	if (rc != CKR_OK) {
		fprintf(stderr, "C_DecryptInit failed: 0x%x\n", rc);
		return rc;
	}
	rc = funcs->C_DecryptUpdate(sess, value, 7, value, &len);
	if (rc != CKR_OK) {
		fprintf(stderr, "C_DecryptUpdate failed: 0x%x\n", rc);
		return rc;
	}
	printf("len = %d, Decryption buffer:\n", len);
	print_hex(value, sizeof(value));
	rc = funcs->C_DecryptUpdate(sess, 
		value+7, sizeof(data)-7, value+7, &len);
	if (rc != CKR_OK) {
		fprintf(stderr, "C_DecryptUpdate failed: 0x%x\n", rc);
		return rc;
	}
	printf("len = %d, Decryption buffer:\n", len);
	print_hex(value, sizeof(value));
	rc = funcs->C_DecryptFinal(sess, 
		value+sizeof(data), &dummy_len);
	if (rc != CKR_OK) {
		fprintf(stderr, "C_DecryptFinal failed: 0x%x\n", rc);
		return rc;
	}
	len = sizeof(data);
	printf("len = %d, Decrypted data:\n", len);
	print_hex(value, len);
	printf("Decryption buffer:\n");
	print_hex(value, sizeof(value));

	CHECK(value, len, data);

	// Restore partial encryption result
	memcpy(value, buf, 4);
	rc = funcs->C_SetOperationState(sess, state, state_len, 
		keyh, CK_INVALID_HANDLE);
	if (rc != CKR_OK) {
		fprintf(stderr, "C_SetOperationState failed: 0x%x\n", rc);
		goto end;
	}
	rc = funcs->C_EncryptUpdate(sess, 
		data+4, sizeof(data)-4, value+4, &len);
	if (rc != CKR_OK) {
		fprintf(stderr, "C_EncryptUpdate failed: 0x%x\n", rc);
		return rc;
	}
	rc = funcs->C_EncryptFinal(sess, 
		value+sizeof(data), &dummy_len);
	if (rc != CKR_OK) {
		fprintf(stderr, "C_EncryptFinal failed: 0x%x\n", rc);
		return rc;
	}
	len = sizeof(data);
	CHECK(value, len, et1);

	rc = funcs->C_DestroyObject(sess, keyh);
	if (rc != CKR_OK) {
		fprintf(stderr, "C_DestroyObject error, code 0x%X\n", rc);
		return rc;
	}
	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);
	if (rc != CKR_OK) {
		fprintf(stderr, "C_CreateObject failed: 0x%x\n", rc);
		return rc;
	}

	rc = funcs->C_EncryptInit(sess, mechanism_cnt, keyh);
	if (rc == CKR_FUNCTION_NOT_SUPPORTED) {
		fprintf(stderr, 
"Non-standard parameter set OID not supported for CKM_GOST28147\n");
	} else { 
		len = sizeof(value);
		rc = funcs->C_Encrypt(sess, data, sizeof(data), value, &len);
		if (rc != CKR_OK) {
			fprintf(stderr, "C_Encrypt failed: 0x%x\n", rc);
			return rc;
		}

		CHECK(value, len, et2);

		rc = funcs->C_DecryptInit(sess, mechanism_cnt, keyh);
		if (rc != CKR_OK) {
			fprintf(stderr, "C_DecryptInit failed: 0x%x\n", rc);
			return rc;
		}

		rc = funcs->C_Decrypt(sess, value, len, value, &len);
		if (rc != CKR_OK) {
			fprintf(stderr, "C_Decrypt failed: 0x%x\n", rc);
			return rc;
		}

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

	// Теперь то же самое, но с использованием сессионного объекта 
	// параметров домена с нестандартным OID.
	dp_cnt[70] = 0x00;  // CNT mode 

	rc = funcs->C_CreateObject(sess, 
		t_dp_cnt, sizeof(t_dp_cnt)/sizeof(CK_ATTRIBUTE),
		&hDP_CNT);
	if (rc != CKR_OK) {
		fprintf(stderr, "C_CreateObject failed: 0x%x\n", rc);
		return rc;
	}

	key_template[2].pValue = oid_cnt;
	key_template[2].ulValueLen = sizeof(oid_cnt);
	rc = funcs->C_CreateObject(sess, 
		key_template, sizeof(key_template)/sizeof(CK_ATTRIBUTE), &keyh);
	if (rc != CKR_OK) {
		fprintf(stderr, "C_CreateObject failed: 0x%x\n", rc);
		return rc;
	}

	rc = funcs->C_EncryptInit(sess, mechanism, keyh);
	if (rc == CKR_FUNCTION_NOT_SUPPORTED) {
		fprintf(stderr, 
"Non-standard parameter set OID not supported for CKM_GOST28147\n");
	} else { 
		len = sizeof(value);
		rc = funcs->C_Encrypt(sess, data, sizeof(data), value, &len);
		if (rc != CKR_OK) {
			fprintf(stderr, "C_Encrypt failed: 0x%x\n", rc);
			return rc;
		}

		CHECK(value, len, et2);

		rc = funcs->C_DecryptInit(sess, mechanism, keyh);
		if (rc != CKR_OK) {
			fprintf(stderr, "C_DecryptInit failed: 0x%x\n", rc);
			return rc;
		}

		rc = funcs->C_Decrypt(sess, value, len, value, &len);
		if (rc != CKR_OK) {
			fprintf(stderr, "C_Decrypt failed: 0x%x\n", rc);
			return rc;
		}

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


	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);
	if (rc != CKR_OK) {
		fprintf(stderr, "C_CreateObject failed: 0x%x\n", rc);
		return rc;
	}

	rc = funcs->C_EncryptInit(sess, mechanism_cnt, keyh);
	if (rc != CKR_OK) {
		fprintf(stderr, "C_EncryptInit failed: 0x%x\n", rc);
		return rc;
	}

	len = sizeof(value);
	rc = funcs->C_Encrypt(sess, data, sizeof(data), value, &len);
	if (rc != CKR_OK) {
		fprintf(stderr, "C_Encrypt failed: 0x%x\n", rc);
		return rc;
	}

	CHECK(value, len, et3);

	rc = funcs->C_DecryptInit(sess, mechanism_cnt, keyh);
	if (rc != CKR_OK) {
		fprintf(stderr, "C_DecryptInit failed: 0x%x\n", rc);
		return rc;
	}

	rc = funcs->C_Decrypt(sess, value, len, value, &len);
	if (rc != CKR_OK) {
		fprintf(stderr, "C_Decrypt failed: 0x%x\n", rc);
		return rc;
	}

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


	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);
	if (rc != CKR_OK) {
		fprintf(stderr, "C_CreateObject failed: 0x%x\n", rc);
		return rc;
	}

	rc = funcs->C_EncryptInit(sess, mechanism_cnt, keyh);
	if (rc != CKR_OK) {
		fprintf(stderr, "C_EncryptInit failed: 0x%x\n", rc);
		return rc;
	}

	len = sizeof(value);
	rc = funcs->C_Encrypt(sess, data, sizeof(data), value, &len);
	if (rc != CKR_OK) {
		fprintf(stderr, "C_Encrypt failed: 0x%x\n", rc);
		return rc;
	}

	CHECK(value, len, et4);

	rc = funcs->C_DecryptInit(sess, mechanism_cnt, keyh);
	if (rc != CKR_OK) {
		fprintf(stderr, "C_DecryptInit failed: 0x%x\n", rc);
		return rc;
	}

	rc = funcs->C_Decrypt(sess, value, len, value, &len);
	if (rc != CKR_OK) {
		fprintf(stderr, "C_Decrypt failed: 0x%x\n", rc);
		return rc;
	}

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


	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);
	if (rc != CKR_OK) {
		fprintf(stderr, "C_CreateObject failed: 0x%x\n", rc);
		return rc;
	}

	rc = funcs->C_EncryptInit(sess, mechanism_cnt, keyh);
	if (rc != CKR_OK) {
		fprintf(stderr, "C_EncryptInit failed: 0x%x\n", rc);
		return rc;
	}

	len = sizeof(value);
	rc = funcs->C_Encrypt(sess, data, sizeof(data), value, &len);
	if (rc != CKR_OK) {
		fprintf(stderr, "C_Encrypt failed: 0x%x\n", rc);
		return rc;
	}

	CHECK(value, len, et5);

	rc = funcs->C_DecryptInit(sess, mechanism_cnt, keyh);
	if (rc != CKR_OK) {
		fprintf(stderr, "C_DecryptInit failed: 0x%x\n", rc);
		return rc;
	}

	rc = funcs->C_Decrypt(sess, value, len, value, &len);
	if (rc != CKR_OK) {
		fprintf(stderr, "C_Decrypt failed: 0x%x\n", rc);
		return rc;
	}

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

end:
	return rc;
}
