mcp_tools/servers/

system_tools.rs

//! System Tools MCP Server
//!
//! Provides system operations and shell command execution via MCP protocol including:
//! - Shell command execution
//! - Environment variable access
//! - Process management
//! - System information gathering
//! - File system operations

use async_trait::async_trait;
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use std::process::{Command, Stdio};
use std::time::{Duration, SystemTime, UNIX_EPOCH};
use tokio::process::Command as TokioCommand;
use tracing::{debug, info, warn};

use crate::common::{
    BaseServer, McpContent, McpServerBase, McpTool, McpToolRequest, McpToolResponse,
    ServerCapabilities, ServerConfig,
};
use crate::{McpToolsError, Result};

/// System Tools MCP Server
pub struct SystemToolsServer {
    base: BaseServer,
}

/// Command execution result
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct CommandResult {
    pub command: String,
    pub args: Vec<String>,
    pub exit_code: i32,
    pub stdout: String,
    pub stderr: String,
    pub execution_time: u64,
    pub working_directory: String,
}

/// System information
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct SystemInfo {
    pub os: String,
    pub arch: String,
    pub hostname: String,
    pub username: String,
    pub uptime: u64,
    pub cpu_count: usize,
    pub memory_total: u64,
    pub memory_available: u64,
    pub disk_usage: Vec<DiskInfo>,
}

/// Disk usage information
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct DiskInfo {
    pub mount_point: String,
    pub total: u64,
    pub available: u64,
    pub used: u64,
    pub filesystem: String,
}

/// Process information
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ProcessInfo {
    pub pid: u32,
    pub name: String,
    pub cpu_usage: f32,
    pub memory_usage: u64,
    pub status: String,
    pub start_time: u64,
}

/// Environment variable information
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct EnvVarInfo {
    pub name: String,
    pub value: String,
    pub is_system: bool,
}

impl SystemToolsServer {
    pub async fn new(config: ServerConfig) -> Result<Self> {
        let base = BaseServer::new(config).await?;
        Ok(Self { base })
    }

    /// Execute shell command
    async fn execute_command(
        &self,
        command: &str,
        args: &[String],
        working_dir: Option<&str>,
        timeout: Option<u64>,
    ) -> Result<CommandResult> {
        debug!("Executing command: {} {:?}", command, args);

        let start_time = std::time::Instant::now();
        let mut cmd = TokioCommand::new(command);

        // Set arguments
        cmd.args(args);

        // Set working directory
        if let Some(dir) = working_dir {
            cmd.current_dir(dir);
        }

        // Configure stdio
        cmd.stdout(Stdio::piped())
            .stderr(Stdio::piped())
            .stdin(Stdio::null());

        // Set timeout
        let timeout_duration = Duration::from_secs(timeout.unwrap_or(30));

        // Execute command with timeout
        let output = tokio::time::timeout(timeout_duration, cmd.output())
            .await
            .map_err(|_| McpToolsError::Server("Command execution timed out".to_string()))?
            .map_err(|e| McpToolsError::Server(format!("Failed to execute command: {}", e)))?;

        let execution_time = start_time.elapsed().as_millis() as u64;
        let working_directory = working_dir.unwrap_or(".").to_string();

        Ok(CommandResult {
            command: command.to_string(),
            args: args.to_vec(),
            exit_code: output.status.code().unwrap_or(-1),
            stdout: String::from_utf8_lossy(&output.stdout).to_string(),
            stderr: String::from_utf8_lossy(&output.stderr).to_string(),
            execution_time,
            working_directory,
        })
    }

