/*
 * Copyright (c) 2017, Linaro Limited
 * All rights reserved.
 *
 * Redistribution and use in source and binary forms, with or without
 * modification, are permitted provided that the following conditions are met:
 *
 * 1. Redistributions of source code must retain the above copyright notice,
 * this list of conditions and the following disclaimer.
 *
 * 2. Redistributions in binary form must reproduce the above copyright notice,
 * this list of conditions and the following disclaimer in the documentation
 * and/or other materials provided with the distribution.
 *
 * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
 * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
 * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
 * ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE
 * LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
 * CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
 * SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
 * INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
 * CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
 * ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
 * POSSIBILITY OF SUCH DAMAGE.
 */

#include <err.h>
#include <stdio.h>
#include <string.h>

/* OP-TEE TEE client API (built by optee_client) */
#include <tee_client_api.h>

/* TA API: UUID and command IDs */
#include <bliq_storage_ta.h>

/* TEE resources */
struct test_ctx {
	TEEC_Context ctx;
	TEEC_Session sess;
};

void open_tee_session(
	struct test_ctx *ctx,
	char *pin, size_t pin_len)
{
	TEEC_Operation op;
	TEEC_UUID uuid = TA_BLIQ_STORAGE_UUID;
	uint32_t origin;
	TEEC_Result res;

	// Set session PIN
	memset(&op, 0, sizeof(op));
	op.paramTypes = TEEC_PARAM_TYPES(TEEC_MEMREF_TEMP_INPUT, TEEC_NONE, TEEC_NONE, TEEC_NONE);

	op.params[0].tmpref.buffer = pin;
	op.params[0].tmpref.size = pin_len;

	/* Initialize a context connecting us to the TEE */
	res = TEEC_InitializeContext(NULL, &ctx->ctx);
	if (res != TEEC_SUCCESS)
		errx(1, "TEEC_InitializeContext failed with code 0x%x", res);

	/* Open a session with the TA */
	res = TEEC_OpenSession(&ctx->ctx, &ctx->sess, &uuid,
			       TEEC_LOGIN_PUBLIC, NULL, &op, &origin);
	if (res != TEEC_SUCCESS)
		errx(1, "TEEC_Opensession failed with code 0x%x origin 0x%x",
			res, origin);
}

void close_tee_session(struct test_ctx *ctx)
{
	TEEC_CloseSession(&ctx->sess);
	TEEC_FinalizeContext(&ctx->ctx);
}

TEEC_Result create_key(
	struct test_ctx *ctx, 
	char *id)
{
	TEEC_Operation op;
	uint32_t origin;
	TEEC_Result res;
	size_t id_len = strlen(id);

	memset(&op, 0, sizeof(op));
	op.paramTypes = TEEC_PARAM_TYPES(
			TEEC_MEMREF_TEMP_INPUT,
			TEEC_NONE,
			TEEC_NONE,
			TEEC_NONE);

	op.params[0].tmpref.buffer = id;
	op.params[0].tmpref.size = id_len;

	res = TEEC_InvokeCommand(
				&ctx->sess,
				TA_BLIQ_STORAGE_CMD_CREATE_KEY,
				&op, &origin);

	switch (res) {
	case TEEC_SUCCESS:
		break;
	default:
		printf("Command TA_BLIQ_STORAGE_CMD_CREATE_KEY failed: 0x%x / %u\n", res, origin);
	}

	return res;
}

TEEC_Result check_key(
	struct test_ctx *ctx, 
	char *id)
{
	TEEC_Operation op;
	uint32_t origin;
	TEEC_Result res;
	size_t id_len = strlen(id);

	memset(&op, 0, sizeof(op));
	op.paramTypes = TEEC_PARAM_TYPES(
			TEEC_MEMREF_TEMP_INPUT,
			TEEC_NONE,
			TEEC_NONE,
			TEEC_NONE);

	op.params[0].tmpref.buffer = id;
	op.params[0].tmpref.size = id_len;

	res = TEEC_InvokeCommand(
			&ctx->sess,
			TA_BLIQ_STORAGE_CMD_CHECK_KEY,
			&op, &origin);

	switch (res) {
	case TEEC_SUCCESS:
	case TEEC_ERROR_ITEM_NOT_FOUND:
		break;
	default:
		printf("Command TA_BLIQ_STORAGE_CMD_CHECK_KEY failed: 0x%x / %u\n", res, origin);
	}

	return res;
}

TEEC_Result delete_key(
	struct test_ctx *ctx, 
	char *id)
{
	TEEC_Operation op;
	uint32_t origin;
	TEEC_Result res;
	size_t id_len = strlen(id);

	memset(&op, 0, sizeof(op));
	op.paramTypes = TEEC_PARAM_TYPES(
			TEEC_MEMREF_TEMP_INPUT,
			TEEC_NONE,
			TEEC_NONE,
			TEEC_NONE);

	op.params[0].tmpref.buffer = id;
	op.params[0].tmpref.size = id_len;

	res = TEEC_InvokeCommand(
			&ctx->sess,
			TA_BLIQ_STORAGE_CMD_DELETE_KEY,
			&op, &origin);

	switch (res) {
	case TEEC_SUCCESS:
	case TEEC_ERROR_ITEM_NOT_FOUND:
		break;
	default:
		printf("Command TA_BLIQ_STORAGE_CMD_DELETE_KEY failed: 0x%x / %u\n", res, origin);
	}

	return res;
}

