hardpass/

ssh.rs

use std::path::Path;
use std::process::{ExitStatus, Stdio};
use std::time::{Duration, Instant};

use anyhow::{Context, Result, bail};
use tokio::process::Command;
use tokio::time::sleep;

use crate::lock::{lock_file, sibling_lock_path};
use crate::state::SshConfig;

#[derive(Debug)]
pub struct ExecOutput {
    pub status: ExitStatus,
    pub stdout: String,
    pub stderr: String,
}

pub async fn ensure_ssh_key(identity_file: &Path) -> Result<String> {
    let _lock = lock_file(sibling_lock_path(identity_file)).await?;
    if !identity_file.is_file() || !identity_file.with_extension("pub").is_file() {
        if let Some(parent) = identity_file.parent() {
            tokio::fs::create_dir_all(parent).await?;
        }
        let output = Command::new("ssh-keygen")
            .arg("-q")
            .arg("-t")
            .arg("ed25519")
            .arg("-N")
            .arg("")
            .arg("-f")
            .arg(identity_file)
            .arg("-C")
            .arg("hardpass")
            .output()
            .await
            .context("run ssh-keygen")?;
        if !output.status.success() {
            bail!(
                "ssh-keygen failed: {}",
                String::from_utf8_lossy(&output.stderr).trim()
            );
        }
    }
    let public_key = tokio::fs::read_to_string(identity_file.with_extension("pub")).await?;
    Ok(public_key.trim().to_string())
}

pub async fn wait_for_ssh(config: &SshConfig, timeout_secs: u64) -> Result<()> {
    let deadline = Instant::now() + Duration::from_secs(timeout_secs);
    loop {
        match ssh_status(config, &["true"]).await {
            Ok(()) => return Ok(()),
            Err(err) if Instant::now() < deadline => {
                let _ = err;
                sleep(Duration::from_millis(500)).await;
            }
            Err(err) => return Err(err),
        }
    }
}

pub async fn open_session(config: &SshConfig, extra_args: &[String]) -> Result<()> {
    let status = Command::new("ssh")
        .args(common_ssh_args(config, false))
        .args(extra_args)
        .arg(format!("{}@{}", config.user, config.host))
        .stdin(Stdio::inherit())
        .stdout(Stdio::inherit())
        .stderr(Stdio::inherit())
        .status()
        .await
        .context("run ssh")?;
    if status.success() {
        Ok(())
    } else {
        bail!("ssh exited with status {status}");
    }
}

pub async fn exec(config: &SshConfig, command: &[String]) -> Result<()> {
    let output = exec_capture(config, command).await?;
    if output.status.success() {
        print!("{}", output.stdout);
        eprint!("{}", output.stderr);
        Ok(())
    } else {
        if !output.stdout.is_empty() {
            print!("{}", output.stdout);
        }
        if !output.stderr.is_empty() {
            eprint!("{}", output.stderr);
        }
        bail!("remote command exited with status {}", output.status);
    }
}

pub async fn exec_capture(config: &SshConfig, command: &[String]) -> Result<ExecOutput> {
    let remote_command = render_remote_command(command);
    let output = Command::new("ssh")
        .args(common_ssh_args(config, true))
        .arg(format!("{}@{}", config.user, config.host))
        .arg(&remote_command)
        .stdin(Stdio::null())
        .stdout(Stdio::piped())
        .stderr(Stdio::piped())
        .output()
        .await
        .context("run ssh")?;
    Ok(ExecOutput {
        status: output.status,
        stdout: String::from_utf8_lossy(&output.stdout).into_owned(),
        stderr: String::from_utf8_lossy(&output.stderr).into_owned(),
    })
}

pub async fn exec_checked(config: &SshConfig, command: &[String]) -> Result<ExecOutput> {
    let output = exec_capture(config, command).await?;
    if output.status.success() {
        Ok(output)
    } else {
        bail!("remote command exited with status {}", output.status);
    }
}

async fn ssh_status(config: &SshConfig, remote_command: &[&str]) -> Result<()> {
    let status = Command::new("ssh")
        .args(common_ssh_args(config, true))
        .arg("-o")
        .arg("ConnectTimeout=2")
        .arg(format!("{}@{}", config.user, config.host))
        .args(remote_command)
        .stdin(Stdio::null())
        .stdout(Stdio::null())
        .stderr(Stdio::null())
        .status()
        .await
        .context("run ssh readiness probe")?;
    if status.success() {
        Ok(())
    } else {
        bail!("ssh not ready yet")
    }
}

fn common_ssh_args(config: &SshConfig, batch_mode: bool) -> Vec<String> {
    let mut args = vec![
        "-i".to_string(),
        config.identity_file.display().to_string(),
        "-p".to_string(),
        config.port.to_string(),
        "-o".to_string(),
        "StrictHostKeyChecking=no".to_string(),
        "-o".to_string(),
        "UserKnownHostsFile=/dev/null".to_string(),
        "-o".to_string(),
        "IdentitiesOnly=yes".to_string(),
        "-o".to_string(),
        "LogLevel=ERROR".to_string(),
    ];
    if batch_mode {
        args.extend(["-o".to_string(), "BatchMode=yes".to_string()]);
    }
    args
}

fn render_remote_command(command: &[String]) -> String {
    command
        .iter()
        .map(|arg| shell_quote(arg))
        .collect::<Vec<_>>()
        .join(" ")
}

fn shell_quote(arg: &str) -> String {
    format!("'{}'", arg.replace('\'', "'\"'\"'"))
}

#[cfg(test)]
mod tests {
    use super::{render_remote_command, shell_quote};

    #[test]
    fn shell_quote_escapes_single_quotes() {
        assert_eq!(shell_quote("it's ready"), r#"'it'"'"'s ready'"#);
    }

    #[test]
    fn render_remote_command_preserves_shell_script_argument() {
        let command = vec![
            "sh".to_string(),
            "-lc".to_string(),
            "sudo apt-get update".to_string(),
        ];
        assert_eq!(
            render_remote_command(&command),
            "'sh' '-lc' 'sudo apt-get update'"
        );
    }

    #[test]
    fn render_remote_command_preserves_empty_arguments() {
        let command = vec!["printf".to_string(), "".to_string(), "done".to_string()];
        assert_eq!(render_remote_command(&command), "'printf' '' 'done'");
    }
}