wifiscan 0.4.0

use std::collections::{HashMap, HashSet, VecDeque};
use std::path::PathBuf;
use std::sync::atomic::AtomicU64;
use std::sync::{Arc, Mutex};
use std::time::Instant;

use clap::Parser;
use serde::Serialize;

// ── CLI ─────────────────────────────────────────────────────────────────────

#[derive(Parser, Debug)]
#[command(name = "wifiscan", version, about = "Wireless network scanner TUI")]
pub struct Args {
    /// Wireless interface (e.g. wlp9s0). Will be put into monitor mode automatically.
    #[arg(short, long)]
    pub interface: String,

    /// Refresh interval in milliseconds
    #[arg(short, long, default_value_t = 1000)]
    pub refresh: u64,

    /// Sort column: ssid, bssid, signal, channel, seen, encryption, clients
    #[arg(short, long, default_value = "signal")]
    pub sort: String,

    /// Channel to lock on (0 = hop all channels)
    #[arg(short, long, default_value_t = 0)]
    pub channel: u8,

    /// Directory to save handshake pcaps (default: ./handshakes)
    #[arg(long, default_value = "./handshakes")]
    pub handshake_dir: PathBuf,

    /// Disable automatic handshake capture
    #[arg(long, default_value_t = false)]
    pub no_handshakes: bool,

    /// Skip automatic monitor mode setup (use if interface is already in monitor mode)
    #[arg(long, default_value_t = false)]
    pub skip_monitor: bool,

    /// 2.4 GHz channel dwell time in ms (default: 200)
    #[arg(long, default_value_t = 200)]
    pub dwell_2g: u64,

    /// 5 GHz channel dwell time in ms (default: 350)
    #[arg(long, default_value_t = 350)]
    pub dwell_5g: u64,

    /// 6 GHz channel dwell time in ms (default: 400)
    #[arg(long, default_value_t = 400)]
    pub dwell_6g: u64,

    /// Deauth burst count per target (default: 64)
    #[arg(long, default_value_t = 64)]
    pub deauth_burst: u32,

    /// Enable debug logging to ./wifiscan_debug.log
    #[arg(long, default_value_t = false)]
    pub debug: bool,

    /// Dump all captured packets to a pcap file for Wireshark analysis
    #[arg(long)]
    pub dump_pcap: Option<PathBuf>,

    /// Handshake state timeout in seconds (default: 120)
    #[arg(long, default_value_t = 120)]
    pub hs_timeout: u64,

    /// Seconds to dwell on AP channel after deauth waiting for handshake (default: 15)
    #[arg(long, default_value_t = 15)]
    pub deauth_dwell: u64,

    /// Band filter: all, 2g, 5g, 6g (default: all)
    #[arg(long, default_value = "all")]
    pub band: String,

    /// Auto-remove APs not seen for this many seconds (0 = never, default: 0)
    #[arg(long, default_value_t = 0)]
    pub auto_expire: u64,

    /// Minimum signal strength to display (default: -100)
    #[arg(long, default_value_t = -100)]
    pub min_signal: i8,
}

// ── Band Filter ─────────────────────────────────────────────────────────────

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum BandFilter {
    All,
    Band2G,
    Band5G,
    Band6G,
}

impl BandFilter {
    pub fn parse_band(s: &str) -> Self {
        match s.to_lowercase().as_str() {
            "2g" | "2.4g" | "2.4ghz" => Self::Band2G,
            "5g" | "5ghz" => Self::Band5G,
            "6g" | "6ghz" => Self::Band6G,
            _ => Self::All,
        }
    }

    pub fn matches(&self, freq: u16) -> bool {
        match self {
            Self::All => true,
            Self::Band2G => (1..3000).contains(&freq),
            Self::Band5G => (5000..5900).contains(&freq),
            Self::Band6G => freq >= 5900,
        }
    }

    pub fn label(&self) -> &str {
        match self {
            Self::All => "All",
            Self::Band2G => "2.4G",
            Self::Band5G => "5G",
            Self::Band6G => "6G",
        }
    }

    pub fn next(&self) -> Self {
        match self {
            Self::All => Self::Band2G,
            Self::Band2G => Self::Band5G,
            Self::Band5G => Self::Band6G,
            Self::Band6G => Self::All,
        }
    }
}

// ── MAC address as value type ───────────────────────────────────────────────

