mur_core/

machine.rs

//! Multi-Machine — SSH remote execution, agent relay, and machine registry.

use anyhow::{Context, Result};
use chrono::{DateTime, Utc};
use serde::{Deserialize, Serialize};
use std::path::{Path, PathBuf};

// ---------------------------------------------------------------------------
// Machine Registry
// ---------------------------------------------------------------------------

/// A registered remote machine.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Machine {
    /// Unique machine identifier (user-chosen name).
    pub id: String,
    /// Hostname or IP address.
    pub host: String,
    /// SSH port (default: 22).
    #[serde(default = "default_ssh_port")]
    pub port: u16,
    /// SSH username.
    pub user: String,
    /// Path to SSH private key (optional, uses ssh-agent by default).
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub key_path: Option<String>,
    /// Optional tags for grouping machines.
    #[serde(default)]
    pub tags: Vec<String>,
    /// Working directory on the remote machine.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub working_dir: Option<String>,
    /// When this machine was registered.
    pub registered_at: DateTime<Utc>,
    /// Last health check result.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub last_health: Option<HealthStatus>,
}

fn default_ssh_port() -> u16 {
    22
}

/// Health check result for a machine.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct HealthStatus {
    /// Whether the machine is reachable.
    pub reachable: bool,
    /// Round-trip latency in milliseconds.
    pub latency_ms: Option<u64>,
    /// System load (if available).
    pub load: Option<f64>,
    /// Available memory in MB (if available).
    pub available_memory_mb: Option<u64>,
    /// Timestamp of this check.
    pub checked_at: DateTime<Utc>,
    /// Error message if unreachable.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub error: Option<String>,
}

/// Machine registry — persists machine configurations.
pub struct MachineRegistry {
    registry_path: PathBuf,
}

impl MachineRegistry {
    /// Create a new machine registry.
    pub fn new(path: &Path) -> Self {
        Self {
            registry_path: path.to_path_buf(),
        }
    }

    /// Default registry path.
    pub fn default_path() -> PathBuf {
        directories::BaseDirs::new()
            .map(|d| {
                d.home_dir()
                    .join(".mur")
                    .join("commander")
                    .join("machines.json")
            })
            .unwrap_or_else(|| PathBuf::from("/tmp/mur-commander/machines.json"))
    }

    /// List all registered machines.
    pub fn list(&self) -> Result<Vec<Machine>> {
        if !self.registry_path.exists() {
            return Ok(Vec::new());
        }
        let content =
            std::fs::read_to_string(&self.registry_path).context("Reading machine registry")?;
        serde_json::from_str(&content).context("Parsing machine registry")
    }

    /// Add a machine to the registry.
    pub fn add(&self, machine: Machine) -> Result<()> {
        let mut machines = self.list()?;
        if machines.iter().any(|m| m.id == machine.id) {
            anyhow::bail!("Machine '{}' already exists", machine.id);
        }
        machines.push(machine);
        self.save(&machines)
    }

    /// Remove a machine by ID.
    pub fn remove(&self, id: &str) -> Result<bool> {
        let mut machines = self.list()?;
        let before = machines.len();
        machines.retain(|m| m.id != id);
        if machines.len() == before {
            return Ok(false);
        }
        self.save(&machines)?;
        Ok(true)
    }

    /// Get a machine by ID.
    pub fn get(&self, id: &str) -> Result<Option<Machine>> {
        let machines = self.list()?;
        Ok(machines.into_iter().find(|m| m.id == id))
    }

    /// Update a machine's health status.
    pub fn update_health(&self, id: &str, health: HealthStatus) -> Result<()> {
        let mut machines = self.list()?;
        if let Some(machine) = machines.iter_mut().find(|m| m.id == id) {
            machine.last_health = Some(health);
        }
        self.save(&machines)
    }

    /// Save the registry.
    fn save(&self, machines: &[Machine]) -> Result<()> {
        if let Some(parent) = self.registry_path.parent() {
            std::fs::create_dir_all(parent)?;
        }
        let json = serde_json::to_string_pretty(machines)?;
        std::fs::write(&self.registry_path, json).context("Writing machine registry")
    }
}

