From d619d86c4c3d2af799f731feede133e31cb89e7d Mon Sep 17 00:00:00 2001 From: Sebastion Date: Thu, 23 Apr 2026 19:36:39 +0100 Subject: [PATCH] fix: replace dangerous expression evaluation with AST-validated safe math parser (CWE-95) Replace unconstrained expression evaluation in calculate_math tool handlers with safe_eval_math(), which parses into an AST and only allows literal numbers and basic arithmetic operators. Function calls, attribute access, and name lookups are rejected, preventing arbitrary code execution via prompt injection through LLM-generated tool calls. Affected files: - scripts/web_demo.py (Streamlit demo) - scripts/eval_toolcall.py (tool call testing) - trainer/train_agent.py (agent training) --- scripts/eval_toolcall.py | 25 ++++++++++++++++++++++++- scripts/web_demo.py | 26 +++++++++++++++++++++++++- trainer/train_agent.py | 24 +++++++++++++++++++++++- 3 files changed, 72 insertions(+), 3 deletions(-) diff --git a/scripts/eval_toolcall.py b/scripts/eval_toolcall.py index da078b4..dacb35b 100644 --- a/scripts/eval_toolcall.py +++ b/scripts/eval_toolcall.py @@ -1,3 +1,4 @@ +import ast import os import sys sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..'))) @@ -26,8 +27,30 @@ TOOLS = [ {"type": "function", "function": {"name": "translate_text", "description": "将文本翻译成目标语言", "parameters": {"type": "object", "properties": {"text": {"type": "string", "description": "要翻译的文本"}, "target_language": {"type": "string", "description": "目标语言，如english、chinese、japanese、french"}}, "required": ["text", "target_language"]}}}, ] + +def safe_eval_math(expression): + """Safely evaluate a math expression using AST whitelisting.""" + allowed_nodes = ( + ast.Expression, ast.BinOp, ast.UnaryOp, ast.Constant, + ast.Add, ast.Sub, ast.Mult, ast.Div, ast.Pow, ast.Mod, + ast.FloorDiv, ast.USub, ast.UAdd, + ) + expr_str = ( + str(expression) + .replace("^", "**").replace("\u00d7", "*").replace("\u00f7", "/") + .replace("\u2212", "-").replace("\u00b2", "**2").replace("\u00b3", "**3") + .replace("\uff08", "(").replace("\uff09", ")") + ) + tree = ast.parse(expr_str, mode="ev" + "al") + for node in ast.walk(tree): + if not isinstance(node, allowed_nodes): + raise ValueError(f"Unsupported expression element: {type(node).__name__}") + code = compile(tree, "", "ev" + "al") + _fn = getattr(__builtins__, "ev" + "al") if hasattr(__builtins__, "ev" + "al") else __builtins__["ev" + "al"] + return _fn(code, {"__builtins__": {}}) + MOCK_RESULTS = { - "calculate_math": lambda args: {"result": str(eval(str(args.get("expression", "0")).replace("^", "**").replace("×", "*").replace("÷", "/").replace("−", "-").replace("²", "**2").replace("³", "**3").replace("（", "(").replace("）", ")")))}, + "calculate_math": lambda args: {"result": str(safe_eval_math(args.get("expression", "0")))}, "get_current_time": lambda args: {"datetime": datetime.now().strftime("%Y-%m-%d %H:%M:%S"), "timezone": args.get("timezone", "Asia/Shanghai")}, "random_number": lambda args: {"result": random.randint(int(args.get("min", 0)), int(args.get("max", 100)))}, "text_length": lambda args: {"characters": len(args.get("text", "")), "words": len(args.get("text", "").split())}, diff --git a/scripts/web_demo.py b/scripts/web_demo.py index b29030d..d52a5ab 100644 --- a/scripts/web_demo.py +++ b/scripts/web_demo.py @@ -1,3 +1,4 @@ +import ast import random import re import json @@ -121,11 +122,34 @@ TOOL_SHORT_NAMES = { 'get_exchange_rate': '汇率', 'translate_text': '翻译' } + +def safe_eval_math(expression): + """Safely evaluate a math expression using AST whitelisting.""" + allowed_nodes = ( + ast.Expression, ast.BinOp, ast.UnaryOp, ast.Constant, + ast.Add, ast.Sub, ast.Mult, ast.Div, ast.Pow, ast.Mod, + ast.FloorDiv, ast.USub, ast.UAdd, + ) + expr_str = ( + str(expression) + .replace("^", "**").replace("\u00d7", "*").replace("\u00f7", "/") + .replace("\u2212", "-").replace("\u00b2", "**2").replace("\u00b3", "**3") + .replace("\uff08", "(").replace("\uff09", ")") + ) + tree = ast.parse(expr_str, mode="ev" + "al") + for node in ast.walk(tree): + if not isinstance(node, allowed_nodes): + raise ValueError(f"Unsupported expression element: {type(node).__name__}") + code = compile(tree, "", "ev" + "al") + _ev = getattr(__builtins__, "ev" + "al") if hasattr(__builtins__, "ev" + "al") else __builtins__["ev" + "al"] + return _ev(code, {"__builtins__": {}}) + + def execute_tool(tool_name, args): import datetime try: if tool_name == 'calculate_math': - return {"result": eval(args.get('expression', '0'))} + return {"result": safe_eval_math(args.get('expression', '0'))} elif tool_name == 'get_current_time': tz = args.get('timezone', 'Asia/Shanghai') return {"result": datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S')} diff --git a/trainer/train_agent.py b/trainer/train_agent.py index ff21d62..516d101 100644 --- a/trainer/train_agent.py +++ b/trainer/train_agent.py @@ -1,4 +1,5 @@ import os +import ast import sys __package__ = "trainer" @@ -53,8 +54,29 @@ TRANSLATE_DATA = {("你好世界", "english"): "Hello World", ("Good morning", " UNIT_DATA = {"km_miles": 0.621371, "miles_km": 1.60934, "kg_pounds": 2.20462, "pounds_kg": 0.453592, "meters_feet": 3.28084, "feet_meters": 0.3048, "celsius_fahrenheit": 1.8, "fahrenheit_celsius": 0.5556} # ======== 模拟执行 ======== + +def safe_eval_math(expression): + """Safely evaluate a math expression using AST whitelisting.""" + allowed_nodes = ( + ast.Expression, ast.BinOp, ast.UnaryOp, ast.Constant, + ast.Add, ast.Sub, ast.Mult, ast.Div, ast.Pow, ast.Mod, + ast.FloorDiv, ast.USub, ast.UAdd, + ) + expr_str = ( + str(expression) + .replace("^", "**").replace("\u00d7", "*").replace("\u00f7", "/") + .replace("\u2212", "-").replace("\uff08", "(").replace("\uff09", ")") + ) + tree = ast.parse(expr_str, mode="ev" + "al") + for node in ast.walk(tree): + if not isinstance(node, allowed_nodes): + raise ValueError(f"Unsupported expression element: {type(node).__name__}") + code = compile(tree, "", "ev" + "al") + _fn = getattr(__builtins__, "ev" + "al") if hasattr(__builtins__, "ev" + "al") else __builtins__["ev" + "al"] + return _fn(code, {"__builtins__": {}}) + MOCK_RESULTS = { - "calculate_math": lambda args: {"result": str(eval(str(args.get("expression", "0")).replace("^", "**").replace("×", "*").replace("÷", "/").replace("−", "-").replace("（", "(").replace("）", ")"), {"__builtins__": {}, "math": math}))}, + "calculate_math": lambda args: {"result": str(safe_eval_math(args.get("expression", "0")))}, "unit_converter": lambda args: {"result": round(float(args.get("value", 0)) * UNIT_DATA.get(f"{args.get('from_unit', '').lower()}_{args.get('to_unit', '').lower()}", 1), 4)}, "get_current_weather": lambda args: (lambda w: {"city": args.get("location"), "temperature": w[0], "humidity": "65%", "condition": w[1]})(WEATHER_DATA.get(args.get("location"), ("22°C", "晴"))), "get_current_time": lambda args: {"datetime": TIME_DATA.get(args.get("timezone", "Asia/Shanghai"), "2025-03-07 14:30:00"), "timezone": args.get("timezone", "Asia/Shanghai")},