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

//! SSH session manager — connects, authenticates, manages PTY channels.

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::sftp::SftpService;
use crate::ssh::cwd::CwdTracker;
use crate::ssh::host_key::{HostKeyResult, HostKeyStore};

pub enum AuthMethod {
    Password(String),
    Key { private_key_pem: String, passphrase: Option<String> },
}

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

pub struct SshSession {
    pub id: String,
    pub hostname: String,
    pub port: u16,
    pub username: String,
    pub channel: Arc<TokioMutex<Channel<Msg>>>,
    pub handle: Arc<TokioMutex<Handle<SshClient>>>,
    pub cwd_tracker: Option<CwdTracker>,
}

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) => {
                let _ = self.host_key_store.store(&self.hostname, self.port, &key_type, &fingerprint, &raw_key);
                Ok(true)
            }
            Ok(HostKeyResult::Match) => Ok(true),
            Ok(HostKeyResult::Changed) => Ok(false),
            Err(_) => Ok(false),
        }
    }
}

pub struct SshService {
    sessions: DashMap<String, Arc<SshSession>>,
    db: Database,
}

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

    pub async fn connect(&self, app_handle: AppHandle, hostname: &str, port: u16, username: &str, auth: AuthMethod, cols: u32, rows: u32, sftp_service: &SftpService) -> Result<String, String> {
        let session_id = uuid::Uuid::new_v4().to_string();
        let config = Arc::new(russh::client::Config::default());
        let handler = SshClient { host_key_store: HostKeyStore::new(self.db.clone()), hostname: hostname.to_string(), port };

        let mut handle = tokio::time::timeout(std::time::Duration::from_secs(10), client::connect(config, (hostname, port), handler))
            .await
            .map_err(|_| format!("SSH connection to {}:{} timed out after 10s", hostname, port))?
            .map_err(|e| format!("SSH connection to {}:{} failed: {}", hostname, port, e))?;

        let auth_success = match auth {
            AuthMethod::Password(ref password) => {
                tokio::time::timeout(std::time::Duration::from_secs(10), handle.authenticate_password(username, password))
                    .await
                    .map_err(|_| "SSH password authentication timed out after 10s".to_string())?
                    .map_err(|e| format!("SSH authentication error: {}", e))?
            }
            AuthMethod::Key { ref private_key_pem, ref passphrase } => {
                let key = russh::keys::decode_secret_key(private_key_pem, passphrase.as_deref()).map_err(|e| format!("Failed to decode private key: {}", e))?;
                tokio::time::timeout(std::time::Duration::from_secs(10), handle.authenticate_publickey(username, Arc::new(key)))
                    .await
                    .map_err(|_| "SSH key authentication timed out after 10s".to_string())?
                    .map_err(|e| format!("SSH authentication error: {}", e))?
            }
        };

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

        let channel = handle.channel_open_session().await.map_err(|e| format!("Failed to open session channel: {}", e))?;
        channel.request_pty(true, "xterm-256color", cols, rows, 0, 0, &[]).await.map_err(|e| format!("Failed to request PTY: {}", e))?;
        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));
        let cwd_tracker = CwdTracker::new();
        cwd_tracker.start(handle.clone(), app_handle.clone(), session_id.clone());

        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);

        {   let h = handle.lock().await;
            if let Ok(sftp_channel) = h.channel_open_session().await {
                if sftp_channel.request_subsystem(true, "sftp").await.is_ok() {
                    if let Ok(sftp_client) = russh_sftp::client::SftpSession::new(sftp_channel.into_stream()).await {
                        sftp_service.register_client(&session_id, sftp_client);
                    }
                }
            }
        }

        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 _ = app.emit(&format!("ssh:data:{}", sid), encoded);
                    }
                    Some(ChannelMsg::ExtendedData { ref data, .. }) => {
                        let encoded = base64::engine::general_purpose::STANDARD.encode(data.as_ref());
                        let _ = app.emit(&format!("ssh:data:{}", sid), encoded);
                    }
                    Some(ChannelMsg::ExitStatus { exit_status }) => {
                        let _ = app.emit(&format!("ssh:exit:{}", sid), exit_status);
                        break;
                    }
                    Some(ChannelMsg::Close) | None => {
                        let _ = app.emit(&format!("ssh:close:{}", sid), ());
                        break;
                    }
                    _ => {}
                }
            }
        });

        Ok(session_id)
    }

    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 = session.channel.lock().await;
        channel.data(&data[..]).await.map_err(|e| format!("Failed to write to session {}: {}", session_id, e))
    }

    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 = session.channel.lock().await;
        channel.window_change(cols, rows, 0, 0).await.map_err(|e| format!("Failed to resize session {}: {}", session_id, e))
    }

    pub async fn disconnect(&self, session_id: &str, sftp_service: &SftpService) -> Result<(), String> {
        let (_, session) = self.sessions.remove(session_id).ok_or_else(|| format!("Session {} not found", session_id))?;
        {   let channel = session.channel.lock().await; let _ = channel.eof().await; let _ = channel.close().await; }
        {   let handle = session.handle.lock().await; let _ = handle.disconnect(Disconnect::ByApplication, "", "en").await; }
        sftp_service.remove_client(session_id);
        Ok(())
    }

    pub fn get_session(&self, session_id: &str) -> Option<Arc<SshSession>> {
        self.sessions.get(session_id).map(|entry| entry.clone())
    }

    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()
    }
}