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}