
/*
 * (C) Copyright 2018, Linaro Limited
 *
 * SPDX-License-Identifier:	GPL-2.0+
 */

#include <avb_verify.h>
#include <command.h>
#include <image.h>
#include <malloc.h>
#include <mmc.h>
#include <android_image.h>
#include <image-android-sig.h>
#include <avbctl_uboot.h>
#include <debug_impl.h>
#include <utility.h>
#include <generated/bootcommands_autogenerated.h>

#define AVB_BOOTARGS	"avb_bootargs"
#define DEVICESTATE_ISLOCKED_BASE           16
static struct AvbOps *avb_ops;

//MTK code starts
#if defined(CONFIG_ANDROID_AVB_PURELINUX_ENABLE)//For MLR
static const char * const requested_partitions[] = {"KL", "DTB", "dtbo",
					     NULL};
#define BOOT_PARTNAME "KL"
static const char * const requested_partitions_recovery_mode[] = {NULL};
#define RECOVERY_PARTNAME ""
#define RECOVERY_VERIFY_FLAG AVB_SLOT_VERIFY_FLAGS_NONE
#else//For Android
static const char * const requested_partitions[] = {"boot","dtbo",
					     "system",
					     "vendor",
					     "vendor_boot",
					     "init_boot",
					     NULL};
#define BOOT_PARTNAME "boot"

/*Recovery partition depended on Boot header version to verify, check table:
  ----------------------------------------------------------------------
  Android             |          AN-R        |         AN-S
  ----------------------------------------------------------------------
  Boot header version |    V1    |   V2      |  V3        |      V4
  ----------------------------------------------------------------------
  nonAB               | recovery |  recovery | recovery   |  vendor_boot
  ----------------------------------------------------------------------
  AB                  |   boot   |   boot   | vendor_boot |  vendor_boot
  ----------------------------------------------------------------------

When boot to recovery mode, the recovery partition should be verified by AVB
for boot safety and reliability.

On Android devices not using A/B, the `recovery` partition usually isn't updated along
with other partitions and therefore can't be referenced from the main `vbmeta` partition.

So the LIBAVB provided a flag called 'AVB_SLOT_VERIFY_FLAGS_NO_VBMETA_PARTITION' to support it.
When AVB_SLOT_VERIFY_FLAGS_NO_VBMETA_PARTITION has been selected, AVB would not checked the vbmeta.
Instead, AVB would validated each partition in request_partition list.

And we used the flag: 'RECOVERY_VERIFY_FLAG' to use it:
if AVB_SLOT_VERIFY_FLAGS_NO_VBMETA_PARTITION is select by RECOVERY_VERIFY_FLAG,
AVB would not checked vbmeta, and verify each partition which contain in requested_partitions_recovery_mode[].
if AVB_SLOT_VERIFY_FLAGS_NONE is select by RECOVERY_VERIFY_FLAG,
AVB would checked vbmeta, and verify the partitions contain in vbmeta.

The recovery partition is exist or into vendor_boot are depended on Boot header version
And it is influenced the vbmeta should check or not, the case are show as below table:
  ----------------------------------------------------------------------
  Boot header version |    V1      |     V2      |  V3         |      V4
  ----------------------------------------------------------------------
  nonAB               |not checked | not checked | not checked |  checked
  ----------------------------------------------------------------------
  AB                  |  checked   |   checked   | checked     |  checked
  ----------------------------------------------------------------------
*/
#if !defined(CONFIG_ANDROID_AB)
#if defined(ANDROID_BOOT_HEADER_VERSION_V1) || defined(ANDROID_BOOT_HEADER_VERSION_V2) || defined(ANDROID_BOOT_HEADER_VERSION_V3)
static const char * const requested_partitions_recovery_mode[] = {"recovery",
					     NULL};
#define RECOVERY_VERIFY_FLAG AVB_SLOT_VERIFY_FLAGS_NO_VBMETA_PARTITION
#define RECOVERY_PARTNAME "recovery"

