misc: small fix or general refactoring i did not bother commenting

modules/wifi_recon.go

@@ -7,6 +7,7 @@ import (
 	"os"
 	"sort"
 	"strconv"
+	"strings"
 	"time"
 
 	"github.com/evilsocket/bettercap-ng/core"
@@ -27,7 +28,6 @@ type WiFiRecon struct {
 	session.SessionModule
 
 	wifi        *WiFi
-	stats       *WiFiStats
 	handle      *pcap.Handle
 	channel     int
 	client      net.HardwareAddr
@@ -37,7 +37,6 @@ type WiFiRecon struct {
 func NewWiFiRecon(s *session.Session) *WiFiRecon {
 	w := &WiFiRecon{
 		SessionModule: session.NewSessionModule("wifi.recon", s),
-		stats:         NewWiFiStats(),
 		channel:       0,
 		client:        make([]byte, 0),
 		accessPoint:   make([]byte, 0),
@@ -143,23 +142,12 @@ func (w *WiFiRecon) getRow(station *WiFiStation) []string {
 	}
 
 	ssid := station.ESSID()
-
+	encryption := station.Encryption
-	encryption := w.stats.EncryptionOf(station.HW)
 	if encryption == "OPEN" {
 		encryption = core.Green(encryption)
 	}
-
+	sent := humanize.Bytes(station.Sent)
-	sent := ""
+	recvd := humanize.Bytes(station.Received)
-	bytes := w.stats.SentFrom(station.HW)
-	if bytes > 0 {
-		sent = humanize.Bytes(bytes)
-	}
-
-	recvd := ""
-	bytes = w.stats.SentTo(station.HW)
-	if bytes > 0 {
-		recvd = humanize.Bytes(bytes)
-	}
 
 	row := []string{
 		fmt.Sprintf("%d dBm", station.RSSI),
@@ -374,17 +362,21 @@ func (w *WiFiRecon) updateStats(dot11 *layers.Dot11, packet gopacket.Packet) {
 	if dot11.Type.MainType() == layers.Dot11TypeData {
 		bytes := uint64(len(packet.Data()))
 
-		dst := dot11.Address1
+		dst := dot11.Address1.String()
-		src := dot11.Address2
+		if station, found := w.wifi.Stations[dst]; found == true {
+			station.Received += bytes
+		}
 
-		w.stats.CollectReceived(dst, bytes)
+		src := dot11.Address2.String()
-		w.stats.CollectSent(src, bytes)
+		if station, found := w.wifi.Stations[src]; found == true {
+			station.Sent += bytes
+		}
 	}
 
 	if ok, enc := packets.Dot11ParseEncryption(packet, dot11); ok == true {
-		w.stats.ResetEncryption(dot11.Address3)
+		bssid := dot11.Address3.String()
-		for _, e := range enc {
+		if station, found := w.wifi.Stations[bssid]; found == true {
-			w.stats.CollectEncryption(dot11.Address3, e)
+			station.Encryption = strings.Join(enc, ", ")
 		}
 	}
 }

modules/wifi_station.go

@@ -9,6 +9,9 @@ type WiFiStation struct {
 	IsAP       bool
 	Channel    int
 	RSSI       int8
+	Sent       uint64
+	Received   uint64
+	Encryption string
 }
 
 func NewWiFiStation(essid, bssid string, isAp bool, channel int, rssi int8) *WiFiStation {

modules/wifi_stats.go

@@ -1,117 +0,0 @@
-package modules
-
-import (
-	"net"
-	"strings"
-	"sync"
-)
-
-type WiFiStationStats struct {
-	Sent       uint64
-	Received   uint64
-	Encryption map[string]bool
-}
-
-func NewWiFiStationStats(sent uint64, recvd uint64) *WiFiStationStats {
-	return &WiFiStationStats{
-		Sent:       sent,
-		Received:   recvd,
-		Encryption: make(map[string]bool),
-	}
-}
-
-type WiFiStats struct {
-	sync.Mutex
-	stats map[string]*WiFiStationStats
-}
-
-func NewWiFiStats() *WiFiStats {
-	return &WiFiStats{
-		stats: make(map[string]*WiFiStationStats),
-	}
-}
-
-func (s *WiFiStats) CollectSent(station net.HardwareAddr, bytes uint64) {
-	s.Lock()
-	defer s.Unlock()
-
-	bssid := station.String()
-	if sstats, found := s.stats[bssid]; found == true {
-		sstats.Sent += bytes
-	} else {
-		s.stats[bssid] = NewWiFiStationStats(bytes, 0)
-	}
-}
-
-func (s *WiFiStats) CollectReceived(station net.HardwareAddr, bytes uint64) {
-	s.Lock()
-	defer s.Unlock()
-
-	bssid := station.String()
-	if sstats, found := s.stats[bssid]; found == true {
-		sstats.Received += bytes
-	} else {
-		s.stats[bssid] = NewWiFiStationStats(0, bytes)
-	}
-}
-
-func (s *WiFiStats) ResetEncryption(station net.HardwareAddr) {
-	s.Lock()
-	defer s.Unlock()
-
-	bssid := station.String()
-	if sstats, found := s.stats[bssid]; found == true {
-		sstats.Encryption = make(map[string]bool)
-	}
-}
-
-func (s *WiFiStats) CollectEncryption(station net.HardwareAddr, enc string) {
-	s.Lock()
-	defer s.Unlock()
-
-	bssid := station.String()
-	if sstats, found := s.stats[bssid]; found == true {
-		sstats.Encryption[enc] = true
-	} else {
-		stats := NewWiFiStationStats(0, 0)
-		stats.Encryption[enc] = true
-		s.stats[bssid] = stats
-	}
-}
-
-func (s *WiFiStats) SentFrom(station net.HardwareAddr) uint64 {
-	s.Lock()
-	defer s.Unlock()
-
-	bssid := station.String()
-	if sstats, found := s.stats[bssid]; found == true {
-		return sstats.Sent
-	}
-	return 0
-}
-
-func (s *WiFiStats) SentTo(station net.HardwareAddr) uint64 {
-	s.Lock()
-	defer s.Unlock()
-
-	bssid := station.String()
-	if sstats, found := s.stats[bssid]; found == true {
-		return sstats.Received
-	}
-	return 0
-}
-
-func (s *WiFiStats) EncryptionOf(station net.HardwareAddr) string {
-	s.Lock()
-	defer s.Unlock()
-
-	bssid := station.String()
-	if sstats, found := s.stats[bssid]; found == true {
-		unique := make([]string, 0)
-		for key := range sstats.Encryption {
-			unique = append(unique, key)
-		}
-		return strings.Join(unique, ", ")
-	}
-	return ""
-}