#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub struct MacAddr(pub [u8; 6]);

impl MacAddr {
    pub const BROADCAST: Self = Self([0xff; 6]);
    pub const ZERO: Self = Self([0; 6]);

    pub fn from_bytes(b: &[u8]) -> Option<Self> {
        if b.len() < 6 { return None; }
        let mut a = [0u8; 6];
        a.copy_from_slice(&b[..6]);
        Some(Self(a))
    }

    pub fn from_str_hex(s: &str) -> Option<Self> {
        let parts: Vec<&str> = s.split(':').collect();
        if parts.len() != 6 { return None; }
        let mut a = [0u8; 6];
        for (i, p) in parts.iter().enumerate() {
            a[i] = u8::from_str_radix(p, 16).ok()?;
        }
        Some(Self(a))
    }

    pub fn is_broadcast(&self) -> bool { *self == Self::BROADCAST }
    pub fn is_zero(&self) -> bool { *self == Self::ZERO }

    pub fn oui_prefix(&self) -> String {
        format!("{:02X}{:02X}{:02X}", self.0[0], self.0[1], self.0[2])
    }

    /// Check if MAC is locally administered (randomized)
    pub fn is_randomized(&self) -> bool {
        (self.0[0] & 0x02) != 0
    }
}

impl std::fmt::Display for MacAddr {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(f, "{:02x}:{:02x}:{:02x}:{:02x}:{:02x}:{:02x}",
            self.0[0], self.0[1], self.0[2], self.0[3], self.0[4], self.0[5])
    }
}

// ── AP / Client Data ────────────────────────────────────────────────────────

#[derive(Debug, Clone)]
pub struct ClientInfo {
    pub mac: String,
    pub signal_dbm: i8,
    pub last_seen: Instant,
    pub first_seen: Instant,
    pub probed_ssids: Vec<String>,
    pub data_count: u64,
    pub vendor: String,
    pub is_randomized: bool,
}

impl ClientInfo {
    pub fn new(mac: String) -> Self {
        let is_randomized = MacAddr::from_str_hex(&mac)
            .map(|m| m.is_randomized())
            .unwrap_or(false);
        let now = Instant::now();
        Self {
            mac,
            signal_dbm: -100,
            last_seen: now,
            first_seen: now,
            probed_ssids: Vec::new(),
            data_count: 0,
            vendor: String::new(),
            is_randomized,
        }
    }

    pub fn age_secs(&self) -> u64 {
        self.last_seen.elapsed().as_secs()
    }

    pub fn uptime_secs(&self) -> u64 {
        self.first_seen.elapsed().as_secs()
    }
}

/// Encryption details parsed from RSN/WPA IEs.
#[derive(Debug, Clone, Serialize)]
pub struct EncryptionInfo {
    pub display: String,
    pub has_rsn: bool,
    pub has_wpa: bool,
    pub has_wep: bool,
    pub wpa3_sae: bool,
    pub pmf_required: bool,
    pub pmf_capable: bool,
    pub akm_suites: Vec<u32>,
    pub pairwise_ciphers: Vec<u32>,
}

impl Default for EncryptionInfo {
    fn default() -> Self {
        Self {
            display: "Open".to_string(),
            has_rsn: false,
            has_wpa: false,
            has_wep: false,
            wpa3_sae: false,
            pmf_required: false,
            pmf_capable: false,
            akm_suites: Vec::new(),
            pairwise_ciphers: Vec::new(),
        }
    }
}

impl EncryptionInfo {
    pub fn update_display(&mut self) {
        if self.wpa3_sae && self.has_rsn {
            if self.akm_suites.iter().any(|&a| a == 2 || a == 1) {
                self.display = "WPA3/WPA2".to_string();
            } else {
                self.display = "WPA3".to_string();
            }
        } else if self.has_rsn && self.has_wpa {
            self.display = "WPA2/WPA".to_string();
        } else if self.has_rsn {
            self.display = "WPA2".to_string();
        } else if self.has_wpa {
            self.display = "WPA".to_string();
        } else if self.has_wep {
            self.display = "WEP".to_string();
        } else {
            self.display = "Open".to_string();
        }
        if self.pmf_required {
            self.display.push_str(" [PMF]");
        }
    }

