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

//! Tauri commands for built-in tools: ping, traceroute, WoL, keygen, passgen.

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

use crate::AppState;

// ── Ping ─────────────────────────────────────────────────────────────────────

#[derive(Debug, Serialize)]
#[serde(rename_all = "camelCase")]
pub struct PingResult {
    pub target: String,
    pub output: String,
}

/// Ping a host through an SSH session's exec channel.
#[tauri::command]
pub async fn tool_ping(
    session_id: String,
    target: String,
    count: Option<u32>,
    state: State<'_, AppState>,
) -> Result<PingResult, String> {
    let session = state.ssh.get_session(&session_id)
        .ok_or_else(|| format!("SSH session {} not found", session_id))?;
    let n = count.unwrap_or(4);
    let cmd = format!("ping -c {} {} 2>&1", n, target);
    let output = exec_on_session(&session.handle, &cmd).await?;
    Ok(PingResult { target, output })
}

/// Traceroute through an SSH session's exec channel.
#[tauri::command]
pub async fn tool_traceroute(
    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!("traceroute {} 2>&1 || tracert {} 2>&1", target, target);
    exec_on_session(&session.handle, &cmd).await
}

// ── Wake on LAN ──────────────────────────────────────────────────────────────

/// Send a Wake-on-LAN magic packet through an SSH session.
/// The remote host broadcasts the WoL packet on its local network.
#[tauri::command]
pub async fn tool_wake_on_lan(
    session_id: String,
    mac_address: 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))?;

    // Build WoL magic packet as a shell one-liner using python or perl (widely available)
    let mac_clean = mac_address.replace([':', '-'], "");
    if mac_clean.len() != 12 || !mac_clean.chars().all(|c| c.is_ascii_hexdigit()) {
        return Err(format!("Invalid MAC address: {}", mac_address));
    }

    let cmd = format!(
        r#"python3 -c "
import socket, struct
mac = bytes.fromhex('{mac_clean}')
pkt = b'\xff'*6 + mac*16
s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
s.setsockopt(socket.SOL_SOCKET, socket.SO_BROADCAST, 1)
s.sendto(pkt, ('255.255.255.255', 9))
s.close()
print('WoL packet sent to {mac_address}')
" 2>&1 || echo "python3 not available — install python3 on remote host for WoL""#
    );

    exec_on_session(&session.handle, &cmd).await
}

// ── SSH Key Generator ────────────────────────────────────────────────────────

#[derive(Debug, Serialize)]
#[serde(rename_all = "camelCase")]
pub struct GeneratedKey {
    pub private_key: String,
    pub public_key: String,
    pub fingerprint: String,
    pub key_type: String,
}

/// Generate an SSH key pair locally (no SSH session needed).
#[tauri::command]
pub fn tool_generate_ssh_key(
    key_type: String,
    comment: Option<String>,
) -> Result<GeneratedKey, String> {
    tool_generate_ssh_key_inner(&key_type, comment)
}

pub fn tool_generate_ssh_key_inner(
    key_type: &str,
    comment: Option<String>,
) -> Result<GeneratedKey, String> {
    use ssh_key::{Algorithm, HashAlg, LineEnding};

    let comment_str = comment.unwrap_or_else(|| "wraith-generated".to_string());

    let algorithm = match key_type.to_lowercase().as_str() {
        "ed25519" => Algorithm::Ed25519,
        "rsa" | "rsa-2048" => Algorithm::Rsa { hash: Some(ssh_key::HashAlg::Sha256) },
        "rsa-4096" => Algorithm::Rsa { hash: Some(ssh_key::HashAlg::Sha256) },
        _ => return Err(format!("Unsupported key type: {}. Use ed25519 or rsa", key_type)),
    };

    let private_key = ssh_key::PrivateKey::random(&mut ssh_key::rand_core::OsRng, algorithm)
        .map_err(|e| format!("Key generation failed: {}", e))?;

    let private_pem = private_key.to_openssh(LineEnding::LF)
        .map_err(|e| format!("Failed to encode private key: {}", e))?;

    let public_key = private_key.public_key();
    let public_openssh = public_key.to_openssh()
        .map_err(|e| format!("Failed to encode public key: {}", e))?;

    let fingerprint = public_key.fingerprint(HashAlg::Sha256).to_string();

    Ok(GeneratedKey {
        private_key: private_pem.to_string(),
        public_key: format!("{} {}", public_openssh, comment_str),
        fingerprint,
        key_type: key_type.to_lowercase(),
    })
}

// ── Password Generator ───────────────────────────────────────────────────────

/// Generate a cryptographically secure random password.
#[tauri::command]
pub fn tool_generate_password(
    length: Option<usize>,
    uppercase: Option<bool>,
    lowercase: Option<bool>,
    digits: Option<bool>,
    symbols: Option<bool>,
) -> Result<String, String> {
    tool_generate_password_inner(length, uppercase, lowercase, digits, symbols)
}

pub fn tool_generate_password_inner(
    length: Option<usize>,
    uppercase: Option<bool>,
    lowercase: Option<bool>,
    digits: Option<bool>,
    symbols: Option<bool>,
) -> Result<String, String> {
    use rand::Rng;

    let len = length.unwrap_or(20).max(4).min(128);
    let use_upper = uppercase.unwrap_or(true);
    let use_lower = lowercase.unwrap_or(true);
    let use_digits = digits.unwrap_or(true);
    let use_symbols = symbols.unwrap_or(true);

    let mut charset = String::new();
    if use_upper { charset.push_str("ABCDEFGHIJKLMNOPQRSTUVWXYZ"); }
    if use_lower { charset.push_str("abcdefghijklmnopqrstuvwxyz"); }
    if use_digits { charset.push_str("0123456789"); }
    if use_symbols { charset.push_str("!@#$%^&*()-_=+[]{}|;:,.<>?"); }

    if charset.is_empty() {
        return Err("At least one character class must be enabled".to_string());
    }

    let chars: Vec<char> = charset.chars().collect();
    let mut rng = rand::rng();
    let password: String = (0..len)
        .map(|_| chars[rng.random_range(0..chars.len())])
        .collect();

    Ok(password)
}

// ── 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)
}