From 536d4cb9e6081249ed8b53f2f54854943505ea68 Mon Sep 17 00:00:00 2001 From: royroyee Date: Thu, 2 May 2024 14:47:00 +0900 Subject: [PATCH] Add --env option in websh command --- README.md | 9 +++++-- api/event/event.go | 3 ++- api/event/types.go | 17 +++++++------- cmd/websh/websh.go | 24 ++++++++++++++++--- cmd/websh/websh_test.go | 52 +++++++++++++++++++++++++++++++++++++---- 5 files changed, 86 insertions(+), 19 deletions(-) diff --git a/README.md b/README.md index 98cecab..38c07c2 100644 --- a/README.md +++ b/README.md @@ -176,15 +176,20 @@ $ alpacon websh -r [SERVER NAME] $ alpacon websh -u [USER NAME] -g [GROUP NAME] [SERVER NAME] ``` + #### Execute a command Execute a command directly on a server and retrieve the output: ```bash $ alpacon websh [SERVER NAME] [COMMAND] -$ alpacon websh -u [USER NAME] -g [GROUP NAME] [COMMAND] +$ alpacon websh -u [USER NAME] -g [GROUP NAME] [SERVER NAME] [COMMAND] + +$ alpacon websh --username=[USER NAME] --groupname=[GROUP NAME] [SERVER NAME] [COMMAND] -$ alpacon websh --username=[USER NAME] --groupname=[GROUP NAME] [COMMAND] +$ alpacon websh --env="KEY1=VALUE1" --env="KEY2=VALUE2" [SERVER NAME] [COMMAND] ``` +- Note: All flags must be placed before the `[SERVER NAME]`. + #### Share your terminal You can share the current terminal to others via a temporary link: diff --git a/api/event/event.go b/api/event/event.go index 80d86b7..fc36f51 100644 --- a/api/event/event.go +++ b/api/event/event.go @@ -62,7 +62,7 @@ func GetEventList(ac *client.AlpaconClient, pageSize int, serverName string, use return eventList, nil } -func RunCommand(ac *client.AlpaconClient, serverName, command string, username, groupname string) (string, error) { +func RunCommand(ac *client.AlpaconClient, serverName, command string, username, groupname string, env map[string]string) (string, error) { serverID, err := server.GetServerIDByName(ac, serverName) if err != nil { return "", err @@ -71,6 +71,7 @@ func RunCommand(ac *client.AlpaconClient, serverName, command string, username, commandRequest := &CommandRequest{ Shell: "system", Line: command, + Env: env, Username: username, Groupname: groupname, Server: serverID, diff --git a/api/event/types.go b/api/event/types.go index b351828..cfff8cd 100644 --- a/api/event/types.go +++ b/api/event/types.go @@ -29,14 +29,15 @@ type EventDetails struct { } type CommandRequest struct { - Shell string `json:"shell"` - Line string `json:"line"` - Data string `json:"data"` - Username string `json:"username"` - Groupname string `json:"groupname"` - ScheduledAt *time.Time `json:"scheduled_at"` - Server string `json:"server"` - RunAfter []string `json:"run_after"` + Shell string `json:"shell"` + Line string `json:"line"` + Env map[string]string `json:"env"` + Data string `json:"data"` + Username string `json:"username"` + Groupname string `json:"groupname"` + ScheduledAt *time.Time `json:"scheduled_at"` + Server string `json:"server"` + RunAfter []string `json:"run_after"` } type CommandResponse struct { diff --git a/cmd/websh/websh.go b/cmd/websh/websh.go index 81eb5b0..85f3c05 100644 --- a/cmd/websh/websh.go +++ b/cmd/websh/websh.go @@ -63,6 +63,8 @@ var WebshCmd = &cobra.Command{ share, readOnly bool ) + env := make(map[string]string) + for i := 0; i < len(args); i++ { switch { case args[i] == "-r" || args[i] == "--root": @@ -80,6 +82,8 @@ var WebshCmd = &cobra.Command{ url, i = extractValue(args, i) case strings.HasPrefix(args[i], "-p") || strings.HasPrefix(args[i], "--password"): password, i = extractValue(args, i) + case strings.HasPrefix(args[i], "--env"): + i = extractEnvValue(args, i, env) case strings.HasPrefix(args[i], "--read-only"): var value string value, i = extractValue(args, i) @@ -120,7 +124,7 @@ var WebshCmd = &cobra.Command{ websh.OpenNewTerminal(alpaconClient, session) } else if len(commandArgs) > 0 { command := strings.Join(commandArgs, " ") - result, err := event.RunCommand(alpaconClient, serverName, command, username, groupname) + result, err := event.RunCommand(alpaconClient, serverName, command, username, groupname, env) if err != nil { utils.CliError("Failed to run the command on the '%s' server: %s.", serverName, err) } @@ -136,12 +140,26 @@ var WebshCmd = &cobra.Command{ } func extractValue(args []string, i int) (string, int) { - if strings.Contains(args[i], "=") { + if strings.Contains(args[i], "=") { // --username=admins parts := strings.SplitN(args[i], "=", 2) return parts[1], i } - if i+1 < len(args) { + if i+1 < len(args) { // --username admin return args[i+1], i + 1 } return "", i } + +func extractEnvValue(args []string, i int, env map[string]string) int { + envString := strings.TrimPrefix(args[i], "--env=") + envString = strings.Trim(envString, "\"") + + parts := strings.SplitN(envString, "=", 2) + if len(parts) == 2 { + env[parts[0]] = parts[1] + } else { + utils.CliError("Invalid format for --env. Expected '--env=KEY=VALUE'.") + } + + return i +} diff --git a/cmd/websh/websh_test.go b/cmd/websh/websh_test.go index 94eefa6..4514e55 100644 --- a/cmd/websh/websh_test.go +++ b/cmd/websh/websh_test.go @@ -14,6 +14,7 @@ func TestCommandParsing(t *testing.T) { expectUsername string expectGroupname string expectServerName string + expectEnv map[string]string expectCommandArgs []string expectShare bool expectJoin bool @@ -26,6 +27,7 @@ func TestCommandParsing(t *testing.T) { args: []string{"-r", "prod-server", "df", "-h"}, expectUsername: "root", expectGroupname: "", + expectEnv: map[string]string{}, expectServerName: "prod-server", expectCommandArgs: []string{"df", "-h"}, }, @@ -34,6 +36,7 @@ func TestCommandParsing(t *testing.T) { args: []string{"-u", "admin", "-g", "sysadmin", "update-server", "sudo", "apt-get", "update"}, expectUsername: "admin", expectGroupname: "sysadmin", + expectEnv: map[string]string{}, expectServerName: "update-server", expectCommandArgs: []string{"sudo", "apt-get", "update"}, }, @@ -42,6 +45,7 @@ func TestCommandParsing(t *testing.T) { args: []string{"deploy-server", "docker-compose", "-f", "/home/admin/deploy/docker-compose.yml", "up", "-d"}, expectUsername: "", expectGroupname: "", + expectEnv: map[string]string{}, expectServerName: "deploy-server", expectCommandArgs: []string{"docker-compose", "-f", "/home/admin/deploy/docker-compose.yml", "up", "-d"}, }, @@ -50,6 +54,7 @@ func TestCommandParsing(t *testing.T) { args: []string{"file-server", "ls", "-l", "/var/www"}, expectUsername: "", expectGroupname: "", + expectEnv: map[string]string{}, expectServerName: "file-server", expectCommandArgs: []string{"ls", "-l", "/var/www"}, }, @@ -58,6 +63,7 @@ func TestCommandParsing(t *testing.T) { args: []string{"-r", "df", "-h"}, expectUsername: "root", expectGroupname: "", + expectEnv: map[string]string{}, expectServerName: "df", expectCommandArgs: []string(nil), }, @@ -66,6 +72,7 @@ func TestCommandParsing(t *testing.T) { args: []string{"-x", "unknown-server", "echo", "Hello World"}, expectUsername: "", expectGroupname: "", + expectEnv: map[string]string{}, expectServerName: "-x", expectCommandArgs: []string{"unknown-server", "echo", "Hello World"}, }, @@ -74,6 +81,7 @@ func TestCommandParsing(t *testing.T) { args: []string{"--username=admin", "--groupname=sysadmin", "multi-flag-server", "uptime"}, expectUsername: "admin", expectGroupname: "sysadmin", + expectEnv: map[string]string{}, expectServerName: "multi-flag-server", expectCommandArgs: []string{"uptime"}, }, @@ -82,6 +90,7 @@ func TestCommandParsing(t *testing.T) { args: []string{"--username", "admin", "server-name", "--fake-flag", "value"}, expectUsername: "admin", expectGroupname: "", + expectEnv: map[string]string{}, expectServerName: "server-name", expectCommandArgs: []string{"--fake-flag", "value"}, }, @@ -90,6 +99,7 @@ func TestCommandParsing(t *testing.T) { args: []string{"-g=sysadmin", "server-name", "echo", "hello world"}, expectUsername: "", expectGroupname: "sysadmin", + expectEnv: map[string]string{}, expectServerName: "server-name", expectCommandArgs: []string{"echo", "hello world"}, }, @@ -98,6 +108,7 @@ func TestCommandParsing(t *testing.T) { args: []string{"-rh"}, expectUsername: "", expectGroupname: "", + expectEnv: map[string]string{}, expectServerName: "-rh", expectCommandArgs: nil, }, @@ -106,6 +117,7 @@ func TestCommandParsing(t *testing.T) { args: []string{"-u", "user", "-x", "unknown-flag", "server-name", "cmd"}, expectUsername: "user", expectGroupname: "", + expectEnv: map[string]string{}, expectServerName: "-x", expectCommandArgs: []string{"unknown-flag", "server-name", "cmd"}, }, @@ -114,6 +126,7 @@ func TestCommandParsing(t *testing.T) { args: []string{"-u", "user", "-g", "group", "-x", "server-name", "cmd"}, expectUsername: "user", expectGroupname: "group", + expectEnv: map[string]string{}, expectServerName: "-x", expectCommandArgs: []string{"server-name", "cmd"}, }, @@ -122,6 +135,7 @@ func TestCommandParsing(t *testing.T) { args: []string{"server-name", "-u", "user", "cmd", "-g", "group"}, expectUsername: "user", expectGroupname: "", + expectEnv: map[string]string{}, expectServerName: "server-name", expectCommandArgs: []string{"cmd", "-g", "group"}, }, @@ -130,6 +144,7 @@ func TestCommandParsing(t *testing.T) { args: []string{"server-name", "-u", "user", "cmd", "-g", "group"}, expectUsername: "user", expectGroupname: "", + expectEnv: map[string]string{}, expectServerName: "server-name", expectCommandArgs: []string{"cmd", "-g", "group"}, }, @@ -138,6 +153,7 @@ func TestCommandParsing(t *testing.T) { args: []string{"test-server", "--share"}, expectUsername: "", expectGroupname: "", + expectEnv: map[string]string{}, expectServerName: "test-server", expectCommandArgs: nil, expectShare: true, @@ -151,6 +167,7 @@ func TestCommandParsing(t *testing.T) { args: []string{"join", "--url", "http://localhost:3000/websh/join?session=abcd", "--password", "1234"}, expectUsername: "", expectGroupname: "", + expectEnv: map[string]string{}, expectServerName: "join", expectCommandArgs: nil, expectShare: false, @@ -164,6 +181,7 @@ func TestCommandParsing(t *testing.T) { args: []string{"test-server", "--share", "--read-only"}, expectUsername: "", expectGroupname: "", + expectEnv: map[string]string{}, expectServerName: "test-server", expectCommandArgs: nil, expectShare: true, @@ -177,6 +195,7 @@ func TestCommandParsing(t *testing.T) { args: []string{"test-server", "--share", "--read-only=True"}, expectUsername: "", expectGroupname: "", + expectEnv: map[string]string{}, expectServerName: "test-server", expectCommandArgs: nil, expectShare: true, @@ -190,6 +209,7 @@ func TestCommandParsing(t *testing.T) { args: []string{"--share", "join", "--url", "http://localhost:3000/websh/join?session=abcd"}, expectUsername: "", expectGroupname: "", + expectEnv: map[string]string{}, expectServerName: "join", expectCommandArgs: nil, expectShare: true, @@ -198,11 +218,29 @@ func TestCommandParsing(t *testing.T) { expectUrl: "http://localhost:3000/websh/join?session=abcd", expectPassword: "", }, + { + testName: "SingleEnvVariable", + args: []string{"--env=KEY1=value1", "server-name", "cmd"}, + expectUsername: "", + expectGroupname: "", + expectEnv: map[string]string{"KEY1": "value1"}, + expectServerName: "server-name", + expectCommandArgs: []string{"cmd"}, + }, + { + testName: "MultipleEnvVariables", + args: []string{"--env=KEY1=value1", "--env=KEY2=value2", "server-name", "cmd"}, + expectUsername: "", + expectGroupname: "", + expectEnv: map[string]string{"KEY1": "value1", "KEY2": "value2"}, + expectServerName: "server-name", + expectCommandArgs: []string{"cmd"}, + }, } for _, tc := range tests { t.Run(tc.testName, func(t *testing.T) { - username, groupname, serverName, commandArgs, share, join, readOnly, url, password := executeTestCommand(tc.args) + username, groupname, serverName, commandArgs, share, join, readOnly, url, password, env := executeTestCommand(tc.args) assert.Equal(t, tc.expectUsername, username, "Mismatch in username") assert.Equal(t, tc.expectGroupname, groupname, "Mismatch in groupname") @@ -213,17 +251,20 @@ func TestCommandParsing(t *testing.T) { assert.Equal(t, tc.expectReadOnly, readOnly, "Mismatch in read-only flag") assert.Equal(t, tc.expectUrl, url, "Mismatch in URL for joining") assert.Equal(t, tc.expectPassword, password, "Mismatch in password for joining") + assert.Equal(t, tc.expectEnv, env, "Mismatch in env") }) } } -func executeTestCommand(args []string) (string, string, string, []string, bool, bool, bool, string, string) { +func executeTestCommand(args []string) (string, string, string, []string, bool, bool, bool, string, string, map[string]string) { var ( share, join, readOnly bool username, groupname, serverName, url, password string commandArgs []string ) + env := make(map[string]string) + for i := 0; i < len(args); i++ { switch { case args[i] == "-r" || args[i] == "--root": @@ -231,8 +272,7 @@ func executeTestCommand(args []string) (string, string, string, []string, bool, case args[i] == "-s" || args[i] == "--share": share = true case args[i] == "-h" || args[i] == "--help": - return username, groupname, serverName, commandArgs, share, join, readOnly, url, password - + return username, groupname, serverName, commandArgs, share, join, readOnly, url, password, env case strings.HasPrefix(args[i], "-u") || strings.HasPrefix(args[i], "--username"): username, i = extractValue(args, i) case strings.HasPrefix(args[i], "-g") || strings.HasPrefix(args[i], "--groupname"): @@ -241,6 +281,8 @@ func executeTestCommand(args []string) (string, string, string, []string, bool, url, i = extractValue(args, i) case strings.HasPrefix(args[i], "-p") || strings.HasPrefix(args[i], "--password"): password, i = extractValue(args, i) + case strings.HasPrefix(args[i], "--env"): + i = extractEnvValue(args, i, env) case strings.HasPrefix(args[i], "--read-only"): var value string value, i = extractValue(args, i) @@ -265,5 +307,5 @@ func executeTestCommand(args []string) (string, string, string, []string, bool, join = true } - return username, groupname, serverName, commandArgs, share, join, readOnly, url, password + return username, groupname, serverName, commandArgs, share, join, readOnly, url, password, env }