    /// Security score 0-100 for sorting/display
    pub fn security_score(&self) -> u8 {
        let mut score: u8 = 0;
        if self.has_wep { score = 10; }
        if self.has_wpa && !self.has_rsn { score = 30; }
        if self.has_rsn { score = 60; }
        if self.wpa3_sae { score = 80; }
        if self.pmf_required { score += 15; }
        if self.pairwise_ciphers.contains(&4) { score += 5; } // CCMP
        score.min(100)
    }
}

#[derive(Debug, Clone)]
pub struct AccessPoint {
    pub bssid: String,
    pub bssid_mac: MacAddr,
    pub essid: String,
    pub signal_dbm: i8,
    pub noise_dbm: Option<i8>,
    pub channel: u8,
    pub channel_width: u16,
    pub frequency_mhz: u16,
    pub encryption: EncryptionInfo,
    pub beacon_count: u64,
    pub data_count: u64,
    pub first_seen: Instant,
    pub last_seen: Instant,
    pub clients: HashMap<String, ClientInfo>,
    pub handshakes: u32,
    pub pmkid_captured: bool,
    pub vendor: String,
    pub signal_history: VecDeque<i8>,
    pub wifi_generation: Option<u8>,
    pub bss_color: Option<u8>,
    /// Track deauth attempts on this AP
    pub deauth_sent: u32,
    pub last_deauth: Option<Instant>,
}

impl AccessPoint {
    pub fn new(bssid: MacAddr) -> Self {
        let now = Instant::now();
        Self {
            bssid: bssid.to_string(),
            bssid_mac: bssid,
            essid: String::new(),
            signal_dbm: -100,
            noise_dbm: None,
            channel: 0,
            channel_width: 20,
            frequency_mhz: 0,
            encryption: EncryptionInfo::default(),
            beacon_count: 0,
            data_count: 0,
            first_seen: now,
            last_seen: now,
            clients: HashMap::new(),
            handshakes: 0,
            pmkid_captured: false,
            vendor: String::new(),
            signal_history: VecDeque::new(),
            wifi_generation: None,
            bss_color: None,
            deauth_sent: 0,
            last_deauth: None,
        }
    }

    pub fn age_secs(&self) -> u64 {
        self.last_seen.elapsed().as_secs()
    }

    pub fn uptime_secs(&self) -> u64 {
        self.first_seen.elapsed().as_secs()
    }

    pub fn signal_bar(&self) -> &str {
        match self.signal_dbm {
            -30..=0 => "█████",
            -50..=-31 => "████░",
            -60..=-51 => "███░░",
            -70..=-61 => "██░░░",
            -80..=-71 => "█░░░░",
            _ => "░░░░░",
        }
    }

    pub fn signal_color(&self) -> ratatui::style::Color {
        use ratatui::style::Color;
        match self.signal_dbm {
            -30..=0 => Color::Green,
            -50..=-31 => Color::LightGreen,
            -60..=-51 => Color::Yellow,
            -70..=-61 => Color::Rgb(255, 165, 0),
            -80..=-71 => Color::LightRed,
            _ => Color::Red,
        }
    }

    pub fn push_signal(&mut self, sig: i8) {
        self.signal_history.push_back(sig);
        if self.signal_history.len() > 120 {
            self.signal_history.pop_front();
        }
    }

    pub fn wifi_gen_str(&self) -> &str {
        match self.wifi_generation {
            Some(4) => "4",
            Some(5) => "5",
            Some(6) => "6",
            Some(7) => "7",
            _ => "-",
        }
    }

    pub fn band_str(&self) -> &str {
        if self.frequency_mhz >= 5900 { "6G" }
        else if self.frequency_mhz >= 5000 { "5G" }
        else if self.frequency_mhz > 0 { "2G" }
        else { "--" }
    }

    /// SNR in dB if noise floor is available
    pub fn snr(&self) -> Option<i16> {
        self.noise_dbm.map(|n| self.signal_dbm as i16 - n as i16)
    }

    /// Active client count (seen in last 60s)
    pub fn active_clients(&self) -> usize {
        self.clients.values().filter(|c| c.age_secs() < 60).count()
    }

    /// Remove stale clients not seen for `secs`
    pub fn prune_clients(&mut self, secs: u64) {
        self.clients.retain(|_, c| c.age_secs() < secs);
    }
}

// ── Shared State Type Aliases ───────────────────────────────────────────────

