shuttle/

migrate.rs

use serde_json::{json, Value};
use std::collections::HashMap;
use std::path::{Path, PathBuf};

struct SshConfigEntry {
    host_alias: String,
    hostname: Option<String>,
    user: Option<String>,
    port: Option<u16>,
    identity_file: Option<String>,
    proxy_jump: Option<String>,
}

struct SshKeyFile {
    name: String,
    algorithm: String,
    has_public: bool,
    is_pem: bool,
}

/// Scan ~/.ssh/ to produce a migration plan for shuttle.
///
/// SECURITY: This tool intentionally reads ~/.ssh/config, directory listings,
/// and known_hosts line counts. This is an explicit exception to the
/// transfer.rs blocklist which blocks ~/.ssh for scp/sftp operations.
/// The exception is justified because:
/// - No private key file content is ever read (only filenames)
/// - The output is a migration plan, not executable config
/// - Identity file paths are exposed as basenames only, not full paths
/// - known_hosts content is counted, not returned
pub async fn ssh_migrate_analyze(
    _params: &Value,
) -> omegon_extension::Result<Value> {
    let ssh_dir = dirs::home_dir()
        .map(|h| h.join(".ssh"))
        .unwrap_or_else(|| PathBuf::from("/nonexistent"));

    let mut report = json!({
        "ssh_dir_exists": ssh_dir.is_dir(),
    });

    if !ssh_dir.is_dir() {
        report["summary"] = json!("No ~/.ssh directory found. Nothing to migrate.");
        return Ok(report);
    }

    let config_entries = parse_ssh_config(&ssh_dir.join("config"));
    let key_files = scan_key_files(&ssh_dir);
    let known_hosts_count = count_known_hosts(&ssh_dir.join("known_hosts"));

    let keys: Vec<Value> = key_files
        .iter()
        .map(|k| {
            json!({
                "filename": k.name,
                "algorithm": k.algorithm,
                "has_public_key": k.has_public,
                "is_pem": k.is_pem,
            })
        })
        .collect();

    // Build host inventory — redact identity_file paths (only expose the
    // basename, not the full path, to avoid leaking directory structure)
    let hosts: Vec<Value> = config_entries
        .iter()
        .map(|e| {
            let mut entry = json!({
                "alias": sanitize_for_display(&e.host_alias),
            });
            if let Some(ref h) = e.hostname {
                entry["hostname"] = json!(h);
            }
            if let Some(ref u) = e.user {
                entry["user"] = json!(u);
            }
            if let Some(p) = e.port {
                entry["port"] = json!(p);
            }
            if let Some(ref f) = e.identity_file {
                let basename = Path::new(f)
                    .file_name()
                    .and_then(|n| n.to_str())
                    .unwrap_or("(key)");
                entry["identity_file_name"] = json!(basename);
            }
            if e.proxy_jump.is_some() {
                entry["has_proxy_jump"] = json!(true);
                entry["note"] = json!(
                    "ProxyJump is not directly supported by shuttle. \
                     Use ssh_tunnel_open through the bastion instead."
                );
            }
            entry
        })
        .collect();

    // Group hosts by identity file basename to suggest label assignments
    let mut label_groups: HashMap<String, Vec<String>> = HashMap::new();
    for entry in &config_entries {
        let key = entry
            .identity_file
            .as_deref()
            .and_then(|f| Path::new(f).file_name())
            .and_then(|n| n.to_str())
            .unwrap_or("(default)")
            .to_string();
        label_groups
            .entry(key)
            .or_default()
            .push(entry.host_alias.clone());
    }

    let suggested_labels: Vec<Value> = label_groups
        .iter()
        .map(|(key_file, hosts)| {
            let label_suggestion = suggest_label(key_file, hosts);
            json!({
                "current_key_name": key_file,
                "hosts": hosts.iter().map(|h| sanitize_for_display(h)).collect::<Vec<_>>(),
                "suggested_label": label_suggestion,
            })
        })
        .collect();

    // Generate draft hosts.toml with proper TOML escaping
    let mut hosts_toml_lines = vec![
        "# Draft hosts.toml generated by ssh_migrate_analyze".to_string(),
        "# Review and adjust identity_label names to match your trust domains.".to_string(),
        String::new(),
    ];

    for entry in &config_entries {
        if entry.host_alias == "*" || entry.hostname.is_none() {
            continue;
        }
        let alias = sanitize_toml_key(&entry.host_alias);
        let hostname = toml_escape(entry.hostname.as_deref().unwrap());
        let user = toml_escape(entry.user.as_deref().unwrap_or("TODO_SET_USER"));
        let port = entry.port.unwrap_or(22);
        let label = entry
            .identity_file
            .as_deref()
            .and_then(|f| Path::new(f).file_name())
            .and_then(|n| n.to_str())
            .map(|f| suggest_label(f, &[entry.host_alias.clone()]))
            .unwrap_or_else(|| "default".to_string());

        hosts_toml_lines.push(format!("[{alias}]"));
        hosts_toml_lines.push(format!("address = \"{hostname}\""));
        hosts_toml_lines.push(format!("user = \"{user}\""));
        if port != 22 {
            hosts_toml_lines.push(format!("port = {port}"));
        }
        hosts_toml_lines.push(format!("identity_label = \"{}\"", toml_escape(&label)));
        hosts_toml_lines.push("trust_on_first_use = true".to_string());
        hosts_toml_lines.push(String::new());
    }

    // Collect unique labels for keygen commands
    let labels: Vec<String> = label_groups
        .iter()
        .map(|(key_file, hosts)| suggest_label(key_file, hosts))
        .collect::<std::collections::HashSet<_>>()
        .into_iter()
        .collect();

    let keygen_commands: Vec<String> = labels
        .iter()
        .map(|label| {
            format!(
                "shuttle-keygen ~/.config/styrene/identity.key {label}"
            )
        })
        .collect();

    report["key_files"] = json!(keys);
    report["ssh_config_hosts"] = json!(hosts);
    report["known_hosts_count"] = json!(known_hosts_count);
    report["suggested_label_groups"] = json!(suggested_labels);
    report["draft_hosts_toml"] = json!(hosts_toml_lines.join("\n"));
    report["keygen_commands"] = json!(keygen_commands);
    report["migration_steps"] = json!([
        "1. Review the draft hosts.toml — adjust identity_label names to match your trust domains",
        "2. Run each keygen command to export the SSH public keys",
        "3. Add each public key to the corresponding remote server's ~/.ssh/authorized_keys",
        "4. Copy the hosts.toml to ~/.omegon/shuttle/hosts.toml",
        "5. Test each host with ssh_ping",
        "6. Once verified, remove old key entries from ~/.ssh/config for migrated hosts",
    ]);

    Ok(report)
}

