massgen 0.1.2__py3-none-any.whl → 0.1.4__py3-none-any.whl

massgen/tests/test_binary_file_blocking.py

@@ -0,0 +1,274 @@
+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+"""
+Tests for binary file blocking in PathPermissionManager.
+
+These tests ensure that text-based read tools (Read, read_text_file, etc.)
+are blocked from reading binary files (images, videos, audio, etc.) to prevent
+context pollution with binary data.
+"""
+
+from pathlib import Path
+
+import pytest
+
+from massgen.filesystem_manager._base import Permission
+from massgen.filesystem_manager._path_permission_manager import PathPermissionManager
+
+
+@pytest.fixture
+def permission_manager():
+    """Create a PathPermissionManager instance for testing."""
+    manager = PathPermissionManager(
+        context_write_access_enabled=False,
+        enforce_read_before_delete=True,
+    )
+    # Add a workspace path for testing
+    test_workspace = Path("/tmp/test_workspace").resolve()
+    manager.add_path(test_workspace, Permission.WRITE, "workspace")
+    return manager
+
+
+class TestBinaryFileBlocking:
+    """Test suite for binary file blocking functionality."""
+
+    @pytest.mark.asyncio
+    async def test_block_read_image_with_read_tool(self, permission_manager):
+        """Test that Read tool is blocked from reading image files."""
+        tool_name = "Read"
+        tool_args = {"file_path": "/tmp/test_workspace/photo.jpg"}
+
+        allowed, reason = await permission_manager.pre_tool_use_hook(tool_name, tool_args)
+
+        assert not allowed, "Read should be blocked from reading .jpg files"
+        assert reason is not None
+        assert "understand_image" in reason.lower()
+        assert "photo.jpg" in reason
+
+    @pytest.mark.asyncio
+    async def test_block_read_text_file_image(self, permission_manager):
+        """Test that read_text_file (MCP) is blocked from reading image files."""
+        tool_name = "mcp__filesystem__read_text_file"
+        tool_args = {"path": "/tmp/test_workspace/diagram.png"}
+
+        allowed, reason = await permission_manager.pre_tool_use_hook(tool_name, tool_args)
+
+        assert not allowed, "read_text_file should be blocked from reading .png files"
+        assert reason is not None
+        assert "understand_image" in reason.lower()
+
+    @pytest.mark.asyncio
+    async def test_block_read_video(self, permission_manager):
+        """Test that Read tool is blocked from reading video files."""
+        tool_name = "Read"
+        tool_args = {"file_path": "/tmp/test_workspace/demo.mp4"}
+
+        allowed, reason = await permission_manager.pre_tool_use_hook(tool_name, tool_args)
+
+        assert not allowed, "Read should be blocked from reading .mp4 files"
+        assert reason is not None
+        assert "understand_video" in reason.lower()
+
+    @pytest.mark.asyncio
+    async def test_block_read_audio(self, permission_manager):
+        """Test that Read tool is blocked from reading audio files."""
+        tool_name = "Read"
+        tool_args = {"file_path": "/tmp/test_workspace/recording.mp3"}
+
+        allowed, reason = await permission_manager.pre_tool_use_hook(tool_name, tool_args)
+
+        assert not allowed, "Read should be blocked from reading .mp3 files"
+        assert reason is not None
+        assert "audio" in reason.lower()
+
+    @pytest.mark.asyncio
+    async def test_allow_read_text_file(self, permission_manager):
+        """Test that Read tool is allowed to read text files."""
+        tool_name = "Read"
+        tool_args = {"file_path": "/tmp/test_workspace/document.txt"}
+
+        allowed, reason = await permission_manager.pre_tool_use_hook(tool_name, tool_args)
+
+        assert allowed, "Read should be allowed to read .txt files"
+        assert reason is None
+
+    @pytest.mark.asyncio
+    async def test_allow_read_code_file(self, permission_manager):
+        """Test that Read tool is allowed to read code files."""
+        test_cases = [
+            "script.py",
+            "app.js",
+            "component.tsx",
+            "main.go",
+            "app.rs",
+        ]
+
+        for filename in test_cases:
+            tool_name = "Read"
+            tool_args = {"file_path": f"/tmp/test_workspace/{filename}"}
+
+            allowed, reason = await permission_manager.pre_tool_use_hook(tool_name, tool_args)
+
+            assert allowed, f"Read should be allowed to read {filename}"
+            assert reason is None
+
+    @pytest.mark.asyncio
+    async def test_block_all_image_formats(self, permission_manager):
+        """Test that all image formats are blocked."""
+        image_extensions = [".jpg", ".jpeg", ".png", ".gif", ".bmp", ".ico", ".svg", ".webp", ".tiff", ".tif"]
+
+        for ext in image_extensions:
+            tool_name = "Read"
+            tool_args = {"file_path": f"/tmp/test_workspace/image{ext}"}
+
+            allowed, reason = await permission_manager.pre_tool_use_hook(tool_name, tool_args)
+
+            assert not allowed, f"Read should be blocked from reading {ext} files"
+            assert reason is not None
+
+    @pytest.mark.asyncio
+    async def test_block_all_video_formats(self, permission_manager):
+        """Test that all video formats are blocked."""
+        video_extensions = [".mp4", ".avi", ".mov", ".mkv", ".flv", ".wmv", ".webm", ".m4v", ".mpg", ".mpeg"]
+
+        for ext in video_extensions:
+            tool_name = "Read"
+            tool_args = {"file_path": f"/tmp/test_workspace/video{ext}"}
+
+            allowed, reason = await permission_manager.pre_tool_use_hook(tool_name, tool_args)
+
+            assert not allowed, f"Read should be blocked from reading {ext} files"
+            assert reason is not None
+
+    @pytest.mark.asyncio
+    async def test_block_all_audio_formats(self, permission_manager):
+        """Test that all audio formats are blocked."""
+        audio_extensions = [".mp3", ".wav", ".ogg", ".flac", ".aac", ".m4a", ".wma"]
+
+        for ext in audio_extensions:
+            tool_name = "Read"
+            tool_args = {"file_path": f"/tmp/test_workspace/audio{ext}"}
+
+            allowed, reason = await permission_manager.pre_tool_use_hook(tool_name, tool_args)
+
+            assert not allowed, f"Read should be blocked from reading {ext} files"
+            assert reason is not None
+
+    @pytest.mark.asyncio
+    async def test_block_archive_formats(self, permission_manager):
+        """Test that archive formats are blocked."""
+        archive_extensions = [".zip", ".tar", ".gz", ".bz2", ".7z", ".rar", ".xz"]
+
+        for ext in archive_extensions:
+            tool_name = "Read"
+            tool_args = {"file_path": f"/tmp/test_workspace/archive{ext}"}
+
+            allowed, reason = await permission_manager.pre_tool_use_hook(tool_name, tool_args)
+
+            assert not allowed, f"Read should be blocked from reading {ext} files"
+            assert reason is not None
+
+    @pytest.mark.asyncio
+    async def test_block_executable_formats(self, permission_manager):
+        """Test that executable/binary formats are blocked."""
+        binary_extensions = [".exe", ".bin", ".dll", ".so", ".dylib", ".o", ".a", ".pyc", ".class", ".jar"]
+
+        for ext in binary_extensions:
+            tool_name = "Read"
+            tool_args = {"file_path": f"/tmp/test_workspace/binary{ext}"}
+
+            allowed, reason = await permission_manager.pre_tool_use_hook(tool_name, tool_args)
+
+            assert not allowed, f"Read should be blocked from reading {ext} files"
+            assert reason is not None
+
+    @pytest.mark.asyncio
+    async def test_block_old_office_formats(self, permission_manager):
+        """Test that old Office formats are blocked (use understand_file instead)."""
+        old_office_extensions = [".doc", ".xls", ".ppt"]
+
+        for ext in old_office_extensions:
+            tool_name = "Read"
+            tool_args = {"file_path": f"/tmp/test_workspace/document{ext}"}
+
+            allowed, reason = await permission_manager.pre_tool_use_hook(tool_name, tool_args)
+
+            assert not allowed, f"Read should be blocked from reading {ext} files"
+            assert reason is not None
+
+    @pytest.mark.asyncio
+    async def test_block_office_formats(self, permission_manager):
+        """Test that Office document formats are blocked from Read (must use understand_file).
+
+        These are binary formats that should be handled by understand_file tool,
+        which can properly extract text from them using specialized libraries.
+        """
+        office_extensions = [".pdf", ".docx", ".xlsx", ".pptx"]
+
+        for ext in office_extensions:
+            tool_name = "Read"
+            tool_args = {"file_path": f"/tmp/test_workspace/document{ext}"}
+
+            allowed, reason = await permission_manager.pre_tool_use_hook(tool_name, tool_args)
+
+            assert not allowed, f"Read should be blocked from reading {ext} files (use understand_file)"
+            assert reason is not None
+            assert "understand_file" in reason.lower()
+
+    @pytest.mark.asyncio
+    async def test_case_insensitive_extension_check(self, permission_manager):
+        """Test that extension checking is case-insensitive."""
+        test_cases = [
+            "/tmp/test_workspace/PHOTO.JPG",
+            "/tmp/test_workspace/Video.MP4",
+            "/tmp/test_workspace/Audio.MP3",
+            "/tmp/test_workspace/Image.PNG",
+        ]
+
+        for file_path in test_cases:
+            tool_name = "Read"
+            tool_args = {"file_path": file_path}
+
+            allowed, reason = await permission_manager.pre_tool_use_hook(tool_name, tool_args)
+
+            assert not allowed, f"Read should be blocked from reading {file_path} (case-insensitive)"
+            assert reason is not None
+
+    @pytest.mark.asyncio
+    async def test_non_text_read_tools_not_affected(self, permission_manager):
+        """Test that non-text-read tools are not affected by binary file blocking."""
+        # Tools like Write, Edit, Delete should not be affected
+        test_cases = [
+            ("Write", {"file_path": "/tmp/test_workspace/image.jpg"}),
+            ("Edit", {"file_path": "/tmp/test_workspace/video.mp4"}),
+            ("Grep", {"pattern": "test"}),  # No file_path, should pass
+        ]
+
+        for tool_name, tool_args in test_cases:
+            allowed, reason = await permission_manager.pre_tool_use_hook(tool_name, tool_args)
+
+            # These tools have their own validation, but shouldn't be blocked by binary check
+            # (they may be blocked for other reasons like permissions)
+            # The key is that _validate_binary_file_access is not called for these
+            assert isinstance(allowed, bool)  # Should complete without binary file error
+
+    @pytest.mark.asyncio
+    async def test_helpful_error_messages(self, permission_manager):
+        """Test that error messages provide helpful suggestions for blocked binary files."""
+        test_cases = [
+            (".jpg", "understand_image"),
+            (".mp4", "understand_video"),
+            (".mp3", "audio"),
+            (".pdf", "understand_file"),
+            (".docx", "understand_file"),
+        ]
+
+        for ext, expected_suggestion in test_cases:
+            tool_name = "Read"
+            tool_args = {"file_path": f"/tmp/test_workspace/file{ext}"}
+
+            allowed, reason = await permission_manager.pre_tool_use_hook(tool_name, tool_args)
+
+            assert not allowed, f"File with {ext} extension should be blocked"
+            assert reason is not None
+            assert expected_suggestion.lower() in reason.lower(), f"Error message should suggest {expected_suggestion} for {ext} files"

