#include "database.h"

#if defined(CONFIG_SUPPORT_EP2)

#include "ep2_public_key.h"

#include "symmetric.h"

#include "utils.h"

#endif



typedef struct {

	const char* object_id;

	KeyDBId dbid;

	uint32_t type;

	persistance_type persistant;

} DBField;





#undef D

#define D(id, type, persistance) {#id, id, type, persistance},

const DBField DB_FIELDS[] = {

#include "database_ids.h"

};





// WATCH OUT! this one can return without cleanup, so use wisely!

#define DATABASE_GET_DBFIELD(_id, _dbf) do {\

	if((_id) >= NUMBER_OF_DBID || (_id) <= UNKNOWN_DBID) { \

		return TEE_ERROR_DATABASE_ID_UNKNOWN; \

	} \

	(_dbf) = DB_FIELDS[id]; \

} while(0);





// WATCH OUT! if calls the upper one

#define DATABASE_GET_DBFIELD_OF_TYPE(_id, _dbf, _key_type) do {\

	DATABASE_GET_DBFIELD(_id, _dbf); \

	if((_dbf).type != (_key_type)) { \

		return TEE_ERROR_DATABASE_ID_MISMATCH; \

	} \

} while(0);



#define DATABASE_DELETE(object_id) \

	do { \

		TEE_ObjectHandle _delete_oh_ = TEE_HANDLE_NULL; \

		status = database_get_direct((object_id), &_delete_oh_); \

		if (TEE_SUCCESS == status) { \

			TEE_CloseAndDeletePersistentObject(_delete_oh_); \

		} \

		else if (TEE_ERROR_ITEM_NOT_FOUND == status) { status = TEE_SUCCESS; } \

	} while(0);



#define PERMISSIONS (TEE_DATA_FLAG_ACCESS_READ | \

		     TEE_DATA_FLAG_ACCESS_WRITE | \

		     TEE_DATA_FLAG_ACCESS_WRITE_META | \

		     TEE_DATA_FLAG_SHARE_READ)



TEE_Result 

_allocate_and_set_secure_object_id(const char* object_id, uint8_t** secure_object_id, size_t* object_id_len)

{

	TEE_Result status = TEE_SUCCESS;

	*object_id_len = strlen(object_id);



	CHECK_MALLOC(*secure_object_id, *object_id_len + 1, exit);

	memmove(*secure_object_id, object_id, *object_id_len + 1); 



exit:

	return status;

}



TEE_Result

_database_set_worker(const char* object_id, TEE_ObjectHandle oh, uint8_t* data, size_t data_size)

{

	TEE_Result status = TEE_SUCCESS;



	DATABASE_DELETE(object_id);



	uint8_t* object_id_secure = NULL;

	size_t object_id_len = 0;

	CHECK(_allocate_and_set_secure_object_id(object_id, &object_id_secure, &object_id_len), exit2);



	TEE_ObjectHandle poh = TEE_HANDLE_NULL;

	CHECK(TEE_CreatePersistentObject(

					 TEE_STORAGE_PRIVATE, 

					 object_id_secure, object_id_len,

					 TEE_DATA_FLAG_ACCESS_WRITE_META,

					 oh,

					 data, data_size,

					 &poh), exit);



	TEE_CloseObject(poh);

exit:

	free(object_id_secure);

exit2:

	return status;

}





TEE_Result

_database_set_handle(const char* object_id, TEE_ObjectHandle oh)

{

	return _database_set_worker(object_id, oh, NULL, 0);

}





TEE_Result

database_set(KeyDBId id, TEE_ObjectHandle oh)

{

	TEE_Result status = TEE_SUCCESS;



	DBField dbf;

	DATABASE_GET_DBFIELD(id, dbf);

	CHECK(_database_set_handle(dbf.object_id, oh), exit);



exit:

	return status;

}



int 

_database_ids_consistency_check()

{

	KeyDBId id = UNKNOWN_DBID;

	while(id != NUMBER_OF_DBID) {

		if(DB_FIELDS[id].dbid != id) {

			return id;

		}

		id += 1;

	}

	return 0;

}



