From 9ddb1c7b776093d0584848f69bdb4f9a9697be61 Mon Sep 17 00:00:00 2001 From: SecurePatchOps Date: Sun, 28 Jun 2026 21:36:00 +0800 Subject: [PATCH] Harden react agent loop tool execution --- skills/react-agent-loop/SKILL.md | 11 ++- skills/react-agent-loop/scripts/agent_loop.py | 79 ++++++++++++++++--- skills/react-agent-loop/tests/test_loop.py | 46 +++++++++++ 3 files changed, 125 insertions(+), 11 deletions(-) diff --git a/skills/react-agent-loop/SKILL.md b/skills/react-agent-loop/SKILL.md index c278fec7..edd6153a 100644 --- a/skills/react-agent-loop/SKILL.md +++ b/skills/react-agent-loop/SKILL.md @@ -70,14 +70,21 @@ from skills.provider_factory.scripts.provider_factory import ProviderFactory llm = ProviderFactory.from_env("LLM_") tools = ToolRegistry() -tools.register(FileWriterTool()) -tools.register(BashTool()) +tools.register(FileWriterTool(base_dir="workspace")) + +# Registre BashTool apenas para entradas confiaveis e com allowlist explicita. +# tools.register(BashTool(allowed_commands={"ls", "pwd"}, base_dir="workspace")) loop = AgentLoop(llm_client=llm, tool_registry=tools, max_iterations=5) answer = loop.run("Liste os arquivos .py no diretorio atual") print(answer) ``` +`FileWriterTool` rejeita caminhos absolutos, path traversal e sobrescrita por +padrao. `BashTool` nao permite nenhum comando por padrao; integracoes devem +fornecer uma allowlist minima de comandos e manter o `base_dir` restrito ao +workspace da tarefa. + ## Ficheiros - `scripts/agent_loop.py` — implementacao completa (260 linhas) - `AgentLoop` — engine ReAct com hard limit diff --git a/skills/react-agent-loop/scripts/agent_loop.py b/skills/react-agent-loop/scripts/agent_loop.py index ecc0fcd8..2fe0f553 100644 --- a/skills/react-agent-loop/scripts/agent_loop.py +++ b/skills/react-agent-loop/scripts/agent_loop.py @@ -21,8 +21,11 @@ import json import time import logging +import shlex +import subprocess from abc import ABC, abstractmethod from dataclasses import dataclass, field +from pathlib import Path from typing import Any, Protocol logger = logging.getLogger(__name__) @@ -31,6 +34,20 @@ DEFAULT_TOOL_TIMEOUT_S = 60 +def _resolve_relative_path(base_dir: Path, raw_path: str, label: str) -> Path: + path = Path(raw_path) + if path.is_absolute(): + raise ValueError(f"{label} deve ser relativo ao diretorio base") + + resolved = (base_dir / path).resolve() + try: + resolved.relative_to(base_dir) + except ValueError as exc: + raise ValueError(f"{label} fora do diretorio base") from exc + + return resolved + + @dataclass class ToolCall: name: str @@ -112,12 +129,27 @@ def execute(self, call: ToolCall) -> ToolResult: tool = self._tools.get(call.name) if tool is None: return ToolResult.failure(f"Tool desconhecida: {call.name}") + validation_error = self._validate_arguments(tool, call.arguments) + if validation_error: + return ToolResult.failure(validation_error) try: return tool.execute(**call.arguments) except Exception as exc: logger.exception("Tool '%s' lancou excecao", call.name) return ToolResult.failure(str(exc)) + def _validate_arguments(self, tool: BaseTool, arguments: Any) -> str | None: + if not isinstance(arguments, dict): + return f"Argumentos invalidos para tool '{tool.name}': objeto esperado" + + required = tool.parameters_schema.get("required", []) + for name in required: + value = arguments.get(name) + if value is None or (isinstance(value, str) and not value.strip()): + return f"Argumento obrigatorio ausente para tool '{tool.name}': {name}" + + return None + def get_system_prompt_addendum(self) -> str: if not self._tools: return "" @@ -262,13 +294,17 @@ def _execute_tool( class FileWriterTool(BaseTool): """Tool de exemplo: escreve arquivo no filesystem.""" + def __init__(self, base_dir: str | Path = "workspace", overwrite: bool = False) -> None: + self._base_dir = Path(base_dir).resolve() + self._overwrite = overwrite + @property def name(self) -> str: return "write_file" @property def description(self) -> str: - return "Cria ou sobrescreve um arquivo com o conteudo especificado" + return "Cria um arquivo dentro do diretorio base configurado" @property def parameters_schema(self) -> dict[str, Any]: @@ -283,8 +319,10 @@ def parameters_schema(self) -> dict[str, Any]: def execute(self, path: str = "", content: str = "", **_: Any) -> ToolResult: try: - from pathlib import Path - p = Path(path) + p = _resolve_relative_path(self._base_dir, path, "Caminho do arquivo") + if p.exists() and not self._overwrite: + return ToolResult.failure(f"Arquivo '{path}' ja existe") + p.parent.mkdir(parents=True, exist_ok=True) p.write_text(content, encoding="utf-8") return ToolResult(output=f"Arquivo '{path}' criado com sucesso ({len(content)} bytes)") @@ -295,13 +333,21 @@ def execute(self, path: str = "", content: str = "", **_: Any) -> ToolResult: class BashTool(BaseTool): """Tool de exemplo: executa comando shell.""" + def __init__( + self, + allowed_commands: set[str] | None = None, + base_dir: str | Path = ".", + ) -> None: + self._allowed_commands = allowed_commands or set() + self._base_dir = Path(base_dir).resolve() + @property def name(self) -> str: return "bash" @property def description(self) -> str: - return "Executa um comando shell e retorna a saida" + return "Executa um comando permitido sem shell e retorna a saida" @property def parameters_schema(self) -> dict[str, Any]: @@ -315,18 +361,33 @@ def parameters_schema(self) -> dict[str, Any]: } def execute(self, command: str = "", workdir: str = "", **_: Any) -> ToolResult: - import subprocess try: + argv = shlex.split(command) + if not argv: + return ToolResult.failure("Comando vazio") + + executable = argv[0] + if executable not in self._allowed_commands: + return ToolResult.failure(f"Comando nao permitido: {executable}") + + cwd = self._base_dir + if workdir: + cwd = _resolve_relative_path(self._base_dir, workdir, "Diretorio de trabalho") + if not cwd.is_dir(): + return ToolResult.failure(f"Diretorio de trabalho invalido: {workdir or self._base_dir}") + result = subprocess.run( - command, - shell=True, + argv, + shell=False, capture_output=True, text=True, timeout=DEFAULT_TOOL_TIMEOUT_S, - cwd=workdir or None, + cwd=cwd, ) output = result.stdout.strip() or result.stderr.strip() - return ToolResult(output=output, success=result.returncode == 0) + if result.returncode != 0: + return ToolResult(output=output, success=False, error=output) + return ToolResult(output=output) except subprocess.TimeoutExpired: return ToolResult.failure(f"Comando excedeu timeout de {DEFAULT_TOOL_TIMEOUT_S}s") except Exception as exc: diff --git a/skills/react-agent-loop/tests/test_loop.py b/skills/react-agent-loop/tests/test_loop.py index 71fec76e..f8197374 100644 --- a/skills/react-agent-loop/tests/test_loop.py +++ b/skills/react-agent-loop/tests/test_loop.py @@ -103,6 +103,52 @@ def test_execute_tool(self): assert result_unknown.success is False assert "desconhecida" in result_unknown.error.lower() or "Tool desconhecida" in result_unknown.error + def test_execute_tool_validates_required_arguments(self): + tools = ToolRegistry() + tools.register(EchoTool()) + + result = tools.execute(ToolCall(name="echo", arguments={})) + assert result.success is False + assert "obrigatorio" in result.error.lower() + + def test_file_writer_stays_inside_base_dir(self, tmp_path): + tool = FileWriterTool(base_dir=tmp_path) + + result = tool.execute(path="notes/result.txt", content="ok") + assert result.success is True + assert (tmp_path / "notes" / "result.txt").read_text(encoding="utf-8") == "ok" + + traversal = tool.execute(path="../outside.txt", content="bad") + assert traversal.success is False + assert "diretorio base" in traversal.error + + absolute = tool.execute(path=str(tmp_path.parent / "outside.txt"), content="bad") + assert absolute.success is False + assert "relativo" in absolute.error + + def test_file_writer_rejects_overwrite_by_default(self, tmp_path): + tool = FileWriterTool(base_dir=tmp_path) + target = tmp_path / "existing.txt" + target.write_text("original", encoding="utf-8") + + result = tool.execute(path="existing.txt", content="new") + assert result.success is False + assert target.read_text(encoding="utf-8") == "original" + + def test_bash_tool_denies_commands_by_default(self): + result = BashTool().execute(command="echo hello") + + assert result.success is False + assert "nao permitido" in result.error + + def test_bash_tool_runs_allowlisted_command_without_shell(self, tmp_path): + tool = BashTool(allowed_commands={"echo"}, base_dir=tmp_path) + + result = tool.execute(command="echo hello") + + assert result.success is True + assert result.output == "hello" + def test_max_iterations(self): tool_responses = [] for i in range(6):