style: format code for improved readability and consistency across files

refactor: clean up whitespace and comments in mcp_shell_server module for better maintainability
test: simplify test cases in test_server.py and test_shell_executor.py for clarity and conciseness

Author: tumf

mcp_shell_server/__init__.py

@@ -1,5 +1,6 @@
 """MCP Shell Server Package."""
+
 from .server import main
 
 __version__ = "0.1.0"
-__all__ = ["main"]
+__all__ = ["main"]

mcp_shell_server/server.py

@@ -12,8 +12,10 @@
 
 app = Server("mcp-shell-server")
 
+
 class ExecuteToolHandler:
     """Handler for shell command execution"""
+
     name = "execute"
     description = "Execute a shell command"
 
@@ -31,51 +33,46 @@ def get_tool_description(self) -> Tool:
                     "command": {
                         "type": "array",
                         "items": {"type": "string"},
-                        "description": "Command and its arguments as array"
+                        "description": "Command and its arguments as array",
                     },
                     "stdin": {
                         "type": "string",
-                        "description": "Input to be passed to the command via stdin"
-                    }
+                        "description": "Input to be passed to the command via stdin",
+                    },
                 },
-                "required": ["command"]
-            }
+                "required": ["command"],
+            },
         )
 
     async def run_tool(self, arguments: dict) -> Sequence[TextContent]:
         """Execute the shell command with the given arguments"""
         command = arguments.get("command", [])
         stdin = arguments.get("stdin")
-        
+
         if not command:
             raise ValueError("No command provided")
-        
+
         result = await self.executor.execute(command, stdin)
-        
+
         # Raise error if command execution failed
         if result.get("error"):
             raise RuntimeError(result["error"])
 
         # Convert executor result to TextContent sequence
         content: list[TextContent] = []
-        
+
         if result.get("stdout"):
-            content.append(TextContent(
-                type="text",
-                text=result["stdout"]
-            ))
+            content.append(TextContent(type="text", text=result["stdout"]))
         if result.get("stderr"):
-            content.append(TextContent(
-                type="text",
-                text=result["stderr"]
-            ))
-            
+            content.append(TextContent(type="text", text=result["stderr"]))
+
         return content
 
 
 # Initialize tool handlers
 tool_handler = ExecuteToolHandler()
 
+
 @app.list_tools()
 async def list_tools() -> list[Tool]:
     """List available tools."""
@@ -107,10 +104,8 @@ async def main() -> None:
 
         async with stdio_server() as (read_stream, write_stream):
             await app.run(
-                read_stream,
-                write_stream,
-                app.create_initialization_options()
+                read_stream, write_stream, app.create_initialization_options()
             )
     except Exception as e:
         logger.error(f"Server error: {str(e)}")
-        raise
+        raise

mcp_shell_server/shell_executor.py

@@ -1,14 +1,14 @@
 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):
         """
         Initialize the executor. The allowed commands are read from ALLOW_COMMANDS
@@ -19,7 +19,7 @@ def __init__(self):
     def _get_allowed_commands(self) -> set:
         """
         Get the set of allowed commands from environment variable.
-        
+
         Returns:
             set: Set of allowed command names
         """
@@ -29,10 +29,10 @@ def _get_allowed_commands(self) -> set:
     def _clean_command(self, command: List[str]) -> List[str]:
         """
         Clean command by trimming whitespace from each part.
-        
+
         Args:
             command (List[str]): Original command and its arguments
-            
+
         Returns:
             List[str]: Cleaned command with trimmed whitespace
         """
@@ -41,10 +41,10 @@ def _clean_command(self, command: List[str]) -> List[str]:
     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
         """
@@ -53,7 +53,9 @@ def _validate_command(self, command: List[str]) -> None:
 
         allowed_commands = self._get_allowed_commands()
         if not allowed_commands:
-            raise ValueError("No commands are allowed. Please set ALLOW_COMMANDS environment variable.")
+            raise ValueError(
+                "No commands are allowed. Please set ALLOW_COMMANDS environment variable."
+            )
 
         # Clean and check the first command
         cleaned_cmd = command[0].strip()
@@ -68,16 +70,20 @@ def _validate_command(self, command: List[str]) -> None:
                     raise ValueError(f"Unexpected shell operator: {cleaned_arg}")
                 next_cmd = command[i + 1].strip()
                 if next_cmd not in allowed_commands:
-                    raise ValueError(f"Command not allowed after {cleaned_arg}: {next_cmd}")
+                    raise ValueError(
+                        f"Command not allowed after {cleaned_arg}: {next_cmd}"
+                    )
 