#elif defined(ANDROID_BOOT_HEADER_VERSION_V4)
static const char * const * requested_partitions_recovery_mode = requested_partitions;
#define RECOVERY_VERIFY_FLAG AVB_SLOT_VERIFY_FLAGS_NONE
#define RECOVERY_PARTNAME "vendor_boot"

#else
static const char * const * requested_partitions_recovery_mode = requested_partitions;
#define RECOVERY_VERIFY_FLAG AVB_SLOT_VERIFY_FLAGS_NONE
#define RECOVERY_PARTNAME "vendor_boot"
#endif

#else//defined(CONFIG_ANDROID_AB)
static const char * const * requested_partitions_recovery_mode = requested_partitions;
#define RECOVERY_VERIFY_FLAG AVB_SLOT_VERIFY_FLAGS_NONE

#if defined(ANDROID_BOOT_HEADER_VERSION_V1) || defined(ANDROID_BOOT_HEADER_VERSION_V2)
#define RECOVERY_PARTNAME "boot"
#elif defined(ANDROID_BOOT_HEADER_VERSION_V3) || defined(ANDROID_BOOT_HEADER_VERSION_V4)
#define RECOVERY_PARTNAME "vendor_boot"
#else
#define RECOVERY_PARTNAME "vendor_boot"
#endif
#endif

#endif

static void get_verified_parts(const char* const* requested_partitions, AvbPartitionData* part_data, size_t num_loaded_partitions , const char** verified_parts)
{
	int i,j;
	for(i=0; i < num_loaded_partitions; i++)
	{
		verified_parts[i] = NULL;
		for(j=0; requested_partitions[j] != NULL; j++)
		{
			if(!strncmp(part_data[i].partition_name, requested_partitions[j], strlen(part_data[i].partition_name)))
			{
				verified_parts[i] = part_data[i].partition_name;
			}
		}
	}
}

#if (CONFIG_ROLLBACK_INDEX_IN_EFUSE == 1)
#include <program_efuse_rollback_index.h>
#elif (CONFIG_AB_SIDELOAD == 1)
#define DTB_VERSION_STRUCT_MAGIC "avb_version     "
#define DTB_VERSION_STRUCT_MAGIC_LEN 16
typedef struct
{
	uint64_t ta_avb_rollback_indexes[AVB_MAX_NUMBER_OF_ROLLBACK_INDEX_LOCATIONS];
	char magic[DTB_VERSION_STRUCT_MAGIC_LEN]; //magic = "avb_version     "
}avb_version; //272 bytes

static avb_version version;

static int rollbackindexes_to_dtb(AvbSlotVerifyData* AvbData)
{
	int i;
	memset(&version, 0, sizeof(avb_version));
	avb_memcpy(version.magic, DTB_VERSION_STRUCT_MAGIC, DTB_VERSION_STRUCT_MAGIC_LEN);
	for(i=0; i<AVB_MAX_NUMBER_OF_ROLLBACK_INDEX_LOCATIONS; i++)
	{
		version.ta_avb_rollback_indexes[i] = AvbData->rollback_indexes[i];
	}
	UBOOT_DUMP(&version, sizeof(avb_version));
	setup_avb_data_for_dtb((char*)(&version), sizeof(avb_version));
	return 0;
}
#endif
//MTK code ends

int do_avb_init(cmd_tbl_t *cmdtp, int flag, int argc, char * const argv[])
{
	unsigned long mmc_dev;
	u64 start = 0, end = 0;

	start = get_timer(0);
	if (argc != 2)
		return CMD_RET_USAGE;

	mmc_dev = simple_strtoul(argv[1], NULL, 16);

	if (avb_ops)
		avb_ops_free(avb_ops);

	avb_ops = avb_ops_alloc(mmc_dev);
	if (avb_ops) {
		end = get_timer(0);
		UBOOT_BOOTTIME("[avb init][start:%llu][end:%llu][total time:%llu]\n", start, end, end - start);
		return CMD_RET_SUCCESS;
	}

	printf("Failed to initialize avb2\n");

	return CMD_RET_FAILURE;
}

