feat: 支持漏洞关联多条流量证据

This commit is contained in:
ji ruo committed 2026-09-10 13:03:20 +08:00
1 parent 8dae851b9b
commit ec4ef35fff
41 files changed
+4313 -419

No files matched your search

+22
View File
@@ -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 }
+105
View File
@@ -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")
}
}
+2
View File
@@ -17,6 +17,7 @@ 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 {
findingRecorder FindingRecorder
prov llm.Provider
model string
tx *transcript.Store // raw LLM conversation persistence (nil = off)
@@ -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())
}
+2
View File
@@ -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())
}
+3
View File
@@ -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)
}
+10 -7
View File
@@ -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,见上下文"任务: #<id>")
- report_finding 的**入参**(vulnclass / severity / summary / evidence 等)
- report_finding 的**返回**:形如 "finding recorded: <id>" —— 这个 **<id> 就是本漏洞的 finding_id**,是你后面保存报告的句柄。
- report_finding 的**返回**:形如 "finding recorded: <id>" —— 这个 **<id> 是探索节点 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=<finding_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=<Markdown 全文>)** 保存。这是你的最终产物——不写进去等于没做。
1. **取全证据**:用 get_task_node_detail(task_id, id=<node_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=<node_id>, report=<Markdown 全文>, evidence_version=<实际读取的 version>)** 保存;未读取版本时省略 evidence_version,不得猜测。这是你的最终产物——不写进去等于没做。
━━ 报告结构(Markdown,按需裁剪,但证据/复现/修复必须有)━━
- ` + "`## 概述`" + `:一句话说清是什么漏洞、在哪、能造成什么。
+32 -39
View File
@@ -73,6 +73,7 @@ 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 {
findingRecorder FindingRecorder
as *db.AssetStore // asset store (optional; nil = asset tools not available)
cs *db.CompanyStore // company store (optional)
ts *db.ExplorationStore
@@ -1358,53 +1359,43 @@ 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/多个):参数/端点/站点等。一个漏洞影响多处可全填,纯观察可不填。"},
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 文本"),
}, "vulnclass", "severity", "summary"),
func(_ context.Context, in json.RawMessage) (actool.Result, error) {
"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"`
}
_ = 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)
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
}
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)
refs, err := db.NormalizeTrafficRefs(a.TrafficRefs)
if err != nil {
return actool.Errorf(err.Error()), nil
}
if intent > 0 {
_ = t.ts.Link(intent, db.RelYields, id) // chain: intent -> finding
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
}
_, _ = 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)
iid := input.IntentID
if iid <= 0 {
iid = t.ownerNode
}
@@ -1412,12 +1403,14 @@ func (t *ToolSet) addFinding() actool.CoreTool {
} 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
// 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
})
}
+6 -2
View File
@@ -41,6 +41,7 @@ type WebSearchOpts struct {
}
type Worker struct {
findingRecorder FindingRecorder
prov llm.Provider
model string
workDir string
@@ -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)
+404
View File
@@ -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(&current); 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
}
+85
View File
@@ -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)
}
+87
View File
@@ -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)
}
})
}
}
+34 -22
View File
@@ -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
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
}
if err != nil {
return 0, err
return 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)
var nodeID sql.NullInt64
if err := tx.QueryRow(`DELETE FROM findings WHERE id=$1 RETURNING node_id`, id).Scan(&nodeID); err != nil {
return err
}
return 1, nil
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)
+33
View File
@@ -1058,6 +1058,39 @@ CREATE INDEX IF NOT EXISTS idx_findings_status ON findings(status, created_at DE
-- 「按资产」视图靠 asset_ids @> '[<id>]' 反查发现,没有这个 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. 后端日志持久化
-- =====================================================================
+3 -1
View File
@@ -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}},
+7
View File
@@ -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 {
+91
View File
@@ -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。正文按哈希存放在 `<data>/evidence/blobs/<前两位>/<hash>.bin`,与可清理的 `data/traffic` 独立,多个漏洞可共享快照/正文。快照不提供内容更新接口,校验不一致时读取和导出失败。
在原流量写锁下读取完整正文,包括大正文 blob 和旧目录记录。先持久化并校验文件,再用单个 PostgreSQL 事务写入探索节点、意图关系、漏洞、快照和绑定;提交后才通知规划者。失败可能遗留无引用文件,但不产生部分业务记录。
PostgreSQL advisory lock `7337741004` 协调证据文件和 SQL 引用;任务行锁禁止归档排队后的证据修改。恢复从正文安装到元数据提交全程保持证据锁。删除漏洞级联移除绑定。
清理器每小时运行,只回收无引用且非进行中操作的内容,至少延迟 24 小时。普通流量清理不触及证据目录。备份热数据时同时备份 PostgreSQL 和 `data/evidence`。
## 导出与归档
Markdown 包含有序证据清单及版本;JSON 包含元数据;CSV 增加数量和绑定 ID。`md-zip` 保留漏洞 Markdown,并提供:
```text
evidence/<finding_id>/<binding_id>/
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 ./<package> -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 的证书问题。
+401
View File
@@ -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)
}
}
}
}
+333
View File
@@ -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, &note, 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")
}
}
+35
View File
@@ -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()
}
+36 -1
View File
@@ -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()
}
+9
View File
@@ -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)
@@ -428,6 +435,8 @@ func findingFromDB(f *db.DBFinding, assets map[int64]*db.Asset) FindingDTO {
d := FindingDTO{
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,
+425
View File
@@ -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)
})
}
+121
View File
@@ -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
}
+473
View File
@@ -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)
}
}
}
+6 -2
View File
@@ -75,6 +75,7 @@ func (s *Server) orchestrationTools() []actool.CoreTool {
s.toolSearchWorkerTraces(),
s.toolGetTaskNodeDetail(),
s.toolUpdateFindingReport(),
s.toolGetFindingTraffic(),
}
}
@@ -394,9 +395,11 @@ func (s *Server) toolUpdateFindingReport() actool.CoreTool {
objSchema(map[string]any{
"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 {
EvidenceVersion *int64 `json:"evidence_version"`
FindingID json.RawMessage `json:"finding_id"`
Report string `json:"report"`
}
@@ -405,7 +408,7 @@ func (s *Server) toolUpdateFindingReport() actool.CoreTool {
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
}
@@ -458,6 +461,7 @@ func (s *Server) seedOrchestrationTools() {
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 好了(项目尚无旧库,不做迁移)。
+5
View File
@@ -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] "
+32 -21
View File
@@ -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"
@@ -220,6 +220,7 @@ func New(ctx context.Context, m *Manager, skillDir string, dataDir string, keyDi
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
@@ -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")
+20 -2
View File
@@ -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
+3
View File
@@ -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)
+126
View File
@@ -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
}
}
}
@@ -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({
<span className="truncate font-medium">{f.name || f.vulnclass || "未分类"}</span>
)}
<span className="truncate text-xs text-muted-foreground">{f.summary}</span>
<Badge variant="outline">流量证据 {f.traffic_count ?? 0} 条</Badge>
</div>
</TableCell>
<TableCell className="w-52">
@@ -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 (
<div className="flex items-start justify-between gap-3 py-2.5">
<span className="shrink-0 pt-0.5 text-xs text-muted-foreground">
{label}
</span>
<div className="flex min-w-0 flex-col items-end gap-1 text-right text-sm">
{children}
</div>
<span className="shrink-0 pt-0.5 text-xs text-muted-foreground">{label}</span>
<div className="flex min-w-0 flex-col items-end gap-1 text-right text-sm">{children}</div>
</div>
);
}
@@ -123,9 +112,7 @@ function FindingDetailInner() {
if (!finding) {
return (
<div className="flex flex-1 flex-col items-center justify-center gap-3 p-10 text-center">
<p className="text-muted-foreground">
{loaded ? `未找到发现 ${id}` : "加载中…"}
</p>
<p className="text-muted-foreground">{loaded ? `未找到发现 ${id}` : "加载中…"}</p>
{loaded && (
<Button asChild variant="outline">
<Link href="/function/findings">
@@ -140,11 +127,7 @@ function FindingDetailInner() {
const title = finding.name || finding.vulnclass || "未分类";
return (
<Tabs
value={tab}
onValueChange={setTab}
className="flex flex-1 flex-col gap-0"
>
<Tabs value={tab} onValueChange={setTab} className="flex flex-1 flex-col gap-0">
{/* Sticky header */}
<header className="sticky top-0 z-10 flex flex-col gap-2 border-b bg-background/95 px-4 py-2.5 backdrop-blur lg:px-6">
<div className="flex flex-wrap items-center gap-2">
@@ -158,16 +141,12 @@ function FindingDetailInner() {
<h1 className="max-w-md truncate text-sm font-semibold" title={title}>
{title}
</h1>
<code className="rounded bg-muted px-1.5 py-0.5 font-mono text-xs text-muted-foreground">
#{finding.id}
</code>
<code className="rounded bg-muted px-1.5 py-0.5 font-mono text-xs text-muted-foreground">#{finding.id}</code>
<Separator orientation="vertical" className="mx-1 h-4" />
<StatusBadge domain="severity" value={finding.severity} dot />
<StatusBadge domain="finding" value={finding.status} dot />
{finding.inherited && finding.source_task_id && (
<Badge variant="outline">
来源任务 #{finding.source_task_id} · 只读
</Badge>
<Badge variant="outline">来源任务 #{finding.source_task_id} · 只读</Badge>
)}
</div>
<TabsList>
@@ -188,9 +167,7 @@ function FindingDetailInner() {
<CardTitle className="text-sm">摘要</CardTitle>
</CardHeader>
<CardContent>
<p className="text-sm leading-relaxed whitespace-pre-wrap">
{finding.summary || "(无摘要)"}
</p>
<p className="text-sm leading-relaxed whitespace-pre-wrap">{finding.summary || "(无摘要)"}</p>
</CardContent>
</Card>
<Card>
@@ -207,21 +184,29 @@ function FindingDetailInner() {
)}
</CardContent>
</Card>
<FindingTrafficPanel
key={id}
findingId={id}
contextTask={contextTaskId || undefined}
readOnly={finding.inherited}
onChanged={load}
/>
{/* 证据下方:详细报告(Markdown 渲染) */}
<Card>
<CardHeader className="flex-row items-center justify-between">
<CardTitle className="text-sm">详细报告</CardTitle>
{finding.report && (
<CopyButton text={finding.report} successMessage="已复制详细报告" />
)}
{finding.report && <CopyButton text={finding.report} successMessage="已复制详细报告" />}
</CardHeader>
<CardContent>
{finding.report_stale ? (
<Alert>
<AlertDescription>流量证据已变更,详细报告待更新。</AlertDescription>
</Alert>
) : null}
{finding.report ? (
<Markdown text={finding.report} />
) : (
<p className="text-sm text-muted-foreground">
暂无详细报告。
</p>
<p className="text-sm text-muted-foreground">暂无详细报告。</p>
)}
</CardContent>
</Card>
@@ -245,14 +230,8 @@ function FindingDetailInner() {
{finding.inherited ? (
<StatusBadge domain="severity" value={finding.severity} dot />
) : (
<Select
value={finding.severity}
onValueChange={(v) => changeSeverity(v as Severity)}
>
<SelectTrigger
size="sm"
className="h-7 w-auto border-none px-1 shadow-none focus-visible:ring-0"
>
<Select value={finding.severity} onValueChange={(v) => changeSeverity(v as Severity)}>
<SelectTrigger size="sm" className="h-7 w-auto border-none px-1 shadow-none focus-visible:ring-0">
<StatusBadge domain="severity" value={finding.severity} dot />
</SelectTrigger>
<SelectContent position="popper" align="end">
@@ -273,14 +252,8 @@ function FindingDetailInner() {
{finding.inherited ? (
<StatusBadge domain="finding" value={finding.status} dot />
) : (
<Select
value={finding.status}
onValueChange={(v) => changeStatus(v as FindingStatus)}
>
<SelectTrigger
size="sm"
className="h-7 w-auto border-none px-1 shadow-none focus-visible:ring-0"
>
<Select value={finding.status} onValueChange={(v) => changeStatus(v as FindingStatus)}>
<SelectTrigger size="sm" className="h-7 w-auto border-none px-1 shadow-none focus-visible:ring-0">
<StatusBadge domain="finding" value={finding.status} dot />
</SelectTrigger>
<SelectContent position="popper" align="end">
@@ -299,9 +272,7 @@ function FindingDetailInner() {
{/* 漏洞类型 */}
<FieldRow label="漏洞类型">
{finding.vulnclass ? (
<code className="rounded bg-muted px-1.5 py-0.5 font-mono text-xs">
{finding.vulnclass}
</code>
<code className="rounded bg-muted px-1.5 py-0.5 font-mono text-xs">{finding.vulnclass}</code>
) : (
<span className="text-muted-foreground">—</span>
)}
@@ -334,9 +305,7 @@ function FindingDetailInner() {
className="inline-flex max-w-[16rem] items-center gap-1 text-primary hover:underline"
title={finding.task_description}
>
<span className="truncate">
{finding.task_description || `#${finding.task_id}`}
</span>
<span className="truncate">{finding.task_description || `#${finding.task_id}`}</span>
<ArrowUpRightIcon className="size-3 shrink-0" />
</Link>
) : (
@@ -10,13 +10,7 @@ import { toast } from "sonner";
import { StatusBadge } from "@/components/status-badge";
import { Badge } from "@/components/ui/badge";
import { Card, CardContent } from "@/components/ui/card";
import {
Select,
SelectContent,
SelectGroup,
SelectItem,
SelectTrigger,
} from "@/components/ui/select";
import { Select, SelectContent, SelectGroup, SelectItem, SelectTrigger } from "@/components/ui/select";
import { api } from "@/lib/api";
import { useStoredSortPreference } from "@/lib/sort-preference";
import { statusMeta } from "@/lib/status";
@@ -64,10 +58,7 @@ function Row({
className="flex min-w-0 flex-1 items-center gap-3 text-left"
>
<ChevronRightIcon
className={cn(
"size-4 shrink-0 text-muted-foreground transition-transform",
open && "rotate-90",
)}
className={cn("size-4 shrink-0 text-muted-foreground transition-transform", open && "rotate-90")}
/>
<StatusBadge domain="severity" value={f.severity} dot />
<div className="flex min-w-0 flex-1 flex-col gap-1">
@@ -79,11 +70,10 @@ function Row({
</Badge>
)}
</div>
<span className="truncate text-xs text-muted-foreground">
{f.summary}
</span>
<span className="truncate text-xs text-muted-foreground">{f.summary}</span>
</div>
</button>
<Badge variant="outline">流量证据 {f.traffic_count ?? 0} 条</Badge>
{f.assets && f.assets.length > 0 && (
<div className="hidden shrink-0 flex-wrap justify-end gap-1 sm:flex">
{f.assets.slice(0, 2).map((a) => (
@@ -95,22 +85,12 @@ function Row({
{a.label}
</code>
))}
{f.assets.length > 2 && (
<span className="text-xs text-muted-foreground">
+{f.assets.length - 2}
</span>
)}
{f.assets.length > 2 && <span className="text-xs text-muted-foreground">+{f.assets.length - 2}</span>}
</div>
)}
{f.finding_id && !f.inherited ? (
<Select
value={f.status}
onValueChange={(v) => onStatus(f, v as FindingStatus)}
>
<SelectTrigger
size="sm"
className="h-7 w-28 shrink-0 border-none px-1 shadow-none focus-visible:ring-0"
>
<Select value={f.status} onValueChange={(v) => onStatus(f, v as FindingStatus)}>
<SelectTrigger size="sm" className="h-7 w-28 shrink-0 border-none px-1 shadow-none focus-visible:ring-0">
<StatusBadge domain="finding" value={f.status} dot />
</SelectTrigger>
<SelectContent position="popper" align="end">
@@ -146,9 +126,7 @@ function Row({
</div>
{open && (
<div className="bg-muted/30 px-4 pb-4 pl-11">
<div className="mb-1 text-xs font-medium text-muted-foreground">
证据 / PoC
</div>
<div className="mb-1 text-xs font-medium text-muted-foreground">证据 / PoC</div>
<pre className="overflow-auto rounded-md border bg-background p-3 font-mono text-xs whitespace-pre-wrap">
{f.evidence}
</pre>
@@ -187,25 +165,18 @@ export function FindingsTab({ taskId }: { taskId: string }) {
};
}, [taskId]);
const onStatus = React.useCallback(
async (f: Finding, next: FindingStatus) => {
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)),
);
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)),
);
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 }) {
<Row key={f.id} f={f} contextTaskId={taskId} onStatus={onStatus} />
))}
{items.length === 0 && (
<p className="px-4 py-8 text-center text-sm text-muted-foreground">
本任务及直接关联任务暂无确认发现。
</p>
<p className="px-4 py-8 text-center text-sm text-muted-foreground">本任务及直接关联任务暂无确认发现。</p>
)}
</CardContent>
</Card>
+55 -204
View File
@@ -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 (
<>
<span className="text-muted-foreground">{response[1]}</span>
{response[2]}
<span className={statusTone(Number(response[3]))}>{response[3]}</span>
{response[4]}
</>
);
}
const request = /^([A-Z]+)(\s+)(\S+)(\s+)(HTTP\/\S+)$/.exec(line);
if (request) {
return (
<>
<span className="font-semibold text-primary">{request[1]}</span>
{request[2]}
<span className="text-chart-2">{request[3]}</span>
{request[4]}
<span className="text-muted-foreground">{request[5]}</span>
</>
);
}
return line;
}
function HeaderLine({ line }: { line: string }) {
const separator = line.indexOf(":");
if (separator <= 0) return line;
return (
<>
<span className="text-primary">{line.slice(0, separator)}</span>
<span className="text-muted-foreground">:</span>
<span className="text-chart-2">{line.slice(separator + 1)}</span>
</>
);
}
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(
<span key={`${index}-${match[0].length}`} className={className}>
{match[0]}
</span>,
);
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][^>]*>|<!--[\s\S]*?-->/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(
<span key={`${index}-${match[0].length}`} className="text-primary">
{match[0]}
</span>,
);
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 <JsonBody body={body} />;
if (format === "markup") return <MarkupBody body={body} />;
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 <StartLine line={line} />;
if (separator < 0 || index < separator) return <HeaderLine line={line} />;
if (index === separator) return null;
return <HighlightedBody body={line} format={bodyFormat} />;
};
return (
<div className="relative mx-3 my-4 overflow-hidden rounded-md border bg-background shadow-xs">
<div className="absolute top-2 right-2 flex items-center gap-0.5 rounded-md border bg-background/95 p-0.5 shadow-xs">
<Tooltip>
<TooltipTrigger asChild>
<Button
type="button"
variant="ghost"
size="icon-xs"
aria-label={wrapLines ? "关闭自动换行" : "开启自动换行"}
aria-pressed={wrapLines}
onClick={() => setWrapLines((current) => !current)}
>
<WrapTextIcon />
</Button>
</TooltipTrigger>
<TooltipContent side="bottom">{wrapLines ? "关闭自动换行" : "开启自动换行"}</TooltipContent>
</Tooltip>
<Tooltip>
<TooltipTrigger asChild>
<Button
type="button"
variant="ghost"
size="icon-xs"
aria-label={copied ? "已复制报文" : "复制报文"}
onClick={() => void copyPacket()}
>
{copied ? <CheckIcon /> : <CopyIcon />}
</Button>
</TooltipTrigger>
<TooltipContent side="bottom">{copied ? "已复制" : "复制报文"}</TooltipContent>
</Tooltip>
</div>
{/* biome-ignore lint/a11y/useSemanticElements: textarea cannot preserve line numbers and syntax-highlighting markup. */}
<div
role="textbox"
aria-label="HTTP 报文代码"
aria-multiline="true"
aria-readonly="true"
tabIndex={0}
className="max-h-[calc(100vh-15rem)] min-w-0 overflow-auto bg-background py-3 font-mono text-xs outline-none focus-visible:ring-2 focus-visible:ring-inset focus-visible:ring-ring"
onKeyDown={(event) => {
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) => (
<div
key={`${index}-${line}`}
data-line={index + 1}
className="grid min-w-0 grid-cols-[1.5rem_minmax(0,1fr)] leading-relaxed before:sticky before:left-0 before:self-stretch before:border-r before:bg-muted/20 before:px-1 before:text-right before:text-muted-foreground/60 before:content-[attr(data-line)]"
>
<code
className={cn(
"min-h-[1lh] min-w-0 pr-16 pl-1.5 [tab-size:4]",
wrapLines ? "break-words whitespace-pre-wrap" : "whitespace-pre",
)}
>
{renderLine(line, index)}
</code>
</div>
))}
</div>
</div>
);
}
// 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<Set<string>>(() => 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() {
</div>
</div>
<div className="flex flex-wrap items-center gap-2">
<span className="text-sm text-muted-foreground">已选 {selectedFlows.size} 条流量</span>
<Button variant="outline" size="sm" disabled={selectedFlows.size === 0} onClick={() => setLinking(true)}>
关联到漏洞
</Button>
{selectedFlows.size > 0 ? (
<Button variant="ghost" size="sm" onClick={() => setSelectedFlows(new Set())}>
清空选择
</Button>
) : null}
</div>
{/* History table */}
<div className="flex h-[calc(100vh-15rem)] min-h-0 flex-col">
<Card className="flex min-h-0 flex-1 flex-col overflow-hidden py-0">
@@ -630,6 +443,22 @@ export default function TrafficPage() {
<Table>
<TableHeader className="sticky top-0 z-10 bg-card">
<TableRow>
<TableHead className="w-10">
<Checkbox
aria-label="选择本页流量"
checked={exchanges.length > 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;
})
}
/>
</TableHead>
<TableHead className="w-36">时间</TableHead>
<TableHead className="w-44">host</TableHead>
<TableHead className="w-20">方法</TableHead>
@@ -646,6 +475,21 @@ export default function TrafficPage() {
className={cn("cursor-pointer", selected?.id === e.id && "bg-accent hover:bg-accent")}
onClick={() => setSelected(e)}
>
<TableCell>
<Checkbox
aria-label={`选择流量 ${e.id}`}
checked={selectedFlows.has(e.id)}
onClick={(event) => 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;
})
}
/>
</TableCell>
<TableCell className="text-xs text-muted-foreground tabular-nums">{fmtTime(e.ts)}</TableCell>
<TableCell className="font-mono text-xs">{e.host}</TableCell>
<TableCell>
@@ -665,7 +509,7 @@ export default function TrafficPage() {
))}
{exchanges.length === 0 && (
<TableRow>
<TableCell colSpan={7} className="py-12 text-center text-sm text-muted-foreground">
<TableCell colSpan={8} className="py-12 text-center text-sm text-muted-foreground">
{traffic === null ? "加载中…" : "没有匹配的流量。"}
</TableCell>
</TableRow>
@@ -676,6 +520,13 @@ export default function TrafficPage() {
</Card>
</div>
{linking ? (
<LinkTrafficDialog
trafficIds={[...selectedFlows]}
onClose={() => setLinking(false)}
onBound={() => setSelectedFlows(new Set())}
/>
) : null}
<Sheet open={selected !== null} onOpenChange={(open) => !open && setSelected(null)}>
<SheetContent className="w-full! max-w-none! gap-0 p-0 sm:w-[48rem]! sm:max-w-[48rem]!">
{selected && (
@@ -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<TrafficEvidenceRole, string> = {
baseline: "正常对照",
proof: "漏洞证明",
verification: "补充验证",
supporting: "辅助证据",
};
export function FindingTrafficPanel({
findingId,
contextTask,
readOnly,
onChanged,
}: {
findingId: string;
contextTask?: string;
readOnly?: boolean;
onChanged: () => void;
}) {
const [data, setData] = React.useState<FindingTraffic | null>(null);
const [error, setError] = React.useState("");
const [adding, setAdding] = React.useState(false);
const [preview, setPreview] = React.useState<string | null>(null);
const [editing, setEditing] = React.useState<FindingTrafficBinding | null>(null);
const [role, setRole] = React.useState<TrafficEvidenceRole>("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<FindingTraffic>) {
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 (
<>
<Card>
<CardHeader>
<div className="flex flex-wrap items-center justify-between gap-2">
<CardTitle>关联流量 {data ? `(${data.bindings.length})` : ""}</CardTitle>
{!readOnly ? (
<Button variant="outline" size="sm" disabled={!data || busy} onClick={() => setAdding(true)}>
<PlusIcon data-icon="inline-start" />
绑定流量
</Button>
) : (
<Badge variant="outline">继承证据 · 只读</Badge>
)}
</div>
<CardDescription>按复现顺序组织请求与响应。清理原始流量后,已绑定证据仍然保留。</CardDescription>
</CardHeader>
<CardContent className="flex flex-col gap-3">
{error ? (
<Alert variant="destructive">
<AlertDescription>
{error}
<Button variant="link" onClick={() => setReload((n) => n + 1)}>
重试
</Button>
</AlertDescription>
</Alert>
) : null}
{!data ? (
<Skeleton className="h-28 w-full" />
) : data.bindings.length === 0 ? (
<Empty>
<EmptyHeader>
<EmptyTitle>暂无关联流量</EmptyTitle>
<EmptyDescription>可绑定正常对照、漏洞证明和补充验证请求。</EmptyDescription>
</EmptyHeader>
</Empty>
) : (
data.bindings.map((b, index) => (
<div key={b.id} className="flex min-w-0 flex-col gap-2 rounded-md border p-3">
<div className="flex flex-wrap items-center gap-2">
<Badge variant="outline">
{index + 1}. {ROLES[b.role]}
</Badge>
<Badge variant="secondary">
{b.snapshot.method} · {b.snapshot.status}
</Badge>
<span className="text-xs text-muted-foreground">证据 #{b.id}</span>
</div>
<Button
variant="link"
className="h-auto justify-start p-0 text-left whitespace-normal"
onClick={() => setPreview(b.id)}
>
<span className="break-all font-mono text-xs">{b.snapshot.url}</span>
</Button>
{b.note ? <p className="text-sm whitespace-pre-wrap">{b.note}</p> : null}
<div className="flex flex-wrap items-center justify-between gap-2">
<span className="text-xs text-muted-foreground">
{new Date(b.snapshot.captured_at * 1000).toLocaleString("zh-CN")}
</span>
<div className="flex flex-wrap gap-1">
<Button variant="outline" size="sm" onClick={() => setPreview(b.id)}>
查看报文
</Button>
{!readOnly ? (
<>
<Button
variant="ghost"
size="sm"
disabled={busy}
onClick={() => {
setEditing(b);
setRole(b.role);
setNote(b.note);
}}
>
编辑说明
</Button>
<Button
variant="ghost"
size="icon-sm"
aria-label={`上移证据 ${b.id}`}
disabled={busy || index === 0}
onClick={() => move(index, -1)}
>
<ArrowUpIcon />
</Button>
<Button
variant="ghost"
size="icon-sm"
aria-label={`下移证据 ${b.id}`}
disabled={busy || index === data.bindings.length - 1}
onClick={() => move(index, 1)}
>
<ArrowDownIcon />
</Button>
<Button
variant="ghost"
size="sm"
disabled={busy}
onClick={() =>
void mutate(() => api.removeFindingTraffic(findingId, b.id, data.version, contextTask))
}
>
解除绑定
</Button>
</>
) : null}
</div>
</div>
</div>
))
)}
</CardContent>
</Card>
{adding && data ? (
<TrafficPickerDialog
findingId={findingId}
contextTask={contextTask}
bound={data.bindings}
onClose={() => setAdding(false)}
onBound={() => {
setReload((n) => n + 1);
onChanged();
}}
/>
) : null}
<TrafficEvidenceViewer
findingId={findingId}
bindingId={preview}
contextTask={contextTask}
onClose={() => setPreview(null)}
/>
<Dialog
open={editing !== null}
onOpenChange={(open) => {
if (!open && !busy) setEditing(null);
}}
>
<DialogContent>
<DialogHeader>
<DialogTitle>编辑流量证据</DialogTitle>
<DialogDescription>说明这组请求/响应如何支持漏洞结论。</DialogDescription>
</DialogHeader>
<FieldGroup>
<Field>
<FieldLabel htmlFor="evidence-role">用途</FieldLabel>
<Select value={role} onValueChange={(v) => setRole(v as TrafficEvidenceRole)}>
<SelectTrigger id="evidence-role">
<SelectValue />
</SelectTrigger>
<SelectContent>
<SelectGroup>
{Object.entries(ROLES).map(([value, label]) => (
<SelectItem key={value} value={value}>
{label}
</SelectItem>
))}
</SelectGroup>
</SelectContent>
</Select>
</Field>
<Field>
<FieldLabel htmlFor="evidence-note">证据说明</FieldLabel>
<Textarea id="evidence-note" value={note} onChange={(e) => setNote(e.target.value)} />
</Field>
</FieldGroup>
<DialogFooter>
<Button variant="outline" disabled={busy} onClick={() => setEditing(null)}>
取消
</Button>
<Button
disabled={busy || !editing || !data}
onClick={() => {
if (editing && data)
void mutate(() =>
api.editFindingTraffic(findingId, editing.id, data.version, { role, note }, contextTask),
);
}}
>
保存说明
</Button>
</DialogFooter>
</DialogContent>
</Dialog>
</>
);
}
+216
View File
@@ -0,0 +1,216 @@
"use client";
import * as React from "react";
import { CheckIcon, CopyIcon, WrapTextIcon } from "lucide-react";
import { toast } from "sonner";
import { Button } from "@/components/ui/button";
import { Tooltip, TooltipContent, TooltipTrigger } from "@/components/ui/tooltip";
import { cn, copyText } from "@/lib/utils";
function statusTone(status: number) {
if (status >= 500) return "text-red-500";
if (status >= 400) return "text-amber-500";
if (status >= 300) return "text-blue-500";
if (status >= 200) return "text-emerald-500";
return "text-muted-foreground";
}
function StartLine({ line }: { line: string }) {
const response = /^(HTTP\/\S+)(\s+)(\d{3})(.*)$/.exec(line);
if (response) {
return (
<>
<span className="text-muted-foreground">{response[1]}</span>
{response[2]}
<span className={statusTone(Number(response[3]))}>{response[3]}</span>
{response[4]}
</>
);
}
const request = /^([A-Z]+)(\s+)(\S+)(\s+)(HTTP\/\S+)$/.exec(line);
if (request) {
return (
<>
<span className="font-semibold text-primary">{request[1]}</span>
{request[2]}
<span className="text-chart-2">{request[3]}</span>
{request[4]}
<span className="text-muted-foreground">{request[5]}</span>
</>
);
}
return line;
}
function HeaderLine({ line }: { line: string }) {
const separator = line.indexOf(":");
if (separator <= 0) return line;
return (
<>
<span className="text-primary">{line.slice(0, separator)}</span>
<span className="text-muted-foreground">:</span>
<span className="text-chart-2">{line.slice(separator + 1)}</span>
</>
);
}
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(
<span key={`${index}-${match[0].length}`} className={className}>
{match[0]}
</span>,
);
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][^>]*>|<!--[\s\S]*?-->/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(
<span key={`${index}-${match[0].length}`} className="text-primary">
{match[0]}
</span>,
);
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 <JsonBody body={body} />;
if (format === "markup") return <MarkupBody body={body} />;
return body;
}
export 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 <StartLine line={line} />;
if (separator < 0 || index < separator) return <HeaderLine line={line} />;
if (index === separator) return null;
return <HighlightedBody body={line} format={bodyFormat} />;
};
return (
<div className="relative mx-3 my-4 overflow-hidden rounded-md border bg-background shadow-xs">
<div className="absolute top-2 right-2 flex items-center gap-0.5 rounded-md border bg-background/95 p-0.5 shadow-xs">
<Tooltip>
<TooltipTrigger asChild>
<Button
type="button"
variant="ghost"
size="icon-xs"
aria-label={wrapLines ? "关闭自动换行" : "开启自动换行"}
aria-pressed={wrapLines}
onClick={() => setWrapLines((current) => !current)}
>
<WrapTextIcon />
</Button>
</TooltipTrigger>
<TooltipContent side="bottom">{wrapLines ? "关闭自动换行" : "开启自动换行"}</TooltipContent>
</Tooltip>
<Tooltip>
<TooltipTrigger asChild>
<Button
type="button"
variant="ghost"
size="icon-xs"
aria-label={copied ? "已复制报文" : "复制报文"}
onClick={() => void copyPacket()}
>
{copied ? <CheckIcon /> : <CopyIcon />}
</Button>
</TooltipTrigger>
<TooltipContent side="bottom">{copied ? "已复制" : "复制报文"}</TooltipContent>
</Tooltip>
</div>
{/* biome-ignore lint/a11y/useSemanticElements: textarea cannot preserve line numbers and syntax-highlighting markup. */}
<div
role="textbox"
aria-label="HTTP 报文代码"
aria-multiline="true"
aria-readonly="true"
tabIndex={0}
className="max-h-[calc(100vh-15rem)] min-w-0 overflow-auto bg-background py-3 font-mono text-xs outline-none focus-visible:ring-2 focus-visible:ring-inset focus-visible:ring-ring"
onKeyDown={(event) => {
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) => (
<div
key={`${index}-${line}`}
data-line={index + 1}
className="grid min-w-0 grid-cols-[1.5rem_minmax(0,1fr)] leading-relaxed before:sticky before:left-0 before:self-stretch before:border-r before:bg-muted/20 before:px-1 before:text-right before:text-muted-foreground/60 before:content-[attr(data-line)]"
>
<code
className={cn(
"min-h-[1lh] min-w-0 pr-16 pl-1.5 [tab-size:4]",
wrapLines ? "break-words whitespace-pre-wrap" : "whitespace-pre",
)}
>
{renderLine(line, index)}
</code>
</div>
))}
</div>
</div>
);
}
+163
View File
@@ -0,0 +1,163 @@
"use client";
import * as React from "react";
import { toast } from "sonner";
import { Alert, AlertDescription } from "@/components/ui/alert";
import { Button } from "@/components/ui/button";
import {
Dialog,
DialogContent,
DialogDescription,
DialogFooter,
DialogHeader,
DialogTitle,
} from "@/components/ui/dialog";
import { Field, FieldGroup, FieldLabel } from "@/components/ui/field";
import { Input } from "@/components/ui/input";
import { api } from "@/lib/api";
import type { Finding, FindingsPage } from "@/lib/types";
export function LinkTrafficDialog({
trafficIds,
onClose,
onBound,
}: {
trafficIds: string[];
onClose: () => void;
onBound: () => void;
}) {
const [query, setQuery] = React.useState("");
const [page, setPage] = React.useState(1);
const [data, setData] = React.useState<FindingsPage | null>(null);
const [selected, setSelected] = React.useState<Finding | null>(null);
const [error, setError] = React.useState("");
const [busy, setBusy] = React.useState(false);
const [loading, setLoading] = React.useState(true);
React.useEffect(() => {
let active = true;
setLoading(true);
setError("");
const timer = setTimeout(() => {
api
.findingsPage({ page, pageSize: 20, query })
.then((d) => {
if (active) setData(d);
})
.catch((e: Error) => {
if (active) setError(e.message);
})
.finally(() => {
if (active) setLoading(false);
});
}, 250);
return () => {
active = false;
clearTimeout(timer);
};
}, [query, page]);
async function save() {
if (!selected?.finding_id) return;
setBusy(true);
setError("");
try {
await api.bindFindingTraffic(
selected.finding_id,
trafficIds.map((traffic_id) => ({ traffic_id })),
);
toast.success(`已关联 ${trafficIds.length} 条流量到漏洞 #${selected.finding_id}`);
onBound();
onClose();
} catch (e) {
setError((e as Error).message);
} finally {
setBusy(false);
}
}
return (
<Dialog
open
onOpenChange={(open) => {
if (!open && !busy) onClose();
}}
>
<DialogContent className="max-h-[90vh] overflow-y-auto sm:max-w-2xl">
<DialogHeader>
<DialogTitle>关联到漏洞</DialogTitle>
<DialogDescription>将所选 {trafficIds.length} 条流量保存为已有漏洞的证据。</DialogDescription>
</DialogHeader>
<FieldGroup>
<Field>
<FieldLabel htmlFor="link-finding-query">查找漏洞</FieldLabel>
<Input
id="link-finding-query"
placeholder="名称 / 摘要 / 漏洞类别"
value={query}
onChange={(e) => {
setQuery(e.target.value);
setPage(1);
}}
/>
</Field>
</FieldGroup>
{error ? (
<Alert variant="destructive">
<AlertDescription>{error}</AlertDescription>
</Alert>
) : null}
<div className="flex max-h-[40vh] flex-col gap-2 overflow-auto" aria-busy={loading}>
{(data?.items ?? []).map((f) => (
<Button
key={f.finding_id ?? f.id}
variant={selected?.finding_id === f.finding_id ? "secondary" : "outline"}
className="h-auto justify-start p-3 text-left whitespace-normal"
aria-pressed={selected?.finding_id === f.finding_id}
disabled={busy || !f.finding_id || f.inherited}
onClick={() => setSelected(f)}
>
<span className="flex min-w-0 flex-col gap-1">
<span>
#{f.finding_id ?? f.id} · {f.name || f.vulnclass}
</span>
<span className="line-clamp-2 text-xs text-muted-foreground">{f.summary}</span>
</span>
</Button>
))}
{!data?.items.length ? (
<p className="py-6 text-center text-sm text-muted-foreground">
{loading ? "加载中…" : "没有匹配的漏洞,请先登记漏洞"}
</p>
) : null}
</div>
<div className="flex items-center justify-between gap-2">
<Button variant="outline" size="sm" disabled={loading || page <= 1} onClick={() => setPage((p) => p - 1)}>
上一页
</Button>
<span className="text-xs">
第 {page} 页 · 共 {data?.total ?? 0} 条
</span>
<Button
variant="outline"
size="sm"
disabled={loading || page * 20 >= (data?.total ?? 0)}
onClick={() => setPage((p) => p + 1)}
>
下一页
</Button>
</div>
<p className="text-sm">
{selected ? `已选漏洞:#${selected.finding_id} ${selected.name || selected.vulnclass}` : "请选择一个漏洞"}
</p>
<DialogFooter>
<Button variant="outline" disabled={busy} onClick={onClose}>
取消
</Button>
<Button disabled={busy || !selected} onClick={() => void save()}>
{busy ? "保存中…" : "确认关联"}
</Button>
</DialogFooter>
</DialogContent>
</Dialog>
);
}
@@ -0,0 +1,181 @@
"use client";
import * as React from "react";
import { toast } from "sonner";
import { HttpCodeBlock } from "@/components/http-code-block";
import { Alert, AlertDescription } from "@/components/ui/alert";
import { Button } from "@/components/ui/button";
import { Dialog, DialogContent, DialogDescription, DialogHeader, DialogTitle } from "@/components/ui/dialog";
import { Skeleton } from "@/components/ui/skeleton";
import { Tabs, TabsContent, TabsList, TabsTrigger } from "@/components/ui/tabs";
import { api } from "@/lib/api";
import type { FindingTrafficDetail, TrafficDetail } from "@/lib/types";
export function TrafficEvidenceViewer({
findingId,
bindingId,
contextTask,
onClose,
}: {
findingId: string;
bindingId: string | null;
contextTask?: string;
onClose: () => void;
}) {
const [detail, setDetail] = React.useState<FindingTrafficDetail | null>(null);
const [error, setError] = React.useState("");
const [busy, setBusy] = React.useState(false);
React.useEffect(() => {
if (!bindingId) return;
let active = true;
setDetail(null);
setError("");
api
.findingTrafficDetail(findingId, bindingId, contextTask)
.then((d) => {
if (active) setDetail(d);
})
.catch((e: Error) => {
if (active) setError(e.message);
});
return () => {
active = false;
};
}, [findingId, bindingId, contextTask]);
async function more(side: "request" | "response") {
if (!detail || !bindingId) return;
setBusy(true);
try {
const next = await api.findingTrafficBody(findingId, bindingId, side, detail[side].next_offset, contextTask);
setDetail((current) =>
current?.binding.id === bindingId
? { ...current, [side]: { ...next, content: current[side].content + next.content } }
: current,
);
} catch (e) {
toast.error((e as Error).message);
} finally {
setBusy(false);
}
}
return (
<Dialog
open={bindingId !== null}
onOpenChange={(open) => {
if (!open) onClose();
}}
>
<DialogContent className="max-h-[90vh] overflow-y-auto sm:max-w-4xl">
<DialogHeader>
<DialogTitle>流量证据 #{bindingId}</DialogTitle>
<DialogDescription className="break-all">
{detail?.binding.snapshot.url ?? "查看绑定时保存的请求与响应"}
</DialogDescription>
</DialogHeader>
{error ? (
<Alert variant="destructive">
<AlertDescription>{error}</AlertDescription>
</Alert>
) : detail ? (
<Tabs defaultValue="request">
<TabsList>
<TabsTrigger value="request">请求 Request</TabsTrigger>
<TabsTrigger value="response">响应 Response</TabsTrigger>
</TabsList>
{(["request", "response"] as const).map((side) => (
<TabsContent key={side} value={side}>
<div className="flex flex-wrap items-center justify-between gap-2">
<span className="text-xs text-muted-foreground">
正文 {detail[side].total.toLocaleString()} 字节{detail[side].truncated ? " · 当前为预览" : ""}
</span>
<Button
variant="outline"
size="sm"
onClick={() =>
void api
.downloadFindingTrafficBody(findingId, detail.binding.id, side, contextTask)
.catch((e: Error) => toast.error(e.message))
}
>
下载完整{side === "request" ? "请求" : "响应"}正文
</Button>
</div>
<HttpCodeBlock
raw={`${(side === "request" ? detail.binding.snapshot.req_head : detail.binding.snapshot.resp_head)?.trimEnd() ?? ""}\n\n${detail[side].content}`}
/>
{detail[side].truncated && !detail[side].binary ? (
<Button variant="outline" size="sm" disabled={busy} onClick={() => void more(side)}>
加载更多正文
</Button>
) : null}
</TabsContent>
))}
</Tabs>
) : (
<Skeleton className="h-56 w-full" />
)}
</DialogContent>
</Dialog>
);
}
export function CapturedTrafficViewer({ id, onClose }: { id: string | null; onClose: () => void }) {
const [detail, setDetail] = React.useState<TrafficDetail | null>(null);
const [error, setError] = React.useState("");
React.useEffect(() => {
if (!id) return;
let active = true;
setDetail(null);
setError("");
api
.trafficExchange(id)
.then((d) => {
if (active) setDetail(d);
})
.catch((e: Error) => {
if (active) setError(e.message);
});
return () => {
active = false;
};
}, [id]);
return (
<Dialog
open={id !== null}
onOpenChange={(open) => {
if (!open) onClose();
}}
>
<DialogContent className="max-h-[90vh] overflow-y-auto sm:max-w-4xl">
<DialogHeader>
<DialogTitle>预览流量</DialogTitle>
<DialogDescription>流量 ID:{id}。绑定时会保存完整正文。</DialogDescription>
</DialogHeader>
{error ? (
<Alert variant="destructive">
<AlertDescription>{error}</AlertDescription>
</Alert>
) : detail ? (
<Tabs defaultValue="request">
<TabsList>
<TabsTrigger value="request">请求 Request</TabsTrigger>
<TabsTrigger value="response">响应 Response</TabsTrigger>
</TabsList>
<TabsContent value="request">
<HttpCodeBlock raw={detail.req} />
</TabsContent>
<TabsContent value="response">
<HttpCodeBlock raw={detail.resp} />
</TabsContent>
</Tabs>
) : (
<Skeleton className="h-56 w-full" />
)}
</DialogContent>
</Dialog>
);
}
@@ -0,0 +1,277 @@
"use client";
import * as React from "react";
import { toast } from "sonner";
import { CapturedTrafficViewer } from "@/components/traffic-evidence-viewer";
import { Alert, AlertDescription } from "@/components/ui/alert";
import { Badge } from "@/components/ui/badge";
import { Button } from "@/components/ui/button";
import { Checkbox } from "@/components/ui/checkbox";
import {
Dialog,
DialogContent,
DialogDescription,
DialogFooter,
DialogHeader,
DialogTitle,
} from "@/components/ui/dialog";
import { Field, FieldGroup, FieldLabel } from "@/components/ui/field";
import { Input } from "@/components/ui/input";
import { Select, SelectContent, SelectGroup, SelectItem, SelectTrigger, SelectValue } from "@/components/ui/select";
import { Table, TableBody, TableCell, TableHead, TableHeader, TableRow } from "@/components/ui/table";
import { api } from "@/lib/api";
import type { FindingTrafficBinding, TrafficResp } from "@/lib/types";
const METHODS = ["GET", "POST", "PUT", "DELETE", "PATCH", "HEAD", "OPTIONS"];
export function TrafficPickerDialog({
findingId,
contextTask,
bound,
onClose,
onBound,
}: {
findingId: string;
contextTask?: string;
bound: FindingTrafficBinding[];
onClose: () => void;
onBound: () => void;
}) {
const [host, setHost] = React.useState("");
const [method, setMethod] = React.useState("all");
const [query, setQuery] = React.useState("");
const [page, setPage] = React.useState(0);
const [data, setData] = React.useState<TrafficResp | null>(null);
const [selected, setSelected] = React.useState<Set<string>>(() => new Set());
const [preview, setPreview] = React.useState<string | null>(null);
const [loading, setLoading] = React.useState(true);
const [busy, setBusy] = React.useState(false);
const [error, setError] = React.useState("");
const alreadyBound = new Set(bound.map((b) => b.snapshot.source_traffic_id));
React.useEffect(() => {
let active = true;
setLoading(true);
setError("");
const timer = setTimeout(() => {
api
.traffic(page, 25, host, method === "all" ? "" : method, query)
.then((d) => {
if (active) setData(d);
})
.catch((e: Error) => {
if (active) setError(e.message);
})
.finally(() => {
if (active) setLoading(false);
});
}, 250);
return () => {
active = false;
clearTimeout(timer);
};
}, [page, host, method, query]);
const rows = data?.exchanges ?? [];
const selectable = rows.filter((e) => !alreadyBound.has(e.id));
function toggle(id: string, checked: boolean) {
setSelected((previous) => {
const next = new Set(previous);
if (checked) next.add(id);
else next.delete(id);
return next;
});
}
async function save() {
setBusy(true);
setError("");
try {
await api.bindFindingTraffic(
findingId,
[...selected].map((traffic_id) => ({ traffic_id })),
contextTask,
);
toast.success(`已绑定 ${selected.size} 条流量`);
onBound();
onClose();
} catch (e) {
setError((e as Error).message);
} finally {
setBusy(false);
}
}
return (
<>
<Dialog
open
onOpenChange={(open) => {
if (!open && !busy) onClose();
}}
>
<DialogContent className="max-h-[90vh] overflow-y-auto sm:max-w-5xl">
<DialogHeader>
<DialogTitle>绑定流量</DialogTitle>
<DialogDescription>
筛选并多选请求/响应,已选记录会跨页保留。绑定后可设置用途、说明和顺序。
</DialogDescription>
</DialogHeader>
<FieldGroup className="flex flex-col gap-3 sm:flex-row">
<Field>
<FieldLabel htmlFor="evidence-host">目标 host</FieldLabel>
<Input
id="evidence-host"
value={host}
placeholder="域名或 IP"
onChange={(e) => {
setHost(e.target.value);
setPage(0);
}}
/>
</Field>
<Field>
<FieldLabel htmlFor="evidence-method">请求方法</FieldLabel>
<Select
value={method}
onValueChange={(v) => {
setMethod(v);
setPage(0);
}}
>
<SelectTrigger id="evidence-method">
<SelectValue />
</SelectTrigger>
<SelectContent>
<SelectGroup>
<SelectItem value="all">全部方法</SelectItem>
{METHODS.map((m) => (
<SelectItem key={m} value={m}>
{m}
</SelectItem>
))}
</SelectGroup>
</SelectContent>
</Select>
</Field>
<Field>
<FieldLabel htmlFor="evidence-query">关键词</FieldLabel>
<Input
id="evidence-query"
value={query}
placeholder="URL / 正文关键词"
onChange={(e) => {
setQuery(e.target.value);
setPage(0);
}}
/>
</Field>
</FieldGroup>
{error ? (
<Alert variant="destructive">
<AlertDescription>{error}</AlertDescription>
</Alert>
) : null}
<div className="max-h-[45vh] overflow-auto rounded-md border" aria-busy={loading}>
<Table>
<TableHeader>
<TableRow>
<TableHead>
<Checkbox
aria-label="选择本页未绑定流量"
disabled={loading || busy || selectable.length === 0}
checked={selectable.length > 0 && selectable.every((e) => selected.has(e.id))}
onCheckedChange={(checked) =>
setSelected((previous) => {
const next = new Set(previous);
for (const e of selectable) {
if (checked === true) next.add(e.id);
else next.delete(e.id);
}
return next;
})
}
/>
</TableHead>
<TableHead>方法 / URL</TableHead>
<TableHead>时间</TableHead>
<TableHead>状态码</TableHead>
<TableHead>操作</TableHead>
</TableRow>
</TableHeader>
<TableBody>
{rows.map((e) => (
<TableRow key={e.id}>
<TableCell>
<Checkbox
aria-label={`选择流量 ${e.id}`}
checked={selected.has(e.id) || alreadyBound.has(e.id)}
disabled={busy || loading || alreadyBound.has(e.id)}
onCheckedChange={(checked) => toggle(e.id, checked === true)}
/>
</TableCell>
<TableCell>
<span className="font-mono text-xs">
{e.method} {e.url}
</span>
{alreadyBound.has(e.id) ? <Badge variant="secondary">已绑定</Badge> : null}
</TableCell>
<TableCell className="whitespace-nowrap text-xs">
{new Date(e.ts).toLocaleString("zh-CN")}
</TableCell>
<TableCell>{e.status}</TableCell>
<TableCell>
<Button variant="ghost" size="sm" onClick={() => setPreview(e.id)}>
预览
</Button>
</TableCell>
</TableRow>
))}
{!rows.length ? (
<TableRow>
<TableCell colSpan={5} className="py-8 text-center">
{loading ? "加载中…" : "没有匹配的流量"}
</TableCell>
</TableRow>
) : null}
</TableBody>
</Table>
</div>
<div className="flex flex-wrap items-center justify-between gap-2">
<span className="text-sm">
已选 {selected.size} 条 · 共 {data?.total ?? 0} 条
</span>
<div className="flex items-center gap-2">
<Button
variant="outline"
size="sm"
disabled={loading || page === 0}
onClick={() => setPage((p) => p - 1)}
>
上一页
</Button>
<span className="text-xs">第 {page + 1} 页</span>
<Button
variant="outline"
size="sm"
disabled={loading || (page + 1) * 25 >= (data?.total ?? 0)}
onClick={() => setPage((p) => p + 1)}
>
下一页
</Button>
</div>
</div>
<DialogFooter>
<Button variant="outline" disabled={busy} onClick={onClose}>
取消
</Button>
<Button disabled={busy || selected.size === 0} onClick={() => void save()}>
{busy ? "保存中…" : `绑定 ${selected.size} 条流量`}
</Button>
</DialogFooter>
</DialogContent>
</Dialog>
<CapturedTrafficViewer id={preview} onClose={() => setPreview(null)} />
</>
);
}
+72
View File
@@ -28,6 +28,7 @@ import type {
DeleteTaskOptions,
DeleteTaskResult,
Edge,
EvidenceBodyPreview,
Finding,
FindingAssetTree,
FindingDeepenResponse,
@@ -36,6 +37,8 @@ import type {
FindingStats,
FindingStatus,
FindingsPage,
FindingTraffic,
FindingTrafficDetail,
IntentAsset,
InterceptApprovalRow,
InterceptPending,
@@ -79,6 +82,8 @@ import type {
Tool,
ToolStat,
TrafficDetail,
TrafficEvidenceRef,
TrafficEvidenceRole,
TrafficHost,
TrafficResp,
UsageStats,
@@ -526,6 +531,73 @@ export const api = {
get<Finding>(
`/exploration/findings/${id}${contextTaskId ? `?context_task=${encodeURIComponent(contextTaskId)}` : ""}`,
),
findingTraffic: (id: string, contextTask?: string) =>
get<FindingTraffic>(
`/exploration/findings/${id}/traffic${contextTask ? `?context_task=${encodeURIComponent(contextTask)}` : ""}`,
),
bindFindingTraffic: (id: string, traffic_refs: TrafficEvidenceRef[], contextTask?: string) =>
post<FindingTraffic>(
`/exploration/findings/${id}/traffic${contextTask ? `?context_task=${encodeURIComponent(contextTask)}` : ""}`,
{ traffic_refs },
),
editFindingTraffic: (
id: string,
bindingId: string,
version: number,
fields: { role: TrafficEvidenceRole; note: string },
contextTask?: string,
) =>
patch<FindingTraffic>(
`/exploration/findings/${id}/traffic/${bindingId}${contextTask ? `?context_task=${encodeURIComponent(contextTask)}` : ""}`,
{ version, ...fields },
),
removeFindingTraffic: (id: string, bindingId: string, version: number, contextTask?: string) =>
del<FindingTraffic>(
`/exploration/findings/${id}/traffic/${bindingId}${contextTask ? `?context_task=${encodeURIComponent(contextTask)}` : ""}`,
{ version },
),
orderFindingTraffic: (id: string, binding_ids: string[], version: number, contextTask?: string) =>
http<FindingTraffic>(
`/exploration/findings/${id}/traffic/order${contextTask ? `?context_task=${encodeURIComponent(contextTask)}` : ""}`,
{ method: "PUT", body: JSON.stringify({ binding_ids, version }) },
),
findingTrafficDetail: (id: string, bindingId: string, contextTask?: string) =>
get<FindingTrafficDetail>(
`/exploration/findings/${id}/traffic/${bindingId}${contextTask ? `?context_task=${encodeURIComponent(contextTask)}` : ""}`,
),
findingTrafficBody: (
id: string,
bindingId: string,
side: "request" | "response",
offset: number,
contextTask?: string,
) =>
get<EvidenceBodyPreview>(
`/exploration/findings/${id}/traffic/${bindingId}/body?side=${side}&offset=${offset}${contextTask ? `&context_task=${encodeURIComponent(contextTask)}` : ""}`,
),
downloadFindingTrafficBody: async (
id: string,
bindingId: string,
side: "request" | "response",
contextTask?: string,
) => {
const token = getToken();
const response = await fetch(
`/api/exploration/findings/${id}/traffic/${bindingId}/body?side=${side}&download=1${contextTask ? `&context_task=${encodeURIComponent(contextTask)}` : ""}`,
{ headers: token ? { Authorization: `Bearer ${token}` } : {} },
);
if (!response.ok) {
const error = await response.json().catch(() => ({ error: "下载失败" }));
throw new Error(error.error ?? "下载失败");
}
const blob = await response.blob();
const url = URL.createObjectURL(blob);
const a = document.createElement("a");
a.href = url;
a.download = `evidence-${bindingId}-${side}.bin`;
a.click();
setTimeout(() => URL.revokeObjectURL(url), 1000);
},
// 漏洞链路:该漏洞节点回溯到任务初始节点的子图(节点 + 关系)。
findingLineage: (id: string) => get<{ nodes: TaskNode[]; edges: Edge[] }>(`/exploration/findings/${id}/lineage`),
setFindingStatus: (id: string, status: FindingStatus) => patch<Finding>(`/exploration/findings/${id}`, { status }),
+55
View File
@@ -462,6 +462,10 @@ export interface FindingAsset {
}
export interface Finding {
traffic_count?: number;
evidence_version?: number;
report_evidence_version?: number;
report_stale?: boolean;
id: string;
finding_id?: string; // 独立 findings 表的行 id,状态更新的句柄(任务内旧节点可能缺失)
vulnclass: string;
@@ -1246,3 +1250,54 @@ export interface LLMTask {
task_id: string;
count: number;
}
export type TrafficEvidenceRole = "baseline" | "proof" | "verification" | "supporting";
export interface TrafficEvidenceRef {
traffic_id: string;
role?: TrafficEvidenceRole;
note?: string;
}
export interface TrafficEvidenceSnapshot {
id: string;
source_traffic_id: string;
captured_at: number;
url: string;
method: string;
status: number;
content_type: string;
req_head?: string;
resp_head?: string;
req_hash: string;
resp_hash: string;
req_len: number;
resp_len: number;
}
export interface FindingTrafficBinding {
id: string;
finding_id: string;
snapshot_id: string;
role: TrafficEvidenceRole;
note: string;
position: number;
created_at: string;
snapshot: TrafficEvidenceSnapshot;
}
export interface FindingTraffic {
finding_id: string;
version: number;
report_version: number;
bindings: FindingTrafficBinding[];
}
export interface EvidenceBodyPreview {
content: string;
offset: number;
total: number;
next_offset: number;
truncated: boolean;
binary: boolean;
}
export interface FindingTrafficDetail {
binding: FindingTrafficBinding;
request: EvidenceBodyPreview;
response: EvidenceBodyPreview;
}