From 8f12ea9a3ae195ca4329d889c4bc3bad65d30add Mon Sep 17 00:00:00 2001 From: A Wey <8992484+CurrentlyAWey@users.noreply.github.com> Date: Fri, 3 Jul 2026 11:44:21 -0500 Subject: [PATCH] fix: harden state persistence, redirect handling, and consent coverage Defense-in-depth hardening across four tools. Full details shared privately with the maintainers per the project's security policy. - python_repl: persist REPL state as JSON instead of pickle (dill), so loading a persisted state file cannot execute code. The state file is now repl_state.json. - http_request: drop the X-API-Key header on a cross-host redirect (mirroring how requests already strips Authorization), so an api_key credential is not forwarded off the originally-requested host. - use_aws: decide whether an operation is mutative via a read-only allowlist (fail closed) rather than a fixed verb denylist, so impactful operations whose names are not in that list are still covered by the confirmation prompt. - editor: include undo_edit among the commands that require confirmation, since it overwrites a file from its .bak and deletes the backup. Tests updated for the JSON persistence format. --- src/strands_tools/editor.py | 20 ++++++++- src/strands_tools/http_request.py | 28 ++++++++++++- src/strands_tools/python_repl.py | 28 ++++++++----- src/strands_tools/use_aws.py | 34 +++++++++++++++- tests/test_python_repl.py | 67 ++++++++++++------------------- tests/test_use_aws.py | 5 ++- 6 files changed, 124 insertions(+), 58 deletions(-) diff --git a/src/strands_tools/editor.py b/src/strands_tools/editor.py index a70f86da..1f1d3de0 100644 --- a/src/strands_tools/editor.py +++ b/src/strands_tools/editor.py @@ -332,8 +332,11 @@ def editor( # Check if we're in development mode strands_dev = os.environ.get("BYPASS_TOOL_CONSENT", "").lower() == "true" - # For modifying operations, show confirmation dialog unless in BYPASS_TOOL_CONSENT mode - modifying_commands = {"create", "str_replace", "pattern_replace", "insert"} + # For modifying operations, show confirmation dialog unless in BYPASS_TOOL_CONSENT mode. + # NOTE: "undo_edit" MUST be included -- it overwrites `path` from `.bak` and + # deletes that backup, so it is a state-changing operation and requires the same + # consent as the other mutations (omitting it let an un-consented overwrite/delete through). + modifying_commands = {"create", "str_replace", "pattern_replace", "insert", "undo_edit"} needs_confirmation = command in modifying_commands and not strands_dev if needs_confirmation: @@ -430,6 +433,19 @@ def editor( ) console.print(table) + elif command == "undo_edit": + console.print( + Panel( + Text( + f"Restore {path} from {path}.bak and delete the backup.", + style="yellow", + ), + title="[bold yellow]Undo Preview", + border_style="yellow", + box=box.DOUBLE, + ) + ) + # Get user confirmation user_input = get_user_input( f"Do you want to proceed with the {command} operation? [y/*]" diff --git a/src/strands_tools/http_request.py b/src/strands_tools/http_request.py index b90b0221..4ade5d70 100644 --- a/src/strands_tools/http_request.py +++ b/src/strands_tools/http_request.py @@ -234,9 +234,35 @@ def extract_content_from_html(html: str) -> str: return html +class _StrandsSession(requests.Session): + """A Session that drops tool-injected secret headers on a cross-host redirect. + + ``requests`` strips only ``Authorization`` when a redirect changes host + (``Session.rebuild_auth``). Secrets injected into other headers -- notably the + ``X-API-Key`` set for ``auth_type="api_key"`` -- would otherwise be forwarded to + the redirect target, leaking the credential to an arbitrary host (e.g. via an open + redirect on an allowlisted domain) and defeating the ``HTTP_REQUEST_TOKEN_CONFIG`` + first-hop domain allowlist. We extend the same host-change handling to those headers. + """ + + #: Secret-bearing headers this tool may inject that are NOT ``Authorization``. + _SENSITIVE_REDIRECT_HEADERS = ("X-API-Key",) + + def rebuild_auth(self, prepared_request: requests.PreparedRequest, response: requests.Response) -> None: + super().rebuild_auth(prepared_request, response) # strips Authorization on host change + try: + original_host = urlparse(response.request.url).hostname + redirect_host = urlparse(prepared_request.url).hostname + except Exception: + return + if original_host != redirect_host: + for header in self._SENSITIVE_REDIRECT_HEADERS: + prepared_request.headers.pop(header, None) + + def create_session(config: Dict[str, Any]) -> requests.Session: """Create and configure a requests Session object.""" - session = requests.Session() + session = _StrandsSession() if config.get("keep_alive", True): adapter = HTTPAdapter( diff --git a/src/strands_tools/python_repl.py b/src/strands_tools/python_repl.py index da91ca4e..4c206611 100644 --- a/src/strands_tools/python_repl.py +++ b/src/strands_tools/python_repl.py @@ -33,6 +33,7 @@ """ import fcntl +import json import logging import os import pty @@ -50,7 +51,6 @@ from pathlib import Path from typing import Any, Callable, Dict, List, Optional, Type -import dill from rich import box from rich.panel import Panel from rich.syntax import Syntax @@ -182,15 +182,21 @@ def __init__(self) -> None: else: self.persistence_dir = os.path.join(Path.cwd(), "repl_state") os.makedirs(self.persistence_dir, exist_ok=True) - self.state_file = os.path.join(self.persistence_dir, "repl_state.pkl") + # State is persisted as JSON (see load_state/save_state). Deserializing with + # dill/pickle executes arbitrary code embedded in the file (CWE-502), which is a + # remote-code-execution sink when the state file lives at a predictable, + # potentially attacker-writable location (e.g. a current-working-directory-relative + # path in an untrusted workspace). JSON cannot execute code on load, closing that + # sink. The ".json" name also ensures any pre-existing ".pkl" file is never loaded. + self.state_file = os.path.join(self.persistence_dir, "repl_state.json") self.load_state() def load_state(self) -> None: """Load persisted state with reset on failure.""" if os.path.exists(self.state_file): try: - with open(self.state_file, "rb") as f: - saved_state = dill.load(f) + with open(self.state_file, "r", encoding="utf-8") as f: + saved_state = json.load(f) self._namespace.update(saved_state) logger.debug("Successfully loaded REPL state") except Exception as e: @@ -212,20 +218,22 @@ def save_state(self, code: Optional[str] = None) -> None: if code: exec(code, self._namespace) - # Filter namespace for persistence + # Filter namespace for persistence. Only JSON-serialisable values are kept, so + # the load path (json.load) can never execute code. Non-serialisable objects + # (functions, custom classes, ...) are skipped, as they were before for values + # dill could not pickle. save_dict = {} for name, value in self._namespace.items(): if not name.startswith("_"): try: - # Try to pickle the value - dill.dumps(value) + json.dumps(value) save_dict[name] = value - except BaseException: + except (TypeError, ValueError): continue # Save state - with open(self.state_file, "wb") as f: - dill.dump(save_dict, f) + with open(self.state_file, "w", encoding="utf-8") as f: + json.dump(save_dict, f) logger.debug("Successfully saved REPL state") except Exception as e: diff --git a/src/strands_tools/use_aws.py b/src/strands_tools/use_aws.py index 7894c40b..46fb0a79 100644 --- a/src/strands_tools/use_aws.py +++ b/src/strands_tools/use_aws.py @@ -113,6 +113,32 @@ "accept", ] +# Read-only operation-name prefixes. Detecting *mutation* by matching a fixed verb +# denylist fails OPEN: any impactful operation whose name contains none of the verbs +# (e.g. ssm.send_command, lambda.invoke, ecs.execute_command, ec2.run_instances, +# ec2.authorize_security_group_ingress, rds-data.execute_statement, sqs.purge_queue) +# silently skips the confirmation prompt. We instead treat an operation as mutative +# UNLESS its name starts with a known read-only verb -- fail CLOSED. AWS operation +# names are consistently verb-prefixed, so this allowlist is stable across services. +READONLY_OPERATION_PREFIXES = ( + "get_", + "list_", + "describe_", + "head_", + "batch_get_", + "select", + "scan", + "query", + "lookup", + "search", + "retrieve", + "estimate_", + "preview_", + "validate_", + "simulate_", + "generate_presigned", +) + # Operations that return credentials or secrets and require user consent # even though they are not mutative. These operations disclose sensitive # material that should not be exposed to the LLM context without explicit @@ -159,6 +185,7 @@ "masteruserpassword", } + def redact_sensitive_values(obj: Any) -> Any: """Recursively redact values of known sensitive keys from an AWS response. @@ -393,8 +420,11 @@ def use_aws(tool: ToolUse, **kwargs: Any) -> ToolResult: "Invoking: service_name = %s, operation_name = %s, parameters = %s" % (service_name, operation_name, parameters) ) - # Check if the operation is potentially mutative - is_mutative = any(op in operation_name.lower() for op in MUTATIVE_OPERATIONS) + # Check if the operation is potentially mutative. Fail CLOSED: an operation is + # treated as mutative unless its name starts with a known read-only verb. (A verb + # denylist misses impactful ops like ssm.send_command / lambda.invoke / + # ecs.execute_command / ec2.run_instances and skips their confirmation prompt.) + is_mutative = not operation_name.lower().startswith(READONLY_OPERATION_PREFIXES) # Check if the operation returns credentials or secrets is_sensitive = (service_name.lower(), operation_name.lower()) in SENSITIVE_OPERATIONS diff --git a/tests/test_python_repl.py b/tests/test_python_repl.py index 63b1b392..71ad6fbe 100644 --- a/tests/test_python_repl.py +++ b/tests/test_python_repl.py @@ -1,5 +1,6 @@ """Tests for the python_repl tool using the Agent interface.""" +import json import os import sys import tempfile @@ -8,7 +9,6 @@ from pathlib import Path from unittest.mock import patch -import dill import pytest from strands import Agent @@ -37,11 +37,11 @@ def temp_repl_state_dir(): with tempfile.TemporaryDirectory() as tmpdir: original_dir = python_repl.repl_state.persistence_dir python_repl.repl_state.persistence_dir = tmpdir - python_repl.repl_state.state_file = os.path.join(tmpdir, "repl_state.pkl") + python_repl.repl_state.state_file = os.path.join(tmpdir, "repl_state.json") yield tmpdir # Restore original directory python_repl.repl_state.persistence_dir = original_dir - python_repl.repl_state.state_file = os.path.join(original_dir, "repl_state.pkl") + python_repl.repl_state.state_file = os.path.join(original_dir, "repl_state.json") class TestOutputCapture: @@ -167,15 +167,15 @@ def test_save_state_and_load(self, temp_repl_state_dir): """Test saving and loading state from a file.""" # Create a new state file with our own content test_state = {"test_var": "test value"} - state_file_path = os.path.join(temp_repl_state_dir, "repl_state.pkl") - with open(state_file_path, "wb") as f: - dill.dump(test_state, f) + state_file_path = os.path.join(temp_repl_state_dir, "repl_state.json") + with open(state_file_path, "w", encoding="utf-8") as f: + json.dump(test_state, f) # Force loading of our state file repl = python_repl.ReplState() # Explicitly load the state to ensure it picks up our file - with open(state_file_path, "rb") as f: - saved_state = dill.load(f) + with open(state_file_path, "r", encoding="utf-8") as f: + saved_state = json.load(f) repl._namespace.update(saved_state) # Verify our variable is in the namespace @@ -196,50 +196,33 @@ def test_save_state_with_code(self, temp_repl_state_dir): assert "test_var" in repl.get_namespace() assert repl.get_namespace()["test_var"] == "directly saved" - def test_save_state_with_unpicklable_objects(self, temp_repl_state_dir): - """Test saving state with objects that can't be pickled.""" - # Create a clean state + def test_save_state_with_unserializable_objects(self, temp_repl_state_dir): + """Non-JSON-serialisable values are skipped; serialisable ones persist. + + State is persisted as JSON (json.load cannot execute code, unlike dill/pickle), + so save_state filters the namespace to JSON-serialisable values only. + """ repl = python_repl.ReplState() repl.clear_state() - # Patch the dill.dumps function to simulate pickling failures for specific objects - with patch("dill.dumps") as mock_dumps: - # Configure mock to raise for unpicklable but work for regular values - def side_effect(obj): - if isinstance(obj, dict) and "unpicklable" in obj: - # The whole dict can be pickled, but we'll test the filtering logic - return bytes("mocked pickle", "utf-8") - elif obj == "unpicklable_value": - raise TypeError("Cannot pickle 'unpicklable_value'") - return bytes("mocked pickle", "utf-8") - - mock_dumps.side_effect = side_effect - - # Add objects to namespace - repl._namespace["regular"] = "this should be saved" - repl._namespace["unpicklable"] = "unpicklable_value" - - # Save state - repl.save_state() + # A serialisable value and an unserialisable one (functions aren't JSON-serialisable). + repl._namespace["regular"] = "this should be saved" + repl._namespace["unserializable"] = lambda x: x - # Verify that save_dict would only contain pickable objects - mock_calls = mock_dumps.call_args_list - found_unpicklable_rejection = False + repl.save_state() - # dill.dumps will be called for both individual values and the final save_dict - for call in mock_calls: - args = call[0] - if args[0] == "unpicklable_value": - # This should attempt to pickle but raise an exception - found_unpicklable_rejection = True + # Reload the persisted file directly and confirm the filtering. + with open(repl.state_file, "r", encoding="utf-8") as f: + persisted = json.load(f) - assert found_unpicklable_rejection, "The code should attempt to pickle 'unpicklable_value' but fail" + assert persisted.get("regular") == "this should be saved" + assert "unserializable" not in persisted def test_error_during_state_removal(self, temp_repl_state_dir): """Test handling error when removing corrupted state file.""" # Create a corrupted state file - with open(os.path.join(temp_repl_state_dir, "repl_state.pkl"), "wb") as f: - f.write(b"This is not valid pickle data") + with open(os.path.join(temp_repl_state_dir, "repl_state.json"), "w", encoding="utf-8") as f: + f.write("This is not valid JSON data") # Mock os.remove to raise an exception with patch("os.remove", side_effect=PermissionError("Permission denied")): diff --git a/tests/test_use_aws.py b/tests/test_use_aws.py index fdcf5c78..73f4d7c4 100644 --- a/tests/test_use_aws.py +++ b/tests/test_use_aws.py @@ -150,7 +150,10 @@ def test_use_aws_invalid_operation(mock_available_services, mock_available_opera }, } - result = use_aws.use_aws(tool=tool_use) + # Bypass the consent gate so the call reaches operation-name validation. (An unknown + # op name is now treated as mutative -- fail-closed -- and would otherwise prompt.) + with patch.dict("os.environ", {"BYPASS_TOOL_CONSENT": "true"}): + result = use_aws.use_aws(tool=tool_use) assert result["status"] == "error" assert "Invalid AWS operation: invalid_operation" in result["content"][0]["text"]