#!/usr/bin/env python3
"""
Wi-Fi Analyzer - Enhanced Version
Captures beacon frames, maps clients, detects security types,
tracks RSSI, vendor info, and exports results on exit.

Usage:
    sudo python3 wifi_analyzer.py -i wlan0mon
    sudo python3 wifi_analyzer.py -i wlan0mon --channels 1 6 11 --export
"""

import os
import time
import threading
import argparse
import logging
import subprocess
import json
from datetime import datetime
from dataclasses import dataclass, field, asdict
from typing import Dict, Optional

from scapy.all import sniff, Dot11, Dot11Beacon, Dot11Elt, RadioTap

# ---------------------------------------------------------
# Logging Setup
# ---------------------------------------------------------
logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s [%(levelname)s] %(message)s",
    datefmt="%H:%M:%S"
)
log = logging.getLogger(__name__)

# ---------------------------------------------------------
# OUI Vendor Lookup (lightweight, no external dependency)
# Install `manuf` for full DB: pip install manuf
# Falls back gracefully if not installed.
# ---------------------------------------------------------
try:
    from manuf import manuf
    _mac_parser = manuf.MacParser()
    def get_vendor(mac: str) -> str:
        try:
            v = _mac_parser.get_manuf(mac)
            return v if v else "Unknown"
        except Exception:
            return "Unknown"
except ImportError:
    def get_vendor(mac: str) -> str:
        return "N/A (install manuf)"

# ---------------------------------------------------------
# Data Classes
# ---------------------------------------------------------
@dataclass
class Network:
    ssid: str = "<Hidden>"
    channel: int = 0
    crypto: str = "Unknown"
    bandwidth: str = "Unknown"
    rssi: int = -100
    vendor: str = "Unknown"
    first_seen: str = ""
    last_seen: str = ""

@dataclass
class Client:
    connected_to_bssid: str = ""
    ssid: str = ""
    vendor: str = "Unknown"
    last_seen: str = ""

# Thread-safe storage
networks: Dict[str, Network] = {}
clients: Dict[str, Client] = {}
lock = threading.Lock()

# ---------------------------------------------------------
# Argument Parser
# ---------------------------------------------------------
def parse_args():
    parser = argparse.ArgumentParser(
        description="Live Wi-Fi Analyzer using Scapy",
        formatter_class=argparse.RawTextHelpFormatter
    )
    parser.add_argument(
        "-i", "--interface",
        required=True,
        help="Monitor mode interface (e.g., wlan0mon)"
    )
    parser.add_argument(
        "--channels",
        nargs="+",
        type=int,
        default=list(range(1, 14)) + [36, 40, 44, 48, 149, 153, 157, 161],
        help="Channels to hop (default: all 2.4GHz + common 5GHz)"
    )
    parser.add_argument(
        "--hop-interval",
        type=float,
        default=0.5,
        help="Seconds to dwell on each channel (default: 0.5)"
    )
    parser.add_argument(
        "--export",
        action="store_true",
        help="Export results to JSON/CSV on exit"
    )
    return parser.parse_args()

# ---------------------------------------------------------
# Packet Parsing Helpers
# ---------------------------------------------------------
def get_channel(pkt) -> int:
    """Extracts channel from Dot11Elt layer (ID=3)."""
    p = pkt[Dot11Elt]
    while isinstance(p, Dot11Elt):
        if p.ID == 3 and p.info:
            return p.info[0]  # Fixed: Python 3 bytes indexing
        p = p.payload
    return 0


def get_crypto(pkt) -> str:
    """
    Parses RSN/WPA IEs to determine security type.
    Detects: WPA3 (SAE), WPA2, WPA, WEP, Open.
    """
    crypto = set()
    p = pkt[Dot11Elt]

    while isinstance(p, Dot11Elt):
        # RSN Information Element (WPA2 / WPA3)
        if p.ID == 48 and p.info:
            rsn = p.info
            # AKM Suite list starts at offset 8 (after group/pairwise cipher counts)
            # SAE AKM = 00-0F-AC:08 → WPA3
            # PSK AKM = 00-0F-AC:02 / 00-0F-AC:04 → WPA2
            if b'\x00\x0f\xac\x08' in rsn:
                crypto.add("WPA3")
            else:
                crypto.add("WPA2")

        # WPA Information Element (Vendor Specific, Microsoft OUI)
        elif p.ID == 221 and p.info.startswith(b'\x00P\xf2\x01\x01\x00'):
            crypto.add("WPA")

        p = p.payload

    if not crypto:
        cap = pkt.sprintf("{Dot11Beacon:%Dot11Beacon.cap%}")
        if "privacy" in cap.lower():
            crypto.add("WEP")
        else:
            crypto.add("Open")

    return "/".join(sorted(crypto))


