new: http and https proxy modules can now define an onCommand callback to handle custom session commands (closes #182)

modules/http_proxy_base.go

@@ -144,6 +144,13 @@ func (p *HTTPProxy) Configure(address string, proxyPort int, httpPort int, scrip
 
 	log.Debug("Applied redirection %s", p.Redirection.String())
 
+	p.sess.UnkCmdCallback = func(cmd string) bool {
+		if p.Script != nil {
+			return p.Script.OnCommand(cmd)
+		}
+		return false
+	}
+
 	return nil
 }
 
@@ -323,6 +330,8 @@ func (p *HTTPProxy) Stop() error {
 		p.Redirection = nil
 	}
 
+	p.sess.UnkCmdCallback = nil
+
 	if p.isTLS == true {
 		p.isRunning = false
 		p.sniListener.Close()

modules/http_proxy_script.go

@@ -14,6 +14,7 @@ type HttpProxyScript struct {
 	*ProxyScript
 	onRequestScript  *otto.Script
 	onResponseScript *otto.Script
+	onCommandScript  *otto.Script
 }
 
 func LoadHttpProxyScriptSource(path, source string, sess *session.Session) (err error, s *HttpProxyScript) {
@@ -26,9 +27,9 @@ func LoadHttpProxyScriptSource(path, source string, sess *session.Session) (err
 		ProxyScript:      ps,
 		onRequestScript:  nil,
 		onResponseScript: nil,
+		onCommandScript:  nil,
 	}
 
-	// compile call to onRequest if defined
 	if s.hasCallback("onRequest") {
 		s.onRequestScript, err = s.VM.Compile("", "onRequest(req, res)")
 		if err != nil {
@@ -37,7 +38,6 @@ func LoadHttpProxyScriptSource(path, source string, sess *session.Session) (err
 		}
 	}
 
-	// compile call to onResponse if defined
 	if s.hasCallback("onResponse") {
 		s.onResponseScript, err = s.VM.Compile("", "onResponse(req, res)")
 		if err != nil {
@@ -46,6 +46,14 @@ func LoadHttpProxyScriptSource(path, source string, sess *session.Session) (err
 		}
 	}
 
+	if s.hasCallback("onCommand") {
+		s.onCommandScript, err = s.VM.Compile("", "onCommand(cmd)")
+		if err != nil {
+			log.Error("Error while compiling onCommand callback: %s", err)
+			return
+		}
+	}
+
 	return
 }
 
@@ -91,6 +99,13 @@ func (s *HttpProxyScript) doResponseDefines(res *http.Response) (err error, jsre
 	return
 }
 
+func (s *HttpProxyScript) doCommandDefines(cmd string) (err error) {
+	if err = s.VM.Set("cmd", cmd); err != nil {
+		log.Error("Error while defining cmd: %s", err)
+	}
+	return
+}
+
 func (s *HttpProxyScript) OnRequest(original *http.Request) (jsreq *JSRequest, jsres *JSResponse) {
 	var err error
 
@@ -145,3 +160,24 @@ func (s *HttpProxyScript) OnResponse(res *http.Response) (jsreq *JSRequest, jsre
 
 	return nil, nil
 }
+
+func (s *HttpProxyScript) OnCommand(cmd string) bool {
+	if s.onCommandScript != nil {
+		s.Lock()
+		defer s.Unlock()
+
+		if err := s.doCommandDefines(cmd); err != nil {
+			log.Error("Error while running bootstrap onCommand definitions: %s", err)
+			return false
+		}
+
+		if ret, err := s.VM.Run(s.onCommandScript); err != nil {
+			log.Error("Error while executing onCommand callback: %s", err)
+			return false
+		} else if v, err := ret.ToBoolean(); err == nil {
+			return v
+		}
+	}
+
+	return false
+}

session/session.go

@@ -37,6 +37,8 @@ var (
 	reCmdSpaceCleaner = regexp.MustCompile(`^([^\s]+)\s+(.+)$`)
 )
 
+type UnknownCommandCallback func(cmd string) bool
+
 type Session struct {
 	Options   core.Options             `json:"options"`
 	Interface *network.Endpoint        `json:"interface"`
@@ -57,6 +59,8 @@ type Session struct {
 	Modules      []Module         `json:"-"`
 
 	Events *EventPool `json:"-"`
+
+	UnkCmdCallback UnknownCommandCallback `json:"-"`
 }
 
 func ParseCommands(line string) []string {
@@ -135,6 +139,7 @@ func New() (*Session, error) {
 		CoreHandlers:   make([]CommandHandler, 0),
 		Modules:        make([]Module, 0),
 		Events:         nil,
+		UnkCmdCallback: nil,
 	}
 
 	if s.Options, err = core.ParseOptions(); err != nil {
@@ -501,5 +506,9 @@ func (s *Session) Run(line string) error {
 		}
 	}
 
+	if s.UnkCmdCallback != nil && s.UnkCmdCallback(line) == true {
+		return nil
+	}
+
 	return fmt.Errorf("Unknown or invalid syntax \"%s%s%s\", type %shelp%s for the help menu.", core.BOLD, line, core.RESET, core.BOLD, core.RESET)
 }