mirror of
https://github.com/RuoJi6/ARTEX.git
synced 2026-10-27 17:03:49 +08:00
feat: 支持漏洞关联多条流量证据
This commit is contained in:
ji ruo committed
1 parent
8dae851b9b
commit
ec4ef35fff
41 files changed
+4390
-496
No files matched your search
@@ -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 }
|
||||
@@ -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")
|
||||
}
|
||||
}
|
||||
+16
-14
@@ -17,19 +17,20 @@ import (
|
||||
// (→planner) or direct high-priority intents (→frontier). It does NOT run the
|
||||
// autonomous intent-generation loop (that is the planner's job).
|
||||
type MainAgent struct {
|
||||
prov llm.Provider
|
||||
model string
|
||||
tx *transcript.Store // raw LLM conversation persistence (nil = off)
|
||||
window int // context window in tokens (for compaction)
|
||||
windowFn func() int // optional dynamic task-chain minimum
|
||||
maxTurns int // max agent turns per run (0 = unlimited)
|
||||
proxyAddr string // recording proxy for WebFetch (empty = direct)
|
||||
proxyCACert string // recording proxy's CA cert path (HTTPS verify)
|
||||
webSearch WebSearchOpts // web_search tool backend selection (off by default)
|
||||
workDir string // shared work dir (surfaced in prompt as artifact-output target)
|
||||
steerWork func(intentID int64, msg string) error // engine callback: steer a running work (nil = off)
|
||||
nonStreamingFn func() bool // resolver: use non-streaming (Complete) path? (nil = streaming)
|
||||
maxTokensFn func() int // resolver: per-reply output cap (nil/0 = send no cap)
|
||||
findingRecorder FindingRecorder
|
||||
prov llm.Provider
|
||||
model string
|
||||
tx *transcript.Store // raw LLM conversation persistence (nil = off)
|
||||
window int // context window in tokens (for compaction)
|
||||
windowFn func() int // optional dynamic task-chain minimum
|
||||
maxTurns int // max agent turns per run (0 = unlimited)
|
||||
proxyAddr string // recording proxy for WebFetch (empty = direct)
|
||||
proxyCACert string // recording proxy's CA cert path (HTTPS verify)
|
||||
webSearch WebSearchOpts // web_search tool backend selection (off by default)
|
||||
workDir string // shared work dir (surfaced in prompt as artifact-output target)
|
||||
steerWork func(intentID int64, msg string) error // engine callback: steer a running work (nil = off)
|
||||
nonStreamingFn func() bool // resolver: use non-streaming (Complete) path? (nil = streaming)
|
||||
maxTokensFn func() int // resolver: per-reply output cap (nil/0 = send no cap)
|
||||
}
|
||||
|
||||
// SetNonStreaming wires a resolver deciding whether runs use the non-streaming
|
||||
@@ -103,6 +104,7 @@ func mainAgentSystem(goal, dataDir, workDir string) string {
|
||||
// worker/planner sessions — not just the final answer.
|
||||
func (m *MainAgent) Chat(ctx context.Context, taskID int64, as *db.AssetStore, ts *db.ExplorationStore, goal, message string, emit func(db.Activity), notify, resume func(), notifyGoal func([]string)) (string, error) {
|
||||
tsx := NewToolSet(ts, "human")
|
||||
tsx.SetFindingRecorder(m.findingRecorder)
|
||||
if as != nil {
|
||||
tsx.SetAssetStore(as, as.Companies())
|
||||
}
|
||||
@@ -149,7 +151,7 @@ func (m *MainAgent) Chat(ctx context.Context, taskID int64, as *db.AssetStore, t
|
||||
// 命中预算(步数)→ SDK 跑收尾:向用户输出一句进展总结。Prompt 与收尾轮数可后台编辑(默认 10 轮)。
|
||||
Settlement: wrapupSettlement("mainagent", nil),
|
||||
NonStreaming: m.nonStreaming(), // 该 profile 选非流式时走 Provider.Complete
|
||||
MaxTokens: m.maxTokens(), // 0 = 不发上限,由服务端默认值决定
|
||||
MaxTokens: m.maxTokens(), // 0 = 不发上限,由服务端默认值决定
|
||||
}
|
||||
if m.tx != nil { // persist raw human↔AI conversation; one accumulating file per task
|
||||
opts.Transcript = m.tx
|
||||
|
||||
@@ -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())
|
||||
}
|
||||
|
||||
@@ -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
@@ -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，按需裁剪，但证据/复现/修复必须有）━━
|
||||
- ` + "`## 概述`" + `：一句话说清是什么漏洞、在哪、能造成什么。
|
||||
|
||||
+58
-65
@@ -73,11 +73,12 @@ func compactIntents(ns []*db.Node, parentsOf, yieldsOf map[int64][]int64) []map[
|
||||
// ToolSet exposes the PG-backed dual graph (asset + exploration) to an LLM agent.
|
||||
// One ToolSet is created per planner/worker run; per-run signals live here.
|
||||
type ToolSet struct {
|
||||
as *db.AssetStore // asset store (optional; nil = asset tools not available)
|
||||
cs *db.CompanyStore // company store (optional)
|
||||
ts *db.ExplorationStore
|
||||
worker string
|
||||
taskID int64 // PG tasks.id; 0 when unknown (tests / orchestrator cross-task reads)
|
||||
findingRecorder FindingRecorder
|
||||
as *db.AssetStore // asset store (optional; nil = asset tools not available)
|
||||
cs *db.CompanyStore // company store (optional)
|
||||
ts *db.ExplorationStore
|
||||
worker string
|
||||
taskID int64 // PG tasks.id; 0 when unknown (tests / orchestrator cross-task reads)
|
||||
// coverageDisabled mirrors tasks.coverage_enabled=false. Stored inverted so the
|
||||
// zero value (all existing ToolSet constructions) means ENABLED — matching the
|
||||
// DB default (true). When true: graphOverviewData drops the coverage block, the
|
||||
@@ -1358,67 +1359,59 @@ func (t *ToolSet) goalMet() actool.CoreTool {
|
||||
// --- worker write tools ---
|
||||
|
||||
func (t *ToolSet) addFinding() actool.CoreTool {
|
||||
return writeTool("report_finding", "[重要]发现漏洞时必须调用该工具进行记录!记录一个确认的漏洞发现。在任务上下文中 intent_id 必填（当前正在执行的意图 id）；在会话上下文中 intent_id 可不填。",
|
||||
obj(map[string]any{
|
||||
"vulnclass": str("漏洞类（分类，如 SQL Injection / IDOR / XSS）"),
|
||||
"name": str("漏洞名称（具体可读的标题，如『用户中心订单接口存在越权访问』；建议填写，留空时前端回退展示 vulnclass）"),
|
||||
"severity": str("critical|high|medium|low（严重/高/中/低）"),
|
||||
"summary": str("发现摘要"),
|
||||
"intent_id": idp("产生本发现的意图 id（任务上下文必填；会话上下文可不填）"),
|
||||
"asset_ids": map[string]any{"type": "array", "items": map[string]any{"type": "integer"}, "description": "【存在时尽量填写，否则在摘要中必须要写清楚漏洞位置】受影响资产 id（可选，0/1/多个）：参数/端点/站点等。一个漏洞影响多处可全填，纯观察可不填。"},
|
||||
"evidence": str("证据/PoC 文本"),
|
||||
}, "vulnclass", "severity", "summary"),
|
||||
func(_ context.Context, in json.RawMessage) (actool.Result, error) {
|
||||
var a struct {
|
||||
VulnClass, Name, Severity, Summary, Evidence string
|
||||
IntentID json.RawMessage `json:"intent_id"`
|
||||
AssetIDs []json.RawMessage `json:"asset_ids"`
|
||||
return writeTool("report_finding", "记录确认的漏洞。可选 traffic_refs 绑定多条真实 HTTP 请求/响应：先用 traffic_search / traffic_get 核对 ID；提交的流量必须全部绑定成功。TCP 等非 HTTP 漏洞、未采集或无确切匹配记录时可省略 traffic_refs 或传空数组，仍用 evidence 提供命令输出、日志等可验证证据，建议说明未绑定原因。任务上下文传当前 intent_id。", obj(map[string]any{
|
||||
"vulnclass": str("漏洞类别"), "name": str("漏洞名称"), "severity": str("critical|high|medium|low"), "summary": str("发现摘要"),
|
||||
"intent_id": idp("当前任务的意图 id"), "asset_ids": map[string]any{"type": "array", "items": map[string]any{"type": "integer"}, "description": "受影响资产 id"},
|
||||
"evidence": str("证据/PoC 文本"),
|
||||
"traffic_refs": map[string]any{"type": "array", "description": "可选；HTTP/HTTPS 漏洞先检索并逐条核实请求/响应确实支持漏洞结论，再按复现顺序填写真实 ID。TCP 等非 HTTP 漏洞、未采集或找不到确切记录时省略或传 []，不阻止上报；可在 evidence 说明原因并提供其他可验证证据。不要猜测 ID、按域名/时间推定关联或仅为补包重复探测。用途 baseline 正常对照 / proof 漏洞证明 / verification 补充验证 / supporting 辅助证据。",
|
||||
"items": obj(map[string]any{"traffic_id": str("traffic_search 返回的真实流量 ID"), "role": map[string]any{"type": "string", "enum": []string{"baseline", "proof", "verification", "supporting"}}, "note": str("该流量如何支持漏洞结论")}, "traffic_id")},
|
||||
}, "vulnclass", "severity", "summary"), func(ctx context.Context, in json.RawMessage) (actool.Result, error) {
|
||||
var a struct {
|
||||
VulnClass, Name, Severity, Summary, Evidence string
|
||||
IntentID json.RawMessage `json:"intent_id"`
|
||||
AssetIDs []json.RawMessage `json:"asset_ids"`
|
||||
TrafficRefs []db.TrafficRef `json:"traffic_refs"`
|
||||
}
|
||||
if err := json.Unmarshal(in, &a); err != nil {
|
||||
return actool.Errorf(err.Error()), nil
|
||||
}
|
||||
if t.ts == nil {
|
||||
return actool.Errorf("report_finding 需要任务上下文（exploration store 未初始化）"), nil
|
||||
}
|
||||
refs, err := db.NormalizeTrafficRefs(a.TrafficRefs)
|
||||
if err != nil {
|
||||
return actool.Errorf(err.Error()), nil
|
||||
}
|
||||
input := db.RecordFindingInput{TaskID: t.taskID, ExplorationID: t.ts.ID(), IntentID: pid(a.IntentID), VulnClass: a.VulnClass, Name: a.Name, Severity: a.Severity, Summary: a.Summary, Evidence: a.Evidence, Worker: t.worker, AssetIDs: pidList(a.AssetIDs)}
|
||||
var recorded *db.RecordedFinding
|
||||
if t.findingRecorder != nil {
|
||||
recorded, err = t.findingRecorder.Record(ctx, input, refs)
|
||||
} else if len(refs) > 0 {
|
||||
return actool.Errorf("流量证据存储不可用；未登记漏洞"), nil
|
||||
} else {
|
||||
recorded, err = t.ts.RecordFinding(ctx, input)
|
||||
}
|
||||
if err != nil {
|
||||
return actool.Errorf(err.Error()), nil
|
||||
}
|
||||
if t.notifyFinding != nil {
|
||||
iid := input.IntentID
|
||||
if iid <= 0 {
|
||||
iid = t.ownerNode
|
||||
}
|
||||
_ = json.Unmarshal(in, &a)
|
||||
payload := map[string]any{"vulnclass": a.VulnClass, "name": a.Name, "severity": a.Severity, "summary": a.Summary,
|
||||
"evidence": map[string]any{"by": t.worker, "poc": a.Evidence}}
|
||||
var anchors []int64
|
||||
for _, raw := range a.AssetIDs {
|
||||
if p := pid(raw); p > 0 {
|
||||
anchors = append(anchors, p)
|
||||
}
|
||||
}
|
||||
var id int64
|
||||
if t.ts != nil {
|
||||
intent := pid(a.IntentID)
|
||||
if intent > 0 {
|
||||
node, err := t.ts.GetNode(intent)
|
||||
if err != nil || node == nil || node.Kind != db.KindIntent {
|
||||
return actool.Errorf("intent_id 必须是本任务的意图（关联任务意图只读）"), nil
|
||||
}
|
||||
}
|
||||
var err error
|
||||
id, err = t.ts.AddNode(db.KindFinding, payload, 9, "confirmed", t.worker, anchors)
|
||||
if err != nil {
|
||||
return actool.Errorf(err.Error()), nil
|
||||
}
|
||||
if intent > 0 {
|
||||
_ = t.ts.Link(intent, db.RelYields, id) // chain: intent -> finding
|
||||
}
|
||||
_, _ = t.ts.AddStandaloneFinding(t.taskID, id, a.VulnClass, a.Name, a.Severity, a.Summary, a.Evidence, t.worker, anchors)
|
||||
// 确证漏洞落库 → 当场唤醒本任务 planner（不等 worker 收工，debounce 合并）。
|
||||
// 优先带上下文(哪个意图+finding摘要);intent 用工具参数,缺省回退到 owner 意图。
|
||||
if t.notifyFinding != nil {
|
||||
iid := pid(a.IntentID)
|
||||
if iid <= 0 {
|
||||
iid = t.ownerNode
|
||||
}
|
||||
t.notifyFinding(iid, a.Summary)
|
||||
} else if t.notify != nil {
|
||||
t.notify()
|
||||
}
|
||||
} else {
|
||||
// conversation context: no exploration store available, cannot record finding
|
||||
return actool.Errorf("report_finding 需要任务上下文（exploration store 未初始化）"), nil
|
||||
}
|
||||
t.writes.Findings++
|
||||
return actool.Text(fmt.Sprintf("finding recorded: %d", id)), nil
|
||||
})
|
||||
t.notifyFinding(iid, a.Summary)
|
||||
} else if t.notify != nil {
|
||||
t.notify()
|
||||
}
|
||||
t.writes.Findings++
|
||||
// Keep the first line's node-ID contract for existing reporter triggers.
|
||||
for i := range recorded.Traffic.Bindings {
|
||||
recorded.Traffic.Bindings[i].Snapshot.ReqHead = ""
|
||||
recorded.Traffic.Bindings[i].Snapshot.RespHead = ""
|
||||
}
|
||||
raw, _ := json.Marshal(recorded)
|
||||
return actool.Text(fmt.Sprintf("finding recorded: %d\n%s", recorded.NodeID, raw)), nil
|
||||
})
|
||||
}
|
||||
|
||||
// recordFact writes a general exploration RESULT/conclusion (not a vuln, not a
|
||||
|
||||
+17
-13
@@ -41,17 +41,18 @@ type WebSearchOpts struct {
|
||||
}
|
||||
|
||||
type Worker struct {
|
||||
prov llm.Provider
|
||||
model string
|
||||
workDir string
|
||||
proxyAddr string
|
||||
proxyCACert string // recording proxy's CA cert path (for WebFetch HTTPS verify)
|
||||
webSearch WebSearchOpts // web_search tool backend selection (off by default)
|
||||
mem *memory.Store // cross-engagement tradecraft memory (G4)
|
||||
tx *transcript.Store // raw LLM conversation persistence (nil = off)
|
||||
window int // context window in tokens (for compaction)
|
||||
windowFn func() int // optional dynamic task-chain minimum
|
||||
maxTurns int // max agent turns per run (0 = unlimited)
|
||||
findingRecorder FindingRecorder
|
||||
prov llm.Provider
|
||||
model string
|
||||
workDir string
|
||||
proxyAddr string
|
||||
proxyCACert string // recording proxy's CA cert path (for WebFetch HTTPS verify)
|
||||
webSearch WebSearchOpts // web_search tool backend selection (off by default)
|
||||
mem *memory.Store // cross-engagement tradecraft memory (G4)
|
||||
tx *transcript.Store // raw LLM conversation persistence (nil = off)
|
||||
window int // context window in tokens (for compaction)
|
||||
windowFn func() int // optional dynamic task-chain minimum
|
||||
maxTurns int // max agent turns per run (0 = unlimited)
|
||||
// runTimeout is the wall-clock budget for the main exploration of one intent
|
||||
// (0 = unlimited). When it fires, the run is cut and a settlement round is
|
||||
// forced so already-identified facts get written back instead of being lost.
|
||||
@@ -204,7 +205,9 @@ const workerDefaultTmpl = `你是一个 ARTEX 平台授权渗透测试系统的"
|
||||
**边发现边写回**（写进图才算数，脑子/文字里的不算；每得一个结果立刻写，别攒到最后被步数耗尽丢掉）。三种写回，别串图：
|
||||
- **新资产/资源 → insert_assets（资产图）**：子域 / service / endpoint / 指纹 / 凭据 等一切资产【本身】。**这里只登记资产；探索结论/判断不写这里，用 record_fact。**
|
||||
- **探索结论/事实 → record_fact（探索图，传 intent_id）**：都用它。**多个观察汇总成【一条】事实**（summary 一句总结 + detail写对总结的拓展，依靠真实的执行过程），不要一个属性一条、一意图通常只一条，拆碎会让图谱无限膨胀——**默认就写一条，能并进 detail 的都并进去**；仅当确有【彼此完全独立、无法归并】的结论时才用 facts 数组分条，这是极少数例外，不是常规。**只写增量**：只记这次【新得到】的，别把已有事实换措辞重记（只印证已有、无新增就不必记）。**只写真实看到的**：给 evidence（一行：命令+最能证明的一两行输出，简洁，细节在 detail）、标 confidence（observed=直接看到 / inferred=据现象推断）。
|
||||
- **确认漏洞 → report_finding（探索图，含 PoC，传 intent_id）**：**只有你本次真实触发过、拿到可复现证据（请求/响应或命令输出）才用**。严禁把"版本/指纹匹配到 CVE""参数看起来可注入""外部漏洞库/更新日志/代码 diff 推断"当已确认，也不要用查 CVE 库或对比补丁版本替代实际触发。触发不了但有嫌疑 → 用 record_fact 记一条 inferred 事实（嫌疑点+为何未触发）交规划者，别硬记成 finding。
|
||||
- **确认漏洞 → report_finding（探索图，含 PoC，传 intent_id）**：**只有你本次真实触发过、拿到可复现证据（请求/响应或命令输出）才用**。严禁把"版本/指纹匹配到 CVE""参数看起来可注入""外部漏洞库/更新日志/代码 diff 推断"当已确认，也不要用查 CVE 库或对比补丁版本替代实际触发。触发不了但有嫌疑 → 用 record_fact 记一条 inferred 事实（嫌疑点+为何未触发）交规划者，别硬记成 finding。有对应录制流量时，先用 traffic_search / traffic_get 核对真实 ID，再用 traffic_refs 按复现顺序绑定多条请求/响应；域名和时间只用于候选筛选，不推定任务归属。
|
||||
|
||||
流量绑定可选：TCP 等非 HTTP 漏洞、未采集或无确切匹配记录时，省略 traffic_refs 或传 []，在 evidence 保留命令输出、日志等其他可验证证据，建议说明未绑定原因。不要猜测 ID，也不要仅为补包重复探测。
|
||||
|
||||
完成本意图后用一句话总结你做了什么、写回了哪些事实。务实、克制、聚焦这一条意图。`
|
||||
|
||||
@@ -255,7 +258,7 @@ func workerSystem(proxyAddr, caCert, dataDir, runDir string) string {
|
||||
body := renderSystem("worker", workerDefaultTmpl, WorkerVars{ProxyAddr: proxyAddr, DataDir: dataDir, Now: nowStr()})
|
||||
// caCert is present only when the recording MITM is on, which is exactly when
|
||||
// the traffic_* tools are registered — so it gates the traffic-tool note.
|
||||
return body + workerTrafficBlock(caCert != "") + workerArtifactSpec(runDir)
|
||||
return body + workerTrafficBlock(caCert != "") + findingTrafficGuidance + workerArtifactSpec(runDir)
|
||||
}
|
||||
|
||||
// renderIntentTask formats the claimed intent for the worker's launch USER message:
|
||||
@@ -325,6 +328,7 @@ func (w *Worker) ExecuteWithMessage(ctx context.Context, name string, taskID int
|
||||
|
||||
func (w *Worker) execute(ctx context.Context, name string, taskID int64, as *db.AssetStore, ts *db.ExplorationStore, intent *db.Node, hooks harness.HookRunner, emit func(db.Activity), enr EnrichTrigger, notifyFinding func(int64, string), requestID, message string) (harness.TerminalReason, WriteCounts, error) {
|
||||
tsx := NewToolSet(ts, name)
|
||||
tsx.SetFindingRecorder(w.findingRecorder)
|
||||
tsx.SetTaskID(taskID)
|
||||
coverageEnabled := as == nil || as.CoverageEnabled(taskID)
|
||||
tsx.SetCoverageEnabled(coverageEnabled)
|
||||
|
||||
@@ -0,0 +1,404 @@
|
||||
package db
|
||||
|
||||
import (
|
||||
"context"
|
||||
"crypto/sha256"
|
||||
"database/sql"
|
||||
"encoding/hex"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"strings"
|
||||
"time"
|
||||
)
|
||||
|
||||
var (
|
||||
ErrEvidenceConflict = errors.New("流量证据已变更，请刷新后重试")
|
||||
ErrFindingNotFound = errors.New("漏洞不存在")
|
||||
ErrEvidenceNotFound = errors.New("流量证据不存在")
|
||||
)
|
||||
|
||||
// This lock covers the evidence filesystem as well as its SQL references. All
|
||||
// processes sharing the database use it, including readers, exports and GC.
|
||||
const findingEvidenceLockKey int64 = 7337741004
|
||||
|
||||
func (d *DB) WithEvidenceTx(ctx context.Context, fn func(*sql.Tx) error) error {
|
||||
tx, err := d.BeginTx(ctx, nil)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer tx.Rollback()
|
||||
if _, err = tx.ExecContext(ctx, `SELECT pg_advisory_xact_lock($1)`, findingEvidenceLockKey); err != nil {
|
||||
return err
|
||||
}
|
||||
if err = fn(tx); err != nil {
|
||||
return err
|
||||
}
|
||||
return tx.Commit()
|
||||
}
|
||||
|
||||
type TrafficRef struct {
|
||||
TrafficID string `json:"traffic_id"`
|
||||
Role string `json:"role"`
|
||||
Note string `json:"note"`
|
||||
}
|
||||
|
||||
func NormalizeTrafficRefs(refs []TrafficRef) ([]TrafficRef, error) {
|
||||
out := make([]TrafficRef, 0, len(refs))
|
||||
seen := map[string]bool{}
|
||||
for _, ref := range refs {
|
||||
ref.TrafficID = strings.TrimSpace(ref.TrafficID)
|
||||
if ref.TrafficID == "" {
|
||||
return nil, errors.New("traffic_id 不能为空")
|
||||
}
|
||||
if ref.Role == "" {
|
||||
ref.Role = "supporting"
|
||||
}
|
||||
if !ValidTrafficRole(ref.Role) {
|
||||
return nil, fmt.Errorf("无效的流量用途 %q", ref.Role)
|
||||
}
|
||||
if !seen[ref.TrafficID] {
|
||||
out = append(out, ref)
|
||||
seen[ref.TrafficID] = true
|
||||
}
|
||||
}
|
||||
return out, nil
|
||||
}
|
||||
|
||||
func ValidTrafficRole(role string) bool {
|
||||
return role == "baseline" || role == "proof" || role == "verification" || role == "supporting"
|
||||
}
|
||||
|
||||
type TrafficEvidenceSnapshot struct {
|
||||
ID string `json:"id"`
|
||||
SourceTrafficID string `json:"source_traffic_id"`
|
||||
CapturedAt int64 `json:"captured_at"`
|
||||
URL string `json:"url"`
|
||||
Method string `json:"method"`
|
||||
Status int `json:"status"`
|
||||
ContentType string `json:"content_type"`
|
||||
ReqHead string `json:"req_head,omitempty"`
|
||||
RespHead string `json:"resp_head,omitempty"`
|
||||
ReqHash string `json:"req_hash"`
|
||||
RespHash string `json:"resp_hash"`
|
||||
ReqLen int64 `json:"req_len"`
|
||||
RespLen int64 `json:"resp_len"`
|
||||
CreatedAt time.Time `json:"created_at"`
|
||||
}
|
||||
|
||||
type FindingTrafficBinding struct {
|
||||
ID int64 `json:"id,string"`
|
||||
FindingID int64 `json:"finding_id,string"`
|
||||
SnapshotID string `json:"snapshot_id"`
|
||||
Role string `json:"role"`
|
||||
Note string `json:"note"`
|
||||
Position int `json:"position"`
|
||||
CreatedAt time.Time `json:"created_at"`
|
||||
Snapshot TrafficEvidenceSnapshot `json:"snapshot"`
|
||||
}
|
||||
|
||||
type FindingTraffic struct {
|
||||
FindingID int64 `json:"finding_id,string"`
|
||||
Version int64 `json:"version"`
|
||||
ReportVersion int64 `json:"report_version"`
|
||||
Bindings []FindingTrafficBinding `json:"bindings"`
|
||||
}
|
||||
|
||||
type PreparedTrafficEvidence struct {
|
||||
Ref TrafficRef
|
||||
Snapshot TrafficEvidenceSnapshot
|
||||
}
|
||||
|
||||
func TrafficSnapshotID(snapshot TrafficEvidenceSnapshot) string {
|
||||
snapshot.ID = ""
|
||||
snapshot.CreatedAt = time.Time{}
|
||||
raw, _ := json.Marshal(snapshot)
|
||||
sum := sha256.Sum256(raw)
|
||||
return hex.EncodeToString(sum[:])
|
||||
}
|
||||
|
||||
// LockTaskEvidenceTx serializes writes with archive queueing (which locks the
|
||||
// same task row). Once queued, its snapshot must not acquire new evidence.
|
||||
func LockTaskEvidenceTx(tx *sql.Tx, taskID int64) error {
|
||||
if taskID == 0 {
|
||||
return nil
|
||||
}
|
||||
var deleted sql.NullTime
|
||||
if err := tx.QueryRow(`SELECT deleted_at FROM tasks WHERE id=$1 FOR UPDATE`, taskID).Scan(&deleted); err != nil {
|
||||
return err
|
||||
}
|
||||
if deleted.Valid {
|
||||
return ErrTaskArchiveState
|
||||
}
|
||||
var state string
|
||||
err := tx.QueryRow(`SELECT state FROM task_archives WHERE task_id=$1`, taskID).Scan(&state)
|
||||
if err != nil && !errors.Is(err, sql.ErrNoRows) {
|
||||
return err
|
||||
}
|
||||
if err == nil && state != ArchiveFailed {
|
||||
return ErrTaskArchiveState
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func LockFindingEvidenceTx(tx *sql.Tx, findingID int64, version *int64) error {
|
||||
var taskID sql.NullInt64
|
||||
if err := tx.QueryRow(`SELECT task_id FROM findings WHERE id=$1`, findingID).Scan(&taskID); err != nil {
|
||||
if errors.Is(err, sql.ErrNoRows) {
|
||||
return ErrFindingNotFound
|
||||
}
|
||||
return err
|
||||
}
|
||||
if err := LockTaskEvidenceTx(tx, taskID.Int64); err != nil {
|
||||
return err
|
||||
}
|
||||
var current int64
|
||||
if err := tx.QueryRow(`SELECT evidence_version FROM findings WHERE id=$1 FOR UPDATE`, findingID).Scan(¤t); err != nil {
|
||||
if errors.Is(err, sql.ErrNoRows) {
|
||||
return ErrFindingNotFound
|
||||
}
|
||||
return err
|
||||
}
|
||||
if version != nil && *version != current {
|
||||
return ErrEvidenceConflict
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func InsertEvidenceSnapshotTx(tx *sql.Tx, s TrafficEvidenceSnapshot) error {
|
||||
if s.ID != TrafficSnapshotID(s) {
|
||||
return errors.New("证据快照元数据哈希不匹配")
|
||||
}
|
||||
if s.CreatedAt.IsZero() {
|
||||
s.CreatedAt = time.Now().UTC()
|
||||
}
|
||||
_, err := tx.Exec(`INSERT INTO traffic_evidence_snapshots
|
||||
(id,source_traffic_id,captured_at,url,method,status,content_type,req_head,resp_head,req_hash,resp_hash,req_len,resp_len,created_at)
|
||||
VALUES($1,$2,$3,$4,$5,$6,$7,$8,$9,$10,$11,$12,$13,$14) ON CONFLICT(id) DO NOTHING`,
|
||||
s.ID, s.SourceTrafficID, s.CapturedAt, s.URL, s.Method, s.Status, s.ContentType, s.ReqHead, s.RespHead, s.ReqHash, s.RespHash, s.ReqLen, s.RespLen, s.CreatedAt)
|
||||
return err
|
||||
}
|
||||
|
||||
func AddFindingTrafficTx(tx *sql.Tx, findingID int64, prepared []PreparedTrafficEvidence) error {
|
||||
var position int
|
||||
if err := tx.QueryRow(`SELECT COALESCE(MAX(position)+1,0) FROM finding_traffic_bindings WHERE finding_id=$1`, findingID).Scan(&position); err != nil {
|
||||
return err
|
||||
}
|
||||
changed := false
|
||||
for _, item := range prepared {
|
||||
if err := InsertEvidenceSnapshotTx(tx, item.Snapshot); err != nil {
|
||||
return err
|
||||
}
|
||||
if _, err := tx.Exec(`UPDATE traffic_evidence_snapshots SET unreferenced_at=NULL WHERE id=$1`, item.Snapshot.ID); err != nil {
|
||||
return err
|
||||
}
|
||||
res, err := tx.Exec(`INSERT INTO finding_traffic_bindings(finding_id,snapshot_id,role,note,position)
|
||||
VALUES($1,$2,$3,$4,$5) ON CONFLICT(finding_id,snapshot_id) DO NOTHING`, findingID, item.Snapshot.ID, item.Ref.Role, item.Ref.Note, position)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
n, err := res.RowsAffected()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if n > 0 {
|
||||
changed = true
|
||||
position++
|
||||
}
|
||||
}
|
||||
if changed {
|
||||
return bumpEvidenceVersionTx(tx, findingID)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func bumpEvidenceVersionTx(tx *sql.Tx, findingID int64) error {
|
||||
_, err := tx.Exec(`UPDATE findings SET evidence_version=evidence_version+1 WHERE id=$1`, findingID)
|
||||
return err
|
||||
}
|
||||
|
||||
func FindingTrafficTx(tx *sql.Tx, findingID int64) (*FindingTraffic, error) {
|
||||
out := &FindingTraffic{FindingID: findingID, Bindings: []FindingTrafficBinding{}}
|
||||
if err := tx.QueryRow(`SELECT evidence_version,report_evidence_version FROM findings WHERE id=$1`, findingID).Scan(&out.Version, &out.ReportVersion); err != nil {
|
||||
if errors.Is(err, sql.ErrNoRows) {
|
||||
return nil, ErrFindingNotFound
|
||||
}
|
||||
return nil, err
|
||||
}
|
||||
rows, err := tx.Query(`SELECT b.id,b.finding_id,b.snapshot_id,b.role,b.note,b.position,b.created_at,to_jsonb(s)
|
||||
FROM finding_traffic_bindings b JOIN traffic_evidence_snapshots s ON s.id=b.snapshot_id WHERE b.finding_id=$1 ORDER BY b.position,b.id`, findingID)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer rows.Close()
|
||||
for rows.Next() {
|
||||
var b FindingTrafficBinding
|
||||
var raw []byte
|
||||
if err := rows.Scan(&b.ID, &b.FindingID, &b.SnapshotID, &b.Role, &b.Note, &b.Position, &b.CreatedAt, &raw); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if err := json.Unmarshal(raw, &b.Snapshot); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
out.Bindings = append(out.Bindings, b)
|
||||
}
|
||||
return out, rows.Err()
|
||||
}
|
||||
|
||||
func (d *DB) GetFindingTraffic(ctx context.Context, findingID int64) (out *FindingTraffic, err error) {
|
||||
err = d.WithEvidenceTx(ctx, func(tx *sql.Tx) error { var e error; out, e = FindingTrafficTx(tx, findingID); return e })
|
||||
return
|
||||
}
|
||||
|
||||
func (d *DB) EditFindingTraffic(ctx context.Context, findingID, bindingID, version int64, role, note *string, remove bool, order []int64) error {
|
||||
if role != nil && !ValidTrafficRole(*role) {
|
||||
return errors.New("无效的流量用途")
|
||||
}
|
||||
return d.WithEvidenceTx(ctx, func(tx *sql.Tx) error {
|
||||
if err := LockFindingEvidenceTx(tx, findingID, &version); err != nil {
|
||||
return err
|
||||
}
|
||||
if order != nil {
|
||||
current, err := FindingTrafficTx(tx, findingID)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if len(order) != len(current.Bindings) {
|
||||
return ErrEvidenceConflict
|
||||
}
|
||||
ids := map[int64]bool{}
|
||||
for _, b := range current.Bindings {
|
||||
ids[b.ID] = true
|
||||
}
|
||||
for position, id := range order {
|
||||
if !ids[id] {
|
||||
return ErrEvidenceConflict
|
||||
}
|
||||
delete(ids, id)
|
||||
if _, err := tx.Exec(`UPDATE finding_traffic_bindings SET position=$1 WHERE finding_id=$2 AND id=$3`, position, findingID, id); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
} else {
|
||||
var res sql.Result
|
||||
var err error
|
||||
if remove {
|
||||
res, err = tx.Exec(`DELETE FROM finding_traffic_bindings WHERE finding_id=$1 AND id=$2`, findingID, bindingID)
|
||||
} else {
|
||||
res, err = tx.Exec(`UPDATE finding_traffic_bindings SET role=COALESCE($3,role),note=COALESCE($4,note) WHERE finding_id=$1 AND id=$2`, findingID, bindingID, role, note)
|
||||
}
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
n, err := res.RowsAffected()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if n == 0 {
|
||||
return ErrEvidenceNotFound
|
||||
}
|
||||
}
|
||||
return bumpEvidenceVersionTx(tx, findingID)
|
||||
})
|
||||
}
|
||||
|
||||
type RecordFindingInput struct {
|
||||
TaskID, ExplorationID, IntentID int64
|
||||
VulnClass, Name, Severity, Summary, Evidence, Worker string
|
||||
AssetIDs []int64
|
||||
}
|
||||
|
||||
type RecordedFinding struct {
|
||||
FindingID int64 `json:"finding_id,string"`
|
||||
NodeID int64 `json:"finding_node_id,string"`
|
||||
Traffic *FindingTraffic `json:"traffic"`
|
||||
}
|
||||
|
||||
func RecordFindingTx(tx *sql.Tx, in RecordFindingInput, prepared []PreparedTrafficEvidence) (*RecordedFinding, error) {
|
||||
if err := LockTaskEvidenceTx(tx, in.TaskID); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if in.TaskID > 0 {
|
||||
var expID int64
|
||||
if err := tx.QueryRow(`SELECT exploration_id FROM tasks WHERE id=$1`, in.TaskID).Scan(&expID); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if expID != in.ExplorationID {
|
||||
return nil, errors.New("漏洞所属任务与探索记录不匹配")
|
||||
}
|
||||
}
|
||||
if in.IntentID > 0 {
|
||||
var ok bool
|
||||
if err := tx.QueryRow(`SELECT EXISTS(SELECT 1 FROM exploration_nodes WHERE id=$1 AND exploration_id=$2 AND kind='intent')`, in.IntentID, in.ExplorationID).Scan(&ok); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if !ok {
|
||||
return nil, errors.New("intent_id 必须是本任务的意图（关联任务意图只读）")
|
||||
}
|
||||
}
|
||||
payload, _ := json.Marshal(map[string]any{"vulnclass": in.VulnClass, "name": in.Name, "severity": in.Severity, "summary": in.Summary, "evidence": map[string]any{"by": in.Worker, "poc": in.Evidence}})
|
||||
out := &RecordedFinding{}
|
||||
if err := tx.QueryRow(`INSERT INTO exploration_nodes(exploration_id,kind,payload,priority,state,origin)
|
||||
VALUES($1,'finding',$2,9,'confirmed',$3) RETURNING id`, in.ExplorationID, string(payload), in.Worker).Scan(&out.NodeID); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
for _, asset := range in.AssetIDs {
|
||||
if _, err := tx.Exec(`INSERT INTO exploration_anchors(node_id,asset_id) VALUES($1,$2) ON CONFLICT DO NOTHING`, out.NodeID, asset); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
}
|
||||
if in.IntentID > 0 {
|
||||
if _, err := tx.Exec(`INSERT INTO exploration_edges(exploration_id,src_id,rel,dst_id) VALUES($1,$2,$3,$4)`, in.ExplorationID, in.IntentID, RelYields, out.NodeID); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
}
|
||||
assets := in.AssetIDs
|
||||
if assets == nil {
|
||||
assets = []int64{}
|
||||
}
|
||||
raw, _ := json.Marshal(assets)
|
||||
if err := tx.QueryRow(`INSERT INTO findings(task_id,node_id,vulnclass,name,severity,summary,evidence,worker,asset_ids)
|
||||
VALUES(NULLIF($1,0),$2,$3,$4,$5,$6,$7,$8,$9) RETURNING id`, in.TaskID, out.NodeID, in.VulnClass, in.Name, in.Severity, in.Summary, in.Evidence, in.Worker, string(raw)).Scan(&out.FindingID); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if err := AddFindingTrafficTx(tx, out.FindingID, prepared); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
var err error
|
||||
out.Traffic, err = FindingTrafficTx(tx, out.FindingID)
|
||||
return out, err
|
||||
}
|
||||
|
||||
// RecordFinding is the atomic legacy/no-recorder path used by tools and tests.
|
||||
func (s *ExplorationStore) RecordFinding(ctx context.Context, in RecordFindingInput) (out *RecordedFinding, err error) {
|
||||
in.ExplorationID = s.expID
|
||||
err = s.db.WithEvidenceTx(ctx, func(tx *sql.Tx) error { var e error; out, e = RecordFindingTx(tx, in, nil); return e })
|
||||
return
|
||||
}
|
||||
|
||||
func (d *DB) FindingIDByNodeID(nodeID int64) (id int64, err error) {
|
||||
err = d.QueryRow(`SELECT id FROM findings WHERE node_id=$1`, nodeID).Scan(&id)
|
||||
return
|
||||
}
|
||||
|
||||
func (d *DB) SetFindingReportVersionByNodeID(ctx context.Context, nodeID int64, report string, version *int64) (n int64, err error) {
|
||||
err = d.WithEvidenceTx(ctx, func(tx *sql.Tx) error {
|
||||
var id int64
|
||||
if err := tx.QueryRow(`SELECT id FROM findings WHERE node_id=$1`, nodeID).Scan(&id); err != nil {
|
||||
if errors.Is(err, sql.ErrNoRows) {
|
||||
return nil
|
||||
}
|
||||
return err
|
||||
}
|
||||
if err := LockFindingEvidenceTx(tx, id, version); err != nil {
|
||||
return err
|
||||
}
|
||||
res, err := tx.Exec(`UPDATE findings SET report=$2,report_evidence_version=COALESCE($3::bigint,CASE WHEN evidence_version=0 THEN 0 ELSE -1 END) WHERE id=$1`, id, report, version)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
n, err = res.RowsAffected()
|
||||
return err
|
||||
})
|
||||
return
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
+37
-25
@@ -1,6 +1,7 @@
|
||||
package db
|
||||
|
||||
import (
|
||||
"context"
|
||||
"database/sql"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
@@ -14,6 +15,11 @@ import (
|
||||
// DBFinding is a row in the standalone findings table. It persists across task
|
||||
// deletion unless the caller explicitly requests related finding cleanup.
|
||||
type DBFinding struct {
|
||||
TrafficCount int
|
||||
EvidenceVersion int64
|
||||
ReportEvidenceVersion int64
|
||||
TrafficBindings []FindingTrafficBinding // populated only for export
|
||||
|
||||
ID int64
|
||||
TaskID *int64
|
||||
NodeID *int64
|
||||
@@ -97,7 +103,8 @@ func (d *DB) AddFinding(taskID, nodeID int64, vulnclass, name, severity, summary
|
||||
// list query selects, so scanFinding stays in sync across callers.
|
||||
const findingSelectCols = `f.id, f.task_id, f.node_id, f.vulnclass, COALESCE(f.name, ''), f.severity, f.summary,
|
||||
f.evidence, f.worker, f.asset_ids, COALESCE(f.status, 'pending'), f.created_at,
|
||||
COALESCE(t.description, '') AS task_description`
|
||||
COALESCE(t.description, '') AS task_description, f.evidence_version, f.report_evidence_version,
|
||||
(SELECT count(*) FROM finding_traffic_bindings b WHERE b.finding_id=f.id)`
|
||||
|
||||
// scanFindings materializes rows selected via findingSelectCols.
|
||||
func scanFindings(rows interface {
|
||||
@@ -110,7 +117,7 @@ func scanFindings(rows interface {
|
||||
f := &DBFinding{}
|
||||
var aidsJSON string
|
||||
if err := rows.Scan(&f.ID, &f.TaskID, &f.NodeID, &f.VulnClass, &f.Name, &f.Severity,
|
||||
&f.Summary, &f.Evidence, &f.Worker, &aidsJSON, &f.Status, &f.CreatedAt, &f.TaskDescription); err != nil {
|
||||
&f.Summary, &f.Evidence, &f.Worker, &aidsJSON, &f.Status, &f.CreatedAt, &f.TaskDescription, &f.EvidenceVersion, &f.ReportEvidenceVersion, &f.TrafficCount); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
_ = json.Unmarshal([]byte(aidsJSON), &f.AssetIDs)
|
||||
@@ -503,7 +510,7 @@ func (d *DB) ListFindingsForExport(f FindingFilter, ids []int64) ([]*DBFinding,
|
||||
var aidsJSON string
|
||||
if err := rows.Scan(&f.ID, &f.TaskID, &f.NodeID, &f.VulnClass, &f.Name, &f.Severity,
|
||||
&f.Summary, &f.Evidence, &f.Worker, &aidsJSON, &f.Status, &f.CreatedAt,
|
||||
&f.TaskDescription, &f.Report); err != nil {
|
||||
&f.TaskDescription, &f.EvidenceVersion, &f.ReportEvidenceVersion, &f.TrafficCount, &f.Report); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
_ = json.Unmarshal([]byte(aidsJSON), &f.AssetIDs)
|
||||
@@ -631,7 +638,7 @@ func (d *DB) GetFinding(id int64) (*DBFinding, error) {
|
||||
WHERE f.id = $1`, id).Scan(
|
||||
&f.ID, &f.TaskID, &f.NodeID, &f.VulnClass, &f.Name, &f.Severity,
|
||||
&f.Summary, &f.Evidence, &f.Worker, &aidsJSON, &f.Status, &f.CreatedAt,
|
||||
&f.TaskDescription, &f.Report)
|
||||
&f.TaskDescription, &f.EvidenceVersion, &f.ReportEvidenceVersion, &f.TrafficCount, &f.Report)
|
||||
if err == sql.ErrNoRows {
|
||||
return nil, nil
|
||||
}
|
||||
@@ -647,20 +654,27 @@ func (d *DB) GetFinding(id int64) (*DBFinding, error) {
|
||||
// list, the per-task 发现 Tab, and the exploration graph alike. Deleting the node
|
||||
// cascades its edges + node_assets and nulls any activity referencing it. Returns
|
||||
// rows affected (0 = no finding with that id).
|
||||
func (d *DB) DeleteFinding(id int64) (int64, error) {
|
||||
var nodeID *int64
|
||||
err := d.QueryRow(`DELETE FROM findings WHERE id=$1 RETURNING node_id`, id).Scan(&nodeID)
|
||||
if err == sql.ErrNoRows {
|
||||
return 0, nil
|
||||
}
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
if nodeID != nil {
|
||||
// best-effort: the finding is already gone; a stray node must not fail the op.
|
||||
_, _ = d.Exec(`DELETE FROM exploration_nodes WHERE id=$1 AND kind='finding'`, *nodeID)
|
||||
}
|
||||
return 1, nil
|
||||
func (d *DB) DeleteFinding(id int64) (n int64, err error) {
|
||||
err = d.WithEvidenceTx(context.Background(), func(tx *sql.Tx) error {
|
||||
if err := LockFindingEvidenceTx(tx, id, nil); err != nil {
|
||||
if errors.Is(err, ErrFindingNotFound) {
|
||||
return nil
|
||||
}
|
||||
return err
|
||||
}
|
||||
var nodeID sql.NullInt64
|
||||
if err := tx.QueryRow(`DELETE FROM findings WHERE id=$1 RETURNING node_id`, id).Scan(&nodeID); err != nil {
|
||||
return err
|
||||
}
|
||||
if nodeID.Valid {
|
||||
if _, err := tx.Exec(`DELETE FROM exploration_nodes WHERE id=$1 AND kind='finding'`, nodeID.Int64); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
n = 1
|
||||
return nil
|
||||
})
|
||||
return
|
||||
}
|
||||
|
||||
// DeleteFindingsByTask removes all findings rows of a task. The originating
|
||||
@@ -687,11 +701,7 @@ func (d *DB) SetFindingStatus(id int64, status string) (int64, error) {
|
||||
// whose node_id matches — report_finding returns that node id, so an agent tool
|
||||
// can address the finding it just created. Returns rows affected (0 when no row).
|
||||
func (d *DB) SetFindingReportByNodeID(nodeID int64, report string) (int64, error) {
|
||||
res, err := d.Exec(`UPDATE findings SET report=$1 WHERE node_id=$2`, report, nodeID)
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
return res.RowsAffected()
|
||||
return d.SetFindingReportVersionByNodeID(context.Background(), nodeID, report, nil)
|
||||
}
|
||||
|
||||
// setFindingCol updates one text column on the standalone finding row AND mirrors
|
||||
@@ -737,6 +747,8 @@ func (d *DB) SetFindingVulnClass(id int64, vulnclass string) (int64, error) {
|
||||
// FindingMeta is the standalone-row data (id, triage state, anchored assets) the
|
||||
// per-task view grafts onto its exploration-node findings.
|
||||
type FindingMeta struct {
|
||||
TrafficCount int
|
||||
|
||||
ID int64
|
||||
Status string
|
||||
AssetIDs []int64
|
||||
@@ -757,7 +769,7 @@ func (d *DB) FindingMetaByNodeID(taskID int64) (map[int64]FindingMeta, error) {
|
||||
if taskID <= 0 {
|
||||
return out, nil
|
||||
}
|
||||
rows, err := d.Query(`SELECT node_id, id, COALESCE(status,'pending'), asset_ids FROM findings
|
||||
rows, err := d.Query(`SELECT node_id, id, COALESCE(status,'pending'), asset_ids, (SELECT count(*) FROM finding_traffic_bindings b WHERE b.finding_id=findings.id) FROM findings
|
||||
WHERE task_id=$1 AND node_id IS NOT NULL`, taskID)
|
||||
if err != nil {
|
||||
return out, err
|
||||
@@ -767,7 +779,7 @@ func (d *DB) FindingMetaByNodeID(taskID int64) (map[int64]FindingMeta, error) {
|
||||
var nid int64
|
||||
var m FindingMeta
|
||||
var aidsJSON string
|
||||
if err := rows.Scan(&nid, &m.ID, &m.Status, &aidsJSON); err != nil {
|
||||
if err := rows.Scan(&nid, &m.ID, &m.Status, &aidsJSON, &m.TrafficCount); err != nil {
|
||||
return out, err
|
||||
}
|
||||
_ = json.Unmarshal([]byte(aidsJSON), &m.AssetIDs)
|
||||
|
||||
@@ -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
@@ -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}},
|
||||
|
||||
@@ -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 {
|
||||
|
||||
@@ -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 的证书问题。
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,333 @@
|
||||
package evidence
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"crypto/sha256"
|
||||
"encoding/hex"
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"sync"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/Autumn-27/artex/db"
|
||||
"github.com/Autumn-27/artex/traffic"
|
||||
)
|
||||
|
||||
func evidenceFixture(t *testing.T) (*Store, db.RecordFindingInput, string) {
|
||||
t.Helper()
|
||||
dsn := os.Getenv("ARTEX_PG_DSN")
|
||||
if dsn == "" {
|
||||
t.Skip("ARTEX_PG_DSN is required for evidence integration tests")
|
||||
}
|
||||
pg, err := db.Open(dsn)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
t.Cleanup(func() { pg.Close() })
|
||||
task, err := pg.CreateTask("evidence integration "+t.Name(), "local fixtures", nil, 0, 0)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
t.Cleanup(func() { pg.DeleteFindingsByTask(task.ID); pg.DeleteTask(task.ID) })
|
||||
dir := t.TempDir()
|
||||
tr, err := traffic.Open(filepath.Join(dir, "traffic"), ":0")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
t.Cleanup(func() { tr.Close() })
|
||||
store := New(pg, tr, filepath.Join(dir, "evidence"))
|
||||
in := db.RecordFindingInput{TaskID: task.ID, ExplorationID: task.ExplorationID, Worker: "test", VulnClass: "TEST", Name: "Evidence fixture", Severity: "low", Summary: "local test"}
|
||||
return store, in, dir
|
||||
}
|
||||
|
||||
func seedExchange(t *testing.T, s *Store, id string, body []byte, spill bool) {
|
||||
t.Helper()
|
||||
var blob any
|
||||
inline := body
|
||||
if spill {
|
||||
sum := sha256.Sum256(body)
|
||||
hash := hex.EncodeToString(sum[:])
|
||||
blob = hash
|
||||
path := filepath.Join(filepath.Dir(s.Dir), "traffic", "_blobs", "sha256", hash[:2], hash+".bin")
|
||||
if err := os.MkdirAll(filepath.Dir(path), 0o700); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := os.WriteFile(path, body, 0o600); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
inline = []byte("TRUNCATED PREVIEW")
|
||||
}
|
||||
_, err := s.Traffic.DB().Exec(`INSERT INTO exchanges(id,ts,host,method,url_template,url,status,content_type,req_len,resp_len,path) VALUES(?,?,'fixture.local','POST','/test',?,200,'application/octet-stream',2,?,'')`, id, time.Now().Unix(), "https://fixture.local/"+id, len(body))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
_, err = s.Traffic.DB().Exec(`INSERT INTO exchange_bodies(id,req_head,req_body,resp_head,resp_body,resp_blob) VALUES(?,?,?, ?,?,?)`, id, "POST /test HTTP/1.1\nHost: fixture.local\n", []byte("{}"), "HTTP 200\nContent-Type: application/octet-stream\n", inline, blob)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestEvidenceBindingLifecycle(t *testing.T) {
|
||||
s, in, _ := evidenceFixture(t)
|
||||
ctx := context.Background()
|
||||
large := bytes.Repeat([]byte{0, 1, 2, 255, 'a', 'b'}, 180000)
|
||||
for i := 0; i < 3; i++ {
|
||||
seedExchange(t, s, fmt.Sprint(i), large, i == 2)
|
||||
}
|
||||
refs := []db.TrafficRef{{TrafficID: "0", Role: "baseline", Note: "normal"}, {TrafficID: "1", Role: "proof", Note: "proof"}, {TrafficID: "2", Role: "verification", Note: "binary"}}
|
||||
r, err := s.Record(ctx, in, refs)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(r.Traffic.Bindings) != 3 || r.Traffic.Version != 1 {
|
||||
t.Fatalf("record=%+v", r)
|
||||
}
|
||||
again, err := s.Bind(ctx, r.FindingID, []db.TrafficRef{{TrafficID: "1", Note: "must not overwrite"}})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(again.Bindings) != 3 || again.Version != 1 || again.Bindings[1].Note != "proof" {
|
||||
t.Fatalf("duplicate changed bindings: %+v", again)
|
||||
}
|
||||
other, err := s.Record(ctx, in, []db.TrafficRef{refs[2]})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if other.Traffic.Bindings[0].SnapshotID != r.Traffic.Bindings[2].SnapshotID {
|
||||
t.Fatal("snapshot was not shared")
|
||||
}
|
||||
if _, err = s.Traffic.DeleteHost("fixture.local"); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
for _, binding := range r.Traffic.Bindings {
|
||||
if err = s.WithBinding(ctx, r.FindingID, binding.ID, func(b db.FindingTrafficBinding) error {
|
||||
f, _, err := s.OpenBody(b.Snapshot, "response")
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer f.Close()
|
||||
got, err := io.ReadAll(f)
|
||||
if !bytes.Equal(got, large) {
|
||||
t.Fatal("body was truncated/changed")
|
||||
}
|
||||
return err
|
||||
}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
v := again.Version
|
||||
if _, err = s.DB.SetFindingReportVersionByNodeID(ctx, r.NodeID, "report", &v); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
note := "updated"
|
||||
if err = s.DB.EditFindingTraffic(ctx, r.FindingID, again.Bindings[0].ID, v, nil, ¬e, false, nil); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if _, err = s.DB.SetFindingReportVersionByNodeID(ctx, r.NodeID, "stale report", &v); !errors.Is(err, db.ErrEvidenceConflict) {
|
||||
t.Fatalf("stale report: %v", err)
|
||||
}
|
||||
f, err := s.DB.GetFinding(r.FindingID)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if f.Report != "report" || f.EvidenceVersion == f.ReportEvidenceVersion || f.TrafficCount != 3 {
|
||||
t.Fatalf("finding versions: %+v", f)
|
||||
}
|
||||
if err = s.DB.EditFindingTraffic(ctx, r.FindingID, again.Bindings[0].ID, v, nil, nil, true, nil); !errors.Is(err, db.ErrEvidenceConflict) {
|
||||
t.Fatalf("stale delete: %v", err)
|
||||
}
|
||||
if _, err = s.DB.SetFindingReportByNodeID(r.NodeID, "legacy report"); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
f, err = s.DB.GetFinding(r.FindingID)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if f.ReportEvidenceVersion != -1 {
|
||||
t.Fatal("legacy write claimed current evidence")
|
||||
}
|
||||
if _, err = s.DB.DeleteFinding(r.FindingID); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err = s.Collect(ctx, time.Now().Add(48*time.Hour)); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err = s.WithBinding(ctx, other.FindingID, other.Traffic.Bindings[0].ID, func(b db.FindingTrafficBinding) error {
|
||||
f, _, err := s.OpenBody(b.Snapshot, "response")
|
||||
if f != nil {
|
||||
f.Close()
|
||||
}
|
||||
return err
|
||||
}); err != nil {
|
||||
t.Fatal("GC removed shared body:", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestEvidenceAtomicFailuresAndLegacy(t *testing.T) {
|
||||
s, in, dir := evidenceFixture(t)
|
||||
ctx := context.Background()
|
||||
seedExchange(t, s, "valid", []byte("OK"), false)
|
||||
counts := func() (findings, nodes int) {
|
||||
s.DB.QueryRow(`SELECT count(*) FROM findings WHERE task_id=$1`, in.TaskID).Scan(&findings)
|
||||
s.DB.QueryRow(`SELECT count(*) FROM exploration_nodes WHERE exploration_id=$1 AND kind='finding'`, in.ExplorationID).Scan(&nodes)
|
||||
return
|
||||
}
|
||||
for _, kind := range []string{"missing_id", "missing_body", "disk_failure", "database_failure"} {
|
||||
t.Run(kind, func(t *testing.T) {
|
||||
input := in
|
||||
refs := []db.TrafficRef{{TrafficID: "valid"}, {TrafficID: "absent"}}
|
||||
local := *s
|
||||
switch kind {
|
||||
case "missing_body":
|
||||
seedExchange(t, s, "broken", []byte("more than preview"), false)
|
||||
s.Traffic.DB().Exec(`UPDATE exchange_bodies SET resp_body=? WHERE id='broken'`, []byte("x"))
|
||||
refs = []db.TrafficRef{{TrafficID: "valid"}, {TrafficID: "broken"}}
|
||||
case "disk_failure":
|
||||
local.Dir = filepath.Join(dir, "file-not-dir")
|
||||
if err := os.WriteFile(local.Dir, []byte("x"), 0o600); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
refs = refs[:1]
|
||||
case "database_failure":
|
||||
input.AssetIDs = []int64{9223372036854775807}
|
||||
refs = refs[:1]
|
||||
}
|
||||
if _, err := local.Record(ctx, input, refs); err == nil {
|
||||
t.Fatal("expected failure")
|
||||
}
|
||||
if f, n := counts(); f != 0 || n != 0 {
|
||||
t.Fatalf("partial record: findings=%d nodes=%d", f, n)
|
||||
}
|
||||
})
|
||||
}
|
||||
legacy := filepath.Join(dir, "traffic", "legacy")
|
||||
os.MkdirAll(legacy, 0o700)
|
||||
os.WriteFile(filepath.Join(legacy, "request.http"), []byte("GET / HTTP/1.1\nHost: old.local\n"), 0o600)
|
||||
os.WriteFile(filepath.Join(legacy, "response.http"), []byte("HTTP 200\n\nlegacy body"), 0o600)
|
||||
_, err := s.Traffic.DB().Exec(`INSERT INTO exchanges(id,ts,host,method,url_template,url,status,content_type,req_len,resp_len,path) VALUES('old',1,'old.local','GET','/','http://old.local/',200,'text/plain',0,11,'legacy')`)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if _, err = s.Record(ctx, in, []db.TrafficRef{{TrafficID: "old"}}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if _, err = s.Record(ctx, in, nil); err != nil {
|
||||
t.Fatal("legacy no-traffic report failed:", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestEvidenceConcurrentBindAndDelete(t *testing.T) {
|
||||
s, in, _ := evidenceFixture(t)
|
||||
ctx := context.Background()
|
||||
seedExchange(t, s, "race", bytes.Repeat([]byte("large"), 90000), true)
|
||||
r, err := s.Record(ctx, in, nil)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
var wg sync.WaitGroup
|
||||
errs := make(chan error, 10)
|
||||
for i := 0; i < 8; i++ {
|
||||
wg.Add(1)
|
||||
go func() {
|
||||
defer wg.Done()
|
||||
_, err := s.Bind(ctx, r.FindingID, []db.TrafficRef{{TrafficID: "race"}})
|
||||
errs <- err
|
||||
}()
|
||||
}
|
||||
wg.Wait()
|
||||
close(errs)
|
||||
for err := range errs {
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
list, err := s.DB.GetFindingTraffic(ctx, r.FindingID)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(list.Bindings) != 1 || list.Version != 1 {
|
||||
t.Fatalf("concurrent duplicates: %+v", list)
|
||||
}
|
||||
// Race a new capture copy with source deletion: either the full copy commits
|
||||
// or nothing does. An already-bound snapshot remains readable in both cases.
|
||||
wg.Add(2)
|
||||
go func() { defer wg.Done(); s.Bind(ctx, r.FindingID, []db.TrafficRef{{TrafficID: "race"}}) }()
|
||||
go func() { defer wg.Done(); s.Traffic.DeleteHost("fixture.local") }()
|
||||
wg.Wait()
|
||||
if err = s.WithBinding(ctx, r.FindingID, list.Bindings[0].ID, func(b db.FindingTrafficBinding) error {
|
||||
f, _, err := s.OpenBody(b.Snapshot, "response")
|
||||
if f != nil {
|
||||
f.Close()
|
||||
}
|
||||
return err
|
||||
}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestEvidenceGCGraceAndActiveRestore(t *testing.T) {
|
||||
s, in, _ := evidenceFixture(t)
|
||||
ctx := context.Background()
|
||||
seedExchange(t, s, "gc", []byte("unique unreferenced gc body"), false)
|
||||
f, err := s.Record(ctx, in, []db.TrafficRef{{TrafficID: "gc"}})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
snap := f.Traffic.Bindings[0].Snapshot
|
||||
archive := t.TempDir()
|
||||
if err = s.CopySnapshots(ctx, []db.TrafficEvidenceSnapshot{snap}, archive); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if _, err = s.DB.DeleteFinding(f.FindingID); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
now := time.Now()
|
||||
if err = s.Collect(ctx, now); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
path, err := hashPath(s.Dir, snap.RespHash)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err = s.Collect(ctx, now.Add(23*time.Hour)); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if _, err = os.Stat(path); err != nil {
|
||||
t.Fatal("GC grace ignored", err)
|
||||
}
|
||||
if err = s.Collect(ctx, now.Add(25*time.Hour)); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if _, err = os.Stat(path); !os.IsNotExist(err) {
|
||||
t.Fatal("orphan was not collected", err)
|
||||
}
|
||||
entered, release, finished := make(chan struct{}), make(chan struct{}), make(chan error, 1)
|
||||
go func() {
|
||||
finished <- s.WithInstalledSnapshots(ctx, []db.TrafficEvidenceSnapshot{snap}, archive, func() error { close(entered); <-release; return nil })
|
||||
}()
|
||||
<-entered
|
||||
short, cancel := context.WithTimeout(ctx, 100*time.Millisecond)
|
||||
err = s.Collect(short, now.Add(72*time.Hour))
|
||||
cancel()
|
||||
close(release)
|
||||
if err == nil {
|
||||
t.Fatal("GC entered during active restore")
|
||||
}
|
||||
if err = <-finished; err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if _, err = os.Stat(path); err != nil {
|
||||
t.Fatal("active restore lost body", err)
|
||||
}
|
||||
// A portable package cannot substitute metadata while retaining its old ID.
|
||||
bad := snap
|
||||
bad.URL = "http://tampered.local/"
|
||||
if err = s.InstallSnapshots(ctx, []db.TrafficEvidenceSnapshot{bad}, archive); err == nil {
|
||||
t.Fatal("accepted tampered snapshot")
|
||||
}
|
||||
}
|
||||
@@ -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
@@ -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()
|
||||
}
|
||||
+11
-2
@@ -292,6 +292,12 @@ func coverageAssetRefDTO(ref db.AssetRef) CoverageAssetRefDTO {
|
||||
// ---- Finding (frontend "Finding") ----
|
||||
|
||||
type FindingDTO struct {
|
||||
TrafficCount int `json:"traffic_count"`
|
||||
EvidenceVersion int64 `json:"evidence_version"`
|
||||
ReportEvidenceVersion int64 `json:"report_evidence_version"`
|
||||
ReportStale bool `json:"report_stale"`
|
||||
TrafficBindings []db.FindingTrafficBinding `json:"traffic_bindings,omitempty"`
|
||||
|
||||
ID string `json:"id"`
|
||||
FindingID string `json:"finding_id,omitempty"` // standalone findings-table id — the handle for status updates
|
||||
VulnClass string `json:"vulnclass"`
|
||||
@@ -399,6 +405,7 @@ func findingDTOsForOwner(taskID, description string, in []*db.Node, meta map[int
|
||||
if m, ok := meta[n.ID]; ok {
|
||||
d.FindingID = i64s(m.ID)
|
||||
d.Status = m.Status
|
||||
d.TrafficCount = m.TrafficCount
|
||||
d.Assets = findingAssetDTOs(m.AssetIDs, assets)
|
||||
}
|
||||
out = append(out, d)
|
||||
@@ -426,8 +433,10 @@ func findingFromDB(f *db.DBFinding, assets map[int64]*db.Asset) FindingDTO {
|
||||
status = db.FindingPending
|
||||
}
|
||||
d := FindingDTO{
|
||||
ID: i64s(f.ID),
|
||||
FindingID: i64s(f.ID),
|
||||
ID: i64s(f.ID),
|
||||
FindingID: i64s(f.ID),
|
||||
TrafficCount: f.TrafficCount, EvidenceVersion: f.EvidenceVersion, ReportEvidenceVersion: f.ReportEvidenceVersion,
|
||||
ReportStale: f.Report != "" && f.EvidenceVersion != f.ReportEvidenceVersion, TrafficBindings: f.TrafficBindings,
|
||||
VulnClass: f.VulnClass,
|
||||
Name: f.Name,
|
||||
Severity: f.Severity,
|
||||
|
||||
@@ -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)
|
||||
})
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
}
|
||||
+15
-11
@@ -75,6 +75,7 @@ func (s *Server) orchestrationTools() []actool.CoreTool {
|
||||
s.toolSearchWorkerTraces(),
|
||||
s.toolGetTaskNodeDetail(),
|
||||
s.toolUpdateFindingReport(),
|
||||
s.toolGetFindingTraffic(),
|
||||
}
|
||||
}
|
||||
|
||||
@@ -392,20 +393,22 @@ func (s *Server) toolUpdateFindingReport() actool.CoreTool {
|
||||
return wrTool("update_finding_report",
|
||||
"为已登记的漏洞写入/更新【详细报告】(Markdown 全文,整段覆盖旧内容)。finding_id 传 report_finding 返回的那个 id(\"finding recorded: <id>\" 里的数字)。报告建议包含:漏洞概述、影响与危害、复现步骤、证据/PoC、修复建议。",
|
||||
objSchema(map[string]any{
|
||||
"finding_id": map[string]any{"type": "integer", "description": "目标漏洞 id(report_finding 返回的 id)"},
|
||||
"report": strParam("详细报告全文,Markdown 格式"),
|
||||
"finding_id": map[string]any{"type": "integer", "description": "目标漏洞 id(report_finding 返回的 id)"},
|
||||
"report": strParam("详细报告全文,Markdown 格式"),
|
||||
"evidence_version": map[string]any{"type": "integer", "description": "get_finding_traffic 返回的证据 version；用于防止报告覆盖新的证据变更"},
|
||||
}, "finding_id", "report"),
|
||||
func(_ context.Context, in json.RawMessage) (actool.Result, error) {
|
||||
func(ctx context.Context, in json.RawMessage) (actool.Result, error) {
|
||||
var a struct {
|
||||
FindingID json.RawMessage `json:"finding_id"`
|
||||
Report string `json:"report"`
|
||||
EvidenceVersion *int64 `json:"evidence_version"`
|
||||
FindingID json.RawMessage `json:"finding_id"`
|
||||
Report string `json:"report"`
|
||||
}
|
||||
_ = json.Unmarshal(in, &a)
|
||||
nodeID := parseProfileID(a.FindingID) // 复用「数字或数字字符串」解析
|
||||
if nodeID <= 0 {
|
||||
return actool.Errorf("finding_id 无效"), nil
|
||||
}
|
||||
n, err := s.m.pg.SetFindingReportByNodeID(nodeID, a.Report)
|
||||
n, err := s.m.pg.SetFindingReportVersionByNodeID(ctx, nodeID, a.Report, a.EvidenceVersion)
|
||||
if err != nil {
|
||||
return actool.Errorf(err.Error()), nil
|
||||
}
|
||||
@@ -453,11 +456,12 @@ func (s *Server) seedOrchestrationTools() {
|
||||
s.seedWorkerReadToolsUnbind() // list_facts/node_detail/list_companies/跨 work 检索从 worker 默认解绑(一次性)
|
||||
s.seedAutoReportFindingBinding()
|
||||
s.unbindGoalMetDefault()
|
||||
s.reseedGoalsPrompt() // goals 提示词加入「抽操作约束」步 → 旧库追加一版新默认(一次性)
|
||||
s.reseedMainAgentPrompt() // mainagent 提示词加入「目标达成后 add_intent 反问是否建目标」(一次性)
|
||||
s.reseedPlannerPrompt() // planner 提示词:重写「0 意图」正当理由 + 加量化验收核对(一次性)
|
||||
s.reseedWorkerPrompt() // worker 提示词:加否定结论证据门槛(一次性)
|
||||
s.seedReporterAgent() // 预置「报告撰写」agent + 工具绑定 + finding 触发器(一次性)
|
||||
s.reseedGoalsPrompt() // goals 提示词加入「抽操作约束」步 → 旧库追加一版新默认(一次性)
|
||||
s.reseedMainAgentPrompt() // mainagent 提示词加入「目标达成后 add_intent 反问是否建目标」(一次性)
|
||||
s.reseedPlannerPrompt() // planner 提示词:重写「0 意图」正当理由 + 加量化验收核对(一次性)
|
||||
s.reseedWorkerPrompt() // worker 提示词:加否定结论证据门槛(一次性)
|
||||
s.seedReporterAgent() // 预置「报告撰写」agent + 工具绑定 + finding 触发器(一次性)
|
||||
s.seedFindingTrafficTools() // 增加可选证据参数及只读证据工具，保留用户配置
|
||||
// 注：pentest 的默认工具绑定无需迁移——BuiltinToolSeeds 在全新初始化时就把
|
||||
// list_assets/insert_assets/report_finding/list_findings/list_companies 连同
|
||||
// pentest 一起 seed 好了（项目尚无旧库，不做迁移）。
|
||||
|
||||
@@ -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] "
|
||||
|
||||
+36
-25
@@ -1,7 +1,6 @@
|
||||
package server
|
||||
|
||||
import (
|
||||
"archive/zip"
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
@@ -11,6 +10,7 @@ import (
|
||||
"net"
|
||||
"net/http"
|
||||
"net/url"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"regexp"
|
||||
"sort"
|
||||
@@ -217,9 +217,10 @@ func New(ctx context.Context, m *Manager, skillDir string, dataDir string, keyDi
|
||||
}
|
||||
wireAgentAugment(m.pg, s.skillDir, s.hostTools) // 可见 skills/MCP + 流量/编排 host 工具装配进 agent 工具集
|
||||
domainReg := buildDomainReg(m.Assets())
|
||||
wireTools(m.pg, domainReg) // 内置工具表：按 agent 过滤 + 覆盖描述/schema + 注入默认值
|
||||
seedPrompts(m.pg) // 内置 agent 默认提示词正文播种进 agent_prompts(仅空时)
|
||||
s.seedOrchestrationTools() // P2 跨任务编排工具 seed 进 tools 表(可按 agent 绑定)
|
||||
wireTools(m.pg, domainReg) // 内置工具表：按 agent 过滤 + 覆盖描述/schema + 注入默认值
|
||||
seedPrompts(m.pg) // 内置 agent 默认提示词正文播种进 agent_prompts(仅空时)
|
||||
s.seedOrchestrationTools() // P2 跨任务编排工具 seed 进 tools 表(可按 agent 绑定)
|
||||
go s.evidenceStore().RunGC(s.ctx)
|
||||
s.seedPythonInterpreter() // 自定义脚本工具:开机检测 python 解释器入库(仅空时)
|
||||
go newScheduler(s).Run(s.ctx) // P3 触发器调度(定时/finding/目标事件),仅自定义 agent
|
||||
// Fill the tool cache for any enabled MCP that has none yet (notably the
|
||||
@@ -360,7 +361,7 @@ func (s *Server) saveLLMConfig(cfg agent.Config) error {
|
||||
ContextWindowK: cfg.ContextWindowK, ThinkingType: cfg.ThinkingType, ReasoningEffort: cfg.ReasoningEffort, IsDefault: true,
|
||||
Priority: priority, PoolExclude: poolExclude, Streaming: streaming,
|
||||
MaxTokens: maxTokens, MaxTokensField: maxTokensField, SessionHeaderKey: sessionHeaderKey,
|
||||
Retry: retry,
|
||||
Retry: retry,
|
||||
})
|
||||
if err != nil {
|
||||
return err
|
||||
@@ -405,6 +406,7 @@ func (s *Server) buildPlannerWorker(pinID *int64, gProv llm.Provider, gCfg agent
|
||||
// the tools-table binding (default = worker), so worker behavior is unchanged.
|
||||
wProv, wCfg := s.providerForAgent("worker", pinID, gProv, gCfg)
|
||||
wk := agent.NewWorker(wProv, wCfg.Model, s.m.dir, tx, wCfg.CompactionWindow(), s.agentMaxTurns("worker"))
|
||||
wk.SetFindingRecorder(s.evidenceStore())
|
||||
wk.SetRunTimeout(time.Duration(s.agentRunSeconds("worker")) * time.Second)
|
||||
wk.SetProxy(s.m.ProxyAddr(), s.m.ProxyCACert())
|
||||
wk.SetMemory(memory.NewStore(filepath.Join(s.m.dir, "memory")))
|
||||
@@ -414,6 +416,7 @@ func (s *Server) buildPlannerWorker(pinID *int64, gProv llm.Provider, gCfg agent
|
||||
wk.SetMaxTokens(maxTokensResolver(wCfg)) // 单次回复输出上限(0 = 不发)
|
||||
pProv, pCfg := s.providerForAgent("planner", pinID, gProv, gCfg)
|
||||
pl := agent.NewPlanner(pProv, pCfg.Model, s.m.dir, tx, pCfg.CompactionWindow(), s.agentMaxTurns("planner"))
|
||||
pl.SetFindingRecorder(s.evidenceStore())
|
||||
pl.SetKillWork(s.engine.KillWork) // planner kill_work → terminate a running work
|
||||
pl.SetSteerWork(s.engine.SteerWork) // planner steer_work → inject mid-run course-correction
|
||||
pl.SetProxy(s.m.ProxyAddr(), s.m.ProxyCACert()) // WebFetch through the recording proxy
|
||||
@@ -482,6 +485,7 @@ func (s *Server) applyLLM(cfg agent.Config) error {
|
||||
mProv, mCfg := s.providerForAgent("mainagent", nil, prov, cfg)
|
||||
s.cfgMu.Lock()
|
||||
s.mainAgent = agent.NewMainAgent(mProv, mCfg.Model, s.m.dir, tx, mCfg.CompactionWindow(), s.agentMaxTurns("mainagent"))
|
||||
s.mainAgent.SetFindingRecorder(s.evidenceStore())
|
||||
s.mainAgent.SetProxy(s.m.ProxyAddr(), s.m.ProxyCACert()) // WebFetch through the recording proxy
|
||||
s.mainAgent.SetWebSearch(s.webSearchFor("mainagent"))
|
||||
s.mainAgent.SetSteerWork(s.engine.SteerWork) // steer_work：人对运行中 work 实时纠偏
|
||||
@@ -811,6 +815,7 @@ func (s *Server) Handler() http.Handler {
|
||||
mux.HandleFunc("GET /api/exploration/findings/asset-tree", s.findingAssetTree)
|
||||
mux.HandleFunc("GET /api/exploration/findings/stats", s.findingStats)
|
||||
mux.HandleFunc("GET /api/exploration/findings/export", s.findingsExport)
|
||||
s.registerFindingTraffic(mux)
|
||||
mux.HandleFunc("GET /api/exploration/findings/{id}", s.getFinding)
|
||||
mux.HandleFunc("GET /api/exploration/findings/{id}/lineage", s.findingLineage)
|
||||
mux.HandleFunc("POST /api/exploration/findings/{id}/deepen", s.deepenFinding)
|
||||
@@ -2135,6 +2140,16 @@ func (s *Server) findingsExport(w http.ResponseWriter, r *http.Request) {
|
||||
return
|
||||
}
|
||||
|
||||
stage, err := os.MkdirTemp("", "artex-finding-export-")
|
||||
if err != nil {
|
||||
writeErr(w, 500, err.Error())
|
||||
return
|
||||
}
|
||||
defer os.RemoveAll(stage)
|
||||
if err = s.evidenceStore().StageFindingsExport(r.Context(), fs, stage, format == "md-zip"); err != nil {
|
||||
evidenceError(w, err)
|
||||
return
|
||||
}
|
||||
now := time.Now()
|
||||
stamp := now.Format("20060102-150405")
|
||||
setDownload := func(contentType, filename string) {
|
||||
@@ -2147,27 +2162,19 @@ func (s *Server) findingsExport(w http.ResponseWriter, r *http.Request) {
|
||||
setDownload("text/markdown; charset=utf-8", "findings-"+stamp+".md")
|
||||
_, _ = w.Write([]byte(report.FindingsMarkdown(fs, now)))
|
||||
case "md-zip":
|
||||
path := filepath.Join(stage, "findings.zip")
|
||||
if err := buildFindingsEvidenceZip(path, fs, stage, now); err != nil {
|
||||
writeErr(w, 500, err.Error())
|
||||
return
|
||||
}
|
||||
file, err := os.Open(path)
|
||||
if err != nil {
|
||||
writeErr(w, 500, err.Error())
|
||||
return
|
||||
}
|
||||
defer file.Close()
|
||||
setDownload("application/zip", "findings-"+stamp+".zip")
|
||||
zw := zip.NewWriter(w)
|
||||
used := map[string]int{}
|
||||
for _, f := range fs {
|
||||
base := report.FindingFilename(f)
|
||||
name := base
|
||||
// 去重:同名文件追加 -2、-3……
|
||||
if n := used[base]; n > 0 {
|
||||
name = fmt.Sprintf("%s-%d.md", strings.TrimSuffix(base, ".md"), n+1)
|
||||
}
|
||||
used[base]++
|
||||
fw, werr := zw.Create(name)
|
||||
if werr != nil {
|
||||
log.Printf("[findings-export] zip create %s: %v", name, werr)
|
||||
continue
|
||||
}
|
||||
_, _ = fw.Write([]byte(report.SingleFindingMarkdown(f, now)))
|
||||
}
|
||||
if cerr := zw.Close(); cerr != nil {
|
||||
log.Printf("[findings-export] zip close: %v", cerr)
|
||||
}
|
||||
http.ServeContent(w, r, "findings.zip", now, file)
|
||||
case "csv":
|
||||
setDownload("text/csv; charset=utf-8", "findings-"+stamp+".csv")
|
||||
_, _ = w.Write(report.FindingsCSV(fs))
|
||||
@@ -2175,6 +2182,10 @@ func (s *Server) findingsExport(w http.ResponseWriter, r *http.Request) {
|
||||
assets := s.resolveFindingAssets(fs)
|
||||
out := make([]FindingDTO, 0, len(fs))
|
||||
for _, f := range fs {
|
||||
for i := range f.TrafficBindings {
|
||||
f.TrafficBindings[i].Snapshot.ReqHead = ""
|
||||
f.TrafficBindings[i].Snapshot.RespHead = ""
|
||||
}
|
||||
out = append(out, findingFromDB(f, assets))
|
||||
}
|
||||
setDownload("application/json; charset=utf-8", "findings-"+stamp+".json")
|
||||
|
||||
+20
-2
@@ -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
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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) => {
|
||||
if (f.inherited || !f.finding_id || next === f.status) return;
|
||||
const prev = f.status;
|
||||
setFindings((cur) =>
|
||||
cur.map((x) => (x.id === f.id ? { ...x, status: next } : x)),
|
||||
);
|
||||
try {
|
||||
await api.setFindingStatus(f.finding_id, next);
|
||||
toast.success(`已标记为「${statusMeta("finding", next).label}」`);
|
||||
} catch (e) {
|
||||
setFindings((cur) =>
|
||||
cur.map((x) => (x.id === f.id ? { ...x, status: prev } : x)),
|
||||
);
|
||||
toast.error("更新失败：" + (e as Error).message);
|
||||
}
|
||||
},
|
||||
[],
|
||||
);
|
||||
const onStatus = React.useCallback(async (f: Finding, next: FindingStatus) => {
|
||||
if (f.inherited || !f.finding_id || next === f.status) return;
|
||||
const prev = f.status;
|
||||
setFindings((cur) => cur.map((x) => (x.id === f.id ? { ...x, status: next } : x)));
|
||||
try {
|
||||
await api.setFindingStatus(f.finding_id, next);
|
||||
toast.success(`已标记为「${statusMeta("finding", next).label}」`);
|
||||
} catch (e) {
|
||||
setFindings((cur) => cur.map((x) => (x.id === f.id ? { ...x, status: prev } : x)));
|
||||
toast.error("更新失败：" + (e as Error).message);
|
||||
}
|
||||
}, []);
|
||||
|
||||
const items = findings
|
||||
.filter((f) => f.task_id === taskId || f.inherited)
|
||||
@@ -243,9 +214,7 @@ export function FindingsTab({ taskId }: { taskId: string }) {
|
||||
<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>
|
||||
|
||||
@@ -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>
|
||||
</>
|
||||
);
|
||||
}
|
||||
@@ -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>
|
||||
);
|
||||
}
|
||||
@@ -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)} />
|
||||
</>
|
||||
);
|
||||
}
|
||||
@@ -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 }),
|
||||
|
||||
@@ -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;
|
||||
}
|
||||
Reference in new issue
Block a user