feat(server.py): add detailed docstrings to ShellServer class and its methods for better documentation

feat(shell_executor.py): enhance ShellExecutor with command validation and error handling for secure execution
fix(pyproject.toml): update mcp dependency version to ensure compatibility
chore(tests): configure pytest-asyncio for strict mode and improve test structure with event loop fixture
test(shell_executor.py): add tests for command execution, including validation of shell operators and error handling

Author: tumf

mcp_shell_server/server.py

@@ -1,21 +1,38 @@
+import asyncio
 from mcp.server.stdio import stdio_server
 from .shell_executor import ShellExecutor
 
-
 class ShellServer:
+    """
+    MCP server that executes shell commands in a secure manner.
+    Only commands listed in the ALLOW_COMMANDS environment variable can be executed.
+    """
+    
     def __init__(self):
         self.executor = ShellExecutor()
 
     async def handle(self, args: dict) -> dict:
+        """
+        Handle incoming MCP requests to execute shell commands.
+        
+        Args:
+            args (dict): Arguments containing the command to execute and optional stdin
+            
+        Returns:
+            dict: Execution results including stdout, stderr, status code, and execution time
+        """
         command = args.get("command", [])
         stdin = args.get("stdin")
-
+        
         if not command:
-            return {"error": "No command provided", "status": 1}
-
+            return {
+                "error": "No command provided",
+                "status": 1
+            }
+        
         return await self.executor.execute(command, stdin)
 
-
 def main():
+    """Entry point for the MCP shell server"""
     server = ShellServer()
-    stdio_server(server.handle)
+    stdio_server(server.handle)

mcp_shell_server/shell_executor.py

@@ -1,32 +1,74 @@
 import os
 import time
 import asyncio
+import shlex
 from typing import Dict, List, Optional, Any
 
 class ShellExecutor:
+    """
+    Executes shell commands in a secure manner by validating against a whitelist.
+    """
+    
     def __init__(self):
-        # Allow whitespace in ALLOW_COMMANDS and trim each command
+        """
+        Initialize the executor. The allowed commands are read from ALLOW_COMMANDS
+        environment variable during command validation, not at initialization.
+        """
+        pass
+
+    def _get_allowed_commands(self) -> set:
+        """
+        Get the set of allowed commands from environment variable.
+        
+        Returns:
+            set: Set of allowed command names
+        """
         allow_commands = os.environ.get("ALLOW_COMMANDS", "")
-        self.allowed_commands = set(cmd.strip() for cmd in allow_commands.split(",") if cmd.strip())
+        return {cmd.strip() for cmd in allow_commands.split(",") if cmd.strip()}
 
     def _validate_command(self, command: List[str]) -> None:
+        """
+        Validate if the command is allowed to be executed.
+        
+        Args:
+            command (List[str]): Command and its arguments
+            
+        Raises:
+            ValueError: If the command is empty, not allowed, or contains invalid shell operators
+        """
         if not command:
             raise ValueError("Empty command")
 
-        # Check first command
-        if command[0] not in self.allowed_commands:
+        allowed_commands = self._get_allowed_commands()
+        if not allowed_commands:
+            raise ValueError("No commands are allowed. Please set ALLOW_COMMANDS environment variable.")
+
+        if command[0] not in allowed_commands:
             raise ValueError(f"Command not allowed: {command[0]}")
 
-        # Check for shell operators and subsequent commands
-        for arg in command[1:]:
+        # Check for shell operators and validate subsequent commands
+        for i, arg in enumerate(command[1:], start=1):
             if arg in [";", "&&", "||", "|"]:
-                next_cmd_idx = command.index(arg) + 1
-                if next_cmd_idx < len(command):
-                    next_cmd = command[next_cmd_idx]
-                    if next_cmd not in self.allowed_commands:
-                        raise ValueError(f"Command not allowed: {next_cmd}")
+                if i + 1 >= len(command):
+                    raise ValueError(f"Unexpected shell operator: {arg}")
+                next_cmd = command[i + 1]
+                if next_cmd not in allowed_commands:
+                    raise ValueError(f"Command not allowed after {arg}: {next_cmd}")
 
     async def execute(self, command: List[str], stdin: Optional[str] = None) -> Dict[str, Any]:
+        """
+        Execute a shell command with optional stdin input.
+        
+        Args:
+            command (List[str]): Command and its arguments
+            stdin (Optional[str]): Input to be passed to the command via stdin
+            
+        Returns:
+            Dict[str, Any]: Execution result containing stdout, stderr, status code, and execution time.
+                           If error occurs, result contains additional 'error' field.
+        """
+        start_time = time.time()
+
         try:
             self._validate_command(command)
         except ValueError as e:
@@ -35,28 +77,41 @@ async def execute(self, command: List[str], stdin: Optional[str] = None) -> Dict
                 "status": 1,
                 "stdout": "",
                 "stderr": str(e),
-                "execution_time": 0
+                "execution_time": time.time() - start_time
             }
 
