feat(server.py): add directory parameter to command execution for specifying working directory

tumf · tumf · commit 01095380cf3b · 2024-12-11T13:59:37.000+09:00
feat(shell_executor.py): implement directory validation and support for executing commands in specified directory
test(tests): add tests for command execution with directory functionality to ensure correct behavior
diff --git a/mcp_shell_server/server.py b/mcp_shell_server/server.py
@@ -45,6 +45,10 @@ def get_tool_description(self) -> Tool:
                         "type": "string",
                         "description": "Input to be passed to the command via stdin",
                     },
+                    "directory": {
+                        "type": "string",
+                        "description": "Working directory where the command will be executed",
+                    },
                 },
                 "required": ["command"],
             },
@@ -54,11 +58,12 @@ 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")
+        directory = arguments.get("directory")
 
         if not command:
             raise ValueError("No command provided")
 
-        result = await self.executor.execute(command, stdin)
+        result = await self.executor.execute(command, stdin, directory)
 
         # Raise error if command execution failed
         if result.get("error"):
@@ -115,4 +120,4 @@ async def main() -> None:
             )
     except Exception as e:
         logger.error(f"Server error: {str(e)}")
-        raise
+        raise
diff --git a/mcp_shell_server/shell_executor.py b/mcp_shell_server/shell_executor.py
@@ -74,19 +74,40 @@ def _validate_command(self, command: List[str]) -> None:
                         f"Command not allowed after {cleaned_arg}: {next_cmd}"
                     )
 
+    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 or is not accessible
+        """
+        if directory is None:
+            return
+
+        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_allowed_commands(self) -> list[str]:
         """Get the allowed commands"""
         return list(self._get_allowed_commands())
 
     async def execute(
-        self, command: List[str], stdin: Optional[str] = None
+        self, command: List[str], stdin: Optional[str] = None, directory: Optional[str] = None
     ) -> Dict[str, Any]:
         """
-        Execute a shell command with optional stdin input.
+        Execute a shell command with optional stdin input and working directory.
 
         Args:
             command (List[str]): Command and its arguments
             stdin (Optional[str]): Input to be passed to the command via stdin
+            directory (Optional[str]): Working directory for command execution
 
         Returns:
             Dict[str, Any]: Execution result containing stdout, stderr, status code, and execution time.
