linux-mcp-server 0.1.0.dev0__py3-none-any.whl

linux_mcp_server/tools/ssh_executor.py

@@ -0,0 +1,324 @@
+"""SSH executor for remote command execution.
+
+This module provides functionality to execute commands on remote systems via SSH,
+with connection pooling and SSH key discovery. It seamlessly routes commands to
+either local or remote execution based on the provided parameters.
+"""
+
+import asyncio
+import logging
+import os
+import shlex
+import subprocess
+import time
+from pathlib import Path
+from typing import Optional, Tuple
+
+import asyncssh
+
+from ..audit import log_ssh_connect, log_ssh_command
+
+
+logger = logging.getLogger(__name__)
+
+
+def discover_ssh_key() -> Optional[str]:
+    """
+    Discover SSH private key for authentication.
+    
+    Checks in order:
+    1. LINUX_MCP_SSH_KEY_PATH environment variable
+    2. Default locations: ~/.ssh/id_ed25519, ~/.ssh/id_rsa, ~/.ssh/id_ecdsa
+    
+    Returns:
+        Path to SSH private key if found, None otherwise.
+    """
+    logger.debug("Discovering SSH key for authentication")
+    
+    # Check environment variable first
+    env_key = os.getenv("LINUX_MCP_SSH_KEY_PATH")
+    if env_key:
+        logger.debug(f"Checking SSH key from environment: {env_key}")
+        key_path = Path(env_key)
+        if key_path.exists() and key_path.is_file():
+            logger.info(f"Using SSH key from environment: {env_key}")
+            return str(key_path)
+        else:
+            logger.warning(f"SSH key specified in LINUX_MCP_SSH_KEY_PATH not found: {env_key}")
+            return None
+    
+    # Check default locations (prefer modern algorithms)
+    home = Path.home()
+    default_keys = [
+        home / ".ssh" / "id_ed25519",
+        home / ".ssh" / "id_ecdsa",
+        home / ".ssh" / "id_rsa",
+    ]
+    
+    logger.debug(f"Checking default SSH key locations: {[str(k) for k in default_keys]}")
+    
+    for key_path in default_keys:
+        if key_path.exists() and key_path.is_file():
+            logger.info(f"Using SSH key: {key_path}")
+            return str(key_path)
+    
+    logger.warning("No SSH private key found in default locations")
+    return None
+
+
+class SSHConnectionManager:
+    """
+    Manages SSH connections with connection pooling.
+    
+    This class implements a singleton pattern to maintain a pool of SSH connections
+    across the lifetime of the application, improving performance by reusing
+    connections to the same hosts.
+    """
+    
+    _instance = None
+    
+    def __new__(cls):
+        """Implement singleton pattern."""
+        if cls._instance is None:
+            cls._instance = super().__new__(cls)
+            cls._instance._connections = {}
+            cls._instance._ssh_key = discover_ssh_key()
+        return cls._instance
+    
+    async def get_connection(self, host: str, username: str) -> asyncssh.SSHClientConnection:
+        """
+        Get or create an SSH connection to a host.
+        
+        Args:
+            host: Remote host address
+            username: SSH username
+            
+        Returns:
+            SSH connection object
+            
+        Raises:
+            ConnectionError: If connection fails
+        """
+        key = f"{username}@{host}"
+        
+        # Return existing connection if available
+        if key in self._connections:
+            conn = self._connections[key]
+            if not conn.is_closed():
+                # DEBUG level: Log connection reuse and pool state
+                logger.debug(f"SSH_REUSE: {key} | pool_size={len(self._connections)}")
+                # Use audit log with connection reuse info
+                log_ssh_connect(host, username, status="success", reused=True, key_path=self._ssh_key)
+                return conn
+            else:
+                # Connection was closed, remove it
+                logger.debug(f"SSH_POOL: remove_closed_connection | connection={key}")
+                del self._connections[key]
+        
+        # Create new connection
+        # DEBUG level: Log connection attempt before it completes
+        logger.debug(f"SSH_CONNECTING: {key} | key={self._ssh_key or 'none'}")
+        
+        try:
+            connect_kwargs = {
+                "host": host,
+                "username": username,
+                "known_hosts": None,  # Don't verify host keys for now
+            }
+            
+            if self._ssh_key:
+                connect_kwargs["client_keys"] = [self._ssh_key]
+            
+            conn = await asyncssh.connect(**connect_kwargs)
+            self._connections[key] = conn
+            
+            # Log successful connection using audit function
+            log_ssh_connect(host, username, status="success", reused=False, key_path=self._ssh_key)
+            
+            # DEBUG level: Log pool state
+            logger.debug(f"SSH_POOL: add_connection | connections={len(self._connections)}")
+            
+            return conn
+            
+        except asyncssh.PermissionDenied as e:
+            # Use audit log for authentication failure
+            error_msg = str(e)
+            log_ssh_connect(host, username, status="failed", error=f"Permission denied: {error_msg}")
+            raise ConnectionError(f"Authentication failed for {username}@{host}") from e
+        except asyncssh.Error as e:
+            # Use audit log for connection failure
+            error_msg = str(e)
+            log_ssh_connect(host, username, status="failed", error=error_msg)
+            raise ConnectionError(f"Failed to connect to {username}@{host}: {e}") from e
+    
+    async def execute_remote(
+        self,
+        command: list[str],
+        host: str,
+        username: str
+    ) -> Tuple[int, str, str]:
+        """
+        Execute a command on a remote host via SSH.
+        
+        Args:
+            command: Command and arguments to execute
+            host: Remote host address
+            username: SSH username
+            
+        Returns:
+            Tuple of (return_code, stdout, stderr)
+            
+        Raises:
+            ConnectionError: If SSH connection fails
+        """
+        conn = await self.get_connection(host, username)
+        
+        # Build command string with proper shell escaping
+        # Use shlex.quote() to ensure special characters (like \n in printf format) are preserved
+        cmd_str = " ".join(shlex.quote(arg) for arg in command)
+        
+        # Start timing for command execution
+        start_time = time.time()
+        
+        try:
+            result = await conn.run(cmd_str, check=False)
+            
+            return_code = result.exit_status if result.exit_status is not None else 0
+            stdout = result.stdout if result.stdout else ""
+            stderr = result.stderr if result.stderr else ""
+            
+            # Calculate duration
+            duration = time.time() - start_time
+            
+            # Use audit log for command execution
+            log_ssh_command(cmd_str, host, exit_code=return_code, duration=duration)
+            
+            return return_code, stdout, stderr
+            
+        except asyncssh.Error as e:
+            duration = time.time() - start_time
+            logger.error(f"Error executing command on {username}@{host}: {e}", extra={
+                'event': 'REMOTE_EXEC_ERROR',
+                'command': cmd_str,
+                'host': host,
+                'duration': f"{duration:.3f}s",
+                'error': str(e)
+            })
+            raise ConnectionError(f"Failed to execute command on {username}@{host}: {e}") from e
+    
+    async def close_all(self):
+        """Close all SSH connections."""
+        connection_count = len(self._connections)
+        logger.info(f"Closing {connection_count} SSH connections")
+        
+        for key, conn in list(self._connections.items()):
+            try:
+                logger.debug(f"SSH_CLOSE: {key}")
+                conn.close()
+                await conn.wait_closed()
+            except Exception as e:
+                logger.warning(f"Error closing connection to {key}: {e}")
+        
+        self._connections.clear()
+        logger.debug(f"SSH_POOL: cleared | closed_connections={connection_count}")
+
+
+# Global connection manager instance
+_connection_manager = SSHConnectionManager()
+
+
+async def execute_command(
+    command: list[str],
+    host: Optional[str] = None,
+    username: Optional[str] = None,
+    **kwargs
+) -> Tuple[int, str, str]:
+    """
+    Execute a command locally or remotely.
+    
+    This is the main entry point for command execution. It routes the command
+    to either local subprocess execution or remote SSH execution based on
+    whether host/username parameters are provided.
+    
+    Args:
+        command: Command and arguments to execute
+        host: Optional remote host address
+        username: Optional SSH username (required if host is provided)
+        **kwargs: Additional arguments (reserved for future use)
+        
+    Returns:
+        Tuple of (return_code, stdout, stderr)
+        
+    Raises:
+        ValueError: If host is provided without username
+        ConnectionError: If remote connection fails
+        
+    Examples:
+        # Local execution
+        >>> returncode, stdout, stderr = await execute_command(["ls", "-la"])
+        
+        # Remote execution
+        >>> returncode, stdout, stderr = await execute_command(
+        ...     ["ls", "-la"],
+        ...     host="server.example.com",
+        ...     username="admin"
+        ... )
+    """
+    cmd_str = " ".join(command)
+    
+    # Route to remote execution if host is provided
+    if host:
+        if not username:
+            logger.error(f"Host provided without username for command: {cmd_str}")
+            raise ValueError("username is required when host is provided")
+        
+        logger.debug(f"Routing to remote execution: {username}@{host} | command={cmd_str}")
+        return await _connection_manager.execute_remote(command, host, username)
+    
+    # Local execution
+    logger.debug(f"LOCAL_EXEC: {cmd_str}")
+    return await _execute_local(command)
+
+
+async def _execute_local(command: list[str]) -> Tuple[int, str, str]:
+    """
+    Execute a command locally using subprocess.
+    
+    Args:
+        command: Command and arguments to execute
+        
+    Returns:
+        Tuple of (return_code, stdout, stderr)
+    """
+    cmd_str = " ".join(command)
+    start_time = time.time()
+    
+    try:
+        proc = await asyncio.create_subprocess_exec(
+            *command,
+            stdout=subprocess.PIPE,
+            stderr=subprocess.PIPE
+        )
+        stdout_bytes, stderr_bytes = await proc.communicate()
+        
+        return_code = proc.returncode if proc.returncode is not None else 0
+        stdout = stdout_bytes.decode('utf-8', errors='replace')
+        stderr = stderr_bytes.decode('utf-8', errors='replace')
+        
+        duration = time.time() - start_time
+        
+        # DEBUG level: Log local command execution with timing
+        logger.debug(f"LOCAL_EXEC completed: {cmd_str} | exit_code={return_code} | duration={duration:.3f}s")
+        
+        return return_code, stdout, stderr
+        
+    except Exception as e:
+        duration = time.time() - start_time
+        logger.error(f"Error executing local command: {cmd_str}", extra={
+            'event': 'LOCAL_EXEC_ERROR',
+            'command': cmd_str,
+            'duration': f"{duration:.3f}s",
+            'error': str(e)
+        })
+        return 1, "", str(e)
+

