Skip to main content

hardpass/
ssh.rs

1use std::path::Path;
2use std::process::{ExitStatus, Stdio};
3use std::time::{Duration, Instant};
4
5use anyhow::{Context, Result, bail};
6use tokio::process::Command;
7use tokio::time::sleep;
8
9use crate::lock::{lock_file, sibling_lock_path};
10use crate::state::SshConfig;
11
12#[derive(Debug)]
13pub struct ExecOutput {
14    pub status: ExitStatus,
15    pub stdout: String,
16    pub stderr: String,
17}
18
19pub async fn ensure_ssh_key(identity_file: &Path) -> Result<String> {
20    let _lock = lock_file(sibling_lock_path(identity_file)).await?;
21    if !identity_file.is_file() || !identity_file.with_extension("pub").is_file() {
22        if let Some(parent) = identity_file.parent() {
23            tokio::fs::create_dir_all(parent).await?;
24        }
25        let output = Command::new("ssh-keygen")
26            .arg("-q")
27            .arg("-t")
28            .arg("ed25519")
29            .arg("-N")
30            .arg("")
31            .arg("-f")
32            .arg(identity_file)
33            .arg("-C")
34            .arg("hardpass")
35            .output()
36            .await
37            .context("run ssh-keygen")?;
38        if !output.status.success() {
39            bail!(
40                "ssh-keygen failed: {}",
41                String::from_utf8_lossy(&output.stderr).trim()
42            );
43        }
44    }
45    let public_key = tokio::fs::read_to_string(identity_file.with_extension("pub")).await?;
46    Ok(public_key.trim().to_string())
47}
48
49pub async fn wait_for_ssh(config: &SshConfig, timeout_secs: u64) -> Result<()> {
50    let deadline = Instant::now() + Duration::from_secs(timeout_secs);
51    loop {
52        match ssh_status(config, &["true"]).await {
53            Ok(()) => return Ok(()),
54            Err(err) if Instant::now() < deadline => {
55                let _ = err;
56                sleep(Duration::from_millis(500)).await;
57            }
58            Err(err) => return Err(err),
59        }
60    }
61}
62
63pub async fn open_session(config: &SshConfig, extra_args: &[String]) -> Result<()> {
64    let status = Command::new("ssh")
65        .args(common_ssh_args(config, false))
66        .args(extra_args)
67        .arg(format!("{}@{}", config.user, config.host))
68        .stdin(Stdio::inherit())
69        .stdout(Stdio::inherit())
70        .stderr(Stdio::inherit())
71        .status()
72        .await
73        .context("run ssh")?;
74    if status.success() {
75        Ok(())
76    } else {
77        bail!("ssh exited with status {status}");
78    }
79}
80
81pub async fn exec(config: &SshConfig, command: &[String]) -> Result<()> {
82    let output = exec_capture(config, command).await?;
83    if output.status.success() {
84        print!("{}", output.stdout);
85        eprint!("{}", output.stderr);
86        Ok(())
87    } else {
88        if !output.stdout.is_empty() {
89            print!("{}", output.stdout);
90        }
91        if !output.stderr.is_empty() {
92            eprint!("{}", output.stderr);
93        }
94        bail!("remote command exited with status {}", output.status);
95    }
96}
97
98pub async fn exec_capture(config: &SshConfig, command: &[String]) -> Result<ExecOutput> {
99    let output = Command::new("ssh")
100        .args(common_ssh_args(config, true))
101        .arg(format!("{}@{}", config.user, config.host))
102        .args(command)
103        .stdin(Stdio::null())
104        .stdout(Stdio::piped())
105        .stderr(Stdio::piped())
106        .output()
107        .await
108        .context("run ssh")?;
109    Ok(ExecOutput {
110        status: output.status,
111        stdout: String::from_utf8_lossy(&output.stdout).into_owned(),
112        stderr: String::from_utf8_lossy(&output.stderr).into_owned(),
113    })
114}
115
116pub async fn exec_checked(config: &SshConfig, command: &[String]) -> Result<ExecOutput> {
117    let output = exec_capture(config, command).await?;
118    if output.status.success() {
119        Ok(output)
120    } else {
121        bail!("remote command exited with status {}", output.status);
122    }
123}
124
125async fn ssh_status(config: &SshConfig, remote_command: &[&str]) -> Result<()> {
126    let status = Command::new("ssh")
127        .args(common_ssh_args(config, true))
128        .arg("-o")
129        .arg("ConnectTimeout=2")
130        .arg(format!("{}@{}", config.user, config.host))
131        .args(remote_command)
132        .stdin(Stdio::null())
133        .stdout(Stdio::null())
134        .stderr(Stdio::null())
135        .status()
136        .await
137        .context("run ssh readiness probe")?;
138    if status.success() {
139        Ok(())
140    } else {
141        bail!("ssh not ready yet")
142    }
143}
144
145fn common_ssh_args(config: &SshConfig, batch_mode: bool) -> Vec<String> {
146    let mut args = vec![
147        "-i".to_string(),
148        config.identity_file.display().to_string(),
149        "-p".to_string(),
150        config.port.to_string(),
151        "-o".to_string(),
152        "StrictHostKeyChecking=no".to_string(),
153        "-o".to_string(),
154        "UserKnownHostsFile=/dev/null".to_string(),
155        "-o".to_string(),
156        "IdentitiesOnly=yes".to_string(),
157        "-o".to_string(),
158        "LogLevel=ERROR".to_string(),
159    ];
160    if batch_mode {
161        args.extend(["-o".to_string(), "BatchMode=yes".to_string()]);
162    }
163    args
164}