massgen/tests/test_case_studies.md

@@ -1,12 +1,12 @@
 # MassGen Case Study Test Commands
 
-This document contains commands to test all the case studies from `docs/case_studies/` using the three agents default configuration.
+This document contains commands to test all the case studies from `docs/source/examples/case_studies/` using the three agents default configuration.
 
 ## Quick Commands
 
 All tests use the `three_agents_default.yaml` configuration with:
 - **Gemini 2.5 Flash** (web search enabled)
-- **GPT-4o-mini** (web search + code interpreter)  
+- **GPT-4o-mini** (web search + code interpreter)
 - **Grok 3 mini** (web search with citations)
 
 ### 1. Collaborative Creative Writing
@@ -17,8 +17,8 @@ python massgen/cli.py --config massgen/configs/three_agents_default.yaml "Write
 # From tests directory:
 python ../cli.py --config ../configs/three_agents_default.yaml "Write a short story about a robot who discovers music."
 ```
-**Original:** gpt-4o, gemini-2.5-flash, grok-3-mini  
-**Current:** gemini2.5flash, 4omini, grok3mini with builtin tools  
+**Original:** gpt-4o, gemini-2.5-flash, grok-3-mini
+**Current:** gemini2.5flash, 4omini, grok3mini with builtin tools
 
 ### 2. AI News Synthesis
 ```bash
