From 0108e41e20f3c4daf379215c3e53c8d8d260a806 Mon Sep 17 00:00:00 2001
From: evilsocket <evilsocket@gmail.com>
Date: Tue, 27 Feb 2018 04:39:46 +0100
Subject: [PATCH] new: implemented events.waitfor command (closes #91)

---
 modules/events_stream.go | 52 ++++++++++++++++++++++++++++++++++++++++
 1 file changed, 52 insertions(+)

diff --git a/modules/events_stream.go b/modules/events_stream.go
index d8dbfe69..7fefbedf 100644
--- a/modules/events_stream.go
+++ b/modules/events_stream.go
@@ -3,6 +3,7 @@ package modules
 import (
 	"fmt"
 	"strconv"
+	"time"
 
 	"github.com/bettercap/bettercap/core"
 	"github.com/bettercap/bettercap/log"
@@ -12,6 +13,8 @@ import (
 type EventsStream struct {
 	session.SessionModule
 	ignoreList *IgnoreList
+	waitFor    string
+	waitChan   chan *session.Event
 	quit       chan bool
 }
 
@@ -19,6 +22,8 @@ func NewEventsStream(s *session.Session) *EventsStream {
 	stream := &EventsStream{
 		SessionModule: session.NewSessionModule("events.stream", s),
 		quit:          make(chan bool),
+		waitChan:      make(chan *session.Event),
+		waitFor:       "",
 		ignoreList:    NewIgnoreList(),
 	}
 
@@ -45,6 +50,24 @@ func NewEventsStream(s *session.Session) *EventsStream {
 			return stream.Show(limit)
 		}))
 
+	stream.AddHandler(session.NewModuleHandler("events.waitfor TAG TIMEOUT?", `events.waitfor ([^\s]+)([\s\d]*)`,
+		"Show events stream.",
+		func(args []string) error {
+			tag := args[0]
+			timeout := 0
+			if len(args) == 2 {
+				t := core.Trim(args[1])
+				if t != "" {
+					n, err := strconv.Atoi(t)
+					if err != nil {
+						return err
+					}
+					timeout = n
+				}
+			}
+			return stream.startWaitingFor(tag, timeout)
+		}))
+
 	stream.AddHandler(session.NewModuleHandler("events.ignore FILTER", "events.ignore ([^\\s]+)",
 		"Events with an identifier matching this filter will not be shown (use multiple times to add more filters).",
 		func(args []string) error {
@@ -105,6 +128,11 @@ func (s *EventsStream) Start() error {
 			var e session.Event
 			select {
 			case e = <-s.Session.Events.NewEvents:
+				if e.Tag == s.waitFor {
+					s.waitFor = ""
+					s.waitChan <- &e
+				}
+
 				if s.ignoreList.Ignored(e) == false {
 					s.View(e, true)
 				} else {
@@ -137,6 +165,30 @@ func (s *EventsStream) Show(limit int) error {
 	return nil
 }
 
+func (s *EventsStream) startWaitingFor(tag string, timeout int) error {
+	if timeout == 0 {
+		log.Info("Waiting for event %s ...", core.Green(tag))
+	} else {
+		log.Info("Waiting for event %s for %d seconds ...", core.Green(tag), timeout)
+		go func() {
+			time.Sleep(time.Duration(timeout) * time.Second)
+			s.waitFor = ""
+			s.waitChan <- nil
+		}()
+	}
+
+	s.waitFor = tag
+	event := <-s.waitChan
+
+	if event == nil {
+		return fmt.Errorf("'events.waitFor %s %d' timed out.", tag, timeout)
+	} else {
+		log.Debug("Got event: %v", event)
+	}
+
+	return nil
+}
+
 func (s *EventsStream) Stop() error {
 	return s.SetRunning(false, func() {
 		s.quit <- true