#include <inttypes.h>
#include "operations.h"
#include "key_pair.h"

static char* storeNext(
	TEE_ObjectHandle kp,
	uint32_t attribute_id,
	char *curr, const char * const end)
{
	TEE_Result res;
	uint32_t attribute_sz;

	if (curr == NULL)
		return NULL;

	TEE_GetObjectBufferAttribute(kp, attribute_id, NULL, &attribute_sz);
	if (curr + sizeof(uint32_t) + attribute_sz > end)
		return NULL;
	
	EMSG("[INFO] -->> attr_size %u\n", attribute_sz);
	
	// store attribute size
	curr[0] = attribute_sz       & 0xFF;
	curr[1] = attribute_sz >>  8 & 0xFF;
	curr[2] = attribute_sz >> 16 & 0xFF;
	curr[3] = attribute_sz >> 24 & 0xFF;

	curr += sizeof(uint32_t);		    // move pointer

	res = TEE_GetObjectBufferAttribute(kp, attribute_id, curr, &attribute_sz);
	if (res != TEE_SUCCESS)
		return NULL;

	return curr + attribute_sz;
}

static char* readNext(
	TEE_Attribute *key_attr,
	uint32_t attribute_id,
	char *curr, const char * const end)
{
	if (curr > end || curr == NULL)
		return NULL;

	uint32_t attribute_sz = curr[3] <<24 | curr[2] <<16 | curr[1] <<8 | curr[0];
	EMSG("[INFO] <<-- attr_size %u\n", attribute_sz);
	curr += sizeof(uint32_t);			// move pointer

	TEE_InitRefAttribute(key_attr, attribute_id, curr, attribute_sz);

	return curr + attribute_sz;
}

TEE_Result kp_create(
	TEE_ObjectHandle *kp) 
{
	const uint32_t key_type = TEE_TYPE_RSA_KEYPAIR;
	const uint32_t key_size = 1024; // fast calculations

	TEE_Result res;
	TEE_ObjectHandle key = NULL;

	res = TEE_AllocateTransientObject(
			key_type, key_size, 
			&key);
	if (res != TEE_SUCCESS) {
		EMSG("[ERR] kp_create: TEE_AllocateTransientObject(%#" PRIx32 ", %" PRId32 "): %#" PRIx32, key_type, key_size, res);
		goto exit;
	}

	res = TEE_GenerateKey(
			key, key_size,
			NULL, 0);
	if (res != TEE_SUCCESS) {
		EMSG("[ERR] kp_create: TEE_GenerateKey(%" PRId32 "): %#" PRIx32, key_size, res);
		goto exit;
	}

	*kp = key;
	key = NULL;

exit:
	if (key)
		TEE_FreeTransientObject(key);

	return res;
}

TEE_Result kp_serialize(
	TEE_ObjectHandle kp,
	char *serilized, uint32_t *serilized_sz)
{
	char *curr = serilized;
	char *end  = serilized + *serilized_sz;
	
	curr = storeNext(kp, TEE_ATTR_RSA_MODULUS         , curr, end);
	curr = storeNext(kp, TEE_ATTR_RSA_PUBLIC_EXPONENT , curr, end);
	curr = storeNext(kp, TEE_ATTR_RSA_PRIVATE_EXPONENT, curr, end);
	curr = storeNext(kp, TEE_ATTR_RSA_PRIME1          , curr, end);
	curr = storeNext(kp, TEE_ATTR_RSA_PRIME2          , curr, end);
	curr = storeNext(kp, TEE_ATTR_RSA_EXPONENT1       , curr, end);
	curr = storeNext(kp, TEE_ATTR_RSA_EXPONENT2       , curr, end);
	curr = storeNext(kp, TEE_ATTR_RSA_COEFFICIENT     , curr, end);

	if (curr == NULL)
		return TEE_ERROR_OVERFLOW;

	*serilized_sz = curr - serilized;

	return TEE_SUCCESS;
}

