From 56a04ae1a1cf2cde81cd3f33c834a3705f5a738d Mon Sep 17 00:00:00 2001
From: Jouni Malinen <j@w1.fi>
Date: Sun, 7 Feb 2021 11:37:58 +0200
Subject: [PATCH] wlantest: Support TK list for Management frame decryption

Use the TKs from the PTK file (-T command line argument) to try to
decrypt encrypted Management frames if no BSS/STA key can be found based
on addresses.

Signed-off-by: Jouni Malinen <j@w1.fi>
---
 wlantest/rx_data.c  | 13 --------
 wlantest/rx_mgmt.c  | 76 ++++++++++++++++++++++++++++++++++++++-------
 wlantest/wlantest.c | 13 ++++++++
 wlantest/wlantest.h |  2 ++
 4 files changed, 80 insertions(+), 24 deletions(-)

diff --git a/wlantest/rx_data.c b/wlantest/rx_data.c
index b632013a8..aedf9e824 100644
--- a/wlantest/rx_data.c
+++ b/wlantest/rx_data.c
@@ -119,19 +119,6 @@ static void rx_data_process(struct wlantest *wt, const u8 *bssid,
 }
 
 
-static void write_decrypted_note(struct wlantest *wt, const u8 *decrypted,
-				 const u8 *tk, size_t tk_len, int keyid)
-{
-	char tk_hex[65];
-
-	if (!decrypted)
-		return;
-
-	wpa_snprintf_hex(tk_hex, sizeof(tk_hex), tk, tk_len);
-	add_note(wt, MSG_EXCESSIVE, "TK[%d] %s", keyid, tk_hex);
-}
-
-
 static u8 * try_ptk(int pairwise_cipher, struct wpa_ptk *ptk,
 		    const struct ieee80211_hdr *hdr,
 		    const u8 *data, size_t data_len, size_t *decrypted_len)
diff --git a/wlantest/rx_mgmt.c b/wlantest/rx_mgmt.c
index 61c7f6a4c..2b0580752 100644
--- a/wlantest/rx_mgmt.c
+++ b/wlantest/rx_mgmt.c
@@ -2139,6 +2139,56 @@ static int check_bip(struct wlantest *wt, const u8 *data, size_t len)
 }
 
 
