#include "test_common.h"
#include <time.h>


CK_RV init_all(void);
CK_RV find_obj(CK_SESSION_HANDLE sess, CK_ATTRIBUTE *attr, CK_OBJECT_HANDLE *hObj_res);
CK_RV do_obj_test(CK_BBOOL is_priv, CK_OBJECT_CLASS oclass_data);

void fill_rnd(CK_BYTE *buf, CK_ULONG len){
	srand((unsigned int)time(NULL));
	while(len>0){
		len--;
		buf[len]=rand();
	}
}

//TEST_OBJ_MAX_NEW_SIZE<=TEST_OBJ_SIZE
#define TEST_OBJ_SIZE	256
#define TEST_OBJ_MAX_NEW_SIZE	200

#define TEST_ERROR(msg){ \
							printf("Test error (File:%s, Func:%s, L:%d)\n",__FILE__,  __FUNCTION__, __LINE__); \
							printf msg; \
							goto on_error; \
						} 

int main(int argc, char *argv[]){
	CK_RV rc; // Return Code
	CK_ULONG i;

	printf("create_obj 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 = argv[i];
		}
	}
	/* Load the PKCS11 library */
	
	rc=do_obj_test(CK_FALSE, CKO_DATA);
	if (rc!=CKR_OK) TEST_ERROR((""));

	rc=do_obj_test(CK_TRUE, CKO_DATA);
	if (rc!=CKR_OK) TEST_ERROR((""));

	return CKR_OK;
on_error:
	return rc;
}

