fix: pruning wifi stations after a certain time of inactivity

modules/wifi_recon.go

@@ -24,6 +24,8 @@ import (
 	"github.com/olekukonko/tablewriter"
 )
 
+var maxStationTTL = 5 * time.Minute
+
 type WiFiRecon struct {
 	session.SessionModule
 
@@ -141,8 +143,15 @@ func (w *WiFiRecon) getRow(station *network.WiFiStation) []string {
 	if encryption == "OPEN" || encryption == "" {
 		encryption = core.Green("OPEN")
 	}
-	sent := humanize.Bytes(station.Sent)
+	sent := ""
-	recvd := humanize.Bytes(station.Received)
+	if station.Sent > 0 {
+		sent = humanize.Bytes(station.Sent)
+	}
+
+	recvd := ""
+	if station.Received > 0 {
+		recvd = humanize.Bytes(station.Received)
+	}
 
 	if w.isApSelected() {
 		return []string{
@@ -369,17 +378,7 @@ func (w *WiFiRecon) updateStats(dot11 *layers.Dot11, packet gopacket.Packet) {
 	}
 }
 
-func (w *WiFiRecon) Start() error {
+func (w *WiFiRecon) channelHopper() {
-	if w.Running() == true {
-		return session.ErrAlreadyStarted
-	} else if err := w.Configure(); err != nil {
-		return err
-	}
-
-	w.SetRunning(true, func() {
-		// start channel hopper if needed
-		if w.channel == 0 {
-			go func() {
 	log.Info("Channel hopper started.")
 	for w.Running() == true {
 		for channel := 1; channel < 15; channel++ {
@@ -393,8 +392,37 @@ func (w *WiFiRecon) Start() error {
 			}
 		}
 	}
-			}()
+}
+
+func (w *WiFiRecon) stationPruner() {
+	log.Debug("WiFi stations pruner started.")
+	for w.Running() == true {
+		for _, s := range w.Session.WiFi.List() {
+			sinceLastSeen := time.Since(s.LastSeen)
+			if sinceLastSeen > maxStationTTL {
+				log.Debug("Station %s not seen in %s, removing.", s.BSSID(), sinceLastSeen)
+				w.Session.WiFi.Remove(s.BSSID())
 			}
+		}
+		time.Sleep(5 * time.Second)
+	}
+}
+
+func (w *WiFiRecon) Start() error {
+	if w.Running() == true {
+		return session.ErrAlreadyStarted
+	} else if err := w.Configure(); err != nil {
+		return err
+	}
+
+	w.SetRunning(true, func() {
+		// start channel hopper if needed
+		if w.channel == 0 {
+			go w.channelHopper()
+		}
+
+		// start the pruner
+		go w.stationPruner()
 
 		defer w.handle.Close()
 		src := gopacket.NewPacketSource(w.handle, w.handle.LinkType())

network/wifi.go

@@ -5,16 +5,24 @@ import (
 	"time"
 )
 
+type StationNewCallback func(s *WiFiStation)
+type StationLostCallback func(s *WiFiStation)
+
 type WiFi struct {
 	sync.Mutex
 	Interface *Endpoint
 	Stations  map[string]*WiFiStation
+
+	newCb  StationNewCallback
+	lostCb StationLostCallback
 }
 
-func NewWiFi(iface *Endpoint) *WiFi {
+func NewWiFi(iface *Endpoint, newcb StationNewCallback, lostcb StationLostCallback) *WiFi {
 	return &WiFi{
 		Interface: iface,
 		Stations:  make(map[string]*WiFiStation),
+		newCb:     newcb,
+		lostCb:    lostcb,
 	}
 }
 
@@ -33,8 +41,11 @@ func (w *WiFi) Remove(mac string) {
 	w.Lock()
 	defer w.Unlock()
 
-	if _, found := w.Stations[mac]; found {
+	if s, found := w.Stations[mac]; found {
 		delete(w.Stations, mac)
+		if w.lostCb != nil {
+			w.lostCb(s)
+		}
 	}
 }
 
@@ -52,6 +63,10 @@ func (w *WiFi) AddIfNew(ssid, mac string, isAp bool, channel int, rssi int8) *Wi
 	newStation := NewWiFiStation(ssid, mac, isAp, channel, rssi)
 	w.Stations[mac] = newStation
 
+	if w.newCb != nil {
+		w.newCb(newStation)
+	}
+
 	return nil
 }
 

session/session.go

@@ -377,7 +377,13 @@ func (s *Session) Start() error {
 	}
 
 	s.Firewall = firewall.Make(s.Interface)
-	s.WiFi = network.NewWiFi(s.Interface)
+
+	s.WiFi = network.NewWiFi(s.Interface, func(st *network.WiFiStation) {
+		s.Events.Add("wifi.station.new", st)
+	}, func(st *network.WiFiStation) {
+		s.Events.Add("wifi.station.lost", st)
+	})
+
 	s.Lan = network.NewLAN(s.Interface, s.Gateway, func(e *network.Endpoint) {
 		s.Events.Add("endpoint.new", e)
 	}, func(e *network.Endpoint) {