Add support for --no-enrich flag in runtest.sh to skip data enrichment

This commit is contained in:
Yaro Kasear 2025-04-30 13:50:59 -05:00
parent 027890896e
commit 2a227d7222
2 changed files with 107 additions and 6 deletions

83
listener.py Normal file
View file

@ -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

View file

@ -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")