cnctd-service-ssh 0.1.8

//! Interactive shell session management.
//!
//! This module provides functionality for creating and managing persistent
//! interactive shell sessions over SSH. Sessions support:
//!
//! - Full TUI applications (vim, htop, tmux)
//! - Simple REPL-style interactions
//! - Persistence across LLM conversations via remote tmux
//! - Screen state capture for AI agents

pub mod connection;
pub mod registry;
pub mod terminal;
pub mod types;

use crate::operations::lookup_target;
use crate::service_error::ServiceError;
use connection::SshConnection;
use once_cell::sync::Lazy;
use registry::{ShellSession, ShellSessionRegistry};
use std::sync::Arc;
use tracing::info;
use types::*;

/// Global session registry for MCP server usage
static GLOBAL_SESSION_REGISTRY: Lazy<Arc<ShellSessionRegistry>> =
    Lazy::new(|| Arc::new(ShellSessionRegistry::new()));

/// Service for managing interactive shell sessions
pub struct ShellSessionService {
    registry: Arc<ShellSessionRegistry>,
}

impl ShellSessionService {
    /// Create a new shell session service with its own registry
    pub fn new() -> Self {
        Self {
            registry: Arc::new(ShellSessionRegistry::new()),
        }
    }

    /// Create a new shell session service using the global registry
    pub fn global() -> Self {
        Self {
            registry: Arc::clone(&GLOBAL_SESSION_REGISTRY),
        }
    }

    /// Create a new interactive shell session
    pub async fn create(&self, args: ShellSessionCreateArgs) -> Result<ShellSessionCreateResult, ServiceError> {
        create_session_impl(&self.registry, args).await
    }

    /// Write input to a session
    pub async fn write(&self, args: ShellSessionWriteArgs) -> Result<ShellSessionWriteResult, ServiceError> {
        write_session_impl(&self.registry, args).await
    }

    /// Read output from a session
    pub async fn read(&self, args: ShellSessionReadArgs) -> Result<ShellSessionReadResult, ServiceError> {
        read_session_impl(&self.registry, args).await
    }

    /// List sessions
    pub async fn list(&self, args: ShellSessionListArgs) -> Result<ShellSessionListResult, ServiceError> {
        list_sessions_impl(&self.registry, args).await
    }

    /// Reconnect to a disconnected session
    pub async fn reconnect(&self, args: ShellSessionReconnectArgs) -> Result<ShellSessionReconnectResult, ServiceError> {
        reconnect_session_impl(&self.registry, args).await
    }

    /// Resize a session's terminal
    pub async fn resize(&self, args: ShellSessionResizeArgs) -> Result<ShellSessionResizeResult, ServiceError> {
        resize_session_impl(&self.registry, args).await
    }

    /// Close a session
    pub async fn close(&self, args: ShellSessionCloseArgs) -> Result<ShellSessionCloseResult, ServiceError> {
        close_session_impl(&self.registry, args).await
    }
}

impl Default for ShellSessionService {
    fn default() -> Self {
        Self::new()
    }
}

// ============================================================================
// Global API functions (for MCP server)
// ============================================================================

/// Create a new interactive shell session (global registry)
pub async fn shell_session_create(args: ShellSessionCreateArgs) -> Result<ShellSessionCreateResult, ServiceError> {
    create_session_impl(&GLOBAL_SESSION_REGISTRY, args).await
}

/// Write input to a session (global registry)
pub async fn shell_session_write(args: ShellSessionWriteArgs) -> Result<ShellSessionWriteResult, ServiceError> {
    write_session_impl(&GLOBAL_SESSION_REGISTRY, args).await
}

/// Read output from a session (global registry)
pub async fn shell_session_read(args: ShellSessionReadArgs) -> Result<ShellSessionReadResult, ServiceError> {
    read_session_impl(&GLOBAL_SESSION_REGISTRY, args).await
}

/// List sessions (global registry)
pub async fn shell_session_list(args: ShellSessionListArgs) -> Result<ShellSessionListResult, ServiceError> {
    list_sessions_impl(&GLOBAL_SESSION_REGISTRY, args).await
}

/// Reconnect to a session (global registry)
pub async fn shell_session_reconnect(args: ShellSessionReconnectArgs) -> Result<ShellSessionReconnectResult, ServiceError> {
    reconnect_session_impl(&GLOBAL_SESSION_REGISTRY, args).await
}