fn parse_ssh_config(path: &Path) -> Vec<SshConfigEntry> {
    let content = match std::fs::read_to_string(path) {
        Ok(c) => c,
        Err(_) => return Vec::new(),
    };

    let mut entries = Vec::new();
    let mut current: Option<SshConfigEntry> = None;

    for line in content.lines() {
        let line = line.trim();
        if line.is_empty() || line.starts_with('#') {
            continue;
        }

        let (key, value) = match line.split_once(char::is_whitespace) {
            Some((k, v)) => (k.to_lowercase(), v.trim().to_string()),
            None => continue,
        };

        match key.as_str() {
            "host" => {
                if let Some(entry) = current.take() {
                    entries.push(entry);
                }
                current = Some(SshConfigEntry {
                    host_alias: value,
                    hostname: None,
                    user: None,
                    port: None,
                    identity_file: None,
                    proxy_jump: None,
                });
            }
            "hostname" => {
                if let Some(ref mut e) = current {
                    e.hostname = Some(value);
                }
            }
            "user" => {
                if let Some(ref mut e) = current {
                    e.user = Some(value);
                }
            }
            "port" => {
                if let Some(ref mut e) = current {
                    e.port = value.parse().ok();
                }
            }
            "identityfile" => {
                if let Some(ref mut e) = current {
                    e.identity_file = Some(value);
                }
            }
            "proxyjump" | "proxycommand" => {
                if let Some(ref mut e) = current {
                    e.proxy_jump = Some(value);
                }
            }
            _ => {}
        }
    }

    if let Some(entry) = current {
        entries.push(entry);
    }

    entries
}

