cnctd-service-ssh 0.1.8

//! Tool router for the SSH service.
//!
//! Uses the `rmcp` macros exactly like the template. Tools are thin and
//! delegate to `operations` and `sessions`.

use rmcp::{
    ServerHandler,
    handler::server::{router::tool::ToolRouter, wrapper::Parameters},
    model::{CallToolResult, Content, ErrorData as McpError, ServerCapabilities, ServerInfo},
    schemars, tool, tool_handler, tool_router,
};
use serde::Deserialize;
use std::collections::HashMap;

use crate::operations::{self, SshExecArgs, SshRegisterArgs};
use crate::sessions::{self, types::*};

/// Args for registering an SSH target.
#[derive(Debug, Deserialize, schemars::JsonSchema)]
pub struct SshRegisterRequest {
    /// Client-defined identifier for this target (unique key).
    pub id: String,
    /// Hostname or IP of the SSH server.
    pub host: String,
    /// Username to authenticate as.
    pub user: String,
    /// SSH port (default 22).
    #[serde(default = "default_port")]
    pub port: u16,
    /// Optional passphrase for the private key.
    pub key_passphrase: Option<String>,
    /// Path to OpenSSH known_hosts file (default "~/.ssh/known_hosts").
    #[serde(default = "default_known_hosts")]
    pub known_hosts_path: String,
    /// Identifier for crash-recovery log filtering. Use a stable identifier that
    /// survives session restarts. Examples: 'dot:{dot_id}:user:{user_id}' for cnctd.world, 'claude' for Claude Desktop,
    /// '{app_name}:{unique_id}' for other clients. Logs can be filtered
    /// with: grep '"client":"{your_id}"' ~/.cnctd/ssh_exec.jsonl
    #[serde(default)]
    pub client_id: Option<String>,
}

/// Args for executing a command on a registered target.
#[derive(Debug, Deserialize, schemars::JsonSchema)]
pub struct SshExecRequest {
    /// Target id previously registered via `ssh_register_target`.
    pub id: String,
    /// Shell command to execute remotely.
    pub command: String,
    /// Timeout in seconds (default 120).
    #[serde(default = "default_timeout_secs")]
    pub timeout_secs: u64,
    /// Optional context describing why this command is being run (for crash recovery logs).
    #[serde(default)]
    pub context: Option<String>,
}

/// Args for unregistering a target.
#[derive(Debug, Deserialize, schemars::JsonSchema)]
pub struct SshUnregisterRequest {
    /// Target id to remove.
    pub id: String,
}

// =============================================================================
// Shell Session Request Types
// =============================================================================

/// Args for creating an interactive shell session
#[derive(Debug, Deserialize, schemars::JsonSchema)]
pub struct ShellSessionCreateRequest {
    /// Target ID previously registered via `ssh_register_target`
    pub target_id: String,
    /// Optional human-readable name for this session
    pub name: Option<String>,
    /// Initial shell command (default: user's login shell)
    pub shell: Option<String>,
    /// Terminal width in columns (default: 120)
    #[serde(default = "default_cols")]
    pub cols: u16,
    /// Terminal height in rows (default: 40)
    #[serde(default = "default_rows")]
    pub rows: u16,
    /// Client identifier for session association
    pub client_id: Option<String>,
    /// Environment variables to set
    #[serde(default)]
    pub env: HashMap<String, String>,
}

/// Args for writing input to a session
#[derive(Debug, Deserialize, schemars::JsonSchema)]
pub struct ShellSessionWriteRequest {
    /// Session ID
    pub session_id: String,
    /// Input to send (supports escape sequences like \x03 for Ctrl+C)
    pub input: String,
    /// Whether to append a newline after input (default: false)
    #[serde(default)]
    pub newline: bool,
}

/// Args for reading output from a session
#[derive(Debug, Deserialize, schemars::JsonSchema)]
pub struct ShellSessionReadRequest {
    /// Session ID
    pub session_id: String,
    /// Output format: "raw" (buffered output with ANSI codes), "screen" (current terminal state),
    /// "both", or "stripped" (raw with ANSI codes removed)
    #[serde(default = "default_format")]
    pub format: String,
    /// Whether to consume (clear) the buffer after reading (default: true)
    #[serde(default = "default_true")]
    pub consume: bool,
    /// Maximum wait time for new output in milliseconds (0 = no wait)
    #[serde(default)]
    pub wait_ms: u64,
    /// Minimum bytes to wait for (with wait_ms timeout)
    #[serde(default)]
    pub min_bytes: usize,
    /// Wait until this pattern appears in the screen output (e.g., "❯" or "$ ").
    /// Takes precedence over wait_ms if both are specified.
    #[serde(default)]
    pub wait_for_pattern: Option<String>,
    /// Wait until screen output stabilizes (no changes for this many ms).
    /// Useful for waiting until a TUI finishes rendering.
    #[serde(default)]
    pub wait_for_stable_ms: Option<u64>,
}

