Update output,IO; add output support to WMI

5 files changed, 170 insertions, 201 deletions

diff --git a/cmd/args.go b/cmd/args.go
index b2b6fe2..ca61c22 100644
--- a/cmd/args.go
+++ b/cmd/args.go
@@ -1,142 +1,139 @@
 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 := 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")
+	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.PersistentFlags().AddFlagSet(rpcFlags)
 
-  cmd.MarkFlagsMutuallyExclusive("endpoint", "epm-filter")
-  cmd.MarkFlagsMutuallyExclusive("no-epm", "epm-filter")
+	cmd.MarkFlagsMutuallyExclusive("endpoint", "epm-filter")
+	cmd.MarkFlagsMutuallyExclusive("no-epm", "epm-filter")
 }
 
 func registerProcessExecutionArgs(cmd *cobra.Command) {
-  group := pflag.NewFlagSet("Execution", pflag.ExitOnError)
+	group := pflag.NewFlagSet("Execution", pflag.ExitOnError)
 
-  group.StringVarP(&exec.Input.Arguments, "args", "a", "", "Command line arguments")
-  group.StringVarP(&exec.Input.CommandLine, "command", "c", "", "Windows process command line (executable & arguments)")
-  group.StringVarP(&exec.Input.Executable, "executable", "e", "", "Windows executable to invoke")
+	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.PersistentFlags().AddFlagSet(group)
 
-  cmd.MarkFlagsOneRequired("executable", "command")
-  cmd.MarkFlagsMutuallyExclusive("executable", "command")
+	cmd.MarkFlagsOneRequired("executable", "command")
+	cmd.MarkFlagsMutuallyExclusive("executable", "command")
 }
 
 func registerExecutionOutputArgs(cmd *cobra.Command) {
-  group := pflag.NewFlagSet("Output", pflag.ExitOnError)
+	group := pflag.NewFlagSet("Output", pflag.ExitOnError)
 
-  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")
+	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.PersistentFlags().AddFlagSet(group)
+	cmd.PersistentFlags().AddFlagSet(group)
 }
 
 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 = authOpts.WithTarget(context.TODO(), proto, args[0]); err != nil {
-      return fmt.Errorf("failed to parse target: %w", err)
-    }
+		if credential, target, err = authOpts.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 outputPath != "" {
-
-        if exec.Output.Writer, err = os.OpenFile(outputPath, os.O_WRONLY|os.O_CREATE, 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, 0644); err != nil {
+				log.Fatal().Err(err).Msg("Failed to open output file")
+			}
+		}
+		return
+	})...)
 }
