From e52db457d1df6404d9454ed5c8ed3a64c32c4f10 Mon Sep 17 00:00:00 2001 From: mmprotest Date: Sun, 3 May 2026 21:51:59 +1000 Subject: [PATCH] Harden shell command execution preflight and output handling --- tests/test_tools_hardening.py | 50 ++++++++++++++++++++++++++- villani_code/tools.py | 63 +++++++++++++++++++++++++++++++++-- 2 files changed, 110 insertions(+), 3 deletions(-) diff --git a/tests/test_tools_hardening.py b/tests/test_tools_hardening.py index d2f0f020..e35bfb1c 100644 --- a/tests/test_tools_hardening.py +++ b/tests/test_tools_hardening.py @@ -2,7 +2,8 @@ import pytest -from villani_code.tools import GitSimpleInput, LsInput, _safe_path +from villani_code import tools +from villani_code.tools import GitSimpleInput, LsInput, _safe_path, execute_tool def test_safe_path_accepts_in_repo_path(tmp_path: Path) -> None: @@ -29,3 +30,50 @@ def test_model_defaults_do_not_share_list_state() -> None: git_two = GitSimpleInput() git_one.args.append("status") assert git_two.args == [] + + +def test_windows_sanitizer_rejects_head_pipeline(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(tools.sys, "platform", "win32") + result = execute_tool("Bash", {"command": "pytest -q 2>&1 | head -80"}, repo=tmp_path) + assert result["is_error"] is True + assert "Unix-style shell syntax/tooling" in result["content"] + + +def test_windows_sanitizer_rejects_powershell_cmdlet_without_wrapper(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(tools.sys, "platform", "win32") + result = execute_tool("Bash", {"command": "pytest -q 2>&1 | Select-Object -First 80"}, repo=tmp_path) + assert result["is_error"] is True + assert "PowerShell cmdlet used without explicit PowerShell invocation" in result["content"] + + +def test_sanitizer_rejects_large_multiline_python_c(tmp_path: Path) -> None: + source = "print('x')\\n" * 100 + result = execute_tool("Bash", {"command": f'python -c "{source}" > out.py'}, repo=tmp_path) + assert result["is_error"] is True + assert "Do not use shell quoting to write source files" in result["content"] + + +def test_sanitizer_allows_short_python_c(tmp_path: Path) -> None: + result = execute_tool("Bash", {"command": 'python -c "import sys; print(sys.version)"'}, repo=tmp_path) + assert result["is_error"] is False + + +def test_bash_output_truncates_for_model_and_keeps_debug_full(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr( + tools.subprocess, + "run", + lambda *args, **kwargs: type("P", (), {"returncode": 0, "stdout": "a" * 9000, "stderr": ""})(), + ) + events: list[dict[str, object]] = [] + result = execute_tool( + "Bash", + {"command": "echo hi"}, + repo=tmp_path, + debug_callback=lambda _name, payload: events.append(payload), + tool_call_id="t1", + ) + assert result["is_error"] is False + assert "[stdout truncated to 8000 chars]" in result["content"] + finished = events[-1] + assert finished["truncated"] is True + assert len(str(finished["stdout"])) == 9000 diff --git a/villani_code/tools.py b/villani_code/tools.py index 423807ed..c518019a 100644 --- a/villani_code/tools.py +++ b/villani_code/tools.py @@ -2,8 +2,10 @@ import glob import json +import re import shutil import subprocess +import sys from pathlib import Path from typing import Any from urllib.parse import urlparse @@ -113,6 +115,11 @@ class SubmitPlanInput(BaseModel): } DENYLIST = ["rm -rf", "del /s", "format ", "mkfs", "dd if=", "curl ", "wget "] +MODEL_OUTPUT_CHAR_LIMIT = 8000 +INLINE_PYTHON_C_LIMIT = 500 +UNIX_ONLY_WINDOWS_PATTERNS = [r"\|\s*head\b", r"\|\s*tail\b", r"\bgrep\b", r"\bsed\b", r"\bawk\b", r"\bcat\s*>", r"<<\s*\w+", r"\brm\s+-rf\b", r"\btouch\b", r"\bpwd\b"] +POWERSHELL_CMDLETS = [r"\bSelect-Object\b", r"\bGet-Content\b", r"\bSet-Content\b", r"\bOut-File\b", r"\bWhere-Object\b", r"\bForEach-Object\b"] +RECOVERY_HINT = "Use Bash only for simple commands. Use Write/Patch for file edits. Do not use shell quoting to write source files." def _error(message: str) -> dict[str, Any]: @@ -236,6 +243,7 @@ def _run_search(data: SearchInput, repo: Path) -> str: def _run_bash(data: BashInput, repo: Path, unsafe: bool, debug_callback: Any | None = None, tool_call_id: str = "") -> str: + _preflight_shell_command(data.command, platform_name=sys.platform) lowered = data.command.lower() if not unsafe: for bad in DENYLIST: @@ -245,6 +253,11 @@ def _run_bash(data: BashInput, repo: Path, unsafe: bool, debug_callback: Any | N if callable(debug_callback): debug_callback("command_started", {"command": data.command, "cwd": data.cwd, "tool_call_id": tool_call_id}) proc = subprocess.run(data.command, shell=True, cwd=str(cwd), capture_output=True, text=True, timeout=data.timeout_sec) + stdout_model, stdout_truncated = _truncate_for_model(proc.stdout, "stdout") + stderr_model, stderr_truncated = _truncate_for_model(proc.stderr, "stderr") + has_shell_failure = _looks_like_shell_syntax_or_portability_failure(proc.stderr, proc.stdout) + if has_shell_failure: + stderr_model = f"{stderr_model.rstrip()}\n\nHint: {RECOVERY_HINT}\n" if callable(debug_callback): debug_callback( "command_finished", @@ -254,11 +267,57 @@ def _run_bash(data: BashInput, repo: Path, unsafe: bool, debug_callback: Any | N "exit_code": proc.returncode, "stdout": proc.stdout, "stderr": proc.stderr, - "truncated": False, + "truncated": stdout_truncated or stderr_truncated, "tool_call_id": tool_call_id, }, ) - return json.dumps({"command": data.command, "exit_code": proc.returncode, "stdout": proc.stdout, "stderr": proc.stderr}, indent=2) + return json.dumps({"command": data.command, "exit_code": proc.returncode, "stdout": stdout_model, "stderr": stderr_model}, indent=2) + + +def _preflight_shell_command(command: str, platform_name: str) -> None: + lowered = command.lower() + if _is_brittle_inline_source_command(command): + raise ValueError(f"Rejected shell command: brittle inline source writing detected. {RECOVERY_HINT}") + is_windows = platform_name.lower().startswith("win") + if not is_windows: + return + if _contains_any_pattern(command, UNIX_ONLY_WINDOWS_PATTERNS): + raise ValueError(f"Rejected shell command: Unix-style shell syntax/tooling is not portable in Windows cmd.exe. {RECOVERY_HINT}") + if _contains_any_pattern(command, POWERSHELL_CMDLETS) and "powershell" not in lowered: + raise ValueError("Rejected shell command: PowerShell cmdlet used without explicit PowerShell invocation (cmd.exe context).") + + +def _contains_any_pattern(command: str, patterns: list[str]) -> bool: + return any(re.search(pattern, command, flags=re.IGNORECASE) for pattern in patterns) + + +def _is_brittle_inline_source_command(command: str) -> bool: + if "@'" in command or "'@" in command or '@"' in command or '"@' in command: + return True + if re.search(r"\bpython(\d+(\.\d+)*)?\s+-c\b", command, flags=re.IGNORECASE): + if "\n" in command or len(command) > INLINE_PYTHON_C_LIMIT: + return True + return bool(re.search(r"(echo|printf|python\s+-c).{200,}(>|>>)\s*\S+", command, flags=re.IGNORECASE)) + + +def _truncate_for_model(text: str, stream_name: str, limit: int = MODEL_OUTPUT_CHAR_LIMIT) -> tuple[str, bool]: + if len(text) <= limit: + return text, False + return f"{text[:limit]}\n[{stream_name} truncated to {limit} chars]", True + + +def _looks_like_shell_syntax_or_portability_failure(stderr: str, stdout: str) -> bool: + combined = f"{stdout}\n{stderr}".lower() + patterns = [ + "not recognized as an internal or external command", + "is not recognized", + "syntax error", + "unterminated", + "unexpected eof", + "here-string", + "the string is missing the terminator", + ] + return any(p in combined for p in patterns) def _run_write(data: WriteInput, repo: Path, debug_callback: Any | None = None, tool_call_id: str = "") -> str: