Remodel WMI module; Update TODO; Update deps

3 files changed, 133 insertions, 114 deletions

diff --git a/cmd/root.go b/cmd/root.go
index b44c50e..cbc38c0 100644
--- a/cmd/root.go
+++ b/cmd/root.go
@@ -23,7 +23,8 @@ var (
   proxyUrl *url.URL
 
   // Root flags
-  debug bool
+  unsafe bool // not implemented
+  debug  bool
 
   // Generic flags
   command          string
@@ -103,6 +104,7 @@ func init() {
   rootCmd.InitDefaultHelpCmd()
   rootCmd.PersistentFlags().BoolVar(&debug, "debug", false, "Enable debug logging")
   rootCmd.PersistentFlags().StringVarP(&proxyStr, "proxy", "x", "", "Proxy URL")
+  rootCmd.PersistentFlags().BoolVar(&unsafe, "unsafe", false, "[NOT IMPLEMENTED] Don't ask for permission to run unsafe actions")
 
   authOpts = &adauth.Options{Debug: log.Debug().Msgf}
   authOpts.RegisterFlags(rootCmd.PersistentFlags())
diff --git a/cmd/scmr.go b/cmd/scmr.go
index ee2e9dc..11e1379 100644
--- a/cmd/scmr.go
+++ b/cmd/scmr.go
@@ -28,6 +28,7 @@ func scmrCreateCmdInit() {
   scmrCreateCmd.Flags().BoolVar(&scmrNoStart, "no-start", false, "Don't start service")
   scmrCreateCmd.Flags().StringVarP(&executablePath, "executable-path", "f", "", "Full path to a remote Windows executable file")
   scmrCreateCmd.Flags().StringVarP(&executableArgs, "args", "a", "", "Arguments to pass to the executable")
+  scmrCreateCmd.Flags().BoolVarP(&scmrOutput, "output", "O", false, "Fetch program output")
   if err := scmrCreateCmd.MarkFlagRequired("executable-path"); err != nil {
     panic(err)
   }
@@ -58,6 +59,7 @@ var (
   scmrDisplayName  string
   scmrNoDelete     bool
   scmrNoStart      bool
+  scmrOutput       bool
 
   creds  *adauth.Credential
   target *adauth.Target
@@ -104,6 +106,7 @@ References:
       execCfg := &exec.ExecutionConfig{
         ExecutablePath:  executablePath,
         ExecutableArgs:  executableArgs,
+        ReturnOutput:    scmrOutput,
         ExecutionMethod: scmrexec.MethodCreate,
 
         ExecutionMethodConfig: scmrexec.MethodCreateConfig{
diff --git a/cmd/wmi.go b/cmd/wmi.go
index df047a2..74996b7 100644
--- a/cmd/wmi.go
+++ b/cmd/wmi.go
@@ -1,108 +1,111 @@
 package cmd
 
 import (
-	"encoding/json"
-	"fmt"
-	"github.com/FalconOpsLLC/goexec/internal/exec"
-	wmiexec "github.com/FalconOpsLLC/goexec/internal/exec/wmi"
-	"github.com/spf13/cobra"
-	"regexp"
-	"strings"
+  "encoding/json"
+  "fmt"
+  "github.com/FalconOpsLLC/goexec/internal/exec"
+  wmiexec "github.com/FalconOpsLLC/goexec/internal/exec/wmi"
+  "github.com/spf13/cobra"
 )
 
 func wmiCmdInit() {
-	wmiCustomCmdInit()
-	wmiCmd.AddCommand(wmiCustomCmd)
-	wmiProcessCmdInit()
-	wmiCmd.AddCommand(wmiProcessCmd)
+  registerRpcFlags(wmiCmd)
+  wmiCallCmdInit()
+  wmiCmd.AddCommand(wmiCallCmd)
+  wmiProcessCmdInit()
+  wmiCmd.AddCommand(wmiProcessCmd)
 }
 
-func wmiCustomCmdInit() {
-	wmiCustomCmd.Flags().StringVarP(&wmiArgMethod, "method", "m", "", `WMI Method to use in the format CLASS.METHOD (i.e. "Win32_Process.Create")`)
-	wmiCustomCmd.Flags().StringVarP(&wmiArgMethodArgs, "args", "A", "{}", `WMI Method argument(s) in JSON dictionary format (i.e. {"CommandLine":"calc.exe"})`)
-	if err := wmiCustomCmd.MarkFlagRequired("method"); err != nil {
-		panic(err)
-	}
+func wmiCallCmdInit() {
+  wmiCallCmd.Flags().StringVarP(&wmi.Method, "method", "m", "", `WMI Method to call (i.e. "Create")`)
+  wmiCallCmd.Flags().StringVarP(&wmi.Class, "class", "C", "", `WMI class to instantiate (i.e. "Win32_Process")`)
+  wmiCallCmd.Flags().StringVarP(&wmi.Args, "args", "A", "{}", `WMI Method argument(s) in JSON dictionary format (i.e. {"CommandLine":"calc.exe"})`)
+  if err := wmiCallCmd.MarkFlagRequired("method"); err != nil {
+    panic(err)
+  }
 }
 
 func wmiProcessCmdInit() {
-	wmiProcessCmd.Flags().StringVarP(&command, "command", "c", "", "Process command line")
-	wmiProcessCmd.Flags().StringVarP(&workingDirectory, "directory", "d", `C:\`, "Working directory")
-	if err := wmiProcessCmd.MarkFlagRequired("command"); err != nil {
-		panic(err)
-	}
+  wmiProcessCmd.Flags().StringVarP(&command, "command", "c", "", "Process command line")
+  wmiProcessCmd.Flags().StringVarP(&workingDirectory, "directory", "d", `C:\`, "Working directory")
+  if err := wmiProcessCmd.MarkFlagRequired("command"); err != nil {
+    panic(err)
+  }
 }
 
 var (
-	// for custom method
-	wmiArgMethod     string
-	wmiArgMethodArgs string
-
-	wmiClass         string
-	wmiMethod        string
-	wmiMethodArgsMap map[string]any
-	methodRegex      = regexp.MustCompile(`^\w+\.\w+$`)
-
-	wmiCmd = &cobra.Command{
-		Use:   "wmi",
-		Short: "Establish execution via WMI",
-		Args:  cobra.NoArgs,
-	}
-	wmiCustomCmd = &cobra.Command{
-		Use:   "custom",
-		Short: "Execute specified WMI method",
-		Long: `Description:
-  The custom method creates an instance of the specified WMI class (-c),
+  wmi struct {
+    Namespace string // TODO
+    Class     string
+    Method    string
+    Args      string
+  }
+  wmiMethodArgsMap map[string]any
+
+  wmiCmd = &cobra.Command{
+    Use:   "wmi",
+    Short: "Establish execution via WMI",
+    Args:  cobra.NoArgs,
+  }
+  wmiCallCmd = &cobra.Command{
+    Use:   "call",
+    Short: "Execute specified WMI method",
+    Long: `Description:
+  The call method creates an instance of the specified WMI class (-c),
   then calls the provided method (-m) with the provided arguments (-A).
 
 References:
   https://learn.microsoft.com/en-us/windows/win32/wmisdk/wmi-classes
   `,
-		Args: func(cmd *cobra.Command, args []string) (err error) {
-			if err = needsTarget("cifs")(cmd, args); err == nil {
-				if wmiArgMethod != "" && !methodRegex.MatchString(wmiArgMethod) {
-					return fmt.Errorf("invalid CLASS.METHOD syntax: %s", wmiArgMethod)
-				}
-				if err = json.Unmarshal([]byte(wmiArgMethodArgs), &wmiMethodArgsMap); err != nil {
-					err = fmt.Errorf("failed to parse JSON arguments: %w", err)
-				}
-			}
-			return
-		},
-		Run: func(cmd *cobra.Command, args []string) {
-			module := wmiexec.Module{}
-
-			connCfg := &exec.ConnectionConfig{}
-			cleanCfg := &exec.CleanupConfig{}
-
-			parts := strings.SplitN(wmiArgMethod, ".", 2)
-			wmiClass = parts[0]
-			wmiMethod = parts[1]
-
-			execCfg := &exec.ExecutionConfig{
-				ExecutableName:  executable,
-				ExecutableArgs:  executableArgs,
-				ExecutionMethod: wmiexec.MethodCustom,
-				ExecutionMethodConfig: wmiexec.MethodCustomConfig{
-					Class:     wmiClass,
-					Method:    wmiMethod,
-					Arguments: wmiMethodArgsMap,
-				},
-			}
-			if err := module.Connect(log.WithContext(ctx), creds, target, connCfg); err != nil {
-				log.Fatal().Err(err).Msg("Connection failed")
-			} else if err := module.Exec(log.WithContext(ctx), execCfg); err != nil {
-				log.Fatal().Err(err).Msg("Execution failed")
-			} else if err := module.Cleanup(log.WithContext(ctx), cleanCfg); err != nil {
-				log.Error().Err(err).Msg("Cleanup failed")
-			}
-		},
-	}
-
-	wmiProcessCmd = &cobra.Command{
-		Use:   "proc",
-		Short: "Start a Windows process",
-		Long: `Description:
+    Args: needs(needsTarget("cifs"), needsRpcTarget("cifs"), func(cmd *cobra.Command, args []string) (err error) {
+      if err = json.Unmarshal([]byte(wmi.Args), &wmiMethodArgsMap); err != nil {
+        err = fmt.Errorf("parse JSON arguments: %w", err)
+      }
+      return
+    }),
+    Run: func(cmd *cobra.Command, args []string) {
+
+      executor := wmiexec.Module{}
+      cleanCfg := &exec.CleanupConfig{} // TODO
+      connCfg := &exec.ConnectionConfig{
+        ConnectionMethod:       exec.ConnectionMethodDCE,
+        ConnectionMethodConfig: dceConfig,
+      }
+
+      execCfg := &exec.ExecutionConfig{
+        ExecutableName:  executable,
+        ExecutableArgs:  executableArgs,
+        ExecutionMethod: wmiexec.MethodCall,
+        ExecutionMethodConfig: wmiexec.MethodCallConfig{
+          Class:     wmi.Class,
+          Method:    wmi.Method,
+          Arguments: wmiMethodArgsMap,
+        },
+      }
+
+      ctx = log.With().
+        Str("module", "wmi").
+        Str("method", "proc").
+        Logger().WithContext(ctx)
+
+      if err := executor.Connect(ctx, creds, target, connCfg); err != nil {
+        log.Fatal().Err(err).Msg("Connection failed")
+      }
+      defer func() {
+        if err := executor.Cleanup(ctx, cleanCfg); err != nil {
+          log.Error().Err(err).Msg("Cleanup failed")
+        }
+      }()
+      if err := executor.Exec(ctx, execCfg); err != nil {
+        log.Error().Err(err).Msg("Execution failed")
+      }
+    },
+  }
+
+  wmiProcessCmd = &cobra.Command{
+    Use:   "proc",
+    Short: "Start a Windows process",
+    Long: `Description:
   The proc method creates an instance of the Win32_Process WMI class, then
   calls the Win32_Process.Create method with the provided command (-c),
   and optional working directory (-d).
@@ -110,31 +113,42 @@ References:
 References:
   https://learn.microsoft.com/en-us/windows/win32/cimwin32prov/create-method-in-class-win32-process
 `,
-		Args: needsTarget("cifs"),
-		Run: func(cmd *cobra.Command, args []string) {
-			log = log.With().Str("module", "wmi").Logger()
-
-			module := wmiexec.Module{}
-			connCfg := &exec.ConnectionConfig{}
-			cleanCfg := &exec.CleanupConfig{}
-
-			execCfg := &exec.ExecutionConfig{
-				ExecutableName:  executable,
-				ExecutableArgs:  executableArgs,
-				ExecutionMethod: wmiexec.MethodProcess,
-
-				ExecutionMethodConfig: wmiexec.MethodProcessConfig{
-					Command:          command,
-					WorkingDirectory: workingDirectory,
-				},
-			}
-			if err := module.Connect(log.WithContext(ctx), creds, target, connCfg); err != nil {
-				log.Fatal().Err(err).Msg("Connection failed")
-			} else if err := module.Exec(log.WithContext(ctx), execCfg); err != nil {
-				log.Fatal().Err(err).Msg("Execution failed")
-			} else if err := module.Cleanup(log.WithContext(ctx), cleanCfg); err != nil {
-				log.Error().Err(err).Msg("Cleanup failed")
-			}
-		},
-	}
+    Args: needs(needsTarget("cifs"), needsRpcTarget("cifs")),
+    Run: func(cmd *cobra.Command, args []string) {
+
+      executor := wmiexec.Module{}
+      cleanCfg := &exec.CleanupConfig{} // TODO
+      connCfg := &exec.ConnectionConfig{
+        ConnectionMethod:       exec.ConnectionMethodDCE,
+        ConnectionMethodConfig: dceConfig,
+      }
+      execCfg := &exec.ExecutionConfig{
+        ExecutableName:  executable,
+        ExecutableArgs:  executableArgs,
+        ExecutionMethod: wmiexec.MethodProcess,
+
+        ExecutionMethodConfig: wmiexec.MethodProcessConfig{
+          Command:          command,
+          WorkingDirectory: workingDirectory,
+        },
+      }
+
+      ctx = log.With().
+        Str("module", "wmi").
+        Str("method", "proc").
+        Logger().WithContext(ctx)
+
+      if err := executor.Connect(ctx, creds, target, connCfg); err != nil {
+        log.Fatal().Err(err).Msg("Connection failed")
+      }
+      defer func() {
+        if err := executor.Cleanup(ctx, cleanCfg); err != nil {
+          log.Error().Err(err).Msg("Cleanup failed")
+        }
+      }()
+      if err := executor.Exec(ctx, execCfg); err != nil {
+        log.Error().Err(err).Msg("Execution failed")
+      }
+    },
+  }
 )