#include <stdlib.h>
#include <common.h>
#include <command.h>
#include <linux/ctype.h>
#include <tee/ta_vx.h>
#include <tee/ta_eva.h>
#include <tee/ta_helper.h>
#include <tee/ta_vx_helper.h>

static const struct tee_optee_ta_uuid ta_vx_uuid = TA_VX_UUID;

/**
 * Prints VX LOCKED/UNLOCKED status and metadata.
 */
static int print_lock_status(void)
{
	int rc;
	struct ta_context context = { 0 };
	struct tee_param params[1] = { 0 };
	struct tee_shm *shm_buf = NULL;
	const uint32_t nonce_size = TA_VX_NONCE_SIZE;

	rc = ta_open(&ta_vx_uuid, &context);
	if (rc)
		return rc;

	/*
	 * Read and print device state.
	 */
	params[0].attr = TEE_PARAM_ATTR_TYPE_VALUE_OUTPUT;
	rc = ta_invoke(&context, TA_VX_CMD_READ_LOCK_STATE, 1, params);
	if (rc)
		goto out;

	printf("Device state: %s\n",
	       (params[0].u.value.a) ? "UNLOCKED" : "LOCKED");

	/*
	 * Read and print lock state nonce.
	 */
	rc = tee_shm_alloc(context.tee, nonce_size, TEE_SHM_ALLOC, &shm_buf);
	if (rc)
		goto out;
	memset(&params[0], 0, sizeof(params[0]));
	params[0].attr = TEE_PARAM_ATTR_TYPE_MEMREF_INOUT;
	params[0].u.memref.shm = shm_buf;
	params[0].u.memref.size = nonce_size;
	rc = ta_invoke(&context, TA_VX_CMD_GET_LOCK_STATE_NONCE, 1, params);
	if (rc)
		goto out;

	printf("  (nonce: ");
	assert(params[0].u.memref.size == nonce_size);
	uint8_t *nonce = (uint8_t *)params[0].u.memref.shm->addr;
	for (int i = 0; i < nonce_size; i++)
		printf("%02X", nonce[i]);
	printf(")\n");

out:
	tee_shm_free(shm_buf);
	ta_close(&context);
	return rc;
}

static int list_persist_values(void)
{
	int rc;
	struct ta_context context = { 0 };
	struct tee_param params[2] = {0};
	struct tee_shm *shm_buf = NULL;
	const uint32_t buf_sz = 1024;

	rc = ta_open(&ta_vx_uuid, &context);
	if (rc)
		return rc;

	rc = tee_shm_alloc(context.tee, buf_sz, TEE_SHM_ALLOC, &shm_buf);
	if (rc)
		goto out;
	params[0].attr = TEE_PARAM_ATTR_TYPE_VALUE_OUTPUT;
	params[1].attr = TEE_PARAM_ATTR_TYPE_MEMREF_INOUT;
	params[1].u.memref.shm = shm_buf;
	params[1].u.memref.size = buf_sz;
	rc = ta_invoke(&context, TA_VX_CMD_ENUMERATE_PERSIST_VALUES, 2, params);
	if (rc)
		goto out;

	uint32_t count = params[0].u.value.a;
	if (count) {
		printf("Persistent values:");
		char *names = (char *)params[1].u.memref.shm->addr;
		while (count) {
			printf(" %s", names);
			names += (strlen(names) + 1);
			count--;
		}
		printf("\n");
	}

out:
	tee_shm_free(shm_buf);
	ta_close(&context);
	return rc;
}

/**
 * Lists all non-zero rollback indexes.
 */
static int list_rollback_indexes(void)
{
	uint64_t index;

	printf("Rollback indexes:");
	for (uint32_t slot = 0; slot < TA_VX_MAX_ROLLBACK_LOCATIONS; slot++) {
		int rc = ta_vx_read_rollback_index(slot, &index);
		if (rc) {
			printf("\nError reading slot #%u\n", slot);
			return rc;
		}

		if (index)
			printf(" #%u=%llu", slot, index);
	}
	printf("\n");

	return 0;
}

static int do_status(cmd_tbl_t *cmdtp, int flag, int argc, char *const argv[])
{
	int error_count = 0;

	if (print_lock_status()) {
		printf("Error reading lock status.\n");
		error_count++;
	}

	if (list_persist_values()) {
		printf("Error reading named persistent values.\n");
		error_count++;
	}

	if (list_rollback_indexes()) {
		printf("Error reading rollback indexes.\n");
		error_count++;
	}

	return error_count? -EINVAL : 0;
}