    /// Get system information
    async fn get_system_info(&self) -> Result<SystemInfo> {
        debug!("Gathering system information");

        // Get basic system info
        let os = std::env::consts::OS.to_string();
        let arch = std::env::consts::ARCH.to_string();

        // Get hostname (simplified)
        let hostname = std::env::var("COMPUTERNAME")
            .or_else(|_| std::env::var("HOSTNAME"))
            .unwrap_or_else(|_| "unknown".to_string());

        // Get username
        let username = std::env::var("USER")
            .or_else(|_| std::env::var("USERNAME"))
            .unwrap_or_else(|_| "unknown".to_string());

        // Get uptime (simplified - would need platform-specific code)
        let uptime = SystemTime::now()
            .duration_since(UNIX_EPOCH)
            .unwrap_or_default()
            .as_secs();

        // Get CPU count
        let cpu_count = num_cpus::get();

        // Memory info (simplified - would need platform-specific code)
        let memory_total = 8 * 1024 * 1024 * 1024; // 8GB placeholder
        let memory_available = 4 * 1024 * 1024 * 1024; // 4GB placeholder

        // Disk usage (simplified)
        let disk_usage = vec![DiskInfo {
            mount_point: "/".to_string(),
            total: 100 * 1024 * 1024 * 1024,    // 100GB placeholder
            available: 50 * 1024 * 1024 * 1024, // 50GB placeholder
            used: 50 * 1024 * 1024 * 1024,      // 50GB placeholder
            filesystem: "ext4".to_string(),
        }];

        Ok(SystemInfo {
            os,
            arch,
            hostname,
            username,
            uptime,
            cpu_count,
            memory_total,
            memory_available,
            disk_usage,
        })
    }

    /// Get environment variables
    async fn get_environment_variables(&self, filter: Option<&str>) -> Result<Vec<EnvVarInfo>> {
        debug!("Getting environment variables");

        let mut env_vars = Vec::new();

        for (key, value) in std::env::vars() {
            // Apply filter if provided
            if let Some(filter_str) = filter {
                if !key.to_lowercase().contains(&filter_str.to_lowercase()) {
                    continue;
                }
            }

            // Determine if it's a system variable (simplified heuristic)
            let is_system = key.starts_with("SYSTEM")
                || key.starts_with("OS")
                || key.starts_with("PROCESSOR")
                || key == "PATH"
                || key == "HOME"
                || key == "USER"
                || key == "USERNAME";

            env_vars.push(EnvVarInfo {
                name: key,
                value,
                is_system,
            });
        }

        // Sort by name
        env_vars.sort_by(|a, b| a.name.cmp(&b.name));

        Ok(env_vars)
    }

    /// Get running processes (simplified)
    async fn get_processes(&self) -> Result<Vec<ProcessInfo>> {
        debug!("Getting process information");

        // This is a simplified implementation
        // In a real implementation, we would use platform-specific APIs
        let processes = vec![ProcessInfo {
            pid: std::process::id(),
            name: "mcp-tools".to_string(),
            cpu_usage: 1.5,
            memory_usage: 50 * 1024 * 1024, // 50MB
            status: "Running".to_string(),
            start_time: SystemTime::now()
                .duration_since(UNIX_EPOCH)
                .unwrap_or_default()
                .as_secs(),
        }];

        Ok(processes)
    }

    /// Check if command is safe to execute
    fn is_safe_command(&self, command: &str) -> bool {
        // Basic safety check - in production, this would be more comprehensive
        let dangerous_commands = [
            "rm",
            "del",
            "format",
            "fdisk",
            "mkfs",
            "dd",
            "shutdown",
            "reboot",
            "halt",
            "poweroff",
            "sudo",
            "su",
            "passwd",
            "chmod",
            "chown",
            "iptables",
            "ufw",
            "firewall-cmd",
        ];

        !dangerous_commands
            .iter()
            .any(|&dangerous| command.contains(dangerous))
    }
}