int do_avb_read_part(cmd_tbl_t *cmdtp, int flag, int argc, char * const argv[])
{
	const char *part;
	s64 offset;
	size_t bytes, bytes_read = 0;
	void *buffer;

	if (!avb_ops) {
		printf("AVB 2.0 is not initialized, please run 'avb init'\n");
		return CMD_RET_USAGE;
	}

	if (argc != 5)
		return CMD_RET_USAGE;

	part = argv[1];
	offset = simple_strtoul(argv[2], NULL, 16);
	bytes = simple_strtoul(argv[3], NULL, 16);
	buffer = (void *)simple_strtoul(argv[4], NULL, 16);

	if (avb_ops->read_from_partition(avb_ops, part, offset, bytes,
					 buffer, &bytes_read) ==
					 AVB_IO_RESULT_OK) {
		printf("Read %zu bytes\n", bytes_read);
		return CMD_RET_SUCCESS;
	}

	printf("Failed to read from partition\n");

	return CMD_RET_FAILURE;
}

int do_avb_read_part_hex(cmd_tbl_t *cmdtp, int flag, int argc,
			 char *const argv[])
{
	const char *part;
	s64 offset;
	size_t bytes, bytes_read = 0;
	char *buffer;

	if (!avb_ops) {
		printf("AVB 2.0 is not initialized, please run 'avb init'\n");
		return CMD_RET_USAGE;
	}

	if (argc != 4)
		return CMD_RET_USAGE;

	part = argv[1];
	offset = simple_strtoul(argv[2], NULL, 16);
	bytes = simple_strtoul(argv[3], NULL, 16);

	buffer = malloc(bytes);
	if (!buffer) {
		printf("Failed to tlb_allocate buffer for data\n");
		return CMD_RET_FAILURE;
	}
	memset(buffer, 0, bytes);

	if (avb_ops->read_from_partition(avb_ops, part, offset, bytes, buffer,
					 &bytes_read) == AVB_IO_RESULT_OK) {
		printf("Requested %zu, read %zu bytes\n", bytes, bytes_read);
		printf("Data: ");
		for (int i = 0; i < bytes_read; i++)
			printf("%02X", buffer[i]);

		printf("\n");

		free(buffer);
		return CMD_RET_SUCCESS;
	}

	printf("Failed to read from partition\n");

	free(buffer);
	return CMD_RET_FAILURE;
}

int do_avb_write_part(cmd_tbl_t *cmdtp, int flag, int argc, char * const argv[])
{
	const char *part;
	s64 offset;
	size_t bytes;
	void *buffer;

	if (!avb_ops) {
		printf("AVB 2.0 is not initialized, run 'avb init' first\n");
		return CMD_RET_FAILURE;
	}

	if (argc != 5)
		return CMD_RET_USAGE;

	part = argv[1];
	offset = simple_strtoul(argv[2], NULL, 16);
	bytes = simple_strtoul(argv[3], NULL, 16);
	buffer = (void *)simple_strtoul(argv[4], NULL, 16);

	if (avb_ops->write_to_partition(avb_ops, part, offset, bytes, buffer) ==
	    AVB_IO_RESULT_OK) {
		printf("Wrote %zu bytes\n", bytes);
		return CMD_RET_SUCCESS;
	}

	printf("Failed to write in partition\n");

	return CMD_RET_FAILURE;
}

int do_avb_read_rb(cmd_tbl_t *cmdtp, int flag, int argc, char * const argv[])
{
	size_t index;
	u64 rb_idx;

	if (!avb_ops) {
		printf("AVB 2.0 is not initialized, run 'avb init' first\n");
		return CMD_RET_FAILURE;
	}

	if (argc != 2)
		return CMD_RET_USAGE;

	index = (size_t)simple_strtoul(argv[1], NULL, 16);

	if (avb_ops->read_rollback_index(avb_ops, index, &rb_idx) ==
	    AVB_IO_RESULT_OK) {
		printf("Rollback index: %llx\n", rb_idx);
		return CMD_RET_SUCCESS;
	}

	printf("Failed to read rollback index\n");

	return CMD_RET_FAILURE;
}

