xboct-deploy/src/ssh.rs

use std::fs::File;
use std::path::PathBuf;
use std::sync::{Arc, OnceLock};

use anyhow::{Context, Result, bail};
use russh::ChannelMsg;
use russh::client::{self, AuthResult, Handle};
use russh::keys::agent::client::{AgentClient, AgentStream};
use russh::keys::known_hosts;
use russh::keys::{PrivateKeyWithHashAlg, PublicKey, load_secret_key};
use ssh2_config::{HostParams, ParseRule, SshConfig};

pub struct Target {
    pub host_alias: String,
    pub user: Option<String>,
    pub port: Option<u16>,
    pub key_path: Option<PathBuf>,
}

#[derive(Debug, Default)]
struct HostConfig {
    host_name: Option<String>,
    user: Option<String>,
    port: Option<u16>,
    identity_file: Option<Vec<PathBuf>>,
}

impl From<HostParams> for HostConfig {
    fn from(params: HostParams) -> Self {
        Self {
            host_name: params.host_name,
            user: params.user,
            port: params.port,
            identity_file: params.identity_file,
        }
    }
}

#[derive(Debug)]
pub struct Resolved {
    pub host: String,
    pub port: u16,
    pub user: String,
    pub key_path: PathBuf,
}

fn resolve_with(
    target: &Target,
    ssh_cfg: &HostConfig,
    fallback_user: Option<&str>,
) -> Result<Resolved> {
    let host = ssh_cfg
        .host_name
        .clone()
        .unwrap_or_else(|| target.host_alias.clone());
    let port = target.port.or(ssh_cfg.port).unwrap_or(22);
    let user = target
        .user
        .clone()
        .or_else(|| ssh_cfg.user.clone())
        .or_else(|| fallback_user.map(str::to_string))
        .context("cannot determine SSH user: set 'user' in deploy.toml or ~/.ssh/config")?;
    let key_path = target
        .key_path
        .clone()
        .or_else(|| {
            ssh_cfg
                .identity_file
                .as_ref()
                .and_then(|files| files.first().cloned())
        })
        .unwrap_or_else(default_key);
    Ok(Resolved {
        host,
        port,
        user,
        key_path,
    })
}

pub fn resolve(target: &Target) -> Result<Resolved> {
    resolve_with(
        target,
        &query_ssh_config(&target.host_alias),
        env_user().as_deref(),
    )
}

fn default_key() -> PathBuf {
    home::home_dir()
        .unwrap_or_else(|| PathBuf::from("."))
        .join(".ssh")
        .join("id_rsa")
}

fn env_user() -> Option<String> {
    std::env::var("USERNAME")
        .or_else(|_| std::env::var("USER"))
        .ok()
}

fn ssh_config_path() -> Option<PathBuf> {
    let path = home::home_dir()?.join(".ssh").join("config");
    path.is_file().then_some(path)
}

fn query_ssh_config(alias: &str) -> HostConfig {
    let Some(path) = ssh_config_path() else {
        return HostConfig::default();
    };
    let file = match File::open(&path) {
        Ok(file) => file,
        Err(_) => return HostConfig::default(),
    };
    let mut reader = std::io::BufReader::new(file);
    match SshConfig::default().parse(&mut reader, ParseRule::ALLOW_UNKNOWN_FIELDS) {
        Ok(config) => config.query(alias).into(),
        Err(err) => {
            eprintln!("warning: cannot parse {}: {err}", path.display());
            HostConfig::default()
        }
    }
}

pub(crate) struct ClientHandler {
    host: String,
    port: u16,
}

impl client::Handler for ClientHandler {
    type Error = russh::Error;