static int do_fdr(cmd_tbl_t *cmdtp, int flag, int argc, char *const argv[])
{
	int rc = ta_call(&ta_vx_uuid, TA_VX_CMD_REFRESH_LOCK_STATE_NONCE, 0,
			 NULL);
	if (!rc)
		printf("ALL USER DATA HAS BEEN ERASED!\n");

	return rc;
}

static int do_lock(cmd_tbl_t *cmdtp, int flag, int argc, char *const argv[])
{
	return ta_vx_lock();
}

static int do_unlock(cmd_tbl_t *cmdtp, int flag, int argc, char *const argv[])
{
	return ta_vx_unlock();
}

static int do_rbi(cmd_tbl_t *cmdtp, int flag, int argc, char *const argv[])
{
	uint32_t slot;
	uint64_t index;
	int rc;

	if (argc <= 1)
		return list_rollback_indexes();

	slot = simple_strtoul(argv[1], NULL, 0);
	if (argc > 2) {
		/* Update slot with a new index. */
		index = simple_strtoull(argv[2], NULL, 0);
		rc = ta_vx_write_rollback_index(slot, index);
		return rc;
	}

	rc = ta_vx_read_rollback_index(slot, &index);
	if (!rc)
		printf("#%u=%llu\n", slot, index);
	return rc;
}

static void print_as_string(const char *buf, size_t len)
{
	for (size_t i = 0; i < len; i++)
		printf("%c", isprint(buf[i]) ? buf[i] : '.');
	printf("\n");
}

static void print_as_hex(const char *buf, size_t len)
{
	for (size_t i = 0; i < len; i++) {
		if (i % 16 == 0)
			printf("%08zx  ", i);

		printf("%02x ", buf[i]);

		if (i % 16 == 15)
			printf("\n");
		else if (i % 16 == 7)
			printf(" ");
	}
	printf("\n");
}

static int do_get(cmd_tbl_t *cmdtp, int flag, int argc, char *const argv[])
{
	if (argc < 2 || !argv[1][0])
		return -EINVAL;

	/* Maximum number of bytes to read. */
	size_t max_bytes = (argc > 2) ? simple_strtoul(argv[2], NULL, 0) : 128;
	if (max_bytes == 0)
		return -EINVAL;

	const char *name = argv[1];
	char *buf = malloc(max_bytes);
	if (!buf)
		return -ENOMEM;

	size_t bytes_read;
	int rc = ta_vx_read_persistent_value(name, buf, max_bytes, &bytes_read);
	if (rc) {
		free(buf);
		return rc;
	}

	printf("Hex dump (%zu bytes):\n", bytes_read);
	print_as_hex(buf, bytes_read);
	printf("String dump:\n");
	print_as_string(buf, bytes_read);

	free(buf);
	return 0;
}

static int do_set(cmd_tbl_t *cmdtp, int flag, int argc, char *const argv[])
{
	if (argc < 3 || !argv[1][0] || !argv[2][0])
		return -EINVAL;

	const char *name = argv[1];
	const char *value = argv[2];

	return ta_vx_write_persistent_value(name, value, strlen(value) + 1);
}

static int do_unset(cmd_tbl_t *cmdtp, int flag, int argc, char *const argv[])
{
	if (argc < 2 || !argv[1][0])
		return -EINVAL;

	const char *name = argv[1];

	int rc = ta_vx_delete_persistent_value(name);
	if (rc) {
		printf("Error deleting %s\n", name);
		return rc;
	}

	printf("Poof, it's gone!\n");
	return 0;
}

static int do_random(cmd_tbl_t *cmdtp, int flag, int argc, char *const argv[])
{
	size_t bytes = 32;
	if (argc > 1)
		bytes = simple_strtoul(argv[1], NULL, 0);

	if (bytes == 0)
		return -EINVAL;

	uint8_t *buffer = malloc(bytes);
	if (!buffer)
		return -ENOMEM;

	int rc = ta_vx_cprng_draw(buffer, bytes);
	if (rc) {
		free(buffer);
		return rc;
	}

	for (size_t i = 0; i < bytes; i++)
		printf("%02X", buffer[i]);
	printf("\n");

	free(buffer);
	return 0;
}