#if defined(CONFIG_SUPPORT_EP2)

TEE_Result

_database_set_des3_direct(const char* object_id, const uint8_t* data)

{

	MPOS_AUTH_LOG("_database_set_des3_direct+++");

	TEE_Result status = TEE_SUCCESS;



	TEE_ObjectHandle toh = TEE_HANDLE_NULL;

	CHECK(TEE_AllocateTransientObject(TEE_TYPE_DES3, DES3_KEY_SIZE_BITS, &toh), exit);



	CHECK(populate_des_key( data, toh), exit1);



	CHECK(_database_set_handle(object_id, toh), exit1);



exit1:

	TEE_FreeTransientObject(toh);

exit:

	MPOS_AUTH_LOG("_database_set_des3_direct---");

	MPOS_AUTH_LOG("%s: ret=%d", __func__, status);

	return status;

}



// public from here on



TEE_Result

database_set_des3(KeyDBId id, const uint8_t* data)

{

	MPOS_AUTH_LOG("database_set_des3+++");

	DBField dbf;

	MPOS_AUTH_LOG("dbf.dbid = %d", dbf.dbid);

	MPOS_AUTH_LOG("dbf.type = %d", dbf.type);

	MPOS_AUTH_LOG("dbf.persistant = %d", dbf.persistant);

	DATABASE_GET_DBFIELD_OF_TYPE(id, dbf, TEE_TYPE_DES3);

	MPOS_AUTH_LOG("[after] dbf.dbid = %d", dbf.dbid);

	MPOS_AUTH_LOG("[after] dbf.type = %d", dbf.type);

	MPOS_AUTH_LOG("[after] dbf.persistant = %d", dbf.persistant);



	return _database_set_des3_direct(dbf.object_id, data);

}



const char TRM_KEY_OBJECT_ID_HEADER[] = TRM_KEY_DB_PREFIX;

#define TRM_KEY_OBJECT_ID_LEN (sizeof(TRM_KEY_OBJECT_ID_HEADER) + 1 + ACQUIRER_NAME_LENGTH)



TEE_Result 

database_set_trm_key(const char* acq_name, const uint8_t* data)

{

	const char object_id[TRM_KEY_OBJECT_ID_LEN];

	sprintf(((char*)object_id), "%s-%s", TRM_KEY_OBJECT_ID_HEADER, acq_name);

	MPOS_AUTH_LOG("%s, acq_name : %s", __func__, acq_name);

	MPOS_AUTH_LOG("%s, object_id : %s", __func__, object_id);



	return _database_set_des3_direct(object_id, data);

}



TEE_Result 

database_get_trm_key(const char* acq_name, TEE_ObjectHandle* oh)

{

	const char object_id[TRM_KEY_OBJECT_ID_LEN];

	sprintf(((char*)object_id), "%s-%s", TRM_KEY_OBJECT_ID_HEADER, acq_name); // first integer is not needed

	MPOS_AUTH_LOG("%s, object_id : %s", __func__, object_id);



	return database_get_direct(object_id, oh);

}

#endif





TEE_Result 

database_get_direct(const char *object_id, TEE_ObjectHandle *poh) 

{

	MPOS_AUTH_LOG_DEBUG("%s+++, id=%s", __func__, object_id);

	TEE_Result status = TEE_SUCCESS;



	uint8_t *object_id_secure = NULL;

	size_t object_id_len = 0;

	CHECK(_allocate_and_set_secure_object_id(object_id, &object_id_secure, &object_id_len), exit);



	status = TEE_OpenPersistentObject( TEE_STORAGE_PRIVATE,

					   (char *)object_id_secure, object_id_len,

					   PERMISSIONS,

					   poh );



	free(object_id_secure);

exit:

	MPOS_AUTH_LOG_DEBUG("%s---", __func__);

	return status;

}



TEE_Result

database_get(KeyDBId id, TEE_ObjectHandle* persistent_oh)

