sal_net/
ssh.rs

1use std::path::PathBuf;
2use std::process::Stdio;
3use std::time::Duration;
4
5use anyhow::Result;
6use tokio::io::{AsyncReadExt, BufReader};
7use tokio::process::Command;
8
9/// SSH Connection that uses the system's SSH client
10pub struct SshConnection {
11    host: String,
12    port: u16,
13    user: String,
14    identity_file: Option<PathBuf>,
15    timeout: Duration,
16}
17
18impl SshConnection {
19    /// Execute a command over SSH and return its output
20    pub async fn execute(&self, command: &str) -> Result<(i32, String)> {
21        let mut args = Vec::new();
22
23        // Add SSH options
24        args.push("-o".to_string());
25        args.push(format!("ConnectTimeout={}", self.timeout.as_secs()));
26
27        // Don't check host key to avoid prompts
28        args.push("-o".to_string());
29        args.push("StrictHostKeyChecking=no".to_string());
30
31        // Specify port if not default
32        if self.port != 22 {
33            args.push("-p".to_string());
34            args.push(self.port.to_string());
35        }
36
37        // Add identity file if provided
38        if let Some(identity) = &self.identity_file {
39            args.push("-i".to_string());
40            args.push(identity.to_string_lossy().to_string());
41        }
42
43        // Add user and host
44        args.push(format!("{}@{}", self.user, self.host));
45
46        // Add the command to execute
47        args.push(command.to_string());
48
49        // Run the SSH command
50        let mut child = Command::new("ssh")
51            .args(&args)
52            .stdout(Stdio::piped())
53            .stderr(Stdio::piped())
54            .spawn()?;
55
56        // Collect stdout and stderr
57        let stdout = child.stdout.take().unwrap();
58        let stderr = child.stderr.take().unwrap();
59
60        let mut stdout_reader = BufReader::new(stdout);
61        let mut stderr_reader = BufReader::new(stderr);
62
63        let mut output = String::new();
64        stdout_reader.read_to_string(&mut output).await?;
65
66        let mut error_output = String::new();
67        stderr_reader.read_to_string(&mut error_output).await?;
68
69        // If there's error output, append it to the regular output
70        if !error_output.is_empty() {
71            if !output.is_empty() {
72                output.push('\n');
73            }
74            output.push_str(&error_output);
75        }
76
77        // Wait for the command to complete and get exit status
78        let status = child.wait().await?;
79        let code = status.code().unwrap_or(-1);
80
81        Ok((code, output))
82    }
83
84    /// Check if the host is reachable via SSH
85    pub async fn ping(&self) -> Result<bool> {
86        let result = self.execute("echo 'Connection successful'").await?;
87        Ok(result.0 == 0)
88    }
89}
90
91/// Builder for SSH connections
92pub struct SshConnectionBuilder {
93    host: String,
94    port: u16,
95    user: String,
96    identity_file: Option<PathBuf>,
97    timeout: Duration,
98}
99
100impl Default for SshConnectionBuilder {
101    fn default() -> Self {
102        Self::new()
103    }
104}
105
106impl SshConnectionBuilder {
107    pub fn new() -> Self {
108        Self {
109            host: "localhost".to_string(),
110            port: 22,
111            user: "root".to_string(),
112            identity_file: None,
113            timeout: Duration::from_secs(10),
114        }
115    }
116
117    pub fn host<S: Into<String>>(mut self, host: S) -> Self {
118        self.host = host.into();
119        self
120    }
121
122    pub fn port(mut self, port: u16) -> Self {
123        self.port = port;
124        self
125    }
126
127    pub fn user<S: Into<String>>(mut self, user: S) -> Self {
128        self.user = user.into();
129        self
130    }
131
132    pub fn identity_file(mut self, path: PathBuf) -> Self {
133        self.identity_file = Some(path);
134        self
135    }
136
137    pub fn timeout(mut self, timeout: Duration) -> Self {
138        self.timeout = timeout;
139        self
140    }
141
142    pub fn build(self) -> SshConnection {
143        SshConnection {
144            host: self.host,
145            port: self.port,
146            user: self.user,
147            identity_file: self.identity_file,
148            timeout: self.timeout,
149        }
150    }
151}