int do_avb_write_rb(cmd_tbl_t *cmdtp, int flag, int argc, char * const argv[])
{
	size_t index;
	u64 rb_idx;

	if (!avb_ops) {
		printf("AVB 2.0 is not initialized, run 'avb init' first\n");
		return CMD_RET_FAILURE;
	}

	if (argc != 3)
		return CMD_RET_USAGE;

	index = (size_t)simple_strtoul(argv[1], NULL, 16);
	rb_idx = simple_strtoul(argv[2], NULL, 16);

	if (avb_ops->write_rollback_index(avb_ops, index, rb_idx) ==
	    AVB_IO_RESULT_OK)
		return CMD_RET_SUCCESS;

	printf("Failed to write rollback index\n");

	return CMD_RET_FAILURE;
}

int do_avb_get_uuid(cmd_tbl_t *cmdtp, int flag,
		    int argc, char * const argv[])
{
	const char *part;
	char buffer[UUID_STR_LEN + 1];

	if (!avb_ops) {
		printf("AVB 2.0 is not initialized, run 'avb init' first\n");
		return CMD_RET_FAILURE;
	}

	if (argc != 2)
		return CMD_RET_USAGE;

	part = argv[1];

	if (avb_ops->get_unique_guid_for_partition(avb_ops, part, buffer,
						   UUID_STR_LEN + 1) ==
						   AVB_IO_RESULT_OK) {
		printf("'%s' UUID: %s\n", part, buffer);
		return CMD_RET_SUCCESS;
	}

	printf("Failed to read UUID\n");

	return CMD_RET_FAILURE;
}

static int do_remove_key(char *key)
{
	int ret = 0;

	do {
		ret = del_bootargs(key, 0);
		if (ret == -1)
			return ret;
	} while (ret != 1);
/* we should check all duplicate avb_args are have been clean in bootargs :
 * return -1 means something unexpect error occur
 * eq.the source bootargs or the delete key is NULL or zero size, or someting malloc fail, etc
 * return 0 means we find one avb args is exist and delete it success,
 * and then we keep find until all duplicate avb_args are delete.
 * return 1 means we can not find any duplicate avb_args exist and the do_remove_key can stop.
 */

	return ret;
}

static int do_remove_prestored_avb_args(void)
{
	if (do_remove_key("androidboot.vbmeta.device") != 1)
		return -1;
	if (do_remove_key("androidboot.vbmeta.avb_version") != 1)
		return -1;
	if (do_remove_key("androidboot.vbmeta.device_state") != 1)
		return -1;
	if (do_remove_key("androidboot.vbmeta.hash_alg") != 1)
		return -1;
	if (do_remove_key("androidboot.vbmeta.size") != 1)
		return -1;
	if (do_remove_key("androidboot.vbmeta.digest") != 1)
		return -1;
	if (do_remove_key("androidboot.vbmeta.invalidate_on_error") != 1)
		return -1;
	if (do_remove_key("androidboot.veritymode") != 1)
		return -1;
	if (do_remove_key("androidboot.veritymode.managed") != 1)
		return -1;
	if (do_remove_key("androidboot.verifiedbootstate") != 1)
		return -1;

	return 1;
}

