wifiscan 0.4.0

use crate::types::EncryptionInfo;

// ── Frequency / Channel mapping ─────────────────────────────────────────────

pub fn frequency_to_channel(freq: u16) -> u8 {
    match freq {
        2412 => 1,   2417 => 2,   2422 => 3,   2427 => 4,   2432 => 5,
        2437 => 6,   2442 => 7,   2447 => 8,   2452 => 9,   2457 => 10,
        2462 => 11,  2467 => 12,  2472 => 13,  2484 => 14,
        5180 => 36,  5200 => 40,  5220 => 44,  5240 => 48,
        5260 => 52,  5280 => 56,  5300 => 60,  5320 => 64,
        5500 => 100, 5520 => 104, 5540 => 108, 5560 => 112,
        5580 => 116, 5600 => 120, 5620 => 124, 5640 => 128,
        5660 => 132, 5680 => 136, 5700 => 140, 5720 => 144,
        5745 => 149, 5765 => 153, 5785 => 157, 5805 => 161, 5825 => 165,
        f if (5955..=7115).contains(&f) => ((f - 5950) / 5) as u8,
        _ => 0,
    }
}

// ── Frame Control helpers ───────────────────────────────────────────────────

/// Check the HT Control field presence from FC byte 1 bit 7 (Order bit).
pub fn has_ht_control(fc1: u8) -> bool {
    (fc1 & 0x80) != 0
}

/// Compute offset to data body in an 802.11 data frame,
/// accounting for Address4, QoS, and HT Control fields.
pub fn data_body_offset(frame_bytes: &[u8]) -> usize {
    if frame_bytes.len() < 2 {
        return frame_bytes.len();
    }
    let fc0 = frame_bytes[0];
    let fc1 = frame_bytes[1];
    let subtype = (fc0 >> 4) & 0x0F;
    let to_ds = (fc1 & 0x01) != 0;
    let from_ds = (fc1 & 0x02) != 0;

    let mut offset: usize = if to_ds && from_ds { 30 } else { 24 };

    // QoS subtypes (8-15) add 2 bytes
    if subtype >= 8 {
        offset += 2;
    }

    // HT Control (Order bit)
    if has_ht_control(fc1) {
        offset += 4;
    }

    offset
}

/// Extract addresses from raw 802.11 header based on To DS / From DS bits.
/// Returns (bssid_offset, src_offset, dst_offset) into frame_bytes.
pub fn address_offsets(fc1: u8) -> Option<(usize, usize, usize)> {
    let to_ds = (fc1 & 0x01) != 0;
    let from_ds = (fc1 & 0x02) != 0;
    match (to_ds, from_ds) {
        (false, true)  => Some((10, 16, 4)),
        (true, false)  => Some((4, 10, 16)),
        (false, false) => Some((16, 10, 4)),
        (true, true)   => None, // WDS, skip
    }
}

// ── Information Element parsing ─────────────────────────────────────────────

pub fn extract_ssid_from_ies(frame_bytes: &[u8], ie_start: usize) -> Option<String> {
    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 == 0 {
            let ssid_bytes = &frame_bytes[pos + 2..pos + 2 + tag_len];
            if tag_len == 0 || ssid_bytes.iter().all(|&b| b == 0) {
                return Some("<hidden>".to_string());
            }
            return Some(String::from_utf8_lossy(ssid_bytes).to_string());
        }
        pos += 2 + tag_len;
    }
    None
}

pub fn extract_channel_from_ies(frame_bytes: &[u8], ie_start: usize) -> Option<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 == 3 && tag_len == 1 {
            return Some(frame_bytes[pos + 2]);
        }
        pos += 2 + tag_len;
    }
    None
}

/// Parse all IEs and return detailed encryption info, WiFi generation, channel width, BSS color.
pub struct IeInfo {
    pub encryption: EncryptionInfo,
    pub wifi_generation: Option<u8>,
    pub channel_width: u16,
    pub bss_color: Option<u8>,
}