{

	//STATUS_OK();



	DBField dbf;

	DATABASE_GET_DBFIELD(id, dbf);



	return database_get_direct((char*)dbf.object_id, persistent_oh);

}



bool 

database_exists_direct(const char* object_id)

{

	TEE_ObjectHandle toh = TEE_HANDLE_NULL;



	bool result = TEE_SUCCESS == database_get_direct(object_id, &toh);



	if(result) {

		TEE_CloseObject(toh);

	}

	return result;

}



bool 

database_exists(KeyDBId id)

{

	DBField dbf;

	DATABASE_GET_DBFIELD(id, dbf);

	return database_exists_direct(dbf.object_id);

}



TEE_Result

database_remove(KeyDBId id)

{

	TEE_Result status = TEE_SUCCESS;



	DBField dbf;

	DATABASE_GET_DBFIELD(id, dbf);

	DATABASE_DELETE(dbf.object_id);



	return status;

}



TEE_Result

database_remove_direct(const char* object_id)

{

	TEE_Result status = TEE_SUCCESS;

	DATABASE_DELETE(object_id);

	return status;

}



#if defined(CONFIG_SUPPORT_EP2)

TEE_Result

database_store_pk_direct(uint8_t* exponent, size_t exponent_size, uint8_t* modulus, size_t modulus_size, const char* name)

{

	TEE_Result status = TEE_SUCCESS;



	TEE_Attribute attr[2];

	TEE_InitRefAttribute(&attr[0], TEE_ATTR_RSA_MODULUS, modulus, modulus_size);

	TEE_InitRefAttribute(&attr[1], TEE_ATTR_RSA_PUBLIC_EXPONENT, exponent, exponent_size);



	TEE_ObjectHandle oh = TEE_HANDLE_NULL;

	size_t key_size = ceil_power_of_2(modulus_size);



	CHECK(TEE_AllocateTransientObject(TEE_TYPE_RSA_PUBLIC_KEY, key_size * 8, &oh), exit1);

	CHECK(TEE_PopulateTransientObject(oh, attr, 2), exit);

	CHECK(_database_set_handle(name, oh), exit);



exit:

	TEE_FreeTransientObject(oh);

exit1:

	return status;

}





TEE_Result

database_store_pk(const public_key* pk)

{

	TEE_Result status = TEE_SUCCESS;



	// get id for database

	char key_id[PKDBID_SIZE + 1];

	pk_get_database_id_from_pk(pk, key_id);



	// modulus as an array

	uint8_t exponent[4] = {

		(pk->exponent >> 24) & 0xFF, //NOLINT

		(pk->exponent >> 16) & 0xFF, //NOLINT

		(pk->exponent >> 8) & 0xFF,  //NOLINT

		(pk->exponent) & 0xFF        //NOLINT

	};



	MPOS_AUTH_LOG("%s, key id for stored pk : %s", __func__, key_id);

	CHECK(database_store_pk_direct(exponent, sizeof(exponent), (uint8_t*)pk->modulus, pk->modulus_length, key_id), exit);



exit:

	return status;

}





TEE_Result

database_store_key_pair(KeyDBId id, 

			uint8_t* modulus, 

			size_t modulus_length, 

			uint8_t* public_exponent, 

			size_t pubexp_length, 

			uint8_t* private_exponent, 

			size_t privexp_length)

{

	TEE_Result status = TEE_SUCCESS;



	DBField dbf;

	DATABASE_GET_DBFIELD_OF_TYPE(id, dbf, TEE_TYPE_RSA_KEYPAIR);



	TEE_Attribute attr[3];

	TEE_InitRefAttribute(&attr[0], TEE_ATTR_RSA_MODULUS, modulus, modulus_length);

	TEE_InitRefAttribute(&attr[1], TEE_ATTR_RSA_PUBLIC_EXPONENT, public_exponent, pubexp_length);

	TEE_InitRefAttribute(&attr[2], TEE_ATTR_RSA_PRIVATE_EXPONENT, private_exponent, privexp_length);



	TEE_ObjectHandle oh = TEE_HANDLE_NULL;

	size_t key_size = ceil_power_of_2(modulus_length);



	CHECK(TEE_AllocateTransientObject(TEE_TYPE_RSA_KEYPAIR, key_size * 8, &oh), exit1);

	CHECK(TEE_PopulateTransientObject(oh, attr, 3), exit);

	CHECK(_database_set_handle(dbf.object_id, oh), exit);



exit:

	TEE_FreeTransientObject(oh);

exit1:

	return status;

}