/// Resize a session (global registry)
pub async fn shell_session_resize(args: ShellSessionResizeArgs) -> Result<ShellSessionResizeResult, ServiceError> {
    resize_session_impl(&GLOBAL_SESSION_REGISTRY, args).await
}

/// Close a session (global registry)
pub async fn shell_session_close(args: ShellSessionCloseArgs) -> Result<ShellSessionCloseResult, ServiceError> {
    close_session_impl(&GLOBAL_SESSION_REGISTRY, args).await
}

// ============================================================================
// Implementation functions
// ============================================================================

async fn create_session_impl(
    registry: &ShellSessionRegistry,
    args: ShellSessionCreateArgs,
) -> Result<ShellSessionCreateResult, ServiceError> {
    // Look up the target configuration
    let target = lookup_target(&args.target_id).await?;

    // Generate session ID
    let session_id = uuid::Uuid::new_v4().to_string();
    let tmux_session = format!("cnctd-ssh-{}", &session_id[..8]);

    info!(
        "Creating shell session {} for target {} (tmux: {})",
        session_id, args.target_id, tmux_session
    );

    // Build the shell command that will run inside tmux
    // We create a tmux session and then attach to it
    let shell_cmd = args.shell.as_deref();

    // Connect via SSH with PTY
    let connection = SshConnection::connect(
        &target.host,
        target.port,
        &target.user,
        &target.key_path,
        target.key_passphrase.as_deref(),
        args.cols,
        args.rows,
        shell_cmd,
    )
    .await?;

    // Create the session
    let mut session = ShellSession::new(
        session_id.clone(),
        args.target_id.clone(),
        args.name.clone(),
        args.client_id.clone(),
        tmux_session,
        args.cols,
        args.rows,
        connection,
    );

    // Start output reader
    session.start_output_reader();

    // Wait briefly for initial output
    tokio::time::sleep(std::time::Duration::from_millis(100)).await;

    // Get initial screen state
    let screen = session.screen_state().await;
    let info = session.info().await;

    // Add to registry
    registry.add(session).await;

    Ok(ShellSessionCreateResult {
        session_id,
        info,
        screen,
    })
}

async fn write_session_impl(
    registry: &ShellSessionRegistry,
    args: ShellSessionWriteArgs,
) -> Result<ShellSessionWriteResult, ServiceError> {
    let session_lock = registry
        .get(&args.session_id)
        .await
        .ok_or_else(|| ServiceError::NotFound(format!("Session not found: {}", args.session_id)))?;

    let session = session_lock.read().await;

    // Process input - handle escape sequences
    let mut data = process_escape_sequences(&args.input);
    if args.newline {
        data.push(b'\n');
    }

    let bytes_sent = session.write(&data).await?;

    Ok(ShellSessionWriteResult {
        session_id: args.session_id,
        bytes_sent,
    })
}

async fn read_session_impl(
    registry: &ShellSessionRegistry,
    args: ShellSessionReadArgs,
) -> Result<ShellSessionReadResult, ServiceError> {
    let session_lock = registry
        .get(&args.session_id)
        .await
        .ok_or_else(|| ServiceError::NotFound(format!("Session not found: {}", args.session_id)))?;

    let session = session_lock.read().await;

    // Determine effective timeout (use wait_ms as the max timeout for pattern/stable waits)
    let timeout_ms = if args.wait_ms > 0 { args.wait_ms } else { 30000 }; // Default 30s max

    // Wait for pattern if requested (takes precedence)
    let pattern_matched = if let Some(ref pattern) = args.wait_for_pattern {
        Some(session.wait_for_pattern(pattern, timeout_ms).await)
    } else {
        None
    };

    // Wait for stable output if requested (and pattern wasn't requested or already matched)
    let stabilized = if let Some(stable_ms) = args.wait_for_stable_ms {
        // Only wait for stable if we're not waiting for pattern, or pattern was found
        if args.wait_for_pattern.is_none() || pattern_matched == Some(true) {
            Some(session.wait_for_stable(stable_ms, timeout_ms).await)
        } else {
            Some(false)
        }
    } else {
        None
    };

    // If no special waits, use the basic wait_ms/min_bytes
    if args.wait_for_pattern.is_none() && args.wait_for_stable_ms.is_none() && args.wait_ms > 0 {
        session.wait_for_output(args.wait_ms, args.min_bytes).await;
    }

    let state = *session.state.read().await;

    // Get output based on format
    let (raw, screen, buffer_size, truncated) = match args.format {
        OutputFormat::Raw => {
            let (text, remaining, truncated) = session.read(args.consume).await;
            (Some(text), None, remaining, truncated)
        }
        OutputFormat::Stripped => {
            let (text, remaining, truncated) = session.read(args.consume).await;
            let stripped = strip_ansi_codes(&text);
            (Some(stripped), None, remaining, truncated)
        }
        OutputFormat::Screen => {
            let screen = session.screen_state().await;
            // Don't consume buffer when only getting screen
            let (_, remaining, truncated) = session.read(false).await;
            (None, Some(screen), remaining, truncated)
        }
        OutputFormat::Both => {
            let (text, remaining, truncated) = session.read(args.consume).await;
            let screen = session.screen_state().await;
            (Some(text), Some(screen), remaining, truncated)
        }
    };

    Ok(ShellSessionReadResult {
        session_id: args.session_id,
        raw,
        screen,
        buffer_size,
        truncated,
        state,
        pattern_matched,
        stabilized,
    })
}

