#!/usr/bin/env python3
"""Wireless plumbing: interface discovery, monitor mode, card-state restore,
and tracked airodump-ng / aireplay-ng subprocesses."""
import glob
import os
import re
import shutil
import subprocess
import sys
import tempfile
import time

SCAN_NET_SECS = 10
SCAN_CLIENT_SECS = 15

# --- child process registry (so "Stop & Reset" can kill mid-scan/mid-deauth) ---
_procs = set()


def _spawn(cmd, **kw):
    p = subprocess.Popen(cmd, **kw)
    _procs.add(p)
    return p


def _reap(p):
    _procs.discard(p)


def kill_all_children():
    for p in list(_procs):
        try:
            if p.poll() is None:
                p.terminate()
        except Exception:
            pass
    time.sleep(0.5)
    for p in list(_procs):
        try:
            if p.poll() is None:
                p.kill()
        except Exception:
            pass
        _procs.discard(p)


def run(cmd, timeout=30):
    return subprocess.run(cmd, capture_output=True, text=True, timeout=timeout)


def list_wifi_interfaces():
    """[{iface, phy, type, addr, monitor_capable}] from `iw dev`."""
    try:
        out = run(["iw", "dev"]).stdout
    except Exception:
        return []
    ifaces, cur_phy, cur = [], None, None
    for line in out.splitlines():
        m = re.match(r"phy#(\d+)", line)
        if m:
            cur_phy = "phy" + m.group(1)
        m = re.match(r"\s+Interface (\S+)", line)
        if m:
            if cur:
                ifaces.append(cur)
            cur = {"iface": m.group(1), "phy": cur_phy, "type": "?", "addr": "?"}
        m = re.match(r"\s+addr (\S+)", line)
        if m and cur:
            cur["addr"] = m.group(1)
        m = re.match(r"\s+type (\S+)", line)
        if m and cur:
            cur["type"] = m.group(1)
    if cur:
        ifaces.append(cur)
    for it in ifaces:
        it["monitor_capable"] = _phy_has_monitor(it["phy"])
    return ifaces


def _phy_has_monitor(phy):
    if not phy:
        return False
    out = run(["iw", phy, "info"]).stdout
    grab = False
    for line in out.splitlines():
        if "Supported interface modes:" in line:
            grab = True
            continue
        if grab:
            if re.match(r"\s+\*", line):
                if "monitor" in line:
                    return True
            else:
                break
    return False


class CardState:
    def __init__(self, iface):
        self.orig_iface = iface["iface"]
        self.orig_addr = iface["addr"]
        self.nm_was_active = (
            run(["systemctl", "is-active", "NetworkManager"]).stdout.strip() == "active"
        )
        self.mon_iface = None
        self.used_airmon = False


def enable_monitor(state):
    run(["airmon-ng", "check", "kill"])
    before = {i["iface"] for i in list_wifi_interfaces()}
    res = run(["airmon-ng", "start", state.orig_iface])
    state.used_airmon = True
    time.sleep(1)
    after = list_wifi_interfaces()
    new = {i["iface"] for i in after} - before
    if new:
        state.mon_iface = new.pop()
    else:
        m = re.search(r"enabled (?:for \[[^\]]+\]\S+ )?on (?:\[[^\]]+\])?(\S+?)\)", res.stdout)
        if m:
            state.mon_iface = m.group(1)
        else:
            for i in after:
                if i["iface"] == state.orig_iface and i["type"] == "monitor":
                    state.mon_iface = state.orig_iface
    if not state.mon_iface:
        raise RuntimeError("Could not determine monitor interface:\n" + res.stdout + res.stderr)
    return state.mon_iface


def reset_card(state):
    if not state:
        return
    kill_all_children()
    if state.mon_iface and state.used_airmon:
        run(["airmon-ng", "stop", state.mon_iface])
    names = {i["iface"]: i for i in list_wifi_interfaces()}
    if state.orig_iface not in names and state.mon_iface in names:
        run(["iw", state.mon_iface, "set", "type", "managed"])
        run(["ip", "link", "set", state.mon_iface, "name", state.orig_iface])
    names = {i["iface"]: i for i in list_wifi_interfaces()}
    it = names.get(state.orig_iface)
    if it and state.orig_addr not in ("?", None) and it["addr"].lower() != state.orig_addr.lower():
        run(["ip", "link", "set", state.orig_iface, "down"])
        run(["ip", "link", "set", state.orig_iface, "address", state.orig_addr])
        run(["ip", "link", "set", state.orig_iface, "up"])
    if state.nm_was_active:
        run(["systemctl", "start", "NetworkManager"])