@@ -28,8 +28,8 @@ python massgen/cli.py --config massgen/configs/three_agents_default.yaml "find b
 # From tests directory:
 python ../cli.py --config ../configs/three_agents_default.yaml "find big AI news this week"
 ```
-**Original:** gpt-4.1, gemini-2.5-flash, grok-3-mini  
-**Current:** gemini2.5flash, 4omini, grok3mini with web search  
+**Original:** gpt-4.1, gemini-2.5-flash, grok-3-mini
+**Current:** gemini2.5flash, 4omini, grok3mini with web search
 
 ### 3. Grok HLE Cost Estimation
 ```bash
@@ -39,8 +39,8 @@ python massgen/cli.py --config massgen/configs/three_agents_default.yaml "How mu
 # From tests directory:
 python ../cli.py --config ../configs/three_agents_default.yaml "How much does it cost to run HLE benchmark with Grok-4"
 ```
-**Original:** gpt-4o, gemini-2.5-flash, grok-3-mini  
-**Current:** gemini2.5flash, 4omini, grok3mini with web search  
+**Original:** gpt-4o, gemini-2.5-flash, grok-3-mini
+**Current:** gemini2.5flash, 4omini, grok3mini with web search
 
 ### 4. IMO 2025 Winner
 ```bash
@@ -50,8 +50,8 @@ python massgen/cli.py --config massgen/configs/three_agents_default.yaml "Which
 # From tests directory:
 python ../cli.py --config ../configs/three_agents_default.yaml "Which AI won IMO 2025?"
 ```
-**Original:** gemini-2.5-flash, gpt-4.1 (2 agents)  
-**Current:** gemini2.5flash, 4omini, grok3mini (3 agents with web search)  
+**Original:** gemini-2.5-flash, gpt-4.1 (2 agents)
+**Current:** gemini2.5flash, 4omini, grok3mini (3 agents with web search)
 
 ### 5. Stockholm Travel Guide
 ```bash