// ---------------------------------------------------------------------------
// Remote Execution
// ---------------------------------------------------------------------------

/// Result of a remote command execution.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct RemoteExecResult {
    /// Machine that executed the command.
    pub machine_id: String,
    /// Exit code of the command.
    pub exit_code: i32,
    /// Standard output.
    pub stdout: String,
    /// Standard error.
    pub stderr: String,
    /// Execution duration in milliseconds.
    pub duration_ms: u64,
    /// Whether execution succeeded (exit_code == 0).
    pub success: bool,
}

/// Remote executor using SSH (via async process spawning with `ssh` command).
pub struct RemoteExecutor;

impl RemoteExecutor {
    /// Execute a command on a remote machine via SSH.
    pub async fn exec(machine: &Machine, command: &str) -> Result<RemoteExecResult> {
        let start = std::time::Instant::now();

        let mut ssh_args = vec![
            "-o".to_string(),
            "StrictHostKeyChecking=yes".to_string(),
            "-o".to_string(),
            "ConnectTimeout=10".to_string(),
            "-p".to_string(),
            machine.port.to_string(),
        ];

        if let Some(ref key) = machine.key_path {
            ssh_args.push("-i".into());
            ssh_args.push(key.clone());
        }

        let target = format!("{}@{}", machine.user, machine.host);
        ssh_args.push(target);

        // If a working directory is set, cd into it first
        // Escape the command to prevent shell injection on remote
        let escaped_command = shell_escape(command);
        let full_command = if let Some(ref dir) = machine.working_dir {
            format!("cd {} && sh -c {}", shell_escape(dir), escaped_command)
        } else {
            format!("sh -c {}", escaped_command)
        };
        ssh_args.push(full_command);

        let output = tokio::process::Command::new("ssh")
            .args(&ssh_args)
            .output()
            .await
            .context("Spawning SSH process")?;

        let duration = start.elapsed();
        let exit_code = output.status.code().unwrap_or(-1);

        Ok(RemoteExecResult {
            machine_id: machine.id.clone(),
            exit_code,
            stdout: String::from_utf8_lossy(&output.stdout).to_string(),
            stderr: String::from_utf8_lossy(&output.stderr).to_string(),
            duration_ms: duration.as_millis() as u64,
            success: exit_code == 0,
        })
    }

    /// Check if a machine is reachable via SSH.
    pub async fn health_check(machine: &Machine) -> HealthStatus {
        let start = std::time::Instant::now();

        let result = Self::exec(
            machine,
            "echo ok && uptime | awk '{print $NF}' && free -m 2>/dev/null | awk '/Mem:/{print $7}' || echo 0",
        )
        .await;

        let duration = start.elapsed();

        match result {
            Ok(exec_result) if exec_result.success => {
                let lines: Vec<&str> = exec_result.stdout.lines().collect();
                let load = lines
                    .get(1)
                    .and_then(|l| l.trim().parse::<f64>().ok());
                let mem = lines
                    .get(2)
                    .and_then(|l| l.trim().parse::<u64>().ok());

                HealthStatus {
                    reachable: true,
                    latency_ms: Some(duration.as_millis() as u64),
                    load,
                    available_memory_mb: mem,
                    checked_at: Utc::now(),
                    error: None,
                }
            }
            Ok(exec_result) => HealthStatus {
                reachable: false,
                latency_ms: Some(duration.as_millis() as u64),
                load: None,
                available_memory_mb: None,
                checked_at: Utc::now(),
                error: Some(exec_result.stderr),
            },
            Err(e) => HealthStatus {
                reachable: false,
                latency_ms: None,
                load: None,
                available_memory_mb: None,
                checked_at: Utc::now(),
                error: Some(e.to_string()),
            },
        }
    }
}

// ---------------------------------------------------------------------------
// Agent Relay — execute workflows on remote machines
// ---------------------------------------------------------------------------