linux_mcp_server/tools/storage.py

@@ -0,0 +1,358 @@
+"""Storage and hardware tools."""
+
+import asyncio
+import subprocess
+from pathlib import Path
+from typing import Optional
+import psutil
+
+from .validation import validate_positive_int
+from .ssh_executor import execute_command
+from .utils import format_bytes
+
+
+async def list_block_devices(host: Optional[str] = None, username: Optional[str] = None) -> str:
+    """
+    List block devices.
+    
+    Args:
+        host: Optional remote host to connect to
+        username: Optional SSH username (required if host is provided)
+        
+    Returns:
+        Formatted string with block device information
+    """
+    try:
+        # Try using lsblk first (most readable)
+        returncode, stdout, stderr = await execute_command(
+            ["lsblk", "-o", "NAME,SIZE,TYPE,MOUNTPOINT,FSTYPE,MODEL", "--no-pager"],
+            host=host,
+            username=username
+        )
+        
+        if returncode == 0:
+            result = ["=== Block Devices ===\n"]
+            result.append(stdout)
+            
+            # Add disk I/O per-disk stats if available (only for local execution)
+            if not host:
+                try:
+                    disk_io_per_disk = psutil.disk_io_counters(perdisk=True)
+                    if disk_io_per_disk:
+                        result.append("\n=== Disk I/O Statistics (per disk) ===")
+                        for disk, stats in sorted(disk_io_per_disk.items()):
+                            result.append(f"\n{disk}:")
+                            result.append(f"  Read: {format_bytes(stats.read_bytes)}")
+                            result.append(f"  Write: {format_bytes(stats.write_bytes)}")
+                            result.append(f"  Read Count: {stats.read_count}")
+                            result.append(f"  Write Count: {stats.write_count}")
+                except Exception:
+                    pass
+            
+            return "\n".join(result)
+        else:
+            # Fallback to listing partitions with psutil
+            result = ["=== Block Devices (fallback) ===\n"]
+            partitions = psutil.disk_partitions(all=True)
+            
+            for partition in partitions:
+                result.append(f"\nDevice: {partition.device}")
+                result.append(f"  Mountpoint: {partition.mountpoint}")
+                result.append(f"  Filesystem: {partition.fstype}")
+                result.append(f"  Options: {partition.opts}")
+            
+            return "\n".join(result)
+    except FileNotFoundError:
+        # If lsblk is not available, use psutil
+        result = ["=== Block Devices ===\n"]
+        partitions = psutil.disk_partitions(all=True)
+        
+        for partition in partitions:
+            result.append(f"\nDevice: {partition.device}")
+            result.append(f"  Mountpoint: {partition.mountpoint}")
+            result.append(f"  Filesystem: {partition.fstype}")
+            result.append(f"  Options: {partition.opts}")
+        
+        return "\n".join(result)
+    except Exception as e:
+        return f"Error listing block devices: {str(e)}"
+
+
+async def list_directories_by_size(
+    path: str,
+    top_n: int,
+    host: Optional[str] = None,
+    username: Optional[str] = None
+) -> str:
+    """
+    List directories under a specified path sorted by size (largest first).
+    
+    This function uses efficient Linux primitives (du command) to calculate directory
+    sizes, making it much faster than Python-based directory traversal.
+    
+    Args:
+        path: The directory path to analyze
+        top_n: Number of top directories to return (1-1000). Accepts int or float 
+               (floats are truncated to integers)
+        host: Optional remote host to connect to
+        username: Optional SSH username (required if host is provided)
+    
+    Returns:
+        Formatted string with directory sizes, or error message if validation fails
+    
+    Security Features:
+        - Path validation and resolution using pathlib
+        - Command parameters passed as list (not shell string)
+        - Input sanitization for all parameters
+        - Graceful error handling for permission issues
+    """
+    import os
+    from pathlib import Path
+    
+    try:
+        # Validate and normalize top_n parameter
+        top_n, error = validate_positive_int(
+            top_n,
+            param_name="top_n",
+            min_value=1,
+            max_value=1000
+        )
+        if error:
+            return error
+        
+        # For local execution, validate path
+        if not host:
+            try:
+                path_obj = Path(path).resolve(strict=True)
+            except (OSError, RuntimeError):
+                return f"Error: Path does not exist or cannot be resolved: {path}"
+            
+            if not path_obj.is_dir():
+                return f"Error: Path is not a directory: {path}"
+            
+            if not os.access(path_obj, os.R_OK):
+                return f"Error: Permission denied to read directory: {path}"
+            
+            path_str = str(path_obj)
+        else:
+            # For remote execution, use the path as-is
+            path_str = path
+        
+        # Use du command to get directory sizes efficiently
+        returncode, stdout, _ = await execute_command(
+            ["du", "-b", "--max-depth=1", path_str],
+            host=host,
+            username=username
+        )
+        
+        # Parse output - du may return non-zero on permission errors but still give valid data
+        lines = stdout.strip().split('\n')
+        dir_sizes = []
+        
+        for line in lines:
+            if not line:
+                continue
+            parts = line.split('\t', 1)
+            if len(parts) == 2:
+                try:
+                    size = int(parts[0])
+                    dir_path_str = parts[1]
+                    # Skip the parent directory itself
+                    dir_name = Path(dir_path_str).name
+                    if dir_path_str != path_str:
+                        dir_sizes.append((dir_name, size))
+                except (ValueError, IndexError):
+                    continue
+        
+        if not dir_sizes:
+            # Only error if we got no output AND a bad return code
+            if returncode != 0:
+                return f"Error: du command failed and returned no directory data"
+            return f"No subdirectories found in: {path}"
+        
+        # Sort by size (descending) and take top N
+        dir_sizes.sort(key=lambda x: x[1], reverse=True)
+        top_dirs = dir_sizes[:top_n]
+        
+        # Format output
+        result = []
+        result.append(f"=== Top {len(top_dirs)} Largest Directories ===")
+        result.append(f"Path: {path_str}")
+        result.append(f"\nTotal subdirectories found: {len(dir_sizes)}\n")
+        
+        for i, (dir_name, size) in enumerate(top_dirs, 1):
+            result.append(f"{i}. {dir_name}")
+            result.append(f"   Size: {format_bytes(size)}")
+        
+        return "\n".join(result)
+        
+    except Exception as e:
+        return f"Error analyzing directories: {str(e)}"
+
+
+async def list_directories_by_name(
+    path: str,
+    reverse: bool = False,
+    host: Optional[str] = None,
+    username: Optional[str] = None
+) -> str:
+    """
+    List directories under a specified path sorted by name.
+    
+    This function uses efficient Linux primitives (find and sort) to list directories.
+    
+    Args:
+        path: The directory path to analyze
+        reverse: If True, sort in reverse alphabetical order (Z-A)
+        host: Optional remote host to connect to
+        username: Optional SSH username (required if host is provided)
+    
+    Returns:
+        Formatted string with directory names, or error message if validation fails
+    """
+    import os
+    from pathlib import Path
+    
+    try:
+        # For local execution, validate path
+        if not host:
+            try:
+                path_obj = Path(path).resolve(strict=True)
+            except (OSError, RuntimeError):
+                return f"Error: Path does not exist or cannot be resolved: {path}"
+            
+            if not path_obj.is_dir():
+                return f"Error: Path is not a directory: {path}"
+            
+            if not os.access(path_obj, os.R_OK):
+                return f"Error: Permission denied to read directory: {path}"
+            
+            path_str = str(path_obj)
+        else:
+            # For remote execution, use the path as-is
+            path_str = path
+        
+        # Use find to list only immediate subdirectories
+        returncode, stdout, _ = await execute_command(
+            ["find", path_str, "-mindepth", "1", "-maxdepth", "1", "-type", "d", "-printf", "%f\\n"],
+            host=host,
+            username=username
+        )
+        
+        if returncode != 0:
+            return f"Error running find command: command failed with return code {returncode}"
+        
+        # Parse and sort output
+        directories = [line for line in stdout.strip().split('\n') if line]
+        
+        if not directories:
+            return f"No subdirectories found in: {path}"
+        
+        # Sort alphabetically
+        directories.sort(reverse=reverse)
+        
+        # Format output
+        result = []
+        sort_order = "Reverse Alphabetical" if reverse else "Alphabetical"
+        result.append(f"=== Directories ({sort_order}) ===")
+        result.append(f"Path: {path_str}")
+        result.append(f"\nTotal subdirectories found: {len(directories)}\n")
+        
+        for i, dir_name in enumerate(directories, 1):
+            result.append(f"{i}. {dir_name}")
+        
+        return "\n".join(result)
+        
+    except Exception as e:
+        return f"Error listing directories: {str(e)}"
+
+
+async def list_directories_by_modified_date(
+    path: str,
+    newest_first: bool = True,
+    host: Optional[str] = None,
+    username: Optional[str] = None
+) -> str:
+    """
+    List directories under a specified path sorted by modification date.
+    
+    This function uses efficient Linux primitives (find) to list directories with timestamps.
+    
+    Args:
+        path: The directory path to analyze
+        newest_first: If True, show newest first; if False, show oldest first
+        host: Optional remote host to connect to
+        username: Optional SSH username (required if host is provided)
+    
+    Returns:
+        Formatted string with directory names and dates, or error message if validation fails
+    """
+    import os
+    from pathlib import Path
+    from datetime import datetime
+    
+    try:
+        # For local execution, validate path
+        if not host:
+            try:
+                path_obj = Path(path).resolve(strict=True)
+            except (OSError, RuntimeError):
+                return f"Error: Path does not exist or cannot be resolved: {path}"
+            
+            if not path_obj.is_dir():
+                return f"Error: Path is not a directory: {path}"
+            
+            if not os.access(path_obj, os.R_OK):
+                return f"Error: Permission denied to read directory: {path}"
+            
+            path_str = str(path_obj)
+        else:
+            # For remote execution, use the path as-is
+            path_str = path
+        
+        # Use find with modification time
+        returncode, stdout, _ = await execute_command(
+            ["find", path_str, "-mindepth", "1", "-maxdepth", "1", "-type", "d", "-printf", "%T@\\t%f\\n"],
+            host=host,
+            username=username
+        )
+        
+        if returncode != 0:
+            return f"Error running find command: command failed with return code {returncode}"
+        
+        # Parse output
+        directories = []
+        for line in stdout.strip().split('\n'):
+            if line:
+                parts = line.split('\t', 1)
+                if len(parts) == 2:
+                    try:
+                        timestamp = float(parts[0])
+                        dir_name = parts[1]
+                        directories.append((timestamp, dir_name))
+                    except ValueError:
+                        continue
+        
+        if not directories:
+            return f"No subdirectories found in: {path}"
+        
+        # Sort by timestamp
+        directories.sort(key=lambda x: x[0], reverse=newest_first)
+        
+        # Format output
+        result = []
+        sort_order = "Newest First" if newest_first else "Oldest First"
+        result.append(f"=== Directories ({sort_order}) ===")
+        result.append(f"Path: {path_str}")
+        result.append(f"\nTotal subdirectories found: {len(directories)}\n")
+        
+        for i, (timestamp, dir_name) in enumerate(directories, 1):
+            dt = datetime.fromtimestamp(timestamp)
+            result.append(f"{i}. {dir_name}")
+            result.append(f"   Modified: {dt.strftime('%Y-%m-%d %H:%M:%S')}")
+        
+        return "\n".join(result)
+        
+    except Exception as e:
+        return f"Error listing directories: {str(e)}"
+