#include <inttypes.h>
#include <bliq_storage_ta.h>
#include <tee_internal_api.h>
#include <tee_internal_api_extensions.h>
#include "operations.h"
#include "key_pair.h"
#include <string.h>

static TEE_Result ss_read_object(
	char *obj_id, uint32_t obj_id_sz,
	void *buff, uint32_t *buff_sz
)
{
	TEE_Result res;
	TEE_ObjectHandle object = NULL;
	TEE_ObjectInfo object_info;

	res = TEE_OpenPersistentObject(TEE_STORAGE_PRIVATE,
			obj_id, obj_id_sz,
			TEE_DATA_FLAG_ACCESS_READ |
			TEE_DATA_FLAG_SHARE_READ,
			&object);
	if (res != TEE_SUCCESS) {
		EMSG("[ERR] ss_read_object(): Failed to open persistent object, res=0x%x", res);
		goto exit;
	}

	res = TEE_GetObjectInfo1(
			object,
			&object_info);
	if (res != TEE_SUCCESS) {
		EMSG("[ERR] ss_read_object(): Failed to get persistent object ibfo, res=0x%x", res);
		goto exit;
	}

	if (object_info.dataSize > *buff_sz) {
		/*
		 * Provided buffer is too short.
		 * Return the expected size together with status "short buffer"
		 */
		*buff_sz = object_info.dataSize;
		res = TEE_ERROR_SHORT_BUFFER;
		goto exit;
	}

	res = TEE_ReadObjectData(
			object,
			buff, object_info.dataSize,
			buff_sz);

exit:
	if (object)
		TEE_CloseObject(object);
	return res;
}

static TEE_Result get_key(
	Session *session,
	char *obj_id, uint32_t obj_id_sz,
	TEE_ObjectHandle *kp)
{
	TEE_Result res;

	char    *kp_serialized_buff = NULL;
	uint32_t kp_serialized_buff_sz = 2048;

	char    *kp_decrypted_buff = NULL;
	uint32_t kp_decrypted_buff_sz = 2048;

	kp_serialized_buff = TEE_Malloc(kp_serialized_buff_sz, 0);
	if (!kp_serialized_buff)
		return TEE_ERROR_OUT_OF_MEMORY;

	res = ss_read_object(
			obj_id, obj_id_sz,
			kp_serialized_buff, &kp_serialized_buff_sz);
	if (res != TEE_SUCCESS) {
		EMSG("[ERR]: get_key(%s): Unable to read data from sec store: 0x%x", obj_id, res);
		goto exit;
	}

	kp_decrypted_buff = TEE_Malloc(kp_decrypted_buff_sz, 0);
	if (!kp_decrypted_buff)
		return TEE_ERROR_OUT_OF_MEMORY;

	res = so_decrypt(
			session,
			kp_serialized_buff, kp_serialized_buff_sz,
			kp_decrypted_buff, &kp_decrypted_buff_sz);
	if (res != TEE_SUCCESS) {
		EMSG("[ERR]: get_key(%s): Unable to decrypt RSA key pair: 0x%x", obj_id, res);
		goto exit;
	}

	res = kp_deserialize(
			kp_decrypted_buff, kp_decrypted_buff_sz,
			kp);
	if (res != TEE_SUCCESS) {
		EMSG("[ERR]: get_key(%s): Unable to deserialize RSA key pair: 0x%x", obj_id, res);
		goto exit;
	}

exit:
	TEE_Free(kp_serialized_buff);
	TEE_Free(kp_decrypted_buff);
	return res;
}

/*
 * TA API implementation
 */
