From 47230567baa31a6ee98c30af063cee8d217327e1 Mon Sep 17 00:00:00 2001
From: evilsocket <evilsocket@gmail.com>
Date: Sun, 11 Feb 2018 18:41:49 +0100
Subject: [PATCH] new: cpuprofile and memprofile options for easy profiling

---
 .gitignore         |  1 +
 core/options.go    |  4 ++++
 main.go            |  3 ++-
 session/session.go | 27 +++++++++++++++++++++++++++
 4 files changed, 34 insertions(+), 1 deletion(-)

diff --git a/.gitignore b/.gitignore
index 80a3ba02..2c03bfce 100644
--- a/.gitignore
+++ b/.gitignore
@@ -1,5 +1,6 @@
 *.sw*
 *.tar.gz
+*.prof*
 pcaps
 bettercap-ng*.*
 bettercap-ng*
diff --git a/core/options.go b/core/options.go
index 4c1b7f17..8515ff65 100644
--- a/core/options.go
+++ b/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()
diff --git a/main.go b/main.go
index 87d1a3c6..9685a1f3 100644
--- a/main.go
+++ b/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)
 }
diff --git a/session/session.go b/session/session.go
index d6e9afd5..f512da4d 100644
--- a/session/session.go
+++ b/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 {