int do_avb_verify_part(cmd_tbl_t *cmdtp, int flag,
		       int argc, char *const argv[])
{
	AvbSlotVerifyResult slot_result;
	AvbSlotVerifyData *out_data;
	char *cmdline;
	char *extra_args = NULL;
	char *slot_suffix = "";
	enum avb_boot_state boot_state = AVB_RED;
	char *boot_partition_name = BOOT_PARTNAME;
	uint8_t *boot_img_addr = NULL;
	struct andr_img_hdr *boot_hdr;
	char *bootargs = NULL;
	char *newbootargs = NULL;

	bool unlocked = false;
	int res = CMD_RET_FAILURE;
	u64 start = 0, end = 0;

	start = get_timer(0);
	if (!avb_ops) {
		printf("AVB 2.0 is not initialized, run 'avb init' first\n");
		return CMD_RET_FAILURE;
	}

	if (argc < 1 || argc > 2)
		return CMD_RET_USAGE;

	if (argc == 2)
		slot_suffix = argv[1];

	printf("## Android Verified Boot 2.0 version %s\n",
	       avb_version_string());

	if (avb_ops->read_is_device_unlocked(avb_ops, &unlocked) !=
	    AVB_IO_RESULT_OK) {
		printf("Can't determine device lock state.\n");
		return CMD_RET_FAILURE;
	}

//for recovery mode
	if(is_recovery_mode()) {
		boot_partition_name = RECOVERY_PARTNAME;
		slot_result =
		  avb_slot_verify(avb_ops,
				requested_partitions_recovery_mode,
				slot_suffix,
				unlocked | RECOVERY_VERIFY_FLAG,
				AVB_HASHTREE_ERROR_MODE_RESTART_AND_INVALIDATE,
				&out_data);
		goto avb_slot_verify_done;
	}
	slot_result =
		avb_slot_verify(avb_ops,
				requested_partitions,
				slot_suffix,
				unlocked,
				AVB_HASHTREE_ERROR_MODE_RESTART_AND_INVALIDATE,
				&out_data);

avb_slot_verify_done:

	switch (slot_result) {
	case AVB_SLOT_VERIFY_RESULT_OK:
		/* Until we don't have support of changing unlock states, we
		 * assume that we are by default in locked state.
		 * So in this case we can boot only when verification is
		 * successful; we also supply in cmdline GREEN boot state
		 */
		printf("Verification passed successfully\n");
#if 0
		/* export additional bootargs to AVB_BOOTARGS env var */

		extra_args = avb_set_state(avb_ops, AVB_GREEN);
		if (extra_args)
			cmdline = append_cmd_line(out_data->cmdline,
						  extra_args);
		else
			cmdline = out_data->cmdline;

		env_set(AVB_BOOTARGS, cmdline);
		res = CMD_RET_SUCCESS;
#endif
		break;
	case AVB_SLOT_VERIFY_RESULT_ERROR_VERIFICATION:
		printf("Verification failed\n");
		break;
	case AVB_SLOT_VERIFY_RESULT_ERROR_IO:
		printf("I/O error occurred during verification\n");
		break;
	case AVB_SLOT_VERIFY_RESULT_ERROR_OOM:
		printf("OOM error occurred during verification\n");
		break;
	case AVB_SLOT_VERIFY_RESULT_ERROR_INVALID_METADATA:
		printf("Corrupted dm-verity metadata detected\n");
		break;
	case AVB_SLOT_VERIFY_RESULT_ERROR_UNSUPPORTED_VERSION:
		printf("Unsupported version avbtool was used\n");
		break;
	case AVB_SLOT_VERIFY_RESULT_ERROR_ROLLBACK_INDEX:
		printf("Checking rollback index failed\n");
		break;
	case AVB_SLOT_VERIFY_RESULT_ERROR_PUBLIC_KEY_REJECTED:
		printf("Public key was rejected\n");
		break;
	default:
		printf("Unknown error occurred\n");
	}

	if (unlocked == true
	&& (slot_result == AVB_SLOT_VERIFY_RESULT_OK
	|| slot_result == AVB_SLOT_VERIFY_RESULT_ERROR_VERIFICATION
	|| slot_result == AVB_SLOT_VERIFY_RESULT_ERROR_ROLLBACK_INDEX
	|| slot_result == AVB_SLOT_VERIFY_RESULT_ERROR_PUBLIC_KEY_REJECTED)) {
		boot_state = AVB_ORANGE;
		extra_args = avb_set_state(avb_ops, boot_state);
	}
	else if (unlocked == false && slot_result == AVB_SLOT_VERIFY_RESULT_ERROR_PUBLIC_KEY_REJECTED) {
		boot_state = AVB_YELLOW;
		extra_args = avb_set_state(avb_ops, boot_state);
	}
	else if (unlocked == false && slot_result == AVB_SLOT_VERIFY_RESULT_OK) {
		boot_state = AVB_GREEN;
		extra_args = avb_set_state(avb_ops, boot_state);
	}

	if (boot_state == AVB_GREEN || boot_state == AVB_ORANGE) {
		if (extra_args)
			cmdline = append_cmd_line(out_data->cmdline, extra_args);
		else
			cmdline = out_data->cmdline;
		env_set(AVB_BOOTARGS, cmdline);

/* check avb_cmdline to find args are not duplicated at non-volatile dvice,
 * we would remove it before append avb args to make sure the final bootargs are legal
 */
		if (do_remove_prestored_avb_args() != 1) {
			UBOOT_ERROR("try remove pre-stored avb_args fail!\n");
			return CMD_RET_FAILURE;
		}

#if (CONFIG_MTK_ANDROID_HEADER_VERSION > 3)
		bootargs = env_get("bootconfig");
#else
		bootargs = env_get("bootargs");
#endif
		newbootargs = avb_calloc(CONFIG_SYS_CBSIZE);
		if (newbootargs == NULL) {
			return CMD_RET_FAILURE;
		}
		else {
			if(snprintf(newbootargs, CONFIG_SYS_CBSIZE, "%s %s", bootargs, cmdline) < 0) {
				free(newbootargs);
				return CMD_RET_FAILURE;
			}
#if (CONFIG_MTK_ANDROID_HEADER_VERSION > 3)
			env_set("bootconfig", newbootargs);
#else
			env_set("bootargs", newbootargs);
#endif
			free(newbootargs);
		}
		//get boot image loaded address.
		for(int i=0; i < out_data->num_loaded_partitions; i++) {
			if(strncmp(boot_partition_name, out_data->loaded_partitions[i].partition_name, 4) == 0) {
				boot_img_addr = out_data->loaded_partitions[i].data;
				break;
			}
		}

		boot_hdr = (struct andr_img_hdr *)boot_img_addr;
		if(boot_hdr != NULL) {
			if(check_android_boot_state(unlocked, out_data, slot_result, boot_hdr->os_version, avb_ops) != 0)	{
				return CMD_RET_FAILURE;
			}
#if (CONFIG_ROLLBACK_INDEX_IN_EFUSE == 1)
			//For CONFIG_ROLLBACK_INDEX_IN_EFUSE, we need to invoke ramlog TA at all times to activate once-calling lock in TEE kernel
			//despite status of is_pass_rollback_index_info_needed().
			if (rollbackindexes_to_efuse(out_data)) {
				return CMD_RET_FAILURE;
			}
#elif (CONFIG_AB_SIDELOAD == 1)
			//We are going to pack version numbers and send them to dtb as avb_slot_verify doesn't update version at this time
			//in case of is_pass_rollback_index_info_needed() returns true).
			if(is_pass_rollback_index_info_needed()) {
				if(rollbackindexes_to_dtb(out_data)) {
					return CMD_RET_FAILURE;
				}
			}
#endif
			char filtered_cmd_buff[CONFIG_SYS_CBSIZE];
			const char * const *partition_array = requested_partitions;
			const char* bootcmd = UBOOT_BOOTCMD;
#if !defined(CONFIG_ANDROID_AVB_PURELINUX_ENABLE)
//for recovery mode
			if(is_recovery_mode()) {
				partition_array = requested_partitions_recovery_mode;
				bootcmd = UBOOT_REVOCERYCMD;
			}
#endif
			if(strlen(bootcmd) >= sizeof(filtered_cmd_buff)) {
				UBOOT_ERROR("%s:%d Embedded bootcmd is longer than/equal to CONFIG_SYS_CBSIZE, length is %zu.\n", __FUNCTION__, __LINE__, strlen(bootcmd));
				return CMD_RET_FAILURE;
			}
			strncpy(filtered_cmd_buff, bootcmd, strlen(bootcmd));
			filtered_cmd_buff[strlen(bootcmd)] = 0x00;//Padding a null character at the end to ensure it's a null terminated string.
			const char** verified_parts = malloc(out_data->num_loaded_partitions * sizeof(char*));
			if(!verified_parts) {
				UBOOT_ERROR("%s:%d Failed to allocate verified_parts.\n", __FUNCTION__, __LINE__);
				return CMD_RET_FAILURE;
			}
			//We get verified_parts in terms of the coherence between partition_array and out_data->loaded_partitions which indicates which parts have been verified by AVB
			//For example, if partition_array were "boot", "dtbo", "system", "vendor" and "vendor_boot"
			//and out_data->loaded_partitions were "boot", "dtbo" and "vendor_boot"
			//verified_parts would be "boot", "dtbo" and "vendor_boot"
			get_verified_parts(partition_array, out_data->loaded_partitions, out_data->num_loaded_partitions, verified_parts);
			//Then we filter out any command in filtered_cmd_buff which has the occurrance of any partition name in verified_parts
			//For example, if the input of filtered_cmd_buff was "partition read mmc 0 DTB 0x23000000 0x100000;partition read mmc 0 KL 0x20280000 0x2000000; booti 0x20280000 - 0x23000000"
			//and verified_parts were "KL", "DTB" and "dtbo",
			//the output of filtered_cmd_buff would be "booti 0x20280000 - 0x23000000"
			filterBootcmd((const char* const *)verified_parts, out_data->num_loaded_partitions, filtered_cmd_buff, sizeof(filtered_cmd_buff));
			free(verified_parts);
			UBOOT_DEBUG("%s:%d cmd_buff=%s\n", __FUNCTION__, __LINE__, filtered_cmd_buff);
			end = get_timer(0);
			UBOOT_BOOTTIME("[avb verify][start:%llu][end:%llu][total time:%llu]\n", start, end, end - start);
			run_command(filtered_cmd_buff, 0);
		}
	}
	else if (boot_state == AVB_YELLOW) {
		printf("Boot state is: Yellow !: not support and go into Red Boot state\n");
		if(check_android_boot_state(unlocked, out_data, slot_result, 0, avb_ops) != 0)	{
			res = CMD_RET_FAILURE;
		}
	}
	else { //boot_state == AVB_RED
		if(check_android_boot_state(unlocked, out_data, slot_result, 0, avb_ops) != 0)	{
			res = CMD_RET_FAILURE;
		}
	}

	return res;
}

