Merge branch 'develop'

Author: tumf

Makefile

@@ -1,12 +1,16 @@
 .PHONY: test format lint typecheck check install-pre-commit
 .DEFAULT_GOAL := all
 
+install:
+	uv pip install --upgrade '.[dev,test]'
+
 test:
 	uv run pytest
 
 format:
-	uv run black .
 	uv run isort .
+	uv run black .
+	uv run mypy --install-types --non-interactive src/mcp_shell_server tests
 	uv run ruff check --fix .
 
 

README.md

@@ -5,11 +5,12 @@
 
 [![MseeP.ai Security Assessment Badge](https://mseep.net/pr/tumf-mcp-shell-server-badge.png)](https://mseep.ai/app/tumf-mcp-shell-server)
 
-
 A secure shell command execution server implementing the Model Context Protocol (MCP). This server allows remote execution of whitelisted shell commands with support for stdin input.
 
 <a href="https://glama.ai/mcp/servers/rt2d4pbn22"><img width="380" height="200" src="https://glama.ai/mcp/servers/rt2d4pbn22/badge" alt="mcp-shell-server MCP server" /></a>
 
+<a href="https://glama.ai/mcp/servers/rt2d4pbn22"><img width="380" height="200" src="https://glama.ai/mcp/servers/rt2d4pbn22/badge" alt="mcp-shell-server MCP server" /></a>
+
 ## Features
 
 * **Secure Command Execution**: Only whitelisted commands can be executed
@@ -80,6 +81,17 @@ npx -y @smithery/cli install mcp-shell-server --client claude
 ```
 
 ### Manual Installation
+
+### Installing via Smithery
+
+To install Shell Server for Claude Desktop automatically via [Smithery](https://smithery.ai/server/mcp-shell-server):
+
+```bash
+npx -y @smithery/cli install mcp-shell-server --client claude
+```
+
+### Manual Installation
+
 ```bash
 pip install mcp-shell-server
 ```

pyproject.toml

@@ -73,6 +73,12 @@ target-version = ['py311']
 profile = "black" 
 line_length = 88
 
+[tool.mypy]
+error_summary = false
+hide_error_codes = true
+disallow_untyped_defs = false
+check_untyped_defs = false
+
 [tool.hatch.version]
 path = "src/mcp_shell_server/version.py"
 

src/mcp_shell_server/command_validator.py

@@ -3,6 +3,7 @@
 """
 
 import os
+import re
 from typing import Dict, List
 
 
@@ -15,7 +16,8 @@ def __init__(self):
         """
         Initialize the validator.
         """
-        pass
+        # No state; environment variables are read on demand
+        return None
 
     def _get_allowed_commands(self) -> set[str]:
         """Get the set of allowed commands from environment variables"""
@@ -24,14 +26,27 @@ def _get_allowed_commands(self) -> set[str]:
         commands = allow_commands + "," + allowed_commands
         return {cmd.strip() for cmd in commands.split(",") if cmd.strip()}
 
+    def _get_allowed_patterns(self) -> List[re.Pattern]:
+        """Get the list of allowed regex patterns from environment variables"""
+        allow_patterns = os.environ.get("ALLOW_PATTERNS", "")
+        patterns = [
+            pattern.strip() for pattern in allow_patterns.split(",") if pattern.strip()
+        ]
+        return [re.compile(pattern) for pattern in patterns]
+
     def get_allowed_commands(self) -> list[str]:
-        """Get the list of allowed commands from environment variables"""
+        """Public API: return list form of allowed commands"""
         return list(self._get_allowed_commands())
 
     def is_command_allowed(self, command: str) -> bool:
-        """Check if a command is in the allowed list"""
+        """Check if a command is in the allowed list or matches an allowed pattern"""
         cmd = command.strip()
-        return cmd in self._get_allowed_commands()
+        if cmd in self._get_allowed_commands():
+            return True
+        for pattern in self._get_allowed_patterns():
+            if pattern.match(cmd):
+                return True
+        return False
 
     def validate_no_shell_operators(self, cmd: str) -> None:
         """
@@ -92,13 +107,12 @@ def validate_command(self, command: List[str]) -> None:
         if not command:
             raise ValueError("Empty command")
 
-        allowed_commands = self._get_allowed_commands()
-        if not allowed_commands:
+        if not self._get_allowed_commands() and not self._get_allowed_patterns():
             raise ValueError(
                 "No commands are allowed. Please set ALLOW_COMMANDS environment variable."
             )
 
         # Clean and check the first command
         cleaned_cmd = command[0].strip()
-        if cleaned_cmd not in allowed_commands:
+        if not self.is_command_allowed(cleaned_cmd):
             raise ValueError(f"Command not allowed: {cleaned_cmd}")

src/mcp_shell_server/process_manager.py

@@ -87,20 +87,31 @@ async def start_process(
         return process
 
     async def cleanup_processes(
-        self, processes: List[asyncio.subprocess.Process]
+        self, processes: Optional[List[asyncio.subprocess.Process]] = None
     ) -> None:
         """Clean up processes by killing them if they're still running.
 
         Args:
-            processes: List of processes to clean up
+            processes: Optional list of processes to clean up. If None, clean up all tracked processes
         """
+        if processes is None:
+            processes = list(self._processes)
+
         cleanup_tasks = []
         for process in processes:
             if process.returncode is None:
                 try:
-                    # Force kill immediately as required by tests
-                    process.kill()
-                    cleanup_tasks.append(asyncio.create_task(process.wait()))
+                    # First attempt graceful termination
+                    process.terminate()
+                    try:
+                        await asyncio.wait_for(process.wait(), timeout=0.5)
+                    except asyncio.TimeoutError:
+                        # Force kill if termination didn't work
+                        process.kill()
+                        cleanup_tasks.append(asyncio.create_task(process.wait()))
+                except ProcessLookupError:
+                    # Process already terminated
+                    pass
                 except Exception as e:
                     logging.warning(f"Error killing process: {e}")
 

src/mcp_shell_server/server.py

@@ -1,5 +1,6 @@
 import asyncio
 import logging
+import signal
 import traceback
 from collections.abc import Sequence
 from typing import Any
@@ -30,11 +31,21 @@ def get_allowed_commands(self) -> list[str]:
         """Get the allowed commands"""
         return self.executor.validator.get_allowed_commands()
 
+    def get_allowed_patterns(self) -> list[str]:
+        """Get the allowed regex patterns"""
+        return [
+            pattern.pattern
+            for pattern in self.executor.validator._get_allowed_patterns()
+        ]
+
     def get_tool_description(self) -> Tool:
+        """Get the tool description for the execute command"""
+        allowed_commands = ", ".join(self.get_allowed_commands())
+        allowed_patterns = ", ".join(self.get_allowed_patterns())
         """Get the tool description for the execute command"""
         return Tool(
             name=self.name,
-            description=f"{self.description}\nAllowed commands: {', '.join(self.get_allowed_commands())}",
+            description=f"{self.description}\nAllowed commands: {allowed_commands}\nAllowed patterns: {allowed_patterns}",
             inputSchema={
                 "type": "object",
                 "properties": {
@@ -142,13 +153,64 @@ async def call_tool(name: str, arguments: Any) -> Sequence[TextContent]:
 async def main() -> None:
     """Main entry point for the MCP shell server"""
     logger.info(f"Starting MCP shell server v{__version__}")
+
+    # Setup signal handling
+    loop = asyncio.get_running_loop()
+    stop_event = asyncio.Event()
+
+    def handle_signal():
+        if not stop_event.is_set():  # Prevent duplicate handling
+            logger.info("Received shutdown signal, starting cleanup...")
+            stop_event.set()
+
+    # Register signal handlers
+    for sig in (signal.SIGTERM, signal.SIGINT):
+        loop.add_signal_handler(sig, handle_signal)
+
     try:
         from mcp.server.stdio import stdio_server
 
         async with stdio_server() as (read_stream, write_stream):
-            await app.run(
-                read_stream, write_stream, app.create_initialization_options()
+            # Run the server until stop_event is set
+            server_task = asyncio.create_task(
+                app.run(read_stream, write_stream, app.create_initialization_options())
             )
+
+            # Create task for stop event
+            stop_task = asyncio.create_task(stop_event.wait())
+
+            # Wait for either server completion or stop signal
+            done, pending = await asyncio.wait(
+                [server_task, stop_task], return_when=asyncio.FIRST_COMPLETED
+            )
+
+            # Check for exceptions in completed tasks
+            for task in done:
+                try:
+                    await task
+                except Exception:
+                    raise  # Re-raise the exception
+
+            # Cancel any pending tasks
+            for task in pending:
+                task.cancel()
+                try:
+                    await task
+                except asyncio.CancelledError:
+                    pass
+
     except Exception as e:
         logger.error(f"Server error: {str(e)}")
         raise
+    finally:
+        # Cleanup signal handlers
+        for sig in (signal.SIGTERM, signal.SIGINT):
+            loop.remove_signal_handler(sig)
+
+        # Ensure all processes are terminated
+        if hasattr(tool_handler, "executor") and hasattr(
+            tool_handler.executor, "process_manager"
+        ):
+            await tool_handler.executor.process_manager.cleanup_processes()
+
+        logger.info("Server shutdown complete")

tests/test_command_validator.py

@@ -22,7 +22,15 @@ def test_get_allowed_commands(validator, monkeypatch):
     assert set(validator.get_allowed_commands()) == {"cmd1", "cmd2", "cmd3", "cmd4"}
 
 
-def test_is_command_allowed(validator, monkeypatch):
+def test_is_command_allowed_with_patterns(validator, monkeypatch):
+    clear_env(monkeypatch)
+    monkeypatch.setenv("ALLOW_COMMANDS", "allowed_cmd")
+    monkeypatch.setenv("ALLOW_PATTERNS", "^cmd[0-9]+$")
+
+    assert validator.is_command_allowed("allowed_cmd")
+    assert validator.is_command_allowed("cmd123")
+    assert not validator.is_command_allowed("disallowed_cmd")
+    assert not validator.is_command_allowed("cmdabc")
     clear_env(monkeypatch)
     monkeypatch.setenv("ALLOW_COMMANDS", "allowed_cmd")
     assert validator.is_command_allowed("allowed_cmd")

tests/test_process_manager.py

@@ -142,17 +142,18 @@ async def test_cleanup_processes(process_manager):
     # Create mock processes with different states
     running_proc = create_mock_process()
     running_proc.returncode = None
+    # Mock wait to simulate timeout
+    running_proc.wait.side_effect = [asyncio.TimeoutError(), None]
 
     completed_proc = create_mock_process()
     completed_proc.returncode = 0
 
     # Execute cleanup
     await process_manager.cleanup_processes([running_proc, completed_proc])
 
-    # Verify running process was killed and waited for
-    running_proc.kill.assert_called_once()
-    running_proc.wait.assert_awaited_once()
-
+    # Verify running process was terminated first, then killed
+    running_proc.terminate.assert_called_once()
+    assert running_proc.wait.await_count == 2  # wait called for both terminate and kill
     # Verify completed process was not killed or waited for
     completed_proc.kill.assert_not_called()
     completed_proc.wait.assert_not_called()

tests/test_server.py

@@ -1,5 +1,6 @@
 import asyncio
 import os
+import signal
 import tempfile
 
 import pytest
@@ -434,3 +435,68 @@ async def test_environment_variables(monkeypatch, temp_test_dir):
         {"command": ["env"], "directory": temp_test_dir},
     )
     assert len(result) == 1
+
+
+@pytest.mark.asyncio
+async def test_signal_handling(monkeypatch, mocker):
+    """Test signal handling and cleanup during server shutdown"""
+    from mcp_shell_server.server import main
+
+    # Setup mocks
+    mock_read_stream = mocker.AsyncMock()
+    mock_write_stream = mocker.AsyncMock()
+    mock_cleanup_processes = mocker.AsyncMock()
+
+    # Mock process manager
+    class MockExecutor:
+        def __init__(self):
+            self.process_manager = mocker.MagicMock()
+            self.process_manager.cleanup_processes = mock_cleanup_processes
+
+    class MockToolHandler:
+        def __init__(self):
+            self.executor = MockExecutor()
+
+    # Setup server mocks
+    context_manager = mocker.AsyncMock()
+    context_manager.__aenter__ = mocker.AsyncMock(
+        return_value=(mock_read_stream, mock_write_stream)
+    )
+    context_manager.__aexit__ = mocker.AsyncMock()
+    mock_stdio_server = mocker.Mock(return_value=context_manager)
+    mocker.patch("mcp.server.stdio.stdio_server", mock_stdio_server)
+
+    # Mock server run to simulate long-running task
+    async def mock_run(*args):
+        # Wait indefinitely or until cancelled
+        try:
+            await asyncio.sleep(10)
+        except asyncio.CancelledError:
+            pass
+
+    mocker.patch("mcp_shell_server.server.app.run", side_effect=mock_run)
+
+    # Mock tool handler
+    tool_handler = MockToolHandler()
+    mocker.patch("mcp_shell_server.server.tool_handler", tool_handler)
+
+    # Run main in a task so we can simulate signal
+    task = asyncio.create_task(main())
+
+    # Give the server a moment to start
+    await asyncio.sleep(0.1)
+
+    # Simulate SIGINT
+    loop = asyncio.get_running_loop()
+    loop.call_soon(lambda: [h() for h in loop._signal_handlers.get(signal.SIGINT, [])])
+
+    # Wait for main to complete
+    try:
+        await asyncio.wait_for(task, timeout=1.0)
+    except asyncio.TimeoutError:
+        task.cancel()
+        await asyncio.sleep(0.1)
+
+    # Verify cleanup was called
+    mock_cleanup_processes.assert_called_once()
+    context_manager.__aexit__.assert_called_once()