Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
30 changes: 30 additions & 0 deletions app/orchestrator.py
Original file line number Diff line number Diff line change
Expand Up @@ -21,6 +21,7 @@
from sqlalchemy.ext.asyncio import AsyncSession

from app import dedup
from app.tools import intake_screen
from app.agents import collector
from app.agents import intel as intel_lib
from app.agents import playbook_router
Expand Down Expand Up @@ -2262,6 +2263,12 @@ async def _persist_single_finding(self, task_id: str, target_id: str, f: dict) -
return
target_ref = tgt.url or tgt.host
worker_id = tgt.assigned_worker
# 两道闸门之一(Stage-0 预筛):垃圾/半成品不写 Finding 主表。
ok, reason = intake_screen.screen_submission(f)
if not ok:
await self._reject_intake(session, task_id, target_id, f, reason)
await session.commit()
return
duplicate = await self._find_existing_duplicate(session, target_ref, f)
if duplicate:
return
Expand Down Expand Up @@ -2298,6 +2305,24 @@ async def _persist_single_finding(self, task_id: str, target_id: str, f: dict) -
logger.warning("[realtime_persist] target=%s 实时落库失败(整轮 result 仍会兜底)",
target_id[:8], exc_info=True)

async def _reject_intake(self, session, task_id: str, target_id: str, f, reason: str) -> None:
"""被 Stage-0 预筛拦下的垃圾:只记一条 TaskEvent,不写 Finding 主表。

由调用方负责 commit(各入库事务自行收口,避免跨事务状态不一致)。
"""
try:
title = (str(f.get("title") or "")[:40] if isinstance(f, dict) else str(f)[:40])
session.add(TaskEvent(
task_id=task_id,
agent="orchestrator",
kind="intake_reject",
level="info",
message=f"入库预筛拦截({reason}):{title}",
payload={"target_id": target_id, "reason": reason},
))
except Exception:
logger.warning("intake_reject log failed target=%s", target_id[:8], exc_info=True)

async def _heartbeat_target(self, target_id: str) -> None:
timeout_ref = WORKER_IDLE_TIMEOUT if WORKER_IDLE_TIMEOUT > 0 else WORKER_WALL_TIMEOUT
interval = max(5.0, min(TARGET_HEARTBEAT_INTERVAL, max(5.0, timeout_ref / 4)))
Expand Down Expand Up @@ -2481,6 +2506,11 @@ async def _persist_worker_result(self, task_id: str, target_id: str, result: dic

# 落 Finding(含漏洞级去重;DB 唯一索引兜底,逐条 savepoint 容错并发重复)
for f in findings:
# 两道闸门之一(Stage-0 预筛):垃圾/半成品不写 Finding 主表。
ok, reason = intake_screen.screen_submission(f)
if not ok:
await self._reject_intake(session, task_id, target_id, f, reason)
continue
duplicate = await self._find_existing_duplicate(session, target_ref, f)
if duplicate:
continue
Expand Down
51 changes: 51 additions & 0 deletions app/tools/intake_screen.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,51 @@
"""两段式入库的第一道闸门:确定性预筛(Stage-0)。

把明显是垃圾 / 半成品的提交挡在 Finding 主表之外,只记一条 TaskEvent 日志,
随后 Review 作为第二道闸门(AI 初审 + 人工终审)继续把关。

原则:全确定性、无副作用、无 IO,可单测。宁可少拦、绝不误杀真洞——
拿不准的一律放行进评审,交给第二段去裁定。
"""
from __future__ import annotations

EMPTY_SHELL = "empty_shell"
NO_VULN_NEGATION = "no_vuln_negation"

# 表示"这里其实没洞 / 没利用成功"的措辞,命中且无实证即判垃圾。
_NO_VULN_PHRASES = (
"未发现漏洞", "未发现相关漏洞", "未发现可利用", "未发现可利用漏洞",
"无漏洞", "没有漏洞", "未找到漏洞", "不存在漏洞", "不存在该漏洞",
"not vulnerable", "no vulnerability", "vulnerability not found",
"none found", "not exploitable", "no issue found",
)

# 视为"实证"的字段:任一非空即放行进评审,避免误杀。
_EVIDENCE_FIELDS = ("poc", "raw_request", "raw_response", "steps")


def screen_submission(f) -> tuple[bool, str]:
"""对单个待落库 finding 做第一道预筛。

返回 (ok, reason):ok=False 表示应挡在 Finding 主表之外(记日志即可),
ok=True 表示可进入落库(pending_review),交由第二段评审。
"""
if not isinstance(f, dict):
return False, EMPTY_SHELL

def s(key: str) -> str:
return str(f.get(key) or "").strip()

vuln_type = s("vuln_type")
title = s("title")
description = s("description")
haystack = " ".join((title, description)).lower()

has_evidence = any(s(k) for k in _EVIDENCE_FIELDS) or bool(f.get("evidence"))

if not has_evidence and not vuln_type and not title and not description:
return False, EMPTY_SHELL

if not has_evidence and any(p in haystack for p in _NO_VULN_PHRASES):
return False, NO_VULN_NEGATION

return True, ""
58 changes: 58 additions & 0 deletions tests/test_intake_screen.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,58 @@
"""覆盖两段式入库第一道闸门:Stage-0 确定性预筛 screen_submission。

原则是可单测、无副作用,且宁可少拦也不误杀真洞。
"""
import unittest

from app.tools import intake_screen


def _finding(**overrides) -> dict:
base = {
"vuln_type": "sql_injection",
"title": "某接口存在 SQL 注入",
"severity_claimed": "高危",
"description": "参数未过滤导致报错注入",
"steps": ["GET /id=1'"],
"poc": "<script>alert(1)</script>",
"raw_response": "<html>500 error</html>",
"evidence": {"marker": "sql syntax error"},
}
base.update(overrides)
return base


class TestIntakeScreen(unittest.TestCase):
def test_accepts_full_finding_with_poc(self):
ok, reason = intake_screen.screen_submission(_finding())
self.assertTrue(ok)
self.assertEqual(reason, "")

def test_accepts_finding_with_description_and_steps(self):
ok, reason = intake_screen.screen_submission(_finding(poc="", raw_response="", evidence={}))
self.assertTrue(ok, reason)

def test_rejects_fully_empty_shell(self):
ok, reason = intake_screen.screen_submission({})
self.assertFalse(ok)
self.assertEqual(reason, intake_screen.EMPTY_SHELL)

def test_rejects_no_vuln_negation_without_evidence(self):
junk = _finding(poc="", raw_response="", steps=[], evidence={},
vuln_type="", title="未发现漏洞", description="尝试注入无响应")
ok, reason = intake_screen.screen_submission(junk)
self.assertFalse(ok)
self.assertEqual(reason, intake_screen.NO_VULN_NEGATION)

def test_keeps_negation_phrase_when_evidence_present(self):
# 标题里带"未利用成功"措辞,但有完整 poc 实证 → 不误杀,放行给评审。
ok, reason = intake_screen.screen_submission(_finding(title="布尔盲注未利用成功但存在"))
self.assertTrue(ok, reason)

def test_rejects_non_dict(self):
ok, reason = intake_screen.screen_submission(None)
self.assertFalse(ok)


if __name__ == "__main__":
unittest.main()