Spaces:
Sleeping
Sleeping
| """ | |
| Tests for the Docker Executor and Command Safety Checker. | |
| """ | |
| import pytest | |
| from executor.safety import CommandSafetyChecker, SafetyCheckResult | |
| from executor.docker_executor import DockerExecutor, ExecutionResult | |
| class TestCommandSafetyChecker: | |
| """Tests for the command whitelist/blocklist safety system.""" | |
| def setup_method(self): | |
| """Create a fresh checker for each test.""" | |
| self.checker = CommandSafetyChecker() | |
| def test_whitelisted_commands_pass(self): | |
| safe_commands = [ | |
| "pip install flask", | |
| "python app.py", | |
| "ls -la /app", | |
| "cat /app/config.py", | |
| "echo hello", | |
| "grep error log.txt", | |
| "ps aux", | |
| "kill 12345", | |
| "curl http://localhost:5000", | |
| "mkdir -p /app/logs", | |
| "sed -i 's/old/new/' file.txt", | |
| "export DATABASE_URL=postgres://...", | |
| ] | |
| for cmd in safe_commands: | |
| result = self.checker.check(cmd) | |
| assert result.is_safe, f"Expected '{cmd}' to be safe, got: {result.reason}" | |
| def test_blocklisted_commands_blocked(self): | |
| dangerous_commands = [ | |
| "rm -rf /", | |
| "rm -rf /*", | |
| "dd if=/dev/zero of=/dev/sda", | |
| "mkfs.ext4 /dev/sda1", | |
| "chmod 777 /", | |
| ] | |
| for cmd in dangerous_commands: | |
| result = self.checker.check(cmd) | |
| assert result.is_blocked, f"Expected '{cmd}' to be blocked" | |
| assert not result.is_safe | |
| def test_sudo_dangerous_blocked(self): | |
| sudo_commands = [ | |
| "sudo rm -rf /home", | |
| "sudo dd if=/dev/zero of=/dev/sda", | |
| "sudo shutdown now", | |
| ] | |
| for cmd in sudo_commands: | |
| result = self.checker.check(cmd) | |
| assert result.is_blocked, f"Expected '{cmd}' to be blocked" | |
| def test_unknown_commands_rejected(self): | |
| unknown_commands = [ | |
| "custom_script", | |
| "malware_tool", | |
| "nc -l 4444", | |
| ] | |
| for cmd in unknown_commands: | |
| result = self.checker.check(cmd) | |
| assert not result.is_safe, f"Expected '{cmd}' to be rejected" | |
| assert not result.is_whitelisted | |
| def test_empty_command_rejected(self): | |
| result = self.checker.check("") | |
| assert not result.is_safe | |
| def test_pipe_commands_check_first_part(self): | |
| result = self.checker.check("lsof -i:5000 | grep python") | |
| assert result.is_safe | |
| def test_chained_commands(self): | |
| result = self.checker.check("pip install flask && python app.py") | |
| assert result.is_safe | |
| def test_env_var_prefix(self): | |
| result = self.checker.check("DATABASE_URL=test python app.py") | |
| assert result.is_safe | |
| def test_reboot_keyword_in_path_not_blocked(self): | |
| result = self.checker.check("cat /app/reboot_config.py") | |
| assert result.is_safe | |
| def test_shutdown_keyword_in_grep_not_blocked(self): | |
| result = self.checker.check("grep shutdown /var/log/syslog") | |
| assert result.is_safe | |
| def test_halt_keyword_in_filename_not_blocked(self): | |
| result = self.checker.check("python halt_checker.py") | |
| assert result.is_safe | |
| class TestDockerExecutor: | |
| """Tests for the Docker executor (local fallback mode).""" | |
| def setup_method(self): | |
| """Create executor in local fallback mode.""" | |
| self.executor = DockerExecutor(use_local_fallback=True) | |
| self.executor._container_id = "local-fallback" | |
| def test_safe_command_executes(self): | |
| result = self.executor.execute("echo hello world") | |
| assert result.exit_code == 0 | |
| assert "hello world" in result.stdout | |
| assert not result.blocked | |
| def test_dangerous_command_blocked(self): | |
| result = self.executor.execute("rm -rf /") | |
| assert result.blocked | |
| assert "BLOCKED" in result.stderr | |
| def test_unknown_command_blocked(self): | |
| result = self.executor.execute("nonexistent_command_xyz") | |
| assert result.blocked | |
| def test_execution_result_fields(self): | |
| result = self.executor.execute("echo test") | |
| assert isinstance(result.stdout, str) | |
| assert isinstance(result.stderr, str) | |
| assert isinstance(result.exit_code, int) | |
| assert isinstance(result.timed_out, bool) | |
| assert isinstance(result.blocked, bool) | |
| def test_stop_container(self): | |
| """Test that stop_container doesn't crash.""" | |
| self.executor.stop_container() | |
| assert self.executor._container_id is None | |