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

char     iv[] = {
	0x7D, 0x3F, 0x80, 0x6C, 0x35, 0x7E, 0x04, 0xC2,
	0x92, 0x4C, 0x67, 0x30, 0xBA, 0xFA, 0xE0, 0xF7};
uint32_t iv_sz = sizeof(iv);


static TEE_Result sha256(
	char *in,  size_t in_sz,
	char *out, uint32_t *out_sz)
{
	TEE_Result res;
	TEE_OperationHandle hDigest = NULL;

	res = TEE_AllocateOperation( 
			&hDigest,
			TEE_ALG_SHA256, TEE_MODE_DIGEST, 
			0);
	if (res != TEE_SUCCESS) {
		EMSG("[ERR]: TEE_AllocateOperation: 0x%x", res);
		goto exit;
	}

	res = TEE_DigestDoFinal(
			hDigest,
			in, in_sz,
			out, out_sz);
	if (res != TEE_SUCCESS) {
		EMSG("[ERR]: TEE_DigestDoFinal: 0x%x", res);
		goto exit;
	}

exit:
	if (hDigest)
		TEE_FreeOperation(hDigest);

	return res;
}

/*
 * Session Operations
 */
TEE_Result so_init(
	char *pin, size_t pin_sz,
	Session *session)
{
	TEE_Result res;
	TEE_ObjectHandle hSessionKey = NULL;
	TEE_Attribute key_attr = { 0 };
	uint32_t      key_attr_count = 1;

	char     key[32]; // 256bit
	uint32_t key_sz = sizeof(key);

	res = sha256(
			pin, pin_sz, 
			key, &key_sz);
	if (res != TEE_SUCCESS) {
		EMSG("[ERR]: sha256: %d", res);
		goto exit;
	}

	res = TEE_AllocateTransientObject( 
			TEE_TYPE_AES, 256, 
			&hSessionKey);
	if (res != TEE_SUCCESS) {
		EMSG("[ERR]: SessionKey TEE_AllocateTransientObject: 0x%x", res);
		goto exit;
	}

	TEE_InitRefAttribute(&key_attr, TEE_ATTR_SECRET_VALUE, key, key_sz);
	res = TEE_PopulateTransientObject(
			hSessionKey,
			&key_attr, key_attr_count);
	if (res != TEE_SUCCESS) {
		EMSG("[ERR]: SessionKey TEE_PopulateTransientObject: 0x%x", res);
		goto exit;
	}

	res = TEE_AllocateOperation( 
			&session->hSK_encrypt,
			TEE_ALG_AES_CTR, TEE_MODE_ENCRYPT, 256);
	if (res != TEE_SUCCESS) {
		EMSG("[ERR]: SessionKey hSK_encrypt: TEE_AllocateOperation: 0x%x", res);
		goto exit;
	}

	res = TEE_SetOperationKey( 
			session->hSK_encrypt,
			hSessionKey);
	if (res != TEE_SUCCESS) {
		EMSG("[ERR]: SessionKey TEE_SetOperationKey(sk_encrypt): 0x%x", res);
		goto exit;
	}

	res = TEE_AllocateOperation( 
		&session->hSK_decrypt,
		TEE_ALG_AES_CTR, TEE_MODE_DECRYPT, 256);
	if (res != TEE_SUCCESS) {
		EMSG("[ERR]: SessionKey hSK_decrypt: TEE_AllocateOperation: 0x%x", res);
		goto exit;
	}

	res = TEE_SetOperationKey( 
			session->hSK_decrypt,
			hSessionKey);
	if (res != TEE_SUCCESS) {
		EMSG("[ERR]: SessionKey TEE_SetOperationKey(sk_decrypt): 0x%x", res);
		goto exit;
	}

exit:
	if (hSessionKey)
		TEE_FreeTransientObject(hSessionKey);

	return res;
}

void so_free(
	Session *session)
{
	TEE_FreeOperation(session->hSK_encrypt);
	TEE_FreeOperation(session->hSK_decrypt);
}


TEE_Result so_encrypt(
	Session *session,
	void *in,  uint32_t  in_sz,
	void *out, uint32_t *out_sz)
{
	TEE_CipherInit(session->hSK_encrypt, iv, iv_sz);

	return TEE_CipherDoFinal(session->hSK_encrypt,
		in,  in_sz,		// src
		out, out_sz); 	// dest
}

TEE_Result so_decrypt(
	Session *session,
	void *in,  uint32_t  in_sz,
	void *out, uint32_t *out_sz)
{
	TEE_CipherInit(session->hSK_decrypt, iv, iv_sz);

	return TEE_CipherDoFinal(session->hSK_decrypt,
		in,  in_sz,		// src
		out, out_sz); 	// dest
}