From 3c6d9d22445a274aad26005face44f171ba4465d Mon Sep 17 00:00:00 2001 From: Bryan McNulty Date: Fri, 25 Apr 2025 17:13:55 -0500 Subject: Avoid panic when closing log file --- cmd/root.go | 12 ++++++++---- 1 file changed, 8 insertions(+), 4 deletions(-) diff --git a/cmd/root.go b/cmd/root.go index 33b0116..d00ea2a 100644 --- a/cmd/root.go +++ b/cmd/root.go @@ -151,11 +151,15 @@ Authors: FalconOps LLC (@FalconOpsLLC), }, PersistentPostRun: func(cmd *cobra.Command, args []string) { - if err := logFile.Close(); err != nil { - // ... + if logFile != nil { + if err := logFile.Close(); err != nil { + // ... + } } - if err := exec.Input.StageFile.Close(); err != nil { - // ... + if exec.Input != nil && exec.Input.StageFile != nil { + if err := exec.Input.StageFile.Close(); err != nil { + // ... + } } }, } -- cgit v1.2.3 From 69564ac0f9a7cefc4231aa71e331c00fc4dac316 Mon Sep 17 00:00:00 2001 From: Bryan McNulty Date: Fri, 25 Apr 2025 18:58:06 -0500 Subject: Add hidden cpu & memory profiling flags --- cmd/root.go | 67 ++++++++++++++++++++++++++++++++++++++++++++++++++++++++++--- 1 file changed, 64 insertions(+), 3 deletions(-) diff --git a/cmd/root.go b/cmd/root.go index d00ea2a..96173a4 100644 --- a/cmd/root.go +++ b/cmd/root.go @@ -15,6 +15,7 @@ import ( "golang.org/x/term" "io" "os" + "runtime/pprof" ) type flagSet struct { @@ -56,6 +57,7 @@ var ( defaultAuthFlags, defaultLogFlags, defaultNetRpcFlags *flagSet returnCode int + toClose []io.Closer // === IO === stageFilePath string @@ -79,6 +81,13 @@ var ( smbClient smb.Client // =============== + // === Resource profiling === + cpuProfile string + memProfile string + cpuProfileFile io.WriteCloser + memProfileFile io.WriteCloser + // ========================== + exec = goexec.ExecutionIO{ Input: new(goexec.ExecutionInput), Output: new(goexec.ExecutionOutput), @@ -113,6 +122,7 @@ Authors: FalconOps LLC (@FalconOpsLLC), if err != nil { return } + toClose = append(toClose, logFile) logJson = true } if logQuiet { @@ -128,6 +138,29 @@ Authors: FalconOps LLC (@FalconOpsLLC), log = log.Level(logLevel) } + // CPU / memory profiling + { + if cpuProfile != "" { + if cpuProfileFile, err = os.OpenFile(cpuProfile, os.O_WRONLY|os.O_CREATE|os.O_TRUNC, 0644); err != nil { + log.Error().Err(err).Msg("Failed to open CPU profile for writing") + return + } + toClose = append(toClose, cpuProfileFile) + + if err = pprof.StartCPUProfile(cpuProfileFile); err != nil { + log.Error().Err(err).Msg("Failed to start CPU profile") + return + } + } + if memProfile != "" { + if memProfileFile, err = os.OpenFile(memProfile, os.O_WRONLY|os.O_CREATE|os.O_TRUNC, 0644); err != nil { + log.Error().Err(err).Msg("Failed to open memory profile for writing") + return + } + toClose = append(toClose, memProfileFile) + } + } + if proxy != "" { rpcClient.Proxy = proxy smbClient.Proxy = proxy @@ -151,11 +184,26 @@ Authors: FalconOps LLC (@FalconOpsLLC), }, PersistentPostRun: func(cmd *cobra.Command, args []string) { - if logFile != nil { - if err := logFile.Close(); err != nil { - // ... + + if memProfileFile != nil { + if err := pprof.WriteHeapProfile(memProfileFile); err != nil { + log.Error().Err(err).Msg("Failed to write memory profile") + return + } + } + + if cpuProfileFile != nil { + pprof.StopCPUProfile() + } + + for _, c := range toClose { + if c != nil { + if err := c.Close(); err != nil { + log.Warn().Err(err).Msg("Failed to close stream") + } } } + if exec.Input != nil && exec.Input.StageFile != nil { if err := exec.Input.StageFile.Close(); err != nil { // ... @@ -182,6 +230,19 @@ func init() { gssapi.AddMechanism(ssp.KRB5) } + // CPU / Memory profiling + { + rootCmd.PersistentFlags().StringVar(&cpuProfile, "cpu-profile", "", "Write CPU profile to `file`") + rootCmd.PersistentFlags().StringVar(&memProfile, "mem-profile", "", "Write memory profile to `file`") + + if err := rootCmd.PersistentFlags().MarkHidden("cpu-profile"); err != nil { + panic(err) + } + if err := rootCmd.PersistentFlags().MarkHidden("mem-profile"); err != nil { + panic(err) + } + } + // Cobra init { cobra.EnableCommandSorting = false -- cgit v1.2.3 From 8fba4225d1f1cc1fc1b37cc8749044d6be670937 Mon Sep 17 00:00:00 2001 From: Bryan McNulty Date: Fri, 25 Apr 2025 19:01:22 -0500 Subject: Update TODO --- TODO.md | 12 +++++++++--- 1 file changed, 9 insertions(+), 3 deletions(-) diff --git a/TODO.md b/TODO.md index 7afd1da..7972808 100644 --- a/TODO.md +++ b/TODO.md @@ -1,5 +1,7 @@ # TODO +We wanted to make development of this project as transparent as possible, so we've included our development TODOs in this file. + ## TSCH - [X] Clean up TSCH module @@ -8,6 +10,7 @@ - [X] Output - [X] Add `tsch change` - [ ] Serialize XML with default indent level +- [ ] Add --session to `tsch change` ## SCMR @@ -33,15 +36,18 @@ ## Other - [X] Add proxy support - see https://github.com/oiweiwei/go-msrpc/issues/21 -- [ ] Descriptions for all modules and methods +- [X] README +- [X] pprof integration (hidden flag(s)) +- [X] Descriptions for all modules and methods - [ ] Add SMB file transfer interface -- [ ] README -## Bug Fixes +## Bugs - [X] (Fixed) SMB transport for SCMR module - `rpc_s_cannot_support: The requested operation is not supported.` - [X] (Fixed) Proxy - EPM doesn't use the proxy dialer - [X] (Fixed) Kerberos requests don't dial through proxy +- [X] (Fixed) Panic when closing nil log file +- [ ] `scmr change` doesn't revert service cmdline - [ ] Fix SCMR `change` method so that dependencies field isn't permanently overwritten ## Lower Priority -- cgit v1.2.3