static int do_finalize(cmd_tbl_t *cmdtp, int flag, int argc, char *const argv[])
{
	return ta_call(&ta_vx_uuid, TA_VX_CMD_EXIT_BOOTLOADER,
		       /* num_params = */ 0, /* params = */ NULL);
}

static int do_ta2ta(cmd_tbl_t *cmdtp, int flag, int argc, char *const argv[])
{
	const struct tee_optee_ta_uuid ta_eva_uuid = TA_EVA_UUID;
	return ta_call(&ta_eva_uuid, TA_EVA_CMD_RUN_ALL_TESTS,
		       /* num_params = */ 0, /* params = */ NULL);
}

#ifdef CONFIG_TA_VX_TESTS
static int do_test(cmd_tbl_t *cmdtp, int flag, int argc, char *const argv[])
{
	const char *name = (argc > 1) ? argv[1]: "";
	return ta_vx_run_tests(name);
}
#endif

static cmd_tbl_t commands[] = {
	U_BOOT_CMD_MKENT(status, /* maxargs = */ 1, /* repeatable = */ 0,
			 do_status, "", ""),
	U_BOOT_CMD_MKENT(fdr, /* maxargs = */ 1, /* repeatable = */ 0, do_fdr,
			 "", ""),
	U_BOOT_CMD_MKENT(lock, /* maxargs = */ 1, /* repeatable = */ 0, do_lock,
			 "", ""),
	U_BOOT_CMD_MKENT(unlock, /* maxargs = */ 1, /* repeatable = */ 0,
			 do_unlock, "", ""),
	U_BOOT_CMD_MKENT(random, /* maxargs = */ 2, /* repeatable = */ 0,
			 do_random, "", ""),
	U_BOOT_CMD_MKENT(rbi, /* maxargs = */ 3, /* repeatable = */ 0,
			 do_rbi, "", ""),
	U_BOOT_CMD_MKENT(set, /* maxargs = */ 3, /* repeatable = */ 0, do_set,
			 "", ""),
	U_BOOT_CMD_MKENT(unset, /* maxargs = */ 2, /* repeatable = */ 0, do_unset,
			 "", ""),
	U_BOOT_CMD_MKENT(get, /* maxargs = */ 3, /* repeatable = */ 0, do_get,
			 "", ""),
	U_BOOT_CMD_MKENT(finalize, /* maxargs = */ 1, /* repeatable = */ 0,
			 do_finalize, "", ""),
	U_BOOT_CMD_MKENT(ta2ta, /* maxargs = */ 1, /* repeatable = */ 0,
			 do_ta2ta, "", ""),
#ifdef CONFIG_TA_VX_TESTS
	U_BOOT_CMD_MKENT(test, /* maxargs = */ 2, /* repeatable = */ 1, do_test,
			 "", ""),
#endif
};

static int do_vx(cmd_tbl_t *cmdtp, int flag, int argc, char *const argv[])
{
	cmd_tbl_t *cmd;
	int ret;

	if (argc < 2)
		return CMD_RET_USAGE;

	argc--;
	argv++;

	cmd = find_cmd_tbl(argv[0], commands, ARRAY_SIZE(commands));
	if (!cmd || argc > cmd->maxargs)
		return CMD_RET_USAGE;

	ret = cmd->cmd(cmd, flag, argc, argv);
	return cmd_process_error(cmdtp, ret);
}

U_BOOT_CMD(
	vx, /* maxargs = */ 4, /* repeatable = */ 0, do_vx,
	"Verified Execution Operations.",
	"status - Prints VX status.\n"
	"vx fdr - Erases user data.\n"
	"vx lock - Locks the Bootloader.\n"
	"vx unlock - Unlocks the Bootloader.\n"
	"vx rbi [slot [index]] - List/get/update rollback index(es).\n"
	"vx random [bytes=32] - Draw random bytes.\n"
	"vx set <name> <str> - Set/updates a named persistent value.\n"
	"vx unset <name> - Unset/deletes a named persistent value.\n"
	"vx get <name> [max_bytes]- Dumps a named persistent value.\n"
	"vx finalize - Finalizes VX policies.\n"
	"vx ta2ta - Run TA2TA tests (requires EVA TA).\n"
#ifdef CONFIG_TA_VX_TESTS
	"vx test [name] - Run VX TA unit tests.\n"
#endif
	)