Merge pull request #3 from FalconOpsLLC/goexec#dev

2 files changed, 78 insertions, 7 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
diff --git a/cmd/root.go b/cmd/root.go
index 33b0116..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,30 @@ Authors: FalconOps LLC (@FalconOpsLLC),
 		},
 
 		PersistentPostRun: func(cmd *cobra.Command, args []string) {
-			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 err := exec.Input.StageFile.Close(); err != nil {
-				// ...
+
+			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 {
+					// ...
+				}
 			}
 		},
 	}
@@ -178,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