From 163281cc28409722ef0713ed450e9bdac4fa87d8 Mon Sep 17 00:00:00 2001
From: Yaro Kasear <git@kasear.net>
Date: Mon, 14 Apr 2025 13:42:55 -0500
Subject: [PATCH] Refactor packet filtering in analyze_pcap to improve
 timestamp handling and client retrieval

---
 enrich.py  | 38 +++++++++++++++++++++++++-------------
 runtest.sh |  8 ++++----
 2 files changed, 29 insertions(+), 17 deletions(-)

diff --git a/enrich.py b/enrich.py
index 59488fc..8a3fda7 100755
--- a/enrich.py
+++ b/enrich.py
@@ -149,32 +149,44 @@ def get_clients_on_channel(capture, ap_channel, ap_bssid):
     return len(clients)
 
 def analyze_pcap(pcapng_path, start_ts, end_ts, ap_bssid, ap_channel):
-    
     cap = pyshark.FileCapture(
         pcapng_path,
-        display_filter=f'frame.time_epoch >= {start_ts} && frame.time_epoch <= {end_ts}',
         use_json=True,
         include_raw=False
     )
 
-    # Get clients on the specified AP
-    clients_on_ap = get_clients_on_ap(cap, ap_bssid.lower())
+    ap_channel = int(ap_channel)
 
-    # Get clients on the specified channel
+    clients_on_ap = 0
+    clients_on_channel = 0
 
-    clients_on_channel = get_clients_on_channel(cap, ap_channel, ap_bssid.lower())
+    try:
+        # Filter packets manually by timestamp
+        filtered_packets = []
+        for packet in cap:
+            try:
+                frame_time = float(packet.frame_info.time_epoch)
+                if start_ts <= frame_time <= end_ts:
+                    filtered_packets.append(packet)
+            except Exception:
+                continue
 
-    # Placeholder: Logic will be added for:
-    # - APsOnChannel
-    # - CongestionScore
-    # - AvgAPSignal
-    # - StrongestAPSignal
-    # - UnlinkedDevices
+        clients_on_ap = get_clients_on_ap(filtered_packets, ap_bssid)
+        clients_on_channel = get_clients_on_channel(filtered_packets, ap_channel, ap_bssid)
 
-    cap.close()
+        # Placeholder: Logic will be added for:
+        # - APsOnChannel
+        # - CongestionScore
+        # - AvgAPSignal
+        # - StrongestAPSignal
+        # - UnlinkedDevices
+        
+    finally:
+        cap.close()
 
     return clients_on_ap, clients_on_channel, 0, None, None, None, 0
 
+
 def main():
     args = parse_args()
     cap = pyshark.FileCapture(
diff --git a/runtest.sh b/runtest.sh
index 1725fd2..ef6c562 100755
--- a/runtest.sh
+++ b/runtest.sh
@@ -109,6 +109,10 @@ while [ "$COUNTER" -lt "$NUM_TESTS" ]; do
   fi
 done
 
+echo "Stopping kismet..."
+
+sudo systemctl stop kismet
+
 # Let's enrich the data with passive metrics.
 
 echo "Enriching the data..."
@@ -120,10 +124,6 @@ if [ -z "$KISMET_LOG" ] || [ ! -f "$KISMET_LOG" ]; then
   exit 1
 fi
 
-echo "Stopping kismet..."
-
-sudo systemctl stop kismet
-
 python3 $SCRIPT_DIRECTORY/enrich.py --csv $TEST_FILE --pcapng "$KISMET_LOG" --output "$ENRICHED_FILE"
 
 # Final email with attachment