static TEE_Result create_key(
	Session *session,
	uint32_t param_types, 
	TEE_Param params[4])
{
	const uint32_t exp_param_types = TEE_PARAM_TYPES(
			TEE_PARAM_TYPE_MEMREF_INPUT,
			TEE_PARAM_TYPE_NONE,
			TEE_PARAM_TYPE_NONE,
			TEE_PARAM_TYPE_NONE);
	
	TEE_ObjectHandle kp = NULL;
	TEE_ObjectHandle object;
	TEE_Result res;
	char *obj_id = NULL;
	size_t obj_id_sz;
	uint32_t obj_data_flag;

	char    *kp_serialized_buff = NULL;
	uint32_t kp_serialized_buff_sz = 2048;

	char    *kp_encrypted_buff = NULL;
	uint32_t kp_encrypted_buff_sz = 2048;

	/*
	 * Safely get the invocation parameters
	 */
	if (param_types != exp_param_types)
		return TEE_ERROR_BAD_PARAMETERS;

	obj_id_sz = params[0].memref.size;
	obj_id = TEE_Malloc(obj_id_sz, 0);
	if (!obj_id)
		return TEE_ERROR_OUT_OF_MEMORY;

	TEE_MemMove(obj_id, params[0].memref.buffer, obj_id_sz);

	/*
	 * Generate RSA KeyPair
	 */ 
	res = kp_create(&kp);
	if (res != TEE_SUCCESS) {
		EMSG("Unable to create RSA key pair: 0x%08x", res);
		goto exit;
	}

	/*
	 * Serialize RSA KeyPair into external buffer
	 */
	kp_serialized_buff = TEE_Malloc(kp_serialized_buff_sz, 0);
	if (!kp_serialized_buff)
		return TEE_ERROR_OUT_OF_MEMORY;

	res = kp_serialize(
			kp,
			kp_serialized_buff, &kp_serialized_buff_sz);
	if (res != TEE_SUCCESS) {
		EMSG("Unable to serialize RSA key pair: 0x%08x", res);
		goto exit;
	}

	/*
	 * Encrypt serrialized RSA keypair buffer
	 */ 
	kp_encrypted_buff = TEE_Malloc(kp_encrypted_buff_sz, 0);
	if (!kp_encrypted_buff)
		return TEE_ERROR_OUT_OF_MEMORY;

	res = so_encrypt(
			session,
			kp_serialized_buff, kp_serialized_buff_sz,
			kp_encrypted_buff, &kp_encrypted_buff_sz);
	if (res != TEE_SUCCESS) {
		EMSG("Unable to encrypt RSA key pair: 0x%08x", res);
		goto exit;
	}

	/*
	 * Create object in secure storage and fill it with encrypted RSA keypair
	 */ 
	obj_data_flag = 
			TEE_DATA_FLAG_ACCESS_READ |		    /* we can later read the oject */
			TEE_DATA_FLAG_ACCESS_WRITE |		/* we can later write into the object */
			TEE_DATA_FLAG_ACCESS_WRITE_META |	/* we can later destroy or rename the object */
			TEE_DATA_FLAG_OVERWRITE;		    /* destroy existing object of same ID */

	res = TEE_CreatePersistentObject(TEE_STORAGE_PRIVATE,
			obj_id, obj_id_sz,
			obj_data_flag,
			TEE_HANDLE_NULL,
			NULL, 0,		/* we may not fill it right now */
			&object);
	if (res != TEE_SUCCESS) {
		EMSG("TEE_CreatePersistentObject failed 0x%08x", res);
		goto exit;
	}

	res = TEE_WriteObjectData(
			object, 
			kp_encrypted_buff, kp_encrypted_buff_sz);
	if (res != TEE_SUCCESS) {
		EMSG("TEE_WriteObjectData failed 0x%08x", res);
		TEE_CloseAndDeletePersistentObject1(object);
	} else {
		TEE_CloseObject(object);
	}

exit:
	TEE_Free(obj_id);
	TEE_Free(kp_serialized_buff);
	TEE_Free(kp_encrypted_buff);

	if (kp)
		TEE_FreeTransientObject(kp);
	return res;
}

static TEE_Result check_key(
	Session *session,
	uint32_t param_types, 
	TEE_Param params[4])
{
	const uint32_t exp_param_types = TEE_PARAM_TYPES(
			TEE_PARAM_TYPE_MEMREF_INPUT,
			TEE_PARAM_TYPE_NONE,
			TEE_PARAM_TYPE_NONE,
			TEE_PARAM_TYPE_NONE);
	
	TEE_Result res;
	TEE_ObjectHandle kp = NULL;

	char *obj_id;
	uint32_t obj_id_sz;

	/*
	 * Safely get the invocation parameters
	 */
	if (param_types != exp_param_types)
		return TEE_ERROR_BAD_PARAMETERS;

	obj_id_sz = params[0].memref.size;
	obj_id = TEE_Malloc(obj_id_sz, 0);
	if (!obj_id)
		return TEE_ERROR_OUT_OF_MEMORY;

	TEE_MemMove(obj_id, params[0].memref.buffer, obj_id_sz);
	EMSG("[DBG] check_key(): obj_id '%s'", obj_id);

	res = get_key(
			session, 
			obj_id, obj_id_sz,
			&kp);
	if (res != TEE_SUCCESS) {
		EMSG("[ERR] check_key(): Unable to obtain RSA keypair: 0x%x", res);
		goto exit;
	}
		
exit:
	TEE_Free(obj_id);
	if (kp)
		TEE_FreeTransientObject(kp);

	return res;
}