async fn list_sessions_impl(
    registry: &ShellSessionRegistry,
    args: ShellSessionListArgs,
) -> Result<ShellSessionListResult, ServiceError> {
    let sessions = registry
        .list(
            args.target_id.as_deref(),
            args.client_id.as_deref(),
            args.include_disconnected,
        )
        .await;

    Ok(ShellSessionListResult { sessions })
}

async fn reconnect_session_impl(
    _registry: &ShellSessionRegistry,
    _args: ShellSessionReconnectArgs,
) -> Result<ShellSessionReconnectResult, ServiceError> {
    // TODO: Implement reconnection
    // This would involve:
    // 1. Finding the session in the registry
    // 2. Re-establishing the SSH connection
    // 3. Attaching to the existing tmux session
    // 4. Restarting the output reader

    Err(ServiceError::Internal(
        "Reconnection not yet implemented - coming soon".to_string(),
    ))
}

async fn resize_session_impl(
    registry: &ShellSessionRegistry,
    args: ShellSessionResizeArgs,
) -> Result<ShellSessionResizeResult, ServiceError> {
    let session_lock = registry
        .get(&args.session_id)
        .await
        .ok_or_else(|| ServiceError::NotFound(format!("Session not found: {}", args.session_id)))?;

    let mut session = session_lock.write().await;
    session.resize(args.cols, args.rows).await?;

    Ok(ShellSessionResizeResult {
        session_id: args.session_id,
        size: (args.cols, args.rows),
    })
}

async fn close_session_impl(
    registry: &ShellSessionRegistry,
    args: ShellSessionCloseArgs,
) -> Result<ShellSessionCloseResult, ServiceError> {
    let session_lock = registry
        .remove(&args.session_id)
        .await
        .ok_or_else(|| ServiceError::NotFound(format!("Session not found: {}", args.session_id)))?;

    // Take ownership of the session
    let session = match Arc::try_unwrap(session_lock) {
        Ok(rwlock) => rwlock.into_inner(),
        Err(_) => {
            return Err(ServiceError::Internal(
                "Session is still in use".to_string(),
            ))
        }
    };

    let closed = session.close(args.force).await?;

    Ok(ShellSessionCloseResult {
        session_id: args.session_id,
        closed,
    })
}

/// Process escape sequences in input string
/// Converts \xNN hex escapes and \n, \r, \t, etc.
fn process_escape_sequences(input: &str) -> Vec<u8> {
    let mut result = Vec::with_capacity(input.len());
    let mut chars = input.chars().peekable();

    while let Some(c) = chars.next() {
        if c == '\\' {
            match chars.peek() {
                Some('x') => {
                    chars.next(); // consume 'x'
                    let hex: String = chars.by_ref().take(2).collect();
                    if let Ok(byte) = u8::from_str_radix(&hex, 16) {
                        result.push(byte);
                    } else {
                        // Invalid hex, output as-is
                        result.extend_from_slice(b"\\x");
                        result.extend_from_slice(hex.as_bytes());
                    }
                }
                Some('n') => {
                    chars.next();
                    result.push(b'\n');
                }
                Some('r') => {
                    chars.next();
                    result.push(b'\r');
                }
                Some('t') => {
                    chars.next();
                    result.push(b'\t');
                }
                Some('\\') => {
                    chars.next();
                    result.push(b'\\');
                }
                Some('0') => {
                    chars.next();
                    result.push(0);
                }
                _ => {
                    result.push(b'\\');
                }
            }
        } else {
            let mut buf = [0u8; 4];
            let s = c.encode_utf8(&mut buf);
            result.extend_from_slice(s.as_bytes());
        }
    }

    result
}

