refactor: extract command validation logic to CommandValidator class

- Create CommandValidator class to handle command validation
- Move validation methods from ShellExecutor to CommandValidator
- Update tests and maintain 91% code coverage

Author: tumf

src/mcp_shell_server/command_validator.py

@@ -0,0 +1,104 @@
+"""
+Provides validation for shell commands and ensures they are allowed to be executed.
+"""
+
+import os
+from typing import Dict, List
+
+
+class CommandValidator:
+    """
+    Validates shell commands against a whitelist and checks for unsafe operators.
+    """
+
+    def __init__(self):
+        """
+        Initialize the validator.
+        """
+        pass
+
+    def _get_allowed_commands(self) -> set[str]:
+        """Get the set of allowed commands from environment variables"""
+        allow_commands = os.environ.get("ALLOW_COMMANDS", "")
+        allowed_commands = os.environ.get("ALLOWED_COMMANDS", "")
+        commands = allow_commands + "," + allowed_commands
+        return {cmd.strip() for cmd in commands.split(",") if cmd.strip()}
+
+    def get_allowed_commands(self) -> list[str]:
+        """Get the list of allowed commands from environment variables"""
+        return list(self._get_allowed_commands())
+
+    def is_command_allowed(self, command: str) -> bool:
+        """Check if a command is in the allowed list"""
+        cmd = command.strip()
+        return cmd in self._get_allowed_commands()
+
+    def validate_no_shell_operators(self, cmd: str) -> None:
+        """
+        Validate that the command does not contain shell operators.
+
+        Args:
+            cmd (str): Command to validate
+
+        Raises:
+            ValueError: If the command contains shell operators
+        """
+        if cmd in [";", "&&", "||", "|"]:
+            raise ValueError(f"Unexpected shell operator: {cmd}")
+
+    def validate_pipeline(self, commands: List[str]) -> Dict[str, str]:
+        """
+        Validate pipeline command and ensure all parts are allowed.
+
+        Args:
+            commands (List[str]): List of commands to validate
+
+        Returns:
+            Dict[str, str]: Error message if validation fails, empty dict if success
+
+        Raises:
+            ValueError: If validation fails
+        """
+        current_cmd: List[str] = []
+
+        for token in commands:
+            if token == "|":
+                if not current_cmd:
+                    raise ValueError("Empty command before pipe operator")
+                if not self.is_command_allowed(current_cmd[0]):
+                    raise ValueError(f"Command not allowed: {current_cmd[0]}")
+                current_cmd = []
+            elif token in [";", "&&", "||"]:
+                raise ValueError(f"Unexpected shell operator in pipeline: {token}")
+            else:
+                current_cmd.append(token)
+
+        if current_cmd:
+            if not self.is_command_allowed(current_cmd[0]):
+                raise ValueError(f"Command not allowed: {current_cmd[0]}")
+
+        return {}
+
+    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")
+
+        allowed_commands = self._get_allowed_commands()
+        if not allowed_commands:
+            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:
+            raise ValueError(f"Command not allowed: {cleaned_cmd}")

src/mcp_shell_server/server.py

@@ -28,7 +28,7 @@ def __init__(self):
 
     def get_allowed_commands(self) -> list[str]:
         """Get the allowed commands"""
-        return self.executor.get_allowed_commands()
+        return self.executor.validator.get_allowed_commands()
 
     def get_tool_description(self) -> Tool:
         """Get the tool description for the execute command"""

src/mcp_shell_server/shell_executor.py

@@ -5,6 +5,8 @@
 import time
 from typing import IO, Any, Dict, List, Optional, Tuple, Union
 
+from mcp_shell_server.command_validator import CommandValidator
+
 
 class ShellExecutor:
     """
@@ -13,25 +15,9 @@ class ShellExecutor:
 
     def __init__(self):
         """
