opsxshell/src-tauri/src/services/ssh/known_hosts.rs

//! Host key verification against `~/.ssh/known_hosts` with UI confirmation.
//!
//! Unknown / changed keys emit `ssh:host-key-prompt` and block until the UI
//! responds (or the wait times out). Silent TOFU is intentionally not used.

use std::{
    collections::{HashMap, HashSet},
    fs,
    sync::{Mutex, OnceLock},
    time::Duration,
};

use base64::{engine::general_purpose::STANDARD_NO_PAD, Engine};
use serde::Serialize;
use sha2::{Digest, Sha256};
use ssh2::{CheckResult, HostKeyType, KnownHostFileKind, KnownHostKeyFormat, Session};
use tauri::{AppHandle, Emitter};
use uuid::Uuid;

const PROMPT_TIMEOUT: Duration = Duration::from_secs(120);

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum HostKeyDecision {
    AcceptOnce,
    AcceptAlways,
    Reject,
}

impl HostKeyDecision {
    pub fn parse(value: &str) -> Option<Self> {
        match value {
            "accept-once" => Some(Self::AcceptOnce),
            "accept-always" => Some(Self::AcceptAlways),
            "reject" => Some(Self::Reject),
            _ => None,
        }
    }
}

#[derive(Clone, Serialize)]
#[serde(rename_all = "camelCase")]
struct HostKeyPromptEvent {
    request_id: String,
    host: String,
    port: u16,
    key_type: String,
    fingerprint: String,
    is_changed: bool,
    #[serde(skip_serializing_if = "Option::is_none")]
    previous_fingerprint: Option<String>,
}

struct PendingPrompt {
    tx: std::sync::mpsc::Sender<HostKeyDecision>,
}

static APP_HANDLE: OnceLock<AppHandle> = OnceLock::new();
static PENDING: OnceLock<Mutex<HashMap<String, PendingPrompt>>> = OnceLock::new();
/// Process-lifetime accept-once cache so tunnels/reconnects do not re-prompt.
static ONCE_TRUSTED: OnceLock<Mutex<HashSet<String>>> = OnceLock::new();

fn pending_map() -> &'static Mutex<HashMap<String, PendingPrompt>> {
    PENDING.get_or_init(|| Mutex::new(HashMap::new()))
}

fn once_trusted() -> &'static Mutex<HashSet<String>> {
    ONCE_TRUSTED.get_or_init(|| Mutex::new(HashSet::new()))
}

fn once_key(host: &str, port: u16, fingerprint: &str) -> String {
    format!("{host}:{port}:{fingerprint}")
}

/// Bind the app handle once at startup so SSH workers can emit prompts.
pub fn bind_app_handle(app: AppHandle) {
    let _ = APP_HANDLE.set(app);
}

pub fn respond(request_id: &str, decision: HostKeyDecision) -> Result<(), String> {
    let mut pending = pending_map()
        .lock()
        .map_err(|_| "主机密钥确认队列异常".to_string())?;
    let Some(entry) = pending.remove(request_id) else {
        return Err(format!("主机密钥确认请求不存在或已过期: {request_id}"));
    };
    entry
        .tx
        .send(decision)
        .map_err(|_| "主机密钥确认通道已关闭".to_string())
}

