wraith/src-tauri/src/ssh/session.rs

//! SSH session manager — connects, authenticates, manages PTY channels.
//!
//! Each SSH session runs asynchronously via tokio. Terminal stdout is read in a
//! loop and emitted to the frontend via Tauri events (`ssh:data:{session_id}`,
//! base64 encoded). Terminal stdin receives data from the frontend via Tauri
//! commands.
//!
//! Sessions are stored in a `DashMap<String, Arc<SshSession>>`.

use std::sync::Arc;

use async_trait::async_trait;
use base64::Engine;
use dashmap::DashMap;
use log::{debug, error, info, warn};
use russh::client::{self, Handle, Msg};
use russh::{Channel, ChannelMsg, Disconnect};
use serde::Serialize;
use tauri::{AppHandle, Emitter};
use tokio::sync::Mutex as TokioMutex;

use crate::db::Database;
use crate::ssh::cwd::CwdTracker;
use crate::ssh::host_key::{HostKeyResult, HostKeyStore};

// ── auth method ──────────────────────────────────────────────────────────────

/// Authentication method for SSH connections.
pub enum AuthMethod {
    Password(String),
    Key {
        private_key_pem: String,
        passphrase: Option<String>,
    },
}

// ── session info (serializable for frontend) ─────────────────────────────────

#[derive(Debug, Serialize, Clone)]
#[serde(rename_all = "camelCase")]
pub struct SessionInfo {
    pub id: String,
    pub hostname: String,
    pub port: u16,
    pub username: String,
}

// ── SSH session ──────────────────────────────────────────────────────────────

/// Represents a single active SSH session with a PTY channel.
pub struct SshSession {
    pub id: String,
    pub hostname: String,
    pub port: u16,
    pub username: String,
    /// The PTY channel used for interactive shell I/O.
    pub channel: Arc<TokioMutex<Channel<Msg>>>,
    /// Handle to the underlying SSH connection (used for opening new channels).
    pub handle: Arc<TokioMutex<Handle<SshClient>>>,
    /// CWD tracker that polls via a separate exec channel.
    pub cwd_tracker: Option<CwdTracker>,
}

// ── SSH client handler ───────────────────────────────────────────────────────

/// Minimal `russh::client::Handler` implementation.
///
/// Host key verification is done via TOFU in the `HostKeyStore`. The handler
/// stores the verification result so the connect flow can check it after
/// `client::connect` returns.
pub struct SshClient {
    host_key_store: HostKeyStore,
    hostname: String,
    port: u16,
}

#[async_trait]
impl client::Handler for SshClient {
    type Error = russh::Error;

    async fn check_server_key(
        &mut self,
        server_public_key: &ssh_key::PublicKey,
    ) -> Result<bool, Self::Error> {
        let key_type = server_public_key.algorithm().to_string();
        let fingerprint = server_public_key
            .fingerprint(ssh_key::HashAlg::Sha256)
            .to_string();

        let raw_key = server_public_key
            .to_openssh()
            .unwrap_or_default();

        match self
            .host_key_store
            .verify(&self.hostname, self.port, &key_type, &fingerprint)
        {
            Ok(HostKeyResult::New) => {
                info!(
                    "New host key for {}:{} ({}): {}",
                    self.hostname, self.port, key_type, fingerprint
                );
                // TOFU: store the key on first contact.
                if let Err(e) = self.host_key_store.store(
                    &self.hostname,
                    self.port,
                    &key_type,
                    &fingerprint,
                    &raw_key,
                ) {
                    warn!("Failed to store host key: {}", e);
                }
                Ok(true)
            }
            Ok(HostKeyResult::Match) => {
                debug!(
                    "Host key match for {}:{} ({})",
                    self.hostname, self.port, key_type
                );
                Ok(true)
            }
            Ok(HostKeyResult::Changed) => {
                error!(
                    "HOST KEY CHANGED for {}:{} ({})! Expected stored fingerprint, got {}. \
                     Possible man-in-the-middle attack.",
                    self.hostname, self.port, key_type, fingerprint
                );
                // Reject the connection — the frontend should prompt the user
                // to accept the new key and call delete + reconnect.
                Ok(false)
            }
            Err(e) => {
                error!("Host key verification error: {}", e);
                // On DB error, reject to be safe.
                Ok(false)
            }
        }
    }
}

// ── SSH service ──────────────────────────────────────────────────────────────

/// Manages all active SSH sessions.
pub struct SshService {
    sessions: DashMap<String, Arc<SshSession>>,
    db: Database,
}

impl SshService {
    pub fn new(db: Database) -> Self {
        Self {
            sessions: DashMap::new(),
            db,
        }
    }

    /// Establish an SSH connection, authenticate, open a PTY, start a shell,
    /// and begin streaming output to the frontend.
    ///
    /// Returns the session UUID on success.
    pub async fn connect(
        &self,
        app_handle: AppHandle,
        hostname: &str,
        port: u16,
        username: &str,
        auth: AuthMethod,
        cols: u32,
        rows: u32,
    ) -> Result<String, String> {
        let session_id = uuid::Uuid::new_v4().to_string();

        // Build russh client config.
        let config = russh::client::Config::default();
        let config = Arc::new(config);

        // Build our handler with TOFU host key verification.
        let handler = SshClient {
            host_key_store: HostKeyStore::new(self.db.clone()),
            hostname: hostname.to_string(),
            port,
        };

        // Connect to the SSH server.
        let mut handle = client::connect(config, (hostname, port), handler)
            .await
            .map_err(|e| format!("SSH connection to {}:{} failed: {}", hostname, port, e))?;

        // Authenticate.
        let auth_success = match auth {
            AuthMethod::Password(password) => {
                handle
                    .authenticate_password(username, &password)
                    .await
                    .map_err(|e| format!("Password authentication failed: {}", e))?
            }
            AuthMethod::Key {
                private_key_pem,
                passphrase,
            } => {
                let key = russh::keys::decode_secret_key(
                    &private_key_pem,
                    passphrase.as_deref(),
                )
                .map_err(|e| format!("Failed to decode private key: {}", e))?;

                handle
                    .authenticate_publickey(username, Arc::new(key))
                    .await
                    .map_err(|e| format!("Public key authentication failed: {}", e))?
            }
        };

        if !auth_success {
            return Err("Authentication failed: server rejected credentials".to_string());
        }

        // Open a session channel.
        let channel = handle
            .channel_open_session()
            .await
            .map_err(|e| format!("Failed to open session channel: {}", e))?;

        // Request a PTY.
        channel
            .request_pty(
                true,
                "xterm-256color",
                cols,
                rows,
                0, // pix_width
                0, // pix_height
                &[],
            )
            .await
            .map_err(|e| format!("Failed to request PTY: {}", e))?;

        // Start a shell.
        channel
            .request_shell(true)
            .await
            .map_err(|e| format!("Failed to start shell: {}", e))?;

        let handle = Arc::new(TokioMutex::new(handle));
        let channel = Arc::new(TokioMutex::new(channel));

        // Start CWD tracker.
        let cwd_tracker = CwdTracker::new();
        cwd_tracker.start(
            handle.clone(),
            app_handle.clone(),
            session_id.clone(),
        );

        // Build session object.
        let session = Arc::new(SshSession {
            id: session_id.clone(),
            hostname: hostname.to_string(),
            port,
            username: username.to_string(),
            channel: channel.clone(),
            handle: handle.clone(),
            cwd_tracker: Some(cwd_tracker),
        });

        self.sessions.insert(session_id.clone(), session);

        // Spawn the stdout read loop.
        let sid = session_id.clone();
        let chan = channel.clone();
        let app = app_handle.clone();

        tokio::spawn(async move {
            loop {
                let msg = {
                    let mut ch = chan.lock().await;
                    ch.wait().await
                };

                match msg {
                    Some(ChannelMsg::Data { ref data }) => {
                        let encoded = base64::engine::general_purpose::STANDARD
                            .encode(data.as_ref());
                        let event_name = format!("ssh:data:{}", sid);
                        if let Err(e) = app.emit(&event_name, encoded) {
                            error!("Failed to emit SSH data event: {}", e);
                            break;
                        }
                    }
                    Some(ChannelMsg::ExtendedData { ref data, .. }) => {
                        // stderr — emit on the same event channel so the
                        // terminal renders it inline (same as a real terminal).
                        let encoded = base64::engine::general_purpose::STANDARD
                            .encode(data.as_ref());
                        let event_name = format!("ssh:data:{}", sid);
                        if let Err(e) = app.emit(&event_name, encoded) {
                            error!("Failed to emit SSH stderr event: {}", e);
                            break;
                        }
                    }
                    Some(ChannelMsg::ExitStatus { exit_status }) => {
                        info!("SSH session {} exited with status {}", sid, exit_status);
                        let event_name = format!("ssh:exit:{}", sid);
                        let _ = app.emit(&event_name, exit_status);
                        break;
                    }
                    Some(ChannelMsg::Eof) => {
                        debug!("SSH session {} received EOF", sid);
                    }
                    Some(ChannelMsg::Close) => {
                        info!("SSH session {} channel closed", sid);
                        let event_name = format!("ssh:close:{}", sid);
                        let _ = app.emit(&event_name, ());
                        break;
                    }
                    None => {
                        info!("SSH session {} channel stream ended", sid);
                        let event_name = format!("ssh:close:{}", sid);
                        let _ = app.emit(&event_name, ());
                        break;
                    }
                    _ => {
                        // Ignore other channel messages (WindowAdjust, etc.)
                    }
                }
            }
        });

        info!(
            "SSH session {} connected to {}@{}:{}",
            session_id, username, hostname, port
        );

        Ok(session_id)
    }

