From 1120f98be19e0a4153a763cdf0af169d961cbe71 Mon Sep 17 00:00:00 2001
From: evilsocket <evilsocket@gmail.com>
Date: Wed, 6 Feb 2019 07:08:01 +0100
Subject: [PATCH] fix: one beacon packet per handshake is saved in order to
 have both the ESSID and BSSID

---
 modules/wifi_recon.go     |  6 +++++-
 network/wifi.go           |  6 +++---
 network/wifi_ap.go        |  9 +++++++++
 network/wifi_handshake.go | 11 +++++++++++
 4 files changed, 28 insertions(+), 4 deletions(-)

diff --git a/modules/wifi_recon.go b/modules/wifi_recon.go
index 0fa7f64a..93a8a60a 100644
--- a/modules/wifi_recon.go
+++ b/modules/wifi_recon.go
@@ -79,7 +79,11 @@ func (w *WiFiModule) discoverAccessPoints(radiotap *layers.RadioTap, dot11 *laye
 				frequency = int(radiotap.ChannelFrequency)
 			}
 
-			w.Session.WiFi.AddIfNew(ssid, bssid, frequency, radiotap.DBMAntennaSignal)
+			if ap, isNew := w.Session.WiFi.AddIfNew(ssid, bssid, frequency, radiotap.DBMAntennaSignal); !isNew {
+				ap.EachClient(func(mac string, station *network.Station) {
+					station.Handshake.SetBeacon(packet)
+				})
+			}
 		}
 	}
 }
diff --git a/network/wifi.go b/network/wifi.go
index 0d265856..75dcff2a 100644
--- a/network/wifi.go
+++ b/network/wifi.go
@@ -129,7 +129,7 @@ func isBogusMacESSID(essid string) bool {
 	return false
 }
 
-func (w *WiFi) AddIfNew(ssid, mac string, frequency int, rssi int8) *AccessPoint {
+func (w *WiFi) AddIfNew(ssid, mac string, frequency int, rssi int8) (*AccessPoint, bool) {
 	w.Lock()
 	defer w.Unlock()
 
@@ -141,7 +141,7 @@ func (w *WiFi) AddIfNew(ssid, mac string, frequency int, rssi int8) *AccessPoint
 		if !isBogusMacESSID(ssid) {
 			ap.Hostname = ssid
 		}
-		return ap
+		return ap, false
 	}
 
 	newAp := NewAccessPoint(ssid, mac, frequency, rssi)
@@ -151,7 +151,7 @@ func (w *WiFi) AddIfNew(ssid, mac string, frequency int, rssi int8) *AccessPoint
 		w.newCb(newAp)
 	}
 
-	return nil
+	return newAp, true
 }
 
 func (w *WiFi) Get(mac string) (*AccessPoint, bool) {
diff --git a/network/wifi_ap.go b/network/wifi_ap.go
index 9ff5ec4a..e278c542 100644
--- a/network/wifi_ap.go
+++ b/network/wifi_ap.go
@@ -97,6 +97,15 @@ func (ap *AccessPoint) Clients() (list []*Station) {
 	return
 }
 
+func (ap *AccessPoint) EachClient(cb func(mac string, station *Station)) {
+	ap.Lock()
+	defer ap.Unlock()
+
+	for m, station := range ap.clients {
+		cb(m, station)
+	}
+}
+
 func (ap *AccessPoint) NumHandshakes() int {
 	ap.Lock()
 	defer ap.Unlock()
diff --git a/network/wifi_handshake.go b/network/wifi_handshake.go
index c97a046f..9c42aa63 100644
--- a/network/wifi_handshake.go
+++ b/network/wifi_handshake.go
@@ -8,6 +8,7 @@ import (
 type Handshake struct {
 	sync.Mutex
 
+	Beacon        gopacket.Packet
 	Challenges    []gopacket.Packet
 	Responses     []gopacket.Packet
 	Confirmations []gopacket.Packet
@@ -23,6 +24,16 @@ func NewHandshake() *Handshake {
 	}
 }
 
+func (h *Handshake) SetBeacon(pkt gopacket.Packet) {
+	h.Lock()
+	defer h.Unlock()
+
+	if h.Beacon == nil {
+		h.Beacon = pkt
+		h.unsaved = append(h.unsaved, pkt)
+	}
+}
+
 func (h *Handshake) AddFrame(n int, pkt gopacket.Packet) {
 	h.Lock()
 	defer h.Unlock()