-    async def execute(self, command: List[str], stdin: Optional[str] = None) -> Dict[str, Any]:
+    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.
@@ -89,15 +95,15 @@ async def execute(self, command: List[str], stdin: Optional[str] = None) -> Dict
             cleaned_command = self._clean_command(command)
             if not cleaned_command:
                 raise ValueError("Empty command")
-                
+
             self._validate_command(cleaned_command)
         except ValueError as e:
             return {
                 "error": str(e),
                 "status": 1,
                 "stdout": "",
                 "stderr": str(e),
-                "execution_time": time.time() - start_time
+                "execution_time": time.time() - start_time,
             }
 
         try:
@@ -107,7 +113,7 @@ async def execute(self, command: List[str], stdin: Optional[str] = None) -> Dict
                 stdin=asyncio.subprocess.PIPE if stdin else None,
                 stdout=asyncio.subprocess.PIPE,
                 stderr=asyncio.subprocess.PIPE,
-                env={"PATH": os.environ.get("PATH", "")}
+                env={"PATH": os.environ.get("PATH", "")},
             )
 
             stdin_bytes = stdin.encode() if stdin else None
@@ -118,21 +124,21 @@ async def execute(self, command: List[str], stdin: Optional[str] = None) -> Dict
                 "stdout": stdout.decode() if stdout else "",
                 "stderr": stderr.decode() if stderr else "",
                 "status": process.returncode,
-                "execution_time": time.time() - start_time
+                "execution_time": time.time() - start_time,
             }
         except FileNotFoundError:
             return {
                 "error": f"Command not found: {cleaned_command[0]}",
                 "status": 1,
                 "stdout": "",
                 "stderr": f"Command not found: {cleaned_command[0]}",
-                "execution_time": time.time() - start_time
+                "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
-            }
+                "execution_time": time.time() - start_time,
+            }

tests/conftest.py

@@ -1,6 +1,4 @@
-import pytest
-
 # Configure pytest-asyncio
 def pytest_configure(config):
     """Configure pytest-asyncio defaults"""
-    config.option.asyncio_mode = "strict"
+    config.option.asyncio_mode = "strict"

tests/test_server.py

@@ -1,6 +1,6 @@
 import pytest
 from mcp.types import TextContent, Tool
-from mcp_shell_server.server import call_tool, list_tools, tool_handler
+from mcp_shell_server.server import call_tool, list_tools
 
 
 @pytest.mark.asyncio
@@ -22,9 +22,7 @@ async def test_list_tools():
 async def test_call_tool_valid_command(monkeypatch):
     """Test execution of a valid command"""
     monkeypatch.setenv("ALLOW_COMMANDS", "echo")
-    result = await call_tool("execute", {
-        "command": ["echo", "hello world"]
-    })
+    result = await call_tool("execute", {"command": ["echo", "hello world"]})
     assert len(result) == 1
     assert isinstance(result[0], TextContent)
     assert result[0].type == "text"
@@ -35,10 +33,7 @@ async def test_call_tool_valid_command(monkeypatch):
 async def test_call_tool_with_stdin(monkeypatch):
     """Test command execution with stdin"""
     monkeypatch.setenv("ALLOW_COMMANDS", "cat")
-    result = await call_tool("execute", {
-        "command": ["cat"],
-        "stdin": "test input"
-    })
+    result = await call_tool("execute", {"command": ["cat"], "stdin": "test input"})
     assert len(result) == 1
     assert isinstance(result[0], TextContent)
     assert result[0].type == "text"
@@ -50,9 +45,7 @@ async def test_call_tool_invalid_command(monkeypatch):
     """Test execution of an invalid command"""
     monkeypatch.setenv("ALLOW_COMMANDS", "echo")
     with pytest.raises(RuntimeError) as excinfo:
-        await call_tool("execute", {
-            "command": ["invalid_command"]
-        })
+        await call_tool("execute", {"command": ["invalid_command"]})
     assert "Command not allowed: invalid_command" in str(excinfo.value)
 
 
@@ -76,7 +69,5 @@ async def test_call_tool_invalid_arguments():
 async def test_call_tool_empty_command():
     """Test execution with empty command"""
     with pytest.raises(RuntimeError) as excinfo:
-        await call_tool("execute", {
-            "command": []
-        })
-    assert "No command provided" in str(excinfo.value)
+        await call_tool("execute", {"command": []})
+    assert "No command provided" in str(excinfo.value)

tests/test_shell_executor.py

@@ -27,7 +27,7 @@ async def test_stdin_input(executor, monkeypatch):
 async def test_command_with_space_allowed(executor, monkeypatch):
     monkeypatch.setenv("ALLOW_COMMANDS", "cat")
     result = await executor.execute(["cat "], stdin="hello world")
-    assert result["error"] == None
+    assert result["error"] is None
     assert result["stdout"].strip() == "hello world"
     assert result["status"] == 0