    async fn check_server_key(&mut self, server_key: &PublicKey) -> Result<bool, Self::Error> {
        let key = server_key;
        match known_hosts::check_known_hosts(&self.host, self.port, key) {
            Ok(true) => Ok(true),
            Ok(false) => {
                eprintln!(
                    "note: '{}' is not in known_hosts yet - trusting and remembering it (TOFU)",
                    format_host_port(&self.host, self.port)
                );
                if let Err(err) = known_hosts::learn_known_hosts(&self.host, self.port, key) {
                    eprintln!("warning: could not update known_hosts: {err}");
                }
                Ok(true)
            }
            Err(russh::keys::Error::KeyChanged { line }) => {
                eprintln!(
                    "SECURITY WARNING: host key for '{}' changed (known_hosts line {})! \
                     If this is expected (VPS reinstall), remove that line from \
                     ~/.ssh/known_hosts and retry.",
                    format_host_port(&self.host, self.port),
                    line + 1
                );
                Ok(false)
            }
            Err(other) => {
                eprintln!("warning: cannot verify known_hosts entry: {other}");
                Ok(false)
            }
        }
    }
}

fn format_host_port(host: &str, port: u16) -> String {
    if port == 22 {
        host.to_string()
    } else {
        format!("[{host}]:{port}")
    }
}

pub struct Session {
    handle: Handle<ClientHandler>,
}

impl Session {
    pub(crate) fn handle(&mut self) -> &mut Handle<ClientHandler> {
        &mut self.handle
    }
}

pub(crate) fn runtime() -> &'static tokio::runtime::Runtime {
    static RT: OnceLock<tokio::runtime::Runtime> = OnceLock::new();
    RT.get_or_init(|| {
        tokio::runtime::Builder::new_current_thread()
            .enable_all()
            .build()
            .expect("failed to start async runtime for SSH")
    })
}

pub fn connect(resolved: &Resolved) -> Result<Session> {
    println!("connecting...");
    runtime().block_on(async {
        let handler = ClientHandler {
            host: resolved.host.clone(),
            port: resolved.port,
        };
        let mut handle = client::connect(
            Arc::new(client::Config::default()),
            (resolved.host.as_str(), resolved.port),
            handler,
        )
        .await
        .with_context(|| {
            format!(
                "SSH connection/handshake with {} failed",
                format_host_port(&resolved.host, resolved.port)
            )
        })?;
        authenticate(&mut handle, resolved).await?;
        Ok(Session { handle })
    })
}

async fn authenticate(handle: &mut Handle<ClientHandler>, r: &Resolved) -> Result<()> {
    if try_agent_auth(handle, &r.user).await? {
        return Ok(());
    }

    let secret = load_secret_key(&r.key_path, None).map_err(|err| {
        anyhow::anyhow!(
            "cannot load private key {}: {err} (passphrase-protected keys are not supported yet)",
            r.key_path.display()
        )
    })?;
    let best_hash = handle
        .best_supported_rsa_hash()
        .await
        .ok()
        .flatten()
        .flatten();
    let key_alg = PrivateKeyWithHashAlg::new(Arc::new(secret), best_hash);
    match handle
        .authenticate_publickey(&r.user, key_alg)
        .await
        .context("publickey authentication exchange failed")?
    {
        AuthResult::Success => Ok(()),
        AuthResult::Failure { .. } => bail!(
            "SSH auth rejected for '{}' using key {} (is the public part in authorized_keys on the VPS?)",
            r.user,
            r.key_path.display()
        ),
    }
}

async fn try_agent_auth(handle: &mut Handle<ClientHandler>, user: &str) -> Result<bool> {
    let Some(mut agent) = connect_agent().await else {
        return Ok(false);
    };
    let identities = match agent.request_identities().await {
        Ok(identities) => identities,
        Err(_) => return Ok(false),
    };
    for identity in identities {
        let hash = handle
            .best_supported_rsa_hash()
            .await
            .ok()
            .flatten()
            .flatten();
        let public_key = identity.public_key().into_owned();
        match handle
            .authenticate_publickey_with(user, public_key, hash, &mut agent)
            .await
        {
            Ok(AuthResult::Success) => return Ok(true),
            _ => continue,
        }
    }
    Ok(false)
}

type DynAgentClient = AgentClient<Box<dyn AgentStream + Send + Unpin>>;

