launcher/src/host_keys.rs

//! Auto-discover SSH host key fingerprints for every node in the config
//! and rewrite `config.toml` with `host_key = "SHA256:..."` entries.
//!
//! Uses `toml_edit` for round-trip-safe TOML manipulation so formatting
//! and comments in the config file are preserved.

use anyhow::{Context, Result};
use std::fs;

use crate::config::Config;
use crate::ssh::fetch_host_key;

/// Discover host keys and update the config file.  When `cluster_filter`
/// is `Some`, only that cluster is processed.
pub async fn init_host_keys(
    config: &Config,
    cluster_filter: Option<&str>,
    config_path: &str,
) -> Result<()> {
    let updates = collect_updates(config, cluster_filter).await?;

    if updates.is_empty() {
        println!("{}", fluent_i18n::t!("host-key-no-match"));
        return Ok(());
    }

    println!("{}", fluent_i18n::t!("host-key-header"));
    for (name, ip, fp) in &updates {
        println!(
            "  {}",
            fluent_i18n::t!("host-key-line", { "name" => name, "ip" => ip, "key" => fp })
        );
    }

    let content = fs::read_to_string(config_path)
        .with_context(|| format!("Failed to read {}", config_path))?;
    let new_content = update_host_keys_toml(&content, &updates)?;

    if content == new_content {
        println!("\nNo changes needed in {}", config_path);
    } else {
        fs::write(config_path, &new_content)
            .with_context(|| format!("Failed to write {}", config_path))?;
        println!("\nUpdated {}", config_path);
    }
    Ok(())
}

/// Walk the config and return `(node_name, fingerprint)` for every node.
async fn collect_updates(
    config: &Config,
    cluster_filter: Option<&str>,
) -> Result<Vec<(String, String, String)>> {
    let mut out = Vec::new();
    for cluster in &config.clusters {
        if let Some(name) = cluster_filter
            && cluster.name != name
        {
            continue;
        }
        for node in &cluster.nodes {
            match fetch_host_key(config, node).await {
                Ok(fp) => out.push((node.name.clone(), node.ssh_ip.clone(), fp)),
                Err(e) => println!(
                    "{}",
                    fluent_i18n::t!("host-key-error", { "name" => &node.name, "ip" => &node.ssh_ip, "error" => &e.to_string() })
                ),
            }
        }
    }
    Ok(out)
}

/// Parse `content` as TOML, walk the `[[clusters]]` tables, find each
/// node whose `name` matches an entry in `updates`, and set the
/// `host_key` field on it.  Preserves comments and formatting.
fn update_host_keys_toml(content: &str, updates: &[(String, String, String)]) -> Result<String> {
    use toml_edit::{DocumentMut, Value};

    let mut doc = content
        .parse::<DocumentMut>()
        .with_context(|| "config is not valid TOML")?;
    let clusters = doc
        .get_mut("clusters")
        .and_then(|v| v.as_array_of_tables_mut());

    let Some(clusters) = clusters else {
        return Ok(content.to_string());
    };

    for cluster in clusters.iter_mut() {
        let Some(nodes) = cluster.get_mut("nodes") else {
            continue;
        };
        let Some(arr) = nodes.as_array_mut() else {
            continue;
        };
        for item in arr.iter_mut() {
            let Some(table) = item.as_inline_table_mut() else {
                continue;
            };
            let name_val = table.get("name").and_then(|v| v.as_str());
            let Some(name) = name_val else {
                continue;
            };
            if let Some((_n, _ip, fp)) = updates.iter().find(|(n, _, _)| n == name) {
                table.insert("host_key", Value::from(fp.as_str()));
            }
        }
    }

    Ok(doc.to_string())
}

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

    #[test]
    fn adds_host_key_when_missing() {
        let content = r#"
[[clusters]]
name = "c"
nodes = [
    { name = "dgx1.lan", ssh_ip = "10.0.0.1", ib_ip = "1.2.3.4", rank = 0 },
]
"#;
        let updates = vec![(
            "dgx1.lan".to_string(),
            "10.0.0.1".to_string(),
            "SHA256:abc".to_string(),
        )];
        let result = update_host_keys_toml(content, &updates).unwrap();
        assert!(result.contains(r#"host_key = "SHA256:abc""#));
        assert!(result.contains("dgx1.lan"));
    }

    #[test]
    fn replaces_host_key_value() {
        let content = r#"
[[clusters]]
name = "c"
nodes = [
    { name = "x", rank = 0, host_key = "OLD" },
]
"#;
        let updates = vec![("x".to_string(), "10.0.0.1".to_string(), "NEW".to_string())];
        let result = update_host_keys_toml(content, &updates).unwrap();
        assert!(result.contains(r#"host_key = "NEW""#));
        assert!(!result.contains(r#"host_key = "OLD""#));
    }

    #[test]
    fn only_updates_targeted_nodes() {
        let content = r#"
[[clusters]]
name = "c"
nodes = [
    { name = "a", rank = 0 },
    { name = "b", rank = 1 },
]
"#;
        let updates = vec![(
            "a".to_string(),
            "10.0.0.1".to_string(),
            "SHA256:newA".to_string(),
        )];
        let result = update_host_keys_toml(content, &updates).unwrap();
        assert!(
            result.contains(r#"host_key = "SHA256:newA""#),
            "result: {result}"
        );
        // b should NOT have host_key added
        assert!(result.contains(r#"name = "b""#), "result has b: {result}");
        assert!(
            !result.contains("host_key") || result.matches("host_key").count() == 1,
            "only 'a' should have host_key: {result}"
        );
    }

    #[test]
    fn updates_multiple_clusters() {
        let content = r#"
[[clusters]]
name = "c1"
nodes = [
    { name = "n1", rank = 0 },
]

[[clusters]]
name = "c2"
nodes = [
    { name = "n2", rank = 0 },
]
"#;
        let updates = vec![
            ("n1".to_string(), "10.0.0.1".to_string(), "FP1".to_string()),
            ("n2".to_string(), "10.0.0.2".to_string(), "FP2".to_string()),
        ];
        let result = update_host_keys_toml(content, &updates).unwrap();
        assert!(
            result.contains(r#"host_key = "FP1""#),
            "result has FP1: {result}"
        );
        assert!(
            result.contains(r#"host_key = "FP2""#),
            "result has FP2: {result}"
        );
    }

    #[test]
    fn preserves_trailing_newline() {
        let content = "[[clusters]]\nname = \"c\"\nnodes = []\n";
        let result = update_host_keys_toml(content, &[]).unwrap();
        assert_eq!(result, content);
    }

    #[test]
    fn no_host_key_for_absent_update() {
        let content = r#"[[clusters]]
name = "c"
nodes = [{ name = "n1", rank = 0 }]
"#;
        let result = update_host_keys_toml(content, &[]).unwrap();
        assert_eq!(result, content);
    }
}