-        start_time = time.time()
-        
-        process = await asyncio.create_subprocess_exec(
-            *command,
-            stdin=asyncio.subprocess.PIPE if stdin else None,
-            stdout=asyncio.subprocess.PIPE,
-            stderr=asyncio.subprocess.PIPE
-        )
-
-        if stdin:
-            stdout, stderr = await process.communicate(stdin.encode())
-        else:
-            stdout, stderr = await process.communicate()
-
-        execution_time = time.time() - start_time
-
-        return {
-            "stdout": stdout.decode() if stdout else "",
-            "stderr": stderr.decode() if stderr else "",
-            "status": process.returncode,
-            "execution_time": execution_time
-        }
+        try:
+            process = await asyncio.create_subprocess_exec(
+                command[0],
+                *command[1:],
+                stdin=asyncio.subprocess.PIPE if stdin else None,
+                stdout=asyncio.subprocess.PIPE,
+                stderr=asyncio.subprocess.PIPE,
+                env={"PATH": os.environ.get("PATH", "")}
+            )
+
+            stdin_bytes = stdin.encode() if stdin else None
+            stdout, stderr = await process.communicate(input=stdin_bytes)
+
+            return {
+                "stdout": stdout.decode() if stdout else "",
+                "stderr": stderr.decode() if stderr else "",
+                "status": process.returncode,
+                "execution_time": time.time() - start_time
+            }
+        except FileNotFoundError:
+            return {
+                "error": f"Command not found: {command[0]}",
+                "status": 1,
+                "stdout": "",
+                "stderr": f"Command not found: {command[0]}",
+                "execution_time": time.time() - start_time
+            }
+        except Exception as e:
+            return {
+                "error": str(e),
+                "status": 1,
+                "stdout": "",
+                "stderr": str(e),
+                "execution_time": time.time() - start_time
+            }

pyproject.toml

@@ -6,7 +6,7 @@ authors = [
     { name = "tumf" }
 ]
 dependencies = [
-    "mcp>=1.1.1",
+    "mcp>=1.1.0",
 ]
 requires-python = ">=3.11"
 readme = "README.md"
@@ -25,3 +25,7 @@ test = [
 [build-system]
 requires = ["hatchling"]
 build-backend = "hatchling.build"
+
+[tool.pytest.ini_options]
+asyncio_mode = "strict"
+testpaths = "tests"

tests/conftest.py

@@ -1,3 +1,13 @@
 import pytest
+import asyncio
 
-pytest_plugins = ('pytest_asyncio',)
+# Configure pytest-asyncio to use function scope
+def pytest_configure(config):
+    config.option.asyncio_mode = "strict"
+
+@pytest.fixture(scope="function")
+def event_loop():
+    """Create a new event loop for each test case"""
+    loop = asyncio.new_event_loop()
+    yield loop
+    loop.close()

tests/test_shell_executor.py

@@ -1,4 +1,3 @@
-import os
 import pytest
 from mcp_shell_server.shell_executor import ShellExecutor
 
@@ -8,12 +7,10 @@ def executor():
 
 @pytest.mark.asyncio
 async def test_basic_command_execution(executor, monkeypatch):
-    monkeypatch.setenv("ALLOW_COMMANDS", "echo,ls")
+    monkeypatch.setenv("ALLOW_COMMANDS", "echo")
     result = await executor.execute(["echo", "hello"])
     assert result["stdout"].strip() == "hello"
     assert result["status"] == 0
-    assert result["stderr"] == ""
-    assert "execution_time" in result
 
 @pytest.mark.asyncio
 async def test_stdin_input(executor, monkeypatch):
@@ -45,13 +42,15 @@ async def test_command_with_space_in_allow_commands(executor, monkeypatch):
 @pytest.mark.asyncio
 async def test_multiple_commands_with_operator(executor, monkeypatch):
     monkeypatch.setenv("ALLOW_COMMANDS", "echo,ls")
-    result = await executor.execute(["echo", "hello", ";", "ls", "-l"])
-    assert "Command not allowed: ls" in result["error"]
+    result = await executor.execute(["echo", "hello", ";"])
+    assert result["error"] == "Unexpected shell operator: ;"
     assert result["status"] == 1
 
 @pytest.mark.asyncio
-async def test_command_with_error_output(executor, monkeypatch):
-    monkeypatch.setenv("ALLOW_COMMANDS", "ls")
-    result = await executor.execute(["ls", "/nonexistent"])
-    assert result["stderr"] != ""
-    assert result["status"] != 0
+async def test_shell_operators_not_allowed(executor, monkeypatch):
+    monkeypatch.setenv("ALLOW_COMMANDS", "echo,ls")
+    operators = [";", "&&", "||", "|"]
+    for op in operators:
+        result = await executor.execute(["echo", "hello", op])
+        assert result["error"] == f"Unexpected shell operator: {op}"
+        assert result["status"] == 1