diff --git a/cmd/dcom.go b/cmd/dcom.go
index 5bcec78..c671252 100644
--- a/cmd/dcom.go
+++ b/cmd/dcom.go
@@ -1,37 +1,40 @@
 package cmd
 
 import (
-  "context"
-  dcomexec "github.com/FalconOpsLLC/goexec/pkg/goexec/dcom"
-  "github.com/oiweiwei/go-msrpc/ssp/gssapi"
-  "github.com/spf13/cobra"
+	"context"
+	"github.com/FalconOpsLLC/goexec/pkg/goexec"
+	dcomexec "github.com/FalconOpsLLC/goexec/pkg/goexec/dcom"
+	"github.com/oiweiwei/go-msrpc/ssp/gssapi"
+	"github.com/spf13/cobra"
 )
 
 func dcomCmdInit() {
-  registerRpcFlags(dcomCmd)
-  dcomMmcCmdInit()
-  dcomCmd.AddCommand(dcomMmcCmd)
+	registerRpcFlags(dcomCmd)
+	dcomMmcCmdInit()
+	dcomCmd.AddCommand(dcomMmcCmd)
 }
 
 func dcomMmcCmdInit() {
-  dcomMmcCmd.Flags().StringVarP(&dcomMmc.WorkingDirectory, "directory", "d", `C:\`, "Working directory")
-  dcomMmcCmd.Flags().StringVar(&dcomMmc.WindowState, "window", "Minimized", "Window state")
+	dcomMmcCmd.Flags().StringVarP(&dcomMmc.WorkingDirectory, "directory", "d", `C:\`, "Working directory")
+	dcomMmcCmd.Flags().StringVar(&dcomMmc.WindowState, "window", "Minimized", "Window state")
 
-  registerProcessExecutionArgs(dcomMmcCmd)
+	registerProcessExecutionArgs(dcomMmcCmd)
+	registerExecutionOutputArgs(dcomMmcCmd)
 }
 
 var (
-  dcomMmc dcomexec.DcomMmc
-
-  dcomCmd = &cobra.Command{
-    Use:   "dcom",
-    Short: "Establish execution via DCOM",
-    Args:  cobra.NoArgs,
-  }
-  dcomMmcCmd = &cobra.Command{
-    Use:   "mmc [target]",
-    Short: "Establish execution via the DCOM MMC20.Application object",
-    Long: `Description:
+	dcomMmc dcomexec.DcomMmc
+
+	dcomCmd = &cobra.Command{
+		Use:   "dcom",
+		Short: "Establish execution via DCOM",
+		Args:  cobra.NoArgs,
+	}
+
+	dcomMmcCmd = &cobra.Command{
+		Use:   "mmc [target]",
+		Short: "Establish execution via the DCOM MMC20.Application object",
+		Long: `Description:
   The mmc method uses the exposed MMC20.Application object to call Document.ActiveView.ShellExec,
   and ultimately execute system commands.
 
@@ -41,39 +44,19 @@ References:
   https://github.com/fortra/impacket/blob/master/examples/dcomexec.py
   https://learn.microsoft.com/en-us/previous-versions/windows/desktop/mmc/view-executeshellcommand
 `,
-    Args: argsRpcClient("host"),
-    Run: func(cmd *cobra.Command, args []string) {
-      var err error
-
-      ctx := gssapi.NewSecurityContext(context.Background())
-
-      ctx = log.With().
-        Str("module", "dcom").
-        Str("method", "mmc").
-        Logger().
-        WithContext(ctx)
-
-      if err = rpcClient.Connect(ctx); err != nil {
-        log.Fatal().Err(err).Msg("Connection failed")
-      }
-
-      defer func() {
-        closeErr := rpcClient.Close(ctx)
-        if closeErr != nil {
-          log.Error().Err(closeErr).Msg("Failed to close connection")
-        }
-      }()
-
-      if err = dcomMmc.Init(ctx, &rpcClient); err != nil {
-        log.Error().Err(err).Msg("Module initialization failed")
-        returnCode = 1
-        return
-      }
-
-      if err = dcomMmc.Execute(ctx, exec.Input); err != nil {
-        log.Error().Err(err).Msg("Execution failed")
-        returnCode = 1
-      }
-    },
-  }
+		Args: args(
+			argsRpcClient("host"),
+			argsOutput("smb"),
+		),
+		Run: func(cmd *cobra.Command, args []string) {
+			dcomMmc.Dcom.Client = &rpcClient
+			dcomMmc.IO = exec
+
+			ctx := log.WithContext(gssapi.NewSecurityContext(context.TODO()))
+
+			if err := goexec.ExecuteCleanMethod(ctx, &dcomMmc, &exec); err != nil {
+				log.Fatal().Err(err).Msg("Operation failed")
+			}
+		},
+	}
 )
diff --git a/cmd/root.go b/cmd/root.go
index 3d696aa..913a44a 100644
--- a/cmd/root.go
+++ b/cmd/root.go
@@ -54,15 +54,17 @@ var (
 				log = log.Level(zerolog.DebugLevel)
 			}
 
