Add --env option in websh command

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:

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,

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 {

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
+}

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,26 +251,28 @@ 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":
 			username = "root"
 		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
 }