int do_avb_set_verity(cmd_tbl_t *cmdtp, int flag, int argc, char * const argv[])
{
	bool device_is_unlocked = false;
	char *verity_flag = "";

	if (!avb_ops) {
		printf("AVB not initialized, run 'avb init' first\n");
		return CMD_RET_FAILURE;
	}

	if (argc == 2)
		verity_flag = argv[1];

	avb_ops->read_is_device_unlocked(avb_ops, &device_is_unlocked);
	if(device_is_unlocked == true) {
		if(strcmp(verity_flag,"enable") == 0) {
			avb_set_verity(avb_ops,"",true);
		}
		else if(strcmp(verity_flag,"disable") == 0) {
                avb_set_verity(avb_ops,"",false);
		}
		else {
			printf("%s: Error input\n",__FUNCTION__);
			return CMD_RET_FAILURE;
		}

	}
	else {
		printf("%s: Error, device state is locked, don't allow to on/off dm-verity.\n",__FUNCTION__);
		return CMD_RET_FAILURE;
	}

	return CMD_RET_SUCCESS;
}

static int do_avb_set_DeviceState(cmd_tbl_t *cmdtp, int flag, int argc, char * const argv[])
{
	int value;
	char *in_value = "";

	if (!avb_ops) {
		printf("AVB 2.0 is not initialized, run 'avb init' first\n");
		return CMD_RET_FAILURE;
	}

	if (argc != 2)
		return CMD_RET_USAGE;

	in_value = argv[1];
	if( in_value != NULL ) {
		value = (int)simple_strtoul(in_value, NULL, DEVICESTATE_ISLOCKED_BASE);
		if (avb_ops->write_device_status(avb_ops, value) ==  AVB_IO_RESULT_OK) {
			return CMD_RET_SUCCESS;
		}
  }

	printf("Can't write device state.\n");
	return CMD_RET_FAILURE;
}

