crdhcpc 0.1.1

//! WiFi connection management for wireless interfaces
//!
//! This module provides WiFi connection management using wpa_supplicant
//! to establish wireless connections before DHCP operations.

use std::path::Path;
use std::time::Duration;
use tokio::process::Command;
use tokio::time::sleep;
use tracing::{debug, info};
use anyhow::{Context, Result, anyhow};

/// WiFi security/authentication type
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum WifiSecurity {
    /// Open network (no authentication)
    Open,
    /// WEP (deprecated, insecure)
    Wep { key: String },
    /// WPA/WPA2 Personal (PSK)
    WpaPsk { passphrase: String },
    /// WPA/WPA2 Enterprise (802.1X)
    WpaEnterprise {
        identity: String,
        password: String,
        ca_cert: Option<String>,
    },
    /// WPA3 Personal (SAE)
    Wpa3Sae { passphrase: String },
}

/// WiFi network configuration
#[derive(Debug, Clone)]
pub struct WifiConfig {
    /// Network SSID
    pub ssid: String,
    /// Security configuration
    pub security: WifiSecurity,
    /// Hidden network (doesn't broadcast SSID)
    pub hidden: bool,
    /// Connection timeout in seconds
    pub timeout: u64,
}

impl WifiConfig {
    /// Create a new WiFi configuration for WPA/WPA2
    pub fn new_wpa(ssid: impl Into<String>, passphrase: impl Into<String>) -> Self {
        Self {
            ssid: ssid.into(),
            security: WifiSecurity::WpaPsk {
                passphrase: passphrase.into(),
            },
            hidden: false,
            timeout: 30,
        }
    }

    /// Create a new WiFi configuration for open network
    pub fn new_open(ssid: impl Into<String>) -> Self {
        Self {
            ssid: ssid.into(),
            security: WifiSecurity::Open,
            hidden: false,
            timeout: 30,
        }
    }

    /// Generate wpa_supplicant configuration block
    /// Security: All user-supplied values are escaped to prevent config injection
    fn to_wpa_supplicant_config(&self) -> String {
        let mut config = format!("network={{\n");
        config.push_str(&format!("    ssid=\"{}\"\n", escape_wpa_string(&self.ssid)));

        if self.hidden {
            config.push_str("    scan_ssid=1\n");
        }

        match &self.security {
            WifiSecurity::Open => {
                config.push_str("    key_mgmt=NONE\n");
            }
            WifiSecurity::Wep { key } => {
                // WEP keys are typically hex strings, sanitize to hex only
                let safe_key: String = key.chars()
                    .filter(|c| c.is_ascii_hexdigit())
                    .collect();
                config.push_str(&format!("    wep_key0={}\n", safe_key));
                config.push_str("    key_mgmt=NONE\n");
            }
            WifiSecurity::WpaPsk { passphrase } => {
                config.push_str(&format!("    psk=\"{}\"\n", escape_wpa_string(passphrase)));
                config.push_str("    key_mgmt=WPA-PSK\n");
            }
            WifiSecurity::WpaEnterprise { identity, password, ca_cert } => {
                config.push_str(&format!("    identity=\"{}\"\n", escape_wpa_string(identity)));
                config.push_str(&format!("    password=\"{}\"\n", escape_wpa_string(password)));
                if let Some(cert) = ca_cert {
                    // Certificate paths should only contain safe path characters
                    let safe_cert: String = cert.chars()
                        .filter(|c| c.is_ascii_alphanumeric() || *c == '/' || *c == '.' || *c == '-' || *c == '_')
                        .collect();
                    config.push_str(&format!("    ca_cert=\"{}\"\n", safe_cert));
                }
                config.push_str("    key_mgmt=WPA-EAP\n");
                config.push_str("    eap=PEAP\n");
                config.push_str("    phase2=\"auth=MSCHAPV2\"\n");
            }
            WifiSecurity::Wpa3Sae { passphrase } => {
                config.push_str(&format!("    sae_password=\"{}\"\n", escape_wpa_string(passphrase)));
                config.push_str("    key_mgmt=SAE\n");
                config.push_str("    ieee80211w=2\n");
            }
        }

        config.push_str("}\n");
        config
    }
}

/// Escape a string for safe inclusion in wpa_supplicant config
/// Escapes backslashes, double quotes, and removes control characters
fn escape_wpa_string(s: &str) -> String {
    let mut escaped = String::with_capacity(s.len() * 2);
    for c in s.chars() {
        match c {
            '\\' => escaped.push_str("\\\\"),
            '"' => escaped.push_str("\\\""),
            // Remove control characters (potential config injection)
            c if c.is_control() => {}
            c => escaped.push(c),
        }
    }
    escaped
}

/// WiFi connection manager
pub struct WifiManager {
    interface: String,
}