/// Relay for executing workflows across multiple machines.
pub struct AgentRelay;

/// Result of multi-machine execution.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct RelayResult {
    /// Per-machine results.
    pub results: Vec<RemoteExecResult>,
    /// Total duration in milliseconds.
    pub total_duration_ms: u64,
    /// Number of successful executions.
    pub successes: usize,
    /// Number of failed executions.
    pub failures: usize,
}

impl AgentRelay {
    /// Execute a command on multiple machines sequentially.
    pub async fn exec_sequential(
        machines: &[Machine],
        command: &str,
    ) -> Result<RelayResult> {
        let start = std::time::Instant::now();
        let mut results = Vec::new();

        for machine in machines {
            let result = RemoteExecutor::exec(machine, command).await?;
            results.push(result);
        }

        let successes = results.iter().filter(|r| r.success).count();
        let failures = results.len() - successes;

        Ok(RelayResult {
            results,
            total_duration_ms: start.elapsed().as_millis() as u64,
            successes,
            failures,
        })
    }

    /// Execute a command on multiple machines concurrently.
    pub async fn exec_parallel(
        machines: &[Machine],
        command: &str,
    ) -> Result<RelayResult> {
        let start = std::time::Instant::now();
        let mut handles = Vec::new();

        for machine in machines {
            let machine = machine.clone();
            let cmd = command.to_string();
            handles.push(tokio::spawn(async move {
                RemoteExecutor::exec(&machine, &cmd).await
            }));
        }

        let mut results = Vec::new();
        for handle in handles {
            match handle.await {
                Ok(Ok(result)) => results.push(result),
                Ok(Err(e)) => {
                    results.push(RemoteExecResult {
                        machine_id: "unknown".into(),
                        exit_code: -1,
                        stdout: String::new(),
                        stderr: e.to_string(),
                        duration_ms: 0,
                        success: false,
                    });
                }
                Err(e) => {
                    results.push(RemoteExecResult {
                        machine_id: "unknown".into(),
                        exit_code: -1,
                        stdout: String::new(),
                        stderr: e.to_string(),
                        duration_ms: 0,
                        success: false,
                    });
                }
            }
        }

        let successes = results.iter().filter(|r| r.success).count();
        let failures = results.len() - successes;

        Ok(RelayResult {
            results,
            total_duration_ms: start.elapsed().as_millis() as u64,
            successes,
            failures,
        })
    }

    /// Run health checks on all machines.
    pub async fn health_check_all(machines: &[Machine]) -> Vec<(String, HealthStatus)> {
        let mut results = Vec::new();

        for machine in machines {
            let health = RemoteExecutor::health_check(machine).await;
            results.push((machine.id.clone(), health));
        }

        results
    }
}

/// Minimal shell escaping for a path.
fn shell_escape(s: &str) -> String {
    format!("'{}'", s.replace('\'', "'\\''"))
}

// ---------------------------------------------------------------------------
// Result Aggregation
// ---------------------------------------------------------------------------