int do_avb_is_unlocked(cmd_tbl_t *cmdtp, int flag,
		       int argc, char * const argv[])
{
	bool unlock;

	if (!avb_ops) {
		printf("AVB not initialized, run 'avb init' first\n");
		return CMD_RET_FAILURE;
	}

	if (argc != 1) {
		printf("--%s(-1)\n", __func__);
		return CMD_RET_USAGE;
	}

	if (avb_ops->read_is_device_unlocked(avb_ops, &unlock) ==
	    AVB_IO_RESULT_OK) {
		printf("Unlocked = %d\n", unlock);
		return CMD_RET_SUCCESS;
	}

	printf("Can't determine device lock state.\n");

	return CMD_RET_FAILURE;
}

static cmd_tbl_t cmd_avb[] = {
	U_BOOT_CMD_MKENT(init, 2, 0, do_avb_init, "", ""),
#ifdef CONFIG_CMD_AVB_DEBUG_ENABLE
	U_BOOT_CMD_MKENT(read_rb, 2, 0, do_avb_read_rb, "", ""),
	U_BOOT_CMD_MKENT(write_rb, 3, 0, do_avb_write_rb, "", ""),
	U_BOOT_CMD_MKENT(is_unlocked, 1, 0, do_avb_is_unlocked, "", ""),
	U_BOOT_CMD_MKENT(get_uuid, 2, 0, do_avb_get_uuid, "", ""),
	U_BOOT_CMD_MKENT(read_part, 5, 0, do_avb_read_part, "", ""),
	U_BOOT_CMD_MKENT(read_part_hex, 4, 0, do_avb_read_part_hex, "", ""),
	U_BOOT_CMD_MKENT(write_part, 5, 0, do_avb_write_part, "", ""),
#endif
	U_BOOT_CMD_MKENT(set-devicestate, 2, 0, do_avb_set_DeviceState, "", ""),
	U_BOOT_CMD_MKENT(set-verity, 2, 0, do_avb_set_verity, "", ""),
	U_BOOT_CMD_MKENT(verify, 2, 0, do_avb_verify_part, "", ""),
};

