feat(shell_executor.py): add _clean_command method to trim whitespace from command arguments for improved command validation

fix(shell_executor.py): update command validation to use cleaned command and handle empty command case
test(shell_executor.py): add tests for command execution with leading/trailing spaces and empty command handling

Author: tumf

mcp_shell_server/shell_executor.py

@@ -26,6 +26,18 @@ def _get_allowed_commands(self) -> set:
         allow_commands = os.environ.get("ALLOW_COMMANDS", "")
         return {cmd.strip() for cmd in allow_commands.split(",") if cmd.strip()}
 
+    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
+        """
+        return [arg.strip() for arg in command if arg.strip()]
+
     def _validate_command(self, command: List[str]) -> None:
         """
         Validate if the command is allowed to be executed.
@@ -43,17 +55,20 @@ def _validate_command(self, command: List[str]) -> None:
         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]}")
+        # Clean and check the first command
+        cleaned_cmd = command[0].strip()
+        if cleaned_cmd not in allowed_commands:
+            raise ValueError(f"Command not allowed: {cleaned_cmd}")
 
         # Check for shell operators and validate subsequent commands
         for i, arg in enumerate(command[1:], start=1):
-            if arg in [";", "&&", "||", "|"]:
+            cleaned_arg = arg.strip()
+            if cleaned_arg in [";", "&&", "||", "|"]:
                 if i + 1 >= len(command):
-                    raise ValueError(f"Unexpected shell operator: {arg}")
-                next_cmd = command[i + 1]
+                    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 {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]:
         """
@@ -70,7 +85,12 @@ async def execute(self, command: List[str], stdin: Optional[str] = None) -> Dict
         start_time = time.time()
 
         try:
-            self._validate_command(command)
+            # Clean command before validation and execution
+            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),
@@ -82,8 +102,8 @@ async def execute(self, command: List[str], stdin: Optional[str] = None) -> Dict
 
         try:
             process = await asyncio.create_subprocess_exec(
-                command[0],
-                *command[1:],
+                cleaned_command[0],
+                *cleaned_command[1:],
                 stdin=asyncio.subprocess.PIPE if stdin else None,
                 stdout=asyncio.subprocess.PIPE,
                 stderr=asyncio.subprocess.PIPE,
@@ -94,17 +114,18 @@ async def execute(self, command: List[str], stdin: Optional[str] = None) -> Dict
             stdout, stderr = await process.communicate(input=stdin_bytes)
 
             return {
+                "error": None,  # Set error field to None for success case
                 "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]}",
+                "error": f"Command not found: {cleaned_command[0]}",
                 "status": 1,
                 "stdout": "",
-                "stderr": f"Command not found: {command[0]}",
+                "stderr": f"Command not found: {cleaned_command[0]}",
                 "execution_time": time.time() - start_time
             }
         except Exception as e:

tests/test_shell_executor.py

@@ -1,56 +1,73 @@
 import pytest
 from mcp_shell_server.shell_executor import ShellExecutor
 
+
 @pytest.fixture
 def executor():
     return ShellExecutor()
 
+
 @pytest.mark.asyncio
 async def test_basic_command_execution(executor, monkeypatch):
     monkeypatch.setenv("ALLOW_COMMANDS", "echo")
     result = await executor.execute(["echo", "hello"])
     assert result["stdout"].strip() == "hello"
     assert result["status"] == 0
 
+
 @pytest.mark.asyncio
 async def test_stdin_input(executor, monkeypatch):
     monkeypatch.setenv("ALLOW_COMMANDS", "cat")
     result = await executor.execute(["cat"], stdin="hello world")
     assert result["stdout"].strip() == "hello world"
     assert result["status"] == 0
 
+
+@pytest.mark.asyncio
+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["stdout"].strip() == "hello world"
+    assert result["status"] == 0
+
+
 @pytest.mark.asyncio
 async def test_command_not_allowed(executor, monkeypatch):
     monkeypatch.setenv("ALLOW_COMMANDS", "ls")
     result = await executor.execute(["rm", "-rf", "/"])
     assert result["error"] == "Command not allowed: rm"
     assert result["status"] == 1
 
+
 @pytest.mark.asyncio
 async def test_empty_command(executor):
     result = await executor.execute([])
     assert result["error"] == "Empty command"
     assert result["status"] == 1
 
+
 @pytest.mark.asyncio
 async def test_command_with_space_in_allow_commands(executor, monkeypatch):
     monkeypatch.setenv("ALLOW_COMMANDS", "ls, echo ,cat")
     result = await executor.execute(["echo", "test"])
     assert result["stdout"].strip() == "test"
     assert result["status"] == 0
 
+
 @pytest.mark.asyncio
 async def test_multiple_commands_with_operator(executor, monkeypatch):
     monkeypatch.setenv("ALLOW_COMMANDS", "echo,ls")
     result = await executor.execute(["echo", "hello", ";"])
     assert result["error"] == "Unexpected shell operator: ;"
     assert result["status"] == 1
 
+
 @pytest.mark.asyncio
 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
+        assert result["status"] == 1