From c6e8086d108c85a0c917d87e46533245d791e2f0 Mon Sep 17 00:00:00 2001
From: Mathy Vanhoef <Mathy.Vanhoef@nyu.edu>
Date: Fri, 26 Jun 2020 15:50:18 +0400
Subject: [PATCH] fragattack: updated event handling to support more devices

---
 research/fraginternals.py | 23 ++++++++++-------------
 src/ap/ieee802_11.c       |  5 -----
 src/ap/sta_info.c         |  5 +++++
 wpa_supplicant/events.c   |  4 ++++
 4 files changed, 19 insertions(+), 18 deletions(-)

diff --git a/research/fraginternals.py b/research/fraginternals.py
index 0a47e6ac7..371c19bc1 100644
--- a/research/fraginternals.py
+++ b/research/fraginternals.py
@@ -928,7 +928,7 @@ class Authenticator(Daemon):
 	def handle_wpaspy(self, msg):
 		log(DEBUG, "daemon: " + msg)
 
-		if "AP-STA-CONNECTING" in msg:
+		if "AP-STA-NEW" in msg:
 			cmd, clientmac = msg.split()
 			self.add_station(clientmac)
 
@@ -1111,29 +1111,26 @@ class Supplicant(Daemon):
 	def handle_wpaspy(self, msg):
 		log(DEBUG, "daemon: " + msg)
 
-		if "WPA: Key negotiation completed with" in msg:
-			# This get's the current keys
-			self.station.handle_authenticated()
-
-		elif "Trying to authenticate with" in msg:
+		if "Associated with" in msg:
 			# When using a separate interface to inject, switch to correct channel
 			self.follow_channel()
 
-			p = re.compile("Trying to authenticate with (.*) \(SSID")
-			bss = p.search(msg).group(1)
-			self.station.handle_connecting(bss)
-
-		elif "Trying to associate with" in msg:
 			# With the ath9k_htc, injection in mixed managed/monitor only works after
 			# sending the association request. So only perform injection test now.
 			self.injection_test(self.station.bss)
 
+			p = re.compile("Associated with (.*) successfully")
+			bss = p.search(msg).group(1)
+			self.station.handle_connecting(bss)
+
 		elif "EAPOL-TX" in msg:
 			cmd, srcaddr, payload = msg.split()
 			self.station.handle_eapol_tx(bytes.fromhex(payload))
 
-		# This event only occurs with WEP
-		elif "WPA: EAPOL processing complete" in msg:
+		# The "EAPOL processing" event only occurs with WEP
+		if "WPA: Key negotiation completed with" in msg or \
+		   "WPA: EAPOL processing complete" in msg:
+			# This get's the current keys
 			self.station.handle_authenticated()
 
 	def roam(self, station):
diff --git a/src/ap/ieee802_11.c b/src/ap/ieee802_11.c
index 6074a9dab..379a1239d 100644
--- a/src/ap/ieee802_11.c
+++ b/src/ap/ieee802_11.c
@@ -2613,11 +2613,6 @@ static void handle_auth(struct hostapd_data *hapd,
 #endif /* CONFIG_FILS */
 	}
 
-#ifdef CONFIG_TESTING_OPTIONS
-	wpa_msg(hapd->msg_ctx, MSG_INFO, "AP-STA-CONNECTING " MACSTR,
-		 MAC2STR(sta->addr));
-#endif /* CONFIG_TESTING_OPTIONS */
-
  fail:
 	reply_res = send_auth_reply(hapd, sta, mgmt->sa, mgmt->bssid, auth_alg,
 				    auth_transaction + 1, resp, resp_ies,
diff --git a/src/ap/sta_info.c b/src/ap/sta_info.c
index 438e512ce..aa31bb38e 100644
--- a/src/ap/sta_info.c
+++ b/src/ap/sta_info.c
@@ -725,6 +725,11 @@ struct sta_info * ap_sta_add(struct hostapd_data *hapd, const u8 *addr)
 	sta->last_seq_ctrl = WLAN_INVALID_MGMT_SEQ;
 	dl_list_init(&sta->ip6addr);
 
+#ifdef CONFIG_TESTING_OPTIONS
+	wpa_msg(hapd->msg_ctx, MSG_INFO, "AP-STA-NEW " MACSTR,
+		 MAC2STR(sta->addr));
+#endif /* CONFIG_TESTING_OPTIONS */
+
 #ifdef CONFIG_TAXONOMY
 	sta_track_claim_taxonomy_info(hapd->iface, addr,
 				      &sta->probe_ie_taxonomy);
diff --git a/wpa_supplicant/events.c b/wpa_supplicant/events.c
index e3fce8f72..2d7c5b3c1 100644
--- a/wpa_supplicant/events.c
+++ b/wpa_supplicant/events.c
@@ -2959,6 +2959,10 @@ static void wpa_supplicant_event_assoc(struct wpa_supplicant *wpa_s,
 	wpa_s->sme.last_unprot_disconnect.sec = 0;
 #endif /* CONFIG_SME */
 
+#ifdef CONFIG_TESTING_OPTIONS
+	wpa_msg_ctrl(wpa_s, MSG_INFO, "Associated with " MACSTR
+			" successfully", MAC2STR(bssid));
+#endif /* CONFIG_TESTING_OPTIONS */
 	wpa_msg(wpa_s, MSG_INFO, "Associated with " MACSTR, MAC2STR(bssid));
 	if (wpa_s->current_ssid) {
 		/* When using scanning (ap_scan=1), SIM PC/SC interface can be