wraith/src-tauri/src/ssh/monitor.rs

//! Remote system monitoring via SSH exec channels.
//!
//! Periodically runs lightweight system commands over a separate exec channel
//! (same pattern as CWD tracker) and emits stats to the frontend.
//! No agent installation required — uses standard POSIX and platform commands.

use std::sync::Arc;

use russh::client::Handle;
use russh::ChannelMsg;
use serde::Serialize;
use tauri::{AppHandle, Emitter};
use tokio::sync::Mutex as TokioMutex;

use crate::ssh::session::SshClient;

#[derive(Debug, Serialize, Clone, Default)]
#[serde(rename_all = "camelCase")]
pub struct SystemStats {
    pub cpu_percent: f64,
    pub mem_used_mb: u64,
    pub mem_total_mb: u64,
    pub mem_percent: f64,
    pub disk_used_gb: f64,
    pub disk_total_gb: f64,
    pub disk_percent: f64,
    pub net_rx_bytes: u64,
    pub net_tx_bytes: u64,
    pub os_type: String,
}

/// Spawn a background task that polls system stats every 5 seconds.
pub fn start_monitor(
    handle: Arc<TokioMutex<Handle<SshClient>>>,
    app_handle: AppHandle,
    session_id: String,
) {
    tokio::spawn(async move {
        // Brief delay to let the shell start up
        tokio::time::sleep(tokio::time::Duration::from_secs(2)).await;

        loop {
            let stats = collect_stats(&handle).await;

            if let Some(stats) = stats {
                let _ = app_handle.emit(
                    &format!("ssh:monitor:{}", session_id),
                    &stats,
                );
            }

            tokio::time::sleep(tokio::time::Duration::from_secs(5)).await;
        }
    });
}

async fn collect_stats(handle: &Arc<TokioMutex<Handle<SshClient>>>) -> Option<SystemStats> {
    // Single command that works cross-platform: detect OS then gather stats
    let script = r#"
OS=$(uname -s 2>/dev/null || echo "Unknown")
if [ "$OS" = "Linux" ]; then
    CPU=$(grep 'cpu ' /proc/stat | awk '{usage=($2+$4)*100/($2+$4+$5)} END {printf "%.1f", usage}')
    MEM=$(free -m 2>/dev/null | awk '/^Mem:/ {printf "%d %d", $3, $2}')
    DISK=$(df -BG / 2>/dev/null | awk 'NR==2 {gsub("G",""); printf "%s %s", $3, $2}')
    NET=$(cat /proc/net/dev 2>/dev/null | awk '/eth0:|ens|enp|wlan0:/ {gsub(":",""); printf "%s %s", $2, $10; exit}')
    echo "WRAITH_STATS:$OS:$CPU:$MEM:$DISK:$NET"
elif [ "$OS" = "Darwin" ]; then
    CPU=$(ps -A -o %cpu | awk '{s+=$1} END {printf "%.1f", s/4}')
    MEM_PAGES=$(vm_stat 2>/dev/null | awk '/Pages active/ {gsub(/\./,""); print $3}')
    MEM_TOTAL=$(sysctl -n hw.memsize 2>/dev/null | awk '{printf "%d", $1/1048576}')
    MEM_USED=$(echo "$MEM_PAGES" | awk -v t="$MEM_TOTAL" '{printf "%d", $1*4096/1048576}')
    DISK=$(df -g / 2>/dev/null | awk 'NR==2 {printf "%s %s", $3, $2}')
    NET=$(netstat -ib 2>/dev/null | awk '/en0/ && /Link/ {printf "%s %s", $7, $10; exit}')
    echo "WRAITH_STATS:$OS:$CPU:$MEM_USED $MEM_TOTAL:$DISK:$NET"
else
    echo "WRAITH_STATS:$OS:0:0 0:0 0:0 0"
fi
"#;

    let output = exec_command(handle, script).await?;

    for line in output.lines() {
        if let Some(rest) = line.strip_prefix("WRAITH_STATS:") {
            return parse_stats(rest);
        }
    }

    None
}

fn parse_stats(raw: &str) -> Option<SystemStats> {
    let parts: Vec<&str> = raw.split(':').collect();
    if parts.len() < 5 {
        return None;
    }

    let os_type = parts[0].to_string();
    let cpu_percent = parts[1].parse::<f64>().unwrap_or(0.0);

    let mem_parts: Vec<&str> = parts[2].split_whitespace().collect();
    let mem_used = mem_parts.first().and_then(|s| s.parse::<u64>().ok()).unwrap_or(0);
    let mem_total = mem_parts.get(1).and_then(|s| s.parse::<u64>().ok()).unwrap_or(1);
    let mem_percent = if mem_total > 0 { (mem_used as f64 / mem_total as f64) * 100.0 } else { 0.0 };

    let disk_parts: Vec<&str> = parts[3].split_whitespace().collect();
    let disk_used = disk_parts.first().and_then(|s| s.parse::<f64>().ok()).unwrap_or(0.0);
    let disk_total = disk_parts.get(1).and_then(|s| s.parse::<f64>().ok()).unwrap_or(1.0);
    let disk_percent = if disk_total > 0.0 { (disk_used / disk_total) * 100.0 } else { 0.0 };

    let net_parts: Vec<&str> = parts.get(4).unwrap_or(&"0 0").split_whitespace().collect();
    let net_rx = net_parts.first().and_then(|s| s.parse::<u64>().ok()).unwrap_or(0);
    let net_tx = net_parts.get(1).and_then(|s| s.parse::<u64>().ok()).unwrap_or(0);

    Some(SystemStats {
        cpu_percent,
        mem_used_mb: mem_used,
        mem_total_mb: mem_total,
        mem_percent,
        disk_used_gb: disk_used,
        disk_total_gb: disk_total,
        disk_percent,
        net_rx_bytes: net_rx,
        net_tx_bytes: net_tx,
        os_type,
    })
}

async fn exec_command(handle: &Arc<TokioMutex<Handle<SshClient>>>, cmd: &str) -> Option<String> {
    let mut channel = {
        let h = handle.lock().await;
        h.channel_open_session().await.ok()?
    };

    channel.exec(true, cmd).await.ok()?;

    let mut output = String::new();
    loop {
        match channel.wait().await {
            Some(ChannelMsg::Data { ref data }) => {
                if let Ok(text) = std::str::from_utf8(data.as_ref()) {
                    output.push_str(text);
                }
            }
            Some(ChannelMsg::Eof) | Some(ChannelMsg::Close) | None => break,
            Some(ChannelMsg::ExitStatus { .. }) => {}
            _ => {}
        }
    }

    Some(output)
}