new: cpuprofile and memprofile options for easy profiling

.gitignore

@@ -1,5 +1,6 @@
 *.sw*
 *.tar.gz
+*.prof*
 pcaps
 bettercap-ng*.*
 bettercap-ng*

core/options.go

@@ -9,6 +9,8 @@ type Options struct {
 	Silent        *bool
 	NoHistory     *bool
 	Commands      *string
+	CpuProfile    *string
+	MemProfile    *string
 }
 
 func ParseOptions() (Options, error) {
@@ -19,6 +21,8 @@ func ParseOptions() (Options, error) {
 		Silent:        flag.Bool("silent", false, "Suppress all logs which are not errors."),
 		NoHistory:     flag.Bool("no-history", false, "Disable interactive session history file."),
 		Commands:      flag.String("eval", "", "Run one or more commands separated by ; in the interactive session, used to set variables via command line."),
+		CpuProfile:    flag.String("cpuprofile", "", "Write cpu profile `file`."),
+		MemProfile:    flag.String("memprofile", "", "Write memory profile to `file`."),
 	}
 
 	flag.Parse()

main.go

@@ -44,7 +44,6 @@ func main() {
 	if err = sess.Start(); err != nil {
 		log.Fatal("%s", err)
 	}
-	defer sess.Close()
 
 	// Some modules are enabled by default in order
 	// to make the interactive session useful.
@@ -94,6 +93,8 @@ func main() {
 		}
 	}
 
+	sess.Close()
+
 	// Windows requires this otherwise the app never exits ...
 	os.Exit(0)
 }

session/session.go

@@ -6,6 +6,8 @@ import (
 	"fmt"
 	"os"
 	"os/signal"
+	"runtime"
+	"runtime/pprof"
 	"sort"
 	"strings"
 	"syscall"
@@ -79,6 +81,14 @@ func New() (*Session, error) {
 		return nil, err
 	}
 
+	if *s.Options.CpuProfile != "" {
+		if f, err := os.Create(*s.Options.CpuProfile); err != nil {
+			return nil, err
+		} else if err := pprof.StartCPUProfile(f); err != nil {
+			return nil, err
+		}
+	}
+
 	s.Env = NewEnvironment(s)
 	s.Events = NewEventPool(*s.Options.Debug, *s.Options.Silent)
 
@@ -180,6 +190,23 @@ func (s *Session) Close() {
 
 	s.Firewall.Restore()
 	s.Queue.Stop()
+
+	if *s.Options.CpuProfile != "" {
+		pprof.StopCPUProfile()
+	}
+
+	if *s.Options.MemProfile != "" {
+		f, err := os.Create(*s.Options.MemProfile)
+		if err != nil {
+			fmt.Printf("Could not create memory profile: %s\n", err)
+			return
+		}
+		defer f.Close()
+		runtime.GC() // get up-to-date statistics
+		if err := pprof.WriteHeapProfile(f); err != nil {
+			fmt.Printf("Could not write memory profile: %s\n", err)
+		}
+	}
 }
 
 func (s *Session) Register(mod Module) error {