pub fn parse_ies_full(frame_bytes: &[u8], cap_offset: usize) -> IeInfo {
    let mut enc = EncryptionInfo::default();
    let mut wifi_gen: Option<u8> = None;
    let mut chan_width: u16 = 20;
    let mut bss_color: Option<u8> = None;
    let mut has_ht = false;
    let mut has_vht = false;
    let mut has_he = false;

    // Check capability field for WEP
    let cap_off = if frame_bytes.len() >= 2 && has_ht_control(frame_bytes[1]) {
        cap_offset + 4
    } else {
        cap_offset
    };
    if frame_bytes.len() > cap_off + 1 {
        let cap = u16::from_le_bytes([frame_bytes[cap_off], frame_bytes[cap_off + 1]]);
        if cap & 0x0010 != 0 {
            enc.has_wep = true;
        }
    }

    let ie_start = cap_off + 2;
    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; }
        let tag_data = &frame_bytes[pos + 2..pos + 2 + tag_len];

        match tag_id {
            // RSN (WPA2/WPA3)
            48 if tag_len >= 2 => {
                enc.has_rsn = true;
                parse_rsn_ie(tag_data, &mut enc);
            }
            // Vendor specific: WPA IE
            221 if tag_len >= 4 => {
                if tag_data[0..4] == [0x00, 0x50, 0xf2, 0x01] {
                    enc.has_wpa = true;
                }
            }
            // HT Capabilities (802.11n / WiFi 4)
            45 => {
                has_ht = true;
                if tag_len >= 2 {
                    let ht_cap = u16::from_le_bytes([tag_data[0], tag_data[1]]);
                    // Bit 1: 40 MHz supported
                    if ht_cap & 0x0002 != 0 {
                        chan_width = chan_width.max(40);
                    }
                }
            }
            // HT Operation
            61 if tag_len >= 2 => {
                let sta_chan_width = (tag_data[1] & 0x04) != 0;
                if sta_chan_width {
                    chan_width = chan_width.max(40);
                }
            }
            // VHT Capabilities (802.11ac / WiFi 5)
            191 => {
                has_vht = true;
                if tag_len >= 4 {
                    let vht_cap = u32::from_le_bytes([
                        tag_data[0], tag_data[1], tag_data[2], tag_data[3],
                    ]);
                    // Bits 2-3: max channel width
                    match (vht_cap >> 2) & 0x03 {
                        0 => chan_width = chan_width.max(80),
                        1 => chan_width = chan_width.max(160),
                        2 => chan_width = chan_width.max(160), // 80+80
                        _ => {}
                    }
                }
            }
            // VHT Operation
            192 if tag_len >= 1 => {
                match tag_data[0] {
                    1 => chan_width = chan_width.max(80),
                    2 => chan_width = chan_width.max(160),
                    3 => chan_width = chan_width.max(160), // 80+80
                    _ => {}
                }
            }
            // Element ID Extension (255): HE Capabilities, HE Operation (WiFi 6)
            255 if tag_len >= 1 => {
                let ext_id = tag_data[0];
                match ext_id {
                    // HE Capabilities (ext 35)
                    35 => {
                        has_he = true;
                        // HE implies at least 80 MHz on 5 GHz
                        chan_width = chan_width.max(80);
                        // Check for 160 MHz support in HE PHY cap (byte offset 7, bit 3)
                        if tag_len >= 8 && (tag_data[7] & 0x08) != 0 {
                            chan_width = chan_width.max(160);
                        }
                    }
                    // HE Operation (ext 36)
                    36 if tag_len >= 7 => {
                        has_he = true;
                        // BSS Color is in HE Operation Parameters, byte 4, bits 0-5
                        if tag_len >= 5 {
                            let color = tag_data[4] & 0x3F;
                            if color > 0 {
                                bss_color = Some(color);
                            }
                        }
                    }
                    // EHT Capabilities (ext 108) - WiFi 7
                    108 => {
                        wifi_gen = Some(7);
                        chan_width = chan_width.max(160); // EHT supports 320 MHz too
                        if tag_len >= 4 && (tag_data[1] & 0x02) != 0 {
                            chan_width = 320;
                        }
                    }
                    _ => {}
                }
            }
            _ => {}
        }
        pos += 2 + tag_len;
    }

    // Determine WiFi generation if not already set to 7
    if wifi_gen.is_none() {
        wifi_gen = if has_he {
            Some(6)
        } else if has_vht {
            Some(5)
        } else if has_ht {
            Some(4)
        } else {
            None
        };
    }

    enc.update_display();

    IeInfo {
        encryption: enc,
        wifi_generation: wifi_gen,
        channel_width: chan_width,
        bss_color,
    }
}

/// Parse RSN IE contents into EncryptionInfo.
fn parse_rsn_ie(data: &[u8], enc: &mut EncryptionInfo) {
    if data.len() < 2 { return; }
    // Version
    let mut pos: usize = 2; // skip version

    // Group cipher suite
    if pos + 4 > data.len() { return; }
    pos += 4;

    // Pairwise cipher suite count
    if pos + 2 > data.len() { return; }
    let pw_count = u16::from_le_bytes([data[pos], data[pos + 1]]) as usize;
    pos += 2;

    for _ in 0..pw_count {
        if pos + 4 > data.len() { return; }
        let cipher_type = data[pos + 3] as u32;
        enc.pairwise_ciphers.push(cipher_type);
        pos += 4;
    }

    // AKM suite count
    if pos + 2 > data.len() { return; }
    let akm_count = u16::from_le_bytes([data[pos], data[pos + 1]]) as usize;
    pos += 2;

    for _ in 0..akm_count {
        if pos + 4 > data.len() { return; }
        let akm_type = data[pos + 3] as u32;
        enc.akm_suites.push(akm_type);
        // AKM 8 = SAE (WPA3-Personal), AKM 18 = OWE
        if akm_type == 8 || akm_type == 18 {
            enc.wpa3_sae = true;
        }
        pos += 4;
    }

    // RSN capabilities (PMF bits)
    if pos + 2 > data.len() { return; }
    let rsn_cap = u16::from_le_bytes([data[pos], data[pos + 1]]);
    enc.pmf_capable = (rsn_cap & 0x0080) != 0;
    enc.pmf_required = (rsn_cap & 0x0040) != 0;
}

/// Legacy encryption parser for backward compatibility (returns just display string).
pub fn parse_encryption_from_frame(frame_bytes: &[u8], base_cap_offset: usize) -> String {
    let info = parse_ies_full(frame_bytes, base_cap_offset);
    info.encryption.display
}