sal/net/src/ssh.rs

use std::path::PathBuf;
use std::process::Stdio;
use std::time::Duration;

use anyhow::Result;
use tokio::io::{AsyncReadExt, BufReader};
use tokio::process::Command;

/// SSH Connection that uses the system's SSH client
pub struct SshConnection {
    host: String,
    port: u16,
    user: String,
    identity_file: Option<PathBuf>,
    timeout: Duration,
}

impl SshConnection {
    /// Execute a command over SSH and return its output
    pub async fn execute(&self, command: &str) -> Result<(i32, String)> {
        let mut args = Vec::new();

        // Add SSH options
        args.push("-o".to_string());
        args.push(format!("ConnectTimeout={}", self.timeout.as_secs()));

        // Don't check host key to avoid prompts
        args.push("-o".to_string());
        args.push("StrictHostKeyChecking=no".to_string());

        // Specify port if not default
        if self.port != 22 {
            args.push("-p".to_string());
            args.push(self.port.to_string());
        }

        // Add identity file if provided
        if let Some(identity) = &self.identity_file {
            args.push("-i".to_string());
            args.push(identity.to_string_lossy().to_string());
        }

        // Add user and host
        args.push(format!("{}@{}", self.user, self.host));

        // Add the command to execute
        args.push(command.to_string());

        // Run the SSH command
        let mut child = Command::new("ssh")
            .args(&args)
            .stdout(Stdio::piped())
            .stderr(Stdio::piped())
            .spawn()?;

        // Collect stdout and stderr
        let stdout = child.stdout.take().unwrap();
        let stderr = child.stderr.take().unwrap();

        let mut stdout_reader = BufReader::new(stdout);
        let mut stderr_reader = BufReader::new(stderr);

        let mut output = String::new();
        stdout_reader.read_to_string(&mut output).await?;

        let mut error_output = String::new();
        stderr_reader.read_to_string(&mut error_output).await?;

        // If there's error output, append it to the regular output
        if !error_output.is_empty() {
            if !output.is_empty() {
                output.push('\n');
            }
            output.push_str(&error_output);
        }

        // Wait for the command to complete and get exit status
        let status = child.wait().await?;
        let code = status.code().unwrap_or(-1);

        Ok((code, output))
    }

    /// Check if the host is reachable via SSH
    pub async fn ping(&self) -> Result<bool> {
        let result = self.execute("echo 'Connection successful'").await?;
        Ok(result.0 == 0)
    }
}

/// Builder for SSH connections
pub struct SshConnectionBuilder {
    host: String,
    port: u16,
    user: String,
    identity_file: Option<PathBuf>,
    timeout: Duration,
}

impl Default for SshConnectionBuilder {
    fn default() -> Self {
        Self::new()
    }
}

impl SshConnectionBuilder {
    pub fn new() -> Self {
        Self {
            host: "localhost".to_string(),
            port: 22,
            user: "root".to_string(),
            identity_file: None,
            timeout: Duration::from_secs(10),
        }
    }

    pub fn host<S: Into<String>>(mut self, host: S) -> Self {
        self.host = host.into();
        self
    }

    pub fn port(mut self, port: u16) -> Self {
        self.port = port;
        self
    }

    pub fn user<S: Into<String>>(mut self, user: S) -> Self {
        self.user = user.into();
        self
    }

    pub fn identity_file(mut self, path: PathBuf) -> Self {
        self.identity_file = Some(path);
        self
    }

    pub fn timeout(mut self, timeout: Duration) -> Self {
        self.timeout = timeout;
        self
    }

    pub fn build(self) -> SshConnection {
        SshConnection {
            host: self.host,
            port: self.port,
            user: self.user,
            identity_file: self.identity_file,
            timeout: self.timeout,
        }
    }
}