diff --git a/src/strands_tools/calculator.py b/src/strands_tools/calculator.py index 1b5de4fb..bc45cca7 100644 --- a/src/strands_tools/calculator.py +++ b/src/strands_tools/calculator.py @@ -8,9 +8,17 @@ Security: All expressions are validated against an AST allowlist before evaluation. - Only mathematical syntax (operators, function calls, literals, names, containers) - is permitted. Attribute access, subscripts, lambdas, comprehensions, and imports - are rejected, preventing code execution via eval. + Only mathematical syntax (operators, function calls, numeric/boolean literals, + names, containers) is permitted. Attribute access, subscripts, lambdas, + comprehensions, and imports are rejected, preventing code execution via eval. + + String-literal arguments are rejected as a hardening measure. Several safe + SymPy helpers (for example N, simplify, solve) re-parse string arguments through + sympify, which evaluates them with full builtins and escapes the restricted + namespace used for the top-level expression. String literals are therefore only + permitted as positional arguments to the constructors that parse them as a plain + name or numeric literal rather than through sympify (Symbol, symbols, Rational, + Integer, Float); for example Symbol('x') and Rational('1/3') remain supported. Key Features: 1. Expression Evaluation: @@ -269,6 +277,13 @@ def create_error_panel(console: Console, error_message: str) -> None: } +# Constructors whose string arguments are parsed as plain symbol names or numeric +# literals, not re-evaluated through sympify. A string literal is permitted only as +# a positional argument to one of these; anywhere else it is rejected, which blocks +# the sympify-backed re-parse escape (e.g. N("..."), simplify("..."), solve("...")). +_STRING_ARG_CONSTRUCTORS = frozenset({"Symbol", "symbols", "Rational", "Integer", "Float"}) + + def _validate_expression_ast(expr_str: str) -> None: """Reject expressions containing attribute access, subscripts, or other unsafe syntax. @@ -289,9 +304,29 @@ def _validate_expression_ast(expr_str: str) -> None: except SyntaxError as e: raise ValueError(f"Invalid mathematical expression: {e.msg}") from e + # Collect the string-literal nodes that appear as positional arguments to a + # safe string constructor (e.g. Symbol('x')). Those are the only string + # literals we allow; every other string literal is rejected below. + allowed_string_nodes: set[int] = set() + for node in ast.walk(tree): + if isinstance(node, ast.Call) and isinstance(node.func, ast.Name) and node.func.id in _STRING_ARG_CONSTRUCTORS: + for arg in node.args: + if isinstance(arg, ast.Constant) and isinstance(arg.value, str): + allowed_string_nodes.add(id(arg)) + for node in ast.walk(tree): if type(node) not in _ALLOWED_AST_NODES: raise ValueError(f"Invalid mathematical expression: unsupported syntax '{type(node).__name__}'") + # Reject string (and bytes) literals except as positional arguments to + # the safe string constructors above. Several functions in the safe + # locals (e.g. N, simplify, solve) re-parse string arguments through + # SymPy's sympify, which evaluates them with full builtins and escapes + # the restricted namespace used for the top-level expression. Symbol('x') + # and friends only parse their string as a name or numeric literal, so + # they stay allowed while the sympify-backed re-parse remains blocked. + if isinstance(node, ast.Constant) and isinstance(node.value, (str, bytes)): + if id(node) not in allowed_string_nodes: + raise ValueError("Invalid mathematical expression: string literals are not supported") def parse_expression(expr_str: str) -> Any: @@ -768,6 +803,10 @@ def calculator( "sin(pi/2)", "factorial(5)", "Abs(x)", "Eq(x**2, 4)"). For matrix operations, use Matrix() with functions like det(), transpose(), or trace(). + String-literal arguments are rejected as a security hardening measure, + except as positional arguments to the symbol/number constructors that do + not re-parse them through sympify: Symbol('x'), symbols('x y z'), + Rational('1/3'), Integer('5') and Float('3.14') are supported. mode: The calculation mode to use. Options are: - "evaluate": Compute the value of the expression (default) - "solve": Solve an equation or system of equations @@ -803,6 +842,10 @@ def calculator( Notes: - For equation solving, set the expression equal to zero implicitly (x**2 + 1 means x**2 + 1 = 0) + - To solve a system of equations, pass the equations as a comma-separated + expression (e.g. "x + y - 10, x - y - 2"), which parses to a tuple of + expressions. Passing a quoted string list (e.g. "['x + y - 10', ...]") is + not supported because string literals are rejected during validation. - Use 'pi' and 'e' for mathematical constants - The 'wrt' parameter is required for differentiation and integration - Matrix expressions use Python-like syntax: [[1, 2], [3, 4]] diff --git a/tests/test_calculator.py b/tests/test_calculator.py index e848c085..bc101ca3 100644 --- a/tests/test_calculator.py +++ b/tests/test_calculator.py @@ -586,26 +586,35 @@ def test_error_handling_in_calculation_functions(): def test_calculator_tool_with_system_of_equations(): - """Test the calculator tool with a system of equations.""" - # Create a tool use with a system of equations + """Test the calculator tool with a system of equations using the supported syntax. + + The system is passed as a comma-separated expression, which parses to a tuple of + SymPy expressions and is routed to the system solver. This exercises the real + validator and parser (no mocks). + """ from src.strands_tools.calculator import calculator as calc_function - # Mock parse_expression to return a list of expressions - with mock.patch( - "src.strands_tools.calculator.parse_expression", - side_effect=lambda expr: ( - [sp.Symbol("x") + sp.Symbol("y") - 10, sp.Symbol("x") - sp.Symbol("y") - 2] - if expr.startswith("[") - else expr - ), - ): - # This should trigger the system of equations path in calculator function - result = calc_function(expression="['x + y - 10', 'x - y - 2']", mode="solve") + result = calc_function(expression="x + y - 10, x - y - 2", mode="solve") - # Check for success status - assert result["status"] == "success" - # The result should contain the solution - assert "Result:" in result["content"][0]["text"] + assert result["status"] == "success" + text = result["content"][0]["text"] + # Solution is x = 6, y = 4. + assert "6" in text and "4" in text + + +def test_calculator_tool_rejects_string_list_system_syntax(): + """The quoted string-list system syntax is now rejected during validation. + + This syntax was never functional (parse_expr returns raw strings that the solver + cannot process) and string literals are now rejected outright as a hardening + measure, so the tool returns a validation error rather than a spurious success. + """ + from src.strands_tools.calculator import calculator as calc_function + + result = calc_function(expression="['x + y - 10', 'x - y - 2']", mode="solve") + + assert result["status"] == "error" + assert "string literals are not supported" in result["content"][0]["text"] def test_error_handling(agent): @@ -676,6 +685,84 @@ def test_ast_validation_rejects_unsafe_input(payload): parse_expression(payload) +@pytest.mark.parametrize( + "payload", + [ + # String-literal argument to a function that re-parses it through sympify + "N(\"__import__('os').system('id')\")", + "simplify(\"__import__('os').system('id')\")", + "solve(\"__import__('os').getpid()\")", + "integrate(\"__import__('subprocess').run(['id'])\")", + # Plain string literals in various positions + "'abc'", + "Max(1, 'abc')", + "Matrix([['a', 'b'], ['c', 'd']])", + ], +) +def test_ast_validation_rejects_string_literals(payload): + """Verify that string-literal arguments are rejected before reaching sympify.""" + with pytest.raises(ValueError, match="Invalid mathematical expression"): + parse_expression(payload) + + +def test_string_argument_does_not_execute(): + """Verify a string argument to a sympify-backed function does not run code.""" + payload = "N(\"__import__('os').getpid()\")" + with mock.patch("os.getpid") as mock_getpid: + with pytest.raises(ValueError, match="Invalid mathematical expression"): + parse_expression(payload) + mock_getpid.assert_not_called() + + +@pytest.mark.parametrize( + "expression,expected", + [ + ("Symbol('x')", sp.Symbol("x")), + ("symbols('x y z')", (sp.Symbol("x"), sp.Symbol("y"), sp.Symbol("z"))), + ("Rational('1/3')", sp.Rational(1, 3)), + ("Integer('5')", sp.Integer(5)), + ("Float('3.14')", sp.Float("3.14")), + ], +) +def test_string_arg_constructors_are_allowed(expression, expected): + """Symbol/number constructors parse their string as a name/number, not via sympify. + + These are safe (no sympify re-parse) and remain supported so the hardening does + not break the documented string-constructor syntax. + """ + assert parse_expression(expression) == expected + + +def test_symbol_string_arg_does_not_execute(): + """A malicious string passed to Symbol() becomes a symbol name, never executes.""" + malicious = "__import__('os').getpid()" + payload = f"Symbol({malicious!r})" + with mock.patch("os.getpid") as mock_getpid: + result = parse_expression(payload) + mock_getpid.assert_not_called() + # Accepted as a symbol whose name is the literal string; the string is not eval'd. + assert isinstance(result, sp.Symbol) + assert result.name == malicious + + +@pytest.mark.parametrize( + "payload", + [ + # Only positional string args to the safe constructors are allowed; a + # string in keyword position is still rejected. + "Symbol(name='x')", + "Symbol('x', name='y')", + # Bytes literals are rejected everywhere, including in constructor args. + "Symbol(b'x')", + "simplify(b\"__import__('os').getpid()\")", + ], +) +def test_ast_validation_rejects_non_positional_and_bytes_literals(payload): + """Keyword-position strings and bytes literals stay rejected by the allowlist.""" + with pytest.raises(ValueError, match="Invalid mathematical expression"): + parse_expression(payload) + + @pytest.mark.parametrize( "expression", [