Add -u and -g flags to websh command with test cases

README.md

@@ -157,6 +157,17 @@ Access a server's websh terminal:
 $ alpacon websh [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 --username=[USER NAME] --groupname=[GROUP NAME] [COMMAND]
+```
+
+
 #### Identity and Access Management (IAM)
 Efficiently manage user and group resources:
 ```bash

api/event/event.go

@@ -35,10 +35,12 @@ func GetEventList(ac *client.AlpaconClient, pageSize int, serverName string, use
 	if err != nil {
 		return nil, err
 	}
+
 	var response EventListResponse
 	if err = json.Unmarshal(responseBody, &response); err != nil {
 		return nil, err
 	}
+
 	var eventList []EventAttributes
 	for _, event := range response.Results {
 		eventList = append(eventList, EventAttributes{
@@ -53,15 +55,18 @@ func GetEventList(ac *client.AlpaconClient, pageSize int, serverName string, use
 	}
 	return eventList, nil
 }
-func RunCommand(ac *client.AlpaconClient, serverName, command string) (string, error) {
+func RunCommand(ac *client.AlpaconClient, serverName, command string, username, groupname string) (string, error) {
 	serverID, err := server.GetServerIDByName(ac, serverName)
 	if err != nil {
 		return "", err
 	}
+
 	commandRequest := &Command{
 		Shell:       "system", // TODO Support osquery, alpamon
 		Line:        command,
 		Data:        "",
+		Username:    username,
+		Groupname:   groupname,
 		ScheduledAt: nil,
 		Server:      serverID,
 		RunAfter:    []string{},

api/event/types.go

@@ -50,6 +50,8 @@ type Command 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"`

cmd/websh/websh.go

@@ -19,25 +19,58 @@ var WebshCmd = &cobra.Command{
 	`,
 	Example: `
 	// Open a websh terminal for a server
-	alpacon websh [SERVER NAME]
+	alpacon websh [SERVER_NAME]
 
 	// Execute a command directly on a server and retrieve the output
-	alpacon websh [SERVER NAME] [COMMAND]
+	alpacon websh [SERVER_NAME] [COMMAND]
 
-	// Additional examples with flags
+	// Open a websh terminal as a root user
 	alpacon websh -r [SERVER_NAME]
-	alpacon websh [SERVER NAME] --root
+
+	// Run a command as [USER_NAME]/[GROUP_NAME]
+	alpacon websh -u [USER_NAME] -g [GROUP_NAME] [SERVER_NAME] [COMMAND]
+
+	Flags:
+	-r          					   Run the websh terminal as the root user.
+	-u / --username [USER_NAME]        Specify the username under which the command should be executed.
+	-g / --groupname [GROUP_NAME]      Specify the group name under which the command should be executed.
+
+	Note:
+	- All flags (-r, -u, -g) must be placed before the [SERVER_NAME].
+	- The -u (or --username) and -g (or --groupname) flags require an argument specifying the user or group name, respectively.
 	`,
+	DisableFlagParsing: true,
 	Run: func(cmd *cobra.Command, args []string) {
-		root, _ := cmd.Flags().GetBool("root")
+		var (
+			root                            bool
+			username, groupname, serverName string
+			commandArgs                     []string
+		)
 
-		if len(args) < 1 {
-			cmd.Usage()
-			return
+		for i := 0; i < len(args); i++ {
+			switch {
+			case args[i] == "-r" || args[i] == "--root":
+				root = true
+			case args[i] == "-h" || args[i] == "--help":
+				cmd.Help()
+				return
+			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"):
+				groupname, i = extractValue(args, i)
+			default:
+				if serverName == "" {
+					serverName = args[i]
+				} else {
+					commandArgs = append(commandArgs, args[i:]...)
+					i = len(args)
+				}
+			}
 		}
 
-		serverName := args[0]
-		commandArgs := args[1:]
+		if serverName == "" {
+			utils.CliError("Server name is required.")
+		}
 
 		alpaconClient, err := client.NewAlpaconAPIClient()
 		if err != nil {
@@ -46,7 +79,7 @@ var WebshCmd = &cobra.Command{
 
 		if len(commandArgs) > 0 {
 			command := strings.Join(commandArgs, " ")
-			result, err := event.RunCommand(alpaconClient, serverName, command)
+			result, err := event.RunCommand(alpaconClient, serverName, command, username, groupname)
 			if err != nil {
 				utils.CliError("Failed to run the '%s' command on the '%s' server: %s", command, serverName, err)
 			}
@@ -62,7 +95,12 @@ var WebshCmd = &cobra.Command{
 	},
 }
 
-func init() {
-	var root bool
-	WebshCmd.Flags().BoolVarP(&root, "root", "r", false, "Run as root user")
+func extractValue(args []string, i int) (string, int) {
+	if strings.Contains(args[i], "=") {
+		parts := strings.SplitN(args[i], "=", 2)
+		return parts[1], i
+	} else if i+1 < len(args) {
+		return args[i+1], i + 1
+	}
+	return "", i
 }

cmd/websh/websh_test.go

@@ -0,0 +1,187 @@
+package websh
+
+import (
+	"github.com/stretchr/testify/assert"
+	"strings"
+	"testing"
+)
+
+func TestCommandParsing(t *testing.T) {
+	tests := []struct {
+		testName          string
+		args              []string
+		expectRoot        bool
+		expectUsername    string
+		expectGroupname   string
+		expectServerName  string
+		expectCommandArgs []string
+	}{
+		{
+			testName:          "RootAccessToServer",
+			args:              []string{"-r", "prod-server", "df", "-h"},
+			expectRoot:        true,
+			expectUsername:    "",
+			expectGroupname:   "",
+			expectServerName:  "prod-server",
+			expectCommandArgs: []string{"df", "-h"},
+		},
+		{
+			testName:          "ExecuteUpdateAsAdminSysadmin",
+			args:              []string{"-u", "admin", "-g", "sysadmin", "update-server", "sudo", "apt-get", "update"},
+			expectRoot:        false,
+			expectUsername:    "admin",
+			expectGroupname:   "sysadmin",
+			expectServerName:  "update-server",
+			expectCommandArgs: []string{"sudo", "apt-get", "update"},
+		},
+		{
+			testName:          "DockerComposeDeploymentWithFlags",
+			args:              []string{"deploy-server", "docker-compose", "-f", "/home/admin/deploy/docker-compose.yml", "up", "-d"},
+			expectRoot:        false,
+			expectUsername:    "",
+			expectGroupname:   "",
+			expectServerName:  "deploy-server",
+			expectCommandArgs: []string{"docker-compose", "-f", "/home/admin/deploy/docker-compose.yml", "up", "-d"},
+		},
+		{
+			testName:          "VerboseListInFileServer",
+			args:              []string{"file-server", "ls", "-l", "/var/www"},
+			expectRoot:        false,
+			expectUsername:    "",
+			expectGroupname:   "",
+			expectServerName:  "file-server",
+			expectCommandArgs: []string{"ls", "-l", "/var/www"},
+		},
+		{
+			testName:          "MisplacedFlagOrderWithRoot",
+			args:              []string{"-r", "df", "-h"},
+			expectRoot:        true,
+			expectUsername:    "",
+			expectGroupname:   "",
+			expectServerName:  "df",
+			expectCommandArgs: []string(nil),
+		},
+		{
+			testName:          "UnrecognizedFlagWithEchoCommand",
+			args:              []string{"-x", "unknown-server", "echo", "Hello World"},
+			expectRoot:        false,
+			expectUsername:    "",
+			expectGroupname:   "",
+			expectServerName:  "-x",
+			expectCommandArgs: []string{"unknown-server", "echo", "Hello World"},
+		},
+		{
+			testName:          "AdminSysadminAccessToMultiFlagServer",
+			args:              []string{"--username=admin", "--groupname=sysadmin", "multi-flag-server", "uptime"},
+			expectRoot:        false,
+			expectUsername:    "admin",
+			expectGroupname:   "sysadmin",
+			expectServerName:  "multi-flag-server",
+			expectCommandArgs: []string{"uptime"},
+		},
+		{
+			testName:          "CommandLineArgsResembleFlags",
+			args:              []string{"--username", "admin", "server-name", "--fake-flag", "value"},
+			expectRoot:        false,
+			expectUsername:    "admin",
+			expectGroupname:   "",
+			expectServerName:  "server-name",
+			expectCommandArgs: []string{"--fake-flag", "value"},
+		},
+		{
+			testName:          "SysadminGroupWithMixedSyntax",
+			args:              []string{"-g=sysadmin", "server-name", "echo", "hello world"},
+			expectRoot:        false,
+			expectUsername:    "",
+			expectGroupname:   "sysadmin",
+			expectServerName:  "server-name",
+			expectCommandArgs: []string{"echo", "hello world"},
+		},
+		{
+			testName:          "HelpRequestedViaCombinedFlags",
+			args:              []string{"-rh"},
+			expectRoot:        false,
+			expectUsername:    "",
+			expectGroupname:   "",
+			expectServerName:  "-rh",
+			expectCommandArgs: nil,
+		},
+		{
+			testName:          "InvalidUsageDetected",
+			args:              []string{"-u", "user", "-x", "unknown-flag", "server-name", "cmd"},
+			expectRoot:        false,
+			expectUsername:    "user",
+			expectGroupname:   "",
+			expectServerName:  "-x",
+			expectCommandArgs: []string{"unknown-flag", "server-name", "cmd"},
+		},
+		{
+			testName:          "ValidFlagsFollowedByInvalidFlag",
+			args:              []string{"-u", "user", "-g", "group", "-x", "server-name", "cmd"},
+			expectRoot:        false,
+			expectUsername:    "user",
+			expectGroupname:   "group",
+			expectServerName:  "-x",
+			expectCommandArgs: []string{"server-name", "cmd"},
+		},
+		{
+			testName:          "FlagsIntermixedWithCommandArgs",
+			args:              []string{"server-name", "-u", "user", "cmd", "-g", "group"},
+			expectRoot:        false,
+			expectUsername:    "user",
+			expectGroupname:   "",
+			expectServerName:  "server-name",
+			expectCommandArgs: []string{"cmd", "-g", "group"},
+		},
+		{
+			testName:          "FlagsAndCommandArgsIntertwined",
+			args:              []string{"server-name", "-u", "user", "cmd", "-g", "group"},
+			expectRoot:        false,
+			expectUsername:    "user",
+			expectGroupname:   "",
+			expectServerName:  "server-name",
+			expectCommandArgs: []string{"cmd", "-g", "group"},
+		},
+	}
+
+	for _, tc := range tests {
+		t.Run(tc.testName, func(t *testing.T) {
+			root, username, groupname, serverName, commandArgs := executeTestCommand(tc.args)
+
+			assert.Equal(t, tc.expectRoot, root, "Mismatch in root flag")
+			assert.Equal(t, tc.expectUsername, username, "Mismatch in username")
+			assert.Equal(t, tc.expectGroupname, groupname, "Mismatch in groupname")
+			assert.Equal(t, tc.expectServerName, serverName, "Mismatch in server name")
+			assert.Equal(t, tc.expectCommandArgs, commandArgs, "Mismatch in command arguments")
+
+		})
+	}
+}
+
+func executeTestCommand(args []string) (bool, string, string, string, []string) {
+	var root bool
+	var username, groupname, serverName string
+	var commandArgs []string
+
+	for i := 0; i < len(args); i++ {
+		switch {
+		case args[i] == "-r" || args[i] == "--root":
+			root = true
+		case args[i] == "-h" || args[i] == "--help":
+			return root, username, groupname, serverName, commandArgs
+		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"):
+			groupname, i = extractValue(args, i)
+		default:
+			if serverName == "" {
+				serverName = args[i]
+			} else {
+				commandArgs = append(commandArgs, args[i:]...)
+				i = len(args)
+			}
+		}
+	}
+
+	return root, username, groupname, serverName, commandArgs
+}

utils/utils.go

@@ -214,3 +214,11 @@ func BoolPointerToString(value *bool) string {
 	}
 	return "false"
 }
+
+func StringToStringPointer(value string) *string {
+	if value == "" {
+		return nil
+	} else {
+		return &value
+	}
+}