async fn connect_agent() -> Option<DynAgentClient> {
    #[cfg(unix)]
    if let Ok(sock) = std::env::var("SSH_AUTH_SOCK") {
        if let Ok(client) = AgentClient::connect_uds(sock).await {
            return Some(client.dynamic());
        }
    }

    #[cfg(windows)]
    {
        const OPENSSH_PIPE: &str = "\\\\.\\pipe\\openssh-ssh-agent";
        if let Ok(client) = AgentClient::connect_named_pipe(OPENSSH_PIPE).await {
            return Some(client.dynamic());
        }
        if let Ok(client) = AgentClient::connect_pageant().await {
            return Some(client.dynamic());
        }
    }

    None
}

impl Session {
    pub fn exec(&mut self, cmd: &str) -> Result<()> {
        runtime().block_on(async {
            use std::io::Write;

            let mut channel = self
                .handle
                .channel_open_session()
                .await
                .context("failed to open SSH channel")?;
            channel
                .exec(true, cmd)
                .await
                .with_context(|| format!("failed to start '{cmd}' on the remote host"))?;

            let mut status: Option<u32> = None;
            let mut stderr = String::new();
            while let Some(msg) = channel.wait().await {
                match msg {
                    ChannelMsg::Data { data } => {
                        print!("{}", String::from_utf8_lossy(&data));
                        let _ = std::io::stdout().flush();
                    }
                    ChannelMsg::ExtendedData { data, ext: 1 } => {
                        stderr.push_str(&String::from_utf8_lossy(&data));
                    }
                    ChannelMsg::ExitStatus { exit_status } => status = Some(exit_status),
                    _ => {}
                }
            }

            let status =
                status.context("remote command finished without reporting an exit status")?;
            if status != 0 {
                let tail = if stderr.trim().is_empty() {
                    String::new()
                } else {
                    format!("\nstderr:\n{stderr}")
                };
                bail!("{cmd} FAILED (exit code {status}){tail}");
            }
            Ok(())
        })
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    fn target(user: Option<&str>, port: Option<u16>) -> Target {
        Target {
            host_alias: "my-vps".into(),
            user: user.map(str::to_string),
            port,
            key_path: None,
        }
    }

    #[test]
    fn explicit_values_override_ssh_config() {
        let hc = HostConfig {
            host_name: Some("10.0.0.5".into()),
            user: Some("cfg-user".into()),
            port: Some(2222),
            identity_file: Some(vec![PathBuf::from("/cfg/key")]),
        };
        let t = Target {
            host_alias: "my-vps".into(),
            user: Some("toml-user".into()),
            port: Some(2200),
            key_path: Some(PathBuf::from("~/custom_key")),
        };
        let r = resolve_with(&t, &hc, Some("os-user")).unwrap();
        assert_eq!(r.user, "toml-user");
        assert_eq!(r.port, 2200);
        assert!(r.key_path.starts_with("~"));
    }

    #[test]
    fn ssh_config_fills_gaps() {
        let hc = HostConfig {
            host_name: Some("10.0.0.5".into()),
            user: Some("cfg-user".into()),
            port: Some(2222),
            identity_file: Some(vec![PathBuf::from("/cfg/key")]),
        };
        let r = resolve_with(&target(None, None), &hc, Some("os-user")).unwrap();
        assert_eq!(r.host, "10.0.0.5");
        assert_eq!(r.user, "cfg-user");
        assert_eq!(r.port, 2222);
        assert_eq!(r.key_path, PathBuf::from("/cfg/key"));
    }

    #[test]
    fn defaults_when_nothing_configured() {
        let r = resolve_with(&target(None, None), &HostConfig::default(), Some("os-user")).unwrap();
        assert_eq!(r.host, "my-vps");
        assert_eq!(r.user, "os-user");
        assert_eq!(r.port, 22);
        assert!(r.key_path.ends_with(".ssh/id_rsa"));
    }

    #[test]
    fn missing_user_everywhere_is_clear_error() {
        let err = resolve_with(&target(None, None), &HostConfig::default(), None).unwrap_err();
        assert!(err.to_string().contains("set 'user'"), "{err:#}");
    }
}