Major update to the CLI cosmetics

cmd/args.go

@@ -1,139 +1,134 @@
 package cmd
 
 import (
-	"context"
-	"errors"
-	"fmt"
-	"github.com/spf13/cobra"
-	"github.com/spf13/pflag"
-	"os"
+  "context"
+  "errors"
+  "fmt"
+  "github.com/spf13/cobra"
+  "github.com/spf13/pflag"
+  "os"
 )
 
-func registerRpcFlags(cmd *cobra.Command) {
-	rpcFlags := pflag.NewFlagSet("RPC", pflag.ExitOnError)
-
-	rpcFlags.BoolVar(&rpcClient.NoEpm, "no-epm", false, "Do not use EPM to automatically detect endpoints")
-	//rpcFlags.BoolVar(&rpcClient.Options.EpmAuto, "epm-auto", false, "Automatically detect endpoints instead of using the module defaults")
-	rpcFlags.BoolVar(&rpcClient.NoSign, "no-sign", false, "Disable signing on DCE messages")
-	rpcFlags.BoolVar(&rpcClient.NoSeal, "no-seal", false, "Disable packet stub encryption on DCE messages")
-	rpcFlags.StringVar(&rpcClient.Filter, "epm-filter", "", "String binding to filter endpoints returned by EPM")
-	rpcFlags.StringVar(&rpcClient.Endpoint, "endpoint", "", "Explicit RPC endpoint definition")
-
-	cmd.PersistentFlags().AddFlagSet(rpcFlags)
-
-	cmd.MarkFlagsMutuallyExclusive("endpoint", "epm-filter")
-	cmd.MarkFlagsMutuallyExclusive("no-epm", "epm-filter")
+func registerLoggingFlags(fs *pflag.FlagSet) {
+  fs.SortFlags = false
+  fs.BoolVarP(&logDebug, "debug", "D", false, "Enable debug logging")
+  fs.StringVarP(&logOutput, "log-file", "O", "", "Write JSON logging output to `file`")
+  fs.BoolVarP(&logJson, "json", "j", false, "Write logging output in JSON lines")
+  fs.BoolVarP(&logQuiet, "quiet", "q", false, "Disable info logging")
 }
 
-func registerProcessExecutionArgs(cmd *cobra.Command) {
-	group := pflag.NewFlagSet("Execution", pflag.ExitOnError)
+func registerNetworkFlags(fs *pflag.FlagSet) {
+  fs.StringVarP(&proxy, "proxy", "x", "", "Proxy `URI`")
+  fs.StringVar(&rpcClient.Endpoint, "endpoint", "", "Explicit RPC endpoint definition")
+  fs.StringVar(&rpcClient.Filter, "epm-filter", "", "String binding to filter endpoints returned by the RPC endpoint mapper (EPM)")
+  fs.BoolVar(&rpcClient.NoEpm, "no-epm", false, "Do not use EPM to automatically detect RPC endpoints")
+  fs.BoolVar(&rpcClient.NoSign, "no-sign", false, "Disable signing on DCERPC messages")
+  fs.BoolVar(&rpcClient.NoSeal, "no-seal", false, "Disable packet stub encryption on DCERPC messages")
 
-	group.StringVarP(&exec.Input.Arguments, "args", "a", "", "Command line arguments")
-	group.StringVarP(&exec.Input.Command, "command", "c", "", "Windows process command line (executable & arguments)")
-	group.StringVarP(&exec.Input.Executable, "executable", "e", "", "Windows executable to invoke")
-
-	cmd.PersistentFlags().AddFlagSet(group)
-
-	cmd.MarkFlagsOneRequired("executable", "command")
-	cmd.MarkFlagsMutuallyExclusive("executable", "command")
+  //cmd.MarkFlagsMutuallyExclusive("endpoint", "epm-filter")
+  //cmd.MarkFlagsMutuallyExclusive("no-epm", "epm-filter")
 }
 
-func registerExecutionOutputArgs(cmd *cobra.Command) {
-	group := pflag.NewFlagSet("Output", pflag.ExitOnError)
+func registerExecutionFlags(fs *pflag.FlagSet) {
+  fs.StringVarP(&exec.Input.Executable, "executable", "e", "", "Windows executable to invoke")
+  fs.StringVarP(&exec.Input.Arguments, "args", "a", "", "Process command line arguments")
+  fs.StringVarP(&exec.Input.Command, "command", "c", "", "Windows process command line (executable & arguments)")
 
-	group.StringVarP(&outputPath, "output", "o", "", `Fetch execution output to file or "-" for standard output`)
-	group.StringVarP(&outputMethod, "output-method", "m", "smb", "Method to fetch execution output")
-	group.StringVar(&exec.Output.RemotePath, "remote-output", "", "Location to temporarily store output on remote filesystem")
-	group.BoolVar(&exec.Output.NoDelete, "no-delete-output", false, "Preserve output file on remote filesystem")
+  //cmd.MarkFlagsOneRequired("executable", "command")
+  //cmd.MarkFlagsMutuallyExclusive("executable", "command")
+}
 
-	cmd.PersistentFlags().AddFlagSet(group)
+func registerExecutionOutputFlags(fs *pflag.FlagSet) {
+  fs.StringVarP(&outputPath, "out", "o", "", `Fetch execution output to file or "-" for standard output`)
+  fs.StringVarP(&outputMethod, "out-method", "m", "smb", "Method to fetch execution output")
+  fs.StringVar(&exec.Output.RemotePath, "out-remote", "", "Location to temporarily store output on remote filesystem")
+  fs.BoolVar(&exec.Output.NoDelete, "no-delete-out", false, "Preserve output file on remote filesystem")
 }
 
 func args(reqs ...func(*cobra.Command, []string) error) (fn func(*cobra.Command, []string) error) {
+  return func(cmd *cobra.Command, args []string) (err error) {
 
-	return func(cmd *cobra.Command, args []string) (err error) {
-
-		for _, req := range reqs {
-			if err = req(cmd, args); err != nil {
-				return
-			}
-		}
-		return
-	}
+    for _, req := range reqs {
+      if err = req(cmd, args); err != nil {
+        return
+      }
+    }
+    return
+  }
 }
 
 func argsTarget(proto string) func(cmd *cobra.Command, args []string) error {
 
-	return func(cmd *cobra.Command, args []string) (err error) {
+  return func(cmd *cobra.Command, args []string) (err error) {
 
-		if len(args) != 1 {
-			return errors.New("command require exactly one positional argument: [target]")
-		}
+    if len(args) != 1 {
+      return errors.New("command require exactly one positional argument: [target]")
+    }
 
-		if credential, target, err = adAuthOpts.WithTarget(context.TODO(), proto, args[0]); err != nil {
-			return fmt.Errorf("failed to parse target: %w", err)
-		}
+    if credential, target, err = adAuthOpts.WithTarget(context.TODO(), proto, args[0]); err != nil {
+      return fmt.Errorf("failed to parse target: %w", err)
+    }
 
-		if credential == nil {
-			return errors.New("no credentials supplied")
-		}
-		if target == nil {
-			return errors.New("no target supplied")
-		}
-		return
-	}
+    if credential == nil {
+      return errors.New("no credentials supplied")
+    }
+    if target == nil {
+      return errors.New("no target supplied")
+    }
+    return
+  }
 }
 
 func argsSmbClient() func(cmd *cobra.Command, args []string) error {
-	return args(
-		argsTarget("cifs"),
+  return args(
+    argsTarget("cifs"),
 
-		func(_ *cobra.Command, _ []string) error {
+    func(_ *cobra.Command, _ []string) error {
 
-			smbClient.Credential = credential
-			smbClient.Target = target
-			smbClient.Proxy = proxy
+      smbClient.Credential = credential
+      smbClient.Target = target
+      smbClient.Proxy = proxy
 
-			return smbClient.Parse(context.TODO())
-		},
-	)
+      return smbClient.Parse(context.TODO())
+    },
+  )
 }
 
 func argsRpcClient(proto string) func(cmd *cobra.Command, args []string) error {
-	return args(
-		argsTarget(proto),
+  return args(
+    argsTarget(proto),
 
-		func(cmd *cobra.Command, args []string) (err error) {
+    func(cmd *cobra.Command, args []string) (err error) {
 
-			rpcClient.Target = target
-			rpcClient.Credential = credential
-			rpcClient.Proxy = proxy
+      rpcClient.Target = target
+      rpcClient.Credential = credential
+      rpcClient.Proxy = proxy
 
-			return rpcClient.Parse(context.TODO())
-		},
-	)
+      return rpcClient.Parse(context.TODO())
+    },
+  )
 }
 
 func argsOutput(methods ...string) func(cmd *cobra.Command, args []string) error {
 
-	var as []func(*cobra.Command, []string) error
+  var as []func(*cobra.Command, []string) error
 
-	for _, method := range methods {
-		if method == "smb" {
-			as = append(as, argsSmbClient())
-		}
-	}
+  for _, method := range methods {
+    if method == "smb" {
+      as = append(as, argsSmbClient())
+    }
+  }
 
-	return args(append(as, func(*cobra.Command, []string) (err error) {
+  return args(append(as, func(*cobra.Command, []string) (err error) {
 
-		if outputPath != "" {
-			if outputPath == "-" {
-				exec.Output.Writer = os.Stdout
+    if outputPath != "" {
+      if outputPath == "-" {
+        exec.Output.Writer = os.Stdout
 
-			} else if exec.Output.Writer, err = os.OpenFile(outputPath, os.O_WRONLY|os.O_CREATE|os.O_TRUNC, 0644); err != nil {
-				log.Fatal().Err(err).Msg("Failed to open output file")
-			}
-		}
-		return
-	})...)
+      } else if exec.Output.Writer, err = os.OpenFile(outputPath, os.O_WRONLY|os.O_CREATE|os.O_TRUNC, 0644); err != nil {
+        log.Fatal().Err(err).Msg("Failed to open output file")
+      }
+    }
+    return
+  })...)
 }

cmd/dcom.go

@@ -9,17 +9,36 @@ import (
 )
 
 func dcomCmdInit() {
-  registerRpcFlags(dcomCmd)
+  cmdFlags[dcomCmd] = []*flagSet{
+    defaultAuthFlags,
+    defaultLogFlags,
+    defaultNetRpcFlags,
+  }
   dcomMmcCmdInit()
+
+  dcomCmd.PersistentFlags().AddFlagSet(defaultAuthFlags.Flags)
+  dcomCmd.PersistentFlags().AddFlagSet(defaultLogFlags.Flags)
+  dcomCmd.PersistentFlags().AddFlagSet(defaultNetRpcFlags.Flags)
   dcomCmd.AddCommand(dcomMmcCmd)
 }
 
 func dcomMmcCmdInit() {
-  dcomMmcCmd.Flags().StringVarP(&dcomMmc.WorkingDirectory, "directory", "d", `C:\`, "Working directory")
-  dcomMmcCmd.Flags().StringVar(&dcomMmc.WindowState, "window", "Minimized", "Window state")
+  dcomMmcExecFlags := newFlagSet("Execution")
 
-  registerProcessExecutionArgs(dcomMmcCmd)
-  registerExecutionOutputArgs(dcomMmcCmd)
+  registerExecutionFlags(dcomMmcExecFlags.Flags)
+  registerExecutionOutputFlags(dcomMmcExecFlags.Flags)
+
+  dcomMmcExecFlags.Flags.StringVar(&dcomMmc.WorkingDirectory, "directory", `C:\`, "Working `directory`")
+  dcomMmcExecFlags.Flags.StringVar(&dcomMmc.WindowState, "window", "Minimized", "Window state")
+
+  cmdFlags[dcomMmcCmd] = []*flagSet{
+    dcomMmcExecFlags,
+    defaultAuthFlags,
+    defaultLogFlags,
+    defaultNetRpcFlags,
+  }
+
+  dcomMmcCmd.Flags().AddFlagSet(dcomMmcExecFlags.Flags)
 }
 
 var (

cmd/root.go

@@ -17,8 +17,43 @@ import (
   "os"
 )
 
+type flagSet struct {
+  Label string
+  Flags *pflag.FlagSet
+}
+
+const helpTemplate = `Usage:{{if .Runnable}}
+  {{.UseLine}}{{end}}{{if .HasAvailableSubCommands}}
+  {{.CommandPath}} [command] [flags]{{end}}{{if gt (len .Aliases) 0}}
+
+Aliases:
+{{.NameAndAliases}}{{end}}{{if .HasExample}}
+
+Examples:
+{{.Example}}{{end}}{{if .HasAvailableSubCommands}}{{$cmds := .Commands}}{{if eq (len .Groups) 0}}
+
+Available Commands:{{range $cmds}}{{if (or .IsAvailableCommand (eq .Name "help"))}}
+  {{rpad .Name .NamePadding }} {{.Short}}{{end}}{{end}}{{else}}{{range $group := .Groups}}
+
+{{.Title}}{{range $cmds}}{{if (and (eq .GroupID $group.ID) (or .IsAvailableCommand (eq .Name "help")))}}
+  {{rpad .Name .NamePadding }} {{.Short}}{{end}}{{end}}{{end}}{{if not .AllChildCommandsHaveGroup}}
+
+Additional Commands:{{range $cmds}}{{if (and (eq .GroupID "") (or .IsAvailableCommand (eq .Name "help")))}}
+  {{rpad .Name .NamePadding }} {{.Short}}{{end}}{{end}}{{end}}{{end}}{{end}}{{if (ne .Name "completion")}}{{range $_, $v := cmdFlags .}}
+
+{{$v.Label|trimTrailingWhitespaces}}:
+{{flags $v.Flags|trimTrailingWhitespaces}}{{end}}{{end}}{{if .HasHelpSubCommands}}
+
+Additional help topics:{{range .Commands}}{{if .IsAdditionalHelpTopicCommand}}
+{{rpad .CommandPath .CommandPathPadding}} {{.Short}}{{end}}{{end}}{{end}}{{if .HasAvailableSubCommands}}
+
+Use "{{.CommandPath}} [command] --help" for more information about a command.{{end}}
+`
+
 var (
-  flagGroups = map[string]*pflag.FlagSet{}
+  cmdFlags = make(map[*cobra.Command][]*flagSet)
+
+  defaultAuthFlags, defaultLogFlags, defaultNetRpcFlags *flagSet
 
   returnCode   int
   outputMethod string
@@ -105,28 +140,13 @@ var (
   }
 )
 
-func addFlagSet(fs *pflag.FlagSet) {
-  flagGroups[fs.Name()] = fs
-}
-
-func moduleFlags(cmd *cobra.Command, module string) (fs *pflag.FlagSet) {
-  fs, _ = flagGroups[module]
-  return
-}
-
-// Uses the users terminal size or width of 80 if cannot determine users width
-// Based on https://github.com/spf13/cobra/issues/1805#issuecomment-1246192724
-func wrappedFlagUsages(cmd *pflag.FlagSet) string {
-  fd := int(os.Stdout.Fd())
-  width := 80
-
-  // Get the terminal width and dynamically set
-  termWidth, _, err := term.GetSize(fd)
-  if err == nil {
-    width = termWidth
+func newFlagSet(name string) *flagSet {
+  flags := pflag.NewFlagSet(name, pflag.ExitOnError)
+  flags.SortFlags = false
+  return &flagSet{
+    Label: name,
+    Flags: flags,
   }
-
-  return cmd.FlagUsagesWrapped(width - 1)
 }
 
 func init() {
@@ -140,9 +160,21 @@ func init() {
   // Cobra init
   {
     cobra.EnableCommandSorting = false
-
-    rootCmd.InitDefaultVersionFlag()
-    rootCmd.InitDefaultHelpCmd()
+    {
+      defaultNetRpcFlags = newFlagSet("Network")
+      registerNetworkFlags(defaultNetRpcFlags.Flags)
+    }
+    {
+      defaultLogFlags = newFlagSet("Logging")
+      registerLoggingFlags(defaultLogFlags.Flags)
+    }
+    {
+      defaultAuthFlags = newFlagSet("Authentication")
+      adAuthOpts = &adauth.Options{
+        Debug: log.Debug().Msgf,
+      }
+      adAuthOpts.RegisterFlags(defaultAuthFlags.Flags)
+    }
 
     modules := &cobra.Group{
       ID:    "module",
@@ -150,33 +182,26 @@ func init() {
     }
     rootCmd.AddGroup(modules)
 
-    // Logging flags
-    {
-      logOpts := pflag.NewFlagSet("Logging", pflag.ExitOnError)
-      logOpts.BoolVar(&logDebug, "debug", false, "Enable debug logging")
-      logOpts.BoolVar(&logJson, "json", false, "Write logging output in JSON lines")
-      logOpts.BoolVar(&logQuiet, "quiet", false, "Disable info logging")
-      logOpts.StringVarP(&logOutput, "log-file", "O", "", "Write JSON logging output to `file`")
-      rootCmd.PersistentFlags().AddFlagSet(logOpts)
-      flagGroups["Logging"] = logOpts
+    cmdFlags[rootCmd] = []*flagSet{
+      defaultLogFlags,
+      defaultAuthFlags,
     }
 
-    // Global networking flags
-    {
-      netOpts := pflag.NewFlagSet("Network", pflag.ExitOnError)
-      netOpts.StringVarP(&proxy, "proxy", "x", "", "Proxy `URI`")
-      rootCmd.PersistentFlags().AddFlagSet(netOpts)
-    }
-
-    // Authentication flags
-    {
-      adAuthOpts = &adauth.Options{
-        Debug: log.Debug().Msgf,
+    cobra.AddTemplateFunc("flags", func(fs *pflag.FlagSet) string {
+      if width, _, err := term.GetSize(int(os.Stdout.Fd())); err == nil {
+        return fs.FlagUsagesWrapped(width - 1)
       }
-      authOpts := pflag.NewFlagSet("Authentication", pflag.ExitOnError)
-      adAuthOpts.RegisterFlags(authOpts)
-      rootCmd.PersistentFlags().AddFlagSet(authOpts)
-    }
+      return fs.FlagUsagesWrapped(80 - 1)
+    })
+
+    cobra.AddTemplateFunc("cmdFlags", func(cmd *cobra.Command) []*flagSet {
+      return cmdFlags[cmd]
+    })
+
+    rootCmd.InitDefaultVersionFlag()
+    rootCmd.InitDefaultHelpCmd()
+    rootCmd.SetHelpTemplate("{{if (ne .Long \"\")}}{{.Long}}\n{{end}}" + helpTemplate)
+    rootCmd.SetUsageTemplate(helpTemplate)
 
     // Modules init
     {

cmd/scmr.go

@@ -10,43 +10,90 @@ import (
 )
 
 func scmrCmdInit() {
-  registerRpcFlags(scmrCmd)
+  cmdFlags[scmrCmd] = []*flagSet{
+    defaultAuthFlags,
+    defaultLogFlags,
+    defaultNetRpcFlags,
+  }
   scmrCreateCmdInit()
-  scmrCmd.AddCommand(scmrCreateCmd)
   scmrChangeCmdInit()
-  scmrCmd.AddCommand(scmrChangeCmd)
   scmrDeleteCmdInit()
-  scmrCmd.AddCommand(scmrDeleteCmd)
+
+  scmrCmd.PersistentFlags().AddFlagSet(defaultAuthFlags.Flags)
+  scmrCmd.PersistentFlags().AddFlagSet(defaultLogFlags.Flags)
+  scmrCmd.PersistentFlags().AddFlagSet(defaultNetRpcFlags.Flags)
+  scmrCmd.AddCommand(scmrCreateCmd, scmrChangeCmd, scmrDeleteCmd)
 }
 
 func scmrCreateCmdInit() {
-  scmrCreateCmd.Flags().StringVarP(&scmrCreate.DisplayName, "display-name", "n", "", "Display name of service to create")
-  scmrCreateCmd.Flags().StringVarP(&scmrCreate.ServiceName, "service-name", "s", "", "Name of service to create")
-  scmrCreateCmd.Flags().BoolVar(&scmrCreate.NoDelete, "no-delete", false, "Don't delete service after execution")
-  scmrCreateCmd.Flags().BoolVar(&scmrCreate.NoStart, "no-start", false, "Don't start service")
+  scmrCreateFlags := newFlagSet("Service")
 
-  scmrCreateCmd.Flags().StringVarP(&exec.Input.ExecutablePath, "executable-path", "f", "", "Full path to a remote Windows executable")
-  scmrCreateCmd.Flags().StringVarP(&exec.Input.Arguments, "args", "a", "", "Arguments to pass to the executable")
+  scmrCreateFlags.Flags.StringVarP(&scmrCreate.DisplayName, "display-name", "n", "", "Display name of service to create")
+  scmrCreateFlags.Flags.StringVarP(&scmrCreate.ServiceName, "service", "s", "", "Name of service to create")
+  scmrCreateFlags.Flags.BoolVar(&scmrCreate.NoDelete, "no-delete", false, "Don't delete service after execution")
+  scmrCreateFlags.Flags.BoolVar(&scmrCreate.NoStart, "no-start", false, "Don't start service")
 
-  scmrCreateCmd.MarkFlagsMutuallyExclusive("no-delete", "no-start")
+  scmrCreateExecFlags := newFlagSet("Execution")
 
-  if err := scmrCreateCmd.MarkFlagRequired("executable-path"); err != nil {
-    panic(err)
+  // TODO: SCMR output
+  //registerExecutionOutputFlags(scmrCreateExecFlags.Flags)
+
+  scmrCreateExecFlags.Flags.StringVarP(&exec.Input.ExecutablePath, "executable-path", "f", "", "Full path to a remote Windows executable")
+  scmrCreateExecFlags.Flags.StringVarP(&exec.Input.Arguments, "args", "a", "", "Arguments to pass to the executable")
+
+  scmrCreateCmd.Flags().AddFlagSet(scmrCreateFlags.Flags)
+  scmrCreateCmd.Flags().AddFlagSet(scmrCreateExecFlags.Flags)
+
+  cmdFlags[scmrCreateCmd] = []*flagSet{
+    scmrCreateExecFlags,
+    scmrCreateFlags,
+    defaultAuthFlags,
+    defaultLogFlags,
+    defaultNetRpcFlags,
+  }
+
+  // Constraints
+  {
+    scmrCreateCmd.MarkFlagsMutuallyExclusive("no-delete", "no-start")
+    if err := scmrCreateCmd.MarkFlagRequired("executable-path"); err != nil {
+      panic(err)
+    }
   }
 }
 
 func scmrChangeCmdInit() {
-  scmrChangeCmd.Flags().BoolVar(&scmrChange.NoStart, "no-start", false, "Don't start service")
-  scmrChangeCmd.Flags().StringVarP(&scmrChange.ServiceName, "service-name", "s", "", "Name of service to modify")
+  scmrChangeFlags := newFlagSet("Service Control")
 
-  scmrChangeCmd.Flags().StringVarP(&exec.Input.ExecutablePath, "executable-path", "f", "", "Full path to remote Windows executable")
-  scmrChangeCmd.Flags().StringVarP(&exec.Input.Arguments, "args", "a", "", "Arguments to pass to executable")
+  scmrChangeFlags.Flags.StringVarP(&scmrChange.ServiceName, "service-name", "s", "", "Name of service to modify")
+  scmrChangeFlags.Flags.BoolVar(&scmrChange.NoStart, "no-start", false, "Don't start service")
 
-  if err := scmrChangeCmd.MarkFlagRequired("service-name"); err != nil {
-    panic(err)
+  scmrChangeExecFlags := newFlagSet("Execution")
+
+  scmrChangeExecFlags.Flags.StringVarP(&exec.Input.ExecutablePath, "executable-path", "f", "", "Full path to remote Windows executable")
+  scmrChangeExecFlags.Flags.StringVarP(&exec.Input.Arguments, "args", "a", "", "Arguments to pass to executable")
+
+  // TODO: SCMR output
+  //registerExecutionOutputFlags(scmrChangeExecFlags.Flags)
+
+  cmdFlags[scmrChangeCmd] = []*flagSet{
+    scmrChangeFlags,
+    scmrChangeExecFlags,
+    defaultAuthFlags,
+    defaultLogFlags,
+    defaultNetRpcFlags,
   }
-  if err := scmrCreateCmd.MarkFlagRequired("executable-path"); err != nil {
-    panic(err)
+
+  scmrChangeCmd.Flags().AddFlagSet(scmrChangeFlags.Flags)
+  scmrChangeCmd.Flags().AddFlagSet(scmrChangeExecFlags.Flags)
+
+  // Constraints
+  {
+    if err := scmrChangeCmd.MarkFlagRequired("service-name"); err != nil {
+      panic(err)
+    }
+    if err := scmrCreateCmd.MarkFlagRequired("executable-path"); err != nil {
+      panic(err)
+    }
   }
 }
 
@@ -94,7 +141,7 @@ References:
         WithContext(ctx)
 
       if scmrCreate.ServiceName == "" {
-        log.Warn().Msg("No service name was provided. Using a random string")
+        log.Warn().Msg("No service Label was provided. Using a random string")
         scmrCreate.ServiceName = util.RandomString()
       }
 
@@ -103,7 +150,7 @@ References:
       }
 
       if scmrCreate.DisplayName == "" {
-        log.Debug().Msg("No display name specified, using service name as display name")
+        log.Debug().Msg("No display Label specified, using service Label as display Label")
         scmrCreate.DisplayName = scmrCreate.ServiceName
       }
 

cmd/tsch.go

@@ -12,36 +12,71 @@ import (
 )
 
 func tschCmdInit() {
-  registerRpcFlags(tschCmd)
-
+  cmdFlags[tschCmd] = []*flagSet{
+    defaultAuthFlags,
+    defaultLogFlags,
+    defaultNetRpcFlags,
+  }
   tschDemandCmdInit()
-  tschCmd.AddCommand(tschDemandCmd)
-
   tschCreateCmdInit()
-  tschCmd.AddCommand(tschCreateCmd)
+
+  tschCmd.PersistentFlags().AddFlagSet(defaultAuthFlags.Flags)
+  tschCmd.PersistentFlags().AddFlagSet(defaultLogFlags.Flags)
+  tschCmd.PersistentFlags().AddFlagSet(defaultNetRpcFlags.Flags)
+  tschCmd.AddCommand(tschDemandCmd, tschCreateCmd)
 }
 
 func tschDemandCmdInit() {
-  tschDemandCmd.Flags().StringVarP(&tschTask, "task", "t", "", "Name or path of the new task")
-  tschDemandCmd.Flags().Uint32Var(&tschDemand.SessionId, "session", 0, "Hijack existing session given the session ID")
-  tschDemandCmd.Flags().BoolVar(&tschDemand.NoDelete, "no-delete", false, "Don't delete task after execution")
-  tschDemandCmd.Flags().StringVar(&tschDemand.UserSid, "sid", "S-1-5-18", "User SID to impersonate")
+  tschDemandFlags := newFlagSet("Task Scheduler")
 
-  registerProcessExecutionArgs(tschDemandCmd)
-  registerExecutionOutputArgs(tschDemandCmd)
+  tschDemandFlags.Flags.StringVarP(&tschTask, "task", "t", "", "Name or path of the new task")
+  tschDemandFlags.Flags.Uint32Var(&tschDemand.SessionId, "session", 0, "Hijack existing session given the session ID")
+  tschDemandFlags.Flags.StringVar(&tschDemand.UserSid, "sid", "S-1-5-18", "User SID to impersonate")
+  tschDemandFlags.Flags.BoolVar(&tschDemand.NoDelete, "no-delete", false, "Don't delete task after execution")
+
+  tschDemandExecFlags := newFlagSet("Execution")
+
+  registerExecutionFlags(tschDemandExecFlags.Flags)
+  registerExecutionOutputFlags(tschDemandExecFlags.Flags)
+
+  cmdFlags[tschDemandCmd] = []*flagSet{
+    tschDemandFlags,
+    tschDemandExecFlags,
+    defaultAuthFlags,
+    defaultLogFlags,
+    defaultNetRpcFlags,
+  }
+
+  tschDemandCmd.Flags().AddFlagSet(tschDemandFlags.Flags)
+  tschDemandCmd.Flags().AddFlagSet(tschDemandExecFlags.Flags)
 }
 
 func tschCreateCmdInit() {
-  tschCreateCmd.Flags().StringVarP(&tschTask, "task", "t", "", "Name or path of the new task")
-  tschCreateCmd.Flags().DurationVar(&tschCreate.StopDelay, "delay-stop", 5*time.Second, "Delay between task execution and termination. This won't stop the spawned process")
-  tschCreateCmd.Flags().DurationVar(&tschCreate.StartDelay, "start-delay", 5*time.Second, "Delay between task registration and execution")
-  tschCreateCmd.Flags().DurationVar(&tschCreate.DeleteDelay, "delete-delay", 0*time.Second, "Delay between task termination and deletion")
-  tschCreateCmd.Flags().BoolVar(&tschCreate.NoDelete, "no-delete", false, "Don't delete task after execution")
-  tschCreateCmd.Flags().BoolVar(&tschCreate.CallDelete, "call-delete", false, "Directly call SchRpcDelete to delete task")
-  tschCreateCmd.Flags().StringVar(&tschCreate.UserSid, "sid", "S-1-5-18", "User SID to impersonate")
+  tschCreateFlags := newFlagSet("Task Scheduler")
 
-  registerProcessExecutionArgs(tschCreateCmd)
-  registerExecutionOutputArgs(tschCreateCmd)
+  tschCreateFlags.Flags.StringVarP(&tschTask, "task", "t", "", "Name or path of the new task")
+  tschCreateFlags.Flags.DurationVar(&tschCreate.StopDelay, "delay-stop", 5*time.Second, "Delay between task execution and termination. This won't stop the spawned process")
+  tschCreateFlags.Flags.DurationVar(&tschCreate.StartDelay, "start-delay", 5*time.Second, "Delay between task registration and execution")
+  tschCreateFlags.Flags.DurationVar(&tschCreate.DeleteDelay, "delete-delay", 0*time.Second, "Delay between task termination and deletion")
+  tschCreateFlags.Flags.BoolVar(&tschCreate.NoDelete, "no-delete", false, "Don't delete task after execution")
+  tschCreateFlags.Flags.BoolVar(&tschCreate.CallDelete, "call-delete", false, "Directly call SchRpcDelete to delete task")
+  tschCreateFlags.Flags.StringVar(&tschCreate.UserSid, "sid", "S-1-5-18", "User `SID` to impersonate")
+
+  tschCreateExecFlags := newFlagSet("Execution")
+
+  registerExecutionFlags(tschCreateExecFlags.Flags)
+  registerExecutionOutputFlags(tschCreateExecFlags.Flags)
+
+  cmdFlags[tschCreateCmd] = []*flagSet{
+    tschCreateFlags,
+    tschCreateExecFlags,
+    defaultAuthFlags,
+    defaultLogFlags,
+    defaultNetRpcFlags,
+  }
+
+  tschCreateCmd.Flags().AddFlagSet(tschCreateFlags.Flags)
+  tschCreateCmd.Flags().AddFlagSet(tschCreateExecFlags.Flags)
 }
 
 func argsTask(*cobra.Command, []string) error {
@@ -52,7 +87,7 @@ func argsTask(*cobra.Command, []string) error {
   case tschexec.ValidateTaskName(tschTask) == nil:
     tschTask = `\` + tschTask
   default:
-    return fmt.Errorf("invalid task name or path: %q", tschTask)
+    return fmt.Errorf("invalid task Label or path: %q", tschTask)
   }
   return nil
 }

cmd/wmi.go

@@ -3,33 +3,44 @@ package cmd
 import (
   "context"
   "encoding/json"
-  "fmt"
   "github.com/FalconOpsLLC/goexec/pkg/goexec"
   wmiexec "github.com/FalconOpsLLC/goexec/pkg/goexec/wmi"
   "github.com/oiweiwei/go-msrpc/ssp/gssapi"
   "github.com/spf13/cobra"
+  "os"
 )
 
 func wmiCmdInit() {
-  registerRpcFlags(wmiCmd)
-
+  cmdFlags[wmiCmd] = []*flagSet{
+    defaultAuthFlags,
+    defaultLogFlags,
+    defaultNetRpcFlags,
+  }
   wmiCallCmdInit()
-  wmiCmd.AddCommand(wmiCallCmd)
-
   wmiProcCmdInit()
-  wmiCmd.AddCommand(wmiProcCmd)
-}
 
-func wmiCallArgs(_ *cobra.Command, _ []string) error {
-  return json.Unmarshal([]byte(wmiArguments), &wmiCall.Args)
+  wmiCmd.PersistentFlags().AddFlagSet(defaultAuthFlags.Flags)
+  wmiCmd.PersistentFlags().AddFlagSet(defaultLogFlags.Flags)
+  wmiCmd.PersistentFlags().AddFlagSet(defaultNetRpcFlags.Flags)
+  wmiCmd.AddCommand(wmiProcCmd, wmiCallCmd)
 }
 
 func wmiCallCmdInit() {
-  wmiCallCmd.Flags().StringVarP(&wmiCall.Resource, "namespace", "n", "//./root/cimv2", "WMI namespace")
-  wmiCallCmd.Flags().StringVarP(&wmiCall.Class, "class", "C", "", `WMI class to instantiate (i.e. "Win32_Process")`)
-  wmiCallCmd.Flags().StringVarP(&wmiCall.Method, "method", "m", "", `WMI Method to call (i.e. "Create")`)
-  wmiCallCmd.Flags().StringVarP(&wmiArguments, "args", "A", "{}", `WMI Method argument(s) in JSON dictionary format (i.e. {"Command":"calc.exe"})`)
+  wmiCallFlags := newFlagSet("WMI")
 
+  wmiCallFlags.Flags.StringVarP(&wmiCall.Resource, "namespace", "n", "//./root/cimv2", "WMI namespace")
+  wmiCallFlags.Flags.StringVarP(&wmiCall.Class, "class", "C", "", `WMI class to instantiate (i.e. "Win32_Process")`)
+  wmiCallFlags.Flags.StringVarP(&wmiCall.Method, "method", "m", "", `WMI Method to call (i.e. "Create")`)
+  wmiCallFlags.Flags.StringVarP(&wmiArguments, "args", "A", "{}", `WMI Method argument(s) in JSON dictionary format (i.e. {"Command":"calc.exe"})`)
+
+  wmiCallCmd.Flags().AddFlagSet(wmiCallFlags.Flags)
+
+  cmdFlags[wmiCallCmd] = []*flagSet{
+    wmiCallFlags,
+    defaultAuthFlags,
+    defaultLogFlags,
+    defaultNetRpcFlags,
+  }
   if err := wmiCallCmd.MarkFlagRequired("class"); err != nil {
     panic(err)
   }
@@ -39,11 +50,26 @@ func wmiCallCmdInit() {
 }
 
 func wmiProcCmdInit() {
-  wmiProcCmd.Flags().StringVarP(&wmiProc.Resource, "namespace", "n", "//./root/cimv2", "WMI namespace")
-  wmiProcCmd.Flags().StringVarP(&wmiProc.WorkingDirectory, "directory", "d", `C:\`, "Working directory")
+  wmiProcFlags := newFlagSet("WMI")
 
-  registerProcessExecutionArgs(wmiProcCmd)
-  registerExecutionOutputArgs(wmiProcCmd)
+  wmiProcFlags.Flags.StringVarP(&wmiProc.Resource, "namespace", "n", "//./root/cimv2", "WMI namespace")
+  wmiProcFlags.Flags.StringVarP(&wmiProc.WorkingDirectory, "directory", "d", `C:\`, "Working directory")
+
+  wmiProcExecFlags := newFlagSet("Execution")
+
+  registerExecutionFlags(wmiProcExecFlags.Flags)
+  registerExecutionOutputFlags(wmiProcExecFlags.Flags)
+
+  cmdFlags[wmiProcCmd] = []*flagSet{
+    wmiProcExecFlags,
+    wmiProcFlags,
+    defaultAuthFlags,
+    defaultLogFlags,
+    defaultNetRpcFlags,
+  }
+
+  wmiProcCmd.Flags().AddFlagSet(wmiProcFlags.Flags)
+  wmiProcCmd.Flags().AddFlagSet(wmiProcExecFlags.Flags)
 }
 
 var (
@@ -69,43 +95,24 @@ var (
 References:
   https://learn.microsoft.com/en-us/windows/win32/wmisdk/wmi-classes
 `,
-    Args: args(argsRpcClient("host"), wmiCallArgs),
+    Args: args(
+      argsRpcClient("cifs"),
+      func(cmd *cobra.Command, args []string) error {
+        return json.Unmarshal([]byte(wmiArguments), &wmiCall.Args)
+      }),
 
     Run: func(cmd *cobra.Command, args []string) {
-      var err error
+      wmiCall.Client = &rpcClient
+      wmiCall.Out = os.Stdout
 
-      ctx := gssapi.NewSecurityContext(context.Background())
-
-      ctx = log.With().
+      ctx := log.With().
         Str("module", "wmi").
         Str("method", "call").
-        Logger().
-        WithContext(ctx)
+        Logger().WithContext(gssapi.NewSecurityContext(context.Background()))
 
-      if err = rpcClient.Connect(ctx); err != nil {
-        log.Fatal().Err(err).Msg("Connection failed")
+      if err := goexec.ExecuteAuxiliaryMethod(ctx, &wmiCall); err != nil {
+        log.Fatal().Err(err).Msg("Operation failed")
       }
-
-      defer func() {
-        closeErr := rpcClient.Close(ctx)
-        if closeErr != nil {
-          log.Error().Err(closeErr).Msg("Failed to close connection")
-        }
-      }()
-
-      if err = wmiCall.Init(ctx); err != nil {
-        log.Error().Err(err).Msg("Module initialization failed")
-        returnCode = 2
-        return
-      }
-
-      out, err := wmiCall.Call(ctx)
-      if err != nil {
-        log.Error().Err(err).Msg("Call failed")
-        returnCode = 4
-        return
-      }
-      fmt.Println(string(out))
     },
   }
 
@@ -121,8 +128,8 @@ References:
   https://learn.microsoft.com/en-us/windows/win32/cimwin32prov/create-method-in-class-win32-process
 `,
     Args: args(
+      argsRpcClient("cifs"),
       argsOutput("smb"),
-      argsRpcClient("host"),
     ),
 
     Run: func(cmd *cobra.Command, args []string) {

pkg/goexec/method.go

@@ -6,36 +6,54 @@ import (
   "github.com/rs/zerolog"
 )
 
-type Method interface{}
-
-type RemoteMethod interface {
+type Method interface {
   Connect(ctx context.Context) error
   Init(ctx context.Context) error
 }
 
-type RemoteExecuteMethod interface {
-  RemoteMethod
-  Execute(ctx context.Context, io *ExecutionIO) error
-}
-
-type RemoteExecuteCleanMethod interface {
-  RemoteExecuteMethod
+type Clean interface {
   Clean(ctx context.Context) error
 }
 
-func ExecuteMethod(ctx context.Context, module RemoteExecuteMethod, execIO *ExecutionIO) (err error) {
+type CleanMethod interface {
+  Method
+  Clean
+}
 
+type ExecutionMethod interface {
+  Method
+  Execute(ctx context.Context, io *ExecutionIO) error
+}
+
+type CleanExecutionMethod interface {
+  ExecutionMethod
+  Clean
+}
+
+type AuxiliaryMethod interface {
+  Method
+  Call(ctx context.Context) error
+}
+
+type CleanAuxiliaryMethod interface {
+  AuxiliaryMethod
+  Clean
+}
+
+func ExecuteMethod(ctx context.Context, module ExecutionMethod, execIO *ExecutionIO) (err error) {
   log := zerolog.Ctx(ctx)
 
   if err = module.Connect(ctx); err != nil {
     log.Error().Err(err).Msg("Connection failed")
     return fmt.Errorf("connect: %w", err)
   }
+  log.Debug().Msg("Module connected")
 
   if err = module.Init(ctx); err != nil {
     log.Error().Err(err).Msg("Module initialization failed")
     return fmt.Errorf("init module: %w", err)
   }
+  log.Debug().Msg("Module initialized")
 
   if err = module.Execute(ctx, execIO); err != nil {
     log.Error().Err(err).Msg("Execution failed")
@@ -45,8 +63,48 @@ func ExecuteMethod(ctx context.Context, module RemoteExecuteMethod, execIO *Exec
   return
 }
 
-func ExecuteCleanMethod(ctx context.Context, module RemoteExecuteCleanMethod, execIO *ExecutionIO) (err error) {
+func ExecuteAuxiliaryMethod(ctx context.Context, module AuxiliaryMethod) (err error) {
+  log := zerolog.Ctx(ctx)
 
+  if err = module.Connect(ctx); err != nil {
+    log.Error().Err(err).Msg("Connection failed")
+    return fmt.Errorf("connect: %w", err)
+  }
+  log.Debug().Msg("Auxiliary module connected")
+
+  if err = module.Init(ctx); err != nil {
+    log.Error().Err(err).Msg("Module initialization failed")
+    return fmt.Errorf("init module: %w", err)
+  }
+  log.Debug().Msg("Auxiliary module initialized")
+
+  if err = module.Call(ctx); err != nil {
+    log.Error().Err(err).Msg("Auxiliary method failed")
+    return fmt.Errorf("call: %w", err)
+  }
+  log.Debug().Msg("Auxiliary method succeeded")
+
+  return nil
+}
+
+func ExecuteCleanAuxiliaryMethod(ctx context.Context, module CleanAuxiliaryMethod) (err error) {
+  log := zerolog.Ctx(ctx)
+
+  defer func() {
+    if err = module.Clean(ctx); err != nil {
+      log.Error().Err(err).Msg("Module cleanup failed")
+      err = nil
+    }
+  }()
+
+  if err = ExecuteAuxiliaryMethod(ctx, module); err != nil {
+    log.Error().Err(err).Msg("Auxiliary method failed")
+    return fmt.Errorf("execute auxiliary method: %w", err)
+  }
+  return
+}
+
+func ExecuteCleanMethod(ctx context.Context, module CleanExecutionMethod, execIO *ExecutionIO) (err error) {
   log := zerolog.Ctx(ctx)
 
   defer func() {
@@ -61,9 +119,19 @@ func ExecuteCleanMethod(ctx context.Context, module RemoteExecuteCleanMethod, ex
   }
 
   if execIO.Output != nil && execIO.Output.Provider != nil {
-    defer execIO.Output.Provider.Clean(ctx)
+    log.Info().Msg("Collecting output")
 
-    execIO.Output.Provider.GetOutput(ctx, execIO.Output.Writer)
+    defer func() {
+      if cleanErr := execIO.Output.Provider.Clean(ctx); cleanErr != nil {
+        log.Debug().Err(cleanErr).Msg("Output provider cleanup failed")
+      }
+    }()
+
+    if err := execIO.Output.Provider.GetOutput(ctx, execIO.Output.Writer); err != nil {
+      log.Error().Err(err).Msg("Output collection failed")
+      return fmt.Errorf("get output: %w", err)
+    }
+    log.Debug().Msg("Output collection succeeded")
   }
   return
 }

pkg/goexec/wmi/call.go

@@ -3,7 +3,9 @@ package wmiexec
 import (
   "context"
   "encoding/json"
+  "fmt"
   "github.com/rs/zerolog"
+  "io"
 )
 
 type WmiCall struct {
@@ -12,16 +14,25 @@ type WmiCall struct {
   Class  string
   Method string
   Args   map[string]any
+
+  Out io.Writer
 }
 
-func (m *WmiCall) Call(ctx context.Context) (out []byte, err error) {
+func (m *WmiCall) Call(ctx context.Context) (err error) {
   var outMap map[string]any
 
   if outMap, err = m.query(ctx, m.Class, m.Method, m.Args); err != nil {
     return
   }
-  zerolog.Ctx(ctx).Info().Msg("Call succeeded")
+  zerolog.Ctx(ctx).Info().Msg("WMI call successful")
 
-  out, err = json.Marshal(outMap)
+  out, err := json.Marshal(outMap)
+
+  if m.Out != nil {
+    // Write output with a trailing line feed
+    if _, err = m.Out.Write(append(out, 0x0a)); err != nil {
+      return fmt.Errorf("write output: %w", err)
+    }
+  }
   return
 }