@@ -101,6 +122,7 @@ async def execute(
                 raise ValueError("Empty command")
 
             self._validate_command(cleaned_command)
+            self._validate_directory(directory)
         except ValueError as e:
             return {
                 "error": str(e),
@@ -118,6 +140,7 @@ async def execute(
                 stdout=asyncio.subprocess.PIPE,
                 stderr=asyncio.subprocess.PIPE,
                 env={"PATH": os.environ.get("PATH", "")},
+                cwd=directory,  # Set working directory if specified
             )
 
             stdin_bytes = stdin.encode() if stdin else None
@@ -129,6 +152,7 @@ async def execute(
                 "stderr": stderr.decode() if stderr else "",
                 "status": process.returncode,
                 "execution_time": time.time() - start_time,
+                "directory": directory,  # Include working directory in response
             }
         except FileNotFoundError:
             return {
@@ -145,4 +169,4 @@ async def execute(
                 "stdout": "",
                 "stderr": str(e),
                 "execution_time": time.time() - start_time,
-            }
+            }
diff --git a/pyproject.toml b/pyproject.toml
@@ -1,4 +1,3 @@
-
 [project]
 name = "mcp-shell-server"
 version = "0.1.0"
@@ -22,6 +21,7 @@ test = [
     "pytest>=7.4.0",
     "pytest-asyncio>=0.23.0", 
     "pytest-env>=1.1.0",
+    "pytest-cov>=4.1.0",
 ]
 dev = [
     "ruff>=0.0.262",
@@ -62,4 +62,4 @@ target-version = ['py311']
 
 [tool.isort]
 profile = "black" 
-line_length = 88
+line_length = 88
diff --git a/tests/test_server.py b/tests/test_server.py
@@ -1,9 +1,19 @@
+import os
+import tempfile
 import pytest
 from mcp.types import TextContent, Tool
 
 from mcp_shell_server.server import call_tool, list_tools
 
 
+@pytest.fixture
+def temp_test_dir():
+    """Create a temporary directory for testing"""
+    with tempfile.TemporaryDirectory() as tmpdirname:
+        # Return the real path to handle macOS /private/tmp symlink
+        yield os.path.realpath(tmpdirname)
+
+
 @pytest.mark.asyncio
 async def test_list_tools():
     """Test listing of available tools"""
@@ -16,6 +26,7 @@ async def test_list_tools():
     assert tool.inputSchema["type"] == "object"
     assert "command" in tool.inputSchema["properties"]
     assert "stdin" in tool.inputSchema["properties"]
+    assert "directory" in tool.inputSchema["properties"]  # New assertion
     assert tool.inputSchema["required"] == ["command"]
 
 
@@ -72,3 +83,90 @@ async def test_call_tool_empty_command():
     with pytest.raises(RuntimeError) as excinfo:
         await call_tool("execute", {"command": []})
     assert "No command provided" in str(excinfo.value)
+
+
+# New tests for directory functionality
+@pytest.mark.asyncio
+async def test_call_tool_with_directory(temp_test_dir, monkeypatch):
+    """Test command execution in a specific directory"""
+    monkeypatch.setenv("ALLOW_COMMANDS", "pwd")
+    result = await call_tool("execute", {
+        "command": ["pwd"],
+        "directory": temp_test_dir
+    })
+    assert len(result) == 1
+    assert isinstance(result[0], TextContent)
+    assert result[0].type == "text"
+    assert result[0].text.strip() == temp_test_dir
+
+
+@pytest.mark.asyncio
+async def test_call_tool_with_file_operations(temp_test_dir, monkeypatch):
+    """Test file operations in a specific directory"""
+    monkeypatch.setenv("ALLOW_COMMANDS", "ls,cat")
+    
+    # Create a test file
+    test_file = os.path.join(temp_test_dir, "test.txt")
+    with open(test_file, "w") as f:
+        f.write("test content")
+
+    # Test ls command
+    result = await call_tool("execute", {
+        "command": ["ls"],
+        "directory": temp_test_dir
+    })
+    assert "test.txt" in result[0].text
+
+    # Test cat command
+    result = await call_tool("execute", {
+        "command": ["cat", "test.txt"],
+        "directory": temp_test_dir
+    })
+    assert result[0].text.strip() == "test content"
+
+
+@pytest.mark.asyncio
+async def test_call_tool_with_nonexistent_directory(monkeypatch):
+    """Test command execution with a non-existent directory"""
+    monkeypatch.setenv("ALLOW_COMMANDS", "ls")
+    with pytest.raises(RuntimeError) as excinfo:
+        await call_tool("execute", {
+            "command": ["ls"],
+            "directory": "/nonexistent/directory"
+        })
+    assert "Directory does not exist: /nonexistent/directory" in str(excinfo.value)
+
+
+@pytest.mark.asyncio
+async def test_call_tool_with_file_as_directory(temp_test_dir, monkeypatch):
+    """Test command execution with a file specified as directory"""
+    monkeypatch.setenv("ALLOW_COMMANDS", "ls")
+    
+    # Create a test file
+    test_file = os.path.join(temp_test_dir, "test.txt")
+    with open(test_file, "w") as f:
+        f.write("test content")
+
+    with pytest.raises(RuntimeError) as excinfo:
+        await call_tool("execute", {
+            "command": ["ls"],
+            "directory": test_file
+        })
+    assert f"Not a directory: {test_file}" in str(excinfo.value)
+
+
+@pytest.mark.asyncio
+async def test_call_tool_with_nested_directory(temp_test_dir, monkeypatch):
+    """Test command execution in a nested directory"""
+    monkeypatch.setenv("ALLOW_COMMANDS", "pwd,mkdir")
+    
+    # Create a nested directory
+    nested_dir = os.path.join(temp_test_dir, "nested")
+    os.mkdir(nested_dir)
+    nested_real_path = os.path.realpath(nested_dir)
+    
+    result = await call_tool("execute", {
+        "command": ["pwd"],
+        "directory": nested_dir
+    })
+    assert result[0].text.strip() == nested_real_path
diff --git a/tests/test_shell_executor.py b/tests/test_shell_executor.py
@@ -1,4 +1,6 @@
+import os
 import pytest
+import tempfile
 
 from mcp_shell_server.shell_executor import ShellExecutor
 
@@ -8,6 +10,14 @@ def executor():
     return ShellExecutor()
 
 
+@pytest.fixture
+def temp_test_dir():
+    """Create a temporary directory for testing"""
+    with tempfile.TemporaryDirectory() as tmpdirname:
+        # Return the real path to handle macOS /private/tmp symlink
+        yield os.path.realpath(tmpdirname)
+
+
 @pytest.mark.asyncio
 async def test_basic_command_execution(executor, monkeypatch):
     monkeypatch.setenv("ALLOW_COMMANDS", "echo")
@@ -72,3 +82,83 @@ async def test_shell_operators_not_allowed(executor, monkeypatch):
         result = await executor.execute(["echo", "hello", op])
         assert result["error"] == f"Unexpected shell operator: {op}"
         assert result["status"] == 1
+
+
+# New tests for directory functionality
+@pytest.mark.asyncio
+async def test_execute_in_directory(executor, temp_test_dir, monkeypatch):
+    """Test command execution in a specific directory"""
+    monkeypatch.setenv("ALLOW_COMMANDS", "pwd")
+    result = await executor.execute(["pwd"], directory=temp_test_dir)
+    assert result["error"] is None
+    assert result["status"] == 0
+    assert result["stdout"].strip() == temp_test_dir
+
+
+@pytest.mark.asyncio
+async def test_execute_with_file_in_directory(executor, temp_test_dir, monkeypatch):
+    """Test command execution with a file in the specified directory"""
+    monkeypatch.setenv("ALLOW_COMMANDS", "ls,cat")
+    
+    # Create a test file in the temporary directory
+    test_file = os.path.join(temp_test_dir, "test.txt")
+    with open(test_file, "w") as f:
+        f.write("test content")
+
+    # Test ls command
+    result = await executor.execute(["ls"], directory=temp_test_dir)
+    assert "test.txt" in result["stdout"]
+    
+    # Test cat command
+    result = await executor.execute(["cat", "test.txt"], directory=temp_test_dir)
+    assert result["stdout"].strip() == "test content"
+
+
+@pytest.mark.asyncio
+async def test_execute_with_nonexistent_directory(executor, monkeypatch):
+    """Test command execution with a non-existent directory"""
+    monkeypatch.setenv("ALLOW_COMMANDS", "ls")
+    result = await executor.execute(["ls"], directory="/nonexistent/directory")
+    assert result["error"] == "Directory does not exist: /nonexistent/directory"
+    assert result["status"] == 1
+
+
+@pytest.mark.asyncio
+async def test_execute_with_file_as_directory(executor, temp_test_dir, monkeypatch):
+    """Test command execution with a file specified as directory"""
+    monkeypatch.setenv("ALLOW_COMMANDS", "ls")
+    
+    # Create a test file
+    test_file = os.path.join(temp_test_dir, "test.txt")
+    with open(test_file, "w") as f:
+        f.write("test content")
+
+    result = await executor.execute(["ls"], directory=test_file)
+    assert result["error"] == f"Not a directory: {test_file}"
+    assert result["status"] == 1
+
+
+@pytest.mark.asyncio
+async def test_execute_with_no_directory_specified(executor, monkeypatch):
+    """Test command execution without specifying a directory"""
+    monkeypatch.setenv("ALLOW_COMMANDS", "pwd")
+    result = await executor.execute(["pwd"])
+    assert result["error"] is None
+    assert result["status"] == 0
+    assert os.path.exists(result["stdout"].strip())
+
+
+@pytest.mark.asyncio
+async def test_execute_with_nested_directory(executor, temp_test_dir, monkeypatch):
+    """Test command execution in a nested directory"""
+    monkeypatch.setenv("ALLOW_COMMANDS", "pwd,mkdir,ls")
+    
+    # Create a nested directory
+    nested_dir = os.path.join(temp_test_dir, "nested")
+    os.mkdir(nested_dir)
+    nested_real_path = os.path.realpath(nested_dir)
+    
+    result = await executor.execute(["pwd"], directory=nested_dir)
+    assert result["error"] is None
+    assert result["status"] == 0
+    assert result["stdout"].strip() == nested_real_path
diff --git a/uv.lock b/uv.lock