    /// Write data to a session's PTY stdin.
    pub async fn write(&self, session_id: &str, data: &[u8]) -> Result<(), String> {
        let session = self
            .sessions
            .get(session_id)
            .ok_or_else(|| format!("Session {} not found", session_id))?;

        let channel: tokio::sync::MutexGuard<'_, Channel<Msg>> =
            session.channel.lock().await;
        channel
            .data(&data[..])
            .await
            .map_err(|e| format!("Failed to write to session {}: {}", session_id, e))
    }

    /// Resize the PTY window for a session.
    pub async fn resize(
        &self,
        session_id: &str,
        cols: u32,
        rows: u32,
    ) -> Result<(), String> {
        let session = self
            .sessions
            .get(session_id)
            .ok_or_else(|| format!("Session {} not found", session_id))?;

        let channel: tokio::sync::MutexGuard<'_, Channel<Msg>> =
            session.channel.lock().await;
        channel
            .window_change(cols, rows, 0, 0)
            .await
            .map_err(|e| format!("Failed to resize session {}: {}", session_id, e))
    }

    /// Disconnect a session — close the channel and remove it from the map.
    pub async fn disconnect(&self, session_id: &str) -> Result<(), String> {
        let (_, session) = self
            .sessions
            .remove(session_id)
            .ok_or_else(|| format!("Session {} not found", session_id))?;

        // Close the channel gracefully.
        {
            let channel: tokio::sync::MutexGuard<'_, Channel<Msg>> =
                session.channel.lock().await;
            let _ = channel.eof().await;
            let _ = channel.close().await;
        }

        // Disconnect the SSH connection.
        {
            let handle = session.handle.lock().await;
            let _ = handle
                .disconnect(Disconnect::ByApplication, "", "en")
                .await;
        }

        info!("SSH session {} disconnected", session_id);
        Ok(())
    }

    /// Get a reference to a session by ID.
    pub fn get_session(&self, session_id: &str) -> Option<Arc<SshSession>> {
        self.sessions.get(session_id).map(|entry| entry.clone())
    }

    /// List all active sessions (metadata only).
    pub fn list_sessions(&self) -> Vec<SessionInfo> {
        self.sessions
            .iter()
            .map(|entry| {
                let s = entry.value();
                SessionInfo {
                    id: s.id.clone(),
                    hostname: s.hostname.clone(),
                    port: s.port,
                    username: s.username.clone(),
                }
            })
            .collect()
    }
}