pub(super) fn verify_host_key(session: &Session, host: &str, port: u16) -> Result<(), String> {
    let (key, key_type) = session
        .host_key()
        .ok_or_else(|| "服务器未提供主机密钥".to_string())?;
    let fingerprint = sha256_fingerprint(key);
    let key_type_name = host_key_type_name(key_type);

    if once_trusted()
        .lock()
        .map_err(|_| "主机密钥缓存异常".to_string())?
        .contains(&once_key(host, port, &fingerprint))
    {
        return Ok(());
    }

    let Some(home) = dirs::home_dir() else {
        return Err("无法确定用户主目录，不能校验 known_hosts".to_string());
    };
    let ssh_dir = home.join(".ssh");
    fs::create_dir_all(&ssh_dir).map_err(|error| format!("创建 ~/.ssh 失败: {error}"))?;
    let path = ssh_dir.join("known_hosts");
    let mut known = session
        .known_hosts()
        .map_err(|error| format!("初始化 known_hosts 失败: {error}"))?;
    if path.is_file() {
        known
            .read_file(&path, KnownHostFileKind::OpenSSH)
            .map_err(|error| format!("读取 known_hosts 失败: {error}"))?;
    }

    match known.check_port(host, port, key) {
        CheckResult::Match => Ok(()),
        CheckResult::Mismatch => {
            let previous = lookup_stored_fingerprint(host, port);
            let decision = prompt_user(
                host,
                port,
                &key_type_name,
                &fingerprint,
                true,
                previous.as_deref(),
            )?;
            match decision {
                HostKeyDecision::Reject => {
                    Err(format!("主机密钥已变化，已拒绝连接（{host}:{port}）"))
                }
                HostKeyDecision::AcceptOnce => {
                    remember_once(host, port, &fingerprint);
                    Ok(())
                }
                HostKeyDecision::AcceptAlways => {
                    remove_known_host_entries(host, port)?;
                    // Reload after deletion so add does not keep the stale key.
                    let mut known = session
                        .known_hosts()
                        .map_err(|error| format!("初始化 known_hosts 失败: {error}"))?;
                    if path.is_file() {
                        known
                            .read_file(&path, KnownHostFileKind::OpenSSH)
                            .map_err(|error| format!("读取 known_hosts 失败: {error}"))?;
                    }
                    write_known_host(&mut known, &path, host, port, key, key_type)?;
                    remember_once(host, port, &fingerprint);
                    Ok(())
                }
            }
        }
        CheckResult::NotFound | CheckResult::Failure => {
            let decision = prompt_user(host, port, &key_type_name, &fingerprint, false, None)?;
            match decision {
                HostKeyDecision::Reject => Err(format!("用户拒绝了主机密钥（{host}:{port}）")),
                HostKeyDecision::AcceptOnce => {
                    remember_once(host, port, &fingerprint);
                    Ok(())
                }
                HostKeyDecision::AcceptAlways => {
                    write_known_host(&mut known, &path, host, port, key, key_type)?;
                    remember_once(host, port, &fingerprint);
                    Ok(())
                }
            }
        }
    }
}

fn prompt_user(
    host: &str,
    port: u16,
    key_type: &str,
    fingerprint: &str,
    is_changed: bool,
    previous_fingerprint: Option<&str>,
) -> Result<HostKeyDecision, String> {
    let Some(app) = APP_HANDLE.get() else {
        return Err(
            "无法弹出主机密钥确认（应用未就绪）。请稍后重试，或手动将密钥写入 ~/.ssh/known_hosts"
                .to_string(),
        );
    };

    let request_id = Uuid::new_v4().to_string();
    let (tx, rx) = std::sync::mpsc::channel();
    pending_map()
        .lock()
        .map_err(|_| "主机密钥确认队列异常".to_string())?
        .insert(request_id.clone(), PendingPrompt { tx });

    let payload = HostKeyPromptEvent {
        request_id: request_id.clone(),
        host: host.to_string(),
        port,
        key_type: key_type.to_string(),
        fingerprint: fingerprint.to_string(),
        is_changed,
        previous_fingerprint: previous_fingerprint.map(str::to_string),
    };

    if app.emit("ssh:host-key-prompt", payload).is_err() {
        pending_map()
            .lock()
            .ok()
            .map(|mut map| map.remove(&request_id));
        return Err("发送主机密钥确认事件失败".to_string());
    }

    match rx.recv_timeout(PROMPT_TIMEOUT) {
        Ok(decision) => Ok(decision),
        Err(_) => {
            pending_map()
                .lock()
                .ok()
                .map(|mut map| map.remove(&request_id));
            Err(format!(
                "主机密钥确认超时（{} 秒）",
                PROMPT_TIMEOUT.as_secs()
            ))
        }
    }
}

fn remember_once(host: &str, port: u16, fingerprint: &str) {
    if let Ok(mut set) = once_trusted().lock() {
        set.insert(once_key(host, port, fingerprint));
    }
}

fn write_known_host(
    known: &mut ssh2::KnownHosts,
    path: &std::path::Path,
    host: &str,
    port: u16,
    key: &[u8],
    key_type: HostKeyType,
) -> Result<(), String> {
    let host_label = known_host_label(host, port);
    known
        .add(&host_label, key, &host_label, known_host_format(key_type))
        .map_err(|error| format!("记录主机密钥失败: {error}"))?;
    known
        .write_file(path, KnownHostFileKind::OpenSSH)
        .map_err(|error| format!("写入 known_hosts 失败: {error}"))
}