+static u8 * try_tk(struct wpa_ptk *ptk, const u8 *data, size_t len,
+		   size_t *dlen)
+{
+	const struct ieee80211_hdr *hdr;
+	u8 *decrypted, *frame;
+
+	hdr = (const struct ieee80211_hdr *) data;
+	decrypted = ccmp_decrypt(ptk->tk, hdr, data + 24, len - 24, dlen);
+	if (!decrypted)
+		return NULL;
+
+	frame = os_malloc(24 + *dlen);
+	if (frame) {
+		os_memcpy(frame, data, 24);
+		os_memcpy(frame + 24, decrypted, *dlen);
+		*dlen += 24;
+	}
+	os_free(decrypted);
+	return frame;
+}
+
+
+static u8 * mgmt_ccmp_decrypt_tk(struct wlantest *wt, const u8 *data,
+				 size_t len, size_t *dlen)
+{
+	struct wlantest_ptk *ptk;
+	u8 *decrypted;
+	int prev_level = wpa_debug_level;
+	int keyid;
+
+	keyid = data[24 + 3] >> 6;
+
+	wpa_debug_level = MSG_WARNING;
+	dl_list_for_each(ptk, &wt->ptk, struct wlantest_ptk, list) {
+		decrypted = try_tk(&ptk->ptk, data, len, dlen);
+		if (decrypted) {
+			wpa_debug_level = prev_level;
+			add_note(wt, MSG_DEBUG,
+				 "Found TK match from the list of all known TKs");
+			write_decrypted_note(wt, decrypted, ptk->ptk.tk,
+					     ptk->ptk.tk_len, keyid);
+			return decrypted;
+		}
+	}
+	wpa_debug_level = prev_level;
+
+	return NULL;
+}
+
+
 static u8 * mgmt_ccmp_decrypt(struct wlantest *wt, const u8 *data, size_t len,
 			      size_t *dlen)
 {
@@ -2150,17 +2200,6 @@ static u8 * mgmt_ccmp_decrypt(struct wlantest *wt, const u8 *data, size_t len,
 	u8 pn[6], *rsc;
 
 	hdr = (const struct ieee80211_hdr *) data;
-	bss = bss_get(wt, hdr->addr3);
-	if (bss == NULL)
-		return NULL;
-	if (os_memcmp(hdr->addr1, hdr->addr3, ETH_ALEN) == 0)
-		sta = sta_get(bss, hdr->addr2);
-	else
-		sta = sta_get(bss, hdr->addr1);
-	if (sta == NULL || !sta->ptk_set) {
-		add_note(wt, MSG_MSGDUMP, "No PTK known to decrypt the frame");
-		return NULL;
-	}
 
 	if (len < 24 + 4)
 		return NULL;
@@ -2184,6 +2223,21 @@ static u8 * mgmt_ccmp_decrypt(struct wlantest *wt, const u8 *data, size_t len,
 			 MACSTR, keyid, MAC2STR(hdr->addr2));
 	}
 
+	bss = bss_get(wt, hdr->addr3);
+	if (bss == NULL)
+		return mgmt_ccmp_decrypt_tk(wt, data, len, dlen);
+	if (os_memcmp(hdr->addr1, hdr->addr3, ETH_ALEN) == 0)
+		sta = sta_get(bss, hdr->addr2);
+	else
+		sta = sta_get(bss, hdr->addr1);
+	if (sta == NULL || !sta->ptk_set) {
+		decrypted = mgmt_ccmp_decrypt_tk(wt, data, len, dlen);
+		if (!decrypted)
+			add_note(wt, MSG_MSGDUMP,
+				 "No PTK known to decrypt the frame");
+		return decrypted;
+	}
+
 	if (os_memcmp(hdr->addr1, hdr->addr3, ETH_ALEN) == 0)
 		rsc = sta->rsc_tods[16];
 	else
diff --git a/wlantest/wlantest.c b/wlantest/wlantest.c
index 1b8d7147e..62c89e226 100644
--- a/wlantest/wlantest.c
+++ b/wlantest/wlantest.c
@@ -323,6 +323,19 @@ size_t notes_len(struct wlantest *wt, size_t hdrlen)
 }
 
 
+void write_decrypted_note(struct wlantest *wt, const u8 *decrypted,
+			  const u8 *tk, size_t tk_len, int keyid)
+{
+	char tk_hex[65];
+
+	if (!decrypted)
+		return;
+
+	wpa_snprintf_hex(tk_hex, sizeof(tk_hex), tk, tk_len);
+	add_note(wt, MSG_EXCESSIVE, "TK[%d] %s", keyid, tk_hex);
+}
+
+
 int wlantest_relog(struct wlantest *wt)
 {
 	int ret = 0;
diff --git a/wlantest/wlantest.h b/wlantest/wlantest.h
index 0c266f4e9..658a3a06d 100644
--- a/wlantest/wlantest.h
+++ b/wlantest/wlantest.h
@@ -232,6 +232,8 @@ void add_note(struct wlantest *wt, int level, const char *fmt, ...)
 PRINTF_FORMAT(3, 4);
 void clear_notes(struct wlantest *wt);
 size_t notes_len(struct wlantest *wt, size_t hdrlen);
+void write_decrypted_note(struct wlantest *wt, const u8 *decrypted,
+			  const u8 *tk, size_t tk_len, int keyid);
 
 int add_wep(struct wlantest *wt, const char *key);
 int read_cap_file(struct wlantest *wt, const char *fname);