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)
pull/763/head
Sebastion 2026-04-23 19:36:39 +01:00
parent 693fb1ccf1
commit d619d86c4c
No known key found for this signature in database
3 changed files with 72 additions and 3 deletions

View File

@ -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, "<expr>", "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())},

View File

@ -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, "<expr>", "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')}

View File

@ -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, "<expr>", "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")},