TEE_Result

database_get_pk(char* id, TEE_ObjectHandle* persistent_oh)

{

	return database_get_direct(id, persistent_oh);

}

#endif





TEE_Result

database_set_data(KeyDBId id, uint8_t* data, size_t data_size)

{

	DBField dbf;

	DATABASE_GET_DBFIELD_OF_TYPE(id, dbf, TEE_TYPE_DATA);

	return _database_set_worker(dbf.object_id, TEE_HANDLE_NULL, data, data_size);

}



TEE_Result database_set_data_direct(const char* object_id, uint8_t* data, size_t data_size)

{

	return _database_set_worker(object_id, TEE_HANDLE_NULL, data, data_size);

}





TEE_Result

database_append_data(KeyDBId id, uint8_t* data, size_t data_size)

{

	DBField dbf;

	DATABASE_GET_DBFIELD_OF_TYPE(id, dbf, TEE_TYPE_DATA);

	return database_append_data_direct(dbf.object_id, data, data_size);

}



TEE_Result

database_append_data_direct(const char* object_id, uint8_t* data, size_t data_size)

{

	TEE_Result status = TEE_SUCCESS;

	TEE_ObjectHandle oh = TEE_HANDLE_NULL;



	if(database_exists_direct(object_id)) {

		CHECK(database_get_direct(object_id, &oh), exit);



		CHECK(TEE_SeekObjectData(oh, 0, TEE_DATA_SEEK_END), exit1);

		CHECK(TEE_WriteObjectData(oh, data, data_size), exit1);

	}

	else {

		CHECK(database_set_data_direct(object_id, data, data_size), exit);

	}



exit1:

	TEE_CloseObject(oh);

exit:

	return status;

}



TEE_Result

database_get_data_direct(const char* name, uint8_t* data, size_t* output_size, bool check_output_size, size_t offset)

{

	MPOS_AUTH_LOG_DEBUG("%s+++", __func__);

	TEE_Result status = TEE_SUCCESS;



	TEE_ObjectHandle oh = TEE_HANDLE_NULL;

	CHECK(database_get_direct(name, &oh), exit);



	if(offset != 0) {

		CHECK(TEE_SeekObjectData(oh,  offset, TEE_DATA_SEEK_CUR), exit1);

	}



	uint32_t count = 0;

	CHECK(TEE_ReadObjectData(oh, data, *output_size, &count), exit1);

	MPOS_AUTH_LOG("%s, TEE_ReadObjectData, count=%d, output_size=%d", __func__, count, *output_size);

	MPOS_DBG_DUMP("DATA", data, *output_size);



	if(check_output_size && (count != *output_size)) {

		// TODO(WMEO) better error!!

		status = TEE_ERROR_DATABASE_OUTPUT_SIZE_MISMATCH;

		MPOS_AUTH_LOG("%s, size mismatch. count = %d, output_size = %d", __func__, count, *output_size);

	}



	*output_size = count;

exit1:

	TEE_CloseObject(oh);

exit:

	MPOS_AUTH_LOG_DEBUG("%s---", __func__);

	return status;

}





TEE_Result

database_get_data(KeyDBId id, uint8_t* data, size_t* output_size, bool check_output_size, size_t offset)

{

	DBField dbf;

	DATABASE_GET_DBFIELD(id, dbf);

	return database_get_data_direct(dbf.object_id, data, output_size, check_output_size, offset);

}





TEE_Result

database_get_data_size_direct(const char* object_name, size_t* data_size)

