refactor: Add process manager and mock tests

Author: tumf

pyproject.toml

@@ -45,6 +45,7 @@ filterwarnings = [
     "ignore::RuntimeWarning:selectors:",
     "ignore::pytest.PytestUnhandledCoroutineWarning:",
     "ignore::pytest.PytestUnraisableExceptionWarning:",
+    "ignore::DeprecationWarning:pytest_asyncio.plugin:",
 ]
 
 [tool.ruff]

src/mcp_shell_server/command_preprocessor.py

@@ -0,0 +1,131 @@
+import shlex
+from typing import Dict, List, Tuple, Union
+
+
+class CommandPreProcessor:
+    """
+    Pre-processes and validates shell commands before execution
+    """
+
+    def preprocess_command(self, command: List[str]) -> List[str]:
+        """
+        Preprocess the command to handle cases where '|' is attached to a command.
+        """
+        preprocessed_command = []
+        for token in command:
+            if token in ["||", "&&", ";"]:  # Special shell operators
+                preprocessed_command.append(token)
+            elif "|" in token and token != "|":
+                parts = token.split("|")
+                preprocessed_command.extend(
+                    [part.strip() for part in parts if part.strip()]
+                )
+                preprocessed_command.append("|")
+            else:
+                preprocessed_command.append(token)
+        return preprocessed_command
+
+    def clean_command(self, command: List[str]) -> List[str]:
+        """
+        Clean command by trimming whitespace from each part.
+        Removes empty strings but preserves arguments that are meant to be spaces.
+
+        Args:
+            command (List[str]): Original command and its arguments
+
+        Returns:
+            List[str]: Cleaned command
+        """
+        return [arg for arg in command if arg]  # Remove empty strings
+
+    def create_shell_command(self, command: List[str]) -> str:
+        """
+        Create a shell command string from a list of arguments.
+        Handles wildcards and arguments properly.
+        """
+        if not command:
+            return ""
+
+        escaped_args = []
+        for arg in command:
+            if arg.isspace():
+                # Wrap space-only arguments in single quotes
+                escaped_args.append(f"'{arg}'")
+            else:
+                # Properly escape all arguments including those with wildcards
+                escaped_args.append(shlex.quote(arg.strip()))
+
+        return " ".join(escaped_args)
+
+    def split_pipe_commands(self, command: List[str]) -> List[List[str]]:
+        """
+        Split commands by pipe operator into separate commands.
+
+        Args:
+            command (List[str]): Command and its arguments with pipe operators
+
+        Returns:
+            List[List[str]]: List of commands split by pipe operator
+        """
+        commands: List[List[str]] = []
+        current_command: List[str] = []
+
+        for arg in command:
+            if arg.strip() == "|":
+                if current_command:
+                    commands.append(current_command)
+                    current_command = []
+            else:
+                current_command.append(arg)
+
+        if current_command:
+            commands.append(current_command)
+
+        return commands
+
+    def parse_command(
+        self, command: List[str]
+    ) -> Tuple[List[str], Dict[str, Union[None, str, bool]]]:
+        """
+        Parse command and extract redirections.
+        """
+        cmd = []
+        redirects: Dict[str, Union[None, str, bool]] = {
+            "stdin": None,
+            "stdout": None,
+            "stdout_append": False,
+        }
+
+        i = 0
+        while i < len(command):
+            token = command[i]
+
+            # Shell operators check
+            if token in ["|", ";", "&&", "||"]:
+                raise ValueError(f"Unexpected shell operator: {token}")
+
+            # Output redirection
+            if token in [">", ">>"]:
+                if i + 1 >= len(command):
+                    raise ValueError("Missing path for output redirection")
+                if i + 1 < len(command) and command[i + 1] in [">", ">>", "<"]:
+                    raise ValueError("Invalid redirection target: operator found")
+                path = command[i + 1]
+                redirects["stdout"] = path
+                redirects["stdout_append"] = token == ">>"
+                i += 2
+                continue
+
+            # Input redirection
+            if token == "<":
+                if i + 1 >= len(command):
+                    raise ValueError("Missing path for input redirection")
+                path = command[i + 1]
+                redirects["stdin"] = path
+                i += 2
+                continue
+
+            cmd.append(token)
+            i += 1
+
+        return cmd, redirects

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/directory_manager.py

@@ -0,0 +1,52 @@
+import os
+from typing import Optional
+
+
+class DirectoryManager:
+    """
+    Manages directory validation and path operations for shell command execution.
+    """
+
+    def validate_directory(self, directory: Optional[str]) -> None:
+        """
+        Validate if the directory exists and is accessible.
+
+        Args:
+            directory (Optional[str]): Directory path to validate
+
+        Raises:
+            ValueError: If the directory doesn't exist, not absolute or is not accessible
+        """
+        # make directory required
+        if directory is None:
+            raise ValueError("Directory is required")
+
+        # verify directory is absolute path
+        if not os.path.isabs(directory):
+            raise ValueError(f"Directory must be an absolute path: {directory}")
+
+        if not os.path.exists(directory):
+            raise ValueError(f"Directory does not exist: {directory}")
+
+        if not os.path.isdir(directory):
+            raise ValueError(f"Not a directory: {directory}")
+
+        if not os.access(directory, os.R_OK | os.X_OK):
+            raise ValueError(f"Directory is not accessible: {directory}")
+
+    def get_absolute_path(self, path: str, base_directory: Optional[str] = None) -> str:
+        """
+        Get absolute path by joining base directory with path if path is relative.
+
+        Args:
+            path (str): The path to make absolute
+            base_directory (Optional[str]): Base directory to join with relative paths
+
+        Returns:
+            str: Absolute path
+        """
+        if os.path.isabs(path):
+            return path
+        if not base_directory:
+            return os.path.abspath(path)
+        return os.path.join(base_directory, path)