From d47a66ea88c5cfe9a2b5fa2e54ef9446d5a08ecf Mon Sep 17 00:00:00 2001 From: Wenjie Zhang Date: Tue, 3 Feb 2026 13:31:25 +0800 Subject: [PATCH] =?UTF-8?q?fix(security):=20=E7=A7=BB=E9=99=A4=20SQL=20?= =?UTF-8?q?=E6=B3=A8=E9=87=8A=E5=90=8E=E5=86=8D=E8=BF=9B=E8=A1=8C=E5=AE=89?= =?UTF-8?q?=E5=85=A8=E9=AA=8C=E8=AF=81?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - 防止通过注释绕过后续的安全检查逻辑 - 增强 MySQLSecurityChecker 的防护能力 --- src/agents/common/toolkits/mysql/security.py | 6 ++++-- src/agents/reporter/graph.py | 15 +++++++++++++++ 2 files changed, 19 insertions(+), 2 deletions(-) diff --git a/src/agents/common/toolkits/mysql/security.py b/src/agents/common/toolkits/mysql/security.py index b53e16c9..fc1bfe6a 100644 --- a/src/agents/common/toolkits/mysql/security.py +++ b/src/agents/common/toolkits/mysql/security.py @@ -34,8 +34,10 @@ class MySQLSecurityChecker: if not sql: return False - # 标准化SQL - sql_upper = sql.strip().upper() + # 移除SQL注释（-- 和 /* */）后再验证 + sql_clean = re.sub(r'--.*$', '', sql, flags=re.MULTILINE) + sql_clean = re.sub(r'/\*.*?\*/', '', sql_clean) + sql_upper = sql_clean.strip().upper() # 检查是否是允许的操作 if not any(sql_upper.startswith(op) for op in cls.ALLOWED_OPERATIONS): diff --git a/src/agents/reporter/graph.py b/src/agents/reporter/graph.py index c8b82ad9..5b9aa4f8 100644 --- a/src/agents/reporter/graph.py +++ b/src/agents/reporter/graph.py @@ -12,9 +12,24 @@ from src.services.mcp_service import get_mcp_server_names, get_tools_from_all_se from src.utils import logger +PROMPT = """你的任务是根据用户的指令，使用数据库工具和图表绘制工具，构建 SQL 查询报告。 +你需要根据用户的指令，生成相应的 SQL 查询，并将查询结果以报表的形式返回给用户。 +在生成报表时，你可以调用工具生成图表，以更直观地展示数据。 +务必确保生成的 SQL 查询是正确且高效的，以避免对数据库造成不必要的负担。 +在生成报表时，请遵循以下原则： +1. 理解用户的指令，明确报表的需求和目标。 +2. 图表生成工具的返回结果不会默认渲染，需要在最终的报表中以图片形式(markdown格式)嵌入。 +3. 必要时，使用网络检索相关工具补充信息。 +""" + @dataclass(kw_only=True) class ReporterContext(BaseContext): """覆盖 BaseContext，定义数据库报表助手智能体的可配置参数""" + # 覆盖 system_prompt，提供更具体的默认值 + system_prompt: Annotated[str, {"__template_metadata__": {"kind": "prompt"}}] = field( + default=PROMPT, + metadata={"name": "系统提示词", "description": "用来描述智能体的角色和行为"}, + ) mcps: Annotated[list[str], {"__template_metadata__": {"kind": "mcps"}}] = field( default_factory=lambda: ["mcp-server-chart"],