def airodump_scan(mon_iface, secs, bssid=None, channel=None):
    tmpdir = tempfile.mkdtemp(prefix="wifiscan_")
    prefix = os.path.join(tmpdir, "dump")
    cmd = ["airodump-ng", "-w", prefix, "--output-format", "csv", "--write-interval", "1"]
    if bssid:
        cmd += ["--bssid", bssid]
    if channel:
        cmd += ["-c", str(channel)]
    cmd += [mon_iface]
    p = _spawn(cmd, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
    # fixed scan duration is a naive heuristic; longer = more results.
    t0 = time.time()
    while time.time() - t0 < secs and p.poll() is None:
        time.sleep(0.3)
    p.terminate()
    try:
        p.wait(timeout=5)
    except subprocess.TimeoutExpired:
        p.kill()
    _reap(p)
    csvs = sorted(glob.glob(prefix + "*.csv"))
    aps, stations = ([], [])
    if csvs:
        aps, stations = parse_airodump_csv(csvs[-1])
    shutil.rmtree(tmpdir, ignore_errors=True)
    return aps, stations


def parse_airodump_csv(path):
    try:
        with open(path, encoding="utf-8", errors="replace") as f:
            text = f.read()
    except OSError:
        return [], []
    aps, stations, section = [], [], None
    for line in text.splitlines():
        s = line.strip()
        if not s:
            continue
        if s.startswith("BSSID,"):
            section = "ap"
            continue
        if s.startswith("Station MAC,"):
            section = "sta"
            continue
        # naive comma split; an ESSID containing a comma will skew fields.
        parts = [c.strip() for c in line.split(",")]
        if section == "ap" and len(parts) >= 14:
            aps.append({
                "bssid": parts[0], "channel": parts[3], "power": parts[8],
                "privacy": parts[5], "essid": parts[13],
            })
        elif section == "sta" and len(parts) >= 6:
            stations.append({
                "station": parts[0], "power": parts[3],
                "packets": parts[4], "bssid": parts[5],
            })
    return aps, stations


class Deauth:
    def __init__(self):
        self.proc = None

    def start(self, mon_iface, bssid, client):
        cmd = ["aireplay-ng", "--deauth", "0", "-a", bssid, "-c", client, mon_iface]
        self.proc = _spawn(cmd, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)

    def stop(self):
        p = self.proc
        if p and p.poll() is None:
            p.terminate()
            try:
                p.wait(timeout=5)
            except subprocess.TimeoutExpired:
                p.kill()
        if p:
            _reap(p)
        self.proc = None

    def running(self):
        return self.proc is not None and self.proc.poll() is None


def _selftest():
    sample = (
        "BSSID, First time seen, Last time seen, channel, Speed, Privacy, Cipher, "
        "Authentication, Power, # beacons, # IV, LAN IP, ID-length, ESSID, Key\r\n"
        "\r\n"
        "AA:BB:CC:DD:EE:FF, 2024-01-01 00:00:00, 2024-01-01 00:01:00,  6,  130, "
        "WPA2, CCMP, PSK, -40, 10, 0, 0.0.0.0, 6, MyWiFi, \r\n"
        "\r\n"
        "Station MAC, First time seen, Last time seen, Power, # packets, BSSID, "
        "Probed ESSIDs\r\n"
        "\r\n"
        "11:22:33:44:55:66, 2024-01-01 00:00:10, 2024-01-01 00:01:00, -50, 25, "
        "AA:BB:CC:DD:EE:FF, MyWiFi\r\n"
        "99:88:77:66:55:44, 2024-01-01 00:00:10, 2024-01-01 00:01:00, -60, 5, "
        "(not associated), Other\r\n"
    )
    p = os.path.join(tempfile.gettempdir(), "wd_selftest.csv")
    with open(p, "w") as f:
        f.write(sample)
    aps, sts = parse_airodump_csv(p)
    os.remove(p)
    assert len(aps) == 1, aps
    assert aps[0]["essid"] == "MyWiFi" and aps[0]["channel"] == "6", aps
    assert len(sts) == 2, sts
    assoc = [s for s in sts if s["bssid"].lower() == "aa:bb:cc:dd:ee:ff"]
    assert len(assoc) == 1 and assoc[0]["station"] == "11:22:33:44:55:66", sts
    print("selftest ok")


if __name__ == "__main__":
    if "--selftest" in sys.argv:
        _selftest()
        sys.exit(0)