From 2a227d722285314a2fa0772c3eec9f19bb11a007 Mon Sep 17 00:00:00 2001
From: Yaro Kasear <git@kasear.net>
Date: Wed, 30 Apr 2025 13:50:59 -0500
Subject: [PATCH] Add support for --no-enrich flag in runtest.sh to skip data
 enrichment

---
 listener.py | 83 +++++++++++++++++++++++++++++++++++++++++++++++++++++
 runtest.sh  | 30 +++++++++++++++----
 2 files changed, 107 insertions(+), 6 deletions(-)
 create mode 100644 listener.py

diff --git a/listener.py b/listener.py
new file mode 100644
index 0000000..fcf0d39
--- /dev/null
+++ b/listener.py
@@ -0,0 +1,83 @@
+#!/usr/bin/env python3
+import os
+import signal
+import time
+from collections import defaultdict
+from scapy.all import sniff, Dot11, RadioTap
+from dotenv import load_dotenv
+
+# === Load ENV ===
+load_dotenv(dotenv_path=os.path.expanduser("~/wifi_test/settings.env"))
+INTERFACE = os.getenv("INTERFACE", "wlan0")
+
+# === Globals ===
+clients_per_channel = defaultdict(set)
+aps_per_channel = defaultdict(set)
+running = True
+
+def get_channel(pkt):
+    # Convert frequency to channel
+    if not pkt.haslayer(RadioTap):
+        return None
+    try:
+        freq = pkt[RadioTap].ChannelFrequency
+        if 2412 <= freq <= 2472:
+            return (freq - 2407) // 5
+        elif freq == 2484:
+            return 14
+        elif 5180 <= freq <= 5825:
+            return (freq - 5000) // 5
+    except:
+        return None
+    return None
+
+def handle_packet(pkt):
+    if not pkt.haslayer(Dot11):
+        return
+    ch = get_channel(pkt)
+    if ch is None:
+        return
+
+    dot11 = pkt[Dot11]
+
+    if dot11.type == 0 and dot11.subtype in (5, 8):  # Probe Response or Beacon
+        aps_per_channel[ch].add(dot11.addr2)
+    else:
+        if dot11.addr1:
+            clients_per_channel[ch].add(dot11.addr1)
+        if dot11.addr2:
+            clients_per_channel[ch].add(dot11.addr2)
+
+def print_stats():
+    print("\n[+] Live Wi-Fi Stats")
+    for ch in sorted(set(clients_per_channel) | set(aps_per_channel)):
+        c = len(clients_per_channel[ch])
+        a = len(aps_per_channel[ch])
+        print(f"  - Channel {ch}: {c} clients, {a} APs")
+    print("-" * 40)
+
+def stop_sniff(signum, frame):
+    global running
+    print("\n[!] Caught Ctrl+C, exiting...")
+    running = False
+
+# === Main ===
+if __name__ == "__main__":
+    signal.signal(signal.SIGINT, stop_sniff)
+    print(f"[+] Listening on interface {INTERFACE} (press Ctrl+C to stop)")
+
+    # Start sniffing in a background thread
+    sniff_thread = sniff(
+        iface=INTERFACE,
+        prn=handle_packet,
+        store=False,
+        monitor=True,
+        timeout=1
+    )
+
+    try:
+        while running:
+            print_stats()
+            time.sleep(10)
+    except KeyboardInterrupt:
+        pass
diff --git a/runtest.sh b/runtest.sh
index ded8159..31a2553 100755
--- a/runtest.sh
+++ b/runtest.sh
@@ -4,6 +4,20 @@ set -uo pipefail
 trap 'echo "[✖] Execution halted at line $LINENO. Please consult your nearest bash therapist." >&2' ERR
 trap 'echo "[✌️] Script exited cleanly. Have a burrito." >&2' EXIT
 
+NO_ENRICH=0
+for arg in "$@"; do
+  case $arg in
+    --no-enrich)
+      NO_ENRICH=1
+      shift
+      ;;
+    *)
+      echo "Unknown argument: $arg"
+      exit 1
+      ;;
+  esac
+done
+
 IFS=$'\n\t'
 
 # === Start logging with 'script' safely ===
@@ -188,13 +202,17 @@ done
 log "Stopping kismet..."
 sudo systemctl stop kismet
 
-log "Enriching data..."
-KISMET_LOG=$(find "$KISMET_LOG_DIR" -type f -name "*.pcapng" -printf "%T@ %p\n" | sort -n | tail -1 | cut -d' ' -f2-)
-[ ! -f "$KISMET_LOG" ] && die "Packet capture not found."
+if [ "$NO_ENRICH" -eq 0 ]; then
+  log "Enriching data..."
+  KISMET_LOG=$(find "$KISMET_LOG_DIR" -type f -name "*.pcapng" -printf "%T@ %p\n" | sort -n | tail -1 | cut -d' ' -f2-)
+  [ ! -f "$KISMET_LOG" ] && die "Packet capture not found."
 
-SECONDS=0
-python3 "$SCRIPT_DIRECTORY/enrich.py" --csv "$TEST_FILE" --pcapng "$KISMET_LOG" --output "$ENRICHED_FILE"
-log "Enrichment took $SECONDS seconds"
+  SECONDS=0
+  python3 "$SCRIPT_DIRECTORY/enrich.py" --csv "$TEST_FILE" --pcapng "$KISMET_LOG" --output "$ENRICHED_FILE"
+  log "Enrichment took $SECONDS seconds"
+else
+  log "Skipping enrichment as per --no-enrich flag."
+fi
 
 ATTACHMENTS=()
 [ -f "$ENRICHED_FILE" ] && ATTACHMENTS+=("$ENRICHED_FILE")