#!/usr/bin/env python3
import os
import signal
import time
from collections import defaultdict
from scapy.all import sniff, Dot11, RadioTap
from dotenv import dotenv_values

# === Load ENV ===
config = dotenv_values(os.path.expanduser("~/wifi_test/settings.env"))
LISTEN_INTERFACE = config.get("LISTEN_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 {LISTEN_INTERFACE} (press Ctrl+C to stop)")

    # Start sniffing in a background thread
    try:
        while running:
            sniff(
                iface=LISTEN_INTERFACE,
                prn=handle_packet,
                store=False,
                monitor=True,
                timeout=10
            )
            print_stats()
    except KeyboardInterrupt:
        pass