(feat) add --profile flag to dump runtime profiling data

Author: leonjza

.gitignore

@@ -25,3 +25,6 @@ screenshots/
 .glide/
 .DS_Store
 .idea/
+
+# profiles
+profiles/

cmd/root.go

@@ -3,6 +3,11 @@ package cmd
 import (
 	"fmt"
 	"os"
+	"path/filepath"
+	"runtime"
+	"runtime/pprof"
+	"runtime/trace"
+	"time"
 
 	"github.com/sensepost/gowitness/internal/ascii"
 	"github.com/sensepost/gowitness/pkg/log"
@@ -12,6 +17,13 @@ import (
 
 var (
 	opts = &runner.Options{}
+
+	// perf profiling
+	enableProfiling bool
+	profileDir      string
+
+	// hooks to run after command execution
+	postRunHooks []func()
 )
 
 var rootCmd = &cobra.Command{
@@ -28,8 +40,74 @@ var rootCmd = &cobra.Command{
 			log.Debug("debug logging enabled")
 		}
 
+		if enableProfiling {
+			ts := time.Now().Format("20060102-150405")
+			profileDir = filepath.Join("profiles", ts)
+
+			if err := os.MkdirAll(profileDir, 0o755); err != nil {
+				return fmt.Errorf("could not create profile directory %q: %w", profileDir, err)
+			}
+
+			cpuPath := filepath.Join(profileDir, "cpu.pprof")
+			memPath := filepath.Join(profileDir, "mem.pprof")
+			tracePath := filepath.Join(profileDir, "trace.out")
+
+			// cpu
+			cpuFile, err := os.Create(cpuPath)
+			if err != nil {
+				return fmt.Errorf("could not create CPU profile file: %w", err)
+			}
+			if err := pprof.StartCPUProfile(cpuFile); err != nil {
+				_ = cpuFile.Close()
+				return fmt.Errorf("could not start CPU profile: %w", err)
+			}
+			postRunHooks = append(postRunHooks, func() {
+				pprof.StopCPUProfile()
+				_ = cpuFile.Close()
+			})
+
+			// memory
+			postRunHooks = append(postRunHooks, func() {
+				memFile, err := os.Create(memPath)
+				if err != nil {
+					fmt.Fprintf(os.Stderr, "could not create memory profile file: %v\n", err)
+					return
+				}
+				defer memFile.Close()
+
+				runtime.GC() // refresh heap statistics
+
+				if err := pprof.WriteHeapProfile(memFile); err != nil {
+					fmt.Fprintf(os.Stderr, "could not write memory profile: %v\n", err)
+				}
+			})
+
+			// trace
+			traceFile, err := os.Create(tracePath)
+			if err != nil {
+				return fmt.Errorf("could not create trace file: %w", err)
+			}
+			if err := trace.Start(traceFile); err != nil {
+				_ = traceFile.Close()
+				return fmt.Errorf("could not start trace: %w", err)
+			}
+			postRunHooks = append(postRunHooks, func() {
+				trace.Stop()
+				_ = traceFile.Close()
+			})
+
+			// Log where results will be written
+			log.Info(fmt.Sprintf("profiling enabled: writing profiles to %s", profileDir))
+		}
+
 		return nil
 	},
+
+	PersistentPostRun: func(cmd *cobra.Command, args []string) {
+		for _, hook := range postRunHooks {
+			hook()
+		}
+	},
 }
 
 func Execute() {
@@ -61,4 +139,5 @@ func Execute() {
 func init() {
 	rootCmd.PersistentFlags().BoolVarP(&opts.Logging.Debug, "debug-log", "D", false, "Enable debug logging")
 	rootCmd.PersistentFlags().BoolVarP(&opts.Logging.Silence, "quiet", "q", false, "Silence (almost all) logging")
+	rootCmd.PersistentFlags().BoolVar(&enableProfiling, "profile", false, "Enable CPU, memory, and trace profiling (writes to profiles/<timestamp>/)")
 }