fix: properly synchronizing modules stop action with waitGroups

modules/arp_spoof.go

@@ -5,6 +5,7 @@ import (
 	"net"
 	"regexp"
 	"strings"
+	"sync"
 	"time"
 
 	"github.com/bettercap/bettercap/log"
@@ -20,19 +21,19 @@ var macParser = regexp.MustCompile(`([a-fA-F0-9]{1,2}:[a-fA-F0-9]{1,2}:[a-fA-F0-
 
 type ArpSpoofer struct {
 	session.SessionModule
-	done      chan bool
 	addresses []net.IP
 	macs      []net.HardwareAddr
 	ban       bool
+	waitGroup *sync.WaitGroup
 }
 
 func NewArpSpoofer(s *session.Session) *ArpSpoofer {
 	p := &ArpSpoofer{
 		SessionModule: session.NewSessionModule("arp.spoof", s),
-		done:          make(chan bool),
 		addresses:     make([]net.IP, 0),
 		macs:          make([]net.HardwareAddr, 0),
 		ban:           false,
+		waitGroup:     &sync.WaitGroup{},
 	}
 
 	p.AddParam(session.NewStringParameter("arp.spoof.targets", session.ParamSubnet, "", "IP or MAC addresses to spoof, also supports nmap style IP ranges."))
@@ -78,6 +79,9 @@ func (p ArpSpoofer) Author() string {
 }
 
 func (p *ArpSpoofer) sendArp(saddr net.IP, smac net.HardwareAddr, check_running bool, probe bool) {
+	p.waitGroup.Add(1)
+	defer p.waitGroup.Done()
+
 	targets := make(map[string]net.HardwareAddr, 0)
 
 	for _, ip := range p.addresses {
@@ -198,20 +202,21 @@ func (p *ArpSpoofer) Start() error {
 
 		log.Info("ARP spoofer started, probing %d targets.", len(p.addresses)+len(p.macs))
 
+		p.waitGroup.Add(1)
+		defer p.waitGroup.Done()
+
 		for p.Running() {
 			p.sendArp(from, from_hw, true, false)
 			time.Sleep(1 * time.Second)
 		}
-
-		p.done <- true
 	})
 }
 
 func (p *ArpSpoofer) Stop() error {
 	return p.SetRunning(false, func() {
 		log.Info("Waiting for ARP spoofer to stop ...")
-		<-p.done
 		p.unSpoof()
 		p.ban = false
+		p.waitGroup.Wait()
 	})
 }

modules/dhcp6_spoof.go

@@ -6,6 +6,7 @@ import (
 	"fmt"
 	"net"
 	"strings"
+	"sync"
 	"time"
 
 	"github.com/bettercap/bettercap/core"
@@ -30,12 +31,14 @@ type DHCP6Spoofer struct {
 	DUIDRaw    []byte
 	Domains    []string
 	RawDomains []byte
+	waitGroup  *sync.WaitGroup
 }
 
 func NewDHCP6Spoofer(s *session.Session) *DHCP6Spoofer {
 	spoof := &DHCP6Spoofer{
 		SessionModule: session.NewSessionModule("dhcp6.spoof", s),
 		Handle:        nil,
+		waitGroup:     &sync.WaitGroup{},
 	}
 
 	spoof.AddParam(session.NewStringParameter("dhcp6.spoof.domains",
@@ -359,6 +362,9 @@ func (s *DHCP6Spoofer) Start() error {
 	return s.SetRunning(true, func() {
 		defer s.Handle.Close()
 
+		s.waitGroup.Add(1)
+		defer s.waitGroup.Done()
+
 		src := gopacket.NewPacketSource(s.Handle, s.Handle.LinkType())
 		for packet := range src.Packets() {
 			if s.Running() == false {
@@ -373,5 +379,6 @@ func (s *DHCP6Spoofer) Start() error {
 func (s *DHCP6Spoofer) Stop() error {
 	return s.SetRunning(false, func() {
 		s.Handle.Close()
+		s.waitGroup.Wait()
 	})
 }

modules/dns_spoof.go

@@ -5,6 +5,7 @@ import (
 	"fmt"
 	"net"
 	"strings"
+	"sync"
 
 	"github.com/bettercap/bettercap/core"
 	"github.com/bettercap/bettercap/log"
@@ -22,6 +23,7 @@ type DNSSpoofer struct {
 	Domains   []string
 	Address   net.IP
 	All       bool
+	waitGroup *sync.WaitGroup
 }
 
 func NewDNSSpoofer(s *session.Session) *DNSSpoofer {
@@ -29,6 +31,7 @@ func NewDNSSpoofer(s *session.Session) *DNSSpoofer {
 		SessionModule: session.NewSessionModule("dns.spoof", s),
 		Handle:        nil,
 		All:           false,
+		waitGroup:     &sync.WaitGroup{},
 	}
 
 	spoof.AddParam(session.NewStringParameter("dns.spoof.domains",
@@ -267,6 +270,9 @@ func (s *DNSSpoofer) Start() error {
 	}
 
 	return s.SetRunning(true, func() {
+		s.waitGroup.Add(1)
+		defer s.waitGroup.Done()
+
 		defer s.Handle.Close()
 
 		src := gopacket.NewPacketSource(s.Handle, s.Handle.LinkType())
@@ -283,5 +289,6 @@ func (s *DNSSpoofer) Start() error {
 func (s *DNSSpoofer) Stop() error {
 	return s.SetRunning(false, func() {
 		s.Handle.Close()
+		s.waitGroup.Wait()
 	})
 }

modules/net_probe.go

@@ -15,11 +15,13 @@ import (
 type Prober struct {
 	session.SessionModule
 	throttle  int
+	waitGroup *sync.WaitGroup
 }
 
 func NewProber(s *session.Session) *Prober {
 	p := &Prober{
 		SessionModule: session.NewSessionModule("net.probe", s),
+		waitGroup:     &sync.WaitGroup{},
 	}
 
 	p.AddParam(session.NewIntParameter("net.probe.throttle",
@@ -81,6 +83,9 @@ func (p *Prober) Start() error {
 	}
 
 	return p.SetRunning(true, func() {
+		p.waitGroup.Add(1)
+		defer p.waitGroup.Done()
+
 		if p.Session.Interface.IpAddress == network.MonitorModeAddress {
 			log.Info("Interface is in monitor mode, skipping net.probe")
 			return
@@ -98,7 +103,9 @@ func (p *Prober) Start() error {
 
 		for p.Running() {
 			for _, ip := range addresses {
-				if p.Session.Skip(ip) == true {
+				if p.Running() == false {
+					return
+				} else if p.Session.Skip(ip) == true {
 					log.Debug("Skipping address %s from UDP probing.", ip)
 					continue
 				}
@@ -114,5 +121,7 @@ func (p *Prober) Start() error {
 }
 
 func (p *Prober) Stop() error {
-	return p.SetRunning(false, nil)
+	return p.SetRunning(false, func() {
+		p.waitGroup.Wait()
+	})
 }

modules/syn_scan.go

@@ -4,6 +4,7 @@ import (
 	"fmt"
 	"net"
 	"strconv"
+	"sync"
 	"time"
 
 	"github.com/bettercap/bettercap/core"
@@ -25,6 +26,7 @@ type SynScanner struct {
 	addresses []net.IP
 	startPort int
 	endPort   int
+	waitGroup *sync.WaitGroup
 }
 
 func NewSynScanner(s *session.Session) *SynScanner {
@@ -33,6 +35,7 @@ func NewSynScanner(s *session.Session) *SynScanner {
 		addresses:     make([]net.IP, 0),
 		startPort:     0,
 		endPort:       0,
+		waitGroup:     &sync.WaitGroup{},
 	}
 
 	ss.AddHandler(session.NewModuleHandler("syn.scan IP-RANGE START-PORT END-PORT", "syn.scan ([^\\s]+) (\\d+)([\\s\\d]*)",
@@ -97,10 +100,6 @@ func (s *SynScanner) Start() error {
 	return nil
 }
 
-func (s *SynScanner) Stop() error {
-	return nil
-}
-
 func (s *SynScanner) inRange(ip net.IP) bool {
 	for _, a := range s.addresses {
 		if a.Equal(ip) {
@@ -154,6 +153,9 @@ func (s *SynScanner) synScan() error {
 	s.SetRunning(true, func() {
 		defer s.SetRunning(false, nil)
 
+		s.waitGroup.Add(1)
+		defer s.waitGroup.Done()
+
 		naddrs := len(s.addresses)
 		plural := "es"
 		if naddrs == 1 {
@@ -172,6 +174,9 @@ func (s *SynScanner) synScan() error {
 
 		// start sending SYN packets and wait
 		for _, address := range s.addresses {
+			if s.Running() == false {
+				break
+			}
 			mac, err := findMAC(s.Session, address, true)
 			if err != nil {
 				log.Debug("Could not get MAC for %s: %s", address.String(), err)
@@ -179,6 +184,10 @@ func (s *SynScanner) synScan() error {
 			}
 
 			for dstPort := s.startPort; dstPort < s.endPort+1; dstPort++ {
+				if s.Running() == false {
+					break
+				}
+
 				err, raw := packets.NewTCPSyn(s.Session.Interface.IP, s.Session.Interface.HW, address, mac, synSourcePort, dstPort)
 				if err != nil {
 					log.Error("Error creating SYN packet: %s", err)
@@ -199,3 +208,9 @@ func (s *SynScanner) synScan() error {
 
 	return nil
 }
+
+func (s *SynScanner) Stop() error {
+	return s.SetRunning(false, func() {
+		s.waitGroup.Wait()
+	})
+}

modules/wifi_recon.go

@@ -6,6 +6,7 @@ import (
 	"os"
 	"sort"
 	"strconv"
+	"sync"
 	"time"
 
 	"github.com/bettercap/bettercap/core"
@@ -41,6 +42,7 @@ type WiFiRecon struct {
 	ap          *network.AccessPoint
 	stickChan   int
 	skipBroken  bool
+	waitGroup   *sync.WaitGroup
 }
 
 func NewWiFiRecon(s *session.Session) *WiFiRecon {
@@ -51,6 +53,7 @@ func NewWiFiRecon(s *session.Session) *WiFiRecon {
 		hopPeriod:     250 * time.Millisecond,
 		ap:            nil,
 		skipBroken:    true,
+		waitGroup:     &sync.WaitGroup{},
 	}
 
 	w.AddHandler(session.NewModuleHandler("wifi.recon on", "",
@@ -343,7 +346,7 @@ func (w *WiFiRecon) injectPacket(data []byte) {
 }
 
 func (w *WiFiRecon) sendDeauthPacket(ap net.HardwareAddr, client net.HardwareAddr) {
-	for seq := uint16(0); seq < 64; seq++ {
+	for seq := uint16(0); seq < 64 && w.Running(); seq++ {
 		if err, pkt := packets.NewDot11Deauth(ap, client, ap, seq); err != nil {
 			log.Error("Could not create deauth packet: %s", err)
 			continue
@@ -385,6 +388,9 @@ func (w *WiFiRecon) startDeauth(to net.HardwareAddr) error {
 		defer w.handle.Close()
 	}
 
+	w.waitGroup.Add(1)
+	defer w.waitGroup.Done()
+
 	bssid := to.String()
 
 	// are we deauthing every client of a given access point?
@@ -393,6 +399,9 @@ func (w *WiFiRecon) startDeauth(to net.HardwareAddr) error {
 		log.Info("Deauthing %d clients from AP %s ...", len(clients), ap.ESSID())
 		w.onChannel(mhz2chan(ap.Frequency), func() {
 			for _, c := range clients {
+				if w.Running() == false {
+					break
+				}
 				w.sendDeauthPacket(ap.HW, c.HW)
 			}
 		})
@@ -403,7 +412,9 @@ func (w *WiFiRecon) startDeauth(to net.HardwareAddr) error {
 	// search for a client
 	aps := w.Session.WiFi.List()
 	for _, ap := range aps {
-		if c, found := ap.Get(bssid); found == true {
+		if w.Running() == false {
+			break
+		} else if c, found := ap.Get(bssid); found == true {
 			log.Info("Deauthing client %s from AP %s ...", c.HwAddress, ap.ESSID())
 			w.onChannel(mhz2chan(ap.Frequency), func() {
 				w.sendDeauthPacket(ap.HW, c.HW)
@@ -518,6 +529,9 @@ func (w *WiFiRecon) updateStats(dot11 *layers.Dot11, packet gopacket.Packet) {
 }
 
 func (w *WiFiRecon) channelHopper() {
+	w.waitGroup.Add(1)
+	defer w.waitGroup.Done()
+
 	log.Info("Channel hopper started.")
 	for w.Running() == true {
 		delay := w.hopPeriod
@@ -549,6 +563,9 @@ func (w *WiFiRecon) channelHopper() {
 }
 
 func (w *WiFiRecon) stationPruner() {
+	w.waitGroup.Add(1)
+	defer w.waitGroup.Done()
+
 	log.Debug("WiFi stations pruner started.")
 	for w.Running() == true {
 		for _, s := range w.Session.WiFi.List() {
@@ -589,6 +606,9 @@ func (w *WiFiRecon) Start() error {
 		// start the pruner
 		go w.stationPruner()
 
+		w.waitGroup.Add(1)
+		defer w.waitGroup.Done()
+
 		defer w.handle.Close()
 		src := gopacket.NewPacketSource(w.handle, w.handle.LinkType())
 		for packet := range src.Packets() {
@@ -618,5 +638,7 @@ func (w *WiFiRecon) Start() error {
 }
 
 func (w *WiFiRecon) Stop() error {
-	return w.SetRunning(false, nil)
+	return w.SetRunning(false, func() {
+		w.waitGroup.Wait()
+	})
 }