From 4eae6a0686e6dc2d65bcdbbdb400fdf008347f7f Mon Sep 17 00:00:00 2001
From: Riccardo Schirone <sirmy15@gmail.com>
Date: Wed, 20 Jan 2021 10:31:55 +0100
Subject: [PATCH] Add rz_cmd_desc_get_arg API

---
 librz/core/cmd_api.c   | 29 +++++++++++++++++++++++++++++
 librz/include/rz_cmd.h |  1 +
 test/unit/test_cmd.c   | 42 ++++++++++++++++++++++++++++++++++++++++++
 3 files changed, 72 insertions(+)

diff --git a/librz/core/cmd_api.c b/librz/core/cmd_api.c
index 2cde94fd39..0fada796ee 100644
--- a/librz/core/cmd_api.c
+++ b/librz/core/cmd_api.c
@@ -1851,6 +1851,35 @@ RZ_API bool rz_cmd_desc_remove(RzCmd *cmd, RzCmdDesc *cd) {
 	return true;
 }
 
+/**
+ * \brief Get a reference to the i-th argument of a command descriptor.
+ *
+ * Get a reference to the i-th argument of a command. This function is useful
+ * to know which RzCmdDescArg an argument actually belongs to. In particular,
+ * it deals with arguments with special flags like \p RZ_CMD_ARG_FLAG_LAST or
+ * \p RZ_CMD_ARG_FLAG_ARRAY, where even if there is just one RzCmdDescArg,
+ * everything is considered as part of the same RzCmdDescArg.
+ */
+RZ_API const RzCmdDescArg *rz_cmd_desc_get_arg(RzCmd *cmd, const RzCmdDesc *cd, size_t i) {
+	const RzCmdDescArg *arg = cd->help->args;
+	size_t j = 0;
+	while (arg && arg->name) {
+		if (arg->type == RZ_CMD_ARG_TYPE_FAKE) {
+			arg++;
+			continue;
+		}
+		if (i == j) {
+			return arg;
+		}
+		if ((arg->flags & RZ_CMD_ARG_FLAG_LAST) || (arg->flags & RZ_CMD_ARG_FLAG_ARRAY)) {
+			return arg;
+		}
+		arg++;
+		j++;
+	}
+	return NULL;
+}
+
 static void cmd_foreach_cmdname(RzCmd *cmd, RzCmdDesc *cd, RzCmdForeachNameCb cb, void *user) {
 	if (!cd) {
 		return;
diff --git a/librz/include/rz_cmd.h b/librz/include/rz_cmd.h
index 1bf10006ff..00e32aae39 100644
--- a/librz/include/rz_cmd.h
+++ b/librz/include/rz_cmd.h
@@ -488,6 +488,7 @@ RZ_API RzCmdDesc *rz_cmd_desc_get_exec(RzCmdDesc *cd);
 RZ_API bool rz_cmd_desc_has_handler(RzCmdDesc *cd);
 RZ_API bool rz_cmd_desc_remove(RzCmd *cmd, RzCmdDesc *cd);
 RZ_API void rz_cmd_foreach_cmdname(RzCmd *cmd, RzCmdForeachNameCb cb, void *user);
+RZ_API const RzCmdDescArg *rz_cmd_desc_get_arg(RzCmd *cmd, const RzCmdDesc *cd, size_t i);
 
 #define rz_cmd_desc_children_foreach(root, it_cd) rz_pvector_foreach (&root->children, it_cd)
 
diff --git a/test/unit/test_cmd.c b/test/unit/test_cmd.c
index efdfda847a..6fe38adcdf 100644
--- a/test/unit/test_cmd.c
+++ b/test/unit/test_cmd.c
@@ -766,6 +766,47 @@ bool test_arg_flags(void) {
 	mu_end;
 }
 
+bool test_get_arg(void) {
+	RzCmdDescArg z_args[] = {
+		{ .name = "a1", .type = RZ_CMD_ARG_TYPE_STRING },
+		{ .name = "a2", .type = RZ_CMD_ARG_TYPE_CMD, .flags = RZ_CMD_ARG_FLAG_LAST },
+		{ 0 }
+	};
+	RzCmdDescHelp z_help = { 0 };
+	z_help.summary = "z summary";
+	z_help.args = z_args;
+	RzCmdDescArg x_args[] = {
+		{ .name = "b1", .type = RZ_CMD_ARG_TYPE_STRING },
+		{ .name = "b2", .type = RZ_CMD_ARG_TYPE_STRING },
+		{ .name = "b3", .type = RZ_CMD_ARG_TYPE_STRING },
+		{ 0 }
+	};
+	RzCmdDescHelp x_help = { 0 };
+	x_help.summary = "x summary";
+	x_help.args = x_args;
+	RzCmd *cmd = rz_cmd_new (false);
+	RzCmdDesc *root = rz_cmd_get_root (cmd);
+	RzCmdDesc *z_cd = rz_cmd_desc_argv_new (cmd, root, "z", z_last_handler, &z_help);
+	RzCmdDesc *x_cd = rz_cmd_desc_argv_new (cmd, root, "x", x_array_handler, &x_help);
+
+	const RzCmdDescArg *a1 = rz_cmd_desc_get_arg (cmd, z_cd, 0);
+	mu_assert_streq (a1->name, "a1", "0th arg of z is a1");
+	const RzCmdDescArg *a2 = rz_cmd_desc_get_arg (cmd, z_cd, 1);
+	mu_assert_streq (a2->name, "a2", "1th arg of z is a2");
+	const RzCmdDescArg *an = rz_cmd_desc_get_arg (cmd, z_cd, 10);
+	mu_assert_streq (an->name, "a2", "10th arg of z is a2");
+
+	const RzCmdDescArg *b1 = rz_cmd_desc_get_arg (cmd, x_cd, 0);
+	mu_assert_streq (b1->name, "b1", "0th arg of x is b1");
+	const RzCmdDescArg *b2 = rz_cmd_desc_get_arg (cmd, x_cd, 1);
+	mu_assert_streq (b2->name, "b2", "1th arg of x is b2");
+	const RzCmdDescArg *bn = rz_cmd_desc_get_arg (cmd, x_cd, 10);
+	mu_assert_null (bn, "10th arg of x does not exist");
+
+	rz_cmd_free (cmd);
+	mu_end;
+}
+
 int all_tests() {
 	mu_run_test (test_parsed_args_noargs);
 	mu_run_test (test_parsed_args_onearg);
@@ -792,6 +833,7 @@ int all_tests() {
 	mu_run_test (test_double_quoted_arg_escaping);
 	mu_run_test (test_single_quoted_arg_escaping);
 	mu_run_test (test_arg_flags);
+	mu_run_test (test_get_arg);
 	return tests_passed != tests_run;
 }