TEE_Result kp_deserialize(
	char *serilized, uint32_t serilized_sz,
	TEE_ObjectHandle *kp)
{
	TEE_Result res;
	TEE_Attribute key_attr[8];
	uint32_t      key_attr_count = 8;

	TEE_ObjectHandle transient_kp = NULL;
	const uint32_t key_type = TEE_TYPE_RSA_KEYPAIR;
	const uint32_t key_size = 1024; // fast calculations

	char *curr     = serilized;
	char *curr_end = serilized + serilized_sz;

	curr = readNext(&key_attr[0], TEE_ATTR_RSA_MODULUS         , curr, curr_end);
	curr = readNext(&key_attr[1], TEE_ATTR_RSA_PUBLIC_EXPONENT , curr, curr_end);
	curr = readNext(&key_attr[2], TEE_ATTR_RSA_PRIVATE_EXPONENT, curr, curr_end);
	curr = readNext(&key_attr[3], TEE_ATTR_RSA_PRIME1          , curr, curr_end);
	curr = readNext(&key_attr[4], TEE_ATTR_RSA_PRIME2          , curr, curr_end);
	curr = readNext(&key_attr[5], TEE_ATTR_RSA_EXPONENT1       , curr, curr_end);
	curr = readNext(&key_attr[6], TEE_ATTR_RSA_EXPONENT2       , curr, curr_end);
	curr = readNext(&key_attr[7], TEE_ATTR_RSA_COEFFICIENT     , curr, curr_end);

	if (curr == NULL) {
		res = TEE_ERROR_OVERFLOW;
		goto exit;
	}

	res = TEE_AllocateTransientObject(
			key_type, key_size, 
			&transient_kp);
	if (res != TEE_SUCCESS) {
		EMSG("[ERR] kp_deserialize(): TEE_AllocateTransientObject(%#" PRIx32 ", %" PRId32 "): %#" PRIx32, key_type, key_size, res);
		goto exit;
	}
	
	res = TEE_PopulateTransientObject(
			transient_kp,
			key_attr, key_attr_count);
	if (res != TEE_SUCCESS) {
		EMSG("[ERR]: kp_deserialize(): TEE_PopulateTransientObject: 0x%x", res);
		goto exit;
	}

	*kp = transient_kp;
	transient_kp = NULL;

exit:
	if (transient_kp)
		TEE_FreeTransientObject(transient_kp);
	return res;
}

TEE_Result kp_encrypt(
	TEE_ObjectHandle kp,
	void *inbuf,  uint32_t  inbuf_sz,
	void *outbuf, uint32_t *outbuf_sz)
{
	const uint32_t alg = TEE_ALG_RSAES_PKCS1_V1_5;

	TEE_Result res;
	TEE_OperationHandle op = NULL;
	TEE_ObjectInfo key_info;

	res = TEE_GetObjectInfo1(kp, &key_info);
	if (res != TEE_SUCCESS) {
		EMSG("[ERR] kp_encrypt: TEE_GetObjectInfo1: %#" PRIx32, res);
		goto exit;
	}

	res = TEE_AllocateOperation(
			&op, 
			alg, TEE_MODE_ENCRYPT,
			key_info.keySize);
	if (res != TEE_SUCCESS) {
		EMSG("[ERR] kp_encrypt: TEE_AllocateOperation(TEE_MODE_ENCRYPT, %#" PRIx32 ", %" PRId32 "): %#" PRIx32, alg, key_info.keySize, res);
		goto exit;
	}

	res = TEE_SetOperationKey(op, kp);
	if (res != TEE_SUCCESS) {
		EMSG("[ERR] kp_encrypt: TEE_SetOperationKey: %#" PRIx32, res);
		goto exit;
	}

	res = TEE_AsymmetricEncrypt(
			op, 
			NULL, 0, 
			inbuf,  inbuf_sz, 
			outbuf, outbuf_sz);
	if (res != TEE_SUCCESS) {
		EMSG("[ERR] kp_encrypt: TEE_AsymmetricEncrypt(%u, %u): 0x%x", inbuf_sz, *outbuf_sz, res);
		goto exit;
	}

exit:
	if (op)
		TEE_FreeOperation(op);
	
	return res;
}

TEE_Result kp_decrypt(
	TEE_ObjectHandle kp,
	void *inbuf,  uint32_t  inbuf_sz,
	void *outbuf, uint32_t *outbuf_sz)
{
	const uint32_t alg = TEE_ALG_RSAES_PKCS1_V1_5;

	TEE_Result res;
	TEE_OperationHandle op = NULL;
	TEE_ObjectInfo key_info;

	res = TEE_GetObjectInfo1(kp, &key_info);
	if (res != TEE_SUCCESS) {
		EMSG("[ERR] kp_encrypt: TEE_GetObjectInfo1: %#" PRIx32, res);
		goto exit;
	}

	res = TEE_AllocateOperation(
			&op, 
			alg, TEE_MODE_DECRYPT,
			key_info.keySize);
	if (res != TEE_SUCCESS) {
		EMSG("[ERR] kp_encrypt: TEE_AllocateOperation(TEE_MODE_DECRYPT, %#" PRIx32 ", %" PRId32 "): %#" PRIx32, alg, key_info.keySize, res);
		goto exit;
	}

	res = TEE_SetOperationKey(op, kp);
	if (res != TEE_SUCCESS) {
		EMSG("[ERR] kp_encrypt: TEE_SetOperationKey: %#" PRIx32, res);
		goto exit;
	}

	res = TEE_AsymmetricDecrypt(
			op, 
			NULL, 0, 
			inbuf,  inbuf_sz, 
			outbuf, outbuf_sz);
	if (res != TEE_SUCCESS) {
		EMSG("[ERR] kp_encrypt: TEE_AsymmetricEncrypt(%u, %u): 0x%x", inbuf_sz, *outbuf_sz, res);
		goto exit;
	}

exit:
	if (op)
		TEE_FreeOperation(op);

	return res;
}