/// Args for listing sessions
#[derive(Debug, Deserialize, schemars::JsonSchema)]
pub struct ShellSessionListRequest {
    /// Filter by target ID
    pub target_id: Option<String>,
    /// Filter by client ID
    pub client_id: Option<String>,
    /// Include disconnected sessions (default: true)
    #[serde(default = "default_true")]
    pub include_disconnected: bool,
}

/// Args for reconnecting to a session
#[derive(Debug, Deserialize, schemars::JsonSchema)]
pub struct ShellSessionReconnectRequest {
    /// Session ID to reconnect to
    pub session_id: String,
}

/// Args for resizing a session
#[derive(Debug, Deserialize, schemars::JsonSchema)]
pub struct ShellSessionResizeRequest {
    /// Session ID
    pub session_id: String,
    /// New width in columns
    pub cols: u16,
    /// New height in rows
    pub rows: u16,
}

/// Args for closing a session
#[derive(Debug, Deserialize, schemars::JsonSchema)]
pub struct ShellSessionCloseRequest {
    /// Session ID to close
    pub session_id: String,
    /// Forcefully kill even if processes are running (default: false)
    #[serde(default)]
    pub force: bool,
}

/// SSH service with a ToolRouter, matching the template's structure.
#[derive(Clone)]
pub struct SshToolRouter {
    pub tool_router: ToolRouter<Self>,
}

#[tool_router]
impl SshToolRouter {
    /// Construct the service and initialize the tool router.
    pub fn new() -> Self {
        Self {
            tool_router: Self::tool_router(),
        }
    }

    /// Register or replace an SSH target configuration.
    #[tool(description = "Register or replace an SSH target configuration")]
    async fn ssh_register_target(
        &self,
        Parameters(req): Parameters<SshRegisterRequest>,
    ) -> Result<CallToolResult, McpError> {
        let args = SshRegisterArgs {
            id: req.id,
            host: req.host,
            user: req.user,
            port: req.port,
            key_passphrase: req.key_passphrase,
            known_hosts_path: req.known_hosts_path,
            client_id: req.client_id,
        };
        let v = operations::ssh_register(args).await?;
        Ok(CallToolResult::success(vec![Content::json(v)?]))
    }

    /// Execute a command on a registered target; returns stdout, stderr, exit code.
    #[tool(description = "Execute a command on a registered target")]
    async fn ssh_exec(
        &self,
        Parameters(req): Parameters<SshExecRequest>,
    ) -> Result<CallToolResult, McpError> {
        let args = SshExecArgs {
            id: req.id,
            command: req.command,
            timeout_secs: req.timeout_secs,
            context: req.context,
        };
        let v = operations::ssh_exec(args).await?;
        Ok(CallToolResult::success(vec![Content::json(v)?]))
    }

    /// Unregister a target by id.
    #[tool(description = "Unregister a previously registered target")]
    async fn ssh_unregister_target(
        &self,
        Parameters(req): Parameters<SshUnregisterRequest>,
    ) -> Result<CallToolResult, McpError> {
        let v = operations::ssh_unregister(req.id).await?;
        Ok(CallToolResult::success(vec![Content::json(v)?]))
    }

    // =========================================================================
    // Shell Session Tools
    // =========================================================================

    /// Create a new interactive shell session on a registered SSH target
    #[tool(description = "Create a new interactive shell session on a registered SSH target. Sessions are persistent and survive disconnections.")]
    async fn shell_session_create(
        &self,
        Parameters(req): Parameters<ShellSessionCreateRequest>,
    ) -> Result<CallToolResult, McpError> {
        let args = ShellSessionCreateArgs {
            target_id: req.target_id,
            name: req.name,
            shell: req.shell,
            cols: req.cols,
            rows: req.rows,
            client_id: req.client_id,
            env: req.env,
        };
        let result = sessions::shell_session_create(args).await?;
        Ok(CallToolResult::success(vec![Content::json(serde_json::to_value(result).map_err(|e| {
            McpError::internal_error(format!("JSON error: {}", e), None)
        })?)?]))
    }

    /// Send input to an interactive shell session
    #[tool(description = "Send input to an interactive shell session. Use for commands, keystrokes (\\x03 for Ctrl+C), or any terminal input.")]
    async fn shell_session_write(
        &self,
        Parameters(req): Parameters<ShellSessionWriteRequest>,
    ) -> Result<CallToolResult, McpError> {
        let args = ShellSessionWriteArgs {
            session_id: req.session_id,
            input: req.input,
            newline: req.newline,
        };
        let result = sessions::shell_session_write(args).await?;
        Ok(CallToolResult::success(vec![Content::json(serde_json::to_value(result).map_err(|e| {
            McpError::internal_error(format!("JSON error: {}", e), None)
        })?)?]))
    }

