Skip to main content

virtuoso_cli/transport/
ssh.rs

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