/// Validate interface name to prevent command injection
/// Interface names should only contain alphanumeric chars, hyphens, and underscores
fn is_valid_interface_name(name: &str) -> bool {
    !name.is_empty()
        && name.len() <= 15  // Linux interface name limit
        && name.chars().all(|c| c.is_ascii_alphanumeric() || c == '-' || c == '_')
}

impl WifiManager {
    /// Create a new WiFi manager for the specified interface
    /// Returns None if the interface name is invalid
    pub fn new(interface: impl Into<String>) -> Result<Self> {
        let interface = interface.into();
        if !is_valid_interface_name(&interface) {
            return Err(anyhow!("Invalid interface name: {}", interface));
        }
        Ok(Self { interface })
    }

    /// Check if the interface is a wireless interface
    pub async fn is_wireless(&self) -> bool {
        let wireless_path = format!("/sys/class/net/{}/wireless", self.interface);
        Path::new(&wireless_path).exists()
    }

    /// Check if wpa_supplicant is installed
    pub async fn check_wpa_supplicant() -> Result<()> {
        let output = Command::new("which")
            .arg("wpa_supplicant")
            .output()
            .await
            .context("Failed to check for wpa_supplicant")?;

        if output.status.success() {
            Ok(())
        } else {
            Err(anyhow!("wpa_supplicant not found. Install with: sudo apt install wpasupplicant"))
        }
    }

    /// Check if already connected to a WiFi network
    pub async fn is_connected(&self) -> Result<bool> {
        debug!("Checking WiFi connection status for {}", self.interface);

        let output = Command::new("iw")
            .args(&[&self.interface, "link"])
            .output()
            .await
            .context("Failed to check WiFi status (is 'iw' installed?)")?;

        if !output.status.success() {
            return Ok(false);
        }

        let output_str = String::from_utf8_lossy(&output.stdout);

        // Check if connected (output contains "Connected to")
        let connected = output_str.contains("Connected to") ||
                       output_str.contains("SSID:");

        if connected {
            debug!("WiFi interface {} is already connected", self.interface);
        } else {
            debug!("WiFi interface {} is not connected", self.interface);
        }

        Ok(connected)
    }

    /// Get current SSID if connected
    pub async fn get_current_ssid(&self) -> Result<Option<String>> {
        let output = Command::new("iw")
            .args(&[&self.interface, "info"])
            .output()
            .await
            .context("Failed to get WiFi info")?;

        if !output.status.success() {
            return Ok(None);
        }

        let output_str = String::from_utf8_lossy(&output.stdout);

        for line in output_str.lines() {
            if let Some(ssid) = line.trim().strip_prefix("ssid ") {
                return Ok(Some(ssid.trim().to_string()));
            }
        }

        Ok(None)
    }

    /// Connect to a WiFi network using wpa_supplicant
    pub async fn connect(&self, config: &WifiConfig) -> Result<()> {
        info!("Connecting to WiFi network \"{}\" on {}", config.ssid, self.interface);

        // Check if wpa_supplicant is available
        Self::check_wpa_supplicant().await?;

        // Stop any existing wpa_supplicant for this interface
        self.stop_wpa_supplicant().await?;

        // Generate temporary config file
        let config_content = config.to_wpa_supplicant_config();
        let config_path = format!("/tmp/wpa_supplicant_{}.conf", self.interface);

        tokio::fs::write(&config_path, config_content)
            .await
            .context("Failed to write wpa_supplicant config")?;

        debug!("Created wpa_supplicant config at {}", config_path);

        // Start wpa_supplicant in background
        info!("Starting wpa_supplicant for {}", self.interface);
        let output = Command::new("wpa_supplicant")
            .args(&[
                "-B",                          // Background mode
                "-i", &self.interface,         // Interface
                "-c", &config_path,            // Config file
                "-D", "nl80211,wext",          // Driver (try nl80211 first, fallback to wext)
            ])
            .output()
            .await
            .context("Failed to start wpa_supplicant")?;

        if !output.status.success() {
            let stderr = String::from_utf8_lossy(&output.stderr);
            return Err(anyhow!("Failed to start wpa_supplicant: {}", stderr));
        }

        info!("✓ wpa_supplicant started for {}", self.interface);

        // Wait for connection
        info!("Waiting for WiFi connection (timeout: {} seconds)...", config.timeout);
        let start = std::time::Instant::now();
        let timeout = Duration::from_secs(config.timeout);

        while start.elapsed() < timeout {
            if self.is_connected().await? {
                if let Some(ssid) = self.get_current_ssid().await? {
                    info!("✓ Successfully connected to WiFi network \"{}\"", ssid);
                    return Ok(());
                }
            }
            sleep(Duration::from_secs(1)).await;
        }

        // Timeout - cleanup and return error
        self.stop_wpa_supplicant().await?;
        Err(anyhow!("WiFi connection timeout after {} seconds", config.timeout))
    }