    /// Read output from an interactive shell session
    #[tool(description = "Read output from an interactive shell session. Supports raw output, screen state (for TUI apps), stripped (no ANSI codes), or both. Use wait_for_pattern to wait for a prompt, or wait_for_stable_ms to wait for output to stop changing.")]
    async fn shell_session_read(
        &self,
        Parameters(req): Parameters<ShellSessionReadRequest>,
    ) -> Result<CallToolResult, McpError> {
        // Convert format string to OutputFormat enum
        let format = match req.format.to_lowercase().as_str() {
            "raw" => OutputFormat::Raw,
            "screen" => OutputFormat::Screen,
            "both" => OutputFormat::Both,
            "stripped" => OutputFormat::Stripped,
            _ => OutputFormat::Raw, // Default
        };

        let args = ShellSessionReadArgs {
            session_id: req.session_id,
            format,
            consume: req.consume,
            wait_ms: req.wait_ms,
            min_bytes: req.min_bytes,
            wait_for_pattern: req.wait_for_pattern,
            wait_for_stable_ms: req.wait_for_stable_ms,
        };
        let result = sessions::shell_session_read(args).await?;
        Ok(CallToolResult::success(vec![Content::json(serde_json::to_value(result).map_err(|e| {
            McpError::internal_error(format!("JSON error: {}", e), None)
        })?)?]))
    }

    /// List interactive shell sessions
    #[tool(description = "List interactive shell sessions. Can filter by target or client ID.")]
    async fn shell_session_list(
        &self,
        Parameters(req): Parameters<ShellSessionListRequest>,
    ) -> Result<CallToolResult, McpError> {
        let args = ShellSessionListArgs {
            target_id: req.target_id,
            client_id: req.client_id,
            include_disconnected: req.include_disconnected,
        };
        let result = sessions::shell_session_list(args).await?;
        Ok(CallToolResult::success(vec![Content::json(serde_json::to_value(result).map_err(|e| {
            McpError::internal_error(format!("JSON error: {}", e), None)
        })?)?]))
    }

    /// Reconnect to a disconnected shell session
    #[tool(description = "Reconnect to a disconnected shell session. The session must still exist on the remote server.")]
    async fn shell_session_reconnect(
        &self,
        Parameters(req): Parameters<ShellSessionReconnectRequest>,
    ) -> Result<CallToolResult, McpError> {
        let args = ShellSessionReconnectArgs {
            session_id: req.session_id,
        };
        let result = sessions::shell_session_reconnect(args).await?;
        Ok(CallToolResult::success(vec![Content::json(serde_json::to_value(result).map_err(|e| {
            McpError::internal_error(format!("JSON error: {}", e), None)
        })?)?]))
    }

    /// Resize a shell session's terminal
    #[tool(description = "Resize a shell session's terminal dimensions. Important for TUI applications.")]
    async fn shell_session_resize(
        &self,
        Parameters(req): Parameters<ShellSessionResizeRequest>,
    ) -> Result<CallToolResult, McpError> {
        let args = ShellSessionResizeArgs {
            session_id: req.session_id,
            cols: req.cols,
            rows: req.rows,
        };
        let result = sessions::shell_session_resize(args).await?;
        Ok(CallToolResult::success(vec![Content::json(serde_json::to_value(result).map_err(|e| {
            McpError::internal_error(format!("JSON error: {}", e), None)
        })?)?]))
    }

    /// Close an interactive shell session
    #[tool(description = "Close an interactive shell session. Terminates the remote shell.")]
    async fn shell_session_close(
        &self,
        Parameters(req): Parameters<ShellSessionCloseRequest>,
    ) -> Result<CallToolResult, McpError> {
        let args = ShellSessionCloseArgs {
            session_id: req.session_id,
            force: req.force,
        };
        let result = sessions::shell_session_close(args).await?;
        Ok(CallToolResult::success(vec![Content::json(serde_json::to_value(result).map_err(|e| {
            McpError::internal_error(format!("JSON error: {}", e), None)
        })?)?]))
    }
}

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

#[tool_handler]
impl ServerHandler for SshToolRouter {
    fn get_info(&self) -> ServerInfo {
        ServerInfo {
            instructions: Some("cnctd-service-ssh: secure SSH command execution with enforced host key verification and whitelisted keys".into()),
            capabilities: ServerCapabilities::builder().enable_tools().build(),
            ..Default::default()
        }
    }
}

/* ---------- defaults mirror operations.rs for schema consistency ---------- */

fn default_port() -> u16 {
    22
}
fn default_known_hosts() -> String {
    "~/.ssh/known_hosts".into()
}
fn default_timeout_secs() -> u64 {
    120
}

/* ---------- defaults for shell session types ---------- */

fn default_cols() -> u16 {
    120
}
fn default_rows() -> u16 {
    40
}
fn default_true() -> bool {
    true
}
fn default_format() -> String {
    "raw".to_string()
}