|
@@ -0,0 +1,184 @@
|
|
|
|
|
+//! 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"}"#);
|
|
|
|
|
+ }
|
|
|
|
|
+}
|