pub type ApMap = Arc<Mutex<HashMap<String, AccessPoint>>>;
pub type CurrentChannel = Arc<Mutex<u8>>;
pub type HandshakeCount = Arc<Mutex<u32>>;
pub type ChannelStats = Arc<Mutex<HashMap<u8, u64>>>;

// ── Handshake Tracking ──────────────────────────────────────────────────────

#[derive(Debug, Clone, Hash, PartialEq, Eq)]
pub struct HandshakeKey {
    pub bssid: String,
    pub client: String,
}

#[derive(Debug, Clone)]
pub struct HandshakeState {
    pub key: HandshakeKey,
    pub essid: String,
    pub messages: [Option<Vec<u8>>; 4],
    pub replay_counters: [Option<u64>; 4],
    pub beacon_raw: Option<Vec<u8>>,
    pub started: Instant,
    pub saved: bool,
    pub pmkid: Option<Vec<u8>>,
}

impl HandshakeState {
    pub fn new(key: HandshakeKey, essid: String) -> Self {
        Self {
            key,
            essid,
            messages: [None, None, None, None],
            replay_counters: [None; 4],
            beacon_raw: None,
            started: Instant::now(),
            saved: false,
            pmkid: None,
        }
    }

    /// Valid handshake: M1+M2 with matching replay, or M2+M3 with sequential replay.
    pub fn is_complete(&self) -> bool {
        if self.messages[0].is_some() && self.messages[1].is_some() {
            if let (Some(rc1), Some(rc2)) = (self.replay_counters[0], self.replay_counters[1]) {
                if rc1 == rc2 { return true; }
            }
        }
        if self.messages[1].is_some() && self.messages[2].is_some() {
            if let (Some(rc2), Some(rc3)) = (self.replay_counters[1], self.replay_counters[2]) {
                if rc3 == rc2 || rc3 == rc2 + 1 { return true; }
            }
        }
        false
    }

    /// Return which messages we have as a compact string like "M1,M2"
    pub fn progress_str(&self) -> String {
        let mut parts = Vec::new();
        for (i, msg) in self.messages.iter().enumerate() {
            if msg.is_some() {
                parts.push(format!("M{}", i + 1));
            }
        }
        if parts.is_empty() { "none".to_string() } else { parts.join(",") }
    }
}

pub type HandshakeMap = Arc<Mutex<HashMap<HandshakeKey, HandshakeState>>>;

// ── Probe request tracking ──────────────────────────────────────────────────

pub type ProbeMap = Arc<Mutex<HashMap<String, HashSet<String>>>>;

// ── Deauth State ────────────────────────────────────────────────────────────

#[derive(Debug, Clone)]
pub enum DeauthTarget {
    All,
    SingleClient(String),
}

#[derive(Debug, Clone)]
pub enum DeauthState {
    Idle,
    Confirm { ap: AccessPoint, target: DeauthTarget },
    /// Deauth running on background thread. UI polls progress atomics.
    Running {
        ap: AccessPoint,
        progress: std::sync::Arc<crate::deauth::DeauthProgress>,
    },
}

// ── Beacon Flood (Evil Twin) State ───────────────────────────────────────────

#[derive(Debug, Clone)]
pub enum BeaconFloodState {
    Idle,
    Confirm { ap: Box<AccessPoint> },
    Running {
        ssid: String,
        bssid: String,
        channel: u8,
        progress: std::sync::Arc<crate::beacon::BeaconFloodProgress>,
    },
}

// ── Diagnostic Counters ─────────────────────────────────────────────────────

pub static FRAMES_TOTAL: AtomicU64 = AtomicU64::new(0);
pub static FRAMES_DATA: AtomicU64 = AtomicU64::new(0);
pub static FRAMES_EAPOL_CHECKED: AtomicU64 = AtomicU64::new(0);
pub static FRAMES_EAPOL_FOUND: AtomicU64 = AtomicU64::new(0);
pub static FRAMES_EAPOL_MSG: [AtomicU64; 4] = [
    AtomicU64::new(0), AtomicU64::new(0),
    AtomicU64::new(0), AtomicU64::new(0),
];
pub static FRAMES_PROTECTED_SKIP: AtomicU64 = AtomicU64::new(0);
pub static FRAMES_NULL_SKIP: AtomicU64 = AtomicU64::new(0);
pub static FRAMES_PROBE_REQ: AtomicU64 = AtomicU64::new(0);