Executable paths with spaces should be quoted in full command string

4 files changed, 21 insertions, 9 deletions

diff --git a/TODO.md b/TODO.md
index bedc988..d6da144 100644
--- a/TODO.md
+++ b/TODO.md
@@ -23,7 +23,7 @@
 
 - [X] Add DCOM module
 - [X] MMC20.Application method
-- [ ] Output
+- [X] Output
 - [ ] ShellWindows & ShellBrowserWindow
 
 ### WMI
diff --git a/cmd/root.go b/cmd/root.go
index 913a44a..6b5a416 100644
--- a/cmd/root.go
+++ b/cmd/root.go
@@ -54,6 +54,11 @@ var (
 				log = log.Level(zerolog.DebugLevel)
 			}
 
+			if proxy != "" {
+				rpcClient.Proxy = proxy
+				smbClient.Proxy = proxy
+			}
+
 			if outputPath != "" {
 				if outputMethod == "smb" {
 					if exec.Output.RemotePath == "" {
@@ -80,6 +85,7 @@ func init() {
 		rootCmd.InitDefaultHelpCmd()
 		rootCmd.PersistentFlags().BoolVar(&debug, "debug", false, "Enable debug logging")
 		rootCmd.PersistentFlags().BoolVar(&logJson, "log-json", false, "Log in JSON format")
+		rootCmd.PersistentFlags().StringVarP(&proxy, "proxy", "x", "", "Proxy URL")
 
 		dcomCmdInit()
 		rootCmd.AddCommand(dcomCmd)
diff --git a/cmd/wmi.go b/cmd/wmi.go
index 42475f0..b75df12 100644
--- a/cmd/wmi.go
+++ b/cmd/wmi.go
@@ -54,7 +54,7 @@ var (
 
 	wmiCmd = &cobra.Command{
 		Use:   "wmi",
-		Short: "Establish execution via wmi",
+		Short: "Establish execution via WMI",
 		Args:  cobra.NoArgs,
 	}
 
diff --git a/pkg/goexec/io.go b/pkg/goexec/io.go
index 6bfb76e..4fa7cb8 100644
--- a/pkg/goexec/io.go
+++ b/pkg/goexec/io.go
@@ -41,6 +41,13 @@ func (execIO *ExecutionIO) GetOutput(ctx context.Context) (err error) {
 	return nil
 }
 
+func (execIO *ExecutionIO) Clean(ctx context.Context) (err error) {
+	if execIO.Output.Provider != nil {
+		return execIO.Output.Provider.Clean(ctx)
+	}
+	return nil
+}
+
 func (execIO *ExecutionIO) CommandLine() (cmd []string) {
 	if execIO.Output.Provider != nil && execIO.Output.RemotePath != "" {
 		return []string{
@@ -51,13 +58,6 @@ func (execIO *ExecutionIO) CommandLine() (cmd []string) {
 	return execIO.Input.CommandLine()
 }
 
-func (execIO *ExecutionIO) Clean(ctx context.Context) (err error) {
-	if execIO.Output.Provider != nil {
-		return execIO.Output.Provider.Clean(ctx)
-	}
-	return nil
-}
-
 func (execIO *ExecutionIO) String() (cmd string) {
 	return strings.Join(execIO.CommandLine(), " ")
 }
@@ -76,6 +76,12 @@ func (i *ExecutionInput) CommandLine() (cmd []string) {
 	case i.Executable != "":
 		cmd[0] = i.Executable
 	}
+
+	// Ensure that executable paths are quoted
+	if strings.Contains(cmd[0], " ") {
+		cmd[0] = fmt.Sprintf(`%q`, cmd[0])
+	}
+
 	return cmd
 }