wraith/src-tauri/src/commands/tools_commands_r2.rs

//! Tauri commands for Tools Round 2: DNS, Whois, Bandwidth, Subnet Calculator.

use tauri::State;
use serde::Serialize;

use crate::AppState;

// ── DNS Lookup ───────────────────────────────────────────────────────────────

#[tauri::command]
pub async fn tool_dns_lookup(
    session_id: String,
    domain: String,
    record_type: Option<String>,
    state: State<'_, AppState>,
) -> Result<String, String> {
    let session = state.ssh.get_session(&session_id)
        .ok_or_else(|| format!("SSH session {} not found", session_id))?;
    let rtype = record_type.unwrap_or_else(|| "A".to_string());
    let cmd = format!(
        r#"dig {} {} +short 2>/dev/null || nslookup -type={} {} 2>/dev/null || host -t {} {} 2>/dev/null"#,
        domain, rtype, rtype, domain, rtype, domain
    );
    exec_on_session(&session.handle, &cmd).await
}

// ── Whois ────────────────────────────────────────────────────────────────────

#[tauri::command]
pub async fn tool_whois(
    session_id: String,
    target: String,
    state: State<'_, AppState>,
) -> Result<String, String> {
    let session = state.ssh.get_session(&session_id)
        .ok_or_else(|| format!("SSH session {} not found", session_id))?;
    let cmd = format!("whois {} 2>&1 | head -80", target);
    exec_on_session(&session.handle, &cmd).await
}

// ── Bandwidth Test ───────────────────────────────────────────────────────────

#[tauri::command]
pub async fn tool_bandwidth_iperf(
    session_id: String,
    server: String,
    duration: Option<u32>,
    state: State<'_, AppState>,
) -> Result<String, String> {
    let session = state.ssh.get_session(&session_id)
        .ok_or_else(|| format!("SSH session {} not found", session_id))?;
    let dur = duration.unwrap_or(5);
    let cmd = format!(
        "iperf3 -c {} -t {} --json 2>/dev/null || iperf3 -c {} -t {} 2>&1 || echo 'iperf3 not installed — run: apt install iperf3 / brew install iperf3'",
        server, dur, server, dur
    );
    exec_on_session(&session.handle, &cmd).await
}

#[tauri::command]
pub async fn tool_bandwidth_speedtest(
    session_id: String,
    state: State<'_, AppState>,
) -> Result<String, String> {
    let session = state.ssh.get_session(&session_id)
        .ok_or_else(|| format!("SSH session {} not found", session_id))?;
    // Try multiple speedtest tools in order of preference
    let cmd = r#"
if command -v speedtest-cli >/dev/null 2>&1; then
    speedtest-cli --simple 2>&1
elif command -v speedtest >/dev/null 2>&1; then
    speedtest --simple 2>&1
elif command -v curl >/dev/null 2>&1; then
    echo "=== Download speed (curl) ==="
    curl -o /dev/null -w "Download: %{speed_download} bytes/sec (%{size_download} bytes in %{time_total}s)\n" https://speed.cloudflare.com/__down?bytes=25000000 2>/dev/null
    echo "=== Upload speed (curl) ==="
    dd if=/dev/zero bs=1M count=10 2>/dev/null | curl -X POST -o /dev/null -w "Upload: %{speed_upload} bytes/sec (%{size_upload} bytes in %{time_total}s)\n" -d @- https://speed.cloudflare.com/__up 2>/dev/null
else
    echo "No speedtest tool found. Install: pip install speedtest-cli"
fi
"#;
    exec_on_session(&session.handle, cmd).await
}

// ── Subnet Calculator ────────────────────────────────────────────────────────

#[derive(Debug, Serialize)]
#[serde(rename_all = "camelCase")]
pub struct SubnetInfo {
    pub cidr: String,
    pub network: String,
    pub broadcast: String,
    pub netmask: String,
    pub wildcard: String,
    pub first_host: String,
    pub last_host: String,
    pub total_hosts: u64,
    pub usable_hosts: u64,
    pub prefix_length: u8,
    pub class: String,
    pub is_private: bool,
}

