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.

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!("(no nodes matched)");
        return Ok(());
    }

    println!("Discovered host keys:");
    for (name, fp) in &updates {
        println!("  {} -> {}", name, fp);
    }

    let content = fs::read_to_string(config_path)
        .with_context(|| format!("Failed to read {}", config_path))?;
    let new_content = update_host_keys_text(&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)>> {
    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(), fp)),
                Err(e) => eprintln!("warning: failed to fetch key for {}: {}", node.name, e),
            }
        }
    }
    Ok(out)
}

/// Walk `content` line by line and rewrite any node definition whose
/// `name = "..."` matches an entry in `updates`, setting or updating the
/// `host_key` field.
fn update_host_keys_text(content: &str, updates: &[(String, String)]) -> String {
    let mut lines: Vec<String> = content.lines().map(String::from).collect();
    for (name, fp) in updates {
        let marker = format!("name = \"{}\"", name);
        for line in &mut lines {
            if is_node_line(line) && line.contains(&marker) {
                *line = update_node_line(line, fp);
            }
        }
    }
    if content.ends_with('\n') {
        lines.join("\n") + "\n"
    } else {
        lines.join("\n")
    }
}

/// Heuristic: a node definition starts (after indent) with `{`.
fn is_node_line(line: &str) -> bool {
    line.trim_start().starts_with('{')
}

/// Set or replace the `host_key = "..."` field in a single node line.
///
/// If the field is already present, its value is swapped for `fp`.
/// Otherwise the field is inserted before the trailing `}` with a comma
/// if needed.
fn update_node_line(line: &str, fp: &str) -> String {
    let marker = "host_key = \"";
    if let Some(start) = line.find(marker) {
        let value_start = start + marker.len();
        if let Some(rel_end) = line[value_start..].find('"') {
            let value_end = value_start + rel_end;
            let mut result = String::with_capacity(line.len());
            result.push_str(&line[..value_start]);
            result.push_str(fp);
            result.push_str(&line[value_end..]);
            return result;
        }
    }

    // No host_key yet — insert before the closing `}`.
    let close_idx = line.rfind('}').expect("node line must contain '}'");
    let before_close = &line[..close_idx];
    let after_close = &line[close_idx..];
    let trimmed = before_close.trim_end();
    let needs_comma = !trimmed.is_empty() && !trimmed.ends_with(',');
    // Original line either ends with `}` or `},`. Detect which so we can
    // keep the formatting consistent (space before `}` only when a comma
    // follows it).
    let has_trailing_comma = after_close.starts_with("},");
    let mut result = String::with_capacity(line.len() + fp.len() + 20);
    result.push_str(trimmed);
    if needs_comma {
        result.push(',');
    }
    result.push(' ');
    result.push_str(&format!("host_key = \"{}\"", fp));
    if has_trailing_comma {
        result.push(' ');
    }
    result.push('}');
    if has_trailing_comma {
        result.push(',');
    }
    result
}

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

    #[test]
    fn adds_host_key_when_missing() {
        let line =
            r#"    { name = "dgx1.lan", ssh_ip = "10.0.0.1", ib_ip = "1.2.3.4", rank = 0 },"#;
        let result = update_node_line(line, "SHA256:abc");
        assert_eq!(
            result,
            r#"    { name = "dgx1.lan", ssh_ip = "10.0.0.1", ib_ip = "1.2.3.4", rank = 0, host_key = "SHA256:abc" },"#
        );
    }

    #[test]
    fn replaces_host_key_value() {
        let line = r#"    { name = "x", rank = 0, host_key = "OLD" },"#;
        let result = update_node_line(line, "NEW");
        assert_eq!(result, r#"    { name = "x", rank = 0, host_key = "NEW" },"#);
    }

    #[test]
    fn update_host_keys_text_targets_only_named_lines() {
        let content = "\
[[clusters]]
name = \"c\"
nodes = [
    { name = \"a\", rank = 0 },
    { name = \"b\", rank = 1 },
]
";
        let updates = vec![("a".to_string(), "SHA256:newA".to_string())];
        let result = update_host_keys_text(content, &updates);
        // a was updated to include host_key
        assert!(result.contains("name = \"a\", rank = 0, host_key = \"SHA256:newA\""));
        // b was NOT updated — original line preserved, no host_key added
        assert!(result.contains("name = \"b\", rank = 1 }"));
        assert!(!result.contains("name = \"b\", rank = 1, host_key"));
    }

    #[test]
    fn update_node_line_with_no_trailing_comma() {
        // Last field already ends with comma — no extra comma needed.
        let line = r#"    { name = "x", rank = 0 }"#;
        let result = update_node_line(line, "FP");
        assert_eq!(result, r#"    { name = "x", rank = 0, host_key = "FP"}"#);
    }
}