CK_RV do_obj_test(CK_BBOOL is_priv, CK_OBJECT_CLASS oclass){
	CK_RV rc; // Return Code
	CK_FLAGS flags = 0; // Bit mask for what options were passed in
	CK_SESSION_HANDLE sess;
	CK_SLOT_INFO sinfo;
	CK_ULONG i;
	CK_ULONG len, new_len;
	CK_BYTE data_value[TEST_OBJ_SIZE], data_tmp[TEST_OBJ_SIZE], data_new[TEST_OBJ_MAX_NEW_SIZE];
	static CK_BBOOL ltrue = CK_TRUE;
	static CK_BBOOL lfalse = CK_FALSE;
	CK_ATTRIBUTE data_find[] = {
		{CKA_CLASS, &oclass, sizeof(oclass)},
	};
	CK_UTF8CHAR label[] = "data0";
	CK_UTF8CHAR label_new[] = "data_new (big label)";
	CK_CHAR app[] = "lsms11";
	CK_ATTRIBUTE data_create[] = {
		{CKA_VALUE, data_value, sizeof(data_value)},
		{CKA_CLASS, &oclass, sizeof(oclass)},
		{CKA_APPLICATION, app, strlen(app)+1},
		{CKA_TOKEN, &ltrue, sizeof(ltrue)},
		{CKA_LABEL, label, strlen(label)+1},
		{CKA_PRIVATE, &is_priv, sizeof(is_priv)},
	};
	CK_OBJECT_HANDLE hObj = CK_INVALID_HANDLE;
	CK_ATTRIBUTE t[]={
		{CKA_CLASS, &oclass, sizeof(oclass)},
		{CKA_LABEL, label, sizeof(label)},
		{CKA_VALUE, data_value, sizeof(data_value)}
	};
	CK_ATTRIBUTE t1={CKA_VALUE, data_tmp, sizeof(data_tmp)};
	CK_ATTRIBUTE attr_update[]={
		{CKA_LABEL, label_new, sizeof(label_new)},
		{CKA_VALUE, data_new, TEST_OBJ_MAX_NEW_SIZE}
	};

	if (TEST_OBJ_MAX_NEW_SIZE>TEST_OBJ_SIZE) TEST_ERROR(("Bad value of TEST_OBJ_MAX_NEW_SIZE and TEST_OBJ_SIZE\n Must be TEST_OBJ_MAX_NEW_SIZE<=TEST_OBJ_SIZE\n"));
	rc = init_all();
	if (rc != CKR_OK) TEST_ERROR(("init failed\n"));

	srand((unsigned int)time(NULL));
	if (oclass==CKO_PUBLIC_KEY) len=64;
	else len=rand()%TEST_OBJ_SIZE;
	fill_rnd(data_value, len);
	t[2].ulValueLen=len;
	for (i = 0; i<SlotCount; i++) {
		SlotId = SlotList[i];
		rc = funcs->C_GetSlotInfo(SlotId, &sinfo);
		if (sinfo.flags & CKF_TOKEN_PRESENT) {
			flags = CKF_SERIAL_SESSION | CKF_RW_SESSION;
			rc = funcs->C_OpenSession( SlotId, flags, NULL, NULL, &sess );
			if (rc != CKR_OK)
				TEST_ERROR(("C_OpenSession failed, rc = 0x%x\n", rc ));
// На аппаратных токенах логин обязателен для изменений на токене
//			if (is_priv){
				rc = funcs->C_Login(sess, CKU_USER, user_pin, strlen(user_pin));
				if (rc != CKR_OK) {
					printf("C_Login failed, rc = 0x%x\n", rc );
					return rc;
				}
				printf("C_Login (CKU_USER) success\n");
//			}

			data_create[0].ulValueLen=len;
			rc = funcs->C_CreateObject(sess, 
				data_create, 
				sizeof(data_create)/sizeof(CK_ATTRIBUTE), 
				&hObj); 
			if (rc != CKR_OK)
				TEST_ERROR(("C_CreateObject failed, rc = 0x%x\n", rc));
			printf("C_CreateObject success\n");

			//проверяем CKA_VALUE сразу после создания
			rc=funcs->C_GetAttributeValue(sess, hObj, &t1, 1);
			if (rc != CKR_OK)
				TEST_ERROR(("C_GetAttributeValue failed, rc = 0x%x\n", rc ));
			if (memcmp(data_value, data_tmp, len))
				TEST_ERROR(("WRONG_CKA_VALUE\n"));

			//при поиске объектов проверится корректность CKA_VALUE
			rc=find_obj(sess, &t[0], &hObj);
			if (rc!=CKR_OK) TEST_ERROR((""));

			rc = funcs->C_CloseSession(sess);
			if (rc != CKR_OK)
				TEST_ERROR(("C_CloseSession failed, rc = 0x%x\n", rc ));
		}
	}

	//проверяем после перезапуска библиотеки
	rc=cleanup(funcs);
	if (rc != CKR_OK)
		TEST_ERROR(("cleanup failed\n"));
	rc = init_all();
	if (rc != CKR_OK)
		TEST_ERROR(("init failed\n"));

	if (oclass==CKO_PUBLIC_KEY) new_len=64;
	new_len=rand()%TEST_OBJ_MAX_NEW_SIZE;
	fill_rnd(data_new, new_len);
	attr_update[1].ulValueLen=new_len;
	for (i = 0; i<SlotCount; i++) {
		SlotId = SlotList[i];
		rc = funcs->C_GetSlotInfo(SlotId, &sinfo);
		if (!(sinfo.flags & CKF_TOKEN_PRESENT)) continue;

		flags = CKF_SERIAL_SESSION | CKF_RW_SESSION;
		rc = funcs->C_OpenSession(SlotId, flags, NULL, NULL, &sess );
		if (rc != CKR_OK)
			TEST_ERROR(("C_OpenSession failed, rc = 0x%x\n", rc ));

// На аппаратных токенах логин обязателен для изменений на токене
//		if (is_priv){
			rc = funcs->C_Login(sess, CKU_USER, user_pin, strlen(user_pin));
			if (rc != CKR_OK) {
				printf("C_Login failed, rc = 0x%x\n", rc );
				return rc;
			}
			printf("C_Login (CKU_USER) success\n");
//		}

		//при поиске объектов проверится корректность CKA_VALUE
		rc=find_obj(sess, &t[0], &hObj);
		if (rc!=CKR_OK) 
			TEST_ERROR((""));

		////проверяем CKA_VALUE для найденного объекта
		t1.ulValueLen=len;
		rc=funcs->C_GetAttributeValue(sess, hObj, &t1, 1);
		if (rc != CKR_OK) 
			TEST_ERROR(("C_GetAttributeValue failed, rc = 0x%x\n", rc ));
 		if (memcmp(data_value, data_tmp, len)) 
			TEST_ERROR(("WRONG_CKA_VALUE\n"));

		//Изменяем атрибуты объекта (label и value), в т.ч. - длину
		rc=funcs->C_SetAttributeValue(sess, hObj, attr_update, 2);
		if (rc!=CKR_OK) 
			TEST_ERROR(("C_SetAttributeValue failed, rc = 0x%x\n", rc ));
		//проверяем CKA_VALUE сразу после создания
		t1.ulValueLen=new_len;
		rc=funcs->C_GetAttributeValue(sess, hObj, &t1, 1);
		if (rc != CKR_OK) 
			TEST_ERROR(("C_GetAttributeValue failed, rc = 0x%x\n", rc ));
 		if (memcmp(data_new, data_tmp, new_len)) 
			TEST_ERROR(("WRONG_CKA_VALUE\n"));

		rc = funcs->C_CloseSession(sess);
		if (rc != CKR_OK)
			TEST_ERROR(("C_CloseSession failed, rc = 0x%x\n", rc ));
	}

//перезапускаем библиотеку и проверяем новые значения лейбла и value
	rc=cleanup(funcs);
	if (rc != CKR_OK)
		TEST_ERROR(("cleanup failed\n"));
	rc = init_all();
	if (rc != CKR_OK)
		TEST_ERROR(("init failed\n"));

	t[1].pValue=label_new;
	t[1].ulValueLen=sizeof(label_new);
	t[2].pValue=data_new;
	t[2].ulValueLen=new_len;
	for (i = 0; i<SlotCount; i++) {
		SlotId = SlotList[i];
		rc = funcs->C_GetSlotInfo(SlotId, &sinfo);
		if (!(sinfo.flags & CKF_TOKEN_PRESENT)) continue;

		flags = CKF_SERIAL_SESSION | CKF_RW_SESSION;
		rc = funcs->C_OpenSession(SlotId, flags, NULL, NULL, &sess );
		if (rc != CKR_OK)
			TEST_ERROR(("C_OpenSession failed, rc = 0x%x\n", rc ));

// На аппаратных токенах логин обязателен для изменений на токене
//		if (is_priv){
			rc = funcs->C_Login(sess, CKU_USER, user_pin, strlen(user_pin));
			if (rc != CKR_OK) {
				printf("C_Login failed, rc = 0x%x\n", rc );
				return rc;
			}
			printf("C_Login (CKU_USER) success\n");
//		}

		//при поиске объектов проверится корректность CKA_VALUE
		rc=find_obj(sess, &t[0], &hObj);
		if (rc!=CKR_OK) TEST_ERROR((""));

//TODO!!!!
		////проверяем CKA_VALUE для найденного объекта
		rc=funcs->C_GetAttributeValue(sess, hObj, &t1, 1);
		if (rc != CKR_OK) 
			TEST_ERROR(("C_GetAttributeValue failed, rc = 0x%x\n", rc ));
 		if (memcmp(data_new, data_tmp, new_len)) 
			TEST_ERROR(("WRONG_CKA_VALUE\n"));

		//Удаляем объект за ненадобностью
		rc=funcs->C_DestroyObject(sess, hObj);
		if (rc!=CKR_OK) TEST_ERROR(("C_DestroyObject failed, rc = 0x%x\n", rc));

		rc = funcs->C_CloseSession(sess);
		if (rc != CKR_OK) TEST_ERROR(("C_CloseSession failed, rc = 0x%x\n", rc ));
	}
	rc=cleanup(funcs);
	if (rc != CKR_OK)
		TEST_ERROR(("cleanup failed\n"));

	printf("CREATE_OBJ test SUCCESS\n");
	return CKR_OK;

on_error:
	printf("TEST FAILED\n");
	return -1;
}