TEEC_Result encrypt(
	struct test_ctx *ctx, 
	char *id,
	char *data, uint32_t data_len,
	char *encr, uint32_t *encr_len
	)
{
	TEEC_Operation op;
	uint32_t origin;
	TEEC_Result res;
	size_t id_len = strlen(id);

	memset(&op, 0, sizeof(op));
	op.paramTypes = TEEC_PARAM_TYPES(
			TEEC_MEMREF_TEMP_INPUT,
			TEEC_MEMREF_TEMP_INPUT,
			TEEC_MEMREF_TEMP_INOUT,
			TEEC_NONE);

	op.params[0].tmpref.buffer = id;
	op.params[0].tmpref.size = id_len;

	op.params[1].tmpref.buffer = data;
	op.params[1].tmpref.size = data_len;

	op.params[2].tmpref.buffer = encr;
	op.params[2].tmpref.size  = *encr_len;

	res = TEEC_InvokeCommand(&ctx->sess,
				 TA_BLIQ_STORAGE_CMD_ENCRYPT,
				 &op, &origin);

	switch (res) {
	case TEEC_SUCCESS:
		*encr_len = op.params[2].tmpref.size;
		break;
	default:
		printf("Command TA_BLIQ_STORAGE_CMD_ENCRYPT failed: 0x%x / %u\n", res, origin);
	}

	return res;
}

TEEC_Result decrypt(
	struct test_ctx *ctx, 
	char *id,
	char *encr, uint32_t encr_len,
	char *decr, uint32_t *decr_len
	)
{
	TEEC_Operation op;
	uint32_t origin;
	TEEC_Result res;
	size_t id_len = strlen(id);

	memset(&op, 0, sizeof(op));
	op.paramTypes = TEEC_PARAM_TYPES(
			TEEC_MEMREF_TEMP_INPUT,
			TEEC_MEMREF_TEMP_INPUT,
			TEEC_MEMREF_TEMP_INOUT,
			TEEC_NONE);

	op.params[0].tmpref.buffer = id;
	op.params[0].tmpref.size = id_len;

	op.params[1].tmpref.buffer = encr;
	op.params[1].tmpref.size = encr_len;

	op.params[2].tmpref.buffer = decr;
	op.params[2].tmpref.size  = *decr_len;

	res = TEEC_InvokeCommand(&ctx->sess,
				 TA_BLIQ_STORAGE_CMD_DECRYPT,
				 &op, &origin);

	switch (res) {
	case TEEC_SUCCESS:
		*decr_len = op.params[2].tmpref.size;
		break;
	default:
		printf("Command TA_BLIQ_STORAGE_CMD_DECRYPT failed: 0x%x / %u\n", res, origin);
	}

	return res;
}



int main(void)
{
	struct test_ctx ctx;
	char key_id[] = "key#1";
	char pin[6] = {0x00, 0x01, 0x02, 0x03, 0x04, 0x05}; // used to encrypt/decrypt key

	char data[] = "Test string for data encryption";

	char encrypted_data[512];
	uint32_t encrypted_data_sz = sizeof(encrypted_data);

	char decrypted_data[512];
	uint32_t decrypted_data_sz = sizeof(decrypted_data);

	TEEC_Result res;

	printf("Prepare session with the TA\n");
	open_tee_session(&ctx, pin, sizeof(pin));

	/*
	 * Check Create , read it, delete it.
	 */
	printf("\nCheck KeyID \"%s\"\n", key_id);
	res = check_key(
			&ctx, 
			key_id);
	if (res == TEEC_ERROR_ITEM_NOT_FOUND) {
		printf(" - KeyId was not found in TA secure storage, create new one\n");
		
		res = create_key(
				&ctx, 
				key_id);
		if (res != TEEC_SUCCESS)
			errx(1, "Unexpected status when creating key: 0x%x", res);

		printf(" - RSA keypair creation [OK]\n");
		printf(" - Run application one more time to test Encrypt / Decrypt APIs..\n");
	}
	else if (res == TEEC_SUCCESS) {
		printf(" - KeyId was found!\n");
		printf(" - Test Encrypt / Decrypt APIs..\n");

		res = encrypt(
				&ctx, 
				key_id, 
				data, sizeof(data),
				encrypted_data, &encrypted_data_sz);
		if (res != TEEC_SUCCESS)
			errx(1, "Unexpected status during data encryption: 0x%x", res);

		printf("  -- encr_data: [%u] first 4 bytes: %02x:%02x:%02x:%02x\n",
				encrypted_data_sz, encrypted_data[0], encrypted_data[1], encrypted_data[2], encrypted_data[3]);

		res = decrypt(&ctx, key_id, 
				encrypted_data, encrypted_data_sz,
				decrypted_data, &decrypted_data_sz);
		if (res != TEEC_SUCCESS)
			errx(1, "Unexpected status during data decryption: 0x%x", res);

		printf("  -- orig_data: [%lu] '%s'\n  -- decr_data: [%d] '%s'\n", 
				sizeof(data), data,  
				decrypted_data_sz, decrypted_data);

		printf(" - Encrypt / Decrypt APIs test [ok]\n");
		printf(" - Delete RSA keypair from secure storage\n");

		res = delete_key(
			&ctx,
			key_id);
		if (res != TEEC_SUCCESS)
			errx(1, "Unexpected status during key delition: 0x%x", res);
	}

	close_tee_session(&ctx);
	return 0;
}