static TEE_Result delete_key(
	uint32_t param_types,
	TEE_Param params[4])
{
	const uint32_t exp_param_types = TEE_PARAM_TYPES(
				TEE_PARAM_TYPE_MEMREF_INPUT,
				TEE_PARAM_TYPE_NONE,
				TEE_PARAM_TYPE_NONE,
				TEE_PARAM_TYPE_NONE);
	TEE_ObjectHandle object;
	TEE_Result res;
	char *obj_id;
	size_t obj_id_sz;

	/*
	 * Safely get the invocation parameters
	 */
	if (param_types != exp_param_types)
		return TEE_ERROR_BAD_PARAMETERS;

	obj_id_sz = params[0].memref.size;
	obj_id = TEE_Malloc(obj_id_sz, 0);
	if (!obj_id)
		return TEE_ERROR_OUT_OF_MEMORY;

	TEE_MemMove(obj_id, params[0].memref.buffer, obj_id_sz);

	/*
	 * Check object exists and delete it
	 */
	res = TEE_OpenPersistentObject(TEE_STORAGE_PRIVATE,
				obj_id, obj_id_sz,
				TEE_DATA_FLAG_ACCESS_READ |
				TEE_DATA_FLAG_ACCESS_WRITE_META, /* we must be allowed to delete it */
				&object);
	if (res != TEE_SUCCESS) {
		EMSG("[ERR] Unable to open persistent object, res=0x%x", res);
		TEE_Free(obj_id);
		return res;
	}

	TEE_CloseAndDeletePersistentObject1(object);
	TEE_Free(obj_id);

	return res;
}

static TEE_Result encrypt_data(
	Session *session,
	uint32_t param_types, 
	TEE_Param params[4])
{
	const uint32_t exp_param_types = TEE_PARAM_TYPES(
			TEE_PARAM_TYPE_MEMREF_INPUT,
			TEE_PARAM_TYPE_MEMREF_INPUT,
			TEE_PARAM_TYPE_MEMREF_INOUT,
			TEE_PARAM_TYPE_NONE);
	
	TEE_Result res;
	TEE_ObjectHandle kp = NULL;

	char *obj_id;
	uint32_t obj_id_sz;
	
	char *data;
	uint32_t data_sz;

	char *enc_data;
	uint32_t enc_data_sz;

	/*
	 * Safely get the invocation parameters
	 */
	if (param_types != exp_param_types)
		return TEE_ERROR_BAD_PARAMETERS;

	obj_id_sz = params[0].memref.size;
	obj_id = TEE_Malloc(obj_id_sz, 0);
	if (!obj_id)
		return TEE_ERROR_OUT_OF_MEMORY;

	TEE_MemMove(obj_id, params[0].memref.buffer, obj_id_sz);
	EMSG("[DBG] encrypt_data: obj_id '%s'", obj_id);

	data_sz = params[1].memref.size;
	data = TEE_Malloc(data_sz, 0);
	if (!data)
		return TEE_ERROR_OUT_OF_MEMORY;

	TEE_MemMove(data, params[1].memref.buffer, data_sz);
	EMSG("[DBG] encrypt_data: plain_data '%s'", data);

	res = get_key(
			session, 
			obj_id, obj_id_sz,
			&kp);
	if (res != TEE_SUCCESS) {
		EMSG("[ERR] encrypt_data(): Unable to obtain RSA keypair: 0x%x", res);
		goto exit;
	}

	// const char msg[] = "hello from kp_encrypt";
	// EMSG("[DBG] encrypt_data(): Return hardcoded string");

	// TEE_MemMove(params[2].memref.buffer, msg, sizeof(msg));
	// params[2].memref.size = sizeof(msg);
	// return res;

	enc_data    = params[2].memref.buffer;
	enc_data_sz = params[2].memref.size;

	res = kp_encrypt(
			kp,
			data, data_sz,
			enc_data, &enc_data_sz);
	if (res != TEE_SUCCESS) {
		EMSG("[ERR] encrypt_data(): Unable to encryp data: 0x%x", res);
		goto exit;
	}

	// Encryption [OK]
	// update output buffer size
	params[2].memref.size = enc_data_sz;
		
exit:
	TEE_Free(obj_id);
	TEE_Free(data);
	if (kp)
		TEE_FreeTransientObject(kp);

	return res;
}

