From e1a0babffe41a080a58a2174bfa311d81d0c7258 Mon Sep 17 00:00:00 2001 From: Bryan McNulty Date: Fri, 12 Sep 2025 10:18:38 -0500 Subject: [PATCH] Fix context linter err; fix indents --- cmd/args.go | 228 +++++++++++++------------- cmd/tsch.go | 336 +++++++++++++++++++-------------------- pkg/goexec/context.go | 8 + pkg/goexec/io.go | 120 +++++++------- pkg/goexec/method.go | 162 +++++++++---------- pkg/goexec/smb/output.go | 159 +++++++++--------- 6 files changed, 512 insertions(+), 501 deletions(-) create mode 100644 pkg/goexec/context.go diff --git a/cmd/args.go b/cmd/args.go index 7460e20..ada1190 100644 --- a/cmd/args.go +++ b/cmd/args.go @@ -1,43 +1,43 @@ package cmd import ( - "context" - "encoding/json" - "errors" - "fmt" - "os" - "time" + "context" + "encoding/json" + "errors" + "fmt" + "os" + "time" - "github.com/spf13/cobra" - "github.com/spf13/pflag" + "github.com/spf13/cobra" + "github.com/spf13/pflag" ) 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") + 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 registerNetworkFlags(fs *pflag.FlagSet) { - fs.StringVarP(&proxy, "proxy", "x", "", "Proxy `URI`") - fs.StringVarP(&rpcClient.Filter, "epm-filter", "F", "", "String binding to filter endpoints returned by the RPC endpoint mapper (EPM)") - fs.StringVar(&rpcClient.Endpoint, "endpoint", "", "Explicit RPC endpoint definition") - fs.BoolVar(&rpcClient.NoEpm, "no-epm", false, "Don't use EPM to discover RPC endpoints") - fs.BoolVar(&rpcClient.UseEpm, "epm", false, "Use EPM to discover available bindings") - 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") + fs.StringVarP(&proxy, "proxy", "x", "", "Proxy `URI`") + fs.StringVarP(&rpcClient.Filter, "epm-filter", "F", "", "String binding to filter endpoints returned by the RPC endpoint mapper (EPM)") + fs.StringVar(&rpcClient.Endpoint, "endpoint", "", "Explicit RPC endpoint definition") + fs.BoolVar(&rpcClient.NoEpm, "no-epm", false, "Don't use EPM to discover RPC endpoints") + fs.BoolVar(&rpcClient.UseEpm, "epm", false, "Use EPM to discover available bindings") + 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") - if err := fs.MarkHidden("no-epm"); err != nil { - panic(err) - } - if err := fs.MarkDeprecated("no-epm", "use --epm=false instead"); err != nil { - panic(err) - } + if err := fs.MarkHidden("no-epm"); err != nil { + panic(err) + } + if err := fs.MarkDeprecated("no-epm", "use --epm=false instead"); err != nil { + panic(err) + } - //cmd.MarkFlagsMutuallyExclusive("endpoint", "epm-filter") - //cmd.MarkFlagsMutuallyExclusive("no-epm", "epm-filter") + //cmd.MarkFlagsMutuallyExclusive("endpoint", "epm-filter") + //cmd.MarkFlagsMutuallyExclusive("no-epm", "epm-filter") } // FUTURE: automatically stage & execute file @@ -49,127 +49,127 @@ func registerStageFlags(fs *pflag.FlagSet) { */ func registerExecutionFlags(fs *pflag.FlagSet) { - fs.StringVarP(&exec.Input.Executable, "exec", "e", "", "Remote 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)") + fs.StringVarP(&exec.Input.Executable, "exec", "e", "", "Remote 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)") - //cmd.MarkFlagsOneRequired("executable", "command") - //cmd.MarkFlagsMutuallyExclusive("executable", "command") + //cmd.MarkFlagsOneRequired("executable", "command") + //cmd.MarkFlagsMutuallyExclusive("executable", "command") } 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.DurationVar(&exec.Output.Timeout, "out-timeout", time.Second*60, "Output timeout `duration`") - //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") + 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.DurationVar(&exec.Output.Timeout, "out-timeout", time.Second*60, "Output timeout `duration`") + //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 argsAcceptValues(name string, in *string, valid ...string) func(*cobra.Command, []string) error { - return func(*cobra.Command, []string) error { - for _, v := range valid { - if *in == v { - return nil - } - } - if j, err := json.Marshal(valid); err == nil { - return fmt.Errorf("parse %s: %q doesn't match any accepted values: %s", name, *in, string(j)) - } else { - return err - } - } + return func(*cobra.Command, []string) error { + for _, v := range valid { + if *in == v { + return nil + } + } + if j, err := json.Marshal(valid); err == nil { + return fmt.Errorf("parse %s: %q doesn't match any accepted values: %s", name, *in, string(j)) + } else { + return err + } + } } 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, endpoint string) func(cmd *cobra.Command, args []string) error { - return args( - argsTarget(proto), + return args( + argsTarget(proto), - func(cmd *cobra.Command, args []string) (err error) { - switch { - case rpcClient.Endpoint != "": - case endpoint == "": - rpcClient.UseEpm = true - default: - rpcClient.Endpoint = endpoint - } - rpcClient.Target = target - rpcClient.Credential = credential - rpcClient.Proxy = proxy + func(cmd *cobra.Command, args []string) (err error) { + switch { + case rpcClient.Endpoint != "": + case endpoint == "": + rpcClient.UseEpm = true + default: + rpcClient.Endpoint = endpoint + } + 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 + })...) } diff --git a/cmd/tsch.go b/cmd/tsch.go index ae4500c..a9b0d71 100644 --- a/cmd/tsch.go +++ b/cmd/tsch.go @@ -1,232 +1,232 @@ package cmd import ( - "context" - "fmt" - "time" + "context" + "fmt" + "time" - "github.com/FalconOpsLLC/goexec/internal/util" - "github.com/FalconOpsLLC/goexec/pkg/goexec" - tschexec "github.com/FalconOpsLLC/goexec/pkg/goexec/tsch" - "github.com/oiweiwei/go-msrpc/ssp/gssapi" - "github.com/spf13/cobra" + "github.com/FalconOpsLLC/goexec/internal/util" + "github.com/FalconOpsLLC/goexec/pkg/goexec" + tschexec "github.com/FalconOpsLLC/goexec/pkg/goexec/tsch" + "github.com/oiweiwei/go-msrpc/ssp/gssapi" + "github.com/spf13/cobra" ) func tschCmdInit() { - cmdFlags[tschCmd] = []*flagSet{ - defaultAuthFlags, - defaultLogFlags, - defaultNetRpcFlags, - } - tschDemandCmdInit() - tschCreateCmdInit() - tschChangeCmdInit() + cmdFlags[tschCmd] = []*flagSet{ + defaultAuthFlags, + defaultLogFlags, + defaultNetRpcFlags, + } + tschDemandCmdInit() + tschCreateCmdInit() + tschChangeCmdInit() - tschCmd.PersistentFlags().AddFlagSet(defaultAuthFlags.Flags) - tschCmd.PersistentFlags().AddFlagSet(defaultLogFlags.Flags) - tschCmd.PersistentFlags().AddFlagSet(defaultNetRpcFlags.Flags) - tschCmd.AddCommand(tschDemandCmd, tschCreateCmd, tschChangeCmd) + tschCmd.PersistentFlags().AddFlagSet(defaultAuthFlags.Flags) + tschCmd.PersistentFlags().AddFlagSet(defaultLogFlags.Flags) + tschCmd.PersistentFlags().AddFlagSet(defaultNetRpcFlags.Flags) + tschCmd.AddCommand(tschDemandCmd, tschCreateCmd, tschChangeCmd) } func tschDemandCmdInit() { - tschDemandFlags := newFlagSet("Task Scheduler") + tschDemandFlags := newFlagSet("Task Scheduler") - 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") + 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") + tschDemandExecFlags := newFlagSet("Execution") - registerExecutionFlags(tschDemandExecFlags.Flags) - registerExecutionOutputFlags(tschDemandExecFlags.Flags) + registerExecutionFlags(tschDemandExecFlags.Flags) + registerExecutionOutputFlags(tschDemandExecFlags.Flags) - cmdFlags[tschDemandCmd] = []*flagSet{ - tschDemandFlags, - tschDemandExecFlags, - defaultAuthFlags, - defaultLogFlags, - defaultNetRpcFlags, - } + cmdFlags[tschDemandCmd] = []*flagSet{ + tschDemandFlags, + tschDemandExecFlags, + defaultAuthFlags, + defaultLogFlags, + defaultNetRpcFlags, + } - tschDemandCmd.Flags().AddFlagSet(tschDemandFlags.Flags) - tschDemandCmd.Flags().AddFlagSet(tschDemandExecFlags.Flags) - tschDemandCmd.MarkFlagsOneRequired("exec", "command") + tschDemandCmd.Flags().AddFlagSet(tschDemandFlags.Flags) + tschDemandCmd.Flags().AddFlagSet(tschDemandExecFlags.Flags) + tschDemandCmd.MarkFlagsOneRequired("exec", "command") } func tschCreateCmdInit() { - tschCreateFlags := newFlagSet("Task Scheduler") + tschCreateFlags := newFlagSet("Task Scheduler") - 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") + 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") + tschCreateExecFlags := newFlagSet("Execution") - registerExecutionFlags(tschCreateExecFlags.Flags) - registerExecutionOutputFlags(tschCreateExecFlags.Flags) + registerExecutionFlags(tschCreateExecFlags.Flags) + registerExecutionOutputFlags(tschCreateExecFlags.Flags) - cmdFlags[tschCreateCmd] = []*flagSet{ - tschCreateFlags, - tschCreateExecFlags, - defaultAuthFlags, - defaultLogFlags, - defaultNetRpcFlags, - } + cmdFlags[tschCreateCmd] = []*flagSet{ + tschCreateFlags, + tschCreateExecFlags, + defaultAuthFlags, + defaultLogFlags, + defaultNetRpcFlags, + } - tschCreateCmd.Flags().AddFlagSet(tschCreateFlags.Flags) - tschCreateCmd.Flags().AddFlagSet(tschCreateExecFlags.Flags) - tschCreateCmd.MarkFlagsOneRequired("exec", "command") + tschCreateCmd.Flags().AddFlagSet(tschCreateFlags.Flags) + tschCreateCmd.Flags().AddFlagSet(tschCreateExecFlags.Flags) + tschCreateCmd.MarkFlagsOneRequired("exec", "command") } func tschChangeCmdInit() { - tschChangeFlags := newFlagSet("Task Scheduler") + tschChangeFlags := newFlagSet("Task Scheduler") - tschChangeFlags.Flags.StringVarP(&tschChange.TaskPath, "task", "t", "", "Path to existing task") - tschChangeFlags.Flags.BoolVar(&tschChange.NoStart, "no-start", false, "Don't start the task") - tschChangeFlags.Flags.BoolVar(&tschChange.NoRevert, "no-revert", false, "Don't restore the original task definition") + tschChangeFlags.Flags.StringVarP(&tschChange.TaskPath, "task", "t", "", "Path to existing task") + tschChangeFlags.Flags.BoolVar(&tschChange.NoStart, "no-start", false, "Don't start the task") + tschChangeFlags.Flags.BoolVar(&tschChange.NoRevert, "no-revert", false, "Don't restore the original task definition") - tschChangeExecFlags := newFlagSet("Execution") + tschChangeExecFlags := newFlagSet("Execution") - registerExecutionFlags(tschChangeExecFlags.Flags) - registerExecutionOutputFlags(tschChangeExecFlags.Flags) + registerExecutionFlags(tschChangeExecFlags.Flags) + registerExecutionOutputFlags(tschChangeExecFlags.Flags) - cmdFlags[tschChangeCmd] = []*flagSet{ - tschChangeFlags, - tschChangeExecFlags, - defaultAuthFlags, - defaultLogFlags, - defaultNetRpcFlags, - } + cmdFlags[tschChangeCmd] = []*flagSet{ + tschChangeFlags, + tschChangeExecFlags, + defaultAuthFlags, + defaultLogFlags, + defaultNetRpcFlags, + } - tschChangeCmd.Flags().AddFlagSet(tschChangeFlags.Flags) - tschChangeCmd.Flags().AddFlagSet(tschChangeExecFlags.Flags) + tschChangeCmd.Flags().AddFlagSet(tschChangeFlags.Flags) + tschChangeCmd.Flags().AddFlagSet(tschChangeExecFlags.Flags) - // Constraints - { - if err := tschChangeCmd.MarkFlagRequired("task"); err != nil { - panic(err) - } - tschChangeCmd.MarkFlagsOneRequired("exec", "command") - } + // Constraints + { + if err := tschChangeCmd.MarkFlagRequired("task"); err != nil { + panic(err) + } + tschChangeCmd.MarkFlagsOneRequired("exec", "command") + } } 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 Label or path: %q", tschTask) - } - return nil + switch { + case tschTask == "": + tschTask = `\` + util.RandomString() + case tschexec.ValidateTaskPath(tschTask) == nil: + case tschexec.ValidateTaskName(tschTask) == nil: + tschTask = `\` + tschTask + default: + return fmt.Errorf("invalid task Label or path: %q", tschTask) + } + return nil } var ( - tschDemand tschexec.TschDemand - tschCreate tschexec.TschCreate - tschChange tschexec.TschChange + tschDemand tschexec.TschDemand + tschCreate tschexec.TschCreate + tschChange tschexec.TschChange - tschTask string + tschTask string - tschCmd = &cobra.Command{ - Use: "tsch", - Short: "Execute with Windows Task Scheduler (MS-TSCH)", - Long: `Description: + tschCmd = &cobra.Command{ + Use: "tsch", + Short: "Execute with Windows Task Scheduler (MS-TSCH)", + Long: `Description: The tsch module makes use of the Windows Task Scheduler service (MS-TSCH) to spawn processes on the remote target.`, - GroupID: "module", - Args: cobra.NoArgs, - } + GroupID: "module", + Args: cobra.NoArgs, + } - tschDemandCmd = &cobra.Command{ - Use: "demand [target]", - Short: "Register a remote scheduled task and demand immediate start", - Long: `Description: + tschDemandCmd = &cobra.Command{ + Use: "demand [target]", + Short: "Register a remote scheduled task and demand immediate start", + Long: `Description: Similar to the create method, the demand method will call SchRpcRegisterTask, But rather than setting a defined time when the task will start, it will additionally call SchRpcRun to forcefully start the task.`, - Args: args( - argsRpcClient("cifs", "ncacn_np:[atsvc]"), - argsOutput("smb"), - argsTask, - ), + Args: args( + argsRpcClient("cifs", "ncacn_np:[atsvc]"), + argsOutput("smb"), + argsTask, + ), - Run: func(*cobra.Command, []string) { - tschDemand.Client = &rpcClient - tschDemand.TaskPath = tschTask + Run: func(*cobra.Command, []string) { + tschDemand.Client = &rpcClient + tschDemand.TaskPath = tschTask - ctx := log.With(). - Str("module", "tsch"). - Str("method", "demand"). - Logger().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") - } - }, - } - tschCreateCmd = &cobra.Command{ - Use: "create [target]", - Short: "Create a remote scheduled task with an automatic start time", - Long: `Description: + if err := goexec.ExecuteCleanMethod(ctx, &tschDemand, &exec); err != nil { + log.Fatal().Err(err).Msg("Operation failed") + } + }, + } + tschCreateCmd = &cobra.Command{ + Use: "create [target]", + Short: "Create a remote scheduled task with an automatic start time", + Long: `Description: The create method calls SchRpcRegisterTask to register a scheduled task with an automatic start time.This method avoids directly calling SchRpcRun, and can even avoid calling SchRpcDelete by populating the DeleteExpiredTaskAfter Setting.`, - Args: args( - argsRpcClient("cifs", "ncacn_np:[atsvc]"), - argsOutput("smb"), - argsTask, - ), + Args: args( + argsRpcClient("cifs", "ncacn_np:[atsvc]"), + argsOutput("smb"), + argsTask, + ), - Run: func(*cobra.Command, []string) { - tschCreate.Client = &rpcClient - tschCreate.TaskPath = tschTask + Run: func(*cobra.Command, []string) { + tschCreate.Client = &rpcClient + tschCreate.TaskPath = tschTask - ctx := log.With(). - Str("module", "tsch"). - Str("method", "create"). - Logger().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") - } - }, - } - tschChangeCmd = &cobra.Command{ - Use: "change [target]", - Short: "Modify an existing task to spawn an arbitrary process", - Long: `Description: + if err := goexec.ExecuteCleanMethod(ctx, &tschCreate, &exec); err != nil { + log.Fatal().Err(err).Msg("Operation failed") + } + }, + } + tschChangeCmd = &cobra.Command{ + Use: "change [target]", + Short: "Modify an existing task to spawn an arbitrary process", + Long: `Description: The change method calls SchRpcRetrieveTask to fetch the definition of an existing task (-t), then modifies the task definition to spawn a process`, - Args: args( - argsRpcClient("cifs", "ncacn_np:[atsvc]"), - argsOutput("smb"), + Args: args( + argsRpcClient("cifs", "ncacn_np:[atsvc]"), + argsOutput("smb"), - func(*cobra.Command, []string) error { - return tschexec.ValidateTaskPath(tschChange.TaskPath) - }, - ), + func(*cobra.Command, []string) error { + return tschexec.ValidateTaskPath(tschChange.TaskPath) + }, + ), - Run: func(*cobra.Command, []string) { - tschChange.Client = &rpcClient - tschChange.IO = exec + Run: func(*cobra.Command, []string) { + tschChange.Client = &rpcClient + tschChange.IO = exec - ctx := log.With(). - Str("module", "tsch"). - Str("method", "change"). - Logger().WithContext(gssapi.NewSecurityContext(context.TODO())) + ctx := log.With(). + Str("module", "tsch"). + Str("method", "change"). + Logger().WithContext(gssapi.NewSecurityContext(context.TODO())) - if err := goexec.ExecuteCleanMethod(ctx, &tschChange, &exec); err != nil { - log.Fatal().Err(err).Msg("Operation failed") - } - }, - } + if err := goexec.ExecuteCleanMethod(ctx, &tschChange, &exec); err != nil { + log.Fatal().Err(err).Msg("Operation failed") + } + }, + } ) diff --git a/pkg/goexec/context.go b/pkg/goexec/context.go new file mode 100644 index 0000000..7061acb --- /dev/null +++ b/pkg/goexec/context.go @@ -0,0 +1,8 @@ +package goexec + +type ContextOption string + +const ( + ContextOptionOutputTimeout ContextOption = "output.timeout" + ContextOptionOutputPollInterval ContextOption = "output.pollInterval" +) diff --git a/pkg/goexec/io.go b/pkg/goexec/io.go index 67cf886..0b24eb5 100644 --- a/pkg/goexec/io.go +++ b/pkg/goexec/io.go @@ -1,100 +1,100 @@ package goexec import ( - "context" - "fmt" - "io" - "strings" - "time" + "context" + "fmt" + "io" + "strings" + "time" ) type OutputProvider interface { - GetOutput(ctx context.Context, writer io.Writer) (err error) - Clean(ctx context.Context) (err error) + GetOutput(ctx context.Context, writer io.Writer) (err error) + Clean(ctx context.Context) (err error) } type ExecutionIO struct { - Cleaner + Cleaner - Input *ExecutionInput - Output *ExecutionOutput + Input *ExecutionInput + Output *ExecutionOutput } type ExecutionOutput struct { - NoDelete bool - RemotePath string - Timeout time.Duration - Provider OutputProvider - Writer io.WriteCloser + NoDelete bool + RemotePath string + Timeout time.Duration + Provider OutputProvider + Writer io.WriteCloser } type ExecutionInput struct { - StageFile io.ReadCloser - Executable string - ExecutablePath string - Arguments string - Command string + StageFile io.ReadCloser + Executable string + ExecutablePath string + Arguments string + Command string } func (execIO *ExecutionIO) GetOutput(ctx context.Context) (err error) { - if execIO.Output.Provider != nil { - ctx = context.WithValue(ctx, "output.timeout", execIO.Output.Timeout) - return execIO.Output.Provider.GetOutput(ctx, execIO.Output.Writer) - } - return nil + if execIO.Output.Provider != nil { + ctx = context.WithValue(ctx, ContextOptionOutputTimeout, execIO.Output.Timeout) + return execIO.Output.Provider.GetOutput(ctx, execIO.Output.Writer) + } + return nil } func (execIO *ExecutionIO) Clean(ctx context.Context) (err error) { - if execIO.Output.Provider != nil { - return execIO.Output.Provider.Clean(ctx) - } - return nil + 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{ - `C:\Windows\System32\cmd.exe`, - fmt.Sprintf(`/C %s > %s 2>&1`, execIO.Input.String(), execIO.Output.RemotePath), - } - } - return execIO.Input.CommandLine() + if execIO.Output.Provider != nil && execIO.Output.RemotePath != "" { + return []string{ + `C:\Windows\System32\cmd.exe`, + fmt.Sprintf(`/C %s > %s 2>&1`, execIO.Input.String(), execIO.Output.RemotePath), + } + } + return execIO.Input.CommandLine() } func (execIO *ExecutionIO) String() (str string) { - cmd := execIO.CommandLine() - // Ensure that executable paths are quoted - if strings.Contains(cmd[0], " ") { - str = fmt.Sprintf(`%q %s`, cmd[0], strings.Join(cmd[1:], " ")) - } else { - str = strings.Join(cmd, " ") - } - return strings.Trim(str, " \t\n\r") // trim whitespace + cmd := execIO.CommandLine() + // Ensure that executable paths are quoted + if strings.Contains(cmd[0], " ") { + str = fmt.Sprintf(`%q %s`, cmd[0], strings.Join(cmd[1:], " ")) + } else { + str = strings.Join(cmd, " ") + } + return strings.Trim(str, " \t\n\r") // trim whitespace } func (i *ExecutionInput) CommandLine() (cmd []string) { - cmd = make([]string, 2) - cmd[1] = i.Arguments + cmd = make([]string, 2) + cmd[1] = i.Arguments - switch { - case i.Command != "": - copy(cmd, strings.SplitN(i.Command, " ", 2)) - case i.ExecutablePath != "": - cmd[0] = i.ExecutablePath - case i.Executable != "": - cmd[0] = i.Executable - } + switch { + case i.Command != "": + copy(cmd, strings.SplitN(i.Command, " ", 2)) + case i.ExecutablePath != "": + cmd[0] = i.ExecutablePath + case i.Executable != "": + cmd[0] = i.Executable + } - return cmd + return cmd } func (i *ExecutionInput) String() string { - return strings.Join(i.CommandLine(), " ") + return strings.Join(i.CommandLine(), " ") } func (i *ExecutionInput) Reader() (reader io.Reader) { - if i.StageFile != nil { - return i.StageFile - } - return strings.NewReader(i.String()) + if i.StageFile != nil { + return i.StageFile + } + return strings.NewReader(i.String()) } diff --git a/pkg/goexec/method.go b/pkg/goexec/method.go index ea2e90f..7fa6d94 100644 --- a/pkg/goexec/method.go +++ b/pkg/goexec/method.go @@ -1,131 +1,131 @@ package goexec import ( - "context" - "fmt" + "context" + "fmt" - "github.com/rs/zerolog" + "github.com/rs/zerolog" ) type Method interface { - Connect(ctx context.Context) error - Init(ctx context.Context) error + Connect(ctx context.Context) error + Init(ctx context.Context) error } type CleanMethod interface { - Method - Clean + Method + Clean } type ExecutionMethod interface { - Method - Execute(ctx context.Context, io *ExecutionIO) error + Method + Execute(ctx context.Context, io *ExecutionIO) error } type CleanExecutionMethod interface { - ExecutionMethod - Clean + ExecutionMethod + Clean } type AuxiliaryMethod interface { - Method - Call(ctx context.Context) error + Method + Call(ctx context.Context) error } type CleanAuxiliaryMethod interface { - AuxiliaryMethod - Clean + AuxiliaryMethod + Clean } func ExecuteMethod(ctx context.Context, module ExecutionMethod, execIO *ExecutionIO) (err error) { - log := zerolog.Ctx(ctx) + 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.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.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") - return fmt.Errorf("execute: %w", err) - } + if err = module.Execute(ctx, execIO); err != nil { + log.Error().Err(err).Msg("Execution failed") + return fmt.Errorf("execute: %w", err) + } - return + return } func ExecuteAuxiliaryMethod(ctx context.Context, module AuxiliaryMethod) (err error) { - log := zerolog.Ctx(ctx) + 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.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.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") + 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 + return nil } func ExecuteCleanAuxiliaryMethod(ctx context.Context, module CleanAuxiliaryMethod) (err error) { - log := zerolog.Ctx(ctx) + log := zerolog.Ctx(ctx) - defer func() { - if err = module.Clean(ctx); err != nil { - log.Error().Err(err).Msg("Module cleanup failed") - err = nil - } - }() + 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 { - return fmt.Errorf("execute auxiliary method: %w", err) - } - return + if err = ExecuteAuxiliaryMethod(ctx, module); err != nil { + return fmt.Errorf("execute auxiliary method: %w", err) + } + return } func ExecuteCleanMethod(ctx context.Context, module CleanExecutionMethod, execIO *ExecutionIO) (err error) { - log := zerolog.Ctx(ctx) + log := zerolog.Ctx(ctx) - if err = ExecuteMethod(ctx, module, execIO); err != nil { - return - } + if err = ExecuteMethod(ctx, module, execIO); err != nil { + return + } - if err = module.Clean(ctx); err != nil { - log.Error().Err(err).Msg("Module cleanup failed") - err = nil - } + if err = module.Clean(ctx); err != nil { + log.Error().Err(err).Msg("Module cleanup failed") + err = nil + } - if execIO.Output != nil && execIO.Output.Provider != nil { - log.Info().Msg("Collecting output") + if execIO.Output != nil && execIO.Output.Provider != nil { + log.Info().Msg("Collecting output") - defer func() { - if cleanErr := execIO.Clean(ctx); cleanErr != nil { - log.Debug().Err(cleanErr).Msg("Output provider cleanup failed") - } - }() + defer func() { + if cleanErr := execIO.Clean(ctx); cleanErr != nil { + log.Debug().Err(cleanErr).Msg("Output provider cleanup failed") + } + }() - if err := execIO.GetOutput(ctx); err != nil { - log.Error().Err(err).Msg("Output collection failed") - return fmt.Errorf("get output: %w", err) - } - log.Debug().Msg("Output collection succeeded") - } - return + if err := execIO.GetOutput(ctx); err != nil { + log.Error().Err(err).Msg("Output collection failed") + return fmt.Errorf("get output: %w", err) + } + log.Debug().Msg("Output collection succeeded") + } + return } diff --git a/pkg/goexec/smb/output.go b/pkg/goexec/smb/output.go index 7275dbd..cfc4bd9 100644 --- a/pkg/goexec/smb/output.go +++ b/pkg/goexec/smb/output.go @@ -1,102 +1,105 @@ package smb import ( - "context" - "errors" - "io" - "os" - "path/filepath" - "regexp" - "strings" - "time" + "context" + "errors" + "io" + "os" + "path/filepath" + "regexp" + "strings" + "time" - "github.com/FalconOpsLLC/goexec/pkg/goexec" - "github.com/rs/zerolog" + "github.com/FalconOpsLLC/goexec/pkg/goexec" + "github.com/rs/zerolog" ) var ( - DefaultOutputPollInterval = 500 * time.Millisecond - DefaultOutputPollTimeout = 60 * time.Second - pathPrefix = regexp.MustCompile(`^([a-zA-Z]:)?[\\/]*`) + DefaultOutputPollInterval = 500 * time.Millisecond + DefaultOutputPollTimeout = 60 * time.Second + pathPrefix = regexp.MustCompile(`^([a-zA-Z]:)?[\\/]*`) ) type OutputFileFetcher struct { - goexec.Cleaner + goexec.Cleaner - Client *Client + Client *Client - Share string - SharePath string - File string - DeleteOutputFile bool - ForceReconnect bool + Share string + SharePath string + File string + DeleteOutputFile bool + ForceReconnect bool - relativePath string + relativePath string } func (o *OutputFileFetcher) GetOutput(ctx context.Context, writer io.Writer) (err error) { - log := zerolog.Ctx(ctx) - timeout := DefaultOutputPollTimeout - pollInterval := DefaultOutputPollInterval + log := zerolog.Ctx(ctx) + timeout := DefaultOutputPollTimeout + pollInterval := DefaultOutputPollInterval - if v := ctx.Value("output.timeout"); v != nil { - if t, ok := v.(time.Duration); ok { - timeout = t - } - } - if v := ctx.Value("output.pollInterval"); v != nil { - if p, ok := v.(time.Duration); ok { - pollInterval = p - } - } - shp := pathPrefix.ReplaceAllString(strings.ToLower(strings.ReplaceAll(o.SharePath, `\`, "/")), "") - fp := pathPrefix.ReplaceAllString(strings.ToLower(strings.ReplaceAll(o.File, `\`, "/")), "") + if v := ctx.Value(goexec.ContextOptionOutputTimeout); v != nil { + if t, ok := v.(time.Duration); ok { + timeout = t + } + } + if v := ctx.Value(goexec.ContextOptionOutputPollInterval); v != nil { + if p, ok := v.(time.Duration); ok { + pollInterval = p + } + } + shp := pathPrefix.ReplaceAllString(strings.ToLower(strings.ReplaceAll(o.SharePath, `\`, "/")), "") + fp := pathPrefix.ReplaceAllString(strings.ToLower(strings.ReplaceAll(o.File, `\`, "/")), "") - if o.relativePath, err = filepath.Rel(shp, fp); err != nil { - return - } + if o.relativePath, err = filepath.Rel(shp, fp); err != nil { + return + } - log.Info().Str("path", o.relativePath).Msg("Fetching output file") + log.Info().Str("path", o.relativePath).Msg("Fetching output file") - if o.ForceReconnect || !o.Client.connected { - err = o.Client.Connect(ctx) - if err != nil { - return - } - defer o.AddCleaners(o.Client.Close) - } + if o.ForceReconnect || !o.Client.connected { + err = o.Client.Connect(ctx) + if err != nil { + return + } + defer o.AddCleaners(o.Client.Close) + } - if o.ForceReconnect || o.Client.share != o.Share { - err = o.Client.Mount(ctx, o.Share) - if err != nil { - return - } - } + if o.ForceReconnect || o.Client.share != o.Share { + err = o.Client.Mount(ctx, o.Share) + if err != nil { + return + } + } - stopAt := time.Now().Add(timeout) - var reader io.ReadCloser + if reader, err := func() (io.ReadCloser, error) { + for { + select { + case <-ctx.Done(): + return nil, ctx.Err() + case <-time.After(timeout): + return nil, errors.New("execution output timeout") + case <-time.After(pollInterval): + if reader, err := o.Client.mount.OpenFile(o.relativePath, os.O_RDWR, 0); err == nil { + return reader, err + } + } + } + }(); err != nil { + return err + } else { + o.AddCleaners(func(_ context.Context) error { + return reader.Close() + }) + if _, err := io.Copy(writer, reader); err != nil { + if o.DeleteOutputFile { + o.AddCleaners(func(_ context.Context) error { + return o.Client.mount.Remove(o.relativePath) + }) + } + } + } - for { - if time.Now().After(stopAt) { - return errors.New("execution output timeout") - } - if reader, err = o.Client.mount.OpenFile(o.relativePath, os.O_RDWR, 0); err == nil { - break - } - time.Sleep(pollInterval) - } - - if _, err = io.Copy(writer, reader); err != nil { - return - } - - o.AddCleaners(func(_ context.Context) error { return reader.Close() }) - - if o.DeleteOutputFile { - o.AddCleaners(func(_ context.Context) error { - return o.Client.mount.Remove(o.relativePath) - }) - } - - return + return }