wifi_scan 0.7.3

use std::{thread::sleep, time::Duration};

use crate::{misc::get_channel, Error, Result, Wifi, WifiSecurity, WlanScanner};

use neli_wifi::Socket as SocketN;
use netlink_rust::{generic, Protocol, Socket};
use nl80211_rs::{
    self as nl80211,
    information_element::{AuthenticationKeyManagement, InformationElement},
};

pub struct ScanLinux;

impl WlanScanner for ScanLinux {
    /// Returns a list of WiFi hotspots in your area.
    /// Uses `nl80211-rs` and `netlink-rust` crates on Linux.
    /// On Linux, very frequent scans may produce unexpected results on some machines,
    /// scanning requires root privileges and results can be up to 2500ms old.
    fn scan(&mut self) -> Result<Vec<Wifi>> {
        let mut socket_conn = SocketN::connect().map_err(|e| Error::SocketError(e.to_string()))?;

        let interfaces = socket_conn
            .get_interfaces_info()
            .map_err(|e| Error::InterfaceError(e.to_string()))?;

        if interfaces.is_empty() {
            return Err(Error::InterfaceError(
                "No WiFi adapters detected".to_string(),
            ));
        }

        let mut filtered_wifis: Vec<Wifi> = Vec::new();
        let mut all_wifis: Vec<Wifi> = Vec::new();

        let scan_result = std::panic::catch_unwind(trigger_scan);

        // sleep if at least one scan succeeded
        match scan_result {
            Ok(inner) => match inner {
                Ok(_) => sleep(Duration::from_millis(1500)),
                Err(e) => {
                    if e.to_string().contains("Operation not permitted") {
                        return Err(Error::ScanFailed(
                            "Operation not permitted. Try running as root.".to_string(),
                        ));
                    }
                }
            },
            Err(_) => eprintln!("WARNING: Code to trigger WiFi scan panicked"),
        }

        for interface in &interfaces {
            if let Some(index) = interface.index {
                let mut results: Vec<Wifi> = Vec::new();
                let bss_list = socket_conn.get_bss_info(index);
                if let Ok(bss_list) = bss_list {
                    for bss in bss_list {
                        if let Some(seen) = bss.seen_ms_ago {
                            if seen <= 2500 {
                                results.push(Wifi {
                                    mac: match bss.bssid {
                                        Some(bytes) => convert_mac(bytes),
                                        None => String::new(),
                                    },
                                    ssid: match bss.information_elements.clone() {
                                        Some(ie_data) => get_ssid(ie_data),
                                        None => String::new(),
                                    },
                                    channel: match bss.frequency {
                                        Some(frequency) => get_channel(frequency),
                                        None => 0,
                                    },
                                    signal_level: match bss.signal {
                                        Some(signal) => signal / 100,
                                        None => 0,
                                    },
                                    security: match bss.information_elements.clone() {
                                        Some(ie_data) => get_security(ie_data),
                                        None => vec![],
                                    },
                                });
                            }
                        }
                    }
                }

                all_wifis.extend(results);
            }
        }
        for wifi in all_wifis {
            let exists = filtered_wifis.iter().any(|x| x.mac == wifi.mac);
            if !exists {
                filtered_wifis.push(wifi);
            }
        }
        Ok(filtered_wifis)
    }
}

/// Returns Ok() if at least one scan succeeded
fn trigger_scan() -> Result<()> {
    let mut control_socket = match Socket::new(Protocol::Generic) {
        Ok(sock) => sock,
        Err(e) => return Err(Error::SocketError(e.to_string())),
    };

    let family = match generic::Family::from_name(&mut control_socket, "nl80211") {
        Ok(fam) => fam,
        Err(e) => return Err(Error::SocketError(e.to_string())),
    };

    let devices = match nl80211::get_wireless_interfaces(&mut control_socket, &family) {
        Ok(dev) => dev,
        Err(e) => return Err(Error::SocketError(e.to_string())),
    };

    let mut failed_count = 0;
    let mut one_succeeded = false;
    for dev in devices {
        match dev.trigger_scan(&mut control_socket) {
            Ok(_) => {
                one_succeeded = true;
                println!("Triggered scan on: {}", dev.interface_name)
            }
            Err(e) => {
                eprintln!(
                    "WARNING: Failed to trigger scan on {}: {}",
                    dev.interface_name, e
                );
                failed_count += 1;
                if e.to_string().contains("not permitted") {
                    return Err(Error::ScanFailed(
                        "Operation not permitted. Try running as root.".to_string(),
                    ));
                }
            }
        }
    }

    if one_succeeded {
        Ok(())
    } else {
        Err(Error::ScanFailed(format!(
            "Triggering a network scan failed on {} devices.",
            failed_count
        )))
    }
}

fn convert_mac(bytes: Vec<u8>) -> String {
    bytes
        .iter()
        .map(|b| format!("{:02x}", b))
        .collect::<Vec<_>>()
        .join(":")
}

fn get_ssid(ie_data: Vec<u8>) -> String {
    let ie_data: &[u8] = &ie_data;
    match InformationElement::parse_all(ie_data) {
        Ok(ies) => {
            for ie in ies {
                if let InformationElement::Ssid(ssid_ie) = ie {
                    return ssid_ie.ssid;
                }
            }
        }
        Err(_) => return String::new(),
    }

    String::new()
}

fn get_security(ie_data: Vec<u8>) -> Vec<WifiSecurity> {
    let ie_data: &[u8] = &ie_data;
    match InformationElement::parse_all(ie_data) {
        Ok(ies) => {
            for ie in ies {
                if let InformationElement::RobustSecurityNetwork(sec_ie) = ie {
                    let mut securities: Vec<WifiSecurity> = Vec::new();

                    for akm in sec_ie.akms {
                        let security = match akm {
                            AuthenticationKeyManagement::PreSharedKey => {
                                WifiSecurity::Wpa2PersonalPsk
                            }
                            AuthenticationKeyManagement::SimultaneousAuthenticationOfEquals => {
                                WifiSecurity::Wpa3PersonalSae
                            }
                            AuthenticationKeyManagement::PairwiseMasterKeySecurityAssociation => {
                                WifiSecurity::Wpa2EnterpriseEap
                            }
                            AuthenticationKeyManagement::PMKSASha256 => {
                                WifiSecurity::Wpa3EnterpriseEap256
                            }
                            AuthenticationKeyManagement::FastTransitionPMKSA => {
                                WifiSecurity::Wpa2EnterpriseEapFt
                            }
                            AuthenticationKeyManagement::PreSharedKeySha256 => {
                                WifiSecurity::Wpa2PersonalPsk256
                            }
                            AuthenticationKeyManagement::FastTransitionPreSharedKey => {
                                WifiSecurity::Wpa2PersonalPskFt
                            }
                            AuthenticationKeyManagement::FastTransitionSAE => {
                                WifiSecurity::Wpa3PersonalSaeFt
                            }
                            AuthenticationKeyManagement::TunneledDirectLinkSetup => {
                                WifiSecurity::TunneledDirectLinkSetup
                            }
                            _ => WifiSecurity::Unknown,
                        };

                        if security != WifiSecurity::Unknown {
                            securities.push(security);
                        }
                    }

                    if securities.is_empty() {
                        securities.push(WifiSecurity::Unknown);
                    }

                    return securities;
                }
            }
        }
        Err(_) => return vec![WifiSecurity::Unknown],
    }

    vec![WifiSecurity::Open]
}