{

	TEE_Result status = TEE_SUCCESS;



	TEE_ObjectHandle oh = TEE_HANDLE_NULL;

	status = database_get_direct(object_name, &oh);



	if(TEE_ERROR_ITEM_NOT_FOUND == status) {

		*data_size = 0;

		status = TEE_SUCCESS;

		goto exit;

	}

	else {

		CHECK(status, exit);

	}



	TEE_ObjectInfo db_info;

	CHECK(TEE_GetObjectInfo1(oh, &db_info), exit1);



	*data_size = db_info.dataSize;



exit1:

	TEE_CloseObject(oh);

exit:

	return status;

}



TEE_Result

database_get_data_size(KeyDBId id, size_t* data_size)

{

	DBField dbf;

	DATABASE_GET_DBFIELD(id, dbf);

	return database_get_data_size_direct(dbf.object_id, data_size);

}



TEE_Result database_dbid_to_strid(KeyDBId id, const char** keyid, unsigned handle_type)

{

	DBField dbf;

	DATABASE_GET_DBFIELD_OF_TYPE(id, dbf, handle_type);

	*keyid = dbf.object_id;

	return TEE_SUCCESS;

}





TEE_Result database_clean(persistance_type clean_type)

{

	TEE_Result status = TEE_SUCCESS;

	TEE_Result result_status = TEE_SUCCESS;



	ASSERT_TRUE(clean_type != HARDCODED, TEE_ERROR_BAD_PARAMETERS, bad_exit);



	for(KeyDBId dbid = UNKNOWN_DBID + 1; dbid < NUMBER_OF_DBID; dbid++) {

		if(DB_FIELDS[dbid].persistant == clean_type) {

			CHECK(database_remove(dbid), bad_remove); // doing CHECK because it gives you nice traces

bad_remove:

			result_status = result_status == TEE_SUCCESS ? status : result_status; // just return first error if there is any error

		}

	}



	return result_status;

bad_exit:

	return status;

}



TEE_Result database_clean_objects_with_prefix(const char* prefix) 

{

	TEE_Result status = TEE_SUCCESS;



	TEE_ObjectEnumHandle eoh = NULL;

	CHECK(TEE_AllocatePersistentObjectEnumerator(&eoh), exit);

	CHECK(TEE_StartPersistentObjectEnumerator(eoh, TEE_STORAGE_PRIVATE), exit);



	TEE_ObjectInfo oi = { 0 };

	char object_id[TEE_OBJECT_ID_MAX_LEN] = { 0 };

	size_t object_id_size = 0;

	size_t prefix_len = strlen(prefix);



	while(true) {

		status = TEE_GetNextPersistentObject(eoh, &oi, object_id, (uint32_t *)&object_id_size);



		if(TEE_ERROR_ITEM_NOT_FOUND == status) {

			status = TEE_SUCCESS;

			break;

		}

		CHECK(status, exit);



		object_id[object_id_size] = '\0';

		if(object_id_size > prefix_len && 0 == memcmp(object_id, prefix, prefix_len)) {

			DATABASE_DELETE(object_id);

		}

	}



exit:

	TEE_FreePersistentObjectEnumerator(eoh);

	return status;

}





TEE_Result database_clean_full(bool clean_persistant)

{

	TEE_Result status = TEE_SUCCESS;



	status |= database_clean(EP2_TEMPORARY);

	status |= database_clean(GKLP_TEMPORARY);

	status |= database_clean(TEMPORARY);



	if(clean_persistant) {

		status |= database_clean(PERSISTANT);

		status |= database_clean_objects_with_prefix(PUBLIC_KEY_DB_PREFIX);

		status |= database_clean_objects_with_prefix(TRM_KEY_DB_PREFIX);

		status |= database_clean_objects_with_prefix(FILE_STORAGE_DB_PREFIX);

	}



#ifdef CLEAR_ATTESTATION_AT_BOOT

	bool clear_attestation_at_boot = true;

#else

	bool clear_attestation_at_boot = false;

#endif

	if(clear_attestation_at_boot || clean_persistant) {

		status |= database_clean(ATTESTATION);

	}



	return status;

}