-			if outputMethod == "smb" {
-				if exec.Output.RemotePath == "" {
-					exec.Output.RemotePath = util.RandomWindowsTempFile()
-				}
-				exec.Output.Provider = &smb.OutputFileFetcher{
-					Client:           &smbClient,
-					Share:            `C$`,
-					File:             exec.Output.RemotePath,
-					DeleteOutputFile: exec.Output.NoDelete, // TEMP
+			if outputPath != "" {
+				if outputMethod == "smb" {
+					if exec.Output.RemotePath == "" {
+						exec.Output.RemotePath = util.RandomWindowsTempFile()
+					}
+					exec.Output.Provider = &smb.OutputFileFetcher{
+						Client:           &smbClient,
+						Share:            `C$`,
+						File:             exec.Output.RemotePath,
+						DeleteOutputFile: !exec.Output.NoDelete,
+					}
 				}
 			}
 		},
@@ -81,13 +83,10 @@ func init() {
 
 		dcomCmdInit()
 		rootCmd.AddCommand(dcomCmd)
-
 		wmiCmdInit()
 		rootCmd.AddCommand(wmiCmd)
-
 		scmrCmdInit()
 		rootCmd.AddCommand(scmrCmd)
-
 		tschCmdInit()
 		rootCmd.AddCommand(tschCmd)
 	}
diff --git a/cmd/tsch.go b/cmd/tsch.go
index 328adcd..2e8370e 100644
--- a/cmd/tsch.go
+++ b/cmd/tsch.go
@@ -2,6 +2,7 @@ package cmd
 
 import (
 	"context"
+	"fmt"
 	"github.com/FalconOpsLLC/goexec/internal/util"
 	"github.com/FalconOpsLLC/goexec/pkg/goexec"
 	tschexec "github.com/FalconOpsLLC/goexec/pkg/goexec/tsch"
@@ -20,42 +21,19 @@ func tschCmdInit() {
 	tschCmd.AddCommand(tschCreateCmd)
 }
 
-func argsTschTaskIdentifiers(name, path string) error {
-	switch {
-	case path != "":
-		return tschexec.ValidateTaskPath(path)
-	case name != "":
-		return tschexec.ValidateTaskName(name)
-	default:
-	}
-	return nil
-}
-
-func argsTschDemand(_ *cobra.Command, _ []string) error {
-	return argsTschTaskIdentifiers(tschDemand.TaskName, tschDemand.TaskPath)
-}
-
-func argsTschCreate(_ *cobra.Command, _ []string) error {
-	return argsTschTaskIdentifiers(tschCreate.TaskName, tschCreate.TaskPath)
-}
-
 func tschDemandCmdInit() {
-	tschDemandCmd.Flags().StringVarP(&tschDemand.TaskName, "name", "t", "", "Name of task to register")
-	tschDemandCmd.Flags().StringVarP(&tschDemand.TaskPath, "path", "P", "", "Path of task to register")
+	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")
 
 	registerProcessExecutionArgs(tschDemandCmd)
 	registerExecutionOutputArgs(tschDemandCmd)
-
-	tschDemandCmd.MarkFlagsMutuallyExclusive("name", "path")
 }
 
 func tschCreateCmdInit() {
-	tschCreateCmd.Flags().StringVarP(&tschCreate.TaskName, "name", "t", "", "Name of task to register")
-	tschCreateCmd.Flags().StringVarP(&tschCreate.TaskPath, "path", "P", "", "Path of task to register")
-	tschCreateCmd.Flags().DurationVar(&tschCreate.StopDelay, "delay-stop", 5*time.Second, "Delay between task execution and termination. This will not stop the process spawned by the task")
+	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")
@@ -64,14 +42,27 @@ func tschCreateCmdInit() {
 
 	registerProcessExecutionArgs(tschCreateCmd)
 	registerExecutionOutputArgs(tschCreateCmd)
+}
 
-	tschCreateCmd.MarkFlagsMutuallyExclusive("name", "path")
+func argsTask(*cobra.Command, []string) error {
+	switch {
+	case tschTask == "":
+		tschTask = `\` + util.RandomString()
+	case tschexec.ValidateTaskPath(tschTask) == nil:
+	case tschexec.ValidateTaskName(tschTask) == nil:
+		tschTask = `\` + tschTask
+	default:
+		return fmt.Errorf("invalid task name or path: %q", tschTask)
+	}
+	return nil
 }
 
 var (
 	tschDemand tschexec.TschDemand
 	tschCreate tschexec.TschCreate
 
+	tschTask string
+
 	tschCmd = &cobra.Command{
 		Use:   "tsch",
 		Short: "Establish execution via Windows Task Scheduler (MS-TSCH)",
@@ -93,18 +84,18 @@ References:
 		Args: args(
 			argsRpcClient("cifs"),
 			argsOutput("smb"),
-			argsTschDemand,
+			argsTask,
 		),
 
-		Run: func(cmd *cobra.Command, args []string) {
-			tschDemand.Client = &rpcClient
+		Run: func(*cobra.Command, []string) {
 			tschDemand.IO = exec
+			tschDemand.Client = &rpcClient
+			tschDemand.TaskPath = tschTask
 
-			if tschDemand.TaskName == "" && tschDemand.TaskPath == "" {
-				tschDemand.TaskPath = `\` + util.RandomString()
-			}
-
-			ctx := log.WithContext(gssapi.NewSecurityContext(context.TODO()))
+			ctx := log.With().
+				Str("module", "tsch").
+				Str("method", "demand").
+				Logger().WithContext(gssapi.NewSecurityContext(context.TODO()))
 
 			if err := goexec.ExecuteCleanMethod(ctx, &tschDemand, &exec); err != nil {
 				log.Fatal().Err(err).Msg("Operation failed")
@@ -129,18 +120,17 @@ References:
 		Args: args(
 			argsRpcClient("cifs"),
 			argsOutput("smb"),
-			argsTschCreate,
+			argsTask,
 		),
 
-		Run: func(cmd *cobra.Command, args []string) {
+		Run: func(*cobra.Command, []string) {
 			tschCreate.Tsch.Client = &rpcClient
 			tschCreate.IO = exec
 
-			if tschCreate.TaskName == "" && tschDemand.TaskPath == "" {
-				tschCreate.TaskPath = `\` + util.RandomString()
-			}
-
-			ctx := log.WithContext(gssapi.NewSecurityContext(context.TODO()))
+			ctx := log.With().
+				Str("module", "tsch").
+				Str("method", "create").
+				Logger().WithContext(gssapi.NewSecurityContext(context.TODO()))
 
 			if err := goexec.ExecuteCleanMethod(ctx, &tschCreate, &exec); err != nil {
 				log.Fatal().Err(err).Msg("Operation failed")
diff --git a/cmd/wmi.go b/cmd/wmi.go
index 6fde395..42475f0 100644
--- a/cmd/wmi.go
+++ b/cmd/wmi.go
@@ -28,7 +28,7 @@ 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. {"CommandLine":"calc.exe"})`)
+	wmiCallCmd.Flags().StringVarP(&wmiArguments, "args", "A", "{}", `WMI Method argument(s) in JSON dictionary format (i.e. {"Command":"calc.exe"})`)
 
 	if err := wmiCallCmd.MarkFlagRequired("class"); err != nil {
 		panic(err)
@@ -131,7 +131,7 @@ References:
 			ctx := log.With().
 				Str("module", "wmi").
 				Str("method", "proc").
-				Logger().WithContext(gssapi.NewSecurityContext(context.TODO()))
+				Logger().WithContext(gssapi.NewSecurityContext(context.Background()))
 
 			if err := goexec.ExecuteCleanMethod(ctx, &wmiProc, &exec); err != nil {
 				log.Fatal().Err(err).Msg("Operation failed")