CK_RV find_obj(CK_SESSION_HANDLE sess, CK_ATTRIBUTE *attr, CK_OBJECT_HANDLE *hObj_res){
	CK_RV rc; // Return Code
	CK_OBJECT_HANDLE hObj = CK_INVALID_HANDLE;
	CK_ULONG number;

	// Ищем объекты по лейблу, классу и value
	rc = funcs->C_FindObjectsInit(sess, attr, 3); 
	if (rc != CKR_OK) {
		printf("C_FindObjectsFinal failed, rc = 0x%x\n", rc );
		return rc;
	}

	number = 0;
	hObj = CK_INVALID_HANDLE;
	rc = funcs->C_FindObjects(sess, &hObj, 1, &number);
	if (rc != CKR_OK) {
		printf("C_FindObjects failed, rc = 0x%x\n", rc );
		return rc;
	}
	if (number == 0 || hObj == CK_INVALID_HANDLE){
		printf("Object isnt found\n");
		return -1;
	}
	*hObj_res=hObj;
	printf("Object is found\n");
	rc = funcs->C_FindObjectsFinal(sess);
	if (rc != CKR_OK) {
		printf("C_FindObjectsFinal failed, rc = 0x%x\n", rc );
		return rc;
	}
	return CKR_OK;
}

CK_RV init_all(void){
	CK_RV rc; // Return Code
	rc = init(api_path, &funcs, NULL);
	if (rc != CKR_OK) return rc;
	rc = get_slot_list(funcs, &SlotList,  &SlotCount);
	if (rc != CKR_OK) {
		printf("get_slot_list failed, rc = 0x%x\n", rc );
		return rc;
	}
	rc = find_slot_with_token(funcs, SlotList, SlotCount, &SlotId);
	if (rc != CKR_OK) {
		printf("find_slot_with_token failed, rc = 0x%x\n", rc );
		return rc;
	}
	return rc;
}