/// Aggregate results from multi-machine execution.
pub fn aggregate_results(results: &[RemoteExecResult]) -> serde_json::Value {
    let total = results.len();
    let successes = results.iter().filter(|r| r.success).count();
    let failures = total - successes;
    let avg_duration = if total > 0 {
        results.iter().map(|r| r.duration_ms).sum::<u64>() / total as u64
    } else {
        0
    };

    serde_json::json!({
        "total_machines": total,
        "successes": successes,
        "failures": failures,
        "avg_duration_ms": avg_duration,
        "machines": results.iter().map(|r| serde_json::json!({
            "id": r.machine_id,
            "success": r.success,
            "exit_code": r.exit_code,
            "duration_ms": r.duration_ms,
        })).collect::<Vec<_>>(),
    })
}

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

    fn test_registry() -> (TempDir, MachineRegistry) {
        let dir = TempDir::new().unwrap();
        let path = dir.path().join("machines.json");
        (dir, MachineRegistry::new(&path))
    }

    fn test_machine(id: &str) -> Machine {
        Machine {
            id: id.into(),
            host: "192.168.1.100".into(),
            port: 22,
            user: "deploy".into(),
            key_path: None,
            tags: vec!["staging".into()],
            working_dir: Some("/opt/app".into()),
            registered_at: Utc::now(),
            last_health: None,
        }
    }

    #[test]
    fn test_empty_registry() {
        let (_dir, registry) = test_registry();
        assert!(registry.list().unwrap().is_empty());
    }

    #[test]
    fn test_add_machine() {
        let (_dir, registry) = test_registry();
        registry.add(test_machine("web-1")).unwrap();

        let machines = registry.list().unwrap();
        assert_eq!(machines.len(), 1);
        assert_eq!(machines[0].id, "web-1");
    }

    #[test]
    fn test_add_duplicate_machine() {
        let (_dir, registry) = test_registry();
        registry.add(test_machine("web-1")).unwrap();
        let err = registry.add(test_machine("web-1"));
        assert!(err.is_err());
    }

    #[test]
    fn test_remove_machine() {
        let (_dir, registry) = test_registry();
        registry.add(test_machine("web-1")).unwrap();

        assert!(registry.remove("web-1").unwrap());
        assert!(!registry.remove("web-1").unwrap());
        assert!(registry.list().unwrap().is_empty());
    }

    #[test]
    fn test_get_machine() {
        let (_dir, registry) = test_registry();
        registry.add(test_machine("web-1")).unwrap();

        let machine = registry.get("web-1").unwrap();
        assert!(machine.is_some());
        assert_eq!(machine.unwrap().host, "192.168.1.100");

        assert!(registry.get("nonexistent").unwrap().is_none());
    }

    #[test]
    fn test_update_health() {
        let (_dir, registry) = test_registry();
        registry.add(test_machine("web-1")).unwrap();

        let health = HealthStatus {
            reachable: true,
            latency_ms: Some(42),
            load: Some(0.5),
            available_memory_mb: Some(4096),
            checked_at: Utc::now(),
            error: None,
        };
        registry.update_health("web-1", health).unwrap();

        let machine = registry.get("web-1").unwrap().unwrap();
        assert!(machine.last_health.is_some());
        assert!(machine.last_health.unwrap().reachable);
    }

    #[test]
    fn test_aggregate_results() {
        let results = vec![
            RemoteExecResult {
                machine_id: "web-1".into(),
                exit_code: 0,
                stdout: "ok".into(),
                stderr: String::new(),
                duration_ms: 100,
                success: true,
            },
            RemoteExecResult {
                machine_id: "web-2".into(),
                exit_code: 1,
                stdout: String::new(),
                stderr: "error".into(),
                duration_ms: 200,
                success: false,
            },
        ];

        let agg = aggregate_results(&results);
        assert_eq!(agg["total_machines"], 2);
        assert_eq!(agg["successes"], 1);
        assert_eq!(agg["failures"], 1);
        assert_eq!(agg["avg_duration_ms"], 150);
    }

    #[test]
    fn test_shell_escape() {
        assert_eq!(shell_escape("/opt/app"), "'/opt/app'");
        assert_eq!(shell_escape("path with spaces"), "'path with spaces'");
        assert_eq!(
            shell_escape("it's a path"),
            "'it'\\''s a path'"
        );
    }

    #[test]
    fn test_machine_serialization() {
        let machine = test_machine("web-1");
        let json = serde_json::to_string(&machine).unwrap();
        let back: Machine = serde_json::from_str(&json).unwrap();
        assert_eq!(back.id, "web-1");
        assert_eq!(back.port, 22);
    }

    #[test]
    fn test_relay_result_serialization() {
        let result = RelayResult {
            results: vec![],
            total_duration_ms: 500,
            successes: 0,
            failures: 0,
        };
        let json = serde_json::to_string(&result).unwrap();
        assert!(json.contains("500"));
    }
}