fn scan_key_files(ssh_dir: &Path) -> Vec<SshKeyFile> {
    let entries = match std::fs::read_dir(ssh_dir) {
        Ok(e) => e,
        Err(_) => return Vec::new(),
    };

    let mut keys = Vec::new();

    for entry in entries.flatten() {
        let name = entry.file_name().to_string_lossy().to_string();

        if name == "config"
            || name == "known_hosts"
            || name == "known_hosts.old"
            || name == "authorized_keys"
            || name == "agent"
            || name == "environment"
            || name.starts_with('.')
            || name.ends_with(".pub")
        {
            continue;
        }

        let path = entry.path();
        if !path.is_file() {
            continue;
        }

        let algorithm = detect_key_algorithm(&name);
        let has_public = ssh_dir.join(format!("{name}.pub")).exists();
        let is_pem = name.ends_with(".pem");

        keys.push(SshKeyFile {
            name,
            algorithm,
            has_public,
            is_pem,
        });
    }

    keys.sort_by(|a, b| a.name.cmp(&b.name));
    keys
}

fn detect_key_algorithm(filename: &str) -> String {
    if filename.contains("ed25519") {
        "ed25519".to_string()
    } else if filename.contains("ecdsa") {
        "ecdsa".to_string()
    } else if filename.contains("rsa") || filename.ends_with(".pem") {
        "rsa (likely)".to_string()
    } else if filename.contains("dsa") {
        "dsa (legacy)".to_string()
    } else {
        "unknown".to_string()
    }
}

/// Count known_hosts entries without parsing or returning hostname content.
fn count_known_hosts(path: &Path) -> usize {
    let content = match std::fs::read_to_string(path) {
        Ok(c) => c,
        Err(_) => return 0,
    };

    content
        .lines()
        .filter(|l| {
            let l = l.trim();
            !l.is_empty() && !l.starts_with('#')
        })
        .count()
}

/// Sanitize a string for TOML table key use — strip characters that could
/// break TOML syntax or inject new sections.
fn sanitize_toml_key(s: &str) -> String {
    let sanitized: String = s
        .chars()
        .filter(|c| c.is_alphanumeric() || *c == '-' || *c == '_' || *c == '.')
        .collect();
    if sanitized.is_empty() {
        "unnamed-host".to_string()
    } else {
        sanitized
    }
}

/// Escape a string for use inside TOML double-quoted values.
fn toml_escape(s: &str) -> String {
    s.replace('\\', "\\\\")
        .replace('"', "\\\"")
        .replace('\n', "\\n")
        .replace('\r', "\\r")
        .replace('\t', "\\t")
}

/// Sanitize a string for display in JSON output — prevent control characters.
fn sanitize_for_display(s: &str) -> String {
    s.chars()
        .filter(|c| !c.is_control())
        .collect()
}

fn suggest_label(key_file: &str, hosts: &[String]) -> String {
    let basename = Path::new(key_file)
        .file_stem()
        .and_then(|s| s.to_str())
        .unwrap_or("default");

    if basename.contains("prod") {
        return "prod".to_string();
    }
    if basename.contains("staging") || basename.contains("stag") {
        return "staging".to_string();
    }
    if basename.contains("dev") || basename.contains("personal") {
        return "dev".to_string();
    }
    if basename.contains("deploy") || basename.contains("ci") {
        return "deploy".to_string();
    }
    if basename.contains("admin") {
        return "admin".to_string();
    }

    if hosts.len() == 1 {
        let h = &hosts[0];
        if h != "*" {
            return sanitize_toml_key(h).to_lowercase();
        }
    }

    sanitize_toml_key(
        basename
            .trim_start_matches("id_")
    ).to_lowercase()
}