launcher/src/ssh.rs

use anyhow::{Context, Result};
use base64::Engine;
use base64::engine::general_purpose::STANDARD;
use russh::keys::{PrivateKeyWithHashAlg, load_secret_key as russh_load_secret_key};
use russh::*;
use sha2::{Digest, Sha256};
use std::sync::Arc;
use tokio::time::{Duration, sleep};
use tracing::warn;

use crate::config::{Config, Node};

/// Establishes an SSH session to a node using the config credentials.
pub async fn ssh_session(config: &Config, node: &Node) -> Result<client::Handle<VerifyingHandler>> {
    let key = load_secret_key(&config.ssh_key_path)?;
    let ssh_config = russh::client::Config::default();
    let expected_key = Arc::new(std::sync::Mutex::new(node.host_key.clone()));

    if expected_key.lock().unwrap().is_none() {
        warn!(
            "No host_key set for {} — skipping verification (insecure)",
            node.name
        );
    }

    let handler = VerifyingHandler { expected_key };
    let mut session =
        russh::client::connect(Arc::new(ssh_config), (node.ssh_ip.as_str(), 22), handler).await?;

    session
        .authenticate_publickey(
            &config.ssh_user,
            PrivateKeyWithHashAlg::new(Arc::new(key), None),
        )
        .await?;

    Ok(session)
}

/// Runs a command on a node via SSH and returns stdout.
///
/// The remote command's exit status is tracked: if it is non-zero,
/// an Err is returned whose `to_string()` includes the captured output
/// for diagnosis.  Callers that don't care about exit status can ignore
/// the Result and inspect stdout via the returned `String`.
pub async fn run_ssh(
    config: &Config,
    node: &Node,
    command: String,
    dry_run: bool,
) -> Result<String> {
    if dry_run {
        println!("  ssh {}@{}: {}", config.ssh_user, node.ssh_ip, command);
        return Ok(String::new());
    }

    let session = ssh_session(config, node).await?;
    let mut channel = session.channel_open_session().await?;
    channel.exec(true, command.as_bytes()).await?;

    let mut output = String::new();
    let mut exit_status: Option<u32> = None;
    while let Some(msg) = channel.wait().await {
        match msg {
            ChannelMsg::Data { data } => {
                output.push_str(&String::from_utf8_lossy(&data));
            }
            ChannelMsg::ExitStatus { exit_status: s } => {
                exit_status = Some(s);
            }
            _ => {}
        }
    }

    match exit_status {
        Some(0) | None => Ok(output),
        Some(code) => Err(anyhow::anyhow!(
            "ssh {}@{}: command exited with status {}\n  command: {}\n  output:\n{}",
            config.ssh_user,
            node.ssh_ip,
            code,
            command,
            output
        )),
    }
}

/// Runs a command via SSH with exponential backoff retry.
pub async fn run_ssh_with_retry(
    config: &Config,
    node: &Node,
    command: String,
    dry_run: bool,
    max_retries: u32,
) -> Result<String> {
    let mut last_err = None;
    for attempt in 0..=max_retries {
        match run_ssh(config, node, command.clone(), dry_run).await {
            Ok(output) => return Ok(output),
            Err(e) => {
                warn!(
                    "SSH attempt {}/{} to {} failed: {}",
                    attempt + 1,
                    max_retries + 1,
                    node.name,
                    e
                );
                last_err = Some(e);
                if attempt < max_retries {
                    sleep(Duration::from_secs(2u64.pow(attempt))).await;
                }
            }
        }
    }
    Err(last_err.unwrap())
}

/// Runs a command via SSH and streams output to stdout (for follow mode).
pub async fn run_ssh_follow(
    config: &Config,
    node: &Node,
    command: String,
    dry_run: bool,
) -> Result<()> {
    if dry_run {
        println!("  ssh {}@{}: {} -f", config.ssh_user, node.ssh_ip, command);
        return Ok(());
    }

    let session = ssh_session(config, node).await?;
    let mut channel = session.channel_open_session().await?;
    channel.exec(true, command.as_bytes()).await?;

    while let Some(msg) = channel.wait().await {
        if let ChannelMsg::Data { data } = msg {
            print!("{}", String::from_utf8_lossy(&data));
        }
    }
    Ok(())
}

/// Fetches the SSH host key fingerprint from a node.
pub async fn fetch_host_key(config: &Config, node: &Node) -> Result<String> {
    use std::sync::Mutex;

    let captured = Arc::new(Mutex::new(None::<Vec<u8>>));
    let captured_clone = captured.clone();

    struct CaptureHandler {
        captured: Arc<Mutex<Option<Vec<u8>>>>,
    }
    impl client::Handler for CaptureHandler {
        type Error = russh::Error;
        async fn check_server_key(
            &mut self,
            key: &russh::keys::PublicKey,
        ) -> Result<bool, Self::Error> {
            *self.captured.lock().unwrap() = Some(key.to_string().into_bytes());
            Ok(true)
        }
    }

    let key = load_secret_key(&config.ssh_key_path)?;
    let ssh_config = russh::client::Config::default();
    let handler = CaptureHandler {
        captured: captured_clone,
    };
    let mut session =
        russh::client::connect(Arc::new(ssh_config), (node.ssh_ip.as_str(), 22), handler).await?;

    session
        .authenticate_publickey(
            &config.ssh_user,
            PrivateKeyWithHashAlg::new(Arc::new(key), None),
        )
        .await?;

    let blob = captured
        .lock()
        .unwrap()
        .take()
        .context("No server key captured")?;
    let hash = Sha256::digest(&blob);
    let fp = format!("SHA256:{}", STANDARD.encode(hash));
    Ok(fp)
}

/// Loads an SSH private key from the given path.
pub fn load_secret_key(path: &str) -> Result<russh::keys::PrivateKey> {
    russh_load_secret_key(path, None).context("Failed to load SSH key")
}

/// Handler that verifies the server's host key against an expected fingerprint.
pub struct VerifyingHandler {
    pub expected_key: Arc<std::sync::Mutex<Option<String>>>,
}

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

    async fn check_server_key(
        &mut self,
        server_public_key: &russh::keys::PublicKey,
    ) -> Result<bool, Self::Error> {
        let expected = self.expected_key.lock().unwrap().clone();
        if let Some(ref expected_fp) = expected {
            let key_bytes = server_public_key.to_string().into_bytes();
            let hash = Sha256::digest(&key_bytes);
            let fp = format!("SHA256:{}", STANDARD.encode(hash));
            if &fp != expected_fp {
                return Err(russh::Error::Kex);
            }
        }
        Ok(true)
    }
}