314 lines
11 KiB
Python
314 lines
11 KiB
Python
"""task-121 N+1 扫描审计工具。
|
||
|
||
对 backend-java modules 下的 Java Service 源码做机械化扫描：
|
||
识别循环（for/while/forEach/map）内对 MyBatis-Plus 单行查询
|
||
（selectById/selectOne/selectList/selectCount/selectObjs/selectMaps/getById/getOne）
|
||
的调用，产出 N+1 候选清单（位置/对象/影响），供人工审计与 task-122 修复。
|
||
本工具只读，不改任何代码。
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import argparse
|
||
import json
|
||
import re
|
||
import sys
|
||
from dataclasses import dataclass, field
|
||
from pathlib import Path
|
||
|
||
HIGH = "HIGH"
|
||
MEDIUM = "MEDIUM"
|
||
LOW = "LOW"
|
||
IMPACT_ORDER = {HIGH: 0, MEDIUM: 1, LOW: 2}
|
||
|
||
MODULES_ROOT = (
|
||
Path(__file__).resolve().parent.parent / "src/main/java/com/nanri/aiimage/modules"
|
||
)
|
||
|
||
# 单行查询调用（selectBatchIds/selectByIds 天然批量，不匹配此正则）
|
||
QUERY_RE = re.compile(
|
||
r"\.(selectById|selectOne|selectList|selectCount|selectObjs|selectMaps|getById|getOne)\("
|
||
)
|
||
# 循环头：for/while 关键字 + 圆括号跨度，或 .forEach( / .map( 方法调用
|
||
LOOP_RE = re.compile(r"\b(for|while)\s*\(|\.(forEach|map)\s*\(")
|
||
# 返回集合的查询调用（循环遍历这类结果集 → 影响 HIGH）
|
||
COLLECTION_QUERY_RE = re.compile(
|
||
r"(selectList|selectObjs|selectMaps|selectBatchIds|selectByIds|\.list\()"
|
||
)
|
||
ASSIGN_RE = re.compile(r"(\w+)\s*=\s*([^;{}=]+);")
|
||
FIELD_RE = re.compile(r"\b[A-Za-z_]\w*\b")
|
||
|
||
|
||
def _strip_comments(text: str) -> str:
|
||
"""用空格替换注释（保留换行与偏移），简化后续扫描。"""
|
||
out = list(text)
|
||
i, n = 0, len(text)
|
||
while i < n:
|
||
if text.startswith("/*", i):
|
||
j = text.find("*/", i + 2)
|
||
j = n if j < 0 else j + 2
|
||
for k in range(i, j):
|
||
if text[k] != "\n":
|
||
out[k] = " "
|
||
i = j
|
||
elif text.startswith("//", i):
|
||
j = text.find("\n", i)
|
||
j = n if j < 0 else j
|
||
for k in range(i, j):
|
||
out[k] = " "
|
||
i = j
|
||
else:
|
||
i += 1
|
||
return "".join(out)
|
||
|
||
|
||
def _match_span(text: str, open_pos: int) -> int | None:
|
||
"""从 text[open_pos] 的 ( 出发，返回匹配的 ) 偏移（含），不匹配返回 None。"""
|
||
depth = 0
|
||
for i in range(open_pos, len(text)):
|
||
c = text[i]
|
||
if c == "(":
|
||
depth += 1
|
||
elif c == ")":
|
||
depth -= 1
|
||
if depth == 0:
|
||
return i
|
||
return None
|
||
|
||
|
||
def _line_of(text: str, offset: int) -> int:
|
||
return text.count("\n", 0, offset) + 1
|
||
|
||
|
||
class _Frame:
|
||
__slots__ = ("body_depth", "expr", "loop_type", "start")
|
||
|
||
def __init__(self, start: int, body_depth: int, loop_type: str, expr: str):
|
||
self.start = start
|
||
self.body_depth = body_depth
|
||
self.loop_type = loop_type
|
||
self.expr = expr
|
||
|
||
|
||
def _query_results(stripped: str) -> set[str]:
|
||
"""收集被赋值给集合查询结果（selectList 等）的变量名。"""
|
||
names = set()
|
||
for lhs, rhs in ASSIGN_RE.findall(stripped):
|
||
if COLLECTION_QUERY_RE.search(rhs):
|
||
names.add(lhs)
|
||
return names
|
||
|
||
|
||
def _loop_headers(stripped: str) -> list[dict]:
|
||
"""预取全部循环头：(start, end, type, expr)（start 为 ( 偏移，end 为匹配的 ) 偏移）。"""
|
||
headers: list[dict] = []
|
||
pos = 0
|
||
while True:
|
||
m = LOOP_RE.search(stripped, pos)
|
||
if not m:
|
||
break
|
||
paren_pos = m.end() - 1
|
||
if stripped[paren_pos] == "(":
|
||
close = _match_span(stripped, paren_pos)
|
||
if close is not None:
|
||
headers.append(
|
||
{
|
||
"start": paren_pos,
|
||
"end": close,
|
||
"type": m.group(1) or m.group(2),
|
||
"expr": stripped[paren_pos + 1 : close].strip(),
|
||
}
|
||
)
|
||
pos = m.end()
|
||
return headers
|
||
|
||
|
||
def _header_context(stripped: str, h: dict, offset: int) -> tuple[bool, str, str]:
|
||
"""查询落在循环头括号跨度内时的归属判定。
|
||
|
||
返回 (是否循环内, 循环类型, 循环表达式)：
|
||
- lambda 箭头之后 → 循环内（无花括号 lambda 体）
|
||
- 增强 for 的 ` : ` 之后（迭代源 selectList 等）→ 集合装载，非循环内
|
||
- for/while 初始化或条件 → 视为循环内（条件逐次求值）
|
||
"""
|
||
expr = h["expr"]
|
||
if "->" in stripped[h["start"] : h["end"]]:
|
||
arrow = stripped.find("->", h["start"], h["end"])
|
||
if arrow != -1 and offset > arrow:
|
||
return True, h["type"], expr
|
||
return False, h["type"], expr
|
||
if " : " in expr:
|
||
colon = stripped.find(" : ", h["start"], h["end"])
|
||
if colon != -1 and offset > colon:
|
||
return False, h["type"], expr
|
||
return True, h["type"], expr
|
||
return True, h["type"], expr
|
||
|
||
|
||
def scan_text(text: str, filename: str = "fixture.java") -> list[dict]:
|
||
"""扫描单份 Java 源码，返回 N+1 候选清单（按影响降序、文件行号升序）。"""
|
||
stripped = _strip_comments(text)
|
||
query_results = _query_results(stripped)
|
||
headers = _loop_headers(stripped)
|
||
|
||
depth = 0
|
||
frames: list[_Frame] = []
|
||
single: tuple[str, str, int, int] | None = None # (loop_type, expr, start, end)
|
||
findings: list[dict] = []
|
||
|
||
def in_loop(offset: int) -> bool:
|
||
if single is not None and single[2] <= offset < single[3]:
|
||
return True
|
||
if any(f.start <= offset and depth >= f.body_depth for f in frames):
|
||
return True
|
||
for h in reversed(headers):
|
||
if h["start"] <= offset <= h["end"]:
|
||
return _header_context(stripped, h, offset)[0]
|
||
return False
|
||
|
||
def enclosing(offset: int) -> tuple[str, str]:
|
||
if single is not None and single[2] <= offset < single[3]:
|
||
return single[0], single[1]
|
||
for h in reversed(headers):
|
||
if h["start"] <= offset <= h["end"]:
|
||
in_loop, loop_type, expr = _header_context(stripped, h, offset)
|
||
if in_loop:
|
||
return loop_type, expr
|
||
for f in reversed(frames):
|
||
if f.start <= offset:
|
||
return f.loop_type, f.expr
|
||
return "loop", ""
|
||
|
||
i = 0
|
||
n = len(stripped)
|
||
while i < n:
|
||
c = stripped[i]
|
||
if c == "{":
|
||
depth += 1
|
||
i += 1
|
||
continue
|
||
if c == "}":
|
||
depth -= 1
|
||
if frames and depth == frames[-1].body_depth - 1:
|
||
frames.pop()
|
||
i += 1
|
||
continue
|
||
|
||
loop_match = LOOP_RE.match(stripped, i)
|
||
if loop_match:
|
||
loop_type = loop_match.group(1) or loop_match.group(2)
|
||
paren_pos = loop_match.end() - 1
|
||
close = (
|
||
_match_span(stripped, paren_pos) if stripped[paren_pos] == "(" else None
|
||
)
|
||
if close is not None:
|
||
expr = stripped[paren_pos + 1 : close].strip()
|
||
body_start = close + 1
|
||
while body_start < n and stripped[body_start].isspace():
|
||
body_start += 1
|
||
if body_start < n and stripped[body_start] == "{":
|
||
frames.append(_Frame(body_start + 1, depth + 1, loop_type, expr))
|
||
else:
|
||
end = stripped.find(";", body_start)
|
||
single = (loop_type, expr, body_start, n if end < 0 else end)
|
||
# 继续扫描括号跨度内容（无花括号 lambda 体的查询落在其中）
|
||
i = paren_pos + 1
|
||
continue
|
||
i = loop_match.end()
|
||
continue
|
||
|
||
query_match = QUERY_RE.match(stripped, i)
|
||
if query_match:
|
||
if in_loop(i):
|
||
line_start = stripped.rfind("\n", 0, i) + 1
|
||
line_end = stripped.find("\n", i)
|
||
line_end = n if line_end < 0 else line_end
|
||
statement = stripped[line_start:line_end].strip()
|
||
loop_type, expr = enclosing(i)
|
||
impact = _classify(loop_type, [expr], query_results)
|
||
findings.append(
|
||
{
|
||
"file": filename,
|
||
"line": _line_of(text, i),
|
||
"statement": statement,
|
||
"loop_type": loop_type,
|
||
"impact": impact,
|
||
}
|
||
)
|
||
i = query_match.end()
|
||
continue
|
||
|
||
i += 1
|
||
|
||
findings.sort(key=lambda f: (IMPACT_ORDER[f["impact"]], f["line"]))
|
||
return findings
|
||
|
||
|
||
def _classify(loop_type: str, exprs: list[str], query_results: set[str]) -> str:
|
||
"""影响分级：遍历查询结果集 → HIGH；forEach/map lambda → MEDIUM；其余 → LOW。"""
|
||
for expr in exprs:
|
||
if expr and COLLECTION_QUERY_RE.search(expr):
|
||
return HIGH
|
||
if expr and any(t in query_results for t in FIELD_RE.findall(expr)):
|
||
return HIGH
|
||
if loop_type in ("forEach", "map"):
|
||
return MEDIUM
|
||
return LOW
|
||
|
||
|
||
@dataclass
|
||
class Report:
|
||
scanned_files: int = 0
|
||
modules: list[str] = field(default_factory=list)
|
||
findings: list[dict] = field(default_factory=list)
|
||
|
||
def to_dict(self) -> dict:
|
||
return {
|
||
"scanned_files": self.scanned_files,
|
||
"modules": self.modules,
|
||
"findings": self.findings,
|
||
}
|
||
|
||
|
||
def scan_modules(root: Path) -> Report:
|
||
"""扫描 modules 根目录下全部 Java 文件，产出 N+1 候选报告。"""
|
||
report = Report()
|
||
module_set: set[str] = set()
|
||
for java_file in sorted(root.rglob("*.java")):
|
||
report.scanned_files += 1
|
||
rel = java_file.relative_to(root)
|
||
module_set.add(rel.parts[0])
|
||
for finding in scan_text(java_file.read_text(encoding="utf-8"), str(rel)):
|
||
finding["module"] = rel.parts[0]
|
||
report.findings.append(finding)
|
||
report.findings.sort(
|
||
key=lambda f: (IMPACT_ORDER[f["impact"]], f["file"], f["line"])
|
||
)
|
||
report.modules = sorted(module_set)
|
||
return report
|
||
|
||
|
||
def main() -> None:
|
||
parser = argparse.ArgumentParser(description="N+1 扫描审计")
|
||
parser.add_argument("--json", help="输出 JSON 报告路径（默认 stdout）")
|
||
args = parser.parse_args()
|
||
|
||
report = scan_modules(MODULES_ROOT)
|
||
payload = report.to_dict()
|
||
if args.json:
|
||
Path(args.json).write_text(
|
||
json.dumps(payload, ensure_ascii=False, indent=2), encoding="utf-8"
|
||
)
|
||
else:
|
||
print(json.dumps(payload, ensure_ascii=False, indent=2))
|
||
print(
|
||
f"scanned {report.scanned_files} files, "
|
||
f"{len(report.findings)} N+1 candidates "
|
||
f"across {len(report.modules)} modules",
|
||
file=sys.stderr,
|
||
)
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main()
|