"""代码生成器模板安全(注入面)回归测试。
背景:``backend/audit-backend.md`` B1 —— 代码生成模板以 ``autoescape=False`` 渲染,并把
``function_name`` / ``column_comment`` / ``column_name`` 原样写进 ``.py``,产物落在
``backend/app/plugin/**`` 且在应用启动时被 ``app/core/discover.py`` import,等价于「有代码生成
权限 = 服务器任意代码执行」。
本测试固化两条不变量:
1. 恶意输入(引号、换行、闭合序列、内嵌 import/exec)不能变成**可执行代码**;
2. 恶意输入之下,渲染产物**仍是语法合法的 Python**(且正常文本按原义保留)。
"""
import ast
from typing import Any
import pytest
from pydantic import ValidationError
from app.modules.generator.gencode.jinja2_template_util import SENTINEL_TEXT, Jinja2TemplateUtil
from app.modules.generator.gencode.schema import GenTableColumnOutSchema, GenTableOutSchema
from app.modules.generator.gencode.template_safety import (
assert_python_structure,
assert_python_syntax,
assert_text_structure,
escape_template_text,
html_attr,
is_safe_identifier,
reserved_word_reason,
safe_ident,
unescape_template_text,
vue_label,
)
# ── 恶意载荷:覆盖「引号 / 换行 / 闭合序列 / 内嵌可执行语句」四类 ──────────────
EVIL_COMMENT = "'); import os; os.system('id'); #"
EVIL_FUNCTION_NAME = 'x\n"""\nimport os\nos.system("id")\n"""\n'
EVIL_COLUMN_NAME = 'evil"); import os; os.system("id"); #'
EVIL_BACKTICK = "`${alert(1)}`"
def _column(**overrides: Any) -> GenTableColumnOutSchema:
"""构造一个字段模型(默认是普通 varchar 列)。"""
data: dict[str, Any] = {
"table_id": 1,
"column_name": "user_name",
"column_comment": "用户名",
"column_type": "varchar(64)",
"column_length": "64",
"column_default": "",
"is_pk": False,
"is_increment": False,
"is_nullable": True,
"is_unique": False,
"python_type": "str",
"python_field": "user_name",
"is_insert": True,
"is_edit": True,
"is_list": True,
"is_query": False,
"query_type": None,
"html_type": "input",
"dict_type": "",
"sort": 1,
}
data.update(overrides)
return GenTableColumnOutSchema(**data)
def _table(columns: list[GenTableColumnOutSchema] | None = None, **overrides: Any) -> GenTableOutSchema:
"""构造一个输出模型(默认包名/模块名已是工程约定形态)。"""
cols = columns if columns is not None else [_column()]
pk = next((c for c in cols if c.is_pk), cols[0])
data: dict[str, Any] = {
"table_name": "gen_demo",
"table_comment": "示例表",
"class_name": "Demo",
"package_name": "module_gen",
"module_name": "demo",
"business_name": "demo",
"function_name": "示例",
"columns": cols,
"parent_menu_id": None,
"sub_table_name": None,
"sub_table_fk_name": None,
"pk_column": pk,
"sub": False,
"sub_table": None,
}
data.update(overrides)
return GenTableOutSchema(**data)
async def _render_all(table: GenTableOutSchema) -> dict[str, str]:
"""按生产路径(prepare_context + 全部模板)渲染所有产物。"""
env = Jinja2TemplateUtil.get_env()
context = Jinja2TemplateUtil.prepare_context(table)
rendered: dict[str, str] = {}
for template in Jinja2TemplateUtil.get_template_list():
rendered[Jinja2TemplateUtil.get_file_name(template, table)] = await env.get_template(template).render_async(**context)
return rendered
def _python_outputs(rendered: dict[str, str]) -> dict[str, str]:
"""只取 .py 产物。"""
return {name: content for name, content in rendered.items() if name.endswith(".py")}
def _dangerous_ast_nodes(source: str) -> list[str]:
"""返回源码中「可执行注入」的 AST 证据(import os / os.system / exec)。"""
tree = ast.parse(source)
findings: list[str] = []
for node in ast.walk(tree):
if isinstance(node, ast.Import):
findings.extend(f"import {alias.name}" for alias in node.names if alias.name.split(".")[0] == "os")
elif isinstance(node, ast.ImportFrom) and (node.module or "").split(".")[0] == "os":
findings.append(f"from {node.module} import")
elif isinstance(node, ast.Call) and isinstance(node.func, ast.Attribute) and node.func.attr in {"system", "popen", "exec", "eval"}:
findings.append(f".{node.func.attr}() 调用")
return findings
# ── 1. 纯函数级 ────────────────────────────────────────────────────────────
def test_escape_template_text_neutralizes_quote_and_newline() -> None:
"""转义后不再含可闭合字面量/可换行的字符,且保留原字符的语义(\\uXXXX)。"""
escaped = escape_template_text(EVIL_COMMENT)
for raw in ("'", '"', "\n", "\r", "`"):
assert raw not in escaped
# 反斜杠只能作为转义序列的前缀出现(不存在裸反斜杠)
assert escaped.count("\\") == escaped.count("\\u")
assert "\\u0027" in escaped # 单引号被转义
assert "os.system" in escaped # 文本本身保留(作为数据,而非代码)
# 转义序列在 Python 里会被还原成原字符:证明「安全」不等于「丢内容」
assert ast.literal_eval(f"'{escaped}'") == EVIL_COMMENT
def test_escape_template_text_breaks_js_interpolation() -> None:
"""${...} 会被打断,避免在 JS 模板字符串里变成插值表达式。"""
escaped = escape_template_text(EVIL_BACKTICK)
assert "${" not in escaped
assert "\\u0060" in escaped
def test_safe_ident_normalizes_and_preserves_legal_names() -> None:
"""合法标识符原样保留;非法字符被替换;数字开头补下划线;空值取默认。"""
assert safe_ident("user_name") == "user_name"
assert safe_ident("_private") == "_private"
assert safe_ident("a__b") == "a__b"
assert safe_ident("2fa") == "_2fa"
assert safe_ident("") == "entity"
assert safe_ident(None, default="") == ""
# 注入型列名被规范为合法标识符(无法再闭合并列语句)
mangled = safe_ident(EVIL_COLUMN_NAME)
assert is_safe_identifier(mangled)
assert not any(ch in mangled for ch in ("'", '"', " ", ";", "(", ")"))
def test_is_safe_identifier() -> None:
"""白名单判定:合法/非法/超长/空值。"""
assert is_safe_identifier("user_name")
assert is_safe_identifier("_x1")
assert not is_safe_identifier("2fa")
assert not is_safe_identifier("a b")
assert not is_safe_identifier(EVIL_COLUMN_NAME)
assert not is_safe_identifier("")
assert not is_safe_identifier("a" * 65)
def test_vue_label_strips_markup() -> None:
"""Vue 文本节点用:尖括号/花括号被剔除并完成字符串级转义。"""
label = vue_label("{{ evil }}")
assert "<" not in label
assert ">" not in label
assert "{" not in label
assert "script" in label # 文本保留,只是失去了标记语义
def test_assert_python_syntax_raises_for_broken_output() -> None:
"""语法网:破损产物报错并带上文件名与行号。"""
assert_python_syntax("x = 1\n") # 合法不抛
with pytest.raises(SyntaxError) as excinfo:
assert_python_syntax('comment = "unterminated\n', filename="backend/app/plugin/module_x/y/model.py")
assert "model.py" in str(excinfo.value)
# ── 2. 口径层:列名白名单校验 ────────────────────────────────────────────────
def test_column_name_whitelist_accepts_legal_and_rejects_injection() -> None:
"""列名必须能直接充当标识符:注入型/数字开头/含空格一律拒绝。"""
assert GenTableColumnOutSchema(**{**_column().model_dump(), "column_name": "order_no"}).column_name == "order_no"
assert GenTableColumnOutSchema(**{**_column().model_dump(), "column_name": ""}).column_name == ""
for bad in (EVIL_COLUMN_NAME, "2fa", "user name", "a\nb", "a'b"):
with pytest.raises(ValidationError):
GenTableColumnOutSchema(**{**_column().model_dump(), "column_name": bad})
# ── 3. 渲染层:恶意输入不能产出可执行代码,且产物必须语法合法 ─────────────────
async def test_malicious_metadata_cannot_inject_executable_python() -> None:
"""列注释/功能名/列名全为注入载荷时:产物仍可 ast.parse,且不含可执行注入节点。"""
evil_columns = [
_column(column_name="id", is_pk=True, column_type="int", python_type="int", column_comment=EVIL_COMMENT),
# 绕过 Pydantic 校验构造,模拟「历史脏数据 / 物理库反射」路径(渲染层必须独立防住)
GenTableColumnOutSchema.model_construct(**{**_column().model_dump(), "column_name": EVIL_COLUMN_NAME}),
]
table = _table(columns=evil_columns, function_name=EVIL_FUNCTION_NAME, table_comment=EVIL_COMMENT)
rendered = _python_outputs(await _render_all(table))
assert rendered, "应至少渲染出一个 Python 产物"
for name, content in rendered.items():
# ① 语法合法
ast.parse(content)
# ② 无 import os / os.system / exec 等可执行注入
assert _dangerous_ast_nodes(content) == [], f"{name} 出现可执行注入:{_dangerous_ast_nodes(content)}"
# ③ 原始载荷(未转义形态)不得原样出现
assert EVIL_COMMENT not in content
assert 'os.system("id")' not in content
async def test_malicious_inputs_keep_output_syntactically_valid_for_all_python_templates() -> None:
"""全部 .py 模板逐个校验语法(含 model/schema/controller/service/crud/__init__)。"""
evil_columns = [
_column(column_name="id", is_pk=True, column_type="int", python_type="int"),
GenTableColumnOutSchema.model_construct(
**{**_column().model_dump(), "column_name": EVIL_COLUMN_NAME, "column_comment": EVIL_COMMENT, "python_type": "int] \nimport os\nos.system('id')"},
),
]
table = _table(columns=evil_columns, function_name=EVIL_FUNCTION_NAME, class_name="Demo", table_name='gen_demo"); import os; #')
rendered = _python_outputs(await _render_all(table))
for name, content in rendered.items():
assert_python_syntax(content, filename=name)
assert _dangerous_ast_nodes(content) == [], f"{name} 出现可执行注入"
async def test_escaped_comment_preserves_readable_text() -> None:
"""安全化不能丢内容:普通文本(含撇号)在产物里按原义保留。"""
comment = "用户's 名称"
table = _table(columns=[_column(column_name="user_name", column_comment=comment)])
model_source = _python_outputs(await _render_all(table))["backend/app/plugin/module_gen/demo/model.py"]
tree = ast.parse(model_source)
comments: list[Any] = []
for node in ast.walk(tree):
if isinstance(node, ast.Call):
for keyword in node.keywords:
if keyword.arg == "comment":
comments.append(ast.literal_eval(keyword.value))
assert comment in comments, f"注释文本未被原义保留:{comments}"
async def test_vue_template_escapes_label_positions() -> None:
"""Vue 文本节点:子表功能名里的标签/脚本不能改变生成页面的标记结构。"""
sub = _table(columns=[_column(column_name="sub_name")], function_name="")
parent = _table(
columns=[_column(column_name="id", is_pk=True, column_type="int", python_type="int")],
sub=True,
sub_table=sub,
sub_table_name="gen_sub",
sub_table_fk_name="demo_id",
)
rendered = await _render_all(parent)
vue_source = rendered["frontend/web/src/views/module_gen/demo/index.vue"]
assert "\n"
with pytest.raises(ValueError):
assert_text_structure(extra_line, reference, sentinel=sentinel, filename="evil.vue")
breakout = 'const a = "x"; alert(1); //";\n__S__列表\n'
with pytest.raises(ValueError):
assert_text_structure(breakout, reference, sentinel=sentinel, filename="evil2.vue")
async def test_all_templates_pass_structural_gate_with_hostile_text() -> None:
"""端到端:恶意文本下,全部模板(主+子)都能通过与落盘路径相同的结构闸门。"""
sub_cols = [
_column(column_name="id", is_pk=True, column_type="int", python_type="int", column_comment="子主键"),
_column(column_name="sub_name", column_comment="子表'备注' & "),
]
sub = _table(
columns=sub_cols,
table_name="gen_sub",
class_name="Sub",
business_name="sub",
function_name="",
)
main_cols = [_column(column_name="id", is_pk=True, column_type="int", python_type="int")]
parent = _table(
columns=main_cols,
function_name=EVIL_FUNCTION_NAME,
table_comment=EVIL_COMMENT,
sub=True,
sub_table=sub,
sub_table_name="gen_sub",
sub_table_fk_name="demo_id",
)
env = Jinja2TemplateUtil.get_env()
safe = Jinja2TemplateUtil.safe_render_schema(parent)
ctx = Jinja2TemplateUtil.prepare_context(parent)
reference_ctx = Jinja2TemplateUtil.prepare_context(Jinja2TemplateUtil.sentinel_render_schema(safe))
sub_ctx = Jinja2TemplateUtil.prepare_sub_render_context(safe, safe.sub_table)
sub_reference_ctx = Jinja2TemplateUtil.prepare_sub_render_context(
Jinja2TemplateUtil.sentinel_render_schema(safe),
Jinja2TemplateUtil.sentinel_render_schema(safe.sub_table),
)
tokens = Jinja2TemplateUtil.iter_render_dynamic_values(safe)
checked = 0
for template in Jinja2TemplateUtil.get_template_list():
name = Jinja2TemplateUtil.get_file_name(template, parent)
real = await env.get_template(template).render_async(**ctx)
ref = await env.get_template(template).render_async(**reference_ctx)
if name.endswith(".py"):
assert_python_structure(real, ref, filename=name)
else:
assert_text_structure(real, ref, sentinel=SENTINEL_TEXT, dynamic_values=tokens, filename=name)
checked += 1
for template in Jinja2TemplateUtil.get_sub_table_template_list():
name = Jinja2TemplateUtil.get_file_name(template, safe.sub_table)
real = await env.get_template(template).render_async(**sub_ctx)
ref = await env.get_template(template).render_async(**sub_reference_ctx)
assert_python_structure(real, ref, filename=name)
checked += 1
assert checked == 13 # 10 主模板 + 3 子表模板