From 0a14ecebf47f64dc1afe86ad8ddd9b2c170a8954 Mon Sep 17 00:00:00 2001
From: Yaro Kasear <git@kasear.net>
Date: Tue, 15 Apr 2025 08:33:00 -0500
Subject: [PATCH] Refactor get_aps_on_channel function to improve packet
 filtering and BSSID handling

---
 enrich.py | 45 +++++++++++++++++----------------------------
 1 file changed, 17 insertions(+), 28 deletions(-)

diff --git a/enrich.py b/enrich.py
index d9123d1..b7018a4 100755
--- a/enrich.py
+++ b/enrich.py
@@ -155,52 +155,41 @@ def get_aps_on_channel(capture, ap_channel):
         print(f"[!] Could not parse channel number: {ap_channel}")
         return 0
 
+    from_channel_freq = get_channel_from_freq(ap_channel)
+    if not from_channel_freq:
+        print(f"[!] Invalid channel: {ap_channel}")
+        return 0
+
     aps = set()
 
     for packet in capture:
         try:
-            if not hasattr(packet, 'radiotap') or not hasattr(packet, 'wlan'):
+            if 'radiotap' not in packet or 'wlan' not in packet:
                 continue
 
-            packet_freq = getattr(packet.radiotap.channel, 'freq', None)
-            if packet_freq is None:
-                print("[DEBUG] No channel frequency found.")
+            # Match channel frequency
+            freq = int(getattr(packet.radiotap, 'channel_freq', -1))
+            if freq != from_channel_freq:
                 continue
 
-            packet_freq = int(packet_freq)
-            packet_channel = get_channel_from_freq(packet_freq)
-
-            if packet_channel != ap_channel:
-                print(f"[DEBUG] Skipped packet on channel {packet_channel}, looking for {ap_channel}")
+            # Check for beacon or probe response
+            ts_hex = getattr(packet.wlan, 'type_subtype', None)
+            if ts_hex is None:
                 continue
 
-            fc_raw = getattr(packet.wlan, 'fc', None)
-            if not fc_raw:
-                print("[DEBUG] No FC field found.")
-                continue
-
-            fc_int = int(fc_raw, 16)
-            frame_type = (fc_int >> 2) & 0b11
-            subtype = (fc_int >> 4) & 0b1111
-
-            print(f"[DEBUG] FC: {fc_raw}, frame_type: {frame_type}, subtype: {subtype}")
-
-            if frame_type != 0 or subtype not in (5, 8):
-                print("[DEBUG] Not a Beacon or Probe Response.")
+            ts = int(ts_hex, 16)
+            if ts not in (5, 8):  # Probe Response or Beacon
                 continue
 
+            # Grab BSSID
             bssid = getattr(packet.wlan, 'bssid', '').lower()
-            if bssid:
-                print(f"[DEBUG] Adding BSSID: {bssid}")
+            if bssid and bssid != 'ff:ff:ff:ff:ff:ff':
                 aps.add(bssid)
-            else:
-                print("[DEBUG] No BSSID found.")
 
         except Exception as e:
-            print(f"[DEBUG] Failed packet: {e}")
+            print(f"[DEBUG] Packet parse error: {e}")
             continue
 
-    print(f"[DEBUG] Final AP count: {len(aps)} | APs: {sorted(aps)}")
     return len(aps)
 
 def analyze_pcap(pcapng_path, start_ts, end_ts, ap_bssid, ap_channel):