/// Pure Rust subnet calculator — no SSH session needed.
#[tauri::command]
pub fn tool_subnet_calc(cidr: String) -> Result<SubnetInfo, String> {
    tool_subnet_calc_inner(&cidr)
}

pub fn tool_subnet_calc_inner(cidr: &str) -> Result<SubnetInfo, String> {
    let cidr = cidr.to_string();
    let parts: Vec<&str> = cidr.split('/').collect();
    if parts.len() != 2 {
        return Err("Expected CIDR notation: e.g. 192.168.1.0/24".to_string());
    }

    let ip_str = parts[0];
    let prefix: u8 = parts[1].parse()
        .map_err(|_| format!("Invalid prefix length: {}", parts[1]))?;

    if prefix > 32 {
        return Err(format!("Prefix length must be 0-32, got {}", prefix));
    }

    let octets: Vec<u8> = ip_str.split('.')
        .map(|o| o.parse::<u8>())
        .collect::<Result<Vec<_>, _>>()
        .map_err(|_| format!("Invalid IP address: {}", ip_str))?;

    if octets.len() != 4 {
        return Err(format!("Invalid IP address: {}", ip_str));
    }

    let ip: u32 = (octets[0] as u32) << 24
        | (octets[1] as u32) << 16
        | (octets[2] as u32) << 8
        | (octets[3] as u32);

    let mask: u32 = if prefix == 0 { 0 } else { !0u32 << (32 - prefix) };
    let wildcard = !mask;
    let network = ip & mask;
    let broadcast = network | wildcard;
    let first_host = if prefix >= 31 { network } else { network + 1 };
    let last_host = if prefix >= 31 { broadcast } else { broadcast - 1 };
    let total: u64 = 1u64 << (32 - prefix as u64);
    let usable = if prefix >= 31 { total } else { total - 2 };

    let class = match octets[0] {
        0..=127 => "A",
        128..=191 => "B",
        192..=223 => "C",
        224..=239 => "D (Multicast)",
        _ => "E (Reserved)",
    };

    let is_private = matches!(
        (octets[0], octets[1]),
        (10, _) | (172, 16..=31) | (192, 168)
    );

    Ok(SubnetInfo {
        cidr: format!("{}/{}", to_ip(network), prefix),
        network: to_ip(network),
        broadcast: to_ip(broadcast),
        netmask: to_ip(mask),
        wildcard: to_ip(wildcard),
        first_host: to_ip(first_host),
        last_host: to_ip(last_host),
        total_hosts: total,
        usable_hosts: usable,
        prefix_length: prefix,
        class: class.to_string(),
        is_private,
    })
}

fn to_ip(val: u32) -> String {
    format!("{}.{}.{}.{}", val >> 24, (val >> 16) & 0xFF, (val >> 8) & 0xFF, val & 0xFF)
}

// ── Helper ───────────────────────────────────────────────────────────────────

async fn exec_on_session(
    handle: &std::sync::Arc<tokio::sync::Mutex<russh::client::Handle<crate::ssh::session::SshClient>>>,
    cmd: &str,
) -> Result<String, String> {
    let mut channel = {
        let h = handle.lock().await;
        h.channel_open_session().await.map_err(|e| format!("Exec channel failed: {}", e))?
    };
    channel.exec(true, cmd).await.map_err(|e| format!("Exec failed: {}", e))?;
    let mut output = String::new();
    loop {
        match channel.wait().await {
            Some(russh::ChannelMsg::Data { ref data }) => {
                if let Ok(text) = std::str::from_utf8(data.as_ref()) { output.push_str(text); }
            }
            Some(russh::ChannelMsg::Eof) | Some(russh::ChannelMsg::Close) | None => break,
            Some(russh::ChannelMsg::ExitStatus { .. }) => {}
            _ => {}
        }
    }
    Ok(output)
}