Refactor query_metrics to utilize packet-specific data for client and AP counts, enhancing accuracy and performance

enrichment/indexed_capture.py

@@ -140,12 +140,11 @@ class IndexedCapture:
         packets = self.get_packets_in_time_range(start_ts, end_ts)
         print(f"[DEBUG] Packets in window: {len(packets)} between {start_ts} and {end_ts}")
 
-        # Use indexed data instead of recalculating
+        # Instead of global stats, base everything off `packets`
         clients_on_ap = self._count_clients_on_ap(packets, ap_bssid)
-        clients_on_channel = len(self.channel_to_clients.get(ap_channel, []))
+        clients_on_channel = self._count_clients_on_channel(packets, ap_channel)
-        aps_on_channel = len(self.channel_to_aps.get(ap_channel, []))
+        aps_on_channel = self._count_aps_on_channel(packets, ap_channel)
-
+        avg_ap_signal, max_ap_signal = self._calc_signal_stats_window(packets, ap_channel)
-        avg_ap_signal, max_ap_signal = self._calc_signal_stats(ap_channel)
         unlinked_devices = self._count_unlinked_devices(packets, ap_channel)
 
         our_ssid = self.bssid_to_ssid.get(ap_bssid)
@@ -163,6 +162,7 @@ class IndexedCapture:
             num_channels_ssid, packet_count
         )
 
+
     def _count_clients_on_ap(self, packets, ap_bssid):
         clients = defaultdict(int)
         ap_bssid = ap_bssid.lower()
@@ -222,3 +222,92 @@ class IndexedCapture:
         if ssid in self.cisco_ssid_clients:
             return max(self.cisco_ssid_clients[ssid])
         return 0
+    
+    def _count_clients_on_channel(self, packets, ap_channel):
+        clients = set()
+        for packet in packets:
+            try:
+                if 'radiotap' not in packet or 'wlan' not in packet:
+                    continue
+                
+                radio = packet.radiotap
+                wlan = packet.wlan
+    
+                if not hasattr(radio, 'channel') or not hasattr(radio.channel, 'freq'):
+                    continue
+                
+                freq = int(radio.channel.freq)
+                packet_channel = get_channel_from_freq(freq)
+    
+                if packet_channel != ap_channel:
+                    continue
+                
+                sa = getattr(wlan, 'sa', '').lower()
+                da = getattr(wlan, 'da', '').lower()
+    
+                for mac in (sa, da):
+                    if mac and mac != 'ff:ff:ff:ff:ff:ff':
+                        clients.add(mac)
+            except Exception:
+                continue
+            
+        return len(clients)
+    
+    def _count_aps_on_channel(self, packets, ap_channel):
+        aps = set()
+        for packet in packets:
+            try:
+                if 'radiotap' not in packet or 'wlan' not in packet:
+                    continue
+                
+                radio = packet.radiotap
+                wlan = packet.wlan
+    
+                if not hasattr(radio, 'channel') or not hasattr(radio.channel, 'freq'):
+                    continue
+                
+                freq = int(radio.channel.freq)
+                packet_channel = get_channel_from_freq(freq)
+    
+                if packet_channel != ap_channel:
+                    continue
+                
+                subtype = int(getattr(wlan, 'type_subtype', '0'), 16)
+                if subtype not in (5, 8):  # Beacon or Probe Response
+                    continue
+                
+                bssid = getattr(wlan, 'bssid', '').lower()
+                if bssid and bssid != 'ff:ff:ff:ff:ff:ff':
+                    aps.add(bssid)
+    
+            except Exception:
+                continue
+            
+        return len(aps)
+    
+    def _calc_signal_stats_window(self, packets, ap_channel):
+        signals = []
+        for packet in packets:
+            try:
+                if 'radiotap' not in packet or 'wlan' not in packet:
+                    continue
+                
+                radio = packet.radiotap
+                wlan = packet.wlan
+    
+                if not hasattr(radio, 'channel') or not hasattr(radio.channel, 'freq'):
+                    continue
+                
+                freq = int(radio.channel.freq)
+                packet_channel = get_channel_from_freq(freq)
+    
+                if packet_channel != ap_channel:
+                    continue
+                
+                signal = getattr(radio, 'dbm_antsignal', None)
+                if signal is not None:
+                    signals.append(int(signal))
+            except Exception:
+                continue
+            
+        return (mean(signals), max(signals)) if signals else (0, 0)