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

modules/events_view_wifi.go

@@ -78,7 +78,7 @@ func (s *EventsStream) viewWiFiHandshakeEvent(e session.Event) {
 	}
 
 	if hand.PMKID != nil {
-		what = fmt.Sprintf("PMKID (%x)", hand.PMKID)
+		what = "RSN PMKID"
 	}
 
 	fmt.Fprintf(s.output, "[%s] [%s] captured %s -> %s %s to %s\n",

modules/wifi_recon.go

@@ -144,6 +144,7 @@ func allZeros(s []byte) bool {
 }
 
 func (w *WiFiModule) discoverHandshakes(radiotap *layers.RadioTap, dot11 *layers.Dot11, packet gopacket.Packet) {
+	// ref. https://wlan1nde.wordpress.com/2014/10/27/4-way-handshake/
 	if keyLayer := packet.Layer(layers.LayerTypeEAPOLKey); keyLayer != nil {
 		if key := keyLayer.(*layers.EAPOLKey); key.KeyType == layers.EAPOLKeyTypePairwise {
 			staMac := net.HardwareAddr{}
@@ -156,74 +157,87 @@ func (w *WiFiModule) discoverHandshakes(radiotap *layers.RadioTap, dot11 *layers
 				apMac = dot11.Address1
 			}
 
+			// first, locate the AP in our list by its BSSID
 			ap, found := w.Session.WiFi.Get(apMac.String())
 			if !found {
 				log.Warning("could not find AP with BSSID %s", apMac.String())
 				return
 			}
 
-			station := (*network.Station)(nil)
+			// locate the client station, if its BSSID is ours, it means we sent
+			// an association request via wifi.assoc because we're trying to capture
+			// the PMKID from the first EAPOL sent by the AP.
+			// (Reference about PMKID https://hashcat.net/forum/thread-7717.html)
+			// In this case, we need to add ourselves as a client station of the AP
+			// in order to have a consistent association of AP, client and handshakes.
 			staIsUs := bytes.Equal(staMac, w.Session.Interface.HW)
-			if staIsUs {
-				// add a fake station
+			station, found := ap.Get(staMac.String())
+			if !found {
 				station, _ = ap.AddClientIfNew(staMac.String(), ap.Frequency, ap.RSSI)
-				found = true
-			} else {
-				station, found = ap.Get(staMac.String())
 			}
 
-			// ref. https://wlan1nde.wordpress.com/2014/10/27/4-way-handshake/
-			if found {
-				rawPMKID := []byte(nil)
-				if !key.Install && key.KeyACK && !key.KeyMIC {
-					// [1] (ACK) AP is sending ANonce to the client
-					log.Debug("[%s] got frame 1/4 of the %s <-> %s handshake (anonce:%x)",
-						tui.Green("wifi"),
-						apMac,
-						staMac,
-						key.Nonce)
-					// ref. https://hashcat.net/forum/thread-7717.html
-					rawPMKID = station.Handshake.AddAndGetPMKID(packet)
-				} else if !key.Install && !key.KeyACK && key.KeyMIC && !allZeros(key.Nonce) {
-					// [2] (MIC) client is sending SNonce+MIC to the API
-					log.Debug("[%s] got frame 2/4 of the %s <-> %s handshake (snonce:%x mic:%x)",
-						tui.Green("wifi"),
-						apMac,
-						staMac,
-						key.Nonce,
-						key.MIC)
-					station.Handshake.AddFrame(1, packet)
-				} else if key.Install && key.KeyACK && key.KeyMIC {
-					// [3]: (INSTALL+ACK+MIC) AP informs the client that the PTK is installed
-					log.Debug("[%s] got frame 3/4 of the %s <-> %s handshake (mic:%x)",
-						tui.Green("wifi"),
-						apMac,
-						staMac,
-						key.MIC)
-					station.Handshake.AddFrame(2, packet)
+			rawPMKID := []byte(nil)
+			if !key.Install && key.KeyACK && !key.KeyMIC {
+				// [1] (ACK) AP is sending ANonce to the client
+				rawPMKID = station.Handshake.AddAndGetPMKID(packet)
+				PMKID := "without PMKID"
+				if rawPMKID != nil {
+					PMKID = "with PMKID"
 				}
 
-				numUnsaved := station.Handshake.NumUnsaved()
-				doSave := numUnsaved > 0
-				if doSave && w.shakesFile != "" {
-					log.Debug("saving handshake frames to %s", w.shakesFile)
-					if err := w.Session.WiFi.SaveHandshakesTo(w.shakesFile, w.handle.LinkType()); err != nil {
-						log.Error("error while saving handshake frames to %s: %s", w.shakesFile, err)
-					}
-				}
+				log.Debug("[%s] got frame 1/4 of the %s <-> %s handshake (%s) (anonce:%x)",
+					tui.Green("wifi"),
+					apMac,
+					staMac,
+					PMKID,
+					key.Nonce)
+			} else if !key.Install && !key.KeyACK && key.KeyMIC && !allZeros(key.Nonce) {
+				// [2] (MIC) client is sending SNonce+MIC to the API
+				station.Handshake.AddFrame(1, packet)
 
-				if doSave && (rawPMKID != nil || station.Handshake.Complete()) {
-					w.Session.Events.Add("wifi.client.handshake", WiFiHandshakeEvent{
-						File:       w.shakesFile,
-						NewPackets: numUnsaved,
-						AP:         apMac,
-						Station:    staMac,
-						PMKID:      rawPMKID,
-					})
-				}
+				log.Debug("[%s] got frame 2/4 of the %s <-> %s handshake (snonce:%x mic:%x)",
+					tui.Green("wifi"),
+					apMac,
+					staMac,
+					key.Nonce,
+					key.MIC)
+			} else if key.Install && key.KeyACK && key.KeyMIC {
+				// [3]: (INSTALL+ACK+MIC) AP informs the client that the PTK is installed
+				station.Handshake.AddFrame(2, packet)
 
-			} else {
-				log.Warning("EAPOL captured for unknown station %s", staMac.String())
+				log.Debug("[%s] got frame 3/4 of the %s <-> %s handshake (mic:%x)",
+					tui.Green("wifi"),
+					apMac,
+					staMac,
+					key.MIC)
+			}
+
+			// if we have unsaved packets as part of the handshake, save them.
+			numUnsaved := station.Handshake.NumUnsaved()
+			doSave := numUnsaved > 0
+			if doSave && w.shakesFile != "" {
+				log.Debug("saving handshake frames to %s", w.shakesFile)
+				if err := w.Session.WiFi.SaveHandshakesTo(w.shakesFile, w.handle.LinkType()); err != nil {
+					log.Error("error while saving handshake frames to %s: %s", w.shakesFile, err)
+				}
+			}
+
+			// if we had unsaved packets and either the handshake is complete
+			// or it contains the PMKID, generate a new event.
+			if doSave && (rawPMKID != nil || station.Handshake.Complete()) {
+				w.Session.Events.Add("wifi.client.handshake", WiFiHandshakeEvent{
+					File:       w.shakesFile,
+					NewPackets: numUnsaved,
+					AP:         apMac,
+					Station:    staMac,
+					PMKID:      rawPMKID,
+				})
+			}
+
+			// if we added ourselves as a client station but we didn't get any
+			// PMKID, just remove it from the list of clients of this AP.
+			if staIsUs && rawPMKID == nil {
+				ap.RemoveClient(staMac.String())
 			}
 		}
 	}