    /// Stop wpa_supplicant for this interface
    pub async fn stop_wpa_supplicant(&self) -> Result<()> {
        debug!("Stopping wpa_supplicant for {}", self.interface);

        // Kill wpa_supplicant process for this interface
        let _ = Command::new("pkill")
            .args(&["-f", &format!("wpa_supplicant.*{}", self.interface)])
            .output()
            .await;

        // Small delay to ensure process cleanup
        sleep(Duration::from_millis(100)).await;

        Ok(())
    }

    /// Disconnect from WiFi
    pub async fn disconnect(&self) -> Result<()> {
        info!("Disconnecting WiFi on {}", self.interface);
        self.stop_wpa_supplicant().await
    }

    /// Scan for available WiFi networks
    pub async fn scan(&self) -> Result<Vec<WifiNetwork>> {
        info!("Scanning for WiFi networks on {}", self.interface);

        // Trigger scan
        let _ = Command::new("iw")
            .args(&[&self.interface, "scan"])
            .output()
            .await;

        // Small delay for scan to complete
        sleep(Duration::from_millis(500)).await;

        // Get scan results
        let output = Command::new("iw")
            .args(&[&self.interface, "scan"])
            .output()
            .await
            .context("Failed to scan WiFi networks")?;

        if !output.status.success() {
            return Err(anyhow!("WiFi scan failed"));
        }

        let output_str = String::from_utf8_lossy(&output.stdout);
        let networks = parse_scan_results(&output_str);

        info!("Found {} WiFi networks", networks.len());
        Ok(networks)
    }
}

/// WiFi network information from scan
#[derive(Debug, Clone)]
pub struct WifiNetwork {
    pub ssid: String,
    pub signal_strength: i32,  // dBm
    pub frequency: u32,         // MHz
    pub security: Vec<String>,  // Security types (WPA2, WPA3, etc.)
}

/// Parse iw scan results
fn parse_scan_results(output: &str) -> Vec<WifiNetwork> {
    let mut networks = Vec::new();
    let mut current_ssid: Option<String> = None;
    let mut current_signal: i32 = -100;
    let mut current_freq: u32 = 0;
    let mut current_security: Vec<String> = Vec::new();

    for line in output.lines() {
        let line = line.trim();

        if line.starts_with("BSS ") {
            // Save previous network
            if let Some(ssid) = current_ssid.take() {
                networks.push(WifiNetwork {
                    ssid,
                    signal_strength: current_signal,
                    frequency: current_freq,
                    security: current_security.clone(),
                });
            }
            // Reset for new network
            current_signal = -100;
            current_freq = 0;
            current_security.clear();
        } else if let Some(ssid) = line.strip_prefix("SSID: ") {
            current_ssid = Some(ssid.to_string());
        } else if line.starts_with("signal: ") {
            if let Some(signal_str) = line.strip_prefix("signal: ") {
                if let Some(db_pos) = signal_str.find(" dBm") {
                    if let Ok(signal) = signal_str[..db_pos].parse::<i32>() {
                        current_signal = signal;
                    }
                }
            }
        } else if let Some(freq_str) = line.strip_prefix("freq: ") {
            if let Ok(freq) = freq_str.parse::<u32>() {
                current_freq = freq;
            }
        } else if line.contains("WPA") || line.contains("RSN") {
            if line.contains("WPA3") {
                current_security.push("WPA3".to_string());
            } else if line.contains("WPA2") {
                current_security.push("WPA2".to_string());
            } else if line.contains("WPA") {
                current_security.push("WPA".to_string());
            }
        }
    }

    // Save last network
    if let Some(ssid) = current_ssid {
        networks.push(WifiNetwork {
            ssid,
            signal_strength: current_signal,
            frequency: current_freq,
            security: current_security,
        });
    }

    networks
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_wpa_config_generation() {
        let config = WifiConfig::new_wpa("TestNetwork", "password123");
        let wpa_config = config.to_wpa_supplicant_config();

        assert!(wpa_config.contains("ssid=\"TestNetwork\""));
        assert!(wpa_config.contains("psk=\"password123\""));
        assert!(wpa_config.contains("key_mgmt=WPA-PSK"));
    }

    #[test]
    fn test_open_network_config() {
        let config = WifiConfig::new_open("OpenNetwork");
        let wpa_config = config.to_wpa_supplicant_config();

        assert!(wpa_config.contains("ssid=\"OpenNetwork\""));
        assert!(wpa_config.contains("key_mgmt=NONE"));
    }

    #[test]
    fn test_hidden_network() {
        let mut config = WifiConfig::new_wpa("HiddenNet", "secret");
        config.hidden = true;
        let wpa_config = config.to_wpa_supplicant_config();

        assert!(wpa_config.contains("scan_ssid=1"));
    }
}