/// Strip ANSI escape sequences from a string
/// Removes CSI sequences like colors, cursor movement, etc.
fn strip_ansi_codes(input: &str) -> String {
    let mut result = String::with_capacity(input.len());
    let mut chars = input.chars().peekable();

    while let Some(c) = chars.next() {
        if c == '\x1b' {
            // ESC character - start of escape sequence
            if chars.peek() == Some(&'[') {
                chars.next(); // consume '['
                // Skip until we hit a letter (end of CSI sequence)
                while let Some(&ch) = chars.peek() {
                    chars.next();
                    if ch.is_ascii_alphabetic() || ch == '~' {
                        break;
                    }
                }
            } else if chars.peek() == Some(&']') {
                // OSC sequence (operating system command) - skip until BEL or ST
                chars.next(); // consume ']'
                while let Some(&ch) = chars.peek() {
                    chars.next();
                    if ch == '\x07' {
                        break; // BEL
                    }
                    if ch == '\x1b' {
                        if chars.peek() == Some(&'\\') {
                            chars.next(); // consume '\\' for ST
                            break;
                        }
                    }
                }
            } else {
                // Other escape sequences - skip the next character
                chars.next();
            }
        } else if c == '\x0f' || c == '\x0e' {
            // SI/SO - shift in/out, skip
        } else if c.is_control() && c != '\n' && c != '\r' && c != '\t' {
            // Skip other control characters (except newline, carriage return, tab)
        } else {
            result.push(c);
        }
    }

    result
}

/// Get tool definitions for shell session tools
pub fn get_shell_session_tool_definitions() -> Vec<crate::operations::ToolDefinition> {
    use crate::operations::ToolDefinition;
    use schemars::schema_for;

    vec![
        ToolDefinition {
            name: "shell_session_create".to_string(),
            description: "Create a new interactive shell session on a registered SSH target. Sessions are persistent and survive disconnections.".to_string(),
            input_schema: serde_json::to_value(schema_for!(ShellSessionCreateArgs)).unwrap_or_default(),
        },
        ToolDefinition {
            name: "shell_session_write".to_string(),
            description: "Send input to an interactive shell session. Use for commands, keystrokes (\\x03 for Ctrl+C), or any terminal input.".to_string(),
            input_schema: serde_json::to_value(schema_for!(ShellSessionWriteArgs)).unwrap_or_default(),
        },
        ToolDefinition {
            name: "shell_session_read".to_string(),
            description: "Read output from an interactive shell session. Supports raw output, screen state (for TUI apps), or both.".to_string(),
            input_schema: serde_json::to_value(schema_for!(ShellSessionReadArgs)).unwrap_or_default(),
        },
        ToolDefinition {
            name: "shell_session_list".to_string(),
            description: "List interactive shell sessions. Can filter by target or client ID.".to_string(),
            input_schema: serde_json::to_value(schema_for!(ShellSessionListArgs)).unwrap_or_default(),
        },
        ToolDefinition {
            name: "shell_session_reconnect".to_string(),
            description: "Reconnect to a disconnected shell session. The session must still exist on the remote server.".to_string(),
            input_schema: serde_json::to_value(schema_for!(ShellSessionReconnectArgs)).unwrap_or_default(),
        },
        ToolDefinition {
            name: "shell_session_resize".to_string(),
            description: "Resize a shell session's terminal dimensions. Important for TUI applications.".to_string(),
            input_schema: serde_json::to_value(schema_for!(ShellSessionResizeArgs)).unwrap_or_default(),
        },
        ToolDefinition {
            name: "shell_session_close".to_string(),
            description: "Close an interactive shell session. Terminates the remote shell.".to_string(),
            input_schema: serde_json::to_value(schema_for!(ShellSessionCloseArgs)).unwrap_or_default(),
        },
    ]
}