@@ -61,7 +61,7 @@ python massgen/cli.py --config massgen/configs/three_agents_default.yaml "what's
 # From tests directory:
 python ../cli.py --config ../configs/three_agents_default.yaml "what's best to do in Stockholm in October 2025"
 ```
-**Original:** gemini-2.5-flash, gpt-4o (2 agents)  
+**Original:** gemini-2.5-flash, gpt-4o (2 agents)
 **Current:** gemini2.5flash, 4omini, grok3mini with web search for current info
 
 ## Configuration Details
@@ -70,7 +70,7 @@ The `three_agents_default.yaml` configuration provides:
 
 ### Agent Capabilities
 - **gemini2.5flash**: Gemini 2.5 Flash with web search
-- **4omini**: GPT-4o-mini with web search + code interpreter  
+- **4omini**: GPT-4o-mini with web search + code interpreter
 - **grok3mini**: Grok 3 mini with web search and citations
 
 ### UI Features

massgen/tests/test_code_execution.py

@@ -154,6 +154,112 @@ class TestCommandSanitization:
             _sanitize_command(cmd)
 
 
+class TestSudoSanitization:
+    """Test sudo sanitization respects enable_sudo flag."""
+
+    def test_sudo_blocked_by_default(self):
+        """Test that sudo is blocked when enable_sudo=False (default)."""
+        from massgen.filesystem_manager._code_execution_server import _sanitize_command
+
+        sudo_commands = [
+            "sudo apt-get update",
+            "sudo apt-get install -y ffmpeg",
+            "sudo pip install tensorflow",
+            "sudo npm install -g typescript",
+            "sudo chmod 755 file.txt",
+            "echo 'test' && sudo apt update",
+        ]
+
+        for cmd in sudo_commands:
+            with pytest.raises(ValueError, match="sudo.*not allowed"):
+                _sanitize_command(cmd, enable_sudo=False)
+
+    def test_sudo_allowed_when_enabled(self):
+        """Test that sudo is allowed when enable_sudo=True."""
+        from massgen.filesystem_manager._code_execution_server import _sanitize_command
+
+        sudo_commands = [
+            "sudo apt-get update",
+            "sudo apt-get install -y ffmpeg",
+            "sudo pip install tensorflow",
+            "sudo npm install -g typescript",
+            "sudo chown user:group file.txt",  # chown allowed with sudo enabled
+            "sudo chmod 755 file.txt",  # chmod allowed with sudo enabled
+        ]
+
+        for cmd in sudo_commands:
+            # Should not raise when enable_sudo=True
+            _sanitize_command(cmd, enable_sudo=True)
+
+    def test_other_dangerous_patterns_still_blocked_with_sudo(self):
+        """Test that other dangerous patterns are still blocked even with sudo enabled."""
+        from massgen.filesystem_manager._code_execution_server import _sanitize_command
+
+        # These should ALWAYS be blocked, regardless of enable_sudo
+        dangerous_commands = [
+            "sudo rm -rf /",  # Still blocked - root deletion
+            "rm -rf /",  # Still blocked
+            "dd if=/dev/zero of=/dev/sda",  # Still blocked - dd command
+            "sudo dd if=/dev/zero of=/dev/sda",  # Still blocked
+            ":(){ :|:& };:",  # Still blocked - fork bomb
+            "mv file /dev/null",  # Still blocked
+            "sudo mv file /dev/null",  # Still blocked
+            "echo test > /dev/sda1",  # Still blocked - writing to disk
+        ]
+
+        for cmd in dangerous_commands:
+            with pytest.raises(ValueError, match="dangerous|not allowed"):
+                _sanitize_command(cmd, enable_sudo=True)
+
+    def test_su_chown_chmod_blocked_without_sudo_flag(self):
+        """Test that su, chown, chmod are blocked when enable_sudo=False."""
+        from massgen.filesystem_manager._code_execution_server import _sanitize_command
+
+        commands = [
+            "su root",
+            "su - postgres",
+            "chown root:root file.txt",
+            "chmod 777 file.txt",
+            "chmod +x script.sh",
+        ]
+
+        for cmd in commands:
+            with pytest.raises(ValueError, match="not allowed"):
+                _sanitize_command(cmd, enable_sudo=False)
+
+    def test_su_chown_chmod_allowed_with_sudo_flag(self):
+        """Test that su, chown, chmod are allowed when enable_sudo=True (Docker sudo mode)."""
+        from massgen.filesystem_manager._code_execution_server import _sanitize_command
+
+        # In Docker sudo mode, these are safe because they're confined to container
+        commands = [
+            "su postgres",
+            "chown user:group file.txt",
+            "chmod 755 file.txt",
+            "chmod +x script.sh",
+        ]
+
+        for cmd in commands:
+            # Should not raise when enable_sudo=True
+            _sanitize_command(cmd, enable_sudo=True)
+
+    def test_local_mode_blocks_sudo(self):
+        """Test that local mode (non-Docker) blocks sudo commands."""
+        from massgen.filesystem_manager._code_execution_server import _sanitize_command
+
+        # In local mode (enable_sudo=False), sudo should be blocked for safety
+        with pytest.raises(ValueError, match="sudo.*not allowed"):
+            _sanitize_command("sudo apt-get install malicious-package", enable_sudo=False)
+
+    def test_docker_sudo_mode_allows_sudo(self):
+        """Test that Docker sudo mode allows sudo commands."""
+        from massgen.filesystem_manager._code_execution_server import _sanitize_command
+
+        # In Docker mode with enable_sudo=True, sudo should be allowed
+        # (safe because it's inside container)
+        _sanitize_command("sudo apt-get install gh", enable_sudo=True)
+
+
 class TestOutputHandling:
     """Test output capture and size limits."""
 
@@ -674,6 +780,78 @@ class TestDockerExecution:
         # Cleanup
         manager.cleanup("test_context")
 
+    @pytest.mark.docker
+    def test_docker_sudo_enabled_image_selection(self):
+        """Test that enabling sudo automatically selects the sudo image variant."""
+        from massgen.filesystem_manager._docker_manager import DockerManager
+
+        # Test 1: Default image with sudo=False should use regular image
+        manager_no_sudo = DockerManager(enable_sudo=False)
+        assert manager_no_sudo.image == "massgen/mcp-runtime:latest"
+        assert manager_no_sudo.enable_sudo is False
+
+        # Test 2: Default image with sudo=True should auto-switch to sudo variant
+        manager_with_sudo = DockerManager(enable_sudo=True)
+        assert manager_with_sudo.image == "massgen/mcp-runtime-sudo:latest"
+        assert manager_with_sudo.enable_sudo is True
+
+        # Test 3: Custom image with sudo=True should keep custom image
+        manager_custom = DockerManager(
+            image="my-custom-image:latest",
+            enable_sudo=True,
+        )
+        assert manager_custom.image == "my-custom-image:latest"
+        assert manager_custom.enable_sudo is True
+
+    @pytest.mark.docker
+    def test_docker_sudo_functionality(self, tmp_path):
+        """Test that sudo commands work in sudo-enabled container."""
+        from massgen.filesystem_manager._docker_manager import DockerManager
+
+        # Skip if sudo image not built
+        manager = DockerManager(enable_sudo=True)
+        try:
+            manager.ensure_image_exists()
+        except RuntimeError:
+            pytest.skip("Sudo Docker image not built. Run: bash massgen/docker/build.sh --sudo")
+
+        workspace = tmp_path / "workspace_sudo"
+        workspace.mkdir()
+
+        # Create container with sudo enabled
+        manager.create_container(
+            agent_id="test_sudo",
+            workspace_path=workspace,
+        )
+
+        # Test 1: Verify whoami returns 'massgen' (non-root user)
+        result_whoami = manager.exec_command(
+            agent_id="test_sudo",
+            command="whoami",
+        )
+        assert result_whoami["success"] is True
+        assert "massgen" in result_whoami["stdout"]
+
+        # Test 2: Verify sudo whoami returns 'root' (sudo works)
+        result_sudo_whoami = manager.exec_command(
+            agent_id="test_sudo",
+            command="sudo whoami",
+        )
+        assert result_sudo_whoami["success"] is True
+        assert "root" in result_sudo_whoami["stdout"]
+
+        # Test 3: Verify sudo apt-get update works (package installation capability)
+        result_apt = manager.exec_command(
+            agent_id="test_sudo",
+            command="sudo apt-get update",
+            timeout=60,
+        )
+        # This should succeed in sudo image (may fail in network=none, but command should run)
+        assert result_apt["exit_code"] is not None
+
+        # Cleanup
+        manager.cleanup("test_sudo")
+
 
 if __name__ == "__main__":
     pytest.main([__file__, "-v"])