static int do_avb(cmd_tbl_t *cmdtp, int flag, int argc, char * const argv[])
{
	cmd_tbl_t *cp;

	cp = find_cmd_tbl(argv[1], cmd_avb, ARRAY_SIZE(cmd_avb));

	argc--;
	argv++;

	if (!cp || argc > cp->maxargs)
		return CMD_RET_USAGE;

	if (flag == CMD_FLAG_REPEAT)
		return CMD_RET_FAILURE;

	return cp->cmd(cmdtp, flag, argc, argv);
}

U_BOOT_CMD(
	avb, 29, 0, do_avb,
	"Provides commands for testing Android Verified Boot 2.0 functionality",
	"init <dev> - initialize avb2 for <dev>\n"
#ifdef CONFIG_CMD_AVB_DEBUG_ENABLE
	"avb read_rb <num> - read rollback index at location <num>\n"
	"avb write_rb <num> <rb> - write rollback index <rb> to <num>\n"
	"avb is_unlocked - returns unlock status of the device\n"
	"avb get_uuid <partname> - read and print uuid of partition <part>\n"
	"avb read_part <partname> <offset> <num> <addr> - read <num> bytes from\n"
	"    partition <partname> to buffer <addr>\n"
	"avb read_part_hex <partname> <offset> <num> - read <num> bytes from\n"
	"    partition <partname> and print to stdout\n"
	"avb write_part <partname> <offset> <num> <addr> - write <num> bytes to\n"
	"    <partname> by <offset> using data from <addr>\n"
#endif
	"avb verify [slot_suffix] - run verification process using hash data\n"
	"    from vbmeta structure\n"
	"    [slot_suffix] - _a, _b, etc (if vbmeta partition is slotted)\n"
	"avb set-verity <enable/disable>  - AVB2.0 enable/disable dm-verity.\n"
	"avb set-devicestate <0/1>  - AVB2.0 set-devicestate, 0=unlock/1=lock.\n"
	);
