From 761db2518c9a2aac6cb92a59d64ac0db657fca9d Mon Sep 17 00:00:00 2001 From: Brent Edwards Date: Thu, 12 Feb 2026 01:07:44 +0000 Subject: [PATCH] test: harden code sandbox, add path traversal protection, and fix test mocks --- .forgejo/workflows/ci.yml | 6 +- .pre-commit-config.yaml | 4 +- features/edge_case_plan_scenarios.feature | 40 +++++ features/steps/actor_cli_steps.py | 47 +++--- features/steps/edge_case_plan_steps.py | 98 ++++++++++++- .../steps/validation_test_fixture_steps.py | 45 +++++- features/validation_test_fixtures.feature | 68 +++++++++ pyproject.toml | 1 + scripts/run-semgrep.sh | 7 + .../application/services/plan_service.py | 32 ++-- src/cleveragents/reactive/application.py | 6 +- src/cleveragents/reactive/stream_router.py | 138 +++++++++++++++--- 12 files changed, 424 insertions(+), 68 deletions(-) create mode 100755 scripts/run-semgrep.sh diff --git a/.forgejo/workflows/ci.yml b/.forgejo/workflows/ci.yml index f76ebb48c..205a5d4d6 100644 --- a/.forgejo/workflows/ci.yml +++ b/.forgejo/workflows/ci.yml @@ -2,9 +2,9 @@ name: CI on: push: - branches: [main, develop] + branches: [master, develop] pull_request: - branches: [main] + branches: [master] env: UV_VERSION: "0.8.0" @@ -157,7 +157,7 @@ jobs: - name: Run tests with coverage run: | - coverage run --source=src -m behave -q --no-capture || true + coverage run --source=src -m behave -q --no-capture coverage xml -o build/coverage.xml coverage report --fail-under=85 diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 8be770037..52d19c3c2 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -12,8 +12,8 @@ repos: rev: v4.6.0 hooks: - id: no-commit-to-branch - name: Prevent direct commits to main - args: ["--branch", "main"] + name: Prevent direct commits to master + args: ["--branch", "master"] - id: check-yaml name: Validate YAML syntax - id: check-toml diff --git a/features/edge_case_plan_scenarios.feature b/features/edge_case_plan_scenarios.feature index 81c7910c2..d85647ed8 100644 --- a/features/edge_case_plan_scenarios.feature +++ b/features/edge_case_plan_scenarios.feature @@ -122,6 +122,46 @@ Feature: Plan edge case scenarios When I apply the edge case changes Then the deeply nested file should exist + # ────────────────────────────────────────────────── + # Section 2b: Path traversal rejection (H4) + # ────────────────────────────────────────────────── + + Scenario: Apply CREATE change with relative path traversal is rejected + Given I have a temporary test directory for edge case plan service + And I have a Unit of Work instance for edge case plan testing + And I have an edge case PlanService instance + And I have a saved project with current plan for edge cases + And the plan has a CREATE change with a path traversal in file_path + When I try to apply the edge case changes + Then a PlanError should be raised mentioning path traversal + + Scenario: Apply MODIFY change with absolute path outside project is rejected + Given I have a temporary test directory for edge case plan service + And I have a Unit of Work instance for edge case plan testing + And I have an edge case PlanService instance + And I have a saved project with current plan for edge cases + And the plan has a MODIFY change with an absolute path outside the project + When I try to apply the edge case changes + Then a PlanError should be raised mentioning path traversal + + Scenario: Apply DELETE change with path traversal is rejected + Given I have a temporary test directory for edge case plan service + And I have a Unit of Work instance for edge case plan testing + And I have an edge case PlanService instance + And I have a saved project with current plan for edge cases + And the plan has a DELETE change with a path traversal in file_path + When I try to apply the edge case changes + Then a PlanError should be raised mentioning path traversal + + Scenario: Apply MOVE change with path traversal in new_path is rejected + Given I have a temporary test directory for edge case plan service + And I have a Unit of Work instance for edge case plan testing + And I have an edge case PlanService instance + And I have a saved project with current plan for edge cases + And the plan has a MOVE change with a path traversal in new_path + When I try to apply the edge case changes + Then a PlanError should be raised mentioning path traversal + # ────────────────────────────────────────────────── # Section 3: Validation failure chains # ────────────────────────────────────────────────── diff --git a/features/steps/actor_cli_steps.py b/features/steps/actor_cli_steps.py index 36b248f31..561d333a9 100644 --- a/features/steps/actor_cli_steps.py +++ b/features/steps/actor_cli_steps.py @@ -177,10 +177,10 @@ def step_impl(context): @when("I run actor add without config") def step_impl(context): - with patch("cleveragents.application.container.get_container") as mock_container: + with patch("cleveragents.cli.commands.actor._get_services") as mock_get_services: + mock_actor_service = MagicMock() mock_actor_registry = MagicMock() - mock_container.return_value.actor_service.return_value = MagicMock() - mock_container.return_value.actor_registry.return_value = mock_actor_registry + mock_get_services.return_value = (mock_actor_service, mock_actor_registry) context.result = context.runner.invoke(actor_app, ["add", "test-actor"]) context.mock_actor_registry = mock_actor_registry context.expected_error = "Config file is required" @@ -375,9 +375,8 @@ def step_impl(context): @when("I run actor add with malformed option") def step_impl(context): - with patch("cleveragents.application.container.get_container") as mock_container: - mock_container.return_value.actor_service.return_value = MagicMock() - mock_container.return_value.actor_registry.return_value = MagicMock() + with patch("cleveragents.cli.commands.actor._get_services") as mock_get_services: + mock_get_services.return_value = (MagicMock(), MagicMock()) context.result = context.runner.invoke( actor_app, [ @@ -394,9 +393,8 @@ def step_impl(context): @when("I run actor add with empty option key") def step_impl(context): - with patch("cleveragents.application.container.get_container") as mock_container: - mock_container.return_value.actor_service.return_value = MagicMock() - mock_container.return_value.actor_registry.return_value = MagicMock() + with patch("cleveragents.cli.commands.actor._get_services") as mock_get_services: + mock_get_services.return_value = (MagicMock(), MagicMock()) context.result = context.runner.invoke( actor_app, [ @@ -495,9 +493,8 @@ def step_impl(context): @when("I run actor add with missing config path") def step_impl(context): missing_path = Path(tempfile.gettempdir()) / "missing-actor-config.json" - with patch("cleveragents.application.container.get_container") as mock_container: - mock_container.return_value.actor_service.return_value = MagicMock() - mock_container.return_value.actor_registry.return_value = MagicMock() + with patch("cleveragents.cli.commands.actor._get_services") as mock_get_services: + mock_get_services.return_value = (MagicMock(), MagicMock()) context.result = context.runner.invoke( actor_app, [ @@ -521,11 +518,11 @@ def step_impl(context): config_path = Path(handle.name) _register_cleanup(context, config_path) - with patch("cleveragents.application.container.get_container") as mock_container: + with patch("cleveragents.cli.commands.actor._get_services") as mock_get_services: + mock_actor_service = MagicMock() mock_actor_registry = MagicMock() mock_actor_registry.upsert_actor.side_effect = BusinessRuleViolation("invalid") - mock_container.return_value.actor_service.return_value = MagicMock() - mock_container.return_value.actor_registry.return_value = mock_actor_registry + mock_get_services.return_value = (mock_actor_service, mock_actor_registry) context.result = context.runner.invoke( actor_app, @@ -672,14 +669,13 @@ def step_impl(context): @when("I run actor update with validation error") def step_impl(context): - with patch("cleveragents.application.container.get_container") as mock_container: + with patch("cleveragents.cli.commands.actor._get_services") as mock_get_services: mock_actor_service = MagicMock() mock_actor_registry = MagicMock() current_actor = _make_actor(name="local/invalid-update") mock_actor_registry.get_actor.return_value = current_actor mock_actor_registry.upsert_actor.side_effect = ValidationError("bad update") - mock_container.return_value.actor_service.return_value = mock_actor_service - mock_container.return_value.actor_registry.return_value = mock_actor_registry + mock_get_services.return_value = (mock_actor_service, mock_actor_registry) context.result = context.runner.invoke( actor_app, @@ -958,12 +954,11 @@ def step_impl(context): @when("I run actor remove and it fails validation") def step_impl(context): - with patch("cleveragents.application.container.get_container") as mock_container: + with patch("cleveragents.cli.commands.actor._get_services") as mock_get_services: mock_actor_service = MagicMock() mock_actor_registry = MagicMock() mock_actor_registry.remove_actor.side_effect = ValidationError("invalid") - mock_container.return_value.actor_service.return_value = mock_actor_service - mock_container.return_value.actor_registry.return_value = mock_actor_registry + mock_get_services.return_value = (mock_actor_service, mock_actor_registry) context.result = context.runner.invoke(actor_app, ["remove", "local/removable"]) context.mock_actor_registry = mock_actor_registry @@ -1023,12 +1018,11 @@ def step_impl(context): @when("I run actor show with validation error") def step_impl(context): - with patch("cleveragents.application.container.get_container") as mock_container: + with patch("cleveragents.cli.commands.actor._get_services") as mock_get_services: mock_actor_service = MagicMock() mock_actor_registry = MagicMock() mock_actor_registry.get_actor.side_effect = ValidationError("bad") - mock_container.return_value.actor_service.return_value = mock_actor_service - mock_container.return_value.actor_registry.return_value = mock_actor_registry + mock_get_services.return_value = (mock_actor_service, mock_actor_registry) context.result = context.runner.invoke(actor_app, ["show", "local/show"]) context.expected_error = "Error:" @@ -1049,14 +1043,13 @@ def step_impl(context): @when("I run set-default actor with violation") def step_impl(context): - with patch("cleveragents.application.container.get_container") as mock_container: + with patch("cleveragents.cli.commands.actor._get_services") as mock_get_services: mock_actor_service = MagicMock() mock_actor_registry = MagicMock() mock_actor_registry.set_default_actor.side_effect = BusinessRuleViolation( "nope" ) - mock_container.return_value.actor_service.return_value = mock_actor_service - mock_container.return_value.actor_registry.return_value = mock_actor_registry + mock_get_services.return_value = (mock_actor_service, mock_actor_registry) context.result = context.runner.invoke(actor_app, ["set-default", "local/fail"]) context.expected_error = "Error:" diff --git a/features/steps/edge_case_plan_steps.py b/features/steps/edge_case_plan_steps.py index 7b9a3f523..441f77d2a 100644 --- a/features/steps/edge_case_plan_steps.py +++ b/features/steps/edge_case_plan_steps.py @@ -4,6 +4,7 @@ from __future__ import annotations import contextlib import os +import shutil import tempfile from datetime import datetime from pathlib import Path @@ -244,7 +245,9 @@ def step_plan_b_apply(context: Context) -> None: @given("I have a temporary test directory for edge case plan service") def step_create_temp_dir_edge(context: Context) -> None: context.edge_temp_dir = Path(tempfile.mkdtemp(prefix="edge_case_plan_")) - context._cleanup_handlers.append(lambda: contextlib.suppress(Exception) or None) + context._cleanup_handlers.append( + lambda d=context.edge_temp_dir: shutil.rmtree(d, ignore_errors=True) + ) @given("I have a Unit of Work instance for edge case plan testing") @@ -539,6 +542,99 @@ def step_check_deeply_nested(context: Context) -> None: assert file_path.read_text() == "# deeply nested file" +# ────────────────────────────────────────────────── +# Section 2b: Path traversal rejection (H4) +# ────────────────────────────────────────────────── + + +@given("the plan has a CREATE change with a path traversal in file_path") +def step_add_create_traversal(context: Context) -> None: + with context.edge_uow.transaction() as ctx: + change = Change( + id=None, + plan_id=context.edge_plan.id, + file_path="../../etc/traversal_test", + operation=OperationType.CREATE, + original_content=None, + new_content="# malicious payload", + new_path=None, + applied=False, + applied_at=None, + created_at=datetime.now(), + ) + ctx.changes.add(change) + + +@given("the plan has a MODIFY change with an absolute path outside the project") +def step_add_modify_absolute_escape(context: Context) -> None: + with context.edge_uow.transaction() as ctx: + change = Change( + id=None, + plan_id=context.edge_plan.id, + file_path="/tmp/outside_project.py", + operation=OperationType.MODIFY, + original_content=None, + new_content="# escaped content", + new_path=None, + applied=False, + applied_at=None, + created_at=datetime.now(), + ) + ctx.changes.add(change) + + +@given("the plan has a DELETE change with a path traversal in file_path") +def step_add_delete_traversal(context: Context) -> None: + with context.edge_uow.transaction() as ctx: + change = Change( + id=None, + plan_id=context.edge_plan.id, + file_path="subdir/../.././../etc/passwd", + operation=OperationType.DELETE, + original_content=None, + new_content=None, + new_path=None, + applied=False, + applied_at=None, + created_at=datetime.now(), + ) + ctx.changes.add(change) + + +@given("the plan has a MOVE change with a path traversal in new_path") +def step_add_move_traversal_new_path(context: Context) -> None: + # Create a safe source file so the MOVE actually attempts the rename + source = context.edge_temp_dir / "safe_source.py" + source.write_text("# safe source") + + with context.edge_uow.transaction() as ctx: + change = Change( + id=None, + plan_id=context.edge_plan.id, + file_path="safe_source.py", + operation=OperationType.MOVE, + original_content=None, + new_content=None, + new_path="../../escaped.py", + applied=False, + applied_at=None, + created_at=datetime.now(), + ) + ctx.changes.add(change) + + +@then("a PlanError should be raised mentioning path traversal") +def step_check_path_traversal_error(context: Context) -> None: + assert context.edge_exception is not None, "Expected a PlanError" + assert isinstance( + context.edge_exception, PlanError + ), f"Expected PlanError, got {type(context.edge_exception).__name__}" + msg = str(context.edge_exception.message).lower() + assert ( + "path traversal" in msg + ), f"Expected 'path traversal' in error message, got: {context.edge_exception.message}" + + # ────────────────────────────────────────────────── # Section 3: Validation failure chains # ────────────────────────────────────────────────── diff --git a/features/steps/validation_test_fixture_steps.py b/features/steps/validation_test_fixture_steps.py index 3fd806162..82e8565b4 100644 --- a/features/steps/validation_test_fixture_steps.py +++ b/features/steps/validation_test_fixture_steps.py @@ -23,7 +23,11 @@ from cleveragents.domain.models.core.action import ( ArgumentRequirement, ArgumentType, ) -from cleveragents.reactive.stream_router import _validate_code_ast, _validate_lambda_ast +from cleveragents.reactive.stream_router import ( + _compile_restricted_code, + _validate_code_ast, + _validate_lambda_ast, +) # ────────────────────────────────────────────────── # Section 1: AST security validation for code @@ -55,6 +59,45 @@ def step_no_ast_error(context: Context) -> None: assert context.ast_error is None, f"Unexpected error: {context.ast_error}" +# ────────────────────────────────────────────────── +# Section 1b: RestrictedPython - dunder / subscript bypass vectors (C2) +# ────────────────────────────────────────────────── + + +@when('I compile restricted code with "{code}"') +def step_compile_restricted_code(context: Context, code: str) -> None: + actual_code = code.replace("\\n", "\n") + context.ast_error = None + try: + _compile_restricted_code(actual_code) + except StreamRoutingError as exc: + context.ast_error = exc + + +@then("no error should be raised from restricted compilation") +def step_no_restricted_error(context: Context) -> None: + assert context.ast_error is None, f"Unexpected error: {context.ast_error}" + + +@when('I exec restricted code with "{code}"') +def step_exec_restricted_code(context: Context, code: str) -> None: + """Compile *and* execute restricted code; capture any runtime error.""" + from RestrictedPython import safe_globals as _sg + + actual_code = code.replace("\\n", "\n") + context.runtime_error = None + try: + compiled = _compile_restricted_code(actual_code) + exec(compiled, _sg.copy(), {}) # nosec B102 + except (StreamRoutingError, Exception) as exc: + context.runtime_error = exc + + +@then("a runtime error should be raised") +def step_runtime_error_raised(context: Context) -> None: + assert context.runtime_error is not None, "Expected a runtime error" + + # ────────────────────────────────────────────────── # Section 2: Lambda AST validation # ────────────────────────────────────────────────── diff --git a/features/validation_test_fixtures.feature b/features/validation_test_fixtures.feature index 9704a0496..a90c8723b 100644 --- a/features/validation_test_fixtures.feature +++ b/features/validation_test_fixtures.feature @@ -55,6 +55,44 @@ Feature: Validation test fixtures When I validate code ast with "result = 1 + 2" Then no error should be raised from ast validation + # ────────────────────────────────────────────────── + # Section 1b: RestrictedPython – dunder / subscript bypass vectors (C2) + # These verify that _compile_restricted_code blocks the attack patterns + # that the previous hand-rolled AST validator could not catch. + # ────────────────────────────────────────────────── + + Scenario: Reject dunder attribute chain via __class__.__bases__ + When I compile restricted code with "x = ().__class__.__bases__[0].__subclasses__()" + Then a StreamRoutingError should be raised mentioning "Forbidden construct" + + Scenario: Reject dunder attribute access via __class__ + When I compile restricted code with "x = ''.__class__" + Then a StreamRoutingError should be raised mentioning "Forbidden construct" + + Scenario: Reject subscript access to __builtins__ + When I compile restricted code with "x = __builtins__['__import__']('os')" + Then a StreamRoutingError should be raised mentioning "Forbidden" + + Scenario: Reject dunder __init__ attribute access + When I compile restricted code with "x = ().__class__.__bases__[0].__subclasses__()[0].__init__.__globals__" + Then a StreamRoutingError should be raised mentioning "Forbidden construct" + + Scenario: Reject attribute-based getattr bypass at runtime + When I exec restricted code with "f = getattr; f([], '__class__')" + Then a runtime error should be raised + + Scenario: Reject __import__ via name mangling + When I compile restricted code with "f = __import__; f('os')" + Then a StreamRoutingError should be raised mentioning "Forbidden" + + Scenario: Accept safe string operation through restricted compilation + When I compile restricted code with "result = 'hello'.upper()" + Then no error should be raised from restricted compilation + + Scenario: Accept safe arithmetic through restricted compilation + When I compile restricted code with "result = sum(range(10))" + Then no error should be raised from restricted compilation + # ────────────────────────────────────────────────── # Section 2: Invalid code samples - Lambda AST validation # ────────────────────────────────────────────────── @@ -75,6 +113,36 @@ Feature: Validation test fixtures When I validate lambda ast with "print('hello')" Then a StreamRoutingError should be raised mentioning "must be a lambda" + # ────────────────────────────────────────────────── + # Section 2b: Lambda body – forbidden call detection (H2) + # These verify that _validate_lambda_ast walks the lambda body and + # rejects calls to dangerous built-ins that RestrictedPython misses. + # ────────────────────────────────────────────────── + + Scenario: Reject lambda calling compile in body + When I validate lambda ast with "lambda x: compile('x', '', 'exec')" + Then a StreamRoutingError should be raised mentioning "Forbidden call" + + Scenario: Reject lambda calling getattr in body + When I validate lambda ast with "lambda x: getattr(x, '__class__')" + Then a StreamRoutingError should be raised mentioning "Forbidden call" + + Scenario: Reject lambda calling setattr in body + When I validate lambda ast with "lambda x: setattr(x, 'k', 'v')" + Then a StreamRoutingError should be raised mentioning "Forbidden call" + + Scenario: Reject lambda calling eval in body + When I validate lambda ast with "lambda x: eval('1+1')" + Then a StreamRoutingError should be raised mentioning "Forbidden call" + + Scenario: Reject lambda calling __import__ in body + When I validate lambda ast with "lambda x: __import__('os')" + Then a StreamRoutingError should be raised mentioning "Forbidden call" + + Scenario: Accept safe lambda with method call + When I validate lambda ast with "lambda x: x.upper()" + Then no error should be raised from lambda validation + # ────────────────────────────────────────────────── # Section 3: Invalid code samples - Python content sanitization # ────────────────────────────────────────────────── diff --git a/pyproject.toml b/pyproject.toml index 24e213f13..1c3fc702a 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -42,6 +42,7 @@ dependencies = [ "alembic>=1.13.1", "numpy>=2.1.0", "python-ulid>=2.7.0", # ULID generation for plan/action IDs + "RestrictedPython>=7.0", # Secure sandbox for user-supplied code ] diff --git a/scripts/run-semgrep.sh b/scripts/run-semgrep.sh new file mode 100755 index 000000000..81247299b --- /dev/null +++ b/scripts/run-semgrep.sh @@ -0,0 +1,7 @@ +#!/bin/sh +# Run semgrep if installed, skip gracefully if not +if command -v semgrep >/dev/null 2>&1; then + semgrep --config=.semgrep.yml --error --quiet src/ +else + echo "semgrep not installed, skipping (install with: pip install semgrep)" +fi diff --git a/src/cleveragents/application/services/plan_service.py b/src/cleveragents/application/services/plan_service.py index 890188211..8c3b9eb4b 100644 --- a/src/cleveragents/application/services/plan_service.py +++ b/src/cleveragents/application/services/plan_service.py @@ -968,14 +968,28 @@ class PlanService: changes = ctx.changes.get_for_plan(current_plan.id) pending_changes = [c for c in changes if not c.applied] + project_root = project.path.resolve() + + def _safe_resolve(raw_path: str) -> Path: + """Resolve *raw_path* and reject it if it escapes *project_root*.""" + if Path(raw_path).is_absolute(): + resolved = Path(raw_path).resolve() + else: + resolved = (project.path / raw_path).resolve() + if not resolved.is_relative_to(project_root): + raise PlanError( + message=( + f"Path traversal rejected: {raw_path} resolves " + f"outside project directory" + ), + details={"file": raw_path}, + ) + return resolved + applied_count = 0 for change in pending_changes: try: - # Make file paths relative to project path - if Path(change.file_path).is_absolute(): - file_path = Path(change.file_path) - else: - file_path = project.path / change.file_path + file_path = _safe_resolve(change.file_path) if change.operation == OperationType.CREATE: file_path.parent.mkdir(parents=True, exist_ok=True) @@ -987,11 +1001,7 @@ class PlanService: file_path.unlink() elif change.operation == OperationType.MOVE: if change.new_path: - # Handle new_path similarly - if Path(change.new_path).is_absolute(): - new_path = Path(change.new_path) - else: - new_path = project.path / change.new_path + new_path = _safe_resolve(change.new_path) if file_path.exists(): new_path.parent.mkdir(parents=True, exist_ok=True) file_path.rename(new_path) @@ -1000,6 +1010,8 @@ class PlanService: ctx.changes.mark_applied(change.id) applied_count += 1 + except PlanError: + raise # Re-raise path traversal and other PlanErrors directly except Exception as e: # Handle operation as either enum or string operation_str = ( diff --git a/src/cleveragents/reactive/application.py b/src/cleveragents/reactive/application.py index ae6018bf6..15d444841 100644 --- a/src/cleveragents/reactive/application.py +++ b/src/cleveragents/reactive/application.py @@ -94,7 +94,7 @@ class ReactiveCleverAgentsApp: def _make_agent_instance(name: str, agent_cfg: Any) -> Any: tools = agent_cfg.config.get("tools", []) if agent_cfg.config else [] if tools: - return SimpleToolAgent(tools) + return SimpleToolAgent(tools, unsafe=self.unsafe) if agent_cfg.type == "llm": return SimpleLLMAgent(name, agent_cfg.config) return lambda msg: getattr(msg, "content", msg) @@ -407,7 +407,7 @@ class ReactiveCleverAgentsApp: if target_node == "workflow_controller": workflow_hops += 1 - if isinstance(agent, (SimpleToolAgent, SimpleLLMAgent)): + if isinstance(agent, (SimpleToolAgent | SimpleLLMAgent)): result = agent.process(extracted_message, context=context) elif hasattr(agent, "process"): result = agent.process(extracted_message) @@ -438,7 +438,7 @@ class ReactiveCleverAgentsApp: next_actor = node_actor_map.get(next_node, next_node) next_agent = agents.get(next_actor) if next_agent is not None: - if isinstance(next_agent, (SimpleToolAgent, SimpleLLMAgent)): + if isinstance(next_agent, (SimpleToolAgent | SimpleLLMAgent)): result = next_agent.process( current_message, context=context ) diff --git a/src/cleveragents/reactive/stream_router.py b/src/cleveragents/reactive/stream_router.py index 4daa6f7b1..14b8cd27b 100644 --- a/src/cleveragents/reactive/stream_router.py +++ b/src/cleveragents/reactive/stream_router.py @@ -14,6 +14,8 @@ from typing import Any import rx from pydantic import BaseModel, Field +from RestrictedPython import compile_restricted_eval, compile_restricted_exec +from RestrictedPython import safe_globals as _rp_safe_globals from rx import operators as ops from rx.core import Observable as ObservableType # type: ignore[attr-defined] from rx.scheduler.eventloop import AsyncIOScheduler # type: ignore[attr-defined] @@ -89,8 +91,18 @@ class StreamConfig(BaseModel): # pylint: disable=too-many-instance-attributes # --------------------------------------------------------------------------- # AST validation helpers for safe dynamic code execution # --------------------------------------------------------------------------- +# RestrictedPython replaces the hand-rolled AST blocklist. It performs a full +# AST transformation that rejects dunder attribute access, subscript-based +# jail-breaks, exec/eval calls, and more. The compiled bytecode is then +# executed against ``safe_globals`` which supplies only whitelisted builtins. +# --------------------------------------------------------------------------- -# AST node types that are forbidden in user-supplied code blocks. +logger_sr = logging.getLogger(__name__) + + +# AST node types explicitly forbidden *before* RestrictedPython compilation. +# RestrictedPython lets `import` and `global`/`nonlocal` compile but they fail +# at runtime; we reject them early for clear error messages. _FORBIDDEN_AST_NODES: set[type[ast.AST]] = { ast.Import, ast.ImportFrom, @@ -98,41 +110,76 @@ _FORBIDDEN_AST_NODES: set[type[ast.AST]] = { ast.Nonlocal, } -logger_sr = logging.getLogger(__name__) +# Named calls that must be blocked even though RestrictedPython may allow some +# of them (e.g. ``getattr`` and ``setattr`` are in ``safe_builtins``). +_DANGEROUS_CALL_NAMES: set[str] = { + "exec", + "eval", + "compile", + "__import__", + "getattr", + "setattr", +} -def _validate_code_ast(code: str) -> None: - """Parse *code* and reject any forbidden AST constructs. +def _compile_restricted_code(code: str) -> Any: + """Compile *code* via RestrictedPython for safe ``exec()``. - Raises ``StreamRoutingError`` if the code contains imports, global/nonlocal - statements, or calls to dangerous built-ins (exec, eval, compile, __import__). + A lightweight AST pre-check rejects imports, global/nonlocal statements, + and calls to dangerous built-in names. RestrictedPython then performs a + full AST transformation that blocks dunder attribute access, subscript- + based jail-breaks, and more. + + Returns the compiled code object. Raises ``StreamRoutingError`` on + syntax errors or any policy violation. """ + # --- Phase 1: quick AST pre-check for constructs RP doesn't block --- try: tree = ast.parse(code) except SyntaxError as exc: raise StreamRoutingError(f"Invalid code syntax: {exc}") from exc - _DANGEROUS_NAMES = {"exec", "eval", "compile", "__import__", "getattr", "setattr"} - for node in ast.walk(tree): if type(node) in _FORBIDDEN_AST_NODES: raise StreamRoutingError( f"Forbidden construct in tool code: {type(node).__name__}" ) - # Reject calls to dangerous built-in names if ( isinstance(node, ast.Call) and isinstance(node.func, ast.Name) - and node.func.id in _DANGEROUS_NAMES + and node.func.id in _DANGEROUS_CALL_NAMES ): raise StreamRoutingError(f"Forbidden call in tool code: {node.func.id}()") + # --- Phase 2: RestrictedPython compilation (blocks dunder access etc.) --- + result = compile_restricted_exec(code) + if result.errors: + raise StreamRoutingError( + f"Forbidden construct in tool code: {'; '.join(result.errors)}" + ) + if result.code is None: + raise StreamRoutingError("Invalid code syntax: compilation produced no code") + return result.code -def _validate_lambda_ast(fn_str: str) -> None: - """Validate that *fn_str* is a single lambda expression and nothing else. - Raises ``StreamRoutingError`` if the string is not a safe lambda. +def _validate_code_ast(code: str) -> None: + """Parse *code* and reject any forbidden AST constructs. + + Uses RestrictedPython's ``compile_restricted_exec`` to reject dangerous + patterns including dunder attribute access, exec/eval calls, and more. + Raises ``StreamRoutingError`` on any violation. """ + _compile_restricted_code(code) + + +def _compile_restricted_lambda(fn_str: str) -> Any: + """Compile *fn_str* via RestrictedPython for safe ``eval()``. + + Validates that the string is a single lambda expression and returns + the compiled code object. Raises ``StreamRoutingError`` on syntax + errors, non-lambda expressions, or RestrictedPython rejections. + """ + # Pre-check: must be a lambda expression try: tree = ast.parse(fn_str, mode="eval") except SyntaxError as exc: @@ -144,12 +191,51 @@ def _validate_lambda_ast(fn_str: str) -> None: f"got {type(tree.body).__name__}" ) + # Walk the lambda body for dangerous named calls (compile, getattr, etc.) + # that RestrictedPython does not block. Statement-only constructs like + # import/global/nonlocal are syntactically impossible inside a lambda. + for node in ast.walk(tree.body): + if ( + isinstance(node, ast.Call) + and isinstance(node.func, ast.Name) + and node.func.id in _DANGEROUS_CALL_NAMES + ): + raise StreamRoutingError( + f"Forbidden call in transform function: {node.func.id}()" + ) + + result = compile_restricted_eval(fn_str) + if result.errors: + raise StreamRoutingError( + f"Forbidden construct in transform function: {'; '.join(result.errors)}" + ) + if result.code is None: + raise StreamRoutingError( + "Invalid transform function syntax: compilation produced no code" + ) + return result.code + + +def _validate_lambda_ast(fn_str: str) -> None: + """Validate that *fn_str* is a single lambda expression and nothing else. + + Uses RestrictedPython's ``compile_restricted_eval`` to reject dangerous + patterns. Raises ``StreamRoutingError`` if the string is not a safe lambda. + """ + _compile_restricted_lambda(fn_str) + class SimpleToolAgent: """Executes inline tool code blocks from agent config.""" - def __init__(self, tools: list[dict[str, Any]] | None = None): + def __init__( + self, + tools: list[dict[str, Any]] | None = None, + *, + unsafe: bool = False, + ): self.tools = tools or [] + self._unsafe = unsafe def process( self, @@ -164,8 +250,6 @@ class SimpleToolAgent: code = tool.get("code") if isinstance(tool, dict) else None if not code: return content - # Validate AST before execution — reject imports, dangerous calls, etc. - _validate_code_ast(code) local_vars: dict[str, Any] = { "input_data": content, "metadata": meta, @@ -173,7 +257,20 @@ class SimpleToolAgent: "result": None, } try: - exec(code, {}, local_vars) # nosec B102 + if self._unsafe: + # Unsafe mode: allow imports and full builtins for configs + # that have explicitly opted in via the --unsafe flag. + exec( # nosec B102 + compile(code, "", "exec"), + {"__builtins__": __builtins__}, + local_vars, + ) + else: + # Safe mode: compile via RestrictedPython — rejects dunder + # access, exec/eval, imports, etc. + restricted_code = _compile_restricted_code(code) + restricted_globals = _rp_safe_globals.copy() + exec(restricted_code, restricted_globals, local_vars) # nosec B102 except StreamRoutingError: raise # Re-raise validation errors except Exception: @@ -404,10 +501,9 @@ class ReactiveStreamRouter: # pylint: disable=too-many-instance-attributes if "fn" in params: fn_str = params["fn"] try: - # Validate: only lambda expressions allowed - _validate_lambda_ast(fn_str) - compiled = compile(fn_str, "", "eval") - transform_fn = eval(compiled, {"__builtins__": {}}) # nosec B307 + # Compile via RestrictedPython — validates and compiles in one step + restricted_code = _compile_restricted_lambda(fn_str) + transform_fn = eval(restricted_code, _rp_safe_globals.copy()) # nosec B307 return ops.map( lambda msg: ( transform_fn(msg.content)