-        Initialize the executor.
+        Initialize the executor with a command validator.
         """
-        pass  # pragma: no cover
-
-    def _get_allowed_commands(self) -> set[str]:
-        """Get the set of allowed commands from environment variables"""
-        allow_commands = os.environ.get("ALLOW_COMMANDS", "")
-        allowed_commands = os.environ.get("ALLOWED_COMMANDS", "")
-        commands = allow_commands + "," + allowed_commands
-        return {cmd.strip() for cmd in commands.split(",") if cmd.strip()}
-
-    def get_allowed_commands(self) -> list[str]:
-        """Get the list of allowed commands from environment variables"""
-        return list(self._get_allowed_commands())
-
-    def is_command_allowed(self, command: str) -> bool:
-        """Check if a command is in the allowed list"""
-        cmd = command.strip()
-        return cmd in self._get_allowed_commands()
+        self.validator = CommandValidator()
 
     def _validate_redirection_syntax(self, command: List[str]) -> None:
         """
@@ -223,16 +209,7 @@ 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:
-            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:
-            raise ValueError(f"Command not allowed: {cleaned_cmd}")
+        self.validator.validate_command(command)
 
     def _validate_directory(self, directory: Optional[str]) -> None:
         """
@@ -259,34 +236,15 @@ def _validate_directory(self, directory: Optional[str]) -> None:
 
     def _validate_no_shell_operators(self, cmd: str) -> None:
         """Validate that the command does not contain shell operators"""
-        if cmd in [";" "&&", "||", "|"]:
-            raise ValueError(f"Unexpected shell operator: {cmd}")
+        self.validator.validate_no_shell_operators(cmd)
 
     def _validate_pipeline(self, commands: List[str]) -> Dict[str, str]:
         """Validate pipeline command and ensure all parts are allowed
 
         Returns:
             Dict[str, str]: Error message if validation fails, empty dict if success
         """
-        current_cmd: List[str] = []
-
-        for token in commands:
-            if token == "|":
-                if not current_cmd:
-                    raise ValueError("Empty command before pipe operator")
-                if not self.is_command_allowed(current_cmd[0]):
-                    raise ValueError(f"Command not allowed: {current_cmd[0]}")
-                current_cmd = []
-            elif token in [";", "&&", "||"]:
-                raise ValueError(f"Unexpected shell operator in pipeline: {token}")
-            else:
-                current_cmd.append(token)
-
-        if current_cmd:
-            if not self.is_command_allowed(current_cmd[0]):
-                raise ValueError(f"Command not allowed: {current_cmd[0]}")
-
-        return {}
+        return self.validator.validate_pipeline(commands)
 
     def _split_pipe_commands(self, command: List[str]) -> List[List[str]]:
         """
@@ -425,13 +383,15 @@ async def execute(
             # First check for pipe operators and handle pipeline
             if "|" in cleaned_command:
                 try:
-                    # Validate pipeline first
-                    error = self._validate_pipeline(cleaned_command)
-                    if error:
+                    # Validate pipeline first using the validator
+                    try:
+                        self.validator.validate_pipeline(cleaned_command)
+                    except ValueError as e:
                         return {
-                            **error,
+                            "error": str(e),
                             "status": 1,
                             "stdout": "",
+                            "stderr": str(e),
                             "execution_time": time.time() - start_time,
                         }
 
@@ -464,18 +424,20 @@ async def execute(
 
             # Then check for other shell operators
             for token in cleaned_command:
-                if token in [";", "&&", "||"]:
+                try:
+                    self.validator.validate_no_shell_operators(token)
+                except ValueError as e:
                     return {
-                        "error": f"Unexpected shell operator: {token}",
+                        "error": str(e),
                         "status": 1,
                         "stdout": "",
-                        "stderr": f"Unexpected shell operator: {token}",
+                        "stderr": str(e),
                         "execution_time": time.time() - start_time,
                     }
 
             # Single command execution
             cmd, redirects = self._parse_command(cleaned_command)
-            self._validate_command(cmd)
+            self.validator.validate_command(cmd)
 
             # Directory validation
             if directory:

tests/test_command_validator.py

@@ -0,0 +1,74 @@
+"""Test cases for the CommandValidator class."""
+
+import pytest
+
+from mcp_shell_server.command_validator import CommandValidator
+
+
+def clear_env(monkeypatch):
+    monkeypatch.delenv("ALLOW_COMMANDS", raising=False)
+    monkeypatch.delenv("ALLOWED_COMMANDS", raising=False)
+
+
+@pytest.fixture
+def validator():
+    return CommandValidator()
+
+
+def test_get_allowed_commands(validator, monkeypatch):
+    clear_env(monkeypatch)
+    monkeypatch.setenv("ALLOW_COMMANDS", "cmd1,cmd2")
+    monkeypatch.setenv("ALLOWED_COMMANDS", "cmd3,cmd4")
+    assert set(validator.get_allowed_commands()) == {"cmd1", "cmd2", "cmd3", "cmd4"}
+
+
+def test_is_command_allowed(validator, monkeypatch):
+    clear_env(monkeypatch)
+    monkeypatch.setenv("ALLOW_COMMANDS", "allowed_cmd")
+    assert validator.is_command_allowed("allowed_cmd")
+    assert not validator.is_command_allowed("disallowed_cmd")
+
+
+def test_validate_no_shell_operators(validator):
+    validator.validate_no_shell_operators("echo")  # Should not raise
+    with pytest.raises(ValueError, match="Unexpected shell operator"):
+        validator.validate_no_shell_operators(";")
+    with pytest.raises(ValueError, match="Unexpected shell operator"):
+        validator.validate_no_shell_operators("&&")
+
+
+def test_validate_pipeline(validator, monkeypatch):
+    clear_env(monkeypatch)
+    monkeypatch.setenv("ALLOW_COMMANDS", "ls,grep")
+
+    # Valid pipeline
+    validator.validate_pipeline(["ls", "|", "grep", "test"])
+
+    # Empty command before pipe
+    with pytest.raises(ValueError, match="Empty command before pipe operator"):
+        validator.validate_pipeline(["|", "grep", "test"])
+
+    # Command not allowed
+    with pytest.raises(ValueError, match="Command not allowed"):
+        validator.validate_pipeline(["invalid_cmd", "|", "grep", "test"])
+
+
+def test_validate_command(validator, monkeypatch):
+    clear_env(monkeypatch)
+
+    # No allowed commands
+    with pytest.raises(ValueError, match="No commands are allowed"):
+        validator.validate_command(["cmd"])
+
+    monkeypatch.setenv("ALLOW_COMMANDS", "allowed_cmd")
+
+    # Empty command
+    with pytest.raises(ValueError, match="Empty command"):
+        validator.validate_command([])
+
+    # Command not allowed
+    with pytest.raises(ValueError, match="Command not allowed"):
+        validator.validate_command(["disallowed_cmd"])
+
+    # Command allowed
+    validator.validate_command(["allowed_cmd", "-arg"])  # Should not raise

tests/test_shell_executor.py

@@ -233,7 +233,7 @@ async def test_allow_commands_precedence(executor, temp_test_dir, monkeypatch):
     monkeypatch.setenv("ALLOW_COMMANDS", "echo,ls")
     monkeypatch.setenv("ALLOWED_COMMANDS", "echo,cat")
 
-    allowed = executor.get_allowed_commands()
+    allowed = executor.validator.get_allowed_commands()
     assert set(allowed) == {"echo", "ls", "cat"}
 
 
@@ -494,11 +494,11 @@ def test_validate_pipeline(executor, monkeypatch):
     monkeypatch.setenv("ALLOWED_COMMANDS", "echo,grep,cat")
 
     # Test valid pipeline
-    executor._validate_pipeline(["echo", "hello", "|", "grep", "h"])
+    executor.validator.validate_pipeline(["echo", "hello", "|", "grep", "h"])
 
     # Test empty command before pipe
     with pytest.raises(ValueError) as exc:
-        executor._validate_pipeline(["|", "grep", "test"])
+        executor.validator.validate_pipeline(["|", "grep", "test"])
     assert str(exc.value) == "Empty command before pipe operator"
 
     # Test disallowed commands in pipeline

tests/test_shell_executor_error_cases.py

@@ -42,15 +42,15 @@ async def test_no_allowed_commands_validation(monkeypatch):
         ValueError,
         match="No commands are allowed. Please set ALLOW_COMMANDS environment variable.",
     ):
-        executor._validate_command(["any_command"])
+        executor.validator.validate_command(["any_command"])
 
 
 @pytest.mark.asyncio
 async def test_shell_operator_validation():
     """Test validation of shell operators"""
     executor = ShellExecutor()
 
-    operators = [";" "&&", "||", "|"]
+    operators = [";", "&&", "||", "|"]
     for op in operators:
         # Test shell operator validation
         with pytest.raises(ValueError, match=f"Unexpected shell operator: {op}"):