def get_bandwidth(pkt) -> str:
    """Detects HT/VHT/HE capability IEs to infer max bandwidth."""
    p = pkt[Dot11Elt]
    has_ht = has_vht = has_he = False

    while isinstance(p, Dot11Elt):
        if p.ID == 45:    # HT Capabilities → 802.11n / Wi-Fi 4
            has_ht = True
        elif p.ID == 191: # VHT Capabilities → 802.11ac / Wi-Fi 5
            has_vht = True
        elif p.ID == 255: # HE Capabilities → 802.11ax / Wi-Fi 6
            has_he = True
        p = p.payload

    if has_he:  return "Up to 160 MHz (Wi-Fi 6/6E)"
    if has_vht: return "Up to 160 MHz (Wi-Fi 5)"
    if has_ht:  return "Up to 40 MHz  (Wi-Fi 4)"
    return         "20 MHz         (Wi-Fi 3 or older)"


def get_rssi(pkt) -> int:
    """Extracts RSSI from RadioTap header."""
    if pkt.haslayer(RadioTap):
        try:
            return pkt[RadioTap].dBm_AntSignal
        except AttributeError:
            pass
    return -100

# ---------------------------------------------------------
# Packet Handler
# ---------------------------------------------------------
def packet_handler(pkt):
    now = datetime.now().strftime("%H:%M:%S")

    # --- Beacon Frames → Map Access Points ---
    if pkt.haslayer(Dot11Beacon):
        bssid = pkt[Dot11].addr2
        if not bssid:
            return

        try:
            ssid = pkt[Dot11Elt].info.decode(errors="ignore").strip()
            ssid = ssid if ssid and not ssid.startswith("\x00") else "<Hidden>"
        except Exception:
            ssid = "<Error>"

        channel   = get_channel(pkt)
        crypto    = get_crypto(pkt)
        bandwidth = get_bandwidth(pkt)
        rssi      = get_rssi(pkt)
        vendor    = get_vendor(bssid)

        with lock:
            if bssid in networks:
                net = networks[bssid]
                net.rssi      = rssi       # update signal strength
                net.last_seen = now
                net.channel   = channel    # AP may change channel
            else:
                networks[bssid] = Network(
                    ssid=ssid, channel=channel, crypto=crypto,
                    bandwidth=bandwidth, rssi=rssi, vendor=vendor,
                    first_seen=now, last_seen=now
                )

    # --- Data Frames → Map Connected Clients ---
    if pkt.haslayer(Dot11) and pkt.type == 2:
        DS       = pkt.FCfield & 0x3
        to_ds    = bool(DS & 0x01)
        from_ds  = bool(DS & 0x02)

        bssid = client = None

        if to_ds and not from_ds:      # STA → AP
            bssid, client = pkt.addr1, pkt.addr2
        elif not to_ds and from_ds:    # AP → STA
            bssid, client = pkt.addr2, pkt.addr1
        elif not to_ds and not from_ds: # Ad-hoc
            bssid, client = pkt.addr3, pkt.addr2

        if bssid and client:
            with lock:
                if bssid in networks:
                    ssid = networks[bssid].ssid
                    if client not in clients:
                        clients[client] = Client(
                            connected_to_bssid=bssid,
                            ssid=ssid,
                            vendor=get_vendor(client),
                            last_seen=now
                        )
                    else:
                        clients[client].last_seen = now

# ---------------------------------------------------------
# Channel Hopper
# ---------------------------------------------------------
def channel_hopper(interface: str, channels: list, interval: float):
    """Hops across channels so we capture the full spectrum."""
    log.info("Channel hopper started: %d channels, %.1fs interval", len(channels), interval)
    while True:
        for ch in channels:
            subprocess.run(
                ["iwconfig", interface, "channel", str(ch)],
                capture_output=True
            )
            time.sleep(interval)

# ---------------------------------------------------------
# Display Dashboard
# ---------------------------------------------------------
def rssi_bar(rssi: int) -> str:
    """Converts RSSI (dBm) to a visual bar indicator."""
    if rssi >= -50:  return "████  Excellent"
    if rssi >= -60:  return "███   Good"
    if rssi >= -70:  return "██    Fair"
    if rssi >= -80:  return "█     Weak"
    return                   "·     No Signal"


