virtuoso_cli/transport/
ssh.rs

1use crate::error::{Result, VirtuosoError};
2use std::collections::HashMap;
3use std::io::Write;
4use std::process::{Command, Stdio};
5use std::time::Instant;
6
7use crate::models::RemoteTaskResult;
8
9fn shell_quote(s: &str) -> String {
10    shlex::try_quote(s)
11        .unwrap_or(std::borrow::Cow::Borrowed(s))
12        .into_owned()
13}
14
15pub struct SSHRunner {
16    pub host: String,
17    pub user: Option<String>,
18    pub jump_host: Option<String>,
19    pub jump_user: Option<String>,
20    pub ssh_port: Option<u16>,
21    pub ssh_key_path: Option<String>,
22    pub ssh_config_path: Option<String>,
23    pub timeout: u64,
24    pub connect_timeout: u64,
25    pub verbose: bool,
26}
27
28impl SSHRunner {
29    pub fn new(host: &str) -> Self {
30        Self {
31            host: host.into(),
32            user: None,
33            jump_host: None,
34            jump_user: None,
35            ssh_port: None,
36            ssh_key_path: None,
37            ssh_config_path: None,
38            timeout: 30,
39            connect_timeout: 10,
40            verbose: false,
41        }
42    }
43
44    pub fn with_jump(mut self, jump: &str) -> Self {
45        self.jump_host = Some(jump.into());
46        self
47    }
48
49    pub fn with_user(mut self, user: &str) -> Self {
50        self.user = Some(user.into());
51        self
52    }
53
54    pub fn test_connection(&self, timeout: Option<u64>) -> Result<bool> {
55        let _timeout = timeout.unwrap_or(self.connect_timeout);
56        let mut cmd = self.build_ssh_cmd();
57        cmd.arg("exit").arg("0");
58
59        let output = cmd
60            .stdin(Stdio::null())
61            .stdout(Stdio::null())
62            .stderr(Stdio::null())
63            .output()
64            .map_err(|e| VirtuosoError::Ssh(format!("failed to run ssh: {e}")))?;
65
66        Ok(output.status.success())
67    }
68
69    pub fn run_command(&self, command: &str, timeout: Option<u64>) -> Result<RemoteTaskResult> {
70        let _timeout = timeout.unwrap_or(self.timeout);
71        let start = Instant::now();
72
73        let mut cmd = self.build_ssh_cmd();
74        cmd.arg("sh").arg("-s");
75
76        let output = cmd
77            .stdin(Stdio::piped())
78            .stdout(Stdio::piped())
79            .stderr(Stdio::piped())
80            .spawn()
81            .map_err(|e| VirtuosoError::Ssh(format!("failed to spawn ssh: {e}")))?;
82
83        if let Some(mut stdin) = output.stdin.as_ref() {
84            stdin
85                .write_all(command.as_bytes())
86                .map_err(|e| VirtuosoError::Ssh(format!("failed to write command: {e}")))?;
87        }
88
89        let output = output
90            .wait_with_output()
91            .map_err(|e| VirtuosoError::Ssh(format!("ssh failed: {e}")))?;
92
93        let elapsed = start.elapsed().as_secs_f64();
94        let _stdout = String::from_utf8_lossy(&output.stdout).to_string();
95        let stderr_str = String::from_utf8_lossy(&output.stderr).to_string();
96
97        let mut timings = HashMap::new();
98        timings.insert("total".into(), elapsed);
99
100        Ok(RemoteTaskResult {
101            success: output.status.success(),
102            returncode: output.status.code().unwrap_or(-1),
103            stdout: String::from_utf8_lossy(&output.stdout).to_string(),
104            stderr: stderr_str.clone(),
105            remote_dir: None,
106            error: if output.status.success() {
107                None
108            } else {
109                Some(self.summarize_error(&stderr_str))
110            },
111            timings,
112        })
113    }
114
115    pub fn upload(&self, local: &str, remote: &str) -> Result<()> {
116        let _target = self.remote_target();
117
118        let status = Command::new("tar")
119            .arg("cf")
120            .arg("-")
121            .arg("-C")
122            .arg(
123                std::path::Path::new(local)
124                    .parent()
125                    .unwrap_or_else(|| std::path::Path::new(".")),
126            )
127            .arg(std::path::Path::new(local).file_name().unwrap_or_default())
128            .stdout(Stdio::piped())
129            .spawn()
130            .map_err(|e| VirtuosoError::Ssh(format!("tar failed: {e}")))?;
131
132        let tar_stdout = status.stdout.unwrap();
133
134        let remote_dir = std::path::Path::new(remote)
135            .parent()
136            .unwrap_or_else(|| std::path::Path::new("."))
137            .to_string_lossy();
138
139        let mut ssh = self.build_ssh_cmd();
140        let quoted_dir = shell_quote(&remote_dir);
141        ssh.arg("sh").arg("-c").arg(format!(
142            "mkdir -p {quoted_dir} && cd {quoted_dir} && tar xf -"
143        ));
144        ssh.stdin(tar_stdout);
145
146        let output = ssh
147            .stdout(Stdio::null())
148            .stderr(Stdio::piped())
149            .output()
150            .map_err(|e| VirtuosoError::Ssh(format!("ssh upload failed: {e}")))?;
151
152        if !output.status.success() {
153            let stderr = String::from_utf8_lossy(&output.stderr);
154            return Err(VirtuosoError::Ssh(format!("upload failed: {stderr}")));
155        }
156
157        Ok(())
158    }
159
160    pub fn upload_text(&self, text: &str, remote: &str) -> Result<()> {
161        let remote_dir = std::path::Path::new(remote)
162            .parent()
163            .unwrap_or_else(|| std::path::Path::new("."))
164            .to_string_lossy();
165
166        let quoted_dir = shell_quote(&remote_dir);
167        let mkdir = self.run_command(&format!("mkdir -p {quoted_dir}"), None)?;
168        if !mkdir.success {
169            return Err(VirtuosoError::Ssh(format!(
170                "failed to create remote dir: {}",
171                mkdir.stderr
172            )));
173        }
174
175        let mut cmd = self.build_ssh_cmd();
176        let quoted_remote = shell_quote(remote);
177        cmd.arg("sh")
178            .arg("-c")
179            .arg(format!("cat > {quoted_remote}"));
180
181        let output = cmd
182            .stdin(Stdio::piped())
183            .stdout(Stdio::null())
184            .stderr(Stdio::piped())
185            .spawn()
186            .map_err(|e| VirtuosoError::Ssh(format!("ssh failed: {e}")))?;
187
188        if let Some(mut stdin) = output.stdin.as_ref() {
189            stdin
190                .write_all(text.as_bytes())
191                .map_err(|e| VirtuosoError::Ssh(format!("write failed: {e}")))?;
192        }
193
194        let output = output
195            .wait_with_output()
196            .map_err(|e| VirtuosoError::Ssh(format!("upload failed: {e}")))?;
197
198        if !output.status.success() {
199            let stderr = String::from_utf8_lossy(&output.stderr);
200            return Err(VirtuosoError::Ssh(format!("upload failed: {stderr}")));
201        }
202
203        Ok(())
204    }
205
206    pub fn download(&self, remote: &str, local: &str) -> Result<()> {
207        let _target = self.remote_target();
208
209        let local_path = std::path::Path::new(local);
210        if let Some(parent) = local_path.parent() {
211            std::fs::create_dir_all(parent)
212                .map_err(|e| VirtuosoError::Ssh(format!("failed to create local dir: {e}")))?;
213        }
214
215        let mut cmd = self.build_ssh_cmd();
216        cmd.arg("cat").arg(remote);
217
218        let output = cmd
219            .stdout(Stdio::piped())
220            .stderr(Stdio::piped())
221            .output()
222            .map_err(|e| VirtuosoError::Ssh(format!("ssh download failed: {e}")))?;
223
224        if !output.status.success() {
225            let stderr = String::from_utf8_lossy(&output.stderr);
226            return Err(VirtuosoError::Ssh(format!("download failed: {stderr}")));
227        }
228
229        std::fs::write(local, output.stdout)
230            .map_err(|e| VirtuosoError::Ssh(format!("failed to write local file: {e}")))?;
231
232        Ok(())
233    }
234
235    pub fn detect_python(&self) -> Result<Option<String>> {
236        for py in &["python3", "python", "python2.7"] {
237            let result = self.run_command(&format!("which {py} 2>/dev/null"), None)?;
238            if result.success && !result.stdout.trim().is_empty() {
239                return Ok(Some(py.to_string()));
240            }
241        }
242        Ok(None)
243    }
244
245    pub fn detect_arch(&self) -> Result<String> {
246        let result = self.run_command("uname -m", None)?;
247        if result.success {
248            Ok(result.stdout.trim().to_string())
249        } else {
250            Err(VirtuosoError::Ssh(format!(
251                "failed to detect arch: {}",
252                result.stderr
253            )))
254        }
255    }
256
257    pub(crate) fn build_ssh_cmd(&self) -> Command {
258        let mut cmd = Command::new("ssh");
259        cmd.args([
260            "-o",
261            "BatchMode=yes",
262            "-o",
263            "StrictHostKeyChecking=accept-new",
264            "-o",
265            &format!("ConnectTimeout={}", self.connect_timeout),
266        ]);
267
268        // ControlMaster: reuse SSH connections to avoid repeated handshakes
269        let control_dir = dirs::cache_dir()
270            .unwrap_or_else(|| std::path::PathBuf::from("/tmp"))
271            .join("virtuoso_bridge")
272            .join("ssh");
273        let _ = std::fs::create_dir_all(&control_dir);
274        let control_path = control_dir.join("%h-%p-%r");
275        cmd.args([
276            "-o",
277            "ControlMaster=auto",
278            "-o",
279            &format!("ControlPath={}", control_path.display()),
280            "-o",
281            "ControlPersist=600",
282        ]);
283
284        if let Some(port) = self.ssh_port {
285            cmd.arg("-p").arg(port.to_string());
286        }
287        if let Some(ref key) = self.ssh_key_path {
288            cmd.arg("-i").arg(key);
289        }
290        if let Some(ref config) = self.ssh_config_path {
291            cmd.arg("-F").arg(config);
292        }
293        if let Some(ref jump) = self.jump_host {
294            let jump_target = match &self.jump_user {
295                Some(u) => format!("{u}@{jump}"),
296                None => jump.clone(),
297            };
298            cmd.arg("-J").arg(jump_target);
299        }
300
301        cmd.arg(self.remote_target());
302        cmd
303    }
304
305    pub fn remote_target(&self) -> String {
306        match &self.user {
307            Some(u) => format!("{u}@{}", self.host),
308            None => self.host.clone(),
309        }
310    }
311
312    pub(crate) fn summarize_error(&self, stderr: &str) -> String {
313        let lower = stderr.to_lowercase();
314        if lower.contains("connection refused") {
315            "connection refused - check if SSH is running".into()
316        } else if lower.contains("authentication") || lower.contains("permission denied") {
317            "authentication failed - check SSH keys".into()
318        } else if lower.contains("timeout") || lower.contains("timed out") {
319            "connection timed out - check network".into()
320        } else if lower.contains("could not resolve") {
321            "hostname resolution failed - check DNS".into()
322        } else {
323            stderr.lines().take(3).collect::<Vec<_>>().join("; ")
324        }
325    }
326}
virtuoso_cli/transport/ssh.rs

virtuoso_cli/transport/
ssh.rs
