From ec4ef35fff9a2b317ae385b91d7954cef00a3f63 Mon Sep 17 00:00:00 2001 From: ji ruo <79234113+RuoJi6@users.noreply.github.com> Date: Thu, 10 Sep 2026 13:03:20 +0800 Subject: [PATCH 1/5] =?UTF-8?q?feat:=20=E6=94=AF=E6=8C=81=E6=BC=8F?= =?UTF-8?q?=E6=B4=9E=E5=85=B3=E8=81=94=E5=A4=9A=E6=9D=A1=E6=B5=81=E9=87=8F?= =?UTF-8?q?=E8=AF=81=E6=8D=AE?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- agent/finding_recorder.go | 22 + agent/finding_recorder_test.go | 105 ++++ agent/mainagent.go | 30 +- agent/planner.go | 2 + agent/prompt_test.go | 3 + agent/promptcatalog.go | 17 +- agent/tools.go | 123 +++-- agent/worker.go | 30 +- db/finding_traffic.go | 404 +++++++++++++++ db/finding_traffic_archive.go | 85 ++++ db/finding_traffic_archive_test.go | 87 ++++ db/findings.go | 62 ++- db/schema.sql | 33 ++ db/task_archives.go | 4 +- db/task_archives_restore.go | 7 + docs/漏洞流量证据.md | 91 ++++ evidence/store.go | 401 +++++++++++++++ evidence/store_test.go | 333 ++++++++++++ evidence/testmain_test.go | 35 ++ report/findings.go | 37 +- server/dto.go | 13 +- server/finding_traffic.go | 425 ++++++++++++++++ server/finding_traffic_export.go | 121 +++++ server/finding_traffic_test.go | 473 ++++++++++++++++++ server/orchestration.go | 26 +- server/scheduler.go | 5 + server/server.go | 61 ++- server/task_archives.go | 22 +- server/task_llm.go | 3 + traffic/evidence.go | 126 +++++ .../findings/_components/findings-table.tsx | 2 + .../(main)/function/findings/detail/page.tsx | 117 ++--- .../tasks/detail/_tabs/findings-tab.tsx | 73 +-- web/src/app/(main)/function/traffic/page.tsx | 259 ++-------- web/src/components/finding-traffic-panel.tsx | 285 +++++++++++ web/src/components/http-code-block.tsx | 216 ++++++++ web/src/components/link-traffic-dialog.tsx | 163 ++++++ .../components/traffic-evidence-viewer.tsx | 181 +++++++ web/src/components/traffic-picker-dialog.tsx | 277 ++++++++++ web/src/lib/api.ts | 72 +++ web/src/lib/types.ts | 55 ++ 41 files changed, 4390 insertions(+), 496 deletions(-) create mode 100644 agent/finding_recorder.go create mode 100644 agent/finding_recorder_test.go create mode 100644 db/finding_traffic.go create mode 100644 db/finding_traffic_archive.go create mode 100644 db/finding_traffic_archive_test.go create mode 100644 docs/漏洞流量证据.md create mode 100644 evidence/store.go create mode 100644 evidence/store_test.go create mode 100644 evidence/testmain_test.go create mode 100644 server/finding_traffic.go create mode 100644 server/finding_traffic_export.go create mode 100644 server/finding_traffic_test.go create mode 100644 traffic/evidence.go create mode 100644 web/src/components/finding-traffic-panel.tsx create mode 100644 web/src/components/http-code-block.tsx create mode 100644 web/src/components/link-traffic-dialog.tsx create mode 100644 web/src/components/traffic-evidence-viewer.tsx create mode 100644 web/src/components/traffic-picker-dialog.tsx diff --git a/agent/finding_recorder.go b/agent/finding_recorder.go new file mode 100644 index 0000000..8331d5f --- /dev/null +++ b/agent/finding_recorder.go @@ -0,0 +1,22 @@ +package agent + +import ( + "context" + + "github.com/Autumn-27/artex/db" +) + +// FindingRecorder is injected by the host; agents never synthesize or copy +// evidence bodies themselves. Its implementation owns the atomic write. +type FindingRecorder interface { + Record(context.Context, db.RecordFindingInput, []db.TrafficRef) (*db.RecordedFinding, error) +} + +// Tool-use guidance is appended without replacing the user's editable prompt. +// It does not require capture or claim that unavailable traffic tools exist. +const findingTrafficGuidance = "\n\n**漏洞流量证据（可选）**：调用 report_finding 上报漏洞时，如有已查看并确认支持漏洞结论的 HTTP 请求/响应，可用 traffic_refs 按复现顺序绑定真实 ID；域名和时间只作候选筛选，不推定关联。TCP 等非 HTTP 漏洞、未采集或无确切匹配时省略或传 []，在 evidence 保留命令输出、日志等其他可验证证据，建议说明未绑定原因。不要猜测 ID，也不要仅为补包重复探测。" + +func (t *ToolSet) SetFindingRecorder(r FindingRecorder) { t.findingRecorder = r } +func (w *Worker) SetFindingRecorder(r FindingRecorder) { w.findingRecorder = r } +func (p *Planner) SetFindingRecorder(r FindingRecorder) { p.findingRecorder = r } +func (m *MainAgent) SetFindingRecorder(r FindingRecorder) { m.findingRecorder = r } diff --git a/agent/finding_recorder_test.go b/agent/finding_recorder_test.go new file mode 100644 index 0000000..57d8055 --- /dev/null +++ b/agent/finding_recorder_test.go @@ -0,0 +1,105 @@ +package agent + +import ( + "context" + "encoding/json" + "errors" + "fmt" + "strings" + "testing" + + "github.com/Autumn-27/artex/db" + "github.com/Autumn-27/artex/evidence" +) + +type failingFindingRecorder struct{} + +func (failingFindingRecorder) Record(context.Context, db.RecordFindingInput, []db.TrafficRef) (*db.RecordedFinding, error) { + return nil, errors.New("fixture persistence failed") +} + +func TestReportFindingOptionalTrafficWithoutCapture(t *testing.T) { + d := testDB(t) + defer d.Close() + task, err := d.CreateTask("TCP evidence", "fixture", nil, 0, 0) + if err != nil { + t.Fatal(err) + } + defer d.DeleteTask(task.ID) + ts := NewToolSet(nil, "") + ts.ts, ts.taskID = d.Exploration(task.ExplorationID), task.ID + ts.SetFindingRecorder(evidence.New(d, nil, t.TempDir())) + notices := 0 + ts.notifyFinding = func(int64, string) { notices++ } + for _, refs := range []string{"", `,"traffic_refs":[]`, `,"traffic_refs":null`} { + res, err := ts.addFinding().Call(t.Context(), json.RawMessage(`{"vulnclass":"TCP","severity":"low","summary":"verified TCP fixture","evidence":"command output proves the finding"`+refs+`}`), nil) + if err != nil || !strings.HasPrefix(res.Flatten(), "finding recorded: ") { + t.Fatalf("optional traffic rejected: %v %s", err, res.Flatten()) + } + var out db.RecordedFinding + if err := json.Unmarshal([]byte(strings.SplitN(res.Flatten(), "\n", 2)[1]), &out); err != nil { + t.Fatal(err) + } + f, err := d.GetFinding(out.FindingID) + if err != nil || f == nil || len(out.Traffic.Bindings) != 0 || f.Evidence != "command output proves the finding" { + t.Fatalf("lost non-HTTP evidence: %+v %v", f, err) + } + } + if notices != 3 { + t.Fatalf("successful reports notified %d times", notices) + } +} + +func TestReportFindingAtomicContract(t *testing.T) { + d := testDB(t) + defer d.Close() + task, err := d.CreateTask("report contract", "fixture", nil, 0, 0) + if err != nil { + t.Fatal(err) + } + defer d.DeleteTask(task.ID) + ts := NewToolSet(nil, "") + ts.ts = d.Exploration(task.ExplorationID) + ts.taskID = task.ID + ts.worker = "fixture" + notices := 0 + ts.notifyFinding = func(int64, string) { notices++ } + call := func(body string) string { + t.Helper() + res, err := ts.addFinding().Call(t.Context(), json.RawMessage(body), nil) + if err != nil { + t.Fatal(err) + } + return res.Flatten() + } + ts.SetFindingRecorder(failingFindingRecorder{}) + if got := call(`{"vulnclass":"TEST","summary":"fail","severity":"low","traffic_refs":[{"traffic_id":"x"}]}`); !strings.Contains(got, "persistence failed") { + t.Fatal(got) + } + if notices != 0 || ts.writes.Findings != 0 { + t.Fatal("notified before commit") + } + ts.SetFindingRecorder(nil) + got := call(`{"vulnclass":"TEST","summary":"legacy report","severity":"low"}`) + lines := strings.SplitN(got, "\n", 2) + if len(lines) != 2 { + t.Fatal(got) + } + var out db.RecordedFinding + if err = json.Unmarshal([]byte(lines[1]), &out); err != nil { + t.Fatal(err) + } + if lines[0] != fmt.Sprintf("finding recorded: %d", out.NodeID) || out.FindingID <= 0 || notices != 1 { + t.Fatal(got) + } + f, err := d.GetFinding(out.FindingID) + if err != nil || f == nil || f.NodeID == nil || *f.NodeID != out.NodeID { + t.Fatalf("wrong finding/node mapping: %+v %v", f, err) + } + if got = call(`{"vulnclass":"TEST","summary":"no storage","severity":"low","traffic_refs":[{"traffic_id":"x"}]}`); !strings.Contains(got, "未登记") { + t.Fatal(got) + } + if notices != 1 { + t.Fatal("failed tool triggered reporter") + } +} diff --git a/agent/mainagent.go b/agent/mainagent.go index e20b80c..3f6cbf2 100644 --- a/agent/mainagent.go +++ b/agent/mainagent.go @@ -17,19 +17,20 @@ import ( // (→planner) or direct high-priority intents (→frontier). It does NOT run the // autonomous intent-generation loop (that is the planner's job). type MainAgent struct { - prov llm.Provider - model string - tx *transcript.Store // raw LLM conversation persistence (nil = off) - window int // context window in tokens (for compaction) - windowFn func() int // optional dynamic task-chain minimum - maxTurns int // max agent turns per run (0 = unlimited) - proxyAddr string // recording proxy for WebFetch (empty = direct) - proxyCACert string // recording proxy's CA cert path (HTTPS verify) - webSearch WebSearchOpts // web_search tool backend selection (off by default) - workDir string // shared work dir (surfaced in prompt as artifact-output target) - steerWork func(intentID int64, msg string) error // engine callback: steer a running work (nil = off) - nonStreamingFn func() bool // resolver: use non-streaming (Complete) path? (nil = streaming) - maxTokensFn func() int // resolver: per-reply output cap (nil/0 = send no cap) + findingRecorder FindingRecorder + prov llm.Provider + model string + tx *transcript.Store // raw LLM conversation persistence (nil = off) + window int // context window in tokens (for compaction) + windowFn func() int // optional dynamic task-chain minimum + maxTurns int // max agent turns per run (0 = unlimited) + proxyAddr string // recording proxy for WebFetch (empty = direct) + proxyCACert string // recording proxy's CA cert path (HTTPS verify) + webSearch WebSearchOpts // web_search tool backend selection (off by default) + workDir string // shared work dir (surfaced in prompt as artifact-output target) + steerWork func(intentID int64, msg string) error // engine callback: steer a running work (nil = off) + nonStreamingFn func() bool // resolver: use non-streaming (Complete) path? (nil = streaming) + maxTokensFn func() int // resolver: per-reply output cap (nil/0 = send no cap) } // SetNonStreaming wires a resolver deciding whether runs use the non-streaming @@ -103,6 +104,7 @@ func mainAgentSystem(goal, dataDir, workDir string) string { // worker/planner sessions — not just the final answer. func (m *MainAgent) Chat(ctx context.Context, taskID int64, as *db.AssetStore, ts *db.ExplorationStore, goal, message string, emit func(db.Activity), notify, resume func(), notifyGoal func([]string)) (string, error) { tsx := NewToolSet(ts, "human") + tsx.SetFindingRecorder(m.findingRecorder) if as != nil { tsx.SetAssetStore(as, as.Companies()) } @@ -149,7 +151,7 @@ func (m *MainAgent) Chat(ctx context.Context, taskID int64, as *db.AssetStore, t // 命中预算(步数)→ SDK 跑收尾:向用户输出一句进展总结。Prompt 与收尾轮数可后台编辑(默认 10 轮)。 Settlement: wrapupSettlement("mainagent", nil), NonStreaming: m.nonStreaming(), // 该 profile 选非流式时走 Provider.Complete - MaxTokens: m.maxTokens(), // 0 = 不发上限,由服务端默认值决定 + MaxTokens: m.maxTokens(), // 0 = 不发上限,由服务端默认值决定 } if m.tx != nil { // persist raw human↔AI conversation; one accumulating file per task opts.Transcript = m.tx diff --git a/agent/planner.go b/agent/planner.go index aa98aec..ccf1415 100644 --- a/agent/planner.go +++ b/agent/planner.go @@ -20,6 +20,7 @@ import ( // assets, judges whether the task goal is met, and emits 0..N exploration intents // into the frontier. It is the sole intent generator. type Planner struct { + findingRecorder FindingRecorder prov llm.Provider model string tx *transcript.Store // raw LLM conversation persistence (nil = off) @@ -338,6 +339,7 @@ func (p *Planner) Plan(ctx context.Context, taskID int64, as *db.AssetStore, ts // runs in a detached goroutine so it never adds latency to this round. p.compactor.OnPlannerRound(ctx, ts) tsx := NewToolSet(ts, "planner") + tsx.SetFindingRecorder(p.findingRecorder) if as != nil { tsx.SetAssetStore(as, as.Companies()) } diff --git a/agent/prompt_test.go b/agent/prompt_test.go index 549ef1d..7dc5335 100644 --- a/agent/prompt_test.go +++ b/agent/prompt_test.go @@ -57,6 +57,9 @@ func TestRenderSystemOverrideAndFallback(t *testing.T) { if !strings.Contains(recording, "traffic_search") { t.Fatalf("worker while recording should inject trafficTool: %q", recording) } + if !strings.Contains(recording, "traffic_refs") || !strings.Contains(recording, "TCP") { + t.Fatalf("custom worker prompt lost optional finding evidence guidance: %q", recording) + } if !strings.Contains(recording, "中间产物输出规约") { t.Fatalf("worker missing artifact tail: %q", recording) } diff --git a/agent/promptcatalog.go b/agent/promptcatalog.go index bcf0636..1f0bd1b 100644 --- a/agent/promptcatalog.go +++ b/agent/promptcatalog.go @@ -52,7 +52,9 @@ const pentestDefaultTmpl = `你是一个授权渗透测试系统的"独立渗透 - 每得出一个结果**立刻**落地，别攒到最后（会话步数耗尽就全丢；记下来的才算数，活在脑子里的不算）。这些记录也是你抗 compaction 的长期记忆。 - **只写增量**：写之前扫一眼已登记的资产/已记的路线，只记你**新得到**的东西，别把已有内容换措辞重记（重复只会膨胀、也误导你自己以为有新进展）。只是印证已有结论而无新增，就不必再记。 - **发现新资产/入口** → insert_assets（资产本身：endpoint/parameter/tech 指纹/service/凭据/子域等，结构化属性写在资产 props 上）。回看已登记资产用 list_assets，避免重复登记。 -- **确认漏洞** → report_finding（含可复现 PoC）。**只有你在本次运行里真实触发过、拿到可复现证据（请求/响应或命令输出）才用它**；回看已报漏洞用 list_findings。严禁把仅凭版本/CVE 匹配、"看起来可注入"、外部漏洞库/更新日志/代码 diff 推断的东西当已确认漏洞上报。**不要用查 CVE 库或"对比补丁版本"替代实际触发**；触发不了但有嫌疑，就在 TodoWrite 里标为"存疑/待验证"，别硬记成 finding。 +- **确认漏洞** → report_finding（含可复现 PoC）。**只有你在本次运行里真实触发过、拿到可复现证据（请求/响应或命令输出）才用它**；回看已报漏洞用 list_findings。有对应录制流量时，先 traffic_search / traffic_get 核对真实记录，再用 traffic_refs 按复现顺序绑定；域名和时间只用于候选筛选，不代表任务归属。严禁把仅凭版本/CVE 匹配、"看起来可注入"、外部漏洞库/更新日志/代码 diff 推断的东西当已确认漏洞上报。**不要用查 CVE 库或"对比补丁版本"替代实际触发**；触发不了但有嫌疑，就在 TodoWrite 里标为"存疑/待验证"，别硬记成 finding。 + +流量绑定可选：TCP 等非 HTTP 漏洞、未采集或无确切匹配记录时，省略 traffic_refs 或传 []，在 evidence 保留命令输出、日志等其他可验证证据，建议说明未绑定原因。不要猜测 ID，也不要仅为补包重复探测。 ━━ 判定与收尾 ━━ - 随时对照任务目标：已被你**验证过**的成果满足了目标，就据此判定达成并说明依据。判"达成"的前提是心法 4 的自检已通过——没独立复现过的战果不算达成依据。 @@ -77,15 +79,16 @@ const ReporterDefaultPrompt = `你是一个授权渗透测试系统里的**漏 每当有 worker 调用 report_finding 登记了一个漏洞，系统就会用一段【由工具调用触发】的上下文唤起你，其中包含： - **任务 id**（task_id，见上下文"任务: #"） - report_finding 的**入参**（vulnclass / severity / summary / evidence 等） -- report_finding 的**返回**：形如 "finding recorded: " —— 这个 ** 就是本漏洞的 finding_id**，是你后面保存报告的句柄。 +- report_finding 的**返回**：形如 "finding recorded: " —— 这个 ** 是探索节点 ID**，是 get_task_node_detail 和 update_finding_report 使用的旧句柄。返回 JSON 中的 finding_id 则是独立漏洞记录 ID，get_finding_traffic 使用它。 -先从上下文里**准确抽取 finding_id 和 task_id**。抽取不到 finding_id 就不要瞎写，说明情况即可。 +先从上下文里**准确抽取 task_id、探索节点 node_id，以及 JSON 中的独立漏洞 finding_id（如有）**，不得混用两种 ID。抽取不到 node_id 就不要瞎写，说明情况即可。 ━━ 工作步骤 ━━ -1. **取全证据**：用 get_task_node_detail(task_id, id=) 读该漏洞节点的**完整证据/PoC**（触发上下文里的 evidence 可能被截断）。 -2. **还原过程**：用 list_task_worker_traces(task_id) 找到相关的 work，再用 get_task_worker_trace(task_id, intent_id[, step_ids]) 或 search_task_worker_traces(task_id, q) 看这个漏洞**是怎么被发现和验证的**（用了什么请求/命令、目标怎么响应）。必要时 get_task_graph(task_id) 看整体态势、list_task_findings(task_id) 看是否有关联漏洞。 -3. **写报告**：综合以上，写一份结构化 Markdown 报告（见下方模板）。 -4. **保存**：调用 **update_finding_report(finding_id=<那个 id>, report=)** 保存。这是你的最终产物——不写进去等于没做。 +1. **取全证据**：用 get_task_node_detail(task_id, id=) 读该漏洞节点的**完整证据/PoC**（触发上下文里的 evidence 可能被截断）。 +2. **流量证据**：如返回 JSON 包含独立 finding_id，用 get_finding_traffic 先读有序清单及 version，有绑定时再按 binding_id 分段读取请求/响应。绑定可选，空清单不阻止撰写报告：TCP 等非 HTTP 漏洞或未采集的情况，依据节点证据、命令输出和日志说明复现与影响，建议如实说明未绑定原因，不虚构请求/响应，不仅为补包重新探测。报告引用稳定证据编号及用途；仅按真实内容描述。保存报告时传入所读 version 作为 evidence_version；如版本冲突，重新读取并生成，不得直接换版本重试。 +3. **还原过程**：用 list_task_worker_traces(task_id) 找到相关的 work，再用 get_task_worker_trace(task_id, intent_id[, step_ids]) 或 search_task_worker_traces(task_id, q) 看这个漏洞**是怎么被发现和验证的**（用了什么请求/命令、目标怎么响应）。必要时 get_task_graph(task_id) 看整体态势、list_task_findings(task_id) 看是否有关联漏洞。 +4. **写报告**：综合以上，写一份结构化 Markdown 报告（见下方模板）。 +5. **保存**：调用 **update_finding_report(finding_id=, report=, evidence_version=<实际读取的 version>)** 保存；未读取版本时省略 evidence_version，不得猜测。这是你的最终产物——不写进去等于没做。 ━━ 报告结构（Markdown，按需裁剪，但证据/复现/修复必须有）━━ - ` + "`## 概述`" + `：一句话说清是什么漏洞、在哪、能造成什么。 diff --git a/agent/tools.go b/agent/tools.go index 2e89656..ad47d69 100644 --- a/agent/tools.go +++ b/agent/tools.go @@ -73,11 +73,12 @@ func compactIntents(ns []*db.Node, parentsOf, yieldsOf map[int64][]int64) []map[ // ToolSet exposes the PG-backed dual graph (asset + exploration) to an LLM agent. // One ToolSet is created per planner/worker run; per-run signals live here. type ToolSet struct { - as *db.AssetStore // asset store (optional; nil = asset tools not available) - cs *db.CompanyStore // company store (optional) - ts *db.ExplorationStore - worker string - taskID int64 // PG tasks.id; 0 when unknown (tests / orchestrator cross-task reads) + findingRecorder FindingRecorder + as *db.AssetStore // asset store (optional; nil = asset tools not available) + cs *db.CompanyStore // company store (optional) + ts *db.ExplorationStore + worker string + taskID int64 // PG tasks.id; 0 when unknown (tests / orchestrator cross-task reads) // coverageDisabled mirrors tasks.coverage_enabled=false. Stored inverted so the // zero value (all existing ToolSet constructions) means ENABLED — matching the // DB default (true). When true: graphOverviewData drops the coverage block, the @@ -1358,67 +1359,59 @@ func (t *ToolSet) goalMet() actool.CoreTool { // --- worker write tools --- func (t *ToolSet) addFinding() actool.CoreTool { - return writeTool("report_finding", "[重要]发现漏洞时必须调用该工具进行记录!记录一个确认的漏洞发现。在任务上下文中 intent_id 必填（当前正在执行的意图 id）；在会话上下文中 intent_id 可不填。", - obj(map[string]any{ - "vulnclass": str("漏洞类（分类，如 SQL Injection / IDOR / XSS）"), - "name": str("漏洞名称（具体可读的标题，如『用户中心订单接口存在越权访问』；建议填写，留空时前端回退展示 vulnclass）"), - "severity": str("critical|high|medium|low（严重/高/中/低）"), - "summary": str("发现摘要"), - "intent_id": idp("产生本发现的意图 id（任务上下文必填；会话上下文可不填）"), - "asset_ids": map[string]any{"type": "array", "items": map[string]any{"type": "integer"}, "description": "【存在时尽量填写，否则在摘要中必须要写清楚漏洞位置】受影响资产 id（可选，0/1/多个）：参数/端点/站点等。一个漏洞影响多处可全填，纯观察可不填。"}, - "evidence": str("证据/PoC 文本"), - }, "vulnclass", "severity", "summary"), - func(_ context.Context, in json.RawMessage) (actool.Result, error) { - var a struct { - VulnClass, Name, Severity, Summary, Evidence string - IntentID json.RawMessage `json:"intent_id"` - AssetIDs []json.RawMessage `json:"asset_ids"` + return writeTool("report_finding", "记录确认的漏洞。可选 traffic_refs 绑定多条真实 HTTP 请求/响应：先用 traffic_search / traffic_get 核对 ID；提交的流量必须全部绑定成功。TCP 等非 HTTP 漏洞、未采集或无确切匹配记录时可省略 traffic_refs 或传空数组，仍用 evidence 提供命令输出、日志等可验证证据，建议说明未绑定原因。任务上下文传当前 intent_id。", obj(map[string]any{ + "vulnclass": str("漏洞类别"), "name": str("漏洞名称"), "severity": str("critical|high|medium|low"), "summary": str("发现摘要"), + "intent_id": idp("当前任务的意图 id"), "asset_ids": map[string]any{"type": "array", "items": map[string]any{"type": "integer"}, "description": "受影响资产 id"}, + "evidence": str("证据/PoC 文本"), + "traffic_refs": map[string]any{"type": "array", "description": "可选；HTTP/HTTPS 漏洞先检索并逐条核实请求/响应确实支持漏洞结论，再按复现顺序填写真实 ID。TCP 等非 HTTP 漏洞、未采集或找不到确切记录时省略或传 []，不阻止上报；可在 evidence 说明原因并提供其他可验证证据。不要猜测 ID、按域名/时间推定关联或仅为补包重复探测。用途 baseline 正常对照 / proof 漏洞证明 / verification 补充验证 / supporting 辅助证据。", + "items": obj(map[string]any{"traffic_id": str("traffic_search 返回的真实流量 ID"), "role": map[string]any{"type": "string", "enum": []string{"baseline", "proof", "verification", "supporting"}}, "note": str("该流量如何支持漏洞结论")}, "traffic_id")}, + }, "vulnclass", "severity", "summary"), func(ctx context.Context, in json.RawMessage) (actool.Result, error) { + var a struct { + VulnClass, Name, Severity, Summary, Evidence string + IntentID json.RawMessage `json:"intent_id"` + AssetIDs []json.RawMessage `json:"asset_ids"` + TrafficRefs []db.TrafficRef `json:"traffic_refs"` + } + if err := json.Unmarshal(in, &a); err != nil { + return actool.Errorf(err.Error()), nil + } + if t.ts == nil { + return actool.Errorf("report_finding 需要任务上下文（exploration store 未初始化）"), nil + } + refs, err := db.NormalizeTrafficRefs(a.TrafficRefs) + if err != nil { + return actool.Errorf(err.Error()), nil + } + input := db.RecordFindingInput{TaskID: t.taskID, ExplorationID: t.ts.ID(), IntentID: pid(a.IntentID), VulnClass: a.VulnClass, Name: a.Name, Severity: a.Severity, Summary: a.Summary, Evidence: a.Evidence, Worker: t.worker, AssetIDs: pidList(a.AssetIDs)} + var recorded *db.RecordedFinding + if t.findingRecorder != nil { + recorded, err = t.findingRecorder.Record(ctx, input, refs) + } else if len(refs) > 0 { + return actool.Errorf("流量证据存储不可用；未登记漏洞"), nil + } else { + recorded, err = t.ts.RecordFinding(ctx, input) + } + if err != nil { + return actool.Errorf(err.Error()), nil + } + if t.notifyFinding != nil { + iid := input.IntentID + if iid <= 0 { + iid = t.ownerNode } - _ = json.Unmarshal(in, &a) - payload := map[string]any{"vulnclass": a.VulnClass, "name": a.Name, "severity": a.Severity, "summary": a.Summary, - "evidence": map[string]any{"by": t.worker, "poc": a.Evidence}} - var anchors []int64 - for _, raw := range a.AssetIDs { - if p := pid(raw); p > 0 { - anchors = append(anchors, p) - } - } - var id int64 - if t.ts != nil { - intent := pid(a.IntentID) - if intent > 0 { - node, err := t.ts.GetNode(intent) - if err != nil || node == nil || node.Kind != db.KindIntent { - return actool.Errorf("intent_id 必须是本任务的意图（关联任务意图只读）"), nil - } - } - var err error - id, err = t.ts.AddNode(db.KindFinding, payload, 9, "confirmed", t.worker, anchors) - if err != nil { - return actool.Errorf(err.Error()), nil - } - if intent > 0 { - _ = t.ts.Link(intent, db.RelYields, id) // chain: intent -> finding - } - _, _ = t.ts.AddStandaloneFinding(t.taskID, id, a.VulnClass, a.Name, a.Severity, a.Summary, a.Evidence, t.worker, anchors) - // 确证漏洞落库 → 当场唤醒本任务 planner（不等 worker 收工，debounce 合并）。 - // 优先带上下文(哪个意图+finding摘要);intent 用工具参数,缺省回退到 owner 意图。 - if t.notifyFinding != nil { - iid := pid(a.IntentID) - if iid <= 0 { - iid = t.ownerNode - } - t.notifyFinding(iid, a.Summary) - } else if t.notify != nil { - t.notify() - } - } else { - // conversation context: no exploration store available, cannot record finding - return actool.Errorf("report_finding 需要任务上下文（exploration store 未初始化）"), nil - } - t.writes.Findings++ - return actool.Text(fmt.Sprintf("finding recorded: %d", id)), nil - }) + t.notifyFinding(iid, a.Summary) + } else if t.notify != nil { + t.notify() + } + t.writes.Findings++ + // Keep the first line's node-ID contract for existing reporter triggers. + for i := range recorded.Traffic.Bindings { + recorded.Traffic.Bindings[i].Snapshot.ReqHead = "" + recorded.Traffic.Bindings[i].Snapshot.RespHead = "" + } + raw, _ := json.Marshal(recorded) + return actool.Text(fmt.Sprintf("finding recorded: %d\n%s", recorded.NodeID, raw)), nil + }) } // recordFact writes a general exploration RESULT/conclusion (not a vuln, not a diff --git a/agent/worker.go b/agent/worker.go index 1151f47..a8781b5 100644 --- a/agent/worker.go +++ b/agent/worker.go @@ -41,17 +41,18 @@ type WebSearchOpts struct { } type Worker struct { - prov llm.Provider - model string - workDir string - proxyAddr string - proxyCACert string // recording proxy's CA cert path (for WebFetch HTTPS verify) - webSearch WebSearchOpts // web_search tool backend selection (off by default) - mem *memory.Store // cross-engagement tradecraft memory (G4) - tx *transcript.Store // raw LLM conversation persistence (nil = off) - window int // context window in tokens (for compaction) - windowFn func() int // optional dynamic task-chain minimum - maxTurns int // max agent turns per run (0 = unlimited) + findingRecorder FindingRecorder + prov llm.Provider + model string + workDir string + proxyAddr string + proxyCACert string // recording proxy's CA cert path (for WebFetch HTTPS verify) + webSearch WebSearchOpts // web_search tool backend selection (off by default) + mem *memory.Store // cross-engagement tradecraft memory (G4) + tx *transcript.Store // raw LLM conversation persistence (nil = off) + window int // context window in tokens (for compaction) + windowFn func() int // optional dynamic task-chain minimum + maxTurns int // max agent turns per run (0 = unlimited) // runTimeout is the wall-clock budget for the main exploration of one intent // (0 = unlimited). When it fires, the run is cut and a settlement round is // forced so already-identified facts get written back instead of being lost. @@ -204,7 +205,9 @@ const workerDefaultTmpl = `你是一个 ARTEX 平台授权渗透测试系统的" **边发现边写回**（写进图才算数，脑子/文字里的不算；每得一个结果立刻写，别攒到最后被步数耗尽丢掉）。三种写回，别串图： - **新资产/资源 → insert_assets（资产图）**：子域 / service / endpoint / 指纹 / 凭据 等一切资产【本身】。**这里只登记资产；探索结论/判断不写这里，用 record_fact。** - **探索结论/事实 → record_fact（探索图，传 intent_id）**：都用它。**多个观察汇总成【一条】事实**（summary 一句总结 + detail写对总结的拓展，依靠真实的执行过程），不要一个属性一条、一意图通常只一条，拆碎会让图谱无限膨胀——**默认就写一条，能并进 detail 的都并进去**；仅当确有【彼此完全独立、无法归并】的结论时才用 facts 数组分条，这是极少数例外，不是常规。**只写增量**：只记这次【新得到】的，别把已有事实换措辞重记（只印证已有、无新增就不必记）。**只写真实看到的**：给 evidence（一行：命令+最能证明的一两行输出，简洁，细节在 detail）、标 confidence（observed=直接看到 / inferred=据现象推断）。 -- **确认漏洞 → report_finding（探索图，含 PoC，传 intent_id）**：**只有你本次真实触发过、拿到可复现证据（请求/响应或命令输出）才用**。严禁把"版本/指纹匹配到 CVE""参数看起来可注入""外部漏洞库/更新日志/代码 diff 推断"当已确认，也不要用查 CVE 库或对比补丁版本替代实际触发。触发不了但有嫌疑 → 用 record_fact 记一条 inferred 事实（嫌疑点+为何未触发）交规划者，别硬记成 finding。 +- **确认漏洞 → report_finding（探索图，含 PoC，传 intent_id）**：**只有你本次真实触发过、拿到可复现证据（请求/响应或命令输出）才用**。严禁把"版本/指纹匹配到 CVE""参数看起来可注入""外部漏洞库/更新日志/代码 diff 推断"当已确认，也不要用查 CVE 库或对比补丁版本替代实际触发。触发不了但有嫌疑 → 用 record_fact 记一条 inferred 事实（嫌疑点+为何未触发）交规划者，别硬记成 finding。有对应录制流量时，先用 traffic_search / traffic_get 核对真实 ID，再用 traffic_refs 按复现顺序绑定多条请求/响应；域名和时间只用于候选筛选，不推定任务归属。 + +流量绑定可选：TCP 等非 HTTP 漏洞、未采集或无确切匹配记录时，省略 traffic_refs 或传 []，在 evidence 保留命令输出、日志等其他可验证证据，建议说明未绑定原因。不要猜测 ID，也不要仅为补包重复探测。 完成本意图后用一句话总结你做了什么、写回了哪些事实。务实、克制、聚焦这一条意图。` @@ -255,7 +258,7 @@ func workerSystem(proxyAddr, caCert, dataDir, runDir string) string { body := renderSystem("worker", workerDefaultTmpl, WorkerVars{ProxyAddr: proxyAddr, DataDir: dataDir, Now: nowStr()}) // caCert is present only when the recording MITM is on, which is exactly when // the traffic_* tools are registered — so it gates the traffic-tool note. - return body + workerTrafficBlock(caCert != "") + workerArtifactSpec(runDir) + return body + workerTrafficBlock(caCert != "") + findingTrafficGuidance + workerArtifactSpec(runDir) } // renderIntentTask formats the claimed intent for the worker's launch USER message: @@ -325,6 +328,7 @@ func (w *Worker) ExecuteWithMessage(ctx context.Context, name string, taskID int func (w *Worker) execute(ctx context.Context, name string, taskID int64, as *db.AssetStore, ts *db.ExplorationStore, intent *db.Node, hooks harness.HookRunner, emit func(db.Activity), enr EnrichTrigger, notifyFinding func(int64, string), requestID, message string) (harness.TerminalReason, WriteCounts, error) { tsx := NewToolSet(ts, name) + tsx.SetFindingRecorder(w.findingRecorder) tsx.SetTaskID(taskID) coverageEnabled := as == nil || as.CoverageEnabled(taskID) tsx.SetCoverageEnabled(coverageEnabled) diff --git a/db/finding_traffic.go b/db/finding_traffic.go new file mode 100644 index 0000000..a1ffb63 --- /dev/null +++ b/db/finding_traffic.go @@ -0,0 +1,404 @@ +package db + +import ( + "context" + "crypto/sha256" + "database/sql" + "encoding/hex" + "encoding/json" + "errors" + "fmt" + "strings" + "time" +) + +var ( + ErrEvidenceConflict = errors.New("流量证据已变更，请刷新后重试") + ErrFindingNotFound = errors.New("漏洞不存在") + ErrEvidenceNotFound = errors.New("流量证据不存在") +) + +// This lock covers the evidence filesystem as well as its SQL references. All +// processes sharing the database use it, including readers, exports and GC. +const findingEvidenceLockKey int64 = 7337741004 + +func (d *DB) WithEvidenceTx(ctx context.Context, fn func(*sql.Tx) error) error { + tx, err := d.BeginTx(ctx, nil) + if err != nil { + return err + } + defer tx.Rollback() + if _, err = tx.ExecContext(ctx, `SELECT pg_advisory_xact_lock($1)`, findingEvidenceLockKey); err != nil { + return err + } + if err = fn(tx); err != nil { + return err + } + return tx.Commit() +} + +type TrafficRef struct { + TrafficID string `json:"traffic_id"` + Role string `json:"role"` + Note string `json:"note"` +} + +func NormalizeTrafficRefs(refs []TrafficRef) ([]TrafficRef, error) { + out := make([]TrafficRef, 0, len(refs)) + seen := map[string]bool{} + for _, ref := range refs { + ref.TrafficID = strings.TrimSpace(ref.TrafficID) + if ref.TrafficID == "" { + return nil, errors.New("traffic_id 不能为空") + } + if ref.Role == "" { + ref.Role = "supporting" + } + if !ValidTrafficRole(ref.Role) { + return nil, fmt.Errorf("无效的流量用途 %q", ref.Role) + } + if !seen[ref.TrafficID] { + out = append(out, ref) + seen[ref.TrafficID] = true + } + } + return out, nil +} + +func ValidTrafficRole(role string) bool { + return role == "baseline" || role == "proof" || role == "verification" || role == "supporting" +} + +type TrafficEvidenceSnapshot struct { + ID string `json:"id"` + SourceTrafficID string `json:"source_traffic_id"` + CapturedAt int64 `json:"captured_at"` + URL string `json:"url"` + Method string `json:"method"` + Status int `json:"status"` + ContentType string `json:"content_type"` + ReqHead string `json:"req_head,omitempty"` + RespHead string `json:"resp_head,omitempty"` + ReqHash string `json:"req_hash"` + RespHash string `json:"resp_hash"` + ReqLen int64 `json:"req_len"` + RespLen int64 `json:"resp_len"` + CreatedAt time.Time `json:"created_at"` +} + +type FindingTrafficBinding struct { + ID int64 `json:"id,string"` + FindingID int64 `json:"finding_id,string"` + SnapshotID string `json:"snapshot_id"` + Role string `json:"role"` + Note string `json:"note"` + Position int `json:"position"` + CreatedAt time.Time `json:"created_at"` + Snapshot TrafficEvidenceSnapshot `json:"snapshot"` +} + +type FindingTraffic struct { + FindingID int64 `json:"finding_id,string"` + Version int64 `json:"version"` + ReportVersion int64 `json:"report_version"` + Bindings []FindingTrafficBinding `json:"bindings"` +} + +type PreparedTrafficEvidence struct { + Ref TrafficRef + Snapshot TrafficEvidenceSnapshot +} + +func TrafficSnapshotID(snapshot TrafficEvidenceSnapshot) string { + snapshot.ID = "" + snapshot.CreatedAt = time.Time{} + raw, _ := json.Marshal(snapshot) + sum := sha256.Sum256(raw) + return hex.EncodeToString(sum[:]) +} + +// LockTaskEvidenceTx serializes writes with archive queueing (which locks the +// same task row). Once queued, its snapshot must not acquire new evidence. +func LockTaskEvidenceTx(tx *sql.Tx, taskID int64) error { + if taskID == 0 { + return nil + } + var deleted sql.NullTime + if err := tx.QueryRow(`SELECT deleted_at FROM tasks WHERE id=$1 FOR UPDATE`, taskID).Scan(&deleted); err != nil { + return err + } + if deleted.Valid { + return ErrTaskArchiveState + } + var state string + err := tx.QueryRow(`SELECT state FROM task_archives WHERE task_id=$1`, taskID).Scan(&state) + if err != nil && !errors.Is(err, sql.ErrNoRows) { + return err + } + if err == nil && state != ArchiveFailed { + return ErrTaskArchiveState + } + return nil +} + +func LockFindingEvidenceTx(tx *sql.Tx, findingID int64, version *int64) error { + var taskID sql.NullInt64 + if err := tx.QueryRow(`SELECT task_id FROM findings WHERE id=$1`, findingID).Scan(&taskID); err != nil { + if errors.Is(err, sql.ErrNoRows) { + return ErrFindingNotFound + } + return err + } + if err := LockTaskEvidenceTx(tx, taskID.Int64); err != nil { + return err + } + var current int64 + if err := tx.QueryRow(`SELECT evidence_version FROM findings WHERE id=$1 FOR UPDATE`, findingID).Scan(¤t); err != nil { + if errors.Is(err, sql.ErrNoRows) { + return ErrFindingNotFound + } + return err + } + if version != nil && *version != current { + return ErrEvidenceConflict + } + return nil +} + +func InsertEvidenceSnapshotTx(tx *sql.Tx, s TrafficEvidenceSnapshot) error { + if s.ID != TrafficSnapshotID(s) { + return errors.New("证据快照元数据哈希不匹配") + } + if s.CreatedAt.IsZero() { + s.CreatedAt = time.Now().UTC() + } + _, err := tx.Exec(`INSERT INTO traffic_evidence_snapshots +(id,source_traffic_id,captured_at,url,method,status,content_type,req_head,resp_head,req_hash,resp_hash,req_len,resp_len,created_at) +VALUES($1,$2,$3,$4,$5,$6,$7,$8,$9,$10,$11,$12,$13,$14) ON CONFLICT(id) DO NOTHING`, + s.ID, s.SourceTrafficID, s.CapturedAt, s.URL, s.Method, s.Status, s.ContentType, s.ReqHead, s.RespHead, s.ReqHash, s.RespHash, s.ReqLen, s.RespLen, s.CreatedAt) + return err +} + +func AddFindingTrafficTx(tx *sql.Tx, findingID int64, prepared []PreparedTrafficEvidence) error { + var position int + if err := tx.QueryRow(`SELECT COALESCE(MAX(position)+1,0) FROM finding_traffic_bindings WHERE finding_id=$1`, findingID).Scan(&position); err != nil { + return err + } + changed := false + for _, item := range prepared { + if err := InsertEvidenceSnapshotTx(tx, item.Snapshot); err != nil { + return err + } + if _, err := tx.Exec(`UPDATE traffic_evidence_snapshots SET unreferenced_at=NULL WHERE id=$1`, item.Snapshot.ID); err != nil { + return err + } + res, err := tx.Exec(`INSERT INTO finding_traffic_bindings(finding_id,snapshot_id,role,note,position) +VALUES($1,$2,$3,$4,$5) ON CONFLICT(finding_id,snapshot_id) DO NOTHING`, findingID, item.Snapshot.ID, item.Ref.Role, item.Ref.Note, position) + if err != nil { + return err + } + n, err := res.RowsAffected() + if err != nil { + return err + } + if n > 0 { + changed = true + position++ + } + } + if changed { + return bumpEvidenceVersionTx(tx, findingID) + } + return nil +} + +func bumpEvidenceVersionTx(tx *sql.Tx, findingID int64) error { + _, err := tx.Exec(`UPDATE findings SET evidence_version=evidence_version+1 WHERE id=$1`, findingID) + return err +} + +func FindingTrafficTx(tx *sql.Tx, findingID int64) (*FindingTraffic, error) { + out := &FindingTraffic{FindingID: findingID, Bindings: []FindingTrafficBinding{}} + if err := tx.QueryRow(`SELECT evidence_version,report_evidence_version FROM findings WHERE id=$1`, findingID).Scan(&out.Version, &out.ReportVersion); err != nil { + if errors.Is(err, sql.ErrNoRows) { + return nil, ErrFindingNotFound + } + return nil, err + } + rows, err := tx.Query(`SELECT b.id,b.finding_id,b.snapshot_id,b.role,b.note,b.position,b.created_at,to_jsonb(s) +FROM finding_traffic_bindings b JOIN traffic_evidence_snapshots s ON s.id=b.snapshot_id WHERE b.finding_id=$1 ORDER BY b.position,b.id`, findingID) + if err != nil { + return nil, err + } + defer rows.Close() + for rows.Next() { + var b FindingTrafficBinding + var raw []byte + if err := rows.Scan(&b.ID, &b.FindingID, &b.SnapshotID, &b.Role, &b.Note, &b.Position, &b.CreatedAt, &raw); err != nil { + return nil, err + } + if err := json.Unmarshal(raw, &b.Snapshot); err != nil { + return nil, err + } + out.Bindings = append(out.Bindings, b) + } + return out, rows.Err() +} + +func (d *DB) GetFindingTraffic(ctx context.Context, findingID int64) (out *FindingTraffic, err error) { + err = d.WithEvidenceTx(ctx, func(tx *sql.Tx) error { var e error; out, e = FindingTrafficTx(tx, findingID); return e }) + return +} + +func (d *DB) EditFindingTraffic(ctx context.Context, findingID, bindingID, version int64, role, note *string, remove bool, order []int64) error { + if role != nil && !ValidTrafficRole(*role) { + return errors.New("无效的流量用途") + } + return d.WithEvidenceTx(ctx, func(tx *sql.Tx) error { + if err := LockFindingEvidenceTx(tx, findingID, &version); err != nil { + return err + } + if order != nil { + current, err := FindingTrafficTx(tx, findingID) + if err != nil { + return err + } + if len(order) != len(current.Bindings) { + return ErrEvidenceConflict + } + ids := map[int64]bool{} + for _, b := range current.Bindings { + ids[b.ID] = true + } + for position, id := range order { + if !ids[id] { + return ErrEvidenceConflict + } + delete(ids, id) + if _, err := tx.Exec(`UPDATE finding_traffic_bindings SET position=$1 WHERE finding_id=$2 AND id=$3`, position, findingID, id); err != nil { + return err + } + } + } else { + var res sql.Result + var err error + if remove { + res, err = tx.Exec(`DELETE FROM finding_traffic_bindings WHERE finding_id=$1 AND id=$2`, findingID, bindingID) + } else { + res, err = tx.Exec(`UPDATE finding_traffic_bindings SET role=COALESCE($3,role),note=COALESCE($4,note) WHERE finding_id=$1 AND id=$2`, findingID, bindingID, role, note) + } + if err != nil { + return err + } + n, err := res.RowsAffected() + if err != nil { + return err + } + if n == 0 { + return ErrEvidenceNotFound + } + } + return bumpEvidenceVersionTx(tx, findingID) + }) +} + +type RecordFindingInput struct { + TaskID, ExplorationID, IntentID int64 + VulnClass, Name, Severity, Summary, Evidence, Worker string + AssetIDs []int64 +} + +type RecordedFinding struct { + FindingID int64 `json:"finding_id,string"` + NodeID int64 `json:"finding_node_id,string"` + Traffic *FindingTraffic `json:"traffic"` +} + +func RecordFindingTx(tx *sql.Tx, in RecordFindingInput, prepared []PreparedTrafficEvidence) (*RecordedFinding, error) { + if err := LockTaskEvidenceTx(tx, in.TaskID); err != nil { + return nil, err + } + if in.TaskID > 0 { + var expID int64 + if err := tx.QueryRow(`SELECT exploration_id FROM tasks WHERE id=$1`, in.TaskID).Scan(&expID); err != nil { + return nil, err + } + if expID != in.ExplorationID { + return nil, errors.New("漏洞所属任务与探索记录不匹配") + } + } + if in.IntentID > 0 { + var ok bool + if err := tx.QueryRow(`SELECT EXISTS(SELECT 1 FROM exploration_nodes WHERE id=$1 AND exploration_id=$2 AND kind='intent')`, in.IntentID, in.ExplorationID).Scan(&ok); err != nil { + return nil, err + } + if !ok { + return nil, errors.New("intent_id 必须是本任务的意图（关联任务意图只读）") + } + } + payload, _ := json.Marshal(map[string]any{"vulnclass": in.VulnClass, "name": in.Name, "severity": in.Severity, "summary": in.Summary, "evidence": map[string]any{"by": in.Worker, "poc": in.Evidence}}) + out := &RecordedFinding{} + if err := tx.QueryRow(`INSERT INTO exploration_nodes(exploration_id,kind,payload,priority,state,origin) +VALUES($1,'finding',$2,9,'confirmed',$3) RETURNING id`, in.ExplorationID, string(payload), in.Worker).Scan(&out.NodeID); err != nil { + return nil, err + } + for _, asset := range in.AssetIDs { + if _, err := tx.Exec(`INSERT INTO exploration_anchors(node_id,asset_id) VALUES($1,$2) ON CONFLICT DO NOTHING`, out.NodeID, asset); err != nil { + return nil, err + } + } + if in.IntentID > 0 { + if _, err := tx.Exec(`INSERT INTO exploration_edges(exploration_id,src_id,rel,dst_id) VALUES($1,$2,$3,$4)`, in.ExplorationID, in.IntentID, RelYields, out.NodeID); err != nil { + return nil, err + } + } + assets := in.AssetIDs + if assets == nil { + assets = []int64{} + } + raw, _ := json.Marshal(assets) + if err := tx.QueryRow(`INSERT INTO findings(task_id,node_id,vulnclass,name,severity,summary,evidence,worker,asset_ids) +VALUES(NULLIF($1,0),$2,$3,$4,$5,$6,$7,$8,$9) RETURNING id`, in.TaskID, out.NodeID, in.VulnClass, in.Name, in.Severity, in.Summary, in.Evidence, in.Worker, string(raw)).Scan(&out.FindingID); err != nil { + return nil, err + } + if err := AddFindingTrafficTx(tx, out.FindingID, prepared); err != nil { + return nil, err + } + var err error + out.Traffic, err = FindingTrafficTx(tx, out.FindingID) + return out, err +} + +// RecordFinding is the atomic legacy/no-recorder path used by tools and tests. +func (s *ExplorationStore) RecordFinding(ctx context.Context, in RecordFindingInput) (out *RecordedFinding, err error) { + in.ExplorationID = s.expID + err = s.db.WithEvidenceTx(ctx, func(tx *sql.Tx) error { var e error; out, e = RecordFindingTx(tx, in, nil); return e }) + return +} + +func (d *DB) FindingIDByNodeID(nodeID int64) (id int64, err error) { + err = d.QueryRow(`SELECT id FROM findings WHERE node_id=$1`, nodeID).Scan(&id) + return +} + +func (d *DB) SetFindingReportVersionByNodeID(ctx context.Context, nodeID int64, report string, version *int64) (n int64, err error) { + err = d.WithEvidenceTx(ctx, func(tx *sql.Tx) error { + var id int64 + if err := tx.QueryRow(`SELECT id FROM findings WHERE node_id=$1`, nodeID).Scan(&id); err != nil { + if errors.Is(err, sql.ErrNoRows) { + return nil + } + return err + } + if err := LockFindingEvidenceTx(tx, id, version); err != nil { + return err + } + res, err := tx.Exec(`UPDATE findings SET report=$2,report_evidence_version=COALESCE($3::bigint,CASE WHEN evidence_version=0 THEN 0 ELSE -1 END) WHERE id=$1`, id, report, version) + if err != nil { + return err + } + n, err = res.RowsAffected() + return err + }) + return +} diff --git a/db/finding_traffic_archive.go b/db/finding_traffic_archive.go new file mode 100644 index 0000000..8e496af --- /dev/null +++ b/db/finding_traffic_archive.go @@ -0,0 +1,85 @@ +package db + +import ( + "database/sql" + "encoding/json" + "errors" + "fmt" +) + +func ArchiveEvidenceSnapshots(snapshot *TaskArchiveSnapshot) ([]TrafficEvidenceSnapshot, error) { + var out []TrafficEvidenceSnapshot + if rawRowCount(snapshot.Tables["traffic_evidence_snapshots"]) == 0 { + return out, nil + } + if snapshot.FormatVersion < 3 { + return nil, ErrTaskArchiveFormatMismatch + } + err := json.Unmarshal(snapshot.Tables["traffic_evidence_snapshots"], &out) + return out, err +} + +func restoreFindingTrafficTx(tx *sql.Tx, snapshot *TaskArchiveSnapshot) error { + snapshots, err := ArchiveEvidenceSnapshots(snapshot) + if err != nil { + return err + } + allowed := map[string]bool{} + for _, v := range snapshots { + if allowed[v.ID] { + return errors.New("duplicate archived evidence snapshot") + } + allowed[v.ID] = true + if err = InsertEvidenceSnapshotTx(tx, v); err != nil { + return err + } + } + rows, err := decodeArchiveRows(snapshot.Tables["finding_traffic_bindings"]) + if err != nil { + return err + } + if len(rows) > 0 && snapshot.FormatVersion < 3 { + return ErrTaskArchiveFormatMismatch + } + for _, row := range rows { + fid, ok := jsonInt64(row["finding_id"]) + if !ok { + return errors.New("invalid archived evidence finding id") + } + sid, _ := row["snapshot_id"].(string) + role, _ := row["role"].(string) + if !allowed[sid] || !ValidTrafficRole(role) { + return errors.New("invalid archived evidence binding") + } + var owned bool + if err = tx.QueryRow(`SELECT EXISTS(SELECT 1 FROM findings WHERE id=$1 AND task_id=$2)`, fid, snapshot.TaskID).Scan(&owned); err != nil { + return err + } + if !owned { + return fmt.Errorf("evidence references finding outside archived task: %d", fid) + } + } + if len(rows) > 0 { + raw, _ := json.Marshal(rows) + if _, err = tx.Exec(`INSERT INTO finding_traffic_bindings SELECT * FROM json_populate_recordset(NULL::finding_traffic_bindings,$1::json)`, string(raw)); err != nil { + return err + } + _, err = tx.Exec(`UPDATE traffic_evidence_snapshots s SET unreferenced_at=NULL WHERE EXISTS(SELECT 1 FROM finding_traffic_bindings b WHERE b.snapshot_id=s.id)`) + } + return err +} + +func normalizeArchivedFindingVersions(raw json.RawMessage) (json.RawMessage, error) { + rows, err := decodeArchiveRows(raw) + if err != nil { + return nil, err + } + for _, row := range rows { + for _, key := range []string{"evidence_version", "report_evidence_version"} { + if row[key] == nil { + row[key] = 0 + } + } + } + return json.Marshal(rows) +} diff --git a/db/finding_traffic_archive_test.go b/db/finding_traffic_archive_test.go new file mode 100644 index 0000000..2460407 --- /dev/null +++ b/db/finding_traffic_archive_test.go @@ -0,0 +1,87 @@ +package db + +import ( + "encoding/json" + "fmt" + "testing" +) + +func TestFindingEvidenceLockNamespace(t *testing.T) { + for _, reserved := range []int64{7337741001, 7337741002, 7337741003} { + if findingEvidenceLockKey == reserved { + t.Fatal("evidence lock collides with migration/test/company lock") + } + } +} + +func TestFindingTrafficLegacyArchiveDefaults(t *testing.T) { + d, err := Open(testDSN(t)) + if err != nil { + t.Fatal(err) + } + defer d.Close() + if err = d.EnsureLLMRecordsTable(); err != nil { + t.Fatal(err) + } + if err = d.EnsureLLMUsageTable(); err != nil { + t.Fatal(err) + } + for _, version := range []int{1, 2} { + t.Run(fmt.Sprint(version), func(t *testing.T) { + task, err := d.CreateTask("legacy archive evidence defaults", "fixture", nil, 0, 0) + if err != nil { + t.Fatal(err) + } + defer func() { d.Exec(`DELETE FROM task_archives WHERE task_id=$1`, task.ID); d.DeleteTask(task.ID) }() + f, err := d.Exploration(task.ExplorationID).RecordFinding(t.Context(), RecordFindingInput{TaskID: task.ID, ExplorationID: task.ExplorationID, Summary: "legacy", Severity: "low"}) + if err != nil { + t.Fatal(err) + } + if err = d.SetPaused(task.ID, true); err != nil { + t.Fatal(err) + } + archive, err := d.QueueTaskArchive(task.ID) + if err != nil { + t.Fatal(err) + } + if _, err = d.ClaimTaskArchiveJob(t.Context()); err != nil { + t.Fatal(err) + } + snapshot, err := d.SnapshotTaskArchive(task.ID) + if err != nil { + t.Fatal(err) + } + if err = d.CompleteTaskArchive(archive.ID, snapshot, "/tmp/legacy-test.tar.zst", "fixture", 1, 1); err != nil { + t.Fatal(err) + } + snapshot.FormatVersion = version + delete(snapshot.Tables, "finding_traffic_bindings") + delete(snapshot.Tables, "traffic_evidence_snapshots") + var rows []map[string]any + if err = json.Unmarshal(snapshot.Tables["findings"], &rows); err != nil { + t.Fatal(err) + } + for _, row := range rows { + delete(row, "evidence_version") + delete(row, "report_evidence_version") + } + snapshot.Tables["findings"], _ = json.Marshal(rows) + if _, err = d.QueueTaskArchiveRestore(archive.ID); err != nil { + t.Fatal(err) + } + if _, err = d.ClaimTaskArchiveJob(t.Context()); err != nil { + t.Fatal(err) + } + if _, err = d.RestoreTaskArchive(archive.ID, snapshot, 0); err != nil { + t.Fatal(err) + } + got, err := d.GetFindingTraffic(t.Context(), f.FindingID) + if err != nil { + t.Fatal(err) + } + if got.Version != 0 || got.ReportVersion != 0 || len(got.Bindings) != 0 { + t.Fatal(got) + } + }) + } +} diff --git a/db/findings.go b/db/findings.go index 7b1c4fe..e8fec99 100644 --- a/db/findings.go +++ b/db/findings.go @@ -1,6 +1,7 @@ package db import ( + "context" "database/sql" "encoding/json" "errors" @@ -14,6 +15,11 @@ import ( // DBFinding is a row in the standalone findings table. It persists across task // deletion unless the caller explicitly requests related finding cleanup. type DBFinding struct { + TrafficCount int + EvidenceVersion int64 + ReportEvidenceVersion int64 + TrafficBindings []FindingTrafficBinding // populated only for export + ID int64 TaskID *int64 NodeID *int64 @@ -97,7 +103,8 @@ func (d *DB) AddFinding(taskID, nodeID int64, vulnclass, name, severity, summary // list query selects, so scanFinding stays in sync across callers. const findingSelectCols = `f.id, f.task_id, f.node_id, f.vulnclass, COALESCE(f.name, ''), f.severity, f.summary, f.evidence, f.worker, f.asset_ids, COALESCE(f.status, 'pending'), f.created_at, - COALESCE(t.description, '') AS task_description` + COALESCE(t.description, '') AS task_description, f.evidence_version, f.report_evidence_version, + (SELECT count(*) FROM finding_traffic_bindings b WHERE b.finding_id=f.id)` // scanFindings materializes rows selected via findingSelectCols. func scanFindings(rows interface { @@ -110,7 +117,7 @@ func scanFindings(rows interface { f := &DBFinding{} var aidsJSON string if err := rows.Scan(&f.ID, &f.TaskID, &f.NodeID, &f.VulnClass, &f.Name, &f.Severity, - &f.Summary, &f.Evidence, &f.Worker, &aidsJSON, &f.Status, &f.CreatedAt, &f.TaskDescription); err != nil { + &f.Summary, &f.Evidence, &f.Worker, &aidsJSON, &f.Status, &f.CreatedAt, &f.TaskDescription, &f.EvidenceVersion, &f.ReportEvidenceVersion, &f.TrafficCount); err != nil { return nil, err } _ = json.Unmarshal([]byte(aidsJSON), &f.AssetIDs) @@ -503,7 +510,7 @@ func (d *DB) ListFindingsForExport(f FindingFilter, ids []int64) ([]*DBFinding, var aidsJSON string if err := rows.Scan(&f.ID, &f.TaskID, &f.NodeID, &f.VulnClass, &f.Name, &f.Severity, &f.Summary, &f.Evidence, &f.Worker, &aidsJSON, &f.Status, &f.CreatedAt, - &f.TaskDescription, &f.Report); err != nil { + &f.TaskDescription, &f.EvidenceVersion, &f.ReportEvidenceVersion, &f.TrafficCount, &f.Report); err != nil { return nil, err } _ = json.Unmarshal([]byte(aidsJSON), &f.AssetIDs) @@ -631,7 +638,7 @@ func (d *DB) GetFinding(id int64) (*DBFinding, error) { WHERE f.id = $1`, id).Scan( &f.ID, &f.TaskID, &f.NodeID, &f.VulnClass, &f.Name, &f.Severity, &f.Summary, &f.Evidence, &f.Worker, &aidsJSON, &f.Status, &f.CreatedAt, - &f.TaskDescription, &f.Report) + &f.TaskDescription, &f.EvidenceVersion, &f.ReportEvidenceVersion, &f.TrafficCount, &f.Report) if err == sql.ErrNoRows { return nil, nil } @@ -647,20 +654,27 @@ func (d *DB) GetFinding(id int64) (*DBFinding, error) { // list, the per-task 发现 Tab, and the exploration graph alike. Deleting the node // cascades its edges + node_assets and nulls any activity referencing it. Returns // rows affected (0 = no finding with that id). -func (d *DB) DeleteFinding(id int64) (int64, error) { - var nodeID *int64 - err := d.QueryRow(`DELETE FROM findings WHERE id=$1 RETURNING node_id`, id).Scan(&nodeID) - if err == sql.ErrNoRows { - return 0, nil - } - if err != nil { - return 0, err - } - if nodeID != nil { - // best-effort: the finding is already gone; a stray node must not fail the op. - _, _ = d.Exec(`DELETE FROM exploration_nodes WHERE id=$1 AND kind='finding'`, *nodeID) - } - return 1, nil +func (d *DB) DeleteFinding(id int64) (n int64, err error) { + err = d.WithEvidenceTx(context.Background(), func(tx *sql.Tx) error { + if err := LockFindingEvidenceTx(tx, id, nil); err != nil { + if errors.Is(err, ErrFindingNotFound) { + return nil + } + return err + } + var nodeID sql.NullInt64 + if err := tx.QueryRow(`DELETE FROM findings WHERE id=$1 RETURNING node_id`, id).Scan(&nodeID); err != nil { + return err + } + if nodeID.Valid { + if _, err := tx.Exec(`DELETE FROM exploration_nodes WHERE id=$1 AND kind='finding'`, nodeID.Int64); err != nil { + return err + } + } + n = 1 + return nil + }) + return } // DeleteFindingsByTask removes all findings rows of a task. The originating @@ -687,11 +701,7 @@ func (d *DB) SetFindingStatus(id int64, status string) (int64, error) { // whose node_id matches — report_finding returns that node id, so an agent tool // can address the finding it just created. Returns rows affected (0 when no row). func (d *DB) SetFindingReportByNodeID(nodeID int64, report string) (int64, error) { - res, err := d.Exec(`UPDATE findings SET report=$1 WHERE node_id=$2`, report, nodeID) - if err != nil { - return 0, err - } - return res.RowsAffected() + return d.SetFindingReportVersionByNodeID(context.Background(), nodeID, report, nil) } // setFindingCol updates one text column on the standalone finding row AND mirrors @@ -737,6 +747,8 @@ func (d *DB) SetFindingVulnClass(id int64, vulnclass string) (int64, error) { // FindingMeta is the standalone-row data (id, triage state, anchored assets) the // per-task view grafts onto its exploration-node findings. type FindingMeta struct { + TrafficCount int + ID int64 Status string AssetIDs []int64 @@ -757,7 +769,7 @@ func (d *DB) FindingMetaByNodeID(taskID int64) (map[int64]FindingMeta, error) { if taskID <= 0 { return out, nil } - rows, err := d.Query(`SELECT node_id, id, COALESCE(status,'pending'), asset_ids FROM findings + rows, err := d.Query(`SELECT node_id, id, COALESCE(status,'pending'), asset_ids, (SELECT count(*) FROM finding_traffic_bindings b WHERE b.finding_id=findings.id) FROM findings WHERE task_id=$1 AND node_id IS NOT NULL`, taskID) if err != nil { return out, err @@ -767,7 +779,7 @@ func (d *DB) FindingMetaByNodeID(taskID int64) (map[int64]FindingMeta, error) { var nid int64 var m FindingMeta var aidsJSON string - if err := rows.Scan(&nid, &m.ID, &m.Status, &aidsJSON); err != nil { + if err := rows.Scan(&nid, &m.ID, &m.Status, &aidsJSON, &m.TrafficCount); err != nil { return out, err } _ = json.Unmarshal([]byte(aidsJSON), &m.AssetIDs) diff --git a/db/schema.sql b/db/schema.sql index 4248569..15d076b 100644 --- a/db/schema.sql +++ b/db/schema.sql @@ -1058,6 +1058,39 @@ CREATE INDEX IF NOT EXISTS idx_findings_status ON findings(status, created_at DE -- 「按资产」视图靠 asset_ids @> '[]' 反查发现,没有这个 GIN 索引就是全表扫。 CREATE INDEX IF NOT EXISTS idx_findings_asset_ids ON findings USING GIN(asset_ids jsonb_path_ops); +ALTER TABLE findings ADD COLUMN IF NOT EXISTS evidence_version BIGINT NOT NULL DEFAULT 0; +ALTER TABLE findings ADD COLUMN IF NOT EXISTS report_evidence_version BIGINT NOT NULL DEFAULT 0; + +CREATE TABLE IF NOT EXISTS traffic_evidence_snapshots ( + id TEXT PRIMARY KEY, + source_traffic_id TEXT NOT NULL, + captured_at BIGINT NOT NULL, + url TEXT NOT NULL, + method TEXT NOT NULL, + status INTEGER NOT NULL, + content_type TEXT NOT NULL DEFAULT '', + req_head TEXT NOT NULL, + resp_head TEXT NOT NULL, + req_hash TEXT NOT NULL, + resp_hash TEXT NOT NULL, + req_len BIGINT NOT NULL, + resp_len BIGINT NOT NULL, + unreferenced_at TIMESTAMPTZ, + created_at TIMESTAMPTZ NOT NULL DEFAULT now() +); +CREATE TABLE IF NOT EXISTS finding_traffic_bindings ( + id BIGSERIAL PRIMARY KEY, + finding_id BIGINT NOT NULL REFERENCES findings(id) ON DELETE CASCADE, + snapshot_id TEXT NOT NULL REFERENCES traffic_evidence_snapshots(id), + role TEXT NOT NULL DEFAULT 'supporting', + note TEXT NOT NULL DEFAULT '', + position INTEGER NOT NULL, + created_at TIMESTAMPTZ NOT NULL DEFAULT now(), + UNIQUE(finding_id, snapshot_id) +); +CREATE INDEX IF NOT EXISTS idx_finding_traffic_order ON finding_traffic_bindings(finding_id, position, id); +CREATE INDEX IF NOT EXISTS idx_finding_traffic_snapshot ON finding_traffic_bindings(snapshot_id); + -- ===================================================================== -- M. 后端日志持久化 -- ===================================================================== diff --git a/db/task_archives.go b/db/task_archives.go index e38a4cd..b7a1d3c 100644 --- a/db/task_archives.go +++ b/db/task_archives.go @@ -16,7 +16,7 @@ import ( ) const ( - TaskArchiveFormatVersion = 2 + TaskArchiveFormatVersion = 3 TaskArchiveLegacyFormatVersion = 1 TaskArchiveLLMRecordsPath = "database/llm_records.ndjson" ) @@ -587,6 +587,8 @@ func (d *DB) snapshotTaskArchive(taskID int64, llmRecords io.Writer) (*TaskArchi {"task_llm_profiles", `SELECT * FROM task_llm_profiles WHERE task_id=$1 ORDER BY position`, []any{taskID}}, {"task_scope", `SELECT * FROM task_scope WHERE task_id=$1 ORDER BY id`, []any{taskID}}, {"findings", `SELECT * FROM findings WHERE task_id=$1 ORDER BY id`, []any{taskID}}, + {"finding_traffic_bindings", `SELECT b.* FROM finding_traffic_bindings b JOIN findings f ON f.id=b.finding_id WHERE f.task_id=$1 ORDER BY b.finding_id,b.position,b.id`, []any{taskID}}, + {"traffic_evidence_snapshots", `SELECT s.* FROM traffic_evidence_snapshots s WHERE EXISTS(SELECT 1 FROM finding_traffic_bindings b JOIN findings f ON f.id=b.finding_id WHERE b.snapshot_id=s.id AND f.task_id=$1) ORDER BY s.id`, []any{taskID}}, {"llm_records", `SELECT * FROM llm_records WHERE COALESCE(task_id,'')=$1 ORDER BY id`, []any{strconv.FormatInt(taskID, 10)}}, {"llm_usage", `SELECT * FROM llm_usage WHERE COALESCE(task_id,'')=$1 OR exploration_id=$2 ORDER BY id`, []any{strconv.FormatInt(taskID, 10), expID}}, {"skill_usage", `SELECT * FROM skill_usage WHERE task_id=$1 OR exploration_id=$2 ORDER BY id`, []any{taskID, expID}}, diff --git a/db/task_archives_restore.go b/db/task_archives_restore.go index 71e916a..014e7ea 100644 --- a/db/task_archives_restore.go +++ b/db/task_archives_restore.go @@ -233,6 +233,10 @@ WHERE archive.id=$1 FOR UPDATE OF archive,task`, archiveID).Scan(&taskID, &expID if err != nil { return nil, err } + remappedTables["findings"], err = normalizeArchivedFindingVersions(remappedTables["findings"]) + if err != nil { + return nil, err + } // Insert graph rows in foreign-key order. The archived stub has no graph rows, // so an ID conflict signals external corruption and must stop the restore. for _, table := range []string{"exploration_nodes", "exploration_edges", "exploration_anchors", "task_constraints", "activity"} { @@ -260,6 +264,9 @@ WHERE archive.id=$1 FOR UPDATE OF archive,task`, archiveID).Scan(&taskID, &expID return nil, fmt.Errorf("restore %s: %w", table, err) } } + if err := restoreFindingTrafficTx(tx, snapshot); err != nil { + return nil, fmt.Errorf("restore finding traffic: %w", err) + } if streamedLLMRecords != "" { count, err := insertArchiveJSONSequenceRows(tx, "llm_records", llmRecords) if err != nil { diff --git a/docs/漏洞流量证据.md b/docs/漏洞流量证据.md new file mode 100644 index 0000000..bcbf204 --- /dev/null +++ b/docs/漏洞流量证据.md @@ -0,0 +1,91 @@ +# 漏洞多流量证据 + +漏洞详情的「关联流量」支持跨页多选、用途和说明、排序及解除绑定。流量页也可多选记录，一次关联到一个已有漏洞。任务继承的漏洞证据只读，修改需进入来源任务。 + +Agent 在发现漏洞、调用 `report_finding` 时可一并绑定多条流量，报告 Agent 随后读取这些快照撰写报告。绑定可选：TCP 等非 HTTP 漏洞、未采集或找不到确切记录时，省略 `traffic_refs` 或传空数组仍可正常上报；在 `evidence` 保留命令输出、日志等其他可验证证据，建议说明未绑定原因，不新增必填字段。提交的 ID 必须全部有效且正文完整。任何一条失败都会回滚本次业务记录。重复追加同一快照不增加绑定，也不覆盖说明。 + +## Agent 编号与报告版本 + +`report_finding` 新增可选参数，数组顺序即初始证据顺序： + +```json +{ + "traffic_refs": [ + {"traffic_id": "真实流量ID", "role": "baseline", "note": "正常账户请求"}, + {"traffic_id": "另一个真实流量ID", "role": "proof", "note": "复现请求"} + ] +} +``` + +用途为 `baseline`（正常对照）、`proof`（漏洞证明）、`verification`（补充验证）、`supporting`（辅助证据，默认）。先用 `traffic_search` / `traffic_get` 核对真实记录；域名和时间仅用于候选筛选，不推定任务归属。 + +提示词参考 [CyberStrikeAI 的漏洞上报工具指引](https://github.com/RuoJi6/CyberStrikeAI/blob/54d56774b8bd285817d16d48b70a4a5e6e0963f7/internal/app/vulnerability_tools.go)，结合本项目的可选绑定约定，不引入无包原因必填校验。不得猜测 ID，也不应仅为补包重复探测。 + +返回第一行仍为 `finding recorded: <探索节点 ID>`；随后 JSON 提供独立漏洞记录 `finding_id`、探索节点 `finding_node_id` 和绑定摘要。 + +- `get_finding_traffic(finding_id)` 使用**独立漏洞记录 ID**，返回有序清单及 `version`；传 `binding_id`、`side=request|response`、`offset`、`length` 可分段读取，每段最多 8192 字节。 +- `update_finding_report` 的 `finding_id` **继续使用探索节点 ID**。新增 `evidence_version` 填实际读取的版本；生成期间证据变化会拒绝旧版本写入，必须重新读取并生成。 +- 旧报告调用未传版本时，不宣称覆盖已有流量证据。绑定、说明、用途或排序变化后，已有报告提示待更新。 + +启动时仅为旧工具 schema 增加可选属性，并首次给 reporter 绑定只读工具；保留已有提示词、描述、启用状态及其他配置。Worker 运行时附加可选绑定的工具使用指引，不替换用户编辑的提示词。失败的 `report_finding` 不触发报告 Agent。 + +## API + +基础路径：`/api/exploration/findings/{finding_id}/traffic`，使用独立漏洞 ID。沿用认证；`context_task` 校验任务可见性及继承只读。 + +| 方法 / 相对路径 | 请求 / 返回 | +| --- | --- | +| `GET` | 有序摘要、证据版本、报告采用版本 | +| `POST` | `{"traffic_refs":[...]}` 整批追加 | +| `PATCH /{binding_id}` | `{"version":1,"role":"proof","note":"说明"}` | +| `DELETE /{binding_id}` | `{"version":1}` | +| `PUT /order` | `{"version":1,"binding_ids":["2","1"]}`，必须是完整列表 | +| `GET /{binding_id}` | 快照元数据和有界正文预览 | +| `GET /{binding_id}/body` | `side`、`offset`、`length`；`download=1` 下载完整原字节 | + +版本/排序集合冲突、归档中写入返回 `409`；继承写入 `403`；不存在或不属于漏洞的绑定 `404`；流量/附件读取及校验失败明确返回错误。 + +## 存储与迁移 + +启动幂等迁移 PostgreSQL：新增 `traffic_evidence_snapshots`、`finding_traffic_bindings`，以及 `findings.evidence_version` / `report_evidence_version`（默认 0）。不根据历史文字猜测补绑。 + +快照保存原流量 ID、采集时间、URL、方法、状态、请求/响应头、正文长度和 SHA-256。正文按哈希存放在 `/evidence/blobs/<前两位>/.bin`，与可清理的 `data/traffic` 独立，多个漏洞可共享快照/正文。快照不提供内容更新接口，校验不一致时读取和导出失败。 + +在原流量写锁下读取完整正文，包括大正文 blob 和旧目录记录。先持久化并校验文件，再用单个 PostgreSQL 事务写入探索节点、意图关系、漏洞、快照和绑定；提交后才通知规划者。失败可能遗留无引用文件，但不产生部分业务记录。 + +PostgreSQL advisory lock `7337741004` 协调证据文件和 SQL 引用；任务行锁禁止归档排队后的证据修改。恢复从正文安装到元数据提交全程保持证据锁。删除漏洞级联移除绑定。 + +清理器每小时运行，只回收无引用且非进行中操作的内容，至少延迟 24 小时。普通流量清理不触及证据目录。备份热数据时同时备份 PostgreSQL 和 `data/evidence`。 + +## 导出与归档 + +Markdown 包含有序证据清单及版本；JSON 包含元数据；CSV 增加数量和绑定 ID。`md-zip` 保留漏洞 Markdown，并提供： + +```text +evidence/// + manifest.json + request.http + response.http + request.bin + response.bin +``` + +Markdown 通过相对链接引用报文。发送下载前完成附件复制、哈希校验、压缩、磁盘同步及全部 ZIP 条目的 CRC 读取校验；缺失/损坏使整个下载失败。完整附件保留二进制原字节。 + +归档 v3 按漏洞绑定关系收集快照和正文，不依赖原流量或域名。包校验完成后才清理热数据；共享证据继续保留。恢复先校验安装正文，再事务恢复元数据和绑定，支持失败重试。v1/v2 继续可恢复，缺失的新字段显式补为 0。 + +## 验证与边界 + +为每个测试包设置独立、新建的 PostgreSQL 测试库，通过 `ARTEX_PG_DSN` 指定，避免残留任务/模型夹具触发后台运行。运行相关包完整测试，并确认没有配置缺失导致的跳过： + +```sh +# 每个包运行前将 ARTEX_PG_DSN 设为对应的独立测试库；显式配置失败必须报错。 +go test ./ -count=1 +go test -race -p 1 ./evidence ./db ./agent ./server -run 'TestEvidence|TestFindingTraffic|TestFindingEvidence|TestReportFindingAtomicContract|TestTaskArchive' +``` + +前端验证包含 `npx tsc --noEmit`、受影响文件的 Biome 检查、Webpack 构建和 `NEXT_EXPORT=1` 静态导出。使用独立缓存目录，避免覆盖运行中的开发服务。 + +本机端到端验收使用独立端口、受控 HTTP / 域名 HTTPS 目标和临时数据目录，覆盖两个绑定入口、跨页选择、排序/说明、错误提示、继承只读、下载，以及删除原流量后导出、归档、回收热正文、恢复并校验哈希。 + +首版采用全局证据协调锁；大批量绑定/导出或慢速附件下载期间，其他证据操作可能等待。缺少录制记录或完整正文时不能补造证据。本功能不改变捕获开关，也不处理 HTTPS 直接访问 IP 的证书问题。 diff --git a/evidence/store.go b/evidence/store.go new file mode 100644 index 0000000..dd9aa06 --- /dev/null +++ b/evidence/store.go @@ -0,0 +1,401 @@ +// Package evidence preserves finding evidence independently of disposable traffic. +package evidence + +import ( + "context" + "crypto/sha256" + "database/sql" + "encoding/hex" + "errors" + "fmt" + "io" + "io/fs" + "log" + "os" + "path/filepath" + "strings" + "time" + + "github.com/Autumn-27/artex/db" + "github.com/Autumn-27/artex/traffic" +) + +type Store struct { + DB *db.DB + Traffic *traffic.Traffic + Dir string +} + +func New(pg *db.DB, tr *traffic.Traffic, dir string) *Store { + return &Store{DB: pg, Traffic: tr, Dir: dir} +} + +func hashPath(dir, hash string) (string, error) { + if len(hash) != 64 { + return "", errors.New("invalid evidence hash") + } + if _, err := hex.DecodeString(hash); err != nil || strings.ToLower(hash) != hash { + return "", errors.New("invalid evidence hash") + } + return filepath.Join(dir, "blobs", hash[:2], hash+".bin"), nil +} + +func verifyFile(path, hash string, length int64) error { + f, err := os.Open(path) + if err != nil { + return err + } + defer f.Close() + h := sha256.New() + n, err := io.Copy(h, f) + if err != nil { + return err + } + if n != length || hex.EncodeToString(h.Sum(nil)) != hash { + return fmt.Errorf("证据正文校验失败: %s", hash) + } + return nil +} + +// A new body becomes visible only after a durable write. Failed SQL commits may +// leave unreferenced files; GC reaps those after a full day's grace period. +func (s *Store) writeBody(r io.Reader, expectedLength int64, expectedHash string) (hash string, err error) { + stage := filepath.Join(s.Dir, ".staging") + if err = os.MkdirAll(stage, 0o700); err != nil { + return + } + f, err := os.CreateTemp(stage, "body-") + if err != nil { + return "", err + } + defer func() { f.Close(); os.Remove(f.Name()) }() + h := sha256.New() + n, err := io.Copy(io.MultiWriter(f, h), r) + if err != nil { + return "", err + } + hash = hex.EncodeToString(h.Sum(nil)) + if n != expectedLength { + return "", fmt.Errorf("正文不完整: 预期 %d 字节，读取 %d 字节", expectedLength, n) + } + if expectedHash != "" && expectedHash != hash { + return "", errors.New("原始流量正文哈希不匹配") + } + if err = f.Sync(); err != nil { + return "", err + } + if err = f.Close(); err != nil { + return "", err + } + path, err := hashPath(s.Dir, hash) + if err != nil { + return "", err + } + if _, err = os.Stat(path); err == nil { + if err = verifyFile(path, hash, n); err != nil { + return "", err + } + // Refresh the grace period for a restored but not-yet-committed body. + now := time.Now() + return hash, os.Chtimes(path, now, now) + } else if !os.IsNotExist(err) { + return "", err + } + if err = os.MkdirAll(filepath.Dir(path), 0o700); err != nil { + return "", err + } + if err = os.Rename(f.Name(), path); err != nil { + return "", err + } + d, err := os.Open(filepath.Dir(path)) + if err != nil { + return "", err + } + defer d.Close() + return hash, d.Sync() +} + +// StageFindingsExport freezes bindings/report versions and makes private body +// copies before the HTTP response is started. The caller owns and removes dest. +func (s *Store) StageFindingsExport(ctx context.Context, findings []*db.DBFinding, dest string, copyBodies bool) error { + return s.DB.WithEvidenceTx(ctx, func(tx *sql.Tx) error { + var snapshots []db.TrafficEvidenceSnapshot + for _, f := range findings { + list, err := db.FindingTrafficTx(tx, f.ID) + if err != nil { + return err + } + f.TrafficBindings = list.Bindings + f.TrafficCount = len(list.Bindings) + f.EvidenceVersion = list.Version + f.ReportEvidenceVersion = list.ReportVersion + if err = tx.QueryRow(`SELECT report FROM findings WHERE id=$1`, f.ID).Scan(&f.Report); err != nil { + return err + } + for _, b := range list.Bindings { + snapshots = append(snapshots, b.Snapshot) + } + } + if copyBodies { + return s.copySnapshots(snapshots, dest) + } + return nil + }) +} + +func (s *Store) prepare(ctx context.Context, refs []db.TrafficRef) ([]db.PreparedTrafficEvidence, error) { + refs, err := db.NormalizeTrafficRefs(refs) + if err != nil { + return nil, err + } + out := make([]db.PreparedTrafficEvidence, 0, len(refs)) + if len(refs) == 0 { + return out, nil + } + ids := make([]string, len(refs)) + byID := map[string]db.TrafficRef{} + for i, ref := range refs { + ids[i] = ref.TrafficID + byID[ref.TrafficID] = ref + } + err = s.Traffic.ReadEvidence(ctx, ids, func(e traffic.EvidenceExchange) error { + rh, err := s.writeBody(e.Request, e.ReqLen, e.ReqHash) + if err != nil { + return err + } + ph, err := s.writeBody(e.Response, e.RespLen, e.RespHash) + if err != nil { + return err + } + v := db.TrafficEvidenceSnapshot{SourceTrafficID: e.ID, CapturedAt: e.TS, URL: e.URL, Method: e.Method, Status: e.Status, ContentType: e.ContentType, + ReqHead: e.ReqHead, RespHead: e.RespHead, ReqHash: rh, RespHash: ph, ReqLen: e.ReqLen, RespLen: e.RespLen} + v.ID = db.TrafficSnapshotID(v) + out = append(out, db.PreparedTrafficEvidence{Ref: byID[e.ID], Snapshot: v}) + return nil + }) + return out, err +} + +func (s *Store) Record(ctx context.Context, in db.RecordFindingInput, refs []db.TrafficRef) (out *db.RecordedFinding, err error) { + err = s.DB.WithEvidenceTx(ctx, func(tx *sql.Tx) error { + if err := db.LockTaskEvidenceTx(tx, in.TaskID); err != nil { + return err + } + prepared, err := s.prepare(ctx, refs) + if err != nil { + return err + } + out, err = db.RecordFindingTx(tx, in, prepared) + return err + }) + return +} + +func (s *Store) Bind(ctx context.Context, findingID int64, refs []db.TrafficRef) (out *db.FindingTraffic, err error) { + err = s.DB.WithEvidenceTx(ctx, func(tx *sql.Tx) error { + if err := db.LockFindingEvidenceTx(tx, findingID, nil); err != nil { + return err + } + prepared, err := s.prepare(ctx, refs) + if err != nil { + return err + } + if err = db.AddFindingTrafficTx(tx, findingID, prepared); err != nil { + return err + } + out, err = db.FindingTrafficTx(tx, findingID) + return err + }) + return +} + +func (s *Store) WithBinding(ctx context.Context, findingID, bindingID int64, fn func(db.FindingTrafficBinding) error) error { + return s.DB.WithEvidenceTx(ctx, func(tx *sql.Tx) error { + list, err := db.FindingTrafficTx(tx, findingID) + if err != nil { + return err + } + for _, b := range list.Bindings { + if b.ID == bindingID { + return fn(b) + } + } + return db.ErrEvidenceNotFound + }) +} + +// OpenBody must be called while holding WithEvidenceTx/WithBinding, and the +// returned file closed before that callback returns. +func (s *Store) OpenBody(snapshot db.TrafficEvidenceSnapshot, side string) (*os.File, int64, error) { + hash, length := snapshot.ReqHash, snapshot.ReqLen + if side == "response" { + hash, length = snapshot.RespHash, snapshot.RespLen + } else if side != "request" { + return nil, 0, errors.New("side 必须为 request 或 response") + } + path, err := hashPath(s.Dir, hash) + if err != nil { + return nil, 0, err + } + if err = verifyFile(path, hash, length); err != nil { + return nil, 0, err + } + f, err := os.Open(path) + return f, length, err +} + +// CopySnapshots is used by both report downloads and portable task archives. +// The destination owns real copies, never links into either disposable store. +func (s *Store) CopySnapshots(ctx context.Context, snapshots []db.TrafficEvidenceSnapshot, dest string) error { + return s.DB.WithEvidenceTx(ctx, func(*sql.Tx) error { return s.copySnapshots(snapshots, dest) }) +} + +func (s *Store) copySnapshots(snapshots []db.TrafficEvidenceSnapshot, dest string) error { + for _, v := range snapshots { + for _, side := range []string{"request", "response"} { + if err := func() error { + f, length, err := s.OpenBody(v, side) + if err != nil { + return err + } + defer f.Close() + hash := v.ReqHash + if side == "response" { + hash = v.RespHash + } + target, err := hashPath(dest, hash) + if err != nil { + return err + } + if err = os.MkdirAll(filepath.Dir(target), 0o700); err != nil { + return err + } + out, err := os.OpenFile(target, os.O_CREATE|os.O_WRONLY|os.O_EXCL, 0o600) + if os.IsExist(err) { + return verifyFile(target, hash, length) + } + if err != nil { + return err + } + _, copyErr := io.Copy(out, f) + syncErr := out.Sync() + closeErr := out.Close() + return errors.Join(copyErr, syncErr, closeErr) + }(); err != nil { + return err + } + } + } + return nil +} + +func (s *Store) InstallSnapshots(ctx context.Context, snapshots []db.TrafficEvidenceSnapshot, source string) error { + return s.WithInstalledSnapshots(ctx, snapshots, source, func() error { return nil }) +} + +// WithInstalledSnapshots pins installed bodies until the metadata restore finishes. +// The callback must not acquire another evidence advisory lock. +func (s *Store) WithInstalledSnapshots(ctx context.Context, snapshots []db.TrafficEvidenceSnapshot, source string, restore func() error) error { + return s.DB.WithEvidenceTx(ctx, func(*sql.Tx) error { + for _, v := range snapshots { + if v.ID != db.TrafficSnapshotID(v) { + return errors.New("归档证据快照元数据哈希不匹配") + } + for _, body := range []struct { + hash string + length int64 + }{{v.ReqHash, v.ReqLen}, {v.RespHash, v.RespLen}} { + if err := func() error { + path, err := hashPath(source, body.hash) + if err != nil { + return err + } + f, err := os.Open(path) + if err != nil { + return err + } + defer f.Close() + _, err = s.writeBody(f, body.length, body.hash) + return err + }(); err != nil { + return err + } + } + } + return restore() + }) +} + +func (s *Store) Collect(ctx context.Context, now time.Time) error { + return s.DB.WithEvidenceTx(ctx, func(tx *sql.Tx) error { + if _, err := tx.Exec(`UPDATE traffic_evidence_snapshots s SET unreferenced_at=$1 +WHERE unreferenced_at IS NULL AND NOT EXISTS(SELECT 1 FROM finding_traffic_bindings b WHERE b.snapshot_id=s.id)`, now); err != nil { + return err + } + if _, err := tx.Exec(`DELETE FROM traffic_evidence_snapshots s WHERE unreferenced_at<$1 +AND NOT EXISTS(SELECT 1 FROM finding_traffic_bindings b WHERE b.snapshot_id=s.id)`, now.Add(-24*time.Hour)); err != nil { + return err + } + rows, err := tx.Query(`SELECT req_hash FROM traffic_evidence_snapshots UNION SELECT resp_hash FROM traffic_evidence_snapshots`) + if err != nil { + return err + } + refs := map[string]bool{} + for rows.Next() { + var hash string + if err := rows.Scan(&hash); err != nil { + rows.Close() + return err + } + refs[hash] = true + } + err = rows.Err() + rows.Close() + if err != nil { + return err + } + for _, root := range []string{filepath.Join(s.Dir, "blobs"), filepath.Join(s.Dir, ".staging")} { + err := filepath.WalkDir(root, func(path string, entry fs.DirEntry, err error) error { + if os.IsNotExist(err) { + return nil + } + if err != nil { + return err + } + if entry.IsDir() { + return nil + } + if refs[strings.TrimSuffix(entry.Name(), ".bin")] { + return nil + } + info, err := entry.Info() + if err != nil { + return err + } + if now.Sub(info.ModTime()) < 24*time.Hour { + return nil + } + return os.Remove(path) + }) + if err != nil { + return err + } + } + return nil + }) +} + +func (s *Store) RunGC(ctx context.Context) { + timer := time.NewTicker(time.Hour) + defer timer.Stop() + for { + select { + case <-ctx.Done(): + return + case now := <-timer.C: + if err := s.Collect(ctx, now); err != nil && !errors.Is(err, context.Canceled) { + log.Printf("[evidence] cleanup failed: %v", err) + } + } + } +} diff --git a/evidence/store_test.go b/evidence/store_test.go new file mode 100644 index 0000000..8d90847 --- /dev/null +++ b/evidence/store_test.go @@ -0,0 +1,333 @@ +package evidence + +import ( + "bytes" + "context" + "crypto/sha256" + "encoding/hex" + "errors" + "fmt" + "io" + "os" + "path/filepath" + "sync" + "testing" + "time" + + "github.com/Autumn-27/artex/db" + "github.com/Autumn-27/artex/traffic" +) + +func evidenceFixture(t *testing.T) (*Store, db.RecordFindingInput, string) { + t.Helper() + dsn := os.Getenv("ARTEX_PG_DSN") + if dsn == "" { + t.Skip("ARTEX_PG_DSN is required for evidence integration tests") + } + pg, err := db.Open(dsn) + if err != nil { + t.Fatal(err) + } + t.Cleanup(func() { pg.Close() }) + task, err := pg.CreateTask("evidence integration "+t.Name(), "local fixtures", nil, 0, 0) + if err != nil { + t.Fatal(err) + } + t.Cleanup(func() { pg.DeleteFindingsByTask(task.ID); pg.DeleteTask(task.ID) }) + dir := t.TempDir() + tr, err := traffic.Open(filepath.Join(dir, "traffic"), ":0") + if err != nil { + t.Fatal(err) + } + t.Cleanup(func() { tr.Close() }) + store := New(pg, tr, filepath.Join(dir, "evidence")) + in := db.RecordFindingInput{TaskID: task.ID, ExplorationID: task.ExplorationID, Worker: "test", VulnClass: "TEST", Name: "Evidence fixture", Severity: "low", Summary: "local test"} + return store, in, dir +} + +func seedExchange(t *testing.T, s *Store, id string, body []byte, spill bool) { + t.Helper() + var blob any + inline := body + if spill { + sum := sha256.Sum256(body) + hash := hex.EncodeToString(sum[:]) + blob = hash + path := filepath.Join(filepath.Dir(s.Dir), "traffic", "_blobs", "sha256", hash[:2], hash+".bin") + if err := os.MkdirAll(filepath.Dir(path), 0o700); err != nil { + t.Fatal(err) + } + if err := os.WriteFile(path, body, 0o600); err != nil { + t.Fatal(err) + } + inline = []byte("TRUNCATED PREVIEW") + } + _, err := s.Traffic.DB().Exec(`INSERT INTO exchanges(id,ts,host,method,url_template,url,status,content_type,req_len,resp_len,path) VALUES(?,?,'fixture.local','POST','/test',?,200,'application/octet-stream',2,?,'')`, id, time.Now().Unix(), "https://fixture.local/"+id, len(body)) + if err != nil { + t.Fatal(err) + } + _, err = s.Traffic.DB().Exec(`INSERT INTO exchange_bodies(id,req_head,req_body,resp_head,resp_body,resp_blob) VALUES(?,?,?, ?,?,?)`, id, "POST /test HTTP/1.1\nHost: fixture.local\n", []byte("{}"), "HTTP 200\nContent-Type: application/octet-stream\n", inline, blob) + if err != nil { + t.Fatal(err) + } +} + +func TestEvidenceBindingLifecycle(t *testing.T) { + s, in, _ := evidenceFixture(t) + ctx := context.Background() + large := bytes.Repeat([]byte{0, 1, 2, 255, 'a', 'b'}, 180000) + for i := 0; i < 3; i++ { + seedExchange(t, s, fmt.Sprint(i), large, i == 2) + } + refs := []db.TrafficRef{{TrafficID: "0", Role: "baseline", Note: "normal"}, {TrafficID: "1", Role: "proof", Note: "proof"}, {TrafficID: "2", Role: "verification", Note: "binary"}} + r, err := s.Record(ctx, in, refs) + if err != nil { + t.Fatal(err) + } + if len(r.Traffic.Bindings) != 3 || r.Traffic.Version != 1 { + t.Fatalf("record=%+v", r) + } + again, err := s.Bind(ctx, r.FindingID, []db.TrafficRef{{TrafficID: "1", Note: "must not overwrite"}}) + if err != nil { + t.Fatal(err) + } + if len(again.Bindings) != 3 || again.Version != 1 || again.Bindings[1].Note != "proof" { + t.Fatalf("duplicate changed bindings: %+v", again) + } + other, err := s.Record(ctx, in, []db.TrafficRef{refs[2]}) + if err != nil { + t.Fatal(err) + } + if other.Traffic.Bindings[0].SnapshotID != r.Traffic.Bindings[2].SnapshotID { + t.Fatal("snapshot was not shared") + } + if _, err = s.Traffic.DeleteHost("fixture.local"); err != nil { + t.Fatal(err) + } + for _, binding := range r.Traffic.Bindings { + if err = s.WithBinding(ctx, r.FindingID, binding.ID, func(b db.FindingTrafficBinding) error { + f, _, err := s.OpenBody(b.Snapshot, "response") + if err != nil { + return err + } + defer f.Close() + got, err := io.ReadAll(f) + if !bytes.Equal(got, large) { + t.Fatal("body was truncated/changed") + } + return err + }); err != nil { + t.Fatal(err) + } + } + v := again.Version + if _, err = s.DB.SetFindingReportVersionByNodeID(ctx, r.NodeID, "report", &v); err != nil { + t.Fatal(err) + } + note := "updated" + if err = s.DB.EditFindingTraffic(ctx, r.FindingID, again.Bindings[0].ID, v, nil, ¬e, false, nil); err != nil { + t.Fatal(err) + } + if _, err = s.DB.SetFindingReportVersionByNodeID(ctx, r.NodeID, "stale report", &v); !errors.Is(err, db.ErrEvidenceConflict) { + t.Fatalf("stale report: %v", err) + } + f, err := s.DB.GetFinding(r.FindingID) + if err != nil { + t.Fatal(err) + } + if f.Report != "report" || f.EvidenceVersion == f.ReportEvidenceVersion || f.TrafficCount != 3 { + t.Fatalf("finding versions: %+v", f) + } + if err = s.DB.EditFindingTraffic(ctx, r.FindingID, again.Bindings[0].ID, v, nil, nil, true, nil); !errors.Is(err, db.ErrEvidenceConflict) { + t.Fatalf("stale delete: %v", err) + } + if _, err = s.DB.SetFindingReportByNodeID(r.NodeID, "legacy report"); err != nil { + t.Fatal(err) + } + f, err = s.DB.GetFinding(r.FindingID) + if err != nil { + t.Fatal(err) + } + if f.ReportEvidenceVersion != -1 { + t.Fatal("legacy write claimed current evidence") + } + if _, err = s.DB.DeleteFinding(r.FindingID); err != nil { + t.Fatal(err) + } + if err = s.Collect(ctx, time.Now().Add(48*time.Hour)); err != nil { + t.Fatal(err) + } + if err = s.WithBinding(ctx, other.FindingID, other.Traffic.Bindings[0].ID, func(b db.FindingTrafficBinding) error { + f, _, err := s.OpenBody(b.Snapshot, "response") + if f != nil { + f.Close() + } + return err + }); err != nil { + t.Fatal("GC removed shared body:", err) + } +} + +func TestEvidenceAtomicFailuresAndLegacy(t *testing.T) { + s, in, dir := evidenceFixture(t) + ctx := context.Background() + seedExchange(t, s, "valid", []byte("OK"), false) + counts := func() (findings, nodes int) { + s.DB.QueryRow(`SELECT count(*) FROM findings WHERE task_id=$1`, in.TaskID).Scan(&findings) + s.DB.QueryRow(`SELECT count(*) FROM exploration_nodes WHERE exploration_id=$1 AND kind='finding'`, in.ExplorationID).Scan(&nodes) + return + } + for _, kind := range []string{"missing_id", "missing_body", "disk_failure", "database_failure"} { + t.Run(kind, func(t *testing.T) { + input := in + refs := []db.TrafficRef{{TrafficID: "valid"}, {TrafficID: "absent"}} + local := *s + switch kind { + case "missing_body": + seedExchange(t, s, "broken", []byte("more than preview"), false) + s.Traffic.DB().Exec(`UPDATE exchange_bodies SET resp_body=? WHERE id='broken'`, []byte("x")) + refs = []db.TrafficRef{{TrafficID: "valid"}, {TrafficID: "broken"}} + case "disk_failure": + local.Dir = filepath.Join(dir, "file-not-dir") + if err := os.WriteFile(local.Dir, []byte("x"), 0o600); err != nil { + t.Fatal(err) + } + refs = refs[:1] + case "database_failure": + input.AssetIDs = []int64{9223372036854775807} + refs = refs[:1] + } + if _, err := local.Record(ctx, input, refs); err == nil { + t.Fatal("expected failure") + } + if f, n := counts(); f != 0 || n != 0 { + t.Fatalf("partial record: findings=%d nodes=%d", f, n) + } + }) + } + legacy := filepath.Join(dir, "traffic", "legacy") + os.MkdirAll(legacy, 0o700) + os.WriteFile(filepath.Join(legacy, "request.http"), []byte("GET / HTTP/1.1\nHost: old.local\n"), 0o600) + os.WriteFile(filepath.Join(legacy, "response.http"), []byte("HTTP 200\n\nlegacy body"), 0o600) + _, err := s.Traffic.DB().Exec(`INSERT INTO exchanges(id,ts,host,method,url_template,url,status,content_type,req_len,resp_len,path) VALUES('old',1,'old.local','GET','/','http://old.local/',200,'text/plain',0,11,'legacy')`) + if err != nil { + t.Fatal(err) + } + if _, err = s.Record(ctx, in, []db.TrafficRef{{TrafficID: "old"}}); err != nil { + t.Fatal(err) + } + if _, err = s.Record(ctx, in, nil); err != nil { + t.Fatal("legacy no-traffic report failed:", err) + } +} + +func TestEvidenceConcurrentBindAndDelete(t *testing.T) { + s, in, _ := evidenceFixture(t) + ctx := context.Background() + seedExchange(t, s, "race", bytes.Repeat([]byte("large"), 90000), true) + r, err := s.Record(ctx, in, nil) + if err != nil { + t.Fatal(err) + } + var wg sync.WaitGroup + errs := make(chan error, 10) + for i := 0; i < 8; i++ { + wg.Add(1) + go func() { + defer wg.Done() + _, err := s.Bind(ctx, r.FindingID, []db.TrafficRef{{TrafficID: "race"}}) + errs <- err + }() + } + wg.Wait() + close(errs) + for err := range errs { + if err != nil { + t.Fatal(err) + } + } + list, err := s.DB.GetFindingTraffic(ctx, r.FindingID) + if err != nil { + t.Fatal(err) + } + if len(list.Bindings) != 1 || list.Version != 1 { + t.Fatalf("concurrent duplicates: %+v", list) + } + // Race a new capture copy with source deletion: either the full copy commits + // or nothing does. An already-bound snapshot remains readable in both cases. + wg.Add(2) + go func() { defer wg.Done(); s.Bind(ctx, r.FindingID, []db.TrafficRef{{TrafficID: "race"}}) }() + go func() { defer wg.Done(); s.Traffic.DeleteHost("fixture.local") }() + wg.Wait() + if err = s.WithBinding(ctx, r.FindingID, list.Bindings[0].ID, func(b db.FindingTrafficBinding) error { + f, _, err := s.OpenBody(b.Snapshot, "response") + if f != nil { + f.Close() + } + return err + }); err != nil { + t.Fatal(err) + } +} + +func TestEvidenceGCGraceAndActiveRestore(t *testing.T) { + s, in, _ := evidenceFixture(t) + ctx := context.Background() + seedExchange(t, s, "gc", []byte("unique unreferenced gc body"), false) + f, err := s.Record(ctx, in, []db.TrafficRef{{TrafficID: "gc"}}) + if err != nil { + t.Fatal(err) + } + snap := f.Traffic.Bindings[0].Snapshot + archive := t.TempDir() + if err = s.CopySnapshots(ctx, []db.TrafficEvidenceSnapshot{snap}, archive); err != nil { + t.Fatal(err) + } + if _, err = s.DB.DeleteFinding(f.FindingID); err != nil { + t.Fatal(err) + } + now := time.Now() + if err = s.Collect(ctx, now); err != nil { + t.Fatal(err) + } + path, err := hashPath(s.Dir, snap.RespHash) + if err != nil { + t.Fatal(err) + } + if err = s.Collect(ctx, now.Add(23*time.Hour)); err != nil { + t.Fatal(err) + } + if _, err = os.Stat(path); err != nil { + t.Fatal("GC grace ignored", err) + } + if err = s.Collect(ctx, now.Add(25*time.Hour)); err != nil { + t.Fatal(err) + } + if _, err = os.Stat(path); !os.IsNotExist(err) { + t.Fatal("orphan was not collected", err) + } + entered, release, finished := make(chan struct{}), make(chan struct{}), make(chan error, 1) + go func() { + finished <- s.WithInstalledSnapshots(ctx, []db.TrafficEvidenceSnapshot{snap}, archive, func() error { close(entered); <-release; return nil }) + }() + <-entered + short, cancel := context.WithTimeout(ctx, 100*time.Millisecond) + err = s.Collect(short, now.Add(72*time.Hour)) + cancel() + close(release) + if err == nil { + t.Fatal("GC entered during active restore") + } + if err = <-finished; err != nil { + t.Fatal(err) + } + if _, err = os.Stat(path); err != nil { + t.Fatal("active restore lost body", err) + } + // A portable package cannot substitute metadata while retaining its old ID. + bad := snap + bad.URL = "http://tampered.local/" + if err = s.InstallSnapshots(ctx, []db.TrafficEvidenceSnapshot{bad}, archive); err == nil { + t.Fatal("accepted tampered snapshot") + } +} diff --git a/evidence/testmain_test.go b/evidence/testmain_test.go new file mode 100644 index 0000000..9b3bf30 --- /dev/null +++ b/evidence/testmain_test.go @@ -0,0 +1,35 @@ +package evidence + +import ( + "context" + "os" + "testing" + + "github.com/Autumn-27/artex/db" +) + +// Initialize an explicitly configured fresh database before taking the same +// suite lock as db, agent and server. Hold it on one pinned connection. +func TestMain(m *testing.M) { + if os.Getenv("ARTEX_PG_DSN") == "" { + os.Exit(m.Run()) + } + os.Exit(runEvidenceSuite(m)) +} +func runEvidenceSuite(m *testing.M) int { + pg, err := db.Open(os.Getenv("ARTEX_PG_DSN")) + if err != nil { + panic(err) + } + defer pg.Close() + conn, err := pg.Conn(context.Background()) + if err != nil { + panic(err) + } + defer conn.Close() + if _, err = conn.ExecContext(context.Background(), `SELECT pg_advisory_lock(7337741002)`); err != nil { + panic(err) + } + defer conn.ExecContext(context.Background(), `SELECT pg_advisory_unlock(7337741002)`) + return m.Run() +} diff --git a/report/findings.go b/report/findings.go index 265bfd3..a91fa63 100644 --- a/report/findings.go +++ b/report/findings.go @@ -84,6 +84,7 @@ func FindingsMarkdown(fs []*db.DBFinding, generatedAt time.Time) string { b.WriteString(rep) b.WriteString("\n\n") } + b.WriteString(findingTrafficMarkdown(f, false)) b.WriteString("---\n\n") } return b.String() @@ -114,6 +115,7 @@ func SingleFindingMarkdown(f *db.DBFinding, generatedAt time.Time) string { b.WriteString(rep) b.WriteString("\n") } + b.WriteString(findingTrafficMarkdown(f, true)) return b.String() } @@ -147,7 +149,7 @@ func FindingsCSV(fs []*db.DBFinding) []byte { var buf bytes.Buffer buf.WriteString("\xEF\xBB\xBF") // UTF-8 BOM w := csv.NewWriter(&buf) - _ = w.Write([]string{"ID", "名称", "类别", "严重等级", "状态", "所属任务", "发现时间", "概述"}) + _ = w.Write([]string{"ID", "名称", "类别", "严重等级", "状态", "所属任务", "发现时间", "概述", "流量证据数量", "流量证据ID"}) for _, f := range items { _ = w.Write([]string{ fmt.Sprintf("%d", f.ID), @@ -158,8 +160,41 @@ func FindingsCSV(fs []*db.DBFinding) []byte { f.TaskDescription, f.CreatedAt.Format("2006-01-02 15:04:05"), strings.TrimSpace(f.Summary), + fmt.Sprint(len(f.TrafficBindings)), findingTrafficIDs(f), }) } w.Flush() return buf.Bytes() } + +func findingTrafficIDs(f *db.DBFinding) string { + ids := make([]string, 0, len(f.TrafficBindings)) + for _, b := range f.TrafficBindings { + ids = append(ids, fmt.Sprint(b.ID)) + } + return strings.Join(ids, ",") +} + +func findingTrafficMarkdown(f *db.DBFinding, attachments bool) string { + stale := f.Report != "" && f.EvidenceVersion != f.ReportEvidenceVersion + if len(f.TrafficBindings) == 0 && !stale { + return "" + } + var out strings.Builder + out.WriteString("\n## 关联流量证据\n\n") + fmt.Fprintf(&out, "证据版本：%d；绑定数量：%d。\n\n", f.EvidenceVersion, len(f.TrafficBindings)) + if stale { + out.WriteString("证据已变更，详细报告待更新。\n\n") + } + for i, b := range f.TrafficBindings { + fmt.Fprintf(&out, "%d. **证据 #%d · %s** — `%s %s`，状态码 %d\n", i+1, b.ID, b.Role, b.Snapshot.Method, strings.ReplaceAll(b.Snapshot.URL, "`", "%60"), b.Snapshot.Status) + if b.Note != "" { + fmt.Fprintf(&out, " %s\n", strings.ReplaceAll(b.Note, "\n", "\n ")) + } + if attachments { + fmt.Fprintf(&out, " [请求报文](evidence/%d/%d/request.http) · [响应报文](evidence/%d/%d/response.http)\n", f.ID, b.ID, f.ID, b.ID) + } + } + out.WriteString("\n") + return out.String() +} diff --git a/server/dto.go b/server/dto.go index 368c9c0..3799e34 100644 --- a/server/dto.go +++ b/server/dto.go @@ -292,6 +292,12 @@ func coverageAssetRefDTO(ref db.AssetRef) CoverageAssetRefDTO { // ---- Finding (frontend "Finding") ---- type FindingDTO struct { + TrafficCount int `json:"traffic_count"` + EvidenceVersion int64 `json:"evidence_version"` + ReportEvidenceVersion int64 `json:"report_evidence_version"` + ReportStale bool `json:"report_stale"` + TrafficBindings []db.FindingTrafficBinding `json:"traffic_bindings,omitempty"` + ID string `json:"id"` FindingID string `json:"finding_id,omitempty"` // standalone findings-table id — the handle for status updates VulnClass string `json:"vulnclass"` @@ -399,6 +405,7 @@ func findingDTOsForOwner(taskID, description string, in []*db.Node, meta map[int if m, ok := meta[n.ID]; ok { d.FindingID = i64s(m.ID) d.Status = m.Status + d.TrafficCount = m.TrafficCount d.Assets = findingAssetDTOs(m.AssetIDs, assets) } out = append(out, d) @@ -426,8 +433,10 @@ func findingFromDB(f *db.DBFinding, assets map[int64]*db.Asset) FindingDTO { status = db.FindingPending } d := FindingDTO{ - ID: i64s(f.ID), - FindingID: i64s(f.ID), + ID: i64s(f.ID), + FindingID: i64s(f.ID), + TrafficCount: f.TrafficCount, EvidenceVersion: f.EvidenceVersion, ReportEvidenceVersion: f.ReportEvidenceVersion, + ReportStale: f.Report != "" && f.EvidenceVersion != f.ReportEvidenceVersion, TrafficBindings: f.TrafficBindings, VulnClass: f.VulnClass, Name: f.Name, Severity: f.Severity, diff --git a/server/finding_traffic.go b/server/finding_traffic.go new file mode 100644 index 0000000..dfc511b --- /dev/null +++ b/server/finding_traffic.go @@ -0,0 +1,425 @@ +package server + +import ( + "bytes" + "context" + "encoding/json" + "errors" + "fmt" + "io" + "log" + "net/http" + "path/filepath" + "strconv" + "unicode/utf8" + + "github.com/Autumn-27/artex/agent" + "github.com/Autumn-27/artex/db" + "github.com/Autumn-27/artex/evidence" + actool "github.com/Autumn-27/norma/tool" +) + +func (s *Server) evidenceStore() *evidence.Store { + return evidence.New(s.m.pg, s.m.traffic, filepath.Join(s.m.dir, "evidence")) +} + +// Add only new optional properties; preserve edited descriptions, existing +// properties, agent bindings and disabled flags. The one-time flag also keeps +// subsequent user unbinding of the evidence reader intact. +func (s *Server) seedFindingTrafficTools() { + const flag = "finding_traffic_tools_v1" + if v, _, _ := s.m.pg.GetSetting(flag); v == "true" { + return + } + for _, key := range []string{"report_finding", "update_finding_report"} { + var schema any + if key == "update_finding_report" { + schema = s.toolUpdateFindingReport().InputSchema() + } else { + for _, seed := range agent.BuiltinToolSeeds() { + if seed.Key == key { + schema = seed.Schema + break + } + } + } + raw, err := json.Marshal(schema) + if err != nil { + log.Printf("[evidence] tool schema: %v", err) + return + } + var obj map[string]json.RawMessage + if err = json.Unmarshal(raw, &obj); err != nil { + return + } + // BuiltinToolSeeds stores Schema as RawMessage; both forms marshal as JSON. + var properties map[string]json.RawMessage + if err = json.Unmarshal(obj["properties"], &properties); err != nil { + return + } + name := "traffic_refs" + if key == "update_finding_report" { + name = "evidence_version" + } + if len(properties[name]) == 0 { + return + } + _, err = s.m.pg.Exec(`UPDATE tools SET schema=jsonb_set(schema,ARRAY['properties',$2::text],$3::jsonb,true),updated_at=now() +WHERE key=$1 AND system AND NOT(COALESCE(schema->'properties','{}'::jsonb) ? $2)`, key, name, string(properties[name])) + if err != nil { + log.Printf("[evidence] upgrade tool %s: %v", key, err) + return + } + } + if reporter, _ := s.m.pg.GetAgentByKey("reporter"); reporter != nil { + if err := s.m.pg.AddAgentToToolBinding("reporter", []string{"get_finding_traffic"}); err != nil { + return + } + } + _ = s.m.pg.SetSetting(flag, "true") +} + +func (s *Server) registerFindingTraffic(mux *http.ServeMux) { + base := "/api/exploration/findings/{id}/traffic" + mux.HandleFunc("GET "+base, s.getFindingTraffic) + mux.HandleFunc("POST "+base, s.bindFindingTraffic) + mux.HandleFunc("PATCH "+base+"/{binding_id}", s.editFindingTraffic) + mux.HandleFunc("DELETE "+base+"/{binding_id}", s.editFindingTraffic) + mux.HandleFunc("PUT "+base+"/order", s.editFindingTraffic) + mux.HandleFunc("GET "+base+"/{binding_id}", s.getFindingTrafficDetail) + mux.HandleFunc("GET "+base+"/{binding_id}/body", s.getFindingTrafficBody) +} + +func evidenceError(w http.ResponseWriter, err error) { + status := http.StatusUnprocessableEntity + if errors.Is(err, db.ErrEvidenceConflict) || errors.Is(err, db.ErrTaskArchiveState) { + status = http.StatusConflict + } + if errors.Is(err, db.ErrFindingNotFound) || errors.Is(err, db.ErrEvidenceNotFound) { + status = http.StatusNotFound + } + writeErr(w, status, err.Error()) +} + +func (s *Server) findingTrafficAccess(w http.ResponseWriter, r *http.Request, write bool) (int64, bool) { + id, err := strconv.ParseInt(r.PathValue("id"), 10, 64) + if err != nil || id <= 0 { + writeErr(w, 400, "invalid finding id") + return 0, false + } + f, err := s.m.pg.GetFinding(id) + if err != nil { + writeErr(w, 500, err.Error()) + return 0, false + } + if f == nil { + writeErr(w, 404, "finding not found") + return 0, false + } + if taskID := r.URL.Query().Get("context_task"); taskID != "" { + task := s.m.ResolveTask(taskID) + if task == nil { + writeErr(w, 404, "context task not found") + return 0, false + } + _, inherited, allowed := findingProvenanceInTask(task, f.TaskID) + if !allowed { + writeErr(w, 404, "finding not available in task context") + return 0, false + } + if write && inherited { + writeErr(w, 403, "继承漏洞只读，请在来源任务中修改") + return 0, false + } + } + return id, true +} + +func trafficSummary(in *db.FindingTraffic) *db.FindingTraffic { + out := *in + out.Bindings = append([]db.FindingTrafficBinding{}, in.Bindings...) + for i := range out.Bindings { + out.Bindings[i].Snapshot.ReqHead = "" + out.Bindings[i].Snapshot.RespHead = "" + } + return &out +} + +func (s *Server) getFindingTraffic(w http.ResponseWriter, r *http.Request) { + id, ok := s.findingTrafficAccess(w, r, false) + if !ok { + return + } + out, err := s.m.pg.GetFindingTraffic(r.Context(), id) + if err != nil { + evidenceError(w, err) + return + } + writeJSON(w, 200, trafficSummary(out)) +} + +func (s *Server) bindFindingTraffic(w http.ResponseWriter, r *http.Request) { + id, ok := s.findingTrafficAccess(w, r, true) + if !ok { + return + } + var body struct { + Refs []db.TrafficRef `json:"traffic_refs"` + } + if err := json.NewDecoder(http.MaxBytesReader(w, r.Body, 1<<20)).Decode(&body); err != nil { + writeErr(w, 400, "invalid body") + return + } + if len(body.Refs) == 0 { + writeErr(w, 400, "请选择流量") + return + } + out, err := s.evidenceStore().Bind(r.Context(), id, body.Refs) + if err != nil { + evidenceError(w, err) + return + } + writeJSON(w, 200, trafficSummary(out)) +} + +func (s *Server) editFindingTraffic(w http.ResponseWriter, r *http.Request) { + id, ok := s.findingTrafficAccess(w, r, true) + if !ok { + return + } + var body struct { + Version *int64 `json:"version"` + Role *string `json:"role"` + Note *string `json:"note"` + Order []string `json:"binding_ids"` + } + if err := json.NewDecoder(http.MaxBytesReader(w, r.Body, 1<<20)).Decode(&body); err != nil || body.Version == nil { + writeErr(w, 400, "version 和有效请求体必填") + return + } + var order []int64 + bindingID := int64(0) + if r.Method == http.MethodPut { + if body.Order == nil { + writeErr(w, 400, "binding_ids 必填") + return + } + order = []int64{} + for _, raw := range body.Order { + v, err := strconv.ParseInt(raw, 10, 64) + if err != nil || v <= 0 { + writeErr(w, 400, "invalid binding id") + return + } + order = append(order, v) + } + } else { + var err error + bindingID, err = strconv.ParseInt(r.PathValue("binding_id"), 10, 64) + if err != nil || bindingID <= 0 { + writeErr(w, 400, "invalid binding id") + return + } + } + err := s.m.pg.EditFindingTraffic(r.Context(), id, bindingID, *body.Version, body.Role, body.Note, r.Method == http.MethodDelete, order) + if err != nil { + evidenceError(w, err) + return + } + s.getFindingTraffic(w, r) +} + +type evidencePreview struct { + Content string `json:"content"` + Offset int64 `json:"offset"` + Total int64 `json:"total"` + NextOffset int64 `json:"next_offset"` + Truncated bool `json:"truncated"` + Binary bool `json:"binary"` +} + +func readEvidencePreview(store *evidence.Store, snapshot db.TrafficEvidenceSnapshot, side string, offset, length int64) (out evidencePreview, err error) { + if offset < 0 || length < 0 { + return out, errors.New("offset / length 不能为负数") + } + if length == 0 || length > 8192 { + length = 8192 + } + f, total, err := store.OpenBody(snapshot, side) + if err != nil { + return out, err + } + defer f.Close() + if offset > total { + return out, errors.New("offset 超出正文长度") + } + if _, err = f.Seek(offset, io.SeekStart); err != nil { + return out, err + } + raw, err := io.ReadAll(io.LimitReader(f, length)) + if err != nil { + return out, err + } + // Leave an incomplete trailing UTF-8 rune for the next page. Explicit byte + // offsets are still accepted; complete downloads always retain original bytes. + if offset+int64(len(raw)) < total && len(raw) >= utf8.UTFMax { + start := len(raw) - 1 + for start > 0 && raw[start]&0xc0 == 0x80 { + start-- + } + if !utf8.FullRune(raw[start:]) { + raw = raw[:start] + } + } + out = evidencePreview{Offset: offset, Total: total, NextOffset: offset + int64(len(raw)), Truncated: offset+int64(len(raw)) < total, + Binary: bytes.IndexByte(raw, 0) >= 0 || (offset == 0 && !utf8.Valid(raw))} + if out.Binary { + out.Content = fmt.Sprintf("[二进制正文，%d 字节；请下载查看]", total) + } else { + out.Content = string(bytes.ToValidUTF8(raw, []byte("�"))) + } + return out, nil +} + +func (s *Server) getFindingTrafficDetail(w http.ResponseWriter, r *http.Request) { + id, ok := s.findingTrafficAccess(w, r, false) + if !ok { + return + } + bid, err := strconv.ParseInt(r.PathValue("binding_id"), 10, 64) + if err != nil || bid <= 0 { + writeErr(w, 400, "invalid binding id") + return + } + store := s.evidenceStore() + var result any + err = store.WithBinding(r.Context(), id, bid, func(b db.FindingTrafficBinding) error { + req, err := readEvidencePreview(store, b.Snapshot, "request", 0, 8192) + if err != nil { + return err + } + resp, err := readEvidencePreview(store, b.Snapshot, "response", 0, 8192) + if err != nil { + return err + } + result = map[string]any{"binding": b, "request": req, "response": resp} + return nil + }) + if err != nil { + evidenceError(w, err) + return + } + writeJSON(w, 200, result) +} + +func (s *Server) getFindingTrafficBody(w http.ResponseWriter, r *http.Request) { + id, ok := s.findingTrafficAccess(w, r, false) + if !ok { + return + } + bid, err := strconv.ParseInt(r.PathValue("binding_id"), 10, 64) + if err != nil || bid <= 0 { + writeErr(w, 400, "invalid binding id") + return + } + side := r.URL.Query().Get("side") + store := s.evidenceStore() + started := false + err = store.WithBinding(r.Context(), id, bid, func(b db.FindingTrafficBinding) error { + if r.URL.Query().Get("download") == "1" { + f, length, err := store.OpenBody(b.Snapshot, side) + if err != nil { + return err + } + defer f.Close() + w.Header().Set("Content-Type", "application/octet-stream") + w.Header().Set("X-Content-Type-Options", "nosniff") + w.Header().Set("Content-Disposition", fmt.Sprintf("attachment; filename=\"evidence-%d-%s.bin\"", bid, side)) + w.Header().Set("Content-Length", strconv.FormatInt(length, 10)) + started = true + _, err = io.Copy(w, f) + return err + } + offset, length := int64(0), int64(8192) + for name, dst := range map[string]*int64{"offset": &offset, "length": &length} { + if raw := r.URL.Query().Get(name); raw != "" { + v, err := strconv.ParseInt(raw, 10, 64) + if err != nil { + return err + } + *dst = v + } + } + preview, err := readEvidencePreview(store, b.Snapshot, side, offset, length) + if err != nil { + return err + } + started = true + writeJSON(w, 200, preview) + return nil + }) + if err != nil && !started { + evidenceError(w, err) + } +} + +func (s *Server) toolGetFindingTraffic() actool.CoreTool { + return roTool("get_finding_traffic", "读取漏洞已绑定的真实流量证据，不依赖捕获开关。finding_id 使用 report_finding JSON 返回的独立漏洞记录 ID（不是第一行的探索节点 ID）。先不传 binding_id 获取清单及 version；空清单是正常情况，TCP 等非 HTTP 漏洞或未采集时仍可依据文字/命令证据编写报告，不强制绑定。有绑定时按 binding_id、side(request/response)、offset 分段读取正文。写报告时将读取的 version 作为 evidence_version 传给 update_finding_report，后者 finding_id 仍使用探索节点 ID。", + objSchema(map[string]any{"finding_id": strParam("独立漏洞记录 ID"), "binding_id": strParam("清单里的绑定 ID，省略则返回清单"), "side": strParam("request 或 response，默认 response"), "offset": map[string]any{"type": "integer"}, "length": map[string]any{"type": "integer"}}, "finding_id"), + func(ctx context.Context, in json.RawMessage) (actool.Result, error) { + var a struct { + FindingID json.RawMessage `json:"finding_id"` + BindingID json.RawMessage `json:"binding_id"` + Side string `json:"side"` + Offset, Length int64 + } + if err := json.Unmarshal(in, &a); err != nil { + return actool.Errorf(err.Error()), nil + } + id, bid := parseProfileID(a.FindingID), parseProfileID(a.BindingID) + if id <= 0 { + return actool.Errorf("finding_id 无效"), nil + } + if ri := agent.RunInfoFrom(ctx); ri.TaskID > 0 { + f, err := s.m.pg.GetFinding(id) + if err != nil || f == nil { + return actool.Errorf("漏洞不存在"), nil + } + task := s.m.ResolveTask(strconv.FormatInt(ri.TaskID, 10)) + if task == nil { + return actool.Errorf("任务不存在"), nil + } + if _, _, allowed := findingProvenanceInTask(task, f.TaskID); !allowed { + return actool.Errorf("当前任务不可读取该漏洞"), nil + } + } + if len(a.BindingID) == 0 { + list, err := s.m.pg.GetFindingTraffic(ctx, id) + if err != nil { + return actool.Errorf(err.Error()), nil + } + return jsonResult(trafficSummary(list)) + } + if a.Side == "" { + a.Side = "response" + } + store := s.evidenceStore() + var result any + err := store.WithBinding(ctx, id, bid, func(b db.FindingTrafficBinding) error { + preview, err := readEvidencePreview(store, b.Snapshot, a.Side, a.Offset, a.Length) + if err != nil { + return err + } + head := b.Snapshot.ReqHead + if a.Side == "response" { + head = b.Snapshot.RespHead + } + result = map[string]any{"binding_id": strconv.FormatInt(bid, 10), "head": head, "body": preview} + return nil + }) + if err != nil { + return actool.Errorf(err.Error()), nil + } + return jsonResult(result) + }) +} diff --git a/server/finding_traffic_export.go b/server/finding_traffic_export.go new file mode 100644 index 0000000..99dc599 --- /dev/null +++ b/server/finding_traffic_export.go @@ -0,0 +1,121 @@ +package server + +import ( + "archive/zip" + "encoding/json" + "errors" + "fmt" + "io" + "os" + "strings" + "time" + + "github.com/Autumn-27/artex/db" + "github.com/Autumn-27/artex/evidence" + "github.com/Autumn-27/artex/report" +) + +func writeFindingsEvidenceZip(out io.Writer, findings []*db.DBFinding, stage string, now time.Time) error { + zw := zip.NewWriter(out) + if err := writeFindingsZipEntries(zw, findings, stage, now); err != nil { + zw.Close() + return err + } + return zw.Close() +} + +func writeFindingsZipEntries(zw *zip.Writer, findings []*db.DBFinding, stage string, now time.Time) error { + store := evidence.New(nil, nil, stage) // private export copy; not subject to GC + used := map[string]int{} + for _, f := range findings { + base := report.FindingFilename(f) + name := base + if n := used[base]; n > 0 { + name = fmt.Sprintf("%s-%d.md", strings.TrimSuffix(base, ".md"), n+1) + } + used[base]++ + w, err := zw.Create(name) + if err != nil { + return err + } + if _, err = io.WriteString(w, report.SingleFindingMarkdown(f, now)); err != nil { + return err + } + for _, b := range f.TrafficBindings { + prefix := fmt.Sprintf("evidence/%d/%d/", f.ID, b.ID) + w, err := zw.Create(prefix + "manifest.json") + if err != nil { + return err + } + if err = json.NewEncoder(w).Encode(b); err != nil { + return err + } + for _, side := range []string{"request", "response"} { + if err := func() error { + body, _, err := store.OpenBody(b.Snapshot, side) + if err != nil { + return err + } + defer body.Close() + w, err := zw.Create(prefix + side + ".http") + if err != nil { + return err + } + head := b.Snapshot.ReqHead + if side == "response" { + head = b.Snapshot.RespHead + } + if _, err = io.WriteString(w, strings.TrimRight(head, "\r\n")+"\n\n"); err != nil { + return err + } + if _, err = io.Copy(w, body); err != nil { + return err + } + if _, err = body.Seek(0, io.SeekStart); err != nil { + return err + } + w, err = zw.Create(prefix + side + ".bin") + if err != nil { + return err + } + _, err = io.Copy(w, body) + return err + }(); err != nil { + return err + } + } + } + } + return nil +} + +func buildFindingsEvidenceZip(path string, findings []*db.DBFinding, stage string, now time.Time) error { + f, err := os.OpenFile(path, os.O_CREATE|os.O_EXCL|os.O_WRONLY, 0o600) + if err != nil { + return err + } + writeErr := writeFindingsEvidenceZip(f, findings, stage, now) + syncErr := f.Sync() + closeErr := f.Close() + if err := errors.Join(writeErr, syncErr, closeErr); err != nil { + return err + } + // Re-read every compressed entry to validate the finished ZIP, including CRCs, + // before the caller starts an HTTP download response. + archive, err := zip.OpenReader(path) + if err != nil { + return err + } + defer archive.Close() + for _, entry := range archive.File { + r, err := entry.Open() + if err != nil { + return err + } + _, readErr := io.Copy(io.Discard, r) + if err := errors.Join(readErr, r.Close()); err != nil { + return err + } + } + return nil +} diff --git a/server/finding_traffic_test.go b/server/finding_traffic_test.go new file mode 100644 index 0000000..0f82a53 --- /dev/null +++ b/server/finding_traffic_test.go @@ -0,0 +1,473 @@ +package server + +import ( + "archive/zip" + "bytes" + "context" + "encoding/json" + "fmt" + "io" + "net/http/httptest" + "os" + "path/filepath" + "strconv" + "strings" + "testing" + "time" + + "github.com/Autumn-27/artex/agent" + "github.com/Autumn-27/artex/db" + "github.com/Autumn-27/artex/traffic" +) + +func trafficEvidenceServer(t *testing.T) (*Server, *db.RecordedFinding, func(string, string, string) *httptest.ResponseRecorder) { + t.Helper() + m, err := NewManager(t.TempDir(), "") + if err != nil { + t.Fatal(err) + } + t.Cleanup(func() { m.Close() }) + m.traffic, err = traffic.Open(filepath.Join(m.dir, "traffic"), ":0") + if err != nil { + t.Fatal(err) + } + task, err := m.CreateTask("traffic evidence API", "local evidence", nil, 0, 0) + if err != nil { + t.Fatal(err) + } + tid, _ := strconv.ParseInt(task.ID, 10, 64) + t.Cleanup(func() { m.pg.Exec(`DELETE FROM task_archives WHERE task_id=$1`, tid); m.pg.DeleteTask(tid) }) + ctx, cancel := context.WithCancel(context.Background()) + cancel() + s := New(ctx, m, t.TempDir(), t.TempDir(), t.TempDir()) + s.archiveWG.Wait() + token, err := signJWT(s.jwtKey) + if err != nil { + t.Fatal(err) + } + request := func(method, path, body string) *httptest.ResponseRecorder { + t.Helper() + r := httptest.NewRequest(method, path, strings.NewReader(body)) + r.Header.Set("Authorization", "Bearer "+token) + r.Header.Set("Content-Type", "application/json") + w := httptest.NewRecorder() + s.Handler().ServeHTTP(w, r) + return w + } + input := db.RecordFindingInput{TaskID: tid, ExplorationID: task.ExpID, Worker: "test", VulnClass: "TEST", Summary: "bound evidence", Name: "Evidence API", Severity: "low"} + f, err := s.evidenceStore().Record(context.Background(), input, nil) + if err != nil { + t.Fatal(err) + } + return s, f, request +} + +func seedServerEvidenceFlow(t *testing.T, s *Server, id string, body []byte) { + t.Helper() + _, err := s.m.traffic.DB().Exec(`INSERT INTO exchanges(id,ts,host,method,url_template,url,status,content_type,req_len,resp_len,path) VALUES(?,1,'evidence.local','POST','/','http://evidence.local/',200,'application/octet-stream',0,?,'')`, id, len(body)) + if err != nil { + t.Fatal(err) + } + _, err = s.m.traffic.DB().Exec(`INSERT INTO exchange_bodies(id,req_head,resp_head,req_body,resp_body) VALUES(?,'POST / HTTP/1.1','HTTP/1.1 200 OK',?,?)`, id, []byte{}, body) + if err != nil { + t.Fatal(err) + } +} + +func TestFindingTrafficToolUpgradePreservesCustomization(t *testing.T) { + s, _, _ := trafficEvidenceServer(t) + pg := s.m.pg + worker, err := pg.GetAgentByKey("worker") + if err != nil || worker == nil { + t.Fatal("missing worker", err) + } + oldPrompt, _ := pg.CurrentPrompt(worker.ID) + t.Cleanup(func() { pg.SavePrompt(worker.ID, oldPrompt, "restore test fixture", "test") }) + if _, err := pg.SavePrompt(worker.ID, "USER CUSTOM PROMPT", "test", "test"); err != nil { + t.Fatal(err) + } + custom := json.RawMessage(`{"type":"object","properties":{"evidence":{"type":"string","description":"user evidence instructions"}},"required":["evidence"]}`) + for _, key := range []string{"report_finding", "update_finding_report", "get_finding_traffic"} { + old, err := pg.GetTool(key) + if err != nil || old == nil { + t.Fatal("missing tool", key, err) + } + t.Cleanup(func() { + bindings, _ := json.Marshal(old.Agents) + pg.UpdateTool(old.Key, old.Description, old.Schema, bindings, old.Enabled) + }) + if err := pg.UpdateTool(key, "user description", custom, json.RawMessage(`["worker"]`), false); err != nil { + t.Fatal(err) + } + } + if err := pg.SetSetting("finding_traffic_tools_v1", "false"); err != nil { + t.Fatal(err) + } + s.seedFindingTrafficTools() + for key, property := range map[string]string{"report_finding": "traffic_refs", "update_finding_report": "evidence_version"} { + tool, err := pg.GetTool(key) + if err != nil || tool == nil { + t.Fatal(err) + } + var schema struct { + Properties map[string]json.RawMessage `json:"properties"` + Required []string `json:"required"` + } + if err := json.Unmarshal(tool.Schema, &schema); err != nil { + t.Fatal(err) + } + if tool.Enabled || tool.Description != "user description" || len(tool.Agents) != 1 || tool.Agents[0] != "worker" || len(schema.Required) != 1 || schema.Required[0] != "evidence" || len(schema.Properties[property]) == 0 || !strings.Contains(string(schema.Properties["evidence"]), "user evidence instructions") { + t.Fatalf("custom configuration overwritten: %+v", tool) + } + } + reader, _ := pg.GetTool("get_finding_traffic") + if reader.Enabled || !contains(reader.Agents, "reporter") { + t.Fatal("reader should be bound without enabling it") + } + if err := pg.RemoveAgentFromTool("reporter", "get_finding_traffic"); err != nil { + t.Fatal(err) + } + s.seedFindingTrafficTools() + reader, _ = pg.GetTool("get_finding_traffic") + if contains(reader.Agents, "reporter") { + t.Fatal("restart overwrote user's unbinding") + } + if got, _ := pg.CurrentPrompt(worker.ID); got != "USER CUSTOM PROMPT" { + t.Fatal("custom prompt replaced") + } +} + +func TestFindingTrafficAPIAndExport(t *testing.T) { + s, f, req := trafficEvidenceServer(t) + body := bytes.Repeat([]byte{0, 255, 65, 66}, 5000) + for _, id := range []string{"a", "b", "c"} { + seedServerEvidenceFlow(t, s, id, body) + } + base := fmt.Sprintf("/api/exploration/findings/%d/traffic", f.FindingID) + w := req("POST", base, `{"traffic_refs":[{"traffic_id":"a","role":"baseline"},{"traffic_id":"b","role":"proof","note":"proof note"},{"traffic_id":"c"}]}`) + if w.Code != 200 { + t.Fatalf("bind: %d %s", w.Code, w.Body) + } + var list db.FindingTraffic + if err := json.Unmarshal(w.Body.Bytes(), &list); err != nil { + t.Fatal(err) + } + if len(list.Bindings) != 3 || list.Bindings[0].Snapshot.ReqHead != "" { + t.Fatal("not a bounded summary") + } + first, second, third := list.Bindings[0].ID, list.Bindings[1].ID, list.Bindings[2].ID + if w = req("PATCH", fmt.Sprintf("%s/%d", base, first), `{"version":0,"note":"stale"}`); w.Code != 409 { + t.Fatal(w.Code, w.Body) + } + if w = req("PUT", base+"/order", fmt.Sprintf(`{"version":1,"binding_ids":["%d","%d","%d"]}`, third, first, second)); w.Code != 200 { + t.Fatal(w.Code, w.Body) + } + if w = req("PUT", base+"/order", fmt.Sprintf(`{"version":2,"binding_ids":["%d"]}`, first)); w.Code != 409 { + t.Fatal(w.Code, w.Body) + } + if w = req("GET", fmt.Sprintf("%s/%d", base, first), ""); w.Code != 200 || !strings.Contains(w.Body.String(), `"binary":true`) { + t.Fatal(w.Code, w.Body) + } + if w = req("GET", fmt.Sprintf("%s/%d/body?side=response&offset=20001", base, first), ""); w.Code != 422 { + t.Fatal(w.Code, w.Body) + } + if _, err := s.m.traffic.DeleteHost("evidence.local"); err != nil { + t.Fatal(err) + } + if w = req("GET", fmt.Sprintf("%s/%d/body?side=response&download=1", base, first), ""); w.Code != 200 || !bytes.Equal(w.Body.Bytes(), body) { + t.Fatal("download after original removal failed") + } + // IDs belong to the requested finding; cross-finding detail cannot be read. + finding, err := s.m.pg.GetFinding(f.FindingID) + if err != nil { + t.Fatal(err) + } + task, err := s.m.pg.GetTask(*finding.TaskID) + if err != nil { + t.Fatal(err) + } + other, err := s.evidenceStore().Record(t.Context(), db.RecordFindingInput{TaskID: task.ID, ExplorationID: task.ExplorationID, Summary: "other", Severity: "low"}, nil) + if err != nil { + t.Fatal(err) + } + if w = req("GET", fmt.Sprintf("/api/exploration/findings/%d/traffic/%d", other.FindingID, first), ""); w.Code != 404 { + t.Fatal(w.Code, w.Body) + } + exportURL := fmt.Sprintf("/api/exploration/findings/export?format=md-zip&scope=selected&ids=%d", f.FindingID) + w = req("GET", exportURL, "") + if w.Code != 200 { + t.Fatal(w.Code, w.Body) + } + archive, err := zip.NewReader(bytes.NewReader(w.Body.Bytes()), int64(w.Body.Len())) + if err != nil { + t.Fatal(err) + } + binaries := 0 + linked := false + for _, entry := range archive.File { + r, err := entry.Open() + if err != nil { + t.Fatal(err) + } + raw, err := io.ReadAll(r) + r.Close() + if err != nil { + t.Fatal(err) + } + if strings.HasSuffix(entry.Name, "/response.bin") { + binaries++ + if !bytes.Equal(raw, body) { + t.Fatal("export changed binary") + } + } + if strings.HasSuffix(entry.Name, ".md") && bytes.Contains(raw, []byte(fmt.Sprintf("evidence/%d/%d/", f.FindingID, third))) { + linked = true + } + } + if binaries != 3 || !linked { + t.Fatalf("missing attachments/links: %d %v", binaries, linked) + } + // Reports can use snapshots with capture stopped and original rows gone. + result, err := s.toolGetFindingTraffic().Call(context.Background(), json.RawMessage(fmt.Sprintf(`{"finding_id":"%d"}`, f.FindingID)), nil) + if err != nil || !strings.Contains(result.Flatten(), "proof note") { + t.Fatal(result, err) + } + // Corrupt body must fail before an attachment response is sent. + snap := list.Bindings[0].Snapshot + path := filepath.Join(s.m.dir, "evidence", "blobs", snap.RespHash[:2], snap.RespHash+".bin") + if err = os.WriteFile(path, []byte("corrupt"), 0600); err != nil { + t.Fatal(err) + } + if w = req("GET", exportURL, ""); w.Code < 400 || strings.Contains(w.Header().Get("Content-Type"), "zip") { + t.Fatal("corrupt export returned an attachment", w.Code) + } +} + +func TestFindingTrafficArchiveV3RoundTripAndRetry(t *testing.T) { + s, f, _ := trafficEvidenceServer(t) + ctx := context.Background() + s.ctx = ctx + body := []byte("evidence without original traffic\x00\xff") + seedServerEvidenceFlow(t, s, "original", body) + list, err := s.evidenceStore().Bind(ctx, f.FindingID, []db.TrafficRef{{TrafficID: "original", Role: "proof", Note: "keep note"}}) + if err != nil { + t.Fatal(err) + } + finding, err := s.m.pg.GetFinding(f.FindingID) + if err != nil { + t.Fatal(err) + } + tid := *finding.TaskID + sharedTask, err := s.m.CreateTask("shared evidence", "fixture", nil, 0, 0) + if err != nil { + t.Fatal(err) + } + sid, _ := strconv.ParseInt(sharedTask.ID, 10, 64) + defer s.m.pg.DeleteTask(sid) + shared, err := s.evidenceStore().Record(ctx, db.RecordFindingInput{TaskID: sid, ExplorationID: sharedTask.ExpID, Summary: "shared", Severity: "low"}, []db.TrafficRef{{TrafficID: "original"}}) + if err != nil { + t.Fatal(err) + } + if _, err = s.m.traffic.DeleteHost("evidence.local"); err != nil { + t.Fatal(err) + } + if err = s.m.SetTaskPaused(strconv.FormatInt(tid, 10), true); err != nil { + t.Fatal(err) + } + job, err := s.m.pg.QueueTaskArchive(tid) + if err != nil { + t.Fatal(err) + } + if err = s.m.pg.EditFindingTraffic(ctx, f.FindingID, list.Bindings[0].ID, 1, nil, nil, true, nil); err == nil { + t.Fatal("write allowed after archive queued") + } + if err = s.runOneTaskArchiveJob(); err != nil { + t.Fatal(err) + } + if got, _ := s.m.pg.GetFinding(f.FindingID); got != nil { + t.Fatal("archive retained hot finding") + } + if err = s.evidenceStore().Collect(ctx, time.Now().Add(48*time.Hour)); err != nil { + t.Fatal(err) + } + if err = s.evidenceStore().WithBinding(ctx, shared.FindingID, shared.Traffic.Bindings[0].ID, func(b db.FindingTrafficBinding) error { + r, _, err := s.evidenceStore().OpenBody(b.Snapshot, "response") + if r != nil { + r.Close() + } + return err + }); err != nil { + t.Fatal("shared body deleted", err) + } + if _, err = s.m.pg.QueueTaskArchiveRestore(job.ID); err != nil { + t.Fatal(err) + } + // Force a database error after validated body installation; then retry normally. + if _, err = s.m.pg.Exec(`CREATE FUNCTION fail_evidence_restore_test() RETURNS trigger LANGUAGE plpgsql AS $$ BEGIN RAISE EXCEPTION 'restore fixture failure'; END $$`); err != nil { + t.Fatal(err) + } + if _, err = s.m.pg.Exec(`CREATE TRIGGER fail_evidence_restore_test BEFORE INSERT ON finding_traffic_bindings FOR EACH ROW EXECUTE FUNCTION fail_evidence_restore_test()`); err != nil { + t.Fatal(err) + } + defer s.m.pg.Exec(`DROP FUNCTION IF EXISTS fail_evidence_restore_test() CASCADE`) + if err = s.runOneTaskArchiveJob(); err == nil { + t.Fatal("expected restore failure") + } + if got, _ := s.m.pg.GetFinding(f.FindingID); got != nil { + t.Fatal("partial restore") + } + if _, err = s.m.pg.Exec(`DROP FUNCTION fail_evidence_restore_test() CASCADE`); err != nil { + t.Fatal(err) + } + if _, err = s.m.pg.QueueTaskArchiveRestore(job.ID); err != nil { + t.Fatal(err) + } + if err = s.runOneTaskArchiveJob(); err != nil { + t.Fatal(err) + } + restored, err := s.m.pg.GetFindingTraffic(ctx, f.FindingID) + if err != nil { + t.Fatal(err) + } + b := restored.Bindings[0] + if restored.Version != list.Version || b.ID != list.Bindings[0].ID || b.Note != "keep note" || b.Role != "proof" || b.Position != 0 { + t.Fatalf("restored: %+v", restored) + } + if err = s.evidenceStore().WithBinding(ctx, f.FindingID, b.ID, func(b db.FindingTrafficBinding) error { + r, _, err := s.evidenceStore().OpenBody(b.Snapshot, "response") + if err != nil { + return err + } + defer r.Close() + got, err := io.ReadAll(r) + if !bytes.Equal(got, body) { + t.Fatal("restore changed body") + } + return err + }); err != nil { + t.Fatal(err) + } + // The read tool enforces task visibility too. + result, err := s.toolGetFindingTraffic().Call(agent.WithRunInfo(ctx, agent.RunInfo{TaskID: sid}), json.RawMessage(fmt.Sprintf(`{"finding_id":"%d"}`, f.FindingID)), nil) + if err != nil || !strings.Contains(result.Flatten(), "不可读取") { + t.Fatal(result, err) + } +} + +func TestFindingTrafficFailedReportDoesNotTrigger(t *testing.T) { + s, f, _ := trafficEvidenceServer(t) + finding, err := s.m.pg.GetFinding(f.FindingID) + if err != nil { + t.Fatal(err) + } + task, err := s.m.pg.GetTask(*finding.TaskID) + if err != nil { + t.Fatal(err) + } + var last int64 + if err = s.m.pg.QueryRow(`SELECT COALESCE(max(id),0) FROM activity`).Scan(&last); err != nil { + t.Fatal(err) + } + if err = s.m.pg.SetSchedState(schedKeyLastToolCall, strconv.FormatInt(last, 10)); err != nil { + t.Fatal(err) + } + id, err := s.m.pg.Exploration(task.ExplorationID).AppendActivity(db.Activity{Kind: "tool_result", Tool: "report_finding", ToolUseID: "failure-fixture", IsError: true, Detail: "body missing"}) + if err != nil { + t.Fatal(err) + } + // A nil server makes any attempted trigger fail this test immediately. + scheduler := &Scheduler{pg: s.m.pg} + scheduler.fireToolCalls([]*db.AgentTrigger{{OnToolCall: true, ToolNames: []string{"report_finding"}, AgentKey: "reporter"}}) + if got := scheduler.mustState(schedKeyLastToolCall); got != strconv.FormatInt(id, 10) { + t.Fatal("failed finding must still advance watermark", got) + } +} + +func TestFindingTrafficUTF8SegmentsAndInheritedWrites(t *testing.T) { + s, f, req := trafficEvidenceServer(t) + ctx := context.Background() + text := strings.Repeat("中文证据🙂", 2000) + seedServerEvidenceFlow(t, s, "unicode", []byte(text)) + list, err := s.evidenceStore().Bind(ctx, f.FindingID, []db.TrafficRef{{TrafficID: "unicode"}}) + if err != nil { + t.Fatal(err) + } + var joined strings.Builder + offset := int64(0) + for { + var preview evidencePreview + err = s.evidenceStore().WithBinding(ctx, f.FindingID, list.Bindings[0].ID, func(b db.FindingTrafficBinding) error { + var e error + preview, e = readEvidencePreview(s.evidenceStore(), b.Snapshot, "response", offset, 8192) + return e + }) + if err != nil { + t.Fatal(err) + } + if preview.Binary { + t.Fatal("valid UTF-8 classified binary") + } + joined.WriteString(preview.Content) + if !preview.Truncated { + break + } + if preview.NextOffset <= offset { + t.Fatal("preview made no progress") + } + offset = preview.NextOffset + } + if joined.String() != text { + t.Fatal("UTF-8 character split damaged preview") + } + finding, err := s.m.pg.GetFinding(f.FindingID) + if err != nil { + t.Fatal(err) + } + child, err := s.m.CreateTaskWithOptions("inherited evidence API", "fixture", db.TaskCreateOptions{SourceTaskIDs: []int64{*finding.TaskID}}) + if err != nil { + t.Fatal(err) + } + cid, _ := strconv.ParseInt(child.ID, 10, 64) + defer s.m.pg.DeleteTask(cid) + base := fmt.Sprintf("/api/exploration/findings/%d/traffic", f.FindingID) + suffix := "?context_task=" + child.ID + if w := req("GET", base+suffix, ""); w.Code != 200 { + t.Fatal(w.Code, w.Body) + } + for _, method := range []string{"POST", "PATCH", "PUT", "DELETE"} { + path := base + if method == "PATCH" || method == "DELETE" { + path += fmt.Sprintf("/%d", list.Bindings[0].ID) + } + if method == "PUT" { + path += "/order" + } + w := req(method, path+suffix, `{"version":1,"traffic_refs":[{"traffic_id":"unicode"}]}`) + if w.Code != 403 { + t.Fatal(method, w.Code, w.Body) + } + } + if _, err = s.m.traffic.DeleteHost("evidence.local"); err != nil { + t.Fatal(err) + } + result, err := s.toolUpdateFindingReport().Call(ctx, json.RawMessage(fmt.Sprintf(`{"finding_id":%d,"evidence_version":1,"report":"## 证据报告\n\n证据 #%d：已验证完整请求响应"}`, f.NodeID, list.Bindings[0].ID)), nil) + if err != nil || result.IsError { + t.Fatal(result, err) + } + updated, err := s.m.pg.GetFinding(f.FindingID) + if err != nil { + t.Fatal(err) + } + if updated.ReportEvidenceVersion != 1 || updated.EvidenceVersion != 1 { + t.Fatal("report did not record the read evidence version") + } + for _, format := range []string{"md-single", "json", "csv"} { + w := req("GET", fmt.Sprintf("/api/exploration/findings/export?scope=selected&ids=%d&format=%s", f.FindingID, format), "") + if w.Code != 200 { + t.Fatal(format, w.Code, w.Body) + } + if !strings.Contains(w.Body.String(), fmt.Sprint(list.Bindings[0].ID)) { + t.Fatal("export omitted evidence ID", format) + } + } +} diff --git a/server/orchestration.go b/server/orchestration.go index 56ead66..50c2a47 100644 --- a/server/orchestration.go +++ b/server/orchestration.go @@ -75,6 +75,7 @@ func (s *Server) orchestrationTools() []actool.CoreTool { s.toolSearchWorkerTraces(), s.toolGetTaskNodeDetail(), s.toolUpdateFindingReport(), + s.toolGetFindingTraffic(), } } @@ -392,20 +393,22 @@ func (s *Server) toolUpdateFindingReport() actool.CoreTool { return wrTool("update_finding_report", "为已登记的漏洞写入/更新【详细报告】(Markdown 全文,整段覆盖旧内容)。finding_id 传 report_finding 返回的那个 id(\"finding recorded: \" 里的数字)。报告建议包含:漏洞概述、影响与危害、复现步骤、证据/PoC、修复建议。", objSchema(map[string]any{ - "finding_id": map[string]any{"type": "integer", "description": "目标漏洞 id(report_finding 返回的 id)"}, - "report": strParam("详细报告全文,Markdown 格式"), + "finding_id": map[string]any{"type": "integer", "description": "目标漏洞 id(report_finding 返回的 id)"}, + "report": strParam("详细报告全文,Markdown 格式"), + "evidence_version": map[string]any{"type": "integer", "description": "get_finding_traffic 返回的证据 version；用于防止报告覆盖新的证据变更"}, }, "finding_id", "report"), - func(_ context.Context, in json.RawMessage) (actool.Result, error) { + func(ctx context.Context, in json.RawMessage) (actool.Result, error) { var a struct { - FindingID json.RawMessage `json:"finding_id"` - Report string `json:"report"` + EvidenceVersion *int64 `json:"evidence_version"` + FindingID json.RawMessage `json:"finding_id"` + Report string `json:"report"` } _ = json.Unmarshal(in, &a) nodeID := parseProfileID(a.FindingID) // 复用「数字或数字字符串」解析 if nodeID <= 0 { return actool.Errorf("finding_id 无效"), nil } - n, err := s.m.pg.SetFindingReportByNodeID(nodeID, a.Report) + n, err := s.m.pg.SetFindingReportVersionByNodeID(ctx, nodeID, a.Report, a.EvidenceVersion) if err != nil { return actool.Errorf(err.Error()), nil } @@ -453,11 +456,12 @@ func (s *Server) seedOrchestrationTools() { s.seedWorkerReadToolsUnbind() // list_facts/node_detail/list_companies/跨 work 检索从 worker 默认解绑(一次性) s.seedAutoReportFindingBinding() s.unbindGoalMetDefault() - s.reseedGoalsPrompt() // goals 提示词加入「抽操作约束」步 → 旧库追加一版新默认(一次性) - s.reseedMainAgentPrompt() // mainagent 提示词加入「目标达成后 add_intent 反问是否建目标」(一次性) - s.reseedPlannerPrompt() // planner 提示词:重写「0 意图」正当理由 + 加量化验收核对(一次性) - s.reseedWorkerPrompt() // worker 提示词:加否定结论证据门槛(一次性) - s.seedReporterAgent() // 预置「报告撰写」agent + 工具绑定 + finding 触发器(一次性) + s.reseedGoalsPrompt() // goals 提示词加入「抽操作约束」步 → 旧库追加一版新默认(一次性) + s.reseedMainAgentPrompt() // mainagent 提示词加入「目标达成后 add_intent 反问是否建目标」(一次性) + s.reseedPlannerPrompt() // planner 提示词:重写「0 意图」正当理由 + 加量化验收核对(一次性) + s.reseedWorkerPrompt() // worker 提示词:加否定结论证据门槛(一次性) + s.seedReporterAgent() // 预置「报告撰写」agent + 工具绑定 + finding 触发器(一次性) + s.seedFindingTrafficTools() // 增加可选证据参数及只读证据工具，保留用户配置 // 注：pentest 的默认工具绑定无需迁移——BuiltinToolSeeds 在全新初始化时就把 // list_assets/insert_assets/report_finding/list_findings/list_companies 连同 // pentest 一起 seed 好了（项目尚无旧库，不做迁移）。 diff --git a/server/scheduler.go b/server/scheduler.go index 97cd0b2..0c6b1e8 100644 --- a/server/scheduler.go +++ b/server/scheduler.go @@ -305,6 +305,11 @@ func (sc *Scheduler) fireToolCalls(triggers []*db.AgentTrigger) { if len(want) == 0 { continue } + // A failed finding write has no committed finding to report. Keep other + // tool-error triggers available for user-defined automation. + if e.ToolIsErr && e.Tool == "report_finding" { + continue + } errTag := "" if e.ToolIsErr { errTag = "[error] " diff --git a/server/server.go b/server/server.go index 3f31803..fe2e712 100644 --- a/server/server.go +++ b/server/server.go @@ -1,7 +1,6 @@ package server import ( - "archive/zip" "context" "encoding/json" "errors" @@ -11,6 +10,7 @@ import ( "net" "net/http" "net/url" + "os" "path/filepath" "regexp" "sort" @@ -217,9 +217,10 @@ func New(ctx context.Context, m *Manager, skillDir string, dataDir string, keyDi } wireAgentAugment(m.pg, s.skillDir, s.hostTools) // 可见 skills/MCP + 流量/编排 host 工具装配进 agent 工具集 domainReg := buildDomainReg(m.Assets()) - wireTools(m.pg, domainReg) // 内置工具表：按 agent 过滤 + 覆盖描述/schema + 注入默认值 - seedPrompts(m.pg) // 内置 agent 默认提示词正文播种进 agent_prompts(仅空时) - s.seedOrchestrationTools() // P2 跨任务编排工具 seed 进 tools 表(可按 agent 绑定) + wireTools(m.pg, domainReg) // 内置工具表：按 agent 过滤 + 覆盖描述/schema + 注入默认值 + seedPrompts(m.pg) // 内置 agent 默认提示词正文播种进 agent_prompts(仅空时) + s.seedOrchestrationTools() // P2 跨任务编排工具 seed 进 tools 表(可按 agent 绑定) + go s.evidenceStore().RunGC(s.ctx) s.seedPythonInterpreter() // 自定义脚本工具:开机检测 python 解释器入库(仅空时) go newScheduler(s).Run(s.ctx) // P3 触发器调度(定时/finding/目标事件),仅自定义 agent // Fill the tool cache for any enabled MCP that has none yet (notably the @@ -360,7 +361,7 @@ func (s *Server) saveLLMConfig(cfg agent.Config) error { ContextWindowK: cfg.ContextWindowK, ThinkingType: cfg.ThinkingType, ReasoningEffort: cfg.ReasoningEffort, IsDefault: true, Priority: priority, PoolExclude: poolExclude, Streaming: streaming, MaxTokens: maxTokens, MaxTokensField: maxTokensField, SessionHeaderKey: sessionHeaderKey, - Retry: retry, + Retry: retry, }) if err != nil { return err @@ -405,6 +406,7 @@ func (s *Server) buildPlannerWorker(pinID *int64, gProv llm.Provider, gCfg agent // the tools-table binding (default = worker), so worker behavior is unchanged. wProv, wCfg := s.providerForAgent("worker", pinID, gProv, gCfg) wk := agent.NewWorker(wProv, wCfg.Model, s.m.dir, tx, wCfg.CompactionWindow(), s.agentMaxTurns("worker")) + wk.SetFindingRecorder(s.evidenceStore()) wk.SetRunTimeout(time.Duration(s.agentRunSeconds("worker")) * time.Second) wk.SetProxy(s.m.ProxyAddr(), s.m.ProxyCACert()) wk.SetMemory(memory.NewStore(filepath.Join(s.m.dir, "memory"))) @@ -414,6 +416,7 @@ func (s *Server) buildPlannerWorker(pinID *int64, gProv llm.Provider, gCfg agent wk.SetMaxTokens(maxTokensResolver(wCfg)) // 单次回复输出上限(0 = 不发) pProv, pCfg := s.providerForAgent("planner", pinID, gProv, gCfg) pl := agent.NewPlanner(pProv, pCfg.Model, s.m.dir, tx, pCfg.CompactionWindow(), s.agentMaxTurns("planner")) + pl.SetFindingRecorder(s.evidenceStore()) pl.SetKillWork(s.engine.KillWork) // planner kill_work → terminate a running work pl.SetSteerWork(s.engine.SteerWork) // planner steer_work → inject mid-run course-correction pl.SetProxy(s.m.ProxyAddr(), s.m.ProxyCACert()) // WebFetch through the recording proxy @@ -482,6 +485,7 @@ func (s *Server) applyLLM(cfg agent.Config) error { mProv, mCfg := s.providerForAgent("mainagent", nil, prov, cfg) s.cfgMu.Lock() s.mainAgent = agent.NewMainAgent(mProv, mCfg.Model, s.m.dir, tx, mCfg.CompactionWindow(), s.agentMaxTurns("mainagent")) + s.mainAgent.SetFindingRecorder(s.evidenceStore()) s.mainAgent.SetProxy(s.m.ProxyAddr(), s.m.ProxyCACert()) // WebFetch through the recording proxy s.mainAgent.SetWebSearch(s.webSearchFor("mainagent")) s.mainAgent.SetSteerWork(s.engine.SteerWork) // steer_work：人对运行中 work 实时纠偏 @@ -811,6 +815,7 @@ func (s *Server) Handler() http.Handler { mux.HandleFunc("GET /api/exploration/findings/asset-tree", s.findingAssetTree) mux.HandleFunc("GET /api/exploration/findings/stats", s.findingStats) mux.HandleFunc("GET /api/exploration/findings/export", s.findingsExport) + s.registerFindingTraffic(mux) mux.HandleFunc("GET /api/exploration/findings/{id}", s.getFinding) mux.HandleFunc("GET /api/exploration/findings/{id}/lineage", s.findingLineage) mux.HandleFunc("POST /api/exploration/findings/{id}/deepen", s.deepenFinding) @@ -2135,6 +2140,16 @@ func (s *Server) findingsExport(w http.ResponseWriter, r *http.Request) { return } + stage, err := os.MkdirTemp("", "artex-finding-export-") + if err != nil { + writeErr(w, 500, err.Error()) + return + } + defer os.RemoveAll(stage) + if err = s.evidenceStore().StageFindingsExport(r.Context(), fs, stage, format == "md-zip"); err != nil { + evidenceError(w, err) + return + } now := time.Now() stamp := now.Format("20060102-150405") setDownload := func(contentType, filename string) { @@ -2147,27 +2162,19 @@ func (s *Server) findingsExport(w http.ResponseWriter, r *http.Request) { setDownload("text/markdown; charset=utf-8", "findings-"+stamp+".md") _, _ = w.Write([]byte(report.FindingsMarkdown(fs, now))) case "md-zip": + path := filepath.Join(stage, "findings.zip") + if err := buildFindingsEvidenceZip(path, fs, stage, now); err != nil { + writeErr(w, 500, err.Error()) + return + } + file, err := os.Open(path) + if err != nil { + writeErr(w, 500, err.Error()) + return + } + defer file.Close() setDownload("application/zip", "findings-"+stamp+".zip") - zw := zip.NewWriter(w) - used := map[string]int{} - for _, f := range fs { - base := report.FindingFilename(f) - name := base - // 去重:同名文件追加 -2、-3…… - if n := used[base]; n > 0 { - name = fmt.Sprintf("%s-%d.md", strings.TrimSuffix(base, ".md"), n+1) - } - used[base]++ - fw, werr := zw.Create(name) - if werr != nil { - log.Printf("[findings-export] zip create %s: %v", name, werr) - continue - } - _, _ = fw.Write([]byte(report.SingleFindingMarkdown(f, now))) - } - if cerr := zw.Close(); cerr != nil { - log.Printf("[findings-export] zip close: %v", cerr) - } + http.ServeContent(w, r, "findings.zip", now, file) case "csv": setDownload("text/csv; charset=utf-8", "findings-"+stamp+".csv") _, _ = w.Write(report.FindingsCSV(fs)) @@ -2175,6 +2182,10 @@ func (s *Server) findingsExport(w http.ResponseWriter, r *http.Request) { assets := s.resolveFindingAssets(fs) out := make([]FindingDTO, 0, len(fs)) for _, f := range fs { + for i := range f.TrafficBindings { + f.TrafficBindings[i].Snapshot.ReqHead = "" + f.TrafficBindings[i].Snapshot.RespHead = "" + } out = append(out, findingFromDB(f, assets)) } setDownload("application/json; charset=utf-8", "findings-"+stamp+".json") diff --git a/server/task_archives.go b/server/task_archives.go index 30cea69..9d0163f 100644 --- a/server/task_archives.go +++ b/server/task_archives.go @@ -178,6 +178,13 @@ func (s *Server) archiveTask(job *pgdb.TaskArchive) (runErr error) { snapshot.DataCounts["traffic"] = trafficCount } archivePath := taskArchivePath(s.m.dir, job.ID, taskID) + evidenceSnapshots, err := pgdb.ArchiveEvidenceSnapshots(snapshot) + if err != nil { + return err + } + if err = s.evidenceStore().CopySnapshots(s.ctx, evidenceSnapshots, filepath.Join(fileStage.payload, "evidence")); err != nil { + return err + } _ = s.m.pg.UpdateTaskArchiveProgress(job.ID, "compress", 50) originalSize, compressedSize, checksum, err := writeTaskArchivePackage(archivePath, fileStage.payload, snapshot) if err != nil { @@ -270,6 +277,17 @@ func (s *Server) restoreTaskArchive(job *pgdb.TaskArchive) (runErr error) { if snapshot.TaskID != job.TaskID || !pgdb.IsTaskArchiveFormatSupported(snapshot.FormatVersion) { return pgdb.ErrTaskArchiveFormatMismatch } + evidenceSnapshots, err := pgdb.ArchiveEvidenceSnapshots(&snapshot) + if err != nil { + return err + } + return s.evidenceStore().WithInstalledSnapshots(s.ctx, evidenceSnapshots, filepath.Join(extracted, "evidence"), func() error { + return s.restoreTaskArchivePayload(job, &snapshot, extracted) + }) +} + +func (s *Server) restoreTaskArchivePayload(job *pgdb.TaskArchive, snapshot *pgdb.TaskArchiveSnapshot, extracted string) (runErr error) { + var err error var llmRecordsFile *os.File relative, hasStreamedLLMRecords := snapshot.StreamedTables["llm_records"] if len(snapshot.StreamedTables) > 0 && @@ -309,13 +327,13 @@ func (s *Server) restoreTaskArchive(job *pgdb.TaskArchive) (runErr error) { var warnings []string if llmRecordsFile != nil { warnings, err = s.m.pg.RestoreTaskArchiveWithLLMRecords( - job.ID, &snapshot, job.RemainingTimeoutSeconds, llmRecordsFile, + job.ID, snapshot, job.RemainingTimeoutSeconds, llmRecordsFile, ) closeErr := llmRecordsFile.Close() llmRecordsFile = nil err = errors.Join(err, closeErr) } else { - warnings, err = s.m.pg.RestoreTaskArchive(job.ID, &snapshot, job.RemainingTimeoutSeconds) + warnings, err = s.m.pg.RestoreTaskArchive(job.ID, snapshot, job.RemainingTimeoutSeconds) } if err != nil { return err diff --git a/server/task_llm.go b/server/task_llm.go index 05893fb..42c59b4 100644 --- a/server/task_llm.go +++ b/server/task_llm.go @@ -563,6 +563,7 @@ func (s *Server) agentsForTask(t *Task) *taskAgentBundle { tx := transcript.NewStore(filepath.Join(s.m.dir, "transcripts")) window := workerRuntime.CompactionWindow() wk := agent.NewWorker(workerRuntime, "task-router", s.m.dir, tx, window, s.agentMaxTurns("worker")) + wk.SetFindingRecorder(s.evidenceStore()) wk.SetCompactionWindowResolver(workerRuntime.CompactionWindow) wk.SetNonStreaming(workerRuntime.nonStreaming) // 按任务当前激活 profile 的流式开关(每轮读) wk.SetMaxTokens(workerRuntime.maxTokens) // 同上,输出上限也跟随当前激活 profile @@ -572,6 +573,7 @@ func (s *Server) agentsForTask(t *Task) *taskAgentBundle { wk.SetWebSearch(s.webSearchFor("worker")) wk.SetConstraintInject(s.constraintInjectWorker) // 操作约束注入 worker(可配置,默认开) pl := agent.NewPlanner(plannerRuntime, "task-router", s.m.dir, tx, plannerRuntime.CompactionWindow(), s.agentMaxTurns("planner")) + pl.SetFindingRecorder(s.evidenceStore()) pl.SetCompactionWindowResolver(plannerRuntime.CompactionWindow) pl.SetNonStreaming(plannerRuntime.nonStreaming) pl.SetMaxTokens(plannerRuntime.maxTokens) @@ -581,6 +583,7 @@ func (s *Server) agentsForTask(t *Task) *taskAgentBundle { pl.SetWebSearch(s.webSearchFor("planner")) pl.SetConstraintInject(s.constraintInjectPlanner) // 操作约束注入 planner(可配置,默认开) main := agent.NewMainAgent(mainRuntime, "task-router", s.m.dir, tx, mainRuntime.CompactionWindow(), s.agentMaxTurns("mainagent")) + main.SetFindingRecorder(s.evidenceStore()) main.SetCompactionWindowResolver(mainRuntime.CompactionWindow) main.SetNonStreaming(mainRuntime.nonStreaming) main.SetMaxTokens(mainRuntime.maxTokens) diff --git a/traffic/evidence.go b/traffic/evidence.go new file mode 100644 index 0000000..b96a0b0 --- /dev/null +++ b/traffic/evidence.go @@ -0,0 +1,126 @@ +package traffic + +import ( + "bufio" + "bytes" + "context" + "database/sql" + "errors" + "fmt" + "io" + "os" + "path/filepath" + "strings" +) + +// EvidenceExchange exposes complete captured bodies only during the callback. +// Readers are closed before releasing the writer lock, so host deletion cannot +// remove a blob between its index lookup and its evidence copy. +type EvidenceExchange struct { + ID, URL, Method, ContentType string + TS int64 + Status int + ReqHead, RespHead string + ReqLen, RespLen int64 + ReqHash, RespHash string + Request, Response io.Reader +} + +func (t *Traffic) ReadEvidence(ctx context.Context, ids []string, consume func(EvidenceExchange) error) error { + if t == nil { + return errors.New("流量录制存储不可用") + } + t.wmu.Lock() + defer t.wmu.Unlock() + for _, id := range ids { + if err := ctx.Err(); err != nil { + return err + } + if err := t.readEvidence(id, consume); err != nil { + return fmt.Errorf("流量 %s: %w", id, err) + } + } + return nil +} + +func (t *Traffic) readEvidence(id string, consume func(EvidenceExchange) error) error { + e := EvidenceExchange{ID: id} + var legacy string + if err := t.db.QueryRow(`SELECT ts,url,method,status,content_type,req_len,resp_len,path FROM exchanges WHERE id=?`, id). + Scan(&e.TS, &e.URL, &e.Method, &e.Status, &e.ContentType, &e.ReqLen, &e.RespLen, &legacy); err != nil { + return err + } + var req, resp []byte + var reqBlob, respBlob sql.NullString + err := t.db.QueryRow(`SELECT req_head,req_body,req_blob,resp_head,resp_body,resp_blob FROM exchange_bodies WHERE id=?`, id). + Scan(&e.ReqHead, &req, &reqBlob, &e.RespHead, &resp, &respBlob) + if errors.Is(err, sql.ErrNoRows) && legacy != "" { + if !filepath.IsLocal(legacy) { + return errors.New("旧流量路径无效") + } + r, err := os.Open(filepath.Join(t.dir, legacy, "request.http")) + if err != nil { + return err + } + defer r.Close() + p, err := os.Open(filepath.Join(t.dir, legacy, "response.http")) + if err != nil { + return err + } + defer p.Close() + e.ReqHead, e.Request, err = splitLegacyEvidence(r) + if err != nil { + return err + } + e.RespHead, e.Response, err = splitLegacyEvidence(p) + if err != nil { + return err + } + return consume(e) + } + if err != nil { + return err + } + e.Request = bytes.NewReader(req) + e.Response = bytes.NewReader(resp) + if reqBlob.Valid && reqBlob.String != "" { + f, _, err := t.Blob(reqBlob.String) + if err != nil { + return err + } + defer f.Close() + e.Request = f + e.ReqHash = reqBlob.String + } + if respBlob.Valid && respBlob.String != "" { + f, _, err := t.Blob(respBlob.String) + if err != nil { + return err + } + defer f.Close() + e.Response = f + e.RespHash = respBlob.String + } + return consume(e) +} + +func splitLegacyEvidence(r io.Reader) (string, io.Reader, error) { + b := bufio.NewReader(r) + var head strings.Builder + for { + line, err := b.ReadString('\n') + if line == "\n" || line == "\r\n" { + return head.String(), b, nil + } + head.WriteString(line) + if head.Len() > 1<<20 { + return "", nil, errors.New("旧流量报文头过大") + } + if errors.Is(err, io.EOF) { + return head.String(), b, nil + } + if err != nil { + return "", nil, err + } + } +} diff --git a/web/src/app/(main)/function/findings/_components/findings-table.tsx b/web/src/app/(main)/function/findings/_components/findings-table.tsx index 570448b..b331bf6 100644 --- a/web/src/app/(main)/function/findings/_components/findings-table.tsx +++ b/web/src/app/(main)/function/findings/_components/findings-table.tsx @@ -27,6 +27,7 @@ import { AlertDialogTitle, AlertDialogTrigger, } from "@/components/ui/alert-dialog"; +import { Badge } from "@/components/ui/badge"; import { Button } from "@/components/ui/button"; import { Checkbox } from "@/components/ui/checkbox"; import { Input } from "@/components/ui/input"; @@ -205,6 +206,7 @@ export function FindingsTable({ {f.name || f.vulnclass || "未分类"} )} {f.summary} + 流量证据 {f.traffic_count ?? 0} 条 diff --git a/web/src/app/(main)/function/findings/detail/page.tsx b/web/src/app/(main)/function/findings/detail/page.tsx index 7d36f5c..4068850 100644 --- a/web/src/app/(main)/function/findings/detail/page.tsx +++ b/web/src/app/(main)/function/findings/detail/page.tsx @@ -1,32 +1,31 @@ "use client"; import * as React from "react"; + import Link from "next/link"; import { useSearchParams } from "next/navigation"; -import { toast } from "sonner"; -import { ArrowLeftIcon, ArrowUpRightIcon, ShieldAlertIcon } from "lucide-react"; -import { SidebarTrigger } from "@/components/ui/sidebar"; -import { Separator } from "@/components/ui/separator"; +import { ArrowLeftIcon, ArrowUpRightIcon, ShieldAlertIcon } from "lucide-react"; +import { toast } from "sonner"; + +import { CopyButton } from "@/components/copy-button"; +import { FindingTrafficPanel } from "@/components/finding-traffic-panel"; +import { Markdown } from "@/components/markdown"; +import { StatusBadge } from "@/components/status-badge"; +import { Alert, AlertDescription } from "@/components/ui/alert"; import { Badge } from "@/components/ui/badge"; import { Button } from "@/components/ui/button"; import { Card, CardContent, CardHeader, CardTitle } from "@/components/ui/card"; +import { Select, SelectContent, SelectGroup, SelectItem, SelectTrigger } from "@/components/ui/select"; +import { Separator } from "@/components/ui/separator"; +import { SidebarTrigger } from "@/components/ui/sidebar"; import { Tabs, TabsContent, TabsList, TabsTrigger } from "@/components/ui/tabs"; -import { - Select, - SelectContent, - SelectGroup, - SelectItem, - SelectTrigger, -} from "@/components/ui/select"; -import { StatusBadge } from "@/components/status-badge"; -import { CopyButton } from "@/components/copy-button"; -import { Markdown } from "@/components/markdown"; -import { FindingLineageView } from "./lineage"; -import { statusMeta } from "@/lib/status"; import { api } from "@/lib/api"; +import { statusMeta } from "@/lib/status"; import type { Finding, FindingStatus, Severity } from "@/lib/types"; +import { FindingLineageView } from "./lineage"; + const SEVERITIES: Severity[] = ["critical", "high", "medium", "low"]; const FINDING_STATUSES: FindingStatus[] = [ "pending", @@ -44,21 +43,11 @@ function fmtTime(ts: string) { } // FieldRow is one label/value line in the right-hand status panel. -function FieldRow({ - label, - children, -}: { - label: string; - children: React.ReactNode; -}) { +function FieldRow({ label, children }: { label: string; children: React.ReactNode }) { return (
- - {label} - -
- {children} -
+ {label} +
{children}
); } @@ -123,9 +112,7 @@ function FindingDetailInner() { if (!finding) { return (
-

- {loaded ? `未找到发现 ${id}` : "加载中…"} -

+

{loaded ? `未找到发现 ${id}` : "加载中…"}

{loaded && (
+ 流量证据 {f.traffic_count ?? 0} 条 {f.assets && f.assets.length > 0 && (
{f.assets.slice(0, 2).map((a) => ( @@ -95,22 +85,12 @@ function Row({ {a.label} ))} - {f.assets.length > 2 && ( - - +{f.assets.length - 2} - - )} + {f.assets.length > 2 && +{f.assets.length - 2}}
)} {f.finding_id && !f.inherited ? ( - onStatus(f, v as FindingStatus)}> + @@ -146,9 +126,7 @@ function Row({ {open && (
-
- 证据 / PoC -
+
证据 / PoC
             {f.evidence}
@@ -187,25 +165,18 @@ export function FindingsTab({ taskId }: { taskId: string }) { }; }, [taskId]); - const onStatus = React.useCallback( - async (f: Finding, next: FindingStatus) => { - if (f.inherited || !f.finding_id || next === f.status) return; - const prev = f.status; - setFindings((cur) => - cur.map((x) => (x.id === f.id ? { ...x, status: next } : x)), - ); - try { - await api.setFindingStatus(f.finding_id, next); - toast.success(`已标记为「${statusMeta("finding", next).label}」`); - } catch (e) { - setFindings((cur) => - cur.map((x) => (x.id === f.id ? { ...x, status: prev } : x)), - ); - toast.error("更新失败：" + (e as Error).message); - } - }, - [], - ); + const onStatus = React.useCallback(async (f: Finding, next: FindingStatus) => { + if (f.inherited || !f.finding_id || next === f.status) return; + const prev = f.status; + setFindings((cur) => cur.map((x) => (x.id === f.id ? { ...x, status: next } : x))); + try { + await api.setFindingStatus(f.finding_id, next); + toast.success(`已标记为「${statusMeta("finding", next).label}」`); + } catch (e) { + setFindings((cur) => cur.map((x) => (x.id === f.id ? { ...x, status: prev } : x))); + toast.error("更新失败：" + (e as Error).message); + } + }, []); const items = findings .filter((f) => f.task_id === taskId || f.inherited) @@ -243,9 +214,7 @@ export function FindingsTab({ taskId }: { taskId: string }) { ))} {items.length === 0 && ( -

- 本任务及直接关联任务暂无确认发现。 -

+

本任务及直接关联任务暂无确认发现。

)} diff --git a/web/src/app/(main)/function/traffic/page.tsx b/web/src/app/(main)/function/traffic/page.tsx index df0c5a3..c9efafd 100644 --- a/web/src/app/(main)/function/traffic/page.tsx +++ b/web/src/app/(main)/function/traffic/page.tsx @@ -5,19 +5,18 @@ import * as React from "react"; import { ArrowDownWideNarrowIcon, ArrowUpNarrowWideIcon, - CheckIcon, ChevronLeftIcon, ChevronRightIcon, - CopyIcon, ListChecksIcon, Loader2Icon, RadioTowerIcon, SearchIcon, Trash2Icon, - WrapTextIcon, } from "lucide-react"; import { toast } from "sonner"; +import { HttpCodeBlock } from "@/components/http-code-block"; +import { LinkTrafficDialog } from "@/components/link-traffic-dialog"; import { AlertDialog, AlertDialogAction, @@ -41,7 +40,7 @@ import { Tabs, TabsContent, TabsList, TabsTrigger } from "@/components/ui/tabs"; import { Tooltip, TooltipContent, TooltipTrigger } from "@/components/ui/tooltip"; import { api } from "@/lib/api"; import type { TrafficDetail, TrafficExchange, TrafficHost, TrafficResp } from "@/lib/types"; -import { cn, copyText } from "@/lib/utils"; +import { cn } from "@/lib/utils"; function fmtTime(ts: string) { return new Date(ts).toLocaleString("zh-CN", { @@ -90,205 +89,6 @@ function requestWithHost(raw: string, exchange: TrafficExchange): string { return `${raw.slice(0, firstLineEnd + newline.length)}Host: ${host}${newline}${raw.slice(firstLineEnd + newline.length)}`; } -function StartLine({ line }: { line: string }) { - const response = /^(HTTP\/\S+)(\s+)(\d{3})(.*)$/.exec(line); - if (response) { - return ( - <> - {response[1]} - {response[2]} - {response[3]} - {response[4]} - - ); - } - - const request = /^([A-Z]+)(\s+)(\S+)(\s+)(HTTP\/\S+)$/.exec(line); - if (request) { - return ( - <> - {request[1]} - {request[2]} - {request[3]} - {request[4]} - {request[5]} - - ); - } - - return line; -} - -function HeaderLine({ line }: { line: string }) { - const separator = line.indexOf(":"); - if (separator <= 0) return line; - return ( - <> - {line.slice(0, separator)} - : - {line.slice(separator + 1)} - - ); -} - -function JsonBody({ body }: { body: string }) { - const parts: React.ReactNode[] = []; - const tokens = /("(?:\\.|[^"\\])*")(\s*:)?|\b(true|false|null)\b|-?\d+(?:\.\d+)?(?:[eE][+-]?\d+)?/g; - let cursor = 0; - for (const match of body.matchAll(tokens)) { - const index = match.index ?? 0; - if (index > cursor) parts.push(body.slice(cursor, index)); - let className = "text-chart-4"; - if (match[1]) className = match[2] ? "text-primary" : "text-chart-2"; - else if (match[3]) className = "text-chart-3"; - parts.push( - - {match[0]} - , - ); - cursor = index + match[0].length; - } - if (cursor < body.length) parts.push(body.slice(cursor)); - return parts; -} - -function MarkupBody({ body }: { body: string }) { - const parts: React.ReactNode[] = []; - const tags = /<\/?[A-Za-z][^>]*>|/g; - let cursor = 0; - for (const match of body.matchAll(tags)) { - const index = match.index ?? 0; - if (index > cursor) parts.push(body.slice(cursor, index)); - parts.push( - - {match[0]} - , - ); - cursor = index + match[0].length; - } - if (cursor < body.length) parts.push(body.slice(cursor)); - return parts; -} - -type BodyFormat = "json" | "markup" | "plain"; - -function detectBodyFormat(body: string): BodyFormat { - const trimmed = body.trimStart(); - if (trimmed.startsWith("{") || trimmed.startsWith("[")) return "json"; - if (trimmed.startsWith("<")) return "markup"; - return "plain"; -} - -function HighlightedBody({ body, format }: { body: string; format: BodyFormat }) { - if (format === "json") return ; - if (format === "markup") return ; - return body; -} - -function HttpCodeBlock({ raw }: { raw: string }) { - const [wrapLines, setWrapLines] = React.useState(true); - const [copied, setCopied] = React.useState(false); - const value = raw || "（空）"; - const lines = value.replaceAll("\r\n", "\n").split("\n"); - const separator = lines.indexOf(""); - const body = separator >= 0 ? lines.slice(separator + 1).join("\n") : ""; - const bodyFormat = detectBodyFormat(body); - - React.useEffect(() => { - if (!copied) return; - const timer = window.setTimeout(() => setCopied(false), 1500); - return () => window.clearTimeout(timer); - }, [copied]); - - const copyPacket = async () => { - const ok = await copyText(value); - if (ok) { - setCopied(true); - return; - } - toast.error("复制失败，请使用 Ctrl/Cmd+A 后复制"); - }; - - const renderLine = (line: string, index: number) => { - if (index === 0) return ; - if (separator < 0 || index < separator) return ; - if (index === separator) return null; - return ; - }; - - return ( -
-
- - - - - {wrapLines ? "关闭自动换行" : "开启自动换行"} - - - - - - {copied ? "已复制" : "复制报文"} - -
- {/* biome-ignore lint/a11y/useSemanticElements: textarea cannot preserve line numbers and syntax-highlighting markup. */} -
{ - if (!(event.ctrlKey || event.metaKey) || event.key.toLowerCase() !== "a") return; - event.preventDefault(); - const selection = window.getSelection(); - if (!selection) return; - const range = document.createRange(); - range.selectNodeContents(event.currentTarget); - selection.removeAllRanges(); - selection.addRange(range); - }} - > - {lines.map((line, index) => ( -
- - {renderLine(line, index)} - -
- ))} -
-
- ); -} - // Fixed method set (server filters exact-match); avoids deriving options from a // single page, which would only ever list the methods on that page. const METHODS = ["GET", "POST", "PUT", "DELETE", "PATCH", "HEAD", "OPTIONS"]; @@ -296,6 +96,8 @@ const PAGE_SIZES = [25, 50, 100, 200]; type HostCountSortDirection = "asc" | "desc"; export default function TrafficPage() { + const [selectedFlows, setSelectedFlows] = React.useState>(() => new Set()); + const [linking, setLinking] = React.useState(false); const [page, setPage] = React.useState(0); const [size, setSize] = React.useState(50); const [host, setHost] = React.useState(""); // raw host input @@ -623,6 +425,17 @@ export default function TrafficPage() {
+
+ 已选 {selectedFlows.size} 条流量 + + {selectedFlows.size > 0 ? ( + + ) : null} +
{/* History table */}
@@ -630,6 +443,22 @@ export default function TrafficPage() { + + 0 && exchanges.every((e) => selectedFlows.has(e.id))} + onCheckedChange={(checked) => + setSelectedFlows((previous) => { + const next = new Set(previous); + for (const e of exchanges) { + if (checked === true) next.add(e.id); + else next.delete(e.id); + } + return next; + }) + } + /> + 时间 host 方法 @@ -646,6 +475,21 @@ export default function TrafficPage() { className={cn("cursor-pointer", selected?.id === e.id && "bg-accent hover:bg-accent")} onClick={() => setSelected(e)} > + + event.stopPropagation()} + onCheckedChange={(checked) => + setSelectedFlows((previous) => { + const next = new Set(previous); + if (checked === true) next.add(e.id); + else next.delete(e.id); + return next; + }) + } + /> + {fmtTime(e.ts)} {e.host} @@ -665,7 +509,7 @@ export default function TrafficPage() { ))} {exchanges.length === 0 && ( - + {traffic === null ? "加载中…" : "没有匹配的流量。"} @@ -676,6 +520,13 @@ export default function TrafficPage() { + {linking ? ( + setLinking(false)} + onBound={() => setSelectedFlows(new Set())} + /> + ) : null} !open && setSelected(null)}> {selected && ( diff --git a/web/src/components/finding-traffic-panel.tsx b/web/src/components/finding-traffic-panel.tsx new file mode 100644 index 0000000..3c6da1f --- /dev/null +++ b/web/src/components/finding-traffic-panel.tsx @@ -0,0 +1,285 @@ +"use client"; + +import * as React from "react"; + +import { ArrowDownIcon, ArrowUpIcon, PlusIcon } from "lucide-react"; +import { toast } from "sonner"; + +import { TrafficEvidenceViewer } from "@/components/traffic-evidence-viewer"; +import { TrafficPickerDialog } from "@/components/traffic-picker-dialog"; +import { Alert, AlertDescription } from "@/components/ui/alert"; +import { Badge } from "@/components/ui/badge"; +import { Button } from "@/components/ui/button"; +import { Card, CardContent, CardDescription, CardHeader, CardTitle } from "@/components/ui/card"; +import { + Dialog, + DialogContent, + DialogDescription, + DialogFooter, + DialogHeader, + DialogTitle, +} from "@/components/ui/dialog"; +import { Empty, EmptyDescription, EmptyHeader, EmptyTitle } from "@/components/ui/empty"; +import { Field, FieldGroup, FieldLabel } from "@/components/ui/field"; +import { Select, SelectContent, SelectGroup, SelectItem, SelectTrigger, SelectValue } from "@/components/ui/select"; +import { Skeleton } from "@/components/ui/skeleton"; +import { Textarea } from "@/components/ui/textarea"; +import { api } from "@/lib/api"; +import type { FindingTraffic, FindingTrafficBinding, TrafficEvidenceRole } from "@/lib/types"; + +const ROLES: Record = { + baseline: "正常对照", + proof: "漏洞证明", + verification: "补充验证", + supporting: "辅助证据", +}; + +export function FindingTrafficPanel({ + findingId, + contextTask, + readOnly, + onChanged, +}: { + findingId: string; + contextTask?: string; + readOnly?: boolean; + onChanged: () => void; +}) { + const [data, setData] = React.useState(null); + const [error, setError] = React.useState(""); + const [adding, setAdding] = React.useState(false); + const [preview, setPreview] = React.useState(null); + const [editing, setEditing] = React.useState(null); + const [role, setRole] = React.useState("supporting"); + const [note, setNote] = React.useState(""); + const [busy, setBusy] = React.useState(false); + const [reload, setReload] = React.useState(0); + // biome-ignore lint/correctness/useExhaustiveDependencies: retry token intentionally refreshes this request. + React.useEffect(() => { + let active = true; + setData(null); + setError(""); + api + .findingTraffic(findingId, contextTask) + .then((value) => { + if (active) setData(value); + }) + .catch((e: Error) => { + if (active) setError(e.message); + }); + return () => { + active = false; + }; + }, [findingId, contextTask, reload]); + + async function mutate(action: () => Promise) { + setBusy(true); + try { + setData(await action()); + setEditing(null); + onChanged(); + toast.success("流量证据已更新"); + } catch (e) { + toast.error((e as Error).message); + setReload((n) => n + 1); + } finally { + setBusy(false); + } + } + + function move(index: number, delta: number) { + if (!data) return; + const ids = data.bindings.map((b) => b.id); + [ids[index], ids[index + delta]] = [ids[index + delta], ids[index]]; + void mutate(() => api.orderFindingTraffic(findingId, ids, data.version, contextTask)); + } + + return ( + <> + + +
+ 关联流量 {data ? `(${data.bindings.length})` : ""} + {!readOnly ? ( + + ) : ( + 继承证据 · 只读 + )} +
+ 按复现顺序组织请求与响应。清理原始流量后，已绑定证据仍然保留。 + + + {error ? ( + + + {error} + + + + ) : null} + {!data ? ( + + ) : data.bindings.length === 0 ? ( + + + 暂无关联流量 + 可绑定正常对照、漏洞证明和补充验证请求。 + + + ) : ( + data.bindings.map((b, index) => ( +
+
+ + {index + 1}. {ROLES[b.role]} + + + {b.snapshot.method} · {b.snapshot.status} + + 证据 #{b.id} +
+ + {b.note ?

{b.note}

: null} +
+ + {new Date(b.snapshot.captured_at * 1000).toLocaleString("zh-CN")} + +
+ + {!readOnly ? ( + <> + + + + + + ) : null} +
+
+
+ )) + )} + + + {adding && data ? ( + setAdding(false)} + onBound={() => { + setReload((n) => n + 1); + onChanged(); + }} + /> + ) : null} + setPreview(null)} + /> + { + if (!open && !busy) setEditing(null); + }} + > + + + 编辑流量证据 + 说明这组请求/响应如何支持漏洞结论。 + + + + 用途 + + + + 证据说明 +