def display_dashboard(interface: str):
    """Periodically refreshes the terminal with live network/client data."""
    while True:
        os.system("clear")
        print("╔══════════════════════════════════════════════════════════════╗")
        print(f"║  Live Wi-Fi Analyzer  │  Interface: {interface:<24}║")
        print(f"║  {datetime.now().strftime('%Y-%m-%d %H:%M:%S'):<60}║")
        print("╚══════════════════════════════════════════════════════════════╝\n")

        with lock:
            net_snapshot    = dict(networks)
            client_snapshot = dict(clients)

        # --- Networks Table ---
        print("┌─ DETECTED NETWORKS ─────────────────────────────────────────")
        if net_snapshot:
            print(f"  {'BSSID':<19} {'SSID':<24} {'CH':>3}  {'Crypto':<10} {'Signal':<20} {'Vendor':<18} {'BW'}")
            print("  " + "─" * 110)
            for bssid, n in sorted(net_snapshot.items(), key=lambda x: x[1].rssi, reverse=True):
                print(
                    f"  {bssid:<19} {n.ssid[:22]:<24} {n.channel:>3}  "
                    f"{n.crypto:<10} {rssi_bar(n.rssi):<20} {n.vendor[:16]:<18} {n.bandwidth}"
                )
        else:
            print("  Listening for beacons...")

        # --- Clients Table ---
        print(f"\n┌─ CONNECTED CLIENTS ({len(client_snapshot)}) ────────────────────────────────")
        if client_snapshot:
            print(f"  {'Client MAC':<19} {'SSID':<24} {'AP BSSID':<19} {'Vendor':<18} {'Last Seen'}")
            print("  " + "─" * 95)
            for mac, c in client_snapshot.items():
                print(f"  {mac:<19} {c.ssid[:22]:<24} {c.connected_to_bssid:<19} {c.vendor[:16]:<18} {c.last_seen}")
        else:
            print("  Waiting for data frames...")

        print(f"\n  Networks: {len(net_snapshot)}  │  Clients: {len(client_snapshot)}  │  Press Ctrl+C to stop\n")
        time.sleep(2)

# ---------------------------------------------------------
# Export
# ---------------------------------------------------------
def export_results():
    """Exports captured data to JSON and CSV files."""
    ts = datetime.now().strftime("%Y%m%d_%H%M%S")

    with lock:
        net_snapshot    = dict(networks)
        client_snapshot = dict(clients)

    # JSON export
    json_path = f"wifi_scan_{ts}.json"
    with open(json_path, "w") as f:
        json.dump({
            "scan_time": ts,
            "networks": {k: asdict(v) for k, v in net_snapshot.items()},
            "clients":  {k: asdict(v) for k, v in client_snapshot.items()}
        }, f, indent=2)

    # CSV — Networks
    net_csv = f"networks_{ts}.csv"
    with open(net_csv, "w") as f:
        f.write("BSSID,SSID,Channel,Crypto,Bandwidth,RSSI,Vendor,First_Seen,Last_Seen\n")
        for bssid, n in net_snapshot.items():
            f.write(f"{bssid},{n.ssid},{n.channel},{n.crypto},{n.bandwidth},{n.rssi},{n.vendor},{n.first_seen},{n.last_seen}\n")

    # CSV — Clients
    cli_csv = f"clients_{ts}.csv"
    with open(cli_csv, "w") as f:
        f.write("Client_MAC,SSID,Connected_To_BSSID,Vendor,Last_Seen\n")
        for mac, c in client_snapshot.items():
            f.write(f"{mac},{c.ssid},{c.connected_to_bssid},{c.vendor},{c.last_seen}\n")

    log.info("Exported → %s, %s, %s", json_path, net_csv, cli_csv)

# ---------------------------------------------------------
# Main
# ---------------------------------------------------------
def main():
    args = parse_args()

    log.info("Starting Wi-Fi Analyzer on interface: %s", args.interface)

    # Channel hopper thread
    hopper = threading.Thread(
        target=channel_hopper,
        args=(args.interface, args.channels, args.hop_interval),
        daemon=True
    )
    hopper.start()

    # Display dashboard thread
    display = threading.Thread(
        target=display_dashboard,
        args=(args.interface,),
        daemon=True
    )
    display.start()

    try:
        sniff(iface=args.interface, prn=packet_handler, store=0)

    except KeyboardInterrupt:
        os.system("clear")
        log.info("Stopped by user.")

        with lock:
            print(f"\n[Final] {len(networks)} networks, {len(clients)} clients captured.\n")

        if args.export:
            export_results()

    except PermissionError:
        log.error("Permission denied. Run with sudo.")

    except OSError as e:
        log.error("Interface error: %s — Is '%s' in monitor mode?", e, args.interface)


if __name__ == "__main__":
    main()