1 files changed, 26 insertions, 20 deletions

diff --git a/cmd-set-environment.c b/cmd-set-environment.c
index fdd703c5..3075fb07 100644
--- a/cmd-set-environment.c
+++ b/cmd-set-environment.c
@@ -31,57 +31,63 @@ int	cmd_set_environment_exec(struct cmd *, struct cmd_ctx *);
 
 const struct cmd_entry cmd_set_environment_entry = {
 	"set-environment", "setenv",
+	"grt:u", 1, 2,
 	"[-gru] " CMD_TARGET_SESSION_USAGE " name [value]",
-	CMD_ARG12, "gru",
+	0,
 	NULL,
-	cmd_target_parse,
-	cmd_set_environment_exec,
-	cmd_target_free,
-	cmd_target_print
+	NULL,
+	cmd_set_environment_exec
 };
 
 int
 cmd_set_environment_exec(struct cmd *self, struct cmd_ctx *ctx)
 {
-	struct cmd_target_data	*data = self->data;
-	struct session		*s;
-	struct environ		*env;
+	struct args	*args = self->args;
+	struct session	*s;
+	struct environ	*env;
+	const char	*name, *value;
 
-	if (*data->arg == '\0') {
+	name = args->argv[0];
+	if (*name == '\0') {
 		ctx->error(ctx, "empty variable name");
 		return (-1);
 	}
-	if (strchr(data->arg, '=') != NULL) {
+	if (strchr(name, '=') != NULL) {
 		ctx->error(ctx, "variable name contains =");
 		return (-1);
 	}
 
-	if (cmd_check_flag(data->chflags, 'g'))
+	if (args->argc < 1)
+		value = NULL;
+	else
+		value = args->argv[1];
+
+	if (args_has(self->args, 'g'))
 		env = &global_environ;
 	else {
-		if ((s = cmd_find_session(ctx, data->target)) == NULL)
+		if ((s = cmd_find_session(ctx, args_get(args, 't'))) == NULL)
 			return (-1);
 		env = &s->environ;
 	}
 
-	if (cmd_check_flag(data->chflags, 'u')) {
-		if (data->arg2 != NULL) {
+	if (args_has(self->args, 'u')) {
+		if (value != NULL) {
 			ctx->error(ctx, "can't specify a value with -u");
 			return (-1);
 		}
-		environ_unset(env, data->arg);
-	} else if (cmd_check_flag(data->chflags, 'r')) {
-		if (data->arg2 != NULL) {
+		environ_unset(env, name);
+	} else if (args_has(self->args, 'r')) {
+		if (value != NULL) {
 			ctx->error(ctx, "can't specify a value with -r");
 			return (-1);
 		}
-		environ_set(env, data->arg, NULL);
+		environ_set(env, name, NULL);
 	} else {
-		if (data->arg2 == NULL) {
+		if (value == NULL) {
 			ctx->error(ctx, "no value specified");
 			return (-1);
 		}
-		environ_set(env, data->arg, data->arg2);
+		environ_set(env, name, value);
 	}
 
 	return (0);