fn remove_known_host_entries(host: &str, port: u16) -> Result<(), String> {
    let Some(path) = dirs::home_dir().map(|home| home.join(".ssh").join("known_hosts")) else {
        return Ok(());
    };
    if !path.is_file() {
        return Ok(());
    }
    let input = fs::read_to_string(&path).map_err(|error| error.to_string())?;
    let label = known_host_label(host, port);
    let retained = input
        .lines()
        .filter(|line| {
            line.split_whitespace().next().is_none_or(|hosts| {
                !hosts
                    .split(',')
                    .any(|candidate| candidate == label || (port == 22 && candidate == host))
            })
        })
        .collect::<Vec<_>>()
        .join("\n");
    fs::write(
        path,
        if retained.is_empty() {
            retained
        } else {
            format!("{retained}\n")
        },
    )
    .map_err(|error| error.to_string())
}

fn lookup_stored_fingerprint(host: &str, port: u16) -> Option<String> {
    let path = dirs::home_dir()?.join(".ssh").join("known_hosts");
    if !path.is_file() {
        return None;
    }
    let input = fs::read_to_string(path).ok()?;
    let label = known_host_label(host, port);
    for line in input
        .lines()
        .filter(|line| !line.trim_start().starts_with('#'))
    {
        let fields = line.split_whitespace().collect::<Vec<_>>();
        if fields.len() < 3 || fields[0].starts_with('|') {
            continue;
        }
        let matches = fields[0]
            .split(',')
            .any(|candidate| candidate == label || (port == 22 && candidate == host));
        if !matches {
            continue;
        }
        let Ok(key) = STANDARD_NO_PAD.decode(fields[2]) else {
            continue;
        };
        return Some(sha256_fingerprint(&key));
    }
    None
}

fn sha256_fingerprint(key: &[u8]) -> String {
    STANDARD_NO_PAD.encode(Sha256::digest(key))
}

fn host_key_type_name(key_type: HostKeyType) -> String {
    match key_type {
        HostKeyType::Rsa => "ssh-rsa",
        HostKeyType::Dss => "ssh-dss",
        HostKeyType::Ecdsa256 => "ecdsa-sha2-nistp256",
        HostKeyType::Ecdsa384 => "ecdsa-sha2-nistp384",
        HostKeyType::Ecdsa521 => "ecdsa-sha2-nistp521",
        HostKeyType::Ed25519 => "ssh-ed25519",
        HostKeyType::Unknown => "unknown",
    }
    .to_string()
}

fn known_host_label(host: &str, port: u16) -> String {
    if port == 22 {
        host.to_string()
    } else {
        format!("[{host}]:{port}")
    }
}

fn known_host_format(key_type: HostKeyType) -> KnownHostKeyFormat {
    match key_type {
        HostKeyType::Rsa => KnownHostKeyFormat::SshRsa,
        HostKeyType::Dss => KnownHostKeyFormat::SshDss,
        HostKeyType::Ecdsa256 => KnownHostKeyFormat::Ecdsa256,
        HostKeyType::Ecdsa384 => KnownHostKeyFormat::Ecdsa384,
        HostKeyType::Ecdsa521 => KnownHostKeyFormat::Ecdsa521,
        HostKeyType::Ed25519 => KnownHostKeyFormat::Ed25519,
        HostKeyType::Unknown => KnownHostKeyFormat::Unknown,
    }
}

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

    #[test]
    fn formats_non_default_ports_like_openssh() {
        assert_eq!(known_host_label("example.com", 22), "example.com");
        assert_eq!(known_host_label("example.com", 2222), "[example.com]:2222");
    }

    #[test]
    fn parses_host_key_decisions() {
        assert_eq!(
            HostKeyDecision::parse("accept-always"),
            Some(HostKeyDecision::AcceptAlways)
        );
        assert_eq!(HostKeyDecision::parse("nope"), None);
    }

    #[test]
    fn fingerprint_is_unpadded_sha256_base64() {
        let fp = sha256_fingerprint(b"test-key");
        assert!(!fp.contains('='));
        assert!(!fp.is_empty());
    }
}