From 420ab9e26cd453d5f0a5862e2bf68976dfc2488a Mon Sep 17 00:00:00 2001
From: Matrix86 <matrix86@tuxmealux.net>
Date: Mon, 19 Feb 2018 20:53:44 +0100
Subject: [PATCH] Fix on concurrent access to the Stations map by UpdateStats
 method

---
 modules/wifi_recon.go | 6 +++---
 network/wifi.go       | 9 +++++++++
 2 files changed, 12 insertions(+), 3 deletions(-)

diff --git a/modules/wifi_recon.go b/modules/wifi_recon.go
index daf50815..f9b3383b 100644
--- a/modules/wifi_recon.go
+++ b/modules/wifi_recon.go
@@ -345,19 +345,19 @@ func (w *WiFiRecon) updateStats(dot11 *layers.Dot11, packet gopacket.Packet) {
 		bytes := uint64(len(packet.Data()))
 
 		dst := dot11.Address1.String()
-		if station, found := w.Session.WiFi.Stations[dst]; found == true {
+		if station, found := w.Session.WiFi.Get(dst); found == true {
 			station.Received += bytes
 		}
 
 		src := dot11.Address2.String()
-		if station, found := w.Session.WiFi.Stations[src]; found == true {
+		if station, found := w.Session.WiFi.Get(src); found == true {
 			station.Sent += bytes
 		}
 	}
 
 	if ok, enc := packets.Dot11ParseEncryption(packet, dot11); ok == true {
 		bssid := dot11.Address3.String()
-		if station, found := w.Session.WiFi.Stations[bssid]; found == true {
+		if station, found := w.Session.WiFi.Get(bssid); found == true {
 			station.Encryption = strings.Join(enc, ", ")
 		}
 	}
diff --git a/network/wifi.go b/network/wifi.go
index d31510f1..bc2ab654 100644
--- a/network/wifi.go
+++ b/network/wifi.go
@@ -70,6 +70,15 @@ func (w *WiFi) AddIfNew(ssid, mac string, isAp bool, channel int, rssi int8) *St
 	return nil
 }
 
+func (w *WiFi) Get(mac string) (*Station, bool) {
+	w.Lock()
+	defer w.Unlock()
+
+	mac = NormalizeMac(mac)
+	station, found := w.Stations[mac]
+	return station, found
+}
+
 func (w *WiFi) Clear() error {
 	w.Stations = make(map[string]*Station)
 	return nil