static TEE_Result decrypt_data(
	Session *session,
	uint32_t param_types, 
	TEE_Param params[4])
{
	const uint32_t exp_param_types = TEE_PARAM_TYPES(
			TEE_PARAM_TYPE_MEMREF_INPUT,
			TEE_PARAM_TYPE_MEMREF_INPUT,
			TEE_PARAM_TYPE_MEMREF_INOUT,
			TEE_PARAM_TYPE_NONE);
	
	TEE_Result res;
	TEE_ObjectHandle kp = NULL;

	char *obj_id;
	uint32_t obj_id_sz;
	
	char *data;
	uint32_t data_sz;

	char *dec_data;
	uint32_t dec_data_sz;

	/*
	 * Safely get the invocation parameters
	 */
	if (param_types != exp_param_types)
		return TEE_ERROR_BAD_PARAMETERS;

	obj_id_sz = params[0].memref.size;
	obj_id = TEE_Malloc(obj_id_sz, 0);
	if (!obj_id)
		return TEE_ERROR_OUT_OF_MEMORY;

	TEE_MemMove(obj_id, params[0].memref.buffer, obj_id_sz);
	EMSG("[DBG] decrypt_data: obj_id '%s'", obj_id);

	data_sz = params[1].memref.size;
	data = TEE_Malloc(data_sz, 0);
	if (!data)
		return TEE_ERROR_OUT_OF_MEMORY;

	TEE_MemMove(data, params[1].memref.buffer, data_sz);

	res = get_key(
			session, 
			obj_id, obj_id_sz,
			&kp);
	if (res != TEE_SUCCESS) {
		EMSG("[ERR] decrypt_data(): Unable to obtain RSA keypair: 0x%x", res);
		goto exit;
	}

	dec_data    = params[2].memref.buffer;
	dec_data_sz = params[2].memref.size;

	res = kp_decrypt(
			kp,
			data, data_sz,
			dec_data, &dec_data_sz);
	if (res != TEE_SUCCESS) {
		EMSG("[ERR] decrypt_data(): Unable to decryp data: 0x%x", res);
		goto exit;
	}

	// Decryption [OK]
	// update output buffer size
	params[2].memref.size = dec_data_sz;
		
exit:
	TEE_Free(obj_id);
	TEE_Free(data);
	if (kp)
		TEE_FreeTransientObject(kp);

	return res;
}

TEE_Result TA_CreateEntryPoint(void)
{
	/* Nothing to do */
	return TEE_SUCCESS;
}

void TA_DestroyEntryPoint(void)
{
	/* Nothing to do */
}

TEE_Result TA_OpenSessionEntryPoint(
	uint32_t param_types,
	TEE_Param params[4],
	void **ctx)
{
	const uint32_t exp_param_types = TEE_PARAM_TYPES(
				TEE_PARAM_TYPE_MEMREF_INPUT,
				TEE_PARAM_TYPE_NONE,
				TEE_PARAM_TYPE_NONE,
				TEE_PARAM_TYPE_NONE);

	/*
	 * Safely get the invocation parameters
	 */
	if (param_types != exp_param_types)
		return TEE_ERROR_BAD_PARAMETERS;

	Session *session = TEE_Malloc(sizeof (Session), 0);
	if (!session)
		return TEE_ERROR_OUT_OF_MEMORY;

	char  *pin    = params[0].memref.buffer;
	size_t pin_sz = params[0].memref.size;

	*ctx = session;

	return so_init(pin, pin_sz, session);
}

void TA_CloseSessionEntryPoint(void *ctx)
{
	Session *session = ctx;
	if (session == NULL)
		return;

	so_free(session);
	TEE_Free(ctx);
}

TEE_Result TA_InvokeCommandEntryPoint(
	void *session,
	uint32_t command,
	uint32_t param_types,
	TEE_Param params[4])
{
	switch (command) {
	case TA_BLIQ_STORAGE_CMD_CREATE_KEY:
		return create_key(session, param_types, params);
	case TA_BLIQ_STORAGE_CMD_CHECK_KEY:
		return check_key(session, param_types, params);
	case TA_BLIQ_STORAGE_CMD_DELETE_KEY:
		return delete_key(param_types, params);
	case TA_BLIQ_STORAGE_CMD_ENCRYPT:
		return encrypt_data(session, param_types, params);
	case TA_BLIQ_STORAGE_CMD_DECRYPT:
	 	return decrypt_data(session, param_types, params);
	default:
		EMSG("Command ID 0x%x is not supported", command);
		return TEE_ERROR_NOT_SUPPORTED;
	}
}