#[async_trait]
impl McpServerBase for SystemToolsServer {
    async fn get_capabilities(&self) -> Result<ServerCapabilities> {
        let mut capabilities = self.base.get_capabilities().await?;

        // Add System Tools-specific tools
        let system_tools = vec![
            McpTool {
                name: "execute_command".to_string(),
                description: "Execute shell commands with safety checks and timeout".to_string(),
                input_schema: serde_json::json!({
                    "type": "object",
                    "properties": {
                        "command": {
                            "type": "string",
                            "description": "Command to execute"
                        },
                        "args": {
                            "type": "array",
                            "items": {"type": "string"},
                            "description": "Command arguments"
                        },
                        "working_dir": {
                            "type": "string",
                            "description": "Working directory for command execution"
                        },
                        "timeout": {
                            "type": "integer",
                            "description": "Timeout in seconds (default: 30, max: 300)",
                            "minimum": 1,
                            "maximum": 300
                        }
                    },
                    "required": ["command"]
                }),
                category: "system".to_string(),
                requires_permission: true,
                permissions: vec!["system.execute".to_string()],
            },
            McpTool {
                name: "get_system_info".to_string(),
                description: "Get comprehensive system information including OS, hardware, and resource usage".to_string(),
                input_schema: serde_json::json!({
                    "type": "object",
                    "properties": {}
                }),
                category: "system".to_string(),
                requires_permission: false,
                permissions: vec![],
            },
            McpTool {
                name: "get_environment".to_string(),
                description: "Get environment variables with optional filtering".to_string(),
                input_schema: serde_json::json!({
                    "type": "object",
                    "properties": {
                        "filter": {
                            "type": "string",
                            "description": "Filter environment variables by name (case-insensitive substring match)"
                        },
                        "include_system": {
                            "type": "boolean",
                            "description": "Include system environment variables (default: true)"
                        }
                    }
                }),
                category: "system".to_string(),
                requires_permission: false,
                permissions: vec![],
            },
            McpTool {
                name: "get_processes".to_string(),
                description: "Get information about running processes".to_string(),
                input_schema: serde_json::json!({
                    "type": "object",
                    "properties": {
                        "filter": {
                            "type": "string",
                            "description": "Filter processes by name"
                        }
                    }
                }),
                category: "system".to_string(),
                requires_permission: true,
                permissions: vec!["system.processes".to_string()],
            },
        ];

        capabilities.tools = system_tools;
        Ok(capabilities)
    }

    async fn handle_tool_request(&self, request: McpToolRequest) -> Result<McpToolResponse> {
        info!("Handling System Tools request: {}", request.tool);

        match request.tool.as_str() {
            "execute_command" => {
                debug!("Executing shell command");

                let command = request
                    .arguments
                    .get("command")
                    .and_then(|v| v.as_str())
                    .ok_or_else(|| {
                        McpToolsError::Server("Missing 'command' parameter".to_string())
                    })?;

                // Safety check
                if !self.is_safe_command(command) {
                    return Ok(McpToolResponse {
                        id: request.id,
                        content: vec![McpContent::text(
                            "Command rejected for security reasons".to_string(),
                        )],
                        is_error: true,
                        error: Some("Unsafe command detected".to_string()),
                        metadata: HashMap::new(),
                    });
                }

                let args: Vec<String> = request
                    .arguments
                    .get("args")
                    .and_then(|v| v.as_array())
                    .map(|arr| {
                        arr.iter()
                            .filter_map(|v| v.as_str().map(|s| s.to_string()))
                            .collect()
                    })
                    .unwrap_or_else(|| Vec::new());

                let working_dir = request
                    .arguments
                    .get("working_dir")
                    .and_then(|v| v.as_str());

                let timeout = request.arguments.get("timeout").and_then(|v| v.as_u64());

                let result = self
                    .execute_command(command, &args, working_dir, timeout)
                    .await?;

                let content_text = format!(
                    "Command Execution Complete\n\
                    Command: {} {:?}\n\
                    Exit Code: {}\n\
                    Execution Time: {}ms\n\
                    Working Directory: {}\n\n\
                    STDOUT:\n{}\n\n\
                    STDERR:\n{}",
                    result.command,
                    result.args,
                    result.exit_code,
                    result.execution_time,
                    result.working_directory,
                    if result.stdout.is_empty() {
                        "(empty)"
                    } else {
                        &result.stdout
                    },
                    if result.stderr.is_empty() {
                        "(empty)"
                    } else {
                        &result.stderr
                    }
                );

                let mut metadata = HashMap::new();
                metadata.insert("command_result".to_string(), serde_json::to_value(result)?);

                Ok(McpToolResponse {
                    id: request.id,
                    content: vec![McpContent::text(content_text)],
                    is_error: false,
                    error: None,
                    metadata,
                })
            }
            "get_system_info" => {
                debug!("Getting system information");

                let system_info = self.get_system_info().await?;

                let content_text = format!(
                    "System Information\n\
                    OS: {}\n\
                    Architecture: {}\n\
                    Hostname: {}\n\
                    Username: {}\n\
                    CPU Cores: {}\n\
                    Memory Total: {} GB\n\
                    Memory Available: {} GB\n\
                    Uptime: {} seconds",
                    system_info.os,
                    system_info.arch,
                    system_info.hostname,
                    system_info.username,
                    system_info.cpu_count,
                    system_info.memory_total / (1024 * 1024 * 1024),
                    system_info.memory_available / (1024 * 1024 * 1024),
                    system_info.uptime
                );

                let mut metadata = HashMap::new();
                metadata.insert(
                    "system_info".to_string(),
                    serde_json::to_value(system_info)?,
                );

                Ok(McpToolResponse {
                    id: request.id,
                    content: vec![McpContent::text(content_text)],
                    is_error: false,
                    error: None,
                    metadata,
                })
            }
            "get_environment" => {
                debug!("Getting environment variables");

                let filter = request.arguments.get("filter").and_then(|v| v.as_str());

                let env_vars = self.get_environment_variables(filter).await?;

                let content_text = format!(
                    "Environment Variables\n\
                    Total Variables: {}\n\
                    Filter Applied: {}\n\n{}",
                    env_vars.len(),
                    filter.unwrap_or("None"),
                    env_vars
                        .iter()
                        .take(20) // Limit to first 20 for display
                        .map(|var| format!(
                            "{}={}",
                            var.name,
                            if var.value.len() > 50 {
                                format!("{}...", &var.value[..50])
                            } else {
                                var.value.clone()
                            }
                        ))
                        .collect::<Vec<_>>()
                        .join("\n")
                );

                let mut metadata = HashMap::new();
                metadata.insert(
                    "environment_variables".to_string(),
                    serde_json::to_value(env_vars)?,
                );

                Ok(McpToolResponse {
                    id: request.id,
                    content: vec![McpContent::text(content_text)],
                    is_error: false,
                    error: None,
                    metadata,
                })
            }
            "get_processes" => {
                debug!("Getting process information");

                let processes = self.get_processes().await?;

                let content_text = format!(
                    "Running Processes\n\
                    Total Processes: {}\n\n{}",
                    processes.len(),
                    processes
                        .iter()
                        .map(|proc| format!(
                            "PID: {} | Name: {} | CPU: {:.1}% | Memory: {} MB | Status: {}",
                            proc.pid,
                            proc.name,
                            proc.cpu_usage,
                            proc.memory_usage / (1024 * 1024),
                            proc.status
                        ))
                        .collect::<Vec<_>>()
                        .join("\n")
                );

                let mut metadata = HashMap::new();
                metadata.insert("processes".to_string(), serde_json::to_value(processes)?);

                Ok(McpToolResponse {
                    id: request.id,
                    content: vec![McpContent::text(content_text)],
                    is_error: false,
                    error: None,
                    metadata,
                })
            }
            _ => {
                warn!("Unknown System Tools request: {}", request.tool);
                Err(McpToolsError::Server(format!(
                    "Unknown System Tools request: {}",
                    request.tool
                )))
            }
        }
    }

    async fn get_stats(&self) -> Result<crate::common::ServerStats> {
        self.base.get_stats().await
    }

    async fn initialize(&mut self) -> Result<()> {
        info!("Initializing System Tools MCP Server");
        Ok(())
    }

    async fn shutdown(&mut self) -> Result<()> {
        info!("Shutting down System Tools MCP Server");
        Ok(())
    }
}