Merge PR #102: feat: 支持漏洞关联多条流量证据

冲突解决:
- db/schema.sql: 复测表(#103)与流量证据表(#102)均为追加,两者保留
- server/server.go: 启动序列同时保留 seedFindingRetester 与 evidence GC
- web/src/lib/types.ts: InterceptAudit/InterceptDetail(#101) 与流量证据类型(#102)并存
- findings/detail/page.tsx: 去掉 #103 带入的重复 import,仅保留 FindingRetestPanel
- chat/page.tsx: 保留 #103 的 Spinner import
This commit is contained in:
Autumn-27 committed 2026-09-10 10:04:53 -04:00
commit e8efce605c
54 files changed
+5328 -552

No files matched your search

+6 -4
View File
@@ -12,10 +12,11 @@ import (
// UnlockSkill unlocks a named skill's MCPs — hosts call it to rebuild the unlock set
// from history on a resumed session (design doc C2).
type DeferredInfo struct {
Deferred []string // all MCP tool names (schema withheld)
GlobalNames []string // MCP names to list in the system-prompt block
Unlock *actool.UnlockSet // shared call-gate; nil when no MCP tools
UnlockSkill func(skillName string)
FindingGuidance string // derived from the final permitted tools, including DB overrides
Deferred []string // all MCP tool names (schema withheld)
GlobalNames []string // MCP names to list in the system-prompt block
Unlock *actool.UnlockSet // shared call-gate; nil when no MCP tools
UnlockSkill func(skillName string)
}
// ToolAugment, if set, returns the EXTRA tools an agent should see beyond its
@@ -55,5 +56,6 @@ func AugmentTools(ctx context.Context, agentKey string, base []actool.CoreTool)
if ToolResolve != nil {
out = ToolResolve(ctx, agentKey, out)
}
out, def.FindingGuidance = findingWorkflowTools(agentKey, out)
return out, def, cleanup
}
+1
View File
@@ -16,6 +16,7 @@ import (
// skill loads. Returns a single plain segment + boundary 0 when there is no global
// block to add.
func deferredSystem(sysText string, def DeferredInfo) (system []string, boundary int) {
sysText += def.FindingGuidance
block := actool.RenderDeferredToolsBlock(def.GlobalNames)
if block == "" {
return []string{sysText}, 0
+22
View File
@@ -0,0 +1,22 @@
package agent
import (
"context"
"github.com/Autumn-27/artex/db"
)
// FindingRecorder is injected by the host; agents never synthesize or copy
// evidence bodies themselves. Its implementation owns the atomic write.
type FindingRecorder interface {
Record(context.Context, db.RecordFindingInput, []db.TrafficRef) (*db.RecordedFinding, error)
}
// Tool-use guidance is appended without replacing the user's editable prompt.
// It does not require capture or claim that unavailable traffic tools exist.
const findingTrafficGuidance = "\n\n**漏洞流量证据(可选)**:调用 report_finding 上报漏洞时,如有已查看并确认支持漏洞结论的 HTTP 请求/响应,可用 traffic_refs 按复现顺序绑定真实 ID;域名和时间只作候选筛选,不推定关联。TCP 等非 HTTP 漏洞、未采集或无确切匹配时省略或传 [],在 evidence 保留命令输出、日志等其他可验证证据,建议说明未绑定原因。不要猜测 ID,也不要仅为补包重复探测。"
func (t *ToolSet) SetFindingRecorder(r FindingRecorder) { t.findingRecorder = r }
func (w *Worker) SetFindingRecorder(r FindingRecorder) { w.findingRecorder = r }
func (p *Planner) SetFindingRecorder(r FindingRecorder) { p.findingRecorder = r }
func (m *MainAgent) SetFindingRecorder(r FindingRecorder) { m.findingRecorder = r }
+108
View File
@@ -0,0 +1,108 @@
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) {
old := FindingTrafficBindingEnabled
FindingTrafficBindingEnabled = func() bool { return true }
t.Cleanup(func() { FindingTrafficBindingEnabled = old })
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")
}
}
+130
View File
@@ -0,0 +1,130 @@
package agent
import (
"encoding/json"
"fmt"
"github.com/Autumn-27/artex/db"
actool "github.com/Autumn-27/norma/tool"
)
const findingIDGuidance = "\n\n**漏洞编号约定**:finding_id 是独立漏洞记录 ID;finding_node_id 是探索节点 ID。list_findings / list_task_findings / node_detail / get_task_node_detail 的 id 保留为探索节点 ID,应从同一返回的 finding_id 读取独立编号。get_finding_traffic / bind_finding_traffic 用独立 finding_id。旧 update_finding_report 的 finding_id 参数仍传 finding_node_id。不要把 report_finding 第一行的数字用于证据工具,也不要遇到编号错误后猜测其他数字。"
// The server supplies the persisted setting. A missing setting/host is off.
// Consulted at assembly and again on writes so an already-running session
// cannot keep binding after the user switches the feature off.
var FindingTrafficBindingEnabled func() bool
func findingTrafficBindingEnabled() bool {
return FindingTrafficBindingEnabled != nil && FindingTrafficBindingEnabled()
}
const findingTrafficDisabled = "Agent 自动绑定流量已关闭;请在系统设置开启后重试,或使用页面人工绑定。未登记本次携带流量的漏洞。无流量上报可省略 traffic_refs / evidence_hint_id。"
// Applied after ToolResolve: user descriptions and prompts remain intact, while
// all actual reporters (including Planner and custom chat agents) see the same
// API contract. Disabled/unbound tools are never reintroduced here.
func findingWorkflowTools(agentKey string, tools []actool.CoreTool) ([]actool.CoreTool, string) {
if !findingTrafficBindingEnabled() {
out := make([]actool.CoreTool, 0, len(tools))
for _, tool := range tools {
if tool.Name() == "bind_finding_traffic" {
continue
}
if agentKey == "reporter" && (tool.Name() == "traffic_search" || tool.Name() == "traffic_get" || tool.Name() == "traffic_blob") {
continue
}
switch tool.Name() {
case "report_finding", "add_hint", "add_task_hint":
// Work on a copy: toggling back on must restore the original schema.
raw, _ := json.Marshal(tool.InputSchema())
var schema map[string]any
if json.Unmarshal(raw, &schema) == nil {
stripTrafficParameters(schema)
tool = DecorateTool(tool, tool.Description(), schema)
}
}
out = append(out, tool)
}
return out, ""
}
out := append([]actool.CoreTool(nil), tools...)
has := map[string]bool{}
for i, tool := range out {
has[tool.Name()] = true
note := ""
switch tool.Name() {
case "report_finding":
note = "\n默认由报告 Agent 在编写报告前核对并绑定流量。上报者在 evidence 中保留验证命令、关键输出、已有的真实流量 ID 及其用途,供报告 Agent 对照执行记录核实;无需为绑定额外查包。兼容显式即时绑定:traffic_refs 或 evidence_hint_id 可提交已核实的引用,后者读取本任务指定 hint 的结构化引用;任一无效则本次上报全部失败。TCP/无包不需要这些可选参数。返回 finding_id 与 finding_node_id 分别表示独立记录和探索节点。"
case "add_hint", "add_task_hint":
note = "\n交接已确认漏洞时,在对应提示的 traffic_refs 中保留已核实流量的 ID、用途、说明和顺序(单条放顶层,批量放对应 hints 元素),并在 text 中说明它证明的具体漏洞。调用方不能只交接文字而丢弃已有流量引用。未核实的候选不能作为证据传递。"
case "get_finding_traffic", "bind_finding_traffic", "list_findings", "list_task_findings", "node_detail", "get_task_node_detail", "update_finding_report":
note = findingIDGuidance
}
if note != "" {
out[i] = DecorateTool(tool, tool.Description()+note, tool.InputSchema())
}
}
guidance := ""
if has["report_finding"] || has["add_task_hint"] || has["add_hint"] {
guidance = "\n\n**流量证据交接(可选)**:自动绑定默认由报告 Agent 在漏洞入库后、编写报告前完成。上报者应在 evidence 保留验证命令、关键输出、已有真实流量 ID 及其用途,任务中带 intent_id,便于报告 Agent 追溯;不必为了绑定额外查包。Auto / Planner 代为上报时不要丢弃执行者已有的引用。add_hint / add_task_hint 可用 traffic_refs 交接;显式即时绑定仍兼容 report_finding 的 traffic_refs / evidence_hint_id。TCP 或无包时正常登记,不能猜测 ID,也不能仅为补包重复探测。"
if has["add_task_hint"] && !has["add_hint"] {
guidance += "\n平台对话没有任务上下文时,不直接调用 report_finding;通过 add_task_hint 向已有对应任务交接,由任务 Agent 登记,并用 list_task_findings 核对结果。"
}
if has["prove_goal"] || has["goal_met"] {
guidance += "\n判定目标完成前,先完成本次已有证据的上报/交接。不要在证据交接尚未完成时仅因文字漏洞已登记就结束任务、取消 Worker;无包不要求等待或强行抓包。"
}
}
if has["update_finding_report"] && has["bind_finding_traffic"] && has["get_finding_traffic"] {
guidance += "\n\n**报告前自动关联流量(已开启)**:你负责为本次触发的漏洞核对并绑定流量,再撰写报告。先从 report_finding 返回 JSON 或 get_task_node_detail / list_task_findings 取得明确的 finding_id 与 finding_node_id。读取漏洞详情、对应意图的执行记录及已有证据清单,优先使用上报者交接的真实 ID。若本次验证为 HTTP 且流量工具可用,用 traffic_search 筛选候选,再用 traffic_get 逐条核实请求/响应确实支持该漏洞;域名和时间只用于筛选,不证明归属。将确认的证据按复现顺序用 bind_finding_traffic(finding_id, traffic_refs) 关联,选择 baseline / proof / verification / supporting 并说明用途。只能操作本次漏洞,不重复创建漏洞或重新探测目标。绑定成功后重新调用 get_finding_traffic 获取最新 version,读取所需正文,再将实际读取的 version 作为 evidence_version 传给 update_finding_report(其 finding_id 参数仍用 finding_node_id)。已有绑定不必重复追加。TCP、未采集、工具不可用或没有确切匹配时,跳过自动绑定,依据文字/命令证据正常写报告并说明原因,不得为凑齐流量而猜测。绑定失败不宣称成功;保留已有证据并在报告说明未绑定原因。"
}
if guidance != "" || has["get_finding_traffic"] || has["update_finding_report"] {
guidance += findingIDGuidance
}
return out, guidance
}
func stripTrafficParameters(schema map[string]any) {
props, _ := schema["properties"].(map[string]any)
delete(props, "traffic_refs")
delete(props, "evidence_hint_id")
if required, ok := schema["required"].([]any); ok {
kept := required[:0]
for _, key := range required {
if key != "traffic_refs" && key != "evidence_hint_id" {
kept = append(kept, key)
}
}
schema["required"] = kept
}
if hints, ok := props["hints"].(map[string]any); ok {
if items, ok := hints["items"].(map[string]any); ok {
stripTrafficParameters(items)
}
}
}
// HintTrafficSchema is shared by the task-local and cross-task hint tools.
func HintTrafficSchema() map[string]any {
return map[string]any{"type": "array", "description": "可选:已核实且对应本提示中具体漏洞的流量引用,保留顺序;交接后 report_finding 可传 evidence_hint_id 携带这些引用。", "items": obj(map[string]any{"traffic_id": str("真实流量 ID"), "role": str("baseline / proof / verification / supporting"), "note": str("该流量支持什么结论")}, "traffic_id")}
}
func (t *ToolSet) findingRefsFromHint(hintID int64, explicit []db.TrafficRef) ([]db.TrafficRef, error) {
if hintID <= 0 {
return db.NormalizeTrafficRefs(explicit)
}
n, err := t.ts.GetNode(hintID) // local store only: inherited hints cannot supply evidence
if err != nil {
return nil, err
}
if n == nil || n.Kind != db.KindHint {
return nil, fmt.Errorf("evidence_hint_id=%d 必须是本任务的提示节点(继承提示不可直接用于绑定)", hintID)
}
var payload struct {
Refs []db.TrafficRef `json:"traffic_refs"`
}
if err := json.Unmarshal(n.Payload, &payload); err != nil {
return nil, err
}
return db.NormalizeTrafficRefs(append(append([]db.TrafficRef{}, explicit...), payload.Refs...))
}
+68
View File
@@ -0,0 +1,68 @@
package agent
import (
"context"
"encoding/json"
"strings"
"testing"
actool "github.com/Autumn-27/norma/tool"
)
func TestFindingWorkflowSharedAssemblyAndSwitch(t *testing.T) {
oldPolicy, oldResolve, oldAugment := FindingTrafficBindingEnabled, ToolResolve, ToolAugment
t.Cleanup(func() { FindingTrafficBindingEnabled, ToolResolve, ToolAugment = oldPolicy, oldResolve, oldAugment })
ToolAugment = nil
on := false
FindingTrafficBindingEnabled = func() bool { return on }
ts := NewToolSet(nil, "fixture")
ToolResolve = func(_ context.Context, _ string, base []actool.CoreTool) []actool.CoreTool {
out := make([]actool.CoreTool, len(base))
for i, tool := range base {
out[i] = DecorateTool(tool, "CUSTOM DESCRIPTION", tool.InputSchema())
}
return out
}
for _, role := range []string{"worker", "planner", "mainagent", "auto", "pentest", "custom-agent"} {
for _, enabled := range []bool{false, true, false} {
on = enabled
out, def, cleanup := AugmentTools(t.Context(), role, []actool.CoreTool{ts.addFinding(), ts.addHint()})
cleanup()
system, _ := deferredSystem("USER CUSTOM PROMPT", def)
if !strings.HasPrefix(system[0], "USER CUSTOM PROMPT") {
t.Fatal("custom prompt replaced")
}
if strings.Contains(system[0], "traffic_refs") != on {
t.Fatalf("%s: guidance ignored switch: %v", role, on)
}
if on && (!strings.Contains(system[0], "TCP") || !strings.Contains(system[0], "evidence_hint_id")) {
t.Fatal("missing optional/handoff contract")
}
for _, tool := range out {
if !strings.HasPrefix(tool.Description(), "CUSTOM DESCRIPTION") {
t.Fatal("custom description replaced")
}
raw, _ := json.Marshal(tool.InputSchema())
var schema map[string]any
json.Unmarshal(raw, &schema)
props := schema["properties"].(map[string]any)
if (props["traffic_refs"] != nil) != on {
t.Fatalf("%s: binding schema ignored switch", role)
}
if tool.Name() == "add_hint" {
nested := props["hints"].(map[string]any)["items"].(map[string]any)["properties"].(map[string]any)
if (nested["traffic_refs"] != nil) != on {
t.Fatal("nested hint schema ignored switch")
}
}
}
}
}
on = true
ToolResolve = func(context.Context, string, []actool.CoreTool) []actool.CoreTool { return nil }
out, def, cleanup := AugmentTools(t.Context(), "planner", []actool.CoreTool{ts.addFinding()})
defer cleanup()
if len(out) != 0 || def.FindingGuidance != "" {
t.Fatal("reintroduced disabled tool or its guidance")
}
}
+16 -14
View File
@@ -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
+2
View File
@@ -20,6 +20,7 @@ import (
// assets, judges whether the task goal is met, and emits 0..N exploration intents
// into the frontier. It is the sole intent generator.
type Planner struct {
findingRecorder FindingRecorder
prov llm.Provider
model string
tx *transcript.Store // raw LLM conversation persistence (nil = off)
@@ -338,6 +339,7 @@ func (p *Planner) Plan(ctx context.Context, taskID int64, as *db.AssetStore, ts
// runs in a detached goroutine so it never adds latency to this round.
p.compactor.OnPlannerRound(ctx, ts)
tsx := NewToolSet(ts, "planner")
tsx.SetFindingRecorder(p.findingRecorder)
if as != nil {
tsx.SetAssetStore(as, as.Companies())
}
+3
View File
@@ -57,6 +57,9 @@ func TestRenderSystemOverrideAndFallback(t *testing.T) {
if !strings.Contains(recording, "traffic_search") {
t.Fatalf("worker while recording should inject trafficTool: %q", recording)
}
if strings.Contains(recording, "traffic_refs") {
t.Fatalf("worker bypassed shared optional evidence policy: %q", recording)
}
if !strings.Contains(recording, "中间产物输出规约") {
t.Fatalf("worker missing artifact tail: %q", recording)
}
+10 -7
View File
@@ -52,7 +52,9 @@ const pentestDefaultTmpl = `你是一个授权渗透测试系统的"独立渗透
- 每得出一个结果**立刻**落地,别攒到最后(会话步数耗尽就全丢;记下来的才算数,活在脑子里的不算)。这些记录也是你抗 compaction 的长期记忆。
- **只写增量**:写之前扫一眼已登记的资产/已记的路线,只记你**新得到**的东西,别把已有内容换措辞重记(重复只会膨胀、也误导你自己以为有新进展)。只是印证已有结论而无新增,就不必再记。
- **发现新资产/入口** → insert_assets(资产本身:endpoint/parameter/tech 指纹/service/凭据/子域等,结构化属性写在资产 props 上)。回看已登记资产用 list_assets,避免重复登记。
- **确认漏洞** → report_finding(含可复现 PoC)。**只有你在本次运行里真实触发过、拿到可复现证据(请求/响应或命令输出)才用它**;回看已报漏洞用 list_findings。严禁把仅凭版本/CVE 匹配、"看起来可注入"、外部漏洞库/更新日志/代码 diff 推断的东西当已确认漏洞上报。**不要用查 CVE 库或"对比补丁版本"替代实际触发**;触发不了但有嫌疑,就在 TodoWrite 里标为"存疑/待验证",别硬记成 finding。
- **确认漏洞** → report_finding(含可复现 PoC)。**只有你在本次运行里真实触发过、拿到可复现证据(请求/响应或命令输出)才用它**;回看已报漏洞用 list_findings。有对应录制流量时,先 traffic_search / traffic_get 核对真实记录,再用 traffic_refs 按复现顺序绑定;域名和时间只用于候选筛选,不代表任务归属。严禁把仅凭版本/CVE 匹配、"看起来可注入"、外部漏洞库/更新日志/代码 diff 推断的东西当已确认漏洞上报。**不要用查 CVE 库或"对比补丁版本"替代实际触发**;触发不了但有嫌疑,就在 TodoWrite 里标为"存疑/待验证",别硬记成 finding。
流量绑定可选:TCP 等非 HTTP 漏洞、未采集或无确切匹配记录时,省略 traffic_refs 或传 [],在 evidence 保留命令输出、日志等其他可验证证据,建议说明未绑定原因。不要猜测 ID,也不要仅为补包重复探测。
━━ 判定与收尾 ━━
- 随时对照任务目标:已被你**验证过**的成果满足了目标,就据此判定达成并说明依据。判"达成"的前提是心法 4 的自检已通过——没独立复现过的战果不算达成依据。
@@ -77,15 +79,16 @@ const ReporterDefaultPrompt = `你是一个授权渗透测试系统里的**漏
每当有 worker 调用 report_finding 登记了一个漏洞,系统就会用一段【由工具调用触发】的上下文唤起你,其中包含:
- **任务 id**(task_id,见上下文"任务: #<id>")
- report_finding 的**入参**(vulnclass / severity / summary / evidence 等)
- report_finding 的**返回**:形如 "finding recorded: <id>" —— 这个 **<id> 就是本漏洞的 finding_id**,是你后面保存报告的句柄。
- report_finding 的**返回**:形如 "finding recorded: <id>" —— 这个 **<id> 是探索节点 ID**,是 get_task_node_detail 和 update_finding_report 使用的旧句柄。返回 JSON 中的 finding_id 则是独立漏洞记录 ID,get_finding_traffic 使用它。
先从上下文里**准确抽取 finding_id 和 task_id**。抽取不到 finding_id 就不要瞎写,说明情况即可。
先从上下文里**准确抽取 task_id、探索节点 node_id,以及 JSON 中的独立漏洞 finding_id(如有)**,不得混用两种 ID。抽取不到 node_id 就不要瞎写,说明情况即可。
━━ 工作步骤 ━━
1. **取全证据**:用 get_task_node_detail(task_id, id=<finding_id>) 读该漏洞节点的**完整证据/PoC**(触发上下文里的 evidence 可能被截断)。
2. **还原过程**:用 list_task_worker_traces(task_id) 找到相关的 work,再用 get_task_worker_trace(task_id, intent_id[, step_ids]) 或 search_task_worker_traces(task_id, q) 看这个漏洞**是怎么被发现和验证的**(用了什么请求/命令、目标怎么响应)。必要时 get_task_graph(task_id) 看整体态势、list_task_findings(task_id) 看是否有关联漏洞。
3. **写报告**:综合以上,写一份结构化 Markdown 报告(见下方模板)。
4. **保存**:调用 **update_finding_report(finding_id=<那个 id>, report=<Markdown 全文>)** 保存。这是你的最终产物——不写进去等于没做。
1. **取全证据**:用 get_task_node_detail(task_id, id=<node_id>) 读该漏洞节点的**完整证据/PoC**(触发上下文里的 evidence 可能被截断)。
2. **流量证据**:如返回 JSON 包含独立 finding_id,用 get_finding_traffic 先读有序清单及 version,有绑定时再按 binding_id 分段读取请求/响应。绑定可选,空清单不阻止撰写报告:TCP 等非 HTTP 漏洞或未采集的情况,依据节点证据、命令输出和日志说明复现与影响,建议如实说明未绑定原因,不虚构请求/响应,不仅为补包重新探测。报告引用稳定证据编号及用途;仅按真实内容描述。保存报告时传入所读 version 作为 evidence_version;如版本冲突,重新读取并生成,不得直接换版本重试。
3. **还原过程**:用 list_task_worker_traces(task_id) 找到相关的 work,再用 get_task_worker_trace(task_id, intent_id[, step_ids]) 或 search_task_worker_traces(task_id, q) 看这个漏洞**是怎么被发现和验证的**(用了什么请求/命令、目标怎么响应)。必要时 get_task_graph(task_id) 看整体态势、list_task_findings(task_id) 看是否有关联漏洞。
4. **写报告**:综合以上,写一份结构化 Markdown 报告(见下方模板)。
5. **保存**:调用 **update_finding_report(finding_id=<node_id>, report=<Markdown 全文>, evidence_version=<实际读取的 version>)** 保存;未读取版本时省略 evidence_version,不得猜测。这是你的最终产物——不写进去等于没做。
━━ 报告结构(Markdown,按需裁剪,但证据/复现/修复必须有)━━
- ` + "`## 概述`" + `:一句话说清是什么漏洞、在哪、能造成什么。
+111 -72
View File
@@ -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
@@ -398,7 +399,11 @@ func (t *ToolSet) graphOverviewData() map[string]any {
for _, h := range hints {
var p map[string]any
_ = json.Unmarshal(h.Payload, &p)
hsum = append(hsum, map[string]any{"id": h.ID, "state": h.State, "text": p["text"]})
hint := map[string]any{"id": h.ID, "state": h.State, "text": p["text"]}
if findingTrafficBindingEnabled() && p["traffic_refs"] != nil {
hint["traffic_refs"] = p["traffic_refs"]
}
hsum = append(hsum, hint)
}
out["hints"] = hsum
// lineage from the exploration edges: an intent's parents (what it
@@ -1068,6 +1073,9 @@ func (t *ToolSet) listFindings() actool.CoreTool {
obj(map[string]any{}),
func(context.Context, json.RawMessage) (actool.Result, error) {
f, _ := t.ts.ListByKindWithSources(db.KindFinding, 500)
if err := t.ts.PopulateFindingTrafficIDs(f); err != nil {
return actool.Errorf(err.Error()), nil
}
intentOf, _ := t.ts.FindingIntentsWithSources() // finding id -> 产生它的 intent id
taskID := t.taskID
if taskID <= 0 {
@@ -1076,6 +1084,9 @@ func (t *ToolSet) listFindings() actool.CoreTool {
out := make([]map[string]any, 0, len(f))
for _, n := range f {
m := compactFinding(n)
if n.FindingID > 0 {
m["finding_id"], m["finding_node_id"], m["traffic_count"] = n.FindingID, n.ID, n.TrafficCount
}
if n.Inherited {
m["task_id"] = n.SourceTaskID
} else {
@@ -1167,7 +1178,10 @@ func (t *ToolSet) nodeDetail() actool.CoreTool {
if n == nil {
return actool.Errorf(fmt.Sprintf("未找到探索节点 %d。若你想查的是资产,请用 list_assets / asset_neighbors(资产与探索节点是不同的 id 空间,资产 id 不能传给 node_detail)。", id)), nil
}
return jsonResult(n) // full payload incl. detail / evidence
if err := t.ts.PopulateFindingTrafficIDs([]*db.Node{n}); err != nil {
return actool.Errorf(err.Error()), nil
}
return jsonResult(n) // full payload incl. detail / evidence, plus explicit finding IDs
})
}
@@ -1358,67 +1372,79 @@ 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", "记录确认的漏洞,用 evidence 提供命令输出、日志等可验证证据。任务上下文传当前 intent_id。返回的 finding_id 是独立漏洞记录 ID,finding_node_id 是探索节点 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 文本"),
"evidence_hint_id": idp("可选:本任务中对应此漏洞的提示节点 ID,自动携带其结构化 traffic_refs;不能引用继承提示或其他漏洞的提示"),
"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"`
EvidenceHintID json.RawMessage `json:"evidence_hint_id"`
}
if err := json.Unmarshal(in, &a); err != nil {
return actool.Errorf(err.Error()), nil
}
if t.ts == nil {
return actool.Errorf("report_finding 需要任务上下文;平台对话请通过 add_task_hint 向对应任务交接漏洞,并在提示中携带已有的 traffic_refs,由任务 Agent 登记。已登记漏洞可用 bind_finding_traffic 补绑。"), nil
}
if !findingTrafficBindingEnabled() && (len(a.TrafficRefs) > 0 || pid(a.EvidenceHintID) > 0) {
return actool.Errorf(findingTrafficDisabled), nil
}
if len(a.EvidenceHintID) > 0 && pid(a.EvidenceHintID) <= 0 {
return actool.Errorf("evidence_hint_id 必须为有效的提示节点 ID;无交接提示时省略"), nil
}
refs, err := t.findingRefsFromHint(pid(a.EvidenceHintID), 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)
}
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 = ""
}
result := struct {
*db.RecordedFinding
EvidenceStatus string `json:"evidence_status"`
EvidenceNote string `json:"evidence_note,omitempty"`
}{RecordedFinding: recorded, EvidenceStatus: "bound"}
if len(recorded.Traffic.Bindings) == 0 {
result.EvidenceStatus = "not_bound"
result.EvidenceNote = "漏洞已保存,未绑定流量。TCP/无包情形可正常继续;若已有核实的 HTTP 流量,请用可用的 bind_finding_traffic 或漏洞页面补绑,再完成证据交接。不要重复创建漏洞。"
if !findingTrafficBindingEnabled() {
result.EvidenceNote = "漏洞已保存。Agent 自动绑定流量已关闭,可在页面人工关联流量。"
}
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
})
}
raw, _ := json.Marshal(result)
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
@@ -1530,12 +1556,16 @@ func (t *ToolSet) recordFact() actool.CoreTool {
}
type hintItem struct {
Text string `json:"text"`
AssetIDs []json.RawMessage `json:"asset_ids"`
Text string `json:"text"`
AssetIDs []json.RawMessage `json:"asset_ids"`
TrafficRefs []db.TrafficRef `json:"traffic_refs"`
}
// addOneHint 挂一条 hint 节点(active/human)到探索图,可锚定资产,返回 id。
func (t *ToolSet) addOneHint(it hintItem) (int64, error) {
if len(it.TrafficRefs) > 0 && !findingTrafficBindingEnabled() {
return 0, fmt.Errorf("Agent 自动绑定流量已关闭,未保存携带 traffic_refs 的提示;可在系统设置开启,或仅交接文字")
}
if strings.TrimSpace(it.Text) == "" {
return 0, fmt.Errorf("text 不能为空")
}
@@ -1545,7 +1575,15 @@ func (t *ToolSet) addOneHint(it hintItem) (int64, error) {
anchors = append(anchors, tid)
}
}
id, err := t.ts.AddNode(db.KindHint, map[string]any{"text": it.Text}, 0, "active", "human", anchors)
refs, err := db.NormalizeTrafficRefs(it.TrafficRefs)
if err != nil {
return 0, err
}
payload := map[string]any{"text": it.Text}
if len(refs) > 0 {
payload["traffic_refs"] = refs
}
id, err := t.ts.AddNode(db.KindHint, payload, 0, "active", "human", anchors)
if err == nil && t.notify != nil {
t.notify() // wake the planner so the new hint is read promptly (debounced)
}
@@ -1731,9 +1769,10 @@ func (t *ToolSet) addHint() actool.CoreTool {
return writeTool("add_hint", "把人类/主 agent 的战略提示挂到探索图,规划者下次生成意图时会读到它。\n"+
"★优先批量:多条提示放进 hints 数组一次提交(比逐条调用省往返)。返回 ids 数组,与 hints 等长同序(失败项 id=0,详情见 errors)。单条则省略 hints 直接给顶层 text。",
obj(map[string]any{
"hints": map[string]any{"type": "array", "description": "【优先用这个】要新增的提示数组,按顺序处理。每个元素字段同下方顶层字段(text/asset_ids)。返回 ids 与本数组等长、同序。", "items": map[string]any{"type": "object"}},
"text": str("[单条] 提示内容,如'重点挖认证后接口'"),
"asset_ids": map[string]any{"type": "array", "items": map[string]any{"type": "integer"}, "description": "锚定的资产 id(可选,0/1/多个)"},
"hints": map[string]any{"type": "array", "description": "【优先用这个】要新增的提示数组,按顺序处理。每个元素字段同下方顶层字段(text/asset_ids/traffic_refs)。返回 ids 与本数组等长、同序。", "items": obj(map[string]any{"text": str("提示内容"), "asset_ids": map[string]any{"type": "array", "items": map[string]any{"type": "integer"}}, "traffic_refs": HintTrafficSchema()})},
"text": str("[单条] 提示内容,如'重点挖认证后接口'"),
"traffic_refs": HintTrafficSchema(),
"asset_ids": map[string]any{"type": "array", "items": map[string]any{"type": "integer"}, "description": "锚定的资产 id(可选,0/1/多个)"},
}),
func(_ context.Context, in json.RawMessage) (actool.Result, error) {
var a struct {
+17 -12
View File
@@ -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,6 +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.
// Optional finding guidance is added for every role after tool resolution.
return body + workerTrafficBlock(caCert != "") + workerArtifactSpec(runDir)
}
@@ -325,6 +329,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)
+3
View File
@@ -31,6 +31,9 @@ func utf8Clean(s string) string {
// Node is a typed reasoning node (= old task_nodes). kind ∈ goal|intent|finding|hint.
type Node struct {
FindingID int64 `json:"finding_id,omitempty"` // populated by finding-aware reads; never inferred from ID
FindingNodeID int64 `json:"finding_node_id,omitempty"`
TrafficCount int `json:"traffic_count,omitempty"`
ID int64 `json:"id"`
Kind string `json:"kind"`
Payload json.RawMessage `json:"payload"`
+438
View File
@@ -0,0 +1,438 @@
package db
import (
"context"
"crypto/sha256"
"database/sql"
"encoding/hex"
"encoding/json"
"errors"
"fmt"
"strings"
"time"
)
var (
ErrEvidenceConflict = errors.New("流量证据已变更,请刷新后重试")
ErrFindingNotFound = errors.New("漏洞不存在")
ErrEvidenceNotFound = errors.New("流量证据不存在")
)
// This lock covers the evidence filesystem as well as its SQL references. All
// processes sharing the database use it, including readers, exports and GC.
const findingEvidenceLockKey int64 = 7337741004
func (d *DB) WithEvidenceTx(ctx context.Context, fn func(*sql.Tx) error) error {
tx, err := d.BeginTx(ctx, nil)
if err != nil {
return err
}
defer tx.Rollback()
if _, err = tx.ExecContext(ctx, `SELECT pg_advisory_xact_lock($1)`, findingEvidenceLockKey); err != nil {
return err
}
if err = fn(tx); err != nil {
return err
}
return tx.Commit()
}
type TrafficRef struct {
TrafficID string `json:"traffic_id"`
Role string `json:"role"`
Note string `json:"note"`
}
func NormalizeTrafficRefs(refs []TrafficRef) ([]TrafficRef, error) {
out := make([]TrafficRef, 0, len(refs))
seen := map[string]bool{}
for _, ref := range refs {
ref.TrafficID = strings.TrimSpace(ref.TrafficID)
if ref.TrafficID == "" {
return nil, errors.New("traffic_id 不能为空")
}
if ref.Role == "" {
ref.Role = "supporting"
}
if !ValidTrafficRole(ref.Role) {
return nil, fmt.Errorf("无效的流量用途 %q", ref.Role)
}
if !seen[ref.TrafficID] {
out = append(out, ref)
seen[ref.TrafficID] = true
}
}
return out, nil
}
func ValidTrafficRole(role string) bool {
return role == "baseline" || role == "proof" || role == "verification" || role == "supporting"
}
type TrafficEvidenceSnapshot struct {
ID string `json:"id"`
SourceTrafficID string `json:"source_traffic_id"`
CapturedAt int64 `json:"captured_at"`
URL string `json:"url"`
Method string `json:"method"`
Status int `json:"status"`
ContentType string `json:"content_type"`
ReqHead string `json:"req_head,omitempty"`
RespHead string `json:"resp_head,omitempty"`
ReqHash string `json:"req_hash"`
RespHash string `json:"resp_hash"`
ReqLen int64 `json:"req_len"`
RespLen int64 `json:"resp_len"`
CreatedAt time.Time `json:"created_at"`
}
type FindingTrafficBinding struct {
ID int64 `json:"id,string"`
FindingID int64 `json:"finding_id,string"`
SnapshotID string `json:"snapshot_id"`
Role string `json:"role"`
Note string `json:"note"`
Position int `json:"position"`
CreatedAt time.Time `json:"created_at"`
Snapshot TrafficEvidenceSnapshot `json:"snapshot"`
}
type FindingTraffic struct {
FindingID int64 `json:"finding_id,string"`
Version int64 `json:"version"`
ReportVersion int64 `json:"report_version"`
Bindings []FindingTrafficBinding `json:"bindings"`
}
type PreparedTrafficEvidence struct {
Ref TrafficRef
Snapshot TrafficEvidenceSnapshot
}
func TrafficSnapshotID(snapshot TrafficEvidenceSnapshot) string {
snapshot.ID = ""
snapshot.CreatedAt = time.Time{}
raw, _ := json.Marshal(snapshot)
sum := sha256.Sum256(raw)
return hex.EncodeToString(sum[:])
}
// LockTaskEvidenceTx serializes writes with archive queueing (which locks the
// same task row). Once queued, its snapshot must not acquire new evidence.
func LockTaskEvidenceTx(tx *sql.Tx, taskID int64) error {
if taskID == 0 {
return nil
}
var deleted sql.NullTime
if err := tx.QueryRow(`SELECT deleted_at FROM tasks WHERE id=$1 FOR UPDATE`, taskID).Scan(&deleted); err != nil {
return err
}
if deleted.Valid {
return ErrTaskArchiveState
}
var state string
err := tx.QueryRow(`SELECT state FROM task_archives WHERE task_id=$1`, taskID).Scan(&state)
if err != nil && !errors.Is(err, sql.ErrNoRows) {
return err
}
if err == nil && state != ArchiveFailed {
return ErrTaskArchiveState
}
return nil
}
func LockFindingEvidenceTx(tx *sql.Tx, findingID int64, version *int64) error {
var taskID sql.NullInt64
if err := tx.QueryRow(`SELECT task_id FROM findings WHERE id=$1`, findingID).Scan(&taskID); err != nil {
if errors.Is(err, sql.ErrNoRows) {
return ErrFindingNotFound
}
return err
}
if err := LockTaskEvidenceTx(tx, taskID.Int64); err != nil {
return err
}
var current int64
if err := tx.QueryRow(`SELECT evidence_version FROM findings WHERE id=$1 FOR UPDATE`, findingID).Scan(&current); err != nil {
if errors.Is(err, sql.ErrNoRows) {
return ErrFindingNotFound
}
return err
}
if version != nil && *version != current {
return ErrEvidenceConflict
}
return nil
}
func InsertEvidenceSnapshotTx(tx *sql.Tx, s TrafficEvidenceSnapshot) error {
if s.ID != TrafficSnapshotID(s) {
return errors.New("证据快照元数据哈希不匹配")
}
if s.CreatedAt.IsZero() {
s.CreatedAt = time.Now().UTC()
}
_, err := tx.Exec(`INSERT INTO traffic_evidence_snapshots
(id,source_traffic_id,captured_at,url,method,status,content_type,req_head,resp_head,req_hash,resp_hash,req_len,resp_len,created_at)
VALUES($1,$2,$3,$4,$5,$6,$7,$8,$9,$10,$11,$12,$13,$14) ON CONFLICT(id) DO NOTHING`,
s.ID, s.SourceTrafficID, s.CapturedAt, s.URL, s.Method, s.Status, s.ContentType, s.ReqHead, s.RespHead, s.ReqHash, s.RespHash, s.ReqLen, s.RespLen, s.CreatedAt)
return err
}
func AddFindingTrafficTx(tx *sql.Tx, findingID int64, prepared []PreparedTrafficEvidence) error {
var position int
if err := tx.QueryRow(`SELECT COALESCE(MAX(position)+1,0) FROM finding_traffic_bindings WHERE finding_id=$1`, findingID).Scan(&position); err != nil {
return err
}
changed := false
for _, item := range prepared {
if err := InsertEvidenceSnapshotTx(tx, item.Snapshot); err != nil {
return err
}
if _, err := tx.Exec(`UPDATE traffic_evidence_snapshots SET unreferenced_at=NULL WHERE id=$1`, item.Snapshot.ID); err != nil {
return err
}
res, err := tx.Exec(`INSERT INTO finding_traffic_bindings(finding_id,snapshot_id,role,note,position)
VALUES($1,$2,$3,$4,$5) ON CONFLICT(finding_id,snapshot_id) DO NOTHING`, findingID, item.Snapshot.ID, item.Ref.Role, item.Ref.Note, position)
if err != nil {
return err
}
n, err := res.RowsAffected()
if err != nil {
return err
}
if n > 0 {
changed = true
position++
}
}
if changed {
return bumpEvidenceVersionTx(tx, findingID)
}
return nil
}
func bumpEvidenceVersionTx(tx *sql.Tx, findingID int64) error {
_, err := tx.Exec(`UPDATE findings SET evidence_version=evidence_version+1 WHERE id=$1`, findingID)
return err
}
func FindingTrafficTx(tx *sql.Tx, findingID int64) (*FindingTraffic, error) {
out := &FindingTraffic{FindingID: findingID, Bindings: []FindingTrafficBinding{}}
if err := tx.QueryRow(`SELECT evidence_version,report_evidence_version FROM findings WHERE id=$1`, findingID).Scan(&out.Version, &out.ReportVersion); err != nil {
if errors.Is(err, sql.ErrNoRows) {
return nil, ErrFindingNotFound
}
return nil, err
}
rows, err := tx.Query(`SELECT b.id,b.finding_id,b.snapshot_id,b.role,b.note,b.position,b.created_at,to_jsonb(s)
FROM finding_traffic_bindings b JOIN traffic_evidence_snapshots s ON s.id=b.snapshot_id WHERE b.finding_id=$1 ORDER BY b.position,b.id`, findingID)
if err != nil {
return nil, err
}
defer rows.Close()
for rows.Next() {
var b FindingTrafficBinding
var raw []byte
if err := rows.Scan(&b.ID, &b.FindingID, &b.SnapshotID, &b.Role, &b.Note, &b.Position, &b.CreatedAt, &raw); err != nil {
return nil, err
}
if err := json.Unmarshal(raw, &b.Snapshot); err != nil {
return nil, err
}
out.Bindings = append(out.Bindings, b)
}
return out, rows.Err()
}
func (d *DB) GetFindingTraffic(ctx context.Context, findingID int64) (out *FindingTraffic, err error) {
err = d.WithEvidenceTx(ctx, func(tx *sql.Tx) error { var e error; out, e = FindingTrafficTx(tx, findingID); return e })
return
}
func (d *DB) EditFindingTraffic(ctx context.Context, findingID, bindingID, version int64, role, note *string, remove bool, order []int64) error {
if role != nil && !ValidTrafficRole(*role) {
return errors.New("无效的流量用途")
}
return d.WithEvidenceTx(ctx, func(tx *sql.Tx) error {
if err := LockFindingEvidenceTx(tx, findingID, &version); err != nil {
return err
}
if order != nil {
current, err := FindingTrafficTx(tx, findingID)
if err != nil {
return err
}
if len(order) != len(current.Bindings) {
return ErrEvidenceConflict
}
ids := map[int64]bool{}
for _, b := range current.Bindings {
ids[b.ID] = true
}
for position, id := range order {
if !ids[id] {
return ErrEvidenceConflict
}
delete(ids, id)
if _, err := tx.Exec(`UPDATE finding_traffic_bindings SET position=$1 WHERE finding_id=$2 AND id=$3`, position, findingID, id); err != nil {
return err
}
}
} else {
var res sql.Result
var err error
if remove {
res, err = tx.Exec(`DELETE FROM finding_traffic_bindings WHERE finding_id=$1 AND id=$2`, findingID, bindingID)
} else {
res, err = tx.Exec(`UPDATE finding_traffic_bindings SET role=COALESCE($3,role),note=COALESCE($4,note) WHERE finding_id=$1 AND id=$2`, findingID, bindingID, role, note)
}
if err != nil {
return err
}
n, err := res.RowsAffected()
if err != nil {
return err
}
if n == 0 {
return ErrEvidenceNotFound
}
}
return bumpEvidenceVersionTx(tx, findingID)
})
}
type RecordFindingInput struct {
TaskID, ExplorationID, IntentID int64
VulnClass, Name, Severity, Summary, Evidence, Worker string
AssetIDs []int64
}
type RecordedFinding struct {
FindingID int64 `json:"finding_id,string"`
NodeID int64 `json:"finding_node_id,string"`
Traffic *FindingTraffic `json:"traffic"`
}
func RecordFindingTx(tx *sql.Tx, in RecordFindingInput, prepared []PreparedTrafficEvidence) (*RecordedFinding, error) {
if err := LockTaskEvidenceTx(tx, in.TaskID); err != nil {
return nil, err
}
if in.TaskID > 0 {
var expID int64
if err := tx.QueryRow(`SELECT exploration_id FROM tasks WHERE id=$1`, in.TaskID).Scan(&expID); err != nil {
return nil, err
}
if expID != in.ExplorationID {
return nil, errors.New("漏洞所属任务与探索记录不匹配")
}
}
if in.IntentID > 0 {
var ok bool
if err := tx.QueryRow(`SELECT EXISTS(SELECT 1 FROM exploration_nodes WHERE id=$1 AND exploration_id=$2 AND kind='intent')`, in.IntentID, in.ExplorationID).Scan(&ok); err != nil {
return nil, err
}
if !ok {
return nil, errors.New("intent_id 必须是本任务的意图(关联任务意图只读)")
}
}
payload, _ := json.Marshal(map[string]any{"vulnclass": in.VulnClass, "name": in.Name, "severity": in.Severity, "summary": in.Summary, "evidence": map[string]any{"by": in.Worker, "poc": in.Evidence}})
out := &RecordedFinding{}
if err := tx.QueryRow(`INSERT INTO exploration_nodes(exploration_id,kind,payload,priority,state,origin)
VALUES($1,'finding',$2,9,'confirmed',$3) RETURNING id`, in.ExplorationID, string(payload), in.Worker).Scan(&out.NodeID); err != nil {
return nil, err
}
for _, asset := range in.AssetIDs {
if _, err := tx.Exec(`INSERT INTO exploration_anchors(node_id,asset_id) VALUES($1,$2) ON CONFLICT DO NOTHING`, out.NodeID, asset); err != nil {
return nil, err
}
}
if in.IntentID > 0 {
if _, err := tx.Exec(`INSERT INTO exploration_edges(exploration_id,src_id,rel,dst_id) VALUES($1,$2,$3,$4)`, in.ExplorationID, in.IntentID, RelYields, out.NodeID); err != nil {
return nil, err
}
}
assets := in.AssetIDs
if assets == nil {
assets = []int64{}
}
raw, _ := json.Marshal(assets)
if err := tx.QueryRow(`INSERT INTO findings(task_id,node_id,vulnclass,name,severity,summary,evidence,worker,asset_ids)
VALUES(NULLIF($1,0),$2,$3,$4,$5,$6,$7,$8,$9) RETURNING id`, in.TaskID, out.NodeID, in.VulnClass, in.Name, in.Severity, in.Summary, in.Evidence, in.Worker, string(raw)).Scan(&out.FindingID); err != nil {
return nil, err
}
if err := AddFindingTrafficTx(tx, out.FindingID, prepared); err != nil {
return nil, err
}
var err error
out.Traffic, err = FindingTrafficTx(tx, out.FindingID)
return out, err
}
// RecordFinding is the atomic legacy/no-recorder path used by tools and tests.
func (s *ExplorationStore) RecordFinding(ctx context.Context, in RecordFindingInput) (out *RecordedFinding, err error) {
in.ExplorationID = s.expID
err = s.db.WithEvidenceTx(ctx, func(tx *sql.Tx) error { var e error; out, e = RecordFindingTx(tx, in, nil); return e })
return
}
func (d *DB) FindingIDByNodeID(nodeID int64) (id int64, err error) {
err = d.QueryRow(`SELECT id FROM findings WHERE node_id=$1`, nodeID).Scan(&id)
return
}
// PopulateFindingTrafficIDs only enriches already-visible nodes. It performs no
// discovery or ID guessing, and leaves the legacy node ID unchanged.
func (s *ExplorationStore) PopulateFindingTrafficIDs(nodes []*Node) error {
byID := map[int64]*Node{}
var args []any
var placeholders []string
for _, n := range nodes {
if n == nil || n.Kind != KindFinding || byID[n.ID] != nil {
continue
}
byID[n.ID] = n
args = append(args, n.ID)
placeholders = append(placeholders, fmt.Sprintf("$%d", len(args)))
}
if len(args) == 0 {
return nil
}
rows, err := s.db.Query(`SELECT f.id,f.node_id,(SELECT count(*) FROM finding_traffic_bindings b WHERE b.finding_id=f.id) FROM findings f WHERE f.node_id IN (`+strings.Join(placeholders, ",")+`)`, args...)
if err != nil {
return err
}
defer rows.Close()
for rows.Next() {
var findingID, nodeID int64
var count int
if err := rows.Scan(&findingID, &nodeID, &count); err != nil {
return err
}
n := byID[nodeID]
n.FindingID, n.FindingNodeID, n.TrafficCount = findingID, nodeID, count
}
return rows.Err()
}
func (d *DB) SetFindingReportVersionByNodeID(ctx context.Context, nodeID int64, report string, version *int64) (n int64, err error) {
err = d.WithEvidenceTx(ctx, func(tx *sql.Tx) error {
var id int64
if err := tx.QueryRow(`SELECT id FROM findings WHERE node_id=$1`, nodeID).Scan(&id); err != nil {
if errors.Is(err, sql.ErrNoRows) {
return nil
}
return err
}
if err := LockFindingEvidenceTx(tx, id, version); err != nil {
return err
}
res, err := tx.Exec(`UPDATE findings SET report=$2,report_evidence_version=COALESCE($3::bigint,CASE WHEN evidence_version=0 THEN 0 ELSE -1 END) WHERE id=$1`, id, report, version)
if err != nil {
return err
}
n, err = res.RowsAffected()
return err
})
return
}
+85
View File
@@ -0,0 +1,85 @@
package db
import (
"database/sql"
"encoding/json"
"errors"
"fmt"
)
func ArchiveEvidenceSnapshots(snapshot *TaskArchiveSnapshot) ([]TrafficEvidenceSnapshot, error) {
var out []TrafficEvidenceSnapshot
if rawRowCount(snapshot.Tables["traffic_evidence_snapshots"]) == 0 {
return out, nil
}
if snapshot.FormatVersion < 3 {
return nil, ErrTaskArchiveFormatMismatch
}
err := json.Unmarshal(snapshot.Tables["traffic_evidence_snapshots"], &out)
return out, err
}
func restoreFindingTrafficTx(tx *sql.Tx, snapshot *TaskArchiveSnapshot) error {
snapshots, err := ArchiveEvidenceSnapshots(snapshot)
if err != nil {
return err
}
allowed := map[string]bool{}
for _, v := range snapshots {
if allowed[v.ID] {
return errors.New("duplicate archived evidence snapshot")
}
allowed[v.ID] = true
if err = InsertEvidenceSnapshotTx(tx, v); err != nil {
return err
}
}
rows, err := decodeArchiveRows(snapshot.Tables["finding_traffic_bindings"])
if err != nil {
return err
}
if len(rows) > 0 && snapshot.FormatVersion < 3 {
return ErrTaskArchiveFormatMismatch
}
for _, row := range rows {
fid, ok := jsonInt64(row["finding_id"])
if !ok {
return errors.New("invalid archived evidence finding id")
}
sid, _ := row["snapshot_id"].(string)
role, _ := row["role"].(string)
if !allowed[sid] || !ValidTrafficRole(role) {
return errors.New("invalid archived evidence binding")
}
var owned bool
if err = tx.QueryRow(`SELECT EXISTS(SELECT 1 FROM findings WHERE id=$1 AND task_id=$2)`, fid, snapshot.TaskID).Scan(&owned); err != nil {
return err
}
if !owned {
return fmt.Errorf("evidence references finding outside archived task: %d", fid)
}
}
if len(rows) > 0 {
raw, _ := json.Marshal(rows)
if _, err = tx.Exec(`INSERT INTO finding_traffic_bindings SELECT * FROM json_populate_recordset(NULL::finding_traffic_bindings,$1::json)`, string(raw)); err != nil {
return err
}
_, err = tx.Exec(`UPDATE traffic_evidence_snapshots s SET unreferenced_at=NULL WHERE EXISTS(SELECT 1 FROM finding_traffic_bindings b WHERE b.snapshot_id=s.id)`)
}
return err
}
func normalizeArchivedFindingVersions(raw json.RawMessage) (json.RawMessage, error) {
rows, err := decodeArchiveRows(raw)
if err != nil {
return nil, err
}
for _, row := range rows {
for _, key := range []string{"evidence_version", "report_evidence_version"} {
if row[key] == nil {
row[key] = 0
}
}
}
return json.Marshal(rows)
}
+87
View File
@@ -0,0 +1,87 @@
package db
import (
"encoding/json"
"fmt"
"testing"
)
func TestFindingEvidenceLockNamespace(t *testing.T) {
for _, reserved := range []int64{7337741001, 7337741002, 7337741003} {
if findingEvidenceLockKey == reserved {
t.Fatal("evidence lock collides with migration/test/company lock")
}
}
}
func TestFindingTrafficLegacyArchiveDefaults(t *testing.T) {
d, err := Open(testDSN(t))
if err != nil {
t.Fatal(err)
}
defer d.Close()
if err = d.EnsureLLMRecordsTable(); err != nil {
t.Fatal(err)
}
if err = d.EnsureLLMUsageTable(); err != nil {
t.Fatal(err)
}
for _, version := range []int{1, 2} {
t.Run(fmt.Sprint(version), func(t *testing.T) {
task, err := d.CreateTask("legacy archive evidence defaults", "fixture", nil, 0, 0)
if err != nil {
t.Fatal(err)
}
defer func() { d.Exec(`DELETE FROM task_archives WHERE task_id=$1`, task.ID); d.DeleteTask(task.ID) }()
f, err := d.Exploration(task.ExplorationID).RecordFinding(t.Context(), RecordFindingInput{TaskID: task.ID, ExplorationID: task.ExplorationID, Summary: "legacy", Severity: "low"})
if err != nil {
t.Fatal(err)
}
if err = d.SetPaused(task.ID, true); err != nil {
t.Fatal(err)
}
archive, err := d.QueueTaskArchive(task.ID)
if err != nil {
t.Fatal(err)
}
if _, err = d.ClaimTaskArchiveJob(t.Context()); err != nil {
t.Fatal(err)
}
snapshot, err := d.SnapshotTaskArchive(task.ID)
if err != nil {
t.Fatal(err)
}
if err = d.CompleteTaskArchive(archive.ID, snapshot, "/tmp/legacy-test.tar.zst", "fixture", 1, 1); err != nil {
t.Fatal(err)
}
snapshot.FormatVersion = version
delete(snapshot.Tables, "finding_traffic_bindings")
delete(snapshot.Tables, "traffic_evidence_snapshots")
var rows []map[string]any
if err = json.Unmarshal(snapshot.Tables["findings"], &rows); err != nil {
t.Fatal(err)
}
for _, row := range rows {
delete(row, "evidence_version")
delete(row, "report_evidence_version")
}
snapshot.Tables["findings"], _ = json.Marshal(rows)
if _, err = d.QueueTaskArchiveRestore(archive.ID); err != nil {
t.Fatal(err)
}
if _, err = d.ClaimTaskArchiveJob(t.Context()); err != nil {
t.Fatal(err)
}
if _, err = d.RestoreTaskArchive(archive.ID, snapshot, 0); err != nil {
t.Fatal(err)
}
got, err := d.GetFindingTraffic(t.Context(), f.FindingID)
if err != nil {
t.Fatal(err)
}
if got.Version != 0 || got.ReportVersion != 0 || len(got.Bindings) != 0 {
t.Fatal(got)
}
})
}
}
+37 -25
View File
@@ -1,6 +1,7 @@
package db
import (
"context"
"database/sql"
"encoding/json"
"errors"
@@ -14,6 +15,11 @@ import (
// DBFinding is a row in the standalone findings table. It persists across task
// deletion unless the caller explicitly requests related finding cleanup.
type DBFinding struct {
TrafficCount int
EvidenceVersion int64
ReportEvidenceVersion int64
TrafficBindings []FindingTrafficBinding // populated only for export
ID int64
TaskID *int64
NodeID *int64
@@ -98,7 +104,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 {
@@ -111,7 +118,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)
@@ -504,7 +511,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)
@@ -632,7 +639,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
}
@@ -648,20 +655,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
@@ -688,11 +702,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
@@ -738,6 +748,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
@@ -758,7 +770,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
@@ -768,7 +780,7 @@ func (d *DB) FindingMetaByNodeID(taskID int64) (map[int64]FindingMeta, error) {
var nid int64
var m FindingMeta
var aidsJSON string
if err := rows.Scan(&nid, &m.ID, &m.Status, &aidsJSON); err != nil {
if err := rows.Scan(&nid, &m.ID, &m.Status, &aidsJSON, &m.TrafficCount); err != nil {
return out, err
}
_ = json.Unmarshal([]byte(aidsJSON), &m.AssetIDs)
+33
View File
@@ -1095,6 +1095,39 @@ DROP TRIGGER IF EXISTS trg_conversation_retest_delete ON conversations;
CREATE TRIGGER trg_conversation_retest_delete BEFORE DELETE ON conversations
FOR EACH ROW EXECUTE FUNCTION stop_deleted_conversation_retest();
ALTER TABLE findings ADD COLUMN IF NOT EXISTS evidence_version BIGINT NOT NULL DEFAULT 0;
ALTER TABLE findings ADD COLUMN IF NOT EXISTS report_evidence_version BIGINT NOT NULL DEFAULT 0;
CREATE TABLE IF NOT EXISTS traffic_evidence_snapshots (
id TEXT PRIMARY KEY,
source_traffic_id TEXT NOT NULL,
captured_at BIGINT NOT NULL,
url TEXT NOT NULL,
method TEXT NOT NULL,
status INTEGER NOT NULL,
content_type TEXT NOT NULL DEFAULT '',
req_head TEXT NOT NULL,
resp_head TEXT NOT NULL,
req_hash TEXT NOT NULL,
resp_hash TEXT NOT NULL,
req_len BIGINT NOT NULL,
resp_len BIGINT NOT NULL,
unreferenced_at TIMESTAMPTZ,
created_at TIMESTAMPTZ NOT NULL DEFAULT now()
);
CREATE TABLE IF NOT EXISTS finding_traffic_bindings (
id BIGSERIAL PRIMARY KEY,
finding_id BIGINT NOT NULL REFERENCES findings(id) ON DELETE CASCADE,
snapshot_id TEXT NOT NULL REFERENCES traffic_evidence_snapshots(id),
role TEXT NOT NULL DEFAULT 'supporting',
note TEXT NOT NULL DEFAULT '',
position INTEGER NOT NULL,
created_at TIMESTAMPTZ NOT NULL DEFAULT now(),
UNIQUE(finding_id, snapshot_id)
);
CREATE INDEX IF NOT EXISTS idx_finding_traffic_order ON finding_traffic_bindings(finding_id, position, id);
CREATE INDEX IF NOT EXISTS idx_finding_traffic_snapshot ON finding_traffic_bindings(snapshot_id);
-- =====================================================================
-- M. 后端日志持久化
-- =====================================================================
+3 -1
View File
@@ -16,7 +16,7 @@ import (
)
const (
TaskArchiveFormatVersion = 2
TaskArchiveFormatVersion = 3
TaskArchiveLegacyFormatVersion = 1
TaskArchiveLLMRecordsPath = "database/llm_records.ndjson"
)
@@ -587,6 +587,8 @@ func (d *DB) snapshotTaskArchive(taskID int64, llmRecords io.Writer) (*TaskArchi
{"task_llm_profiles", `SELECT * FROM task_llm_profiles WHERE task_id=$1 ORDER BY position`, []any{taskID}},
{"task_scope", `SELECT * FROM task_scope WHERE task_id=$1 ORDER BY id`, []any{taskID}},
{"findings", `SELECT * FROM findings WHERE task_id=$1 ORDER BY id`, []any{taskID}},
{"finding_traffic_bindings", `SELECT b.* FROM finding_traffic_bindings b JOIN findings f ON f.id=b.finding_id WHERE f.task_id=$1 ORDER BY b.finding_id,b.position,b.id`, []any{taskID}},
{"traffic_evidence_snapshots", `SELECT s.* FROM traffic_evidence_snapshots s WHERE EXISTS(SELECT 1 FROM finding_traffic_bindings b JOIN findings f ON f.id=b.finding_id WHERE b.snapshot_id=s.id AND f.task_id=$1) ORDER BY s.id`, []any{taskID}},
{"llm_records", `SELECT * FROM llm_records WHERE COALESCE(task_id,'')=$1 ORDER BY id`, []any{strconv.FormatInt(taskID, 10)}},
{"llm_usage", `SELECT * FROM llm_usage WHERE COALESCE(task_id,'')=$1 OR exploration_id=$2 ORDER BY id`, []any{strconv.FormatInt(taskID, 10), expID}},
{"skill_usage", `SELECT * FROM skill_usage WHERE task_id=$1 OR exploration_id=$2 ORDER BY id`, []any{taskID, expID}},
+7
View File
@@ -233,6 +233,10 @@ WHERE archive.id=$1 FOR UPDATE OF archive,task`, archiveID).Scan(&taskID, &expID
if err != nil {
return nil, err
}
remappedTables["findings"], err = normalizeArchivedFindingVersions(remappedTables["findings"])
if err != nil {
return nil, err
}
// Insert graph rows in foreign-key order. The archived stub has no graph rows,
// so an ID conflict signals external corruption and must stop the restore.
for _, table := range []string{"exploration_nodes", "exploration_edges", "exploration_anchors", "task_constraints", "activity"} {
@@ -260,6 +264,9 @@ WHERE archive.id=$1 FOR UPDATE OF archive,task`, archiveID).Scan(&taskID, &expID
return nil, fmt.Errorf("restore %s: %w", table, err)
}
}
if err := restoreFindingTrafficTx(tx, snapshot); err != nil {
return nil, fmt.Errorf("restore finding traffic: %w", err)
}
if streamedLLMRecords != "" {
count, err := insertArchiveJSONSequenceRows(tx, "llm_records", llmRecords)
if err != nil {
+99
View File
@@ -0,0 +1,99 @@
# 漏洞多流量证据
漏洞详情的「关联流量」支持跨页多选、用途和说明、排序及解除绑定。流量页也可多选记录,一次关联到一个已有漏洞。任务继承的漏洞证据只读,修改需进入来源任务。
系统设置中的「Agent 自动绑定流量」默认关闭,读取和修改接口为 `/api/settings` 的 `agent_traffic_binding`。开启会增加查阅请求/响应、工具调用及提示词带来的 Token 消耗;下一轮 Agent 使用新设置。关闭时隐藏自动绑定参数和补绑工具,不注入自动绑定指引,并拒绝已运行会话新提交的自动绑定。人工绑定、流量捕获、已保存证据读取和导出不受影响。
开启后,默认流程是「发现入库 → 自动触发报告 Agent → 核对并绑定流量 → 按最新证据版本写报告」。上报者在 `evidence` 保留验证命令、关键输出、已有真实流量 ID 及用途;报告 Agent 结合漏洞详情和执行记录,用 `traffic_search` / `traffic_get` 核实后调用 `bind_finding_traffic`,再读取最新 `version` 保存报告。关闭开关时,报告 Agent 不额外获得原始流量检索/读取工具,也不自动绑定;仍能读取人工绑定的快照并生成报告。
兼容已有调用者:`report_finding` 的 `traffic_refs` / `evidence_hint_id` 仍可显式即时绑定。绑定可选:TCP 等非 HTTP 漏洞、未采集或找不到确切记录时,省略引用仍可正常上报和撰写报告;保留命令输出、日志等其他可验证证据,建议说明未绑定原因,不新增必填字段。提交的 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` 填实际读取的版本;生成期间证据变化会拒绝旧版本写入,必须重新读取并生成。
- 旧报告调用未传版本时,不宣称覆盖已有流量证据。绑定、说明、用途或排序变化后,已有报告提示待更新。
开启自动绑定后,`add_hint` / `add_task_hint` 支持在单条提示或批量 `hints` 的每个元素中保存 `traffic_refs`。Planner 代为上报可传 `evidence_hint_id`,明确选取本任务对应提示中的引用,不能引用继承提示。系统不会按域名、时间或浏览记录猜测绑定。提交失败不生成部分漏洞或提前触发报告。
已有漏洞漏绑时可用 `bind_finding_traffic(finding_id, traffic_refs)` 补绑,无需重复登记。`list_findings` / `list_task_findings` / `node_detail` / `get_task_node_detail` 返回明确的 `finding_id` 和 `finding_node_id`;旧 `id` 保持探索节点语义。
启动时仅为旧工具 schema 增加可选属性,原始默认流量工具绑定扩展至报告 Agent;补绑工具默认交给报告 Agent。自定义绑定列表、提示词、描述及启用状态保留。指引在最终工具装配后统一加入:上报角色负责交接已有证据,报告 Agent 负责核对、绑定和写报告。平台对话缺少任务上下文时应通过结构化提示交接给任务 Agent,不直接上报。判定任务完成前应先交接已有证据;无包不强制等待。失败的 `report_finding` 不触发报告 Agent。
## API
基础路径:`/api/exploration/findings/{finding_id}/traffic`,使用独立漏洞 ID。沿用认证;`context_task` 校验任务可见性及继承只读。
| 方法 / 相对路径 | 请求 / 返回 |
| --- | --- |
| `GET` | 有序摘要、证据版本、报告采用版本 |
| `POST` | `{"traffic_refs":[...]}` 整批追加 |
| `PATCH /{binding_id}` | `{"version":1,"role":"proof","note":"说明"}` |
| `DELETE /{binding_id}` | `{"version":1}` |
| `PUT /order` | `{"version":1,"binding_ids":["2","1"]}`,必须是完整列表 |
| `GET /{binding_id}` | 快照元数据和有界正文预览 |
| `GET /{binding_id}/body` | `side`、`offset`、`length`;`download=1` 下载完整原字节 |
版本/排序集合冲突、归档中写入返回 `409`;继承写入 `403`;不存在或不属于漏洞的绑定 `404`;流量/附件读取及校验失败明确返回错误。
## 存储与迁移
启动幂等迁移 PostgreSQL:新增 `traffic_evidence_snapshots`、`finding_traffic_bindings`,以及 `findings.evidence_version` / `report_evidence_version`(默认 0)。不根据历史文字猜测补绑。
快照保存原流量 ID、采集时间、URL、方法、状态、请求/响应头、正文长度和 SHA-256。正文按哈希存放在 `<data>/evidence/blobs/<前两位>/<hash>.bin`,与可清理的 `data/traffic` 独立,多个漏洞可共享快照/正文。快照不提供内容更新接口,校验不一致时读取和导出失败。
在原流量写锁下读取完整正文,包括大正文 blob 和旧目录记录。先持久化并校验文件,再用单个 PostgreSQL 事务写入探索节点、意图关系、漏洞、快照和绑定;提交后才通知规划者。失败可能遗留无引用文件,但不产生部分业务记录。
PostgreSQL advisory lock `7337741004` 协调证据文件和 SQL 引用;任务行锁禁止归档排队后的证据修改。恢复从正文安装到元数据提交全程保持证据锁。删除漏洞级联移除绑定。
清理器每小时运行,只回收无引用且非进行中操作的内容,至少延迟 24 小时。普通流量清理不触及证据目录。备份热数据时同时备份 PostgreSQL 和 `data/evidence`。
## 导出与归档
Markdown 包含有序证据清单及版本;JSON 包含元数据;CSV 增加数量和绑定 ID。`md-zip` 保留漏洞 Markdown,并提供:
```text
evidence/<finding_id>/<binding_id>/
manifest.json
request.http
response.http
request.bin
response.bin
```
Markdown 通过相对链接引用报文。发送下载前完成附件复制、哈希校验、压缩、磁盘同步及全部 ZIP 条目的 CRC 读取校验;缺失/损坏使整个下载失败。完整附件保留二进制原字节。
归档 v3 按漏洞绑定关系收集快照和正文,不依赖原流量或域名。包校验完成后才清理热数据;共享证据继续保留。恢复先校验安装正文,再事务恢复元数据和绑定,支持失败重试。v1/v2 继续可恢复,缺失的新字段显式补为 0。
## 验证与边界
为每个测试包设置独立、新建的 PostgreSQL 测试库,通过 `ARTEX_PG_DSN` 指定,避免残留任务/模型夹具触发后台运行。运行相关包完整测试,并确认没有配置缺失导致的跳过:
```sh
# 每个包运行前将 ARTEX_PG_DSN 设为对应的独立测试库;显式配置失败必须报错。
go test ./<package> -count=1
go test -race -p 1 ./evidence ./db ./agent ./server -run 'TestEvidence|TestFindingTraffic|TestFindingEvidence|TestReportFindingAtomicContract|TestTaskArchive'
```
前端验证包含 `npx tsc --noEmit`、受影响文件的 Biome 检查、Webpack 构建和 `NEXT_EXPORT=1` 静态导出。使用独立缓存目录,避免覆盖运行中的开发服务。
本机端到端验收使用独立端口、受控 HTTP / 域名 HTTPS 目标和临时数据目录,覆盖两个绑定入口、跨页选择、排序/说明、错误提示、继承只读、下载,以及删除原流量后导出、归档、回收热正文、恢复并校验哈希。
首版采用全局证据协调锁;大批量绑定/导出或慢速附件下载期间,其他证据操作可能等待。缺少录制记录或完整正文时不能补造证据。本功能不改变捕获开关,也不处理 HTTPS 直接访问 IP 的证书问题。
+401
View File
@@ -0,0 +1,401 @@
// Package evidence preserves finding evidence independently of disposable traffic.
package evidence
import (
"context"
"crypto/sha256"
"database/sql"
"encoding/hex"
"errors"
"fmt"
"io"
"io/fs"
"log"
"os"
"path/filepath"
"strings"
"time"
"github.com/Autumn-27/artex/db"
"github.com/Autumn-27/artex/traffic"
)
type Store struct {
DB *db.DB
Traffic *traffic.Traffic
Dir string
}
func New(pg *db.DB, tr *traffic.Traffic, dir string) *Store {
return &Store{DB: pg, Traffic: tr, Dir: dir}
}
func hashPath(dir, hash string) (string, error) {
if len(hash) != 64 {
return "", errors.New("invalid evidence hash")
}
if _, err := hex.DecodeString(hash); err != nil || strings.ToLower(hash) != hash {
return "", errors.New("invalid evidence hash")
}
return filepath.Join(dir, "blobs", hash[:2], hash+".bin"), nil
}
func verifyFile(path, hash string, length int64) error {
f, err := os.Open(path)
if err != nil {
return err
}
defer f.Close()
h := sha256.New()
n, err := io.Copy(h, f)
if err != nil {
return err
}
if n != length || hex.EncodeToString(h.Sum(nil)) != hash {
return fmt.Errorf("证据正文校验失败: %s", hash)
}
return nil
}
// A new body becomes visible only after a durable write. Failed SQL commits may
// leave unreferenced files; GC reaps those after a full day's grace period.
func (s *Store) writeBody(r io.Reader, expectedLength int64, expectedHash string) (hash string, err error) {
stage := filepath.Join(s.Dir, ".staging")
if err = os.MkdirAll(stage, 0o700); err != nil {
return
}
f, err := os.CreateTemp(stage, "body-")
if err != nil {
return "", err
}
defer func() { f.Close(); os.Remove(f.Name()) }()
h := sha256.New()
n, err := io.Copy(io.MultiWriter(f, h), r)
if err != nil {
return "", err
}
hash = hex.EncodeToString(h.Sum(nil))
if n != expectedLength {
return "", fmt.Errorf("正文不完整: 预期 %d 字节,读取 %d 字节", expectedLength, n)
}
if expectedHash != "" && expectedHash != hash {
return "", errors.New("原始流量正文哈希不匹配")
}
if err = f.Sync(); err != nil {
return "", err
}
if err = f.Close(); err != nil {
return "", err
}
path, err := hashPath(s.Dir, hash)
if err != nil {
return "", err
}
if _, err = os.Stat(path); err == nil {
if err = verifyFile(path, hash, n); err != nil {
return "", err
}
// Refresh the grace period for a restored but not-yet-committed body.
now := time.Now()
return hash, os.Chtimes(path, now, now)
} else if !os.IsNotExist(err) {
return "", err
}
if err = os.MkdirAll(filepath.Dir(path), 0o700); err != nil {
return "", err
}
if err = os.Rename(f.Name(), path); err != nil {
return "", err
}
d, err := os.Open(filepath.Dir(path))
if err != nil {
return "", err
}
defer d.Close()
return hash, d.Sync()
}
// StageFindingsExport freezes bindings/report versions and makes private body
// copies before the HTTP response is started. The caller owns and removes dest.
func (s *Store) StageFindingsExport(ctx context.Context, findings []*db.DBFinding, dest string, copyBodies bool) error {
return s.DB.WithEvidenceTx(ctx, func(tx *sql.Tx) error {
var snapshots []db.TrafficEvidenceSnapshot
for _, f := range findings {
list, err := db.FindingTrafficTx(tx, f.ID)
if err != nil {
return err
}
f.TrafficBindings = list.Bindings
f.TrafficCount = len(list.Bindings)
f.EvidenceVersion = list.Version
f.ReportEvidenceVersion = list.ReportVersion
if err = tx.QueryRow(`SELECT report FROM findings WHERE id=$1`, f.ID).Scan(&f.Report); err != nil {
return err
}
for _, b := range list.Bindings {
snapshots = append(snapshots, b.Snapshot)
}
}
if copyBodies {
return s.copySnapshots(snapshots, dest)
}
return nil
})
}
func (s *Store) prepare(ctx context.Context, refs []db.TrafficRef) ([]db.PreparedTrafficEvidence, error) {
refs, err := db.NormalizeTrafficRefs(refs)
if err != nil {
return nil, err
}
out := make([]db.PreparedTrafficEvidence, 0, len(refs))
if len(refs) == 0 {
return out, nil
}
ids := make([]string, len(refs))
byID := map[string]db.TrafficRef{}
for i, ref := range refs {
ids[i] = ref.TrafficID
byID[ref.TrafficID] = ref
}
err = s.Traffic.ReadEvidence(ctx, ids, func(e traffic.EvidenceExchange) error {
rh, err := s.writeBody(e.Request, e.ReqLen, e.ReqHash)
if err != nil {
return err
}
ph, err := s.writeBody(e.Response, e.RespLen, e.RespHash)
if err != nil {
return err
}
v := db.TrafficEvidenceSnapshot{SourceTrafficID: e.ID, CapturedAt: e.TS, URL: e.URL, Method: e.Method, Status: e.Status, ContentType: e.ContentType,
ReqHead: e.ReqHead, RespHead: e.RespHead, ReqHash: rh, RespHash: ph, ReqLen: e.ReqLen, RespLen: e.RespLen}
v.ID = db.TrafficSnapshotID(v)
out = append(out, db.PreparedTrafficEvidence{Ref: byID[e.ID], Snapshot: v})
return nil
})
return out, err
}
func (s *Store) Record(ctx context.Context, in db.RecordFindingInput, refs []db.TrafficRef) (out *db.RecordedFinding, err error) {
err = s.DB.WithEvidenceTx(ctx, func(tx *sql.Tx) error {
if err := db.LockTaskEvidenceTx(tx, in.TaskID); err != nil {
return err
}
prepared, err := s.prepare(ctx, refs)
if err != nil {
return err
}
out, err = db.RecordFindingTx(tx, in, prepared)
return err
})
return
}
func (s *Store) Bind(ctx context.Context, findingID int64, refs []db.TrafficRef) (out *db.FindingTraffic, err error) {
err = s.DB.WithEvidenceTx(ctx, func(tx *sql.Tx) error {
if err := db.LockFindingEvidenceTx(tx, findingID, nil); err != nil {
return err
}
prepared, err := s.prepare(ctx, refs)
if err != nil {
return err
}
if err = db.AddFindingTrafficTx(tx, findingID, prepared); err != nil {
return err
}
out, err = db.FindingTrafficTx(tx, findingID)
return err
})
return
}
func (s *Store) WithBinding(ctx context.Context, findingID, bindingID int64, fn func(db.FindingTrafficBinding) error) error {
return s.DB.WithEvidenceTx(ctx, func(tx *sql.Tx) error {
list, err := db.FindingTrafficTx(tx, findingID)
if err != nil {
return err
}
for _, b := range list.Bindings {
if b.ID == bindingID {
return fn(b)
}
}
return db.ErrEvidenceNotFound
})
}
// OpenBody must be called while holding WithEvidenceTx/WithBinding, and the
// returned file closed before that callback returns.
func (s *Store) OpenBody(snapshot db.TrafficEvidenceSnapshot, side string) (*os.File, int64, error) {
hash, length := snapshot.ReqHash, snapshot.ReqLen
if side == "response" {
hash, length = snapshot.RespHash, snapshot.RespLen
} else if side != "request" {
return nil, 0, errors.New("side 必须为 request 或 response")
}
path, err := hashPath(s.Dir, hash)
if err != nil {
return nil, 0, err
}
if err = verifyFile(path, hash, length); err != nil {
return nil, 0, err
}
f, err := os.Open(path)
return f, length, err
}
// CopySnapshots is used by both report downloads and portable task archives.
// The destination owns real copies, never links into either disposable store.
func (s *Store) CopySnapshots(ctx context.Context, snapshots []db.TrafficEvidenceSnapshot, dest string) error {
return s.DB.WithEvidenceTx(ctx, func(*sql.Tx) error { return s.copySnapshots(snapshots, dest) })
}
func (s *Store) copySnapshots(snapshots []db.TrafficEvidenceSnapshot, dest string) error {
for _, v := range snapshots {
for _, side := range []string{"request", "response"} {
if err := func() error {
f, length, err := s.OpenBody(v, side)
if err != nil {
return err
}
defer f.Close()
hash := v.ReqHash
if side == "response" {
hash = v.RespHash
}
target, err := hashPath(dest, hash)
if err != nil {
return err
}
if err = os.MkdirAll(filepath.Dir(target), 0o700); err != nil {
return err
}
out, err := os.OpenFile(target, os.O_CREATE|os.O_WRONLY|os.O_EXCL, 0o600)
if os.IsExist(err) {
return verifyFile(target, hash, length)
}
if err != nil {
return err
}
_, copyErr := io.Copy(out, f)
syncErr := out.Sync()
closeErr := out.Close()
return errors.Join(copyErr, syncErr, closeErr)
}(); err != nil {
return err
}
}
}
return nil
}
func (s *Store) InstallSnapshots(ctx context.Context, snapshots []db.TrafficEvidenceSnapshot, source string) error {
return s.WithInstalledSnapshots(ctx, snapshots, source, func() error { return nil })
}
// WithInstalledSnapshots pins installed bodies until the metadata restore finishes.
// The callback must not acquire another evidence advisory lock.
func (s *Store) WithInstalledSnapshots(ctx context.Context, snapshots []db.TrafficEvidenceSnapshot, source string, restore func() error) error {
return s.DB.WithEvidenceTx(ctx, func(*sql.Tx) error {
for _, v := range snapshots {
if v.ID != db.TrafficSnapshotID(v) {
return errors.New("归档证据快照元数据哈希不匹配")
}
for _, body := range []struct {
hash string
length int64
}{{v.ReqHash, v.ReqLen}, {v.RespHash, v.RespLen}} {
if err := func() error {
path, err := hashPath(source, body.hash)
if err != nil {
return err
}
f, err := os.Open(path)
if err != nil {
return err
}
defer f.Close()
_, err = s.writeBody(f, body.length, body.hash)
return err
}(); err != nil {
return err
}
}
}
return restore()
})
}
func (s *Store) Collect(ctx context.Context, now time.Time) error {
return s.DB.WithEvidenceTx(ctx, func(tx *sql.Tx) error {
if _, err := tx.Exec(`UPDATE traffic_evidence_snapshots s SET unreferenced_at=$1
WHERE unreferenced_at IS NULL AND NOT EXISTS(SELECT 1 FROM finding_traffic_bindings b WHERE b.snapshot_id=s.id)`, now); err != nil {
return err
}
if _, err := tx.Exec(`DELETE FROM traffic_evidence_snapshots s WHERE unreferenced_at<$1
AND NOT EXISTS(SELECT 1 FROM finding_traffic_bindings b WHERE b.snapshot_id=s.id)`, now.Add(-24*time.Hour)); err != nil {
return err
}
rows, err := tx.Query(`SELECT req_hash FROM traffic_evidence_snapshots UNION SELECT resp_hash FROM traffic_evidence_snapshots`)
if err != nil {
return err
}
refs := map[string]bool{}
for rows.Next() {
var hash string
if err := rows.Scan(&hash); err != nil {
rows.Close()
return err
}
refs[hash] = true
}
err = rows.Err()
rows.Close()
if err != nil {
return err
}
for _, root := range []string{filepath.Join(s.Dir, "blobs"), filepath.Join(s.Dir, ".staging")} {
err := filepath.WalkDir(root, func(path string, entry fs.DirEntry, err error) error {
if os.IsNotExist(err) {
return nil
}
if err != nil {
return err
}
if entry.IsDir() {
return nil
}
if refs[strings.TrimSuffix(entry.Name(), ".bin")] {
return nil
}
info, err := entry.Info()
if err != nil {
return err
}
if now.Sub(info.ModTime()) < 24*time.Hour {
return nil
}
return os.Remove(path)
})
if err != nil {
return err
}
}
return nil
})
}
func (s *Store) RunGC(ctx context.Context) {
timer := time.NewTicker(time.Hour)
defer timer.Stop()
for {
select {
case <-ctx.Done():
return
case now := <-timer.C:
if err := s.Collect(ctx, now); err != nil && !errors.Is(err, context.Canceled) {
log.Printf("[evidence] cleanup failed: %v", err)
}
}
}
}
+333
View File
@@ -0,0 +1,333 @@
package evidence
import (
"bytes"
"context"
"crypto/sha256"
"encoding/hex"
"errors"
"fmt"
"io"
"os"
"path/filepath"
"sync"
"testing"
"time"
"github.com/Autumn-27/artex/db"
"github.com/Autumn-27/artex/traffic"
)
func evidenceFixture(t *testing.T) (*Store, db.RecordFindingInput, string) {
t.Helper()
dsn := os.Getenv("ARTEX_PG_DSN")
if dsn == "" {
t.Skip("ARTEX_PG_DSN is required for evidence integration tests")
}
pg, err := db.Open(dsn)
if err != nil {
t.Fatal(err)
}
t.Cleanup(func() { pg.Close() })
task, err := pg.CreateTask("evidence integration "+t.Name(), "local fixtures", nil, 0, 0)
if err != nil {
t.Fatal(err)
}
t.Cleanup(func() { pg.DeleteFindingsByTask(task.ID); pg.DeleteTask(task.ID) })
dir := t.TempDir()
tr, err := traffic.Open(filepath.Join(dir, "traffic"), ":0")
if err != nil {
t.Fatal(err)
}
t.Cleanup(func() { tr.Close() })
store := New(pg, tr, filepath.Join(dir, "evidence"))
in := db.RecordFindingInput{TaskID: task.ID, ExplorationID: task.ExplorationID, Worker: "test", VulnClass: "TEST", Name: "Evidence fixture", Severity: "low", Summary: "local test"}
return store, in, dir
}
func seedExchange(t *testing.T, s *Store, id string, body []byte, spill bool) {
t.Helper()
var blob any
inline := body
if spill {
sum := sha256.Sum256(body)
hash := hex.EncodeToString(sum[:])
blob = hash
path := filepath.Join(filepath.Dir(s.Dir), "traffic", "_blobs", "sha256", hash[:2], hash+".bin")
if err := os.MkdirAll(filepath.Dir(path), 0o700); err != nil {
t.Fatal(err)
}
if err := os.WriteFile(path, body, 0o600); err != nil {
t.Fatal(err)
}
inline = []byte("TRUNCATED PREVIEW")
}
_, err := s.Traffic.DB().Exec(`INSERT INTO exchanges(id,ts,host,method,url_template,url,status,content_type,req_len,resp_len,path) VALUES(?,?,'fixture.local','POST','/test',?,200,'application/octet-stream',2,?,'')`, id, time.Now().Unix(), "https://fixture.local/"+id, len(body))
if err != nil {
t.Fatal(err)
}
_, err = s.Traffic.DB().Exec(`INSERT INTO exchange_bodies(id,req_head,req_body,resp_head,resp_body,resp_blob) VALUES(?,?,?, ?,?,?)`, id, "POST /test HTTP/1.1\nHost: fixture.local\n", []byte("{}"), "HTTP 200\nContent-Type: application/octet-stream\n", inline, blob)
if err != nil {
t.Fatal(err)
}
}
func TestEvidenceBindingLifecycle(t *testing.T) {
s, in, _ := evidenceFixture(t)
ctx := context.Background()
large := bytes.Repeat([]byte{0, 1, 2, 255, 'a', 'b'}, 180000)
for i := 0; i < 3; i++ {
seedExchange(t, s, fmt.Sprint(i), large, i == 2)
}
refs := []db.TrafficRef{{TrafficID: "0", Role: "baseline", Note: "normal"}, {TrafficID: "1", Role: "proof", Note: "proof"}, {TrafficID: "2", Role: "verification", Note: "binary"}}
r, err := s.Record(ctx, in, refs)
if err != nil {
t.Fatal(err)
}
if len(r.Traffic.Bindings) != 3 || r.Traffic.Version != 1 {
t.Fatalf("record=%+v", r)
}
again, err := s.Bind(ctx, r.FindingID, []db.TrafficRef{{TrafficID: "1", Note: "must not overwrite"}})
if err != nil {
t.Fatal(err)
}
if len(again.Bindings) != 3 || again.Version != 1 || again.Bindings[1].Note != "proof" {
t.Fatalf("duplicate changed bindings: %+v", again)
}
other, err := s.Record(ctx, in, []db.TrafficRef{refs[2]})
if err != nil {
t.Fatal(err)
}
if other.Traffic.Bindings[0].SnapshotID != r.Traffic.Bindings[2].SnapshotID {
t.Fatal("snapshot was not shared")
}
if _, err = s.Traffic.DeleteHost("fixture.local"); err != nil {
t.Fatal(err)
}
for _, binding := range r.Traffic.Bindings {
if err = s.WithBinding(ctx, r.FindingID, binding.ID, func(b db.FindingTrafficBinding) error {
f, _, err := s.OpenBody(b.Snapshot, "response")
if err != nil {
return err
}
defer f.Close()
got, err := io.ReadAll(f)
if !bytes.Equal(got, large) {
t.Fatal("body was truncated/changed")
}
return err
}); err != nil {
t.Fatal(err)
}
}
v := again.Version
if _, err = s.DB.SetFindingReportVersionByNodeID(ctx, r.NodeID, "report", &v); err != nil {
t.Fatal(err)
}
note := "updated"
if err = s.DB.EditFindingTraffic(ctx, r.FindingID, again.Bindings[0].ID, v, nil, &note, false, nil); err != nil {
t.Fatal(err)
}
if _, err = s.DB.SetFindingReportVersionByNodeID(ctx, r.NodeID, "stale report", &v); !errors.Is(err, db.ErrEvidenceConflict) {
t.Fatalf("stale report: %v", err)
}
f, err := s.DB.GetFinding(r.FindingID)
if err != nil {
t.Fatal(err)
}
if f.Report != "report" || f.EvidenceVersion == f.ReportEvidenceVersion || f.TrafficCount != 3 {
t.Fatalf("finding versions: %+v", f)
}
if err = s.DB.EditFindingTraffic(ctx, r.FindingID, again.Bindings[0].ID, v, nil, nil, true, nil); !errors.Is(err, db.ErrEvidenceConflict) {
t.Fatalf("stale delete: %v", err)
}
if _, err = s.DB.SetFindingReportByNodeID(r.NodeID, "legacy report"); err != nil {
t.Fatal(err)
}
f, err = s.DB.GetFinding(r.FindingID)
if err != nil {
t.Fatal(err)
}
if f.ReportEvidenceVersion != -1 {
t.Fatal("legacy write claimed current evidence")
}
if _, err = s.DB.DeleteFinding(r.FindingID); err != nil {
t.Fatal(err)
}
if err = s.Collect(ctx, time.Now().Add(48*time.Hour)); err != nil {
t.Fatal(err)
}
if err = s.WithBinding(ctx, other.FindingID, other.Traffic.Bindings[0].ID, func(b db.FindingTrafficBinding) error {
f, _, err := s.OpenBody(b.Snapshot, "response")
if f != nil {
f.Close()
}
return err
}); err != nil {
t.Fatal("GC removed shared body:", err)
}
}
func TestEvidenceAtomicFailuresAndLegacy(t *testing.T) {
s, in, dir := evidenceFixture(t)
ctx := context.Background()
seedExchange(t, s, "valid", []byte("OK"), false)
counts := func() (findings, nodes int) {
s.DB.QueryRow(`SELECT count(*) FROM findings WHERE task_id=$1`, in.TaskID).Scan(&findings)
s.DB.QueryRow(`SELECT count(*) FROM exploration_nodes WHERE exploration_id=$1 AND kind='finding'`, in.ExplorationID).Scan(&nodes)
return
}
for _, kind := range []string{"missing_id", "missing_body", "disk_failure", "database_failure"} {
t.Run(kind, func(t *testing.T) {
input := in
refs := []db.TrafficRef{{TrafficID: "valid"}, {TrafficID: "absent"}}
local := *s
switch kind {
case "missing_body":
seedExchange(t, s, "broken", []byte("more than preview"), false)
s.Traffic.DB().Exec(`UPDATE exchange_bodies SET resp_body=? WHERE id='broken'`, []byte("x"))
refs = []db.TrafficRef{{TrafficID: "valid"}, {TrafficID: "broken"}}
case "disk_failure":
local.Dir = filepath.Join(dir, "file-not-dir")
if err := os.WriteFile(local.Dir, []byte("x"), 0o600); err != nil {
t.Fatal(err)
}
refs = refs[:1]
case "database_failure":
input.AssetIDs = []int64{9223372036854775807}
refs = refs[:1]
}
if _, err := local.Record(ctx, input, refs); err == nil {
t.Fatal("expected failure")
}
if f, n := counts(); f != 0 || n != 0 {
t.Fatalf("partial record: findings=%d nodes=%d", f, n)
}
})
}
legacy := filepath.Join(dir, "traffic", "legacy")
os.MkdirAll(legacy, 0o700)
os.WriteFile(filepath.Join(legacy, "request.http"), []byte("GET / HTTP/1.1\nHost: old.local\n"), 0o600)
os.WriteFile(filepath.Join(legacy, "response.http"), []byte("HTTP 200\n\nlegacy body"), 0o600)
_, err := s.Traffic.DB().Exec(`INSERT INTO exchanges(id,ts,host,method,url_template,url,status,content_type,req_len,resp_len,path) VALUES('old',1,'old.local','GET','/','http://old.local/',200,'text/plain',0,11,'legacy')`)
if err != nil {
t.Fatal(err)
}
if _, err = s.Record(ctx, in, []db.TrafficRef{{TrafficID: "old"}}); err != nil {
t.Fatal(err)
}
if _, err = s.Record(ctx, in, nil); err != nil {
t.Fatal("legacy no-traffic report failed:", err)
}
}
func TestEvidenceConcurrentBindAndDelete(t *testing.T) {
s, in, _ := evidenceFixture(t)
ctx := context.Background()
seedExchange(t, s, "race", bytes.Repeat([]byte("large"), 90000), true)
r, err := s.Record(ctx, in, nil)
if err != nil {
t.Fatal(err)
}
var wg sync.WaitGroup
errs := make(chan error, 10)
for i := 0; i < 8; i++ {
wg.Add(1)
go func() {
defer wg.Done()
_, err := s.Bind(ctx, r.FindingID, []db.TrafficRef{{TrafficID: "race"}})
errs <- err
}()
}
wg.Wait()
close(errs)
for err := range errs {
if err != nil {
t.Fatal(err)
}
}
list, err := s.DB.GetFindingTraffic(ctx, r.FindingID)
if err != nil {
t.Fatal(err)
}
if len(list.Bindings) != 1 || list.Version != 1 {
t.Fatalf("concurrent duplicates: %+v", list)
}
// Race a new capture copy with source deletion: either the full copy commits
// or nothing does. An already-bound snapshot remains readable in both cases.
wg.Add(2)
go func() { defer wg.Done(); s.Bind(ctx, r.FindingID, []db.TrafficRef{{TrafficID: "race"}}) }()
go func() { defer wg.Done(); s.Traffic.DeleteHost("fixture.local") }()
wg.Wait()
if err = s.WithBinding(ctx, r.FindingID, list.Bindings[0].ID, func(b db.FindingTrafficBinding) error {
f, _, err := s.OpenBody(b.Snapshot, "response")
if f != nil {
f.Close()
}
return err
}); err != nil {
t.Fatal(err)
}
}
func TestEvidenceGCGraceAndActiveRestore(t *testing.T) {
s, in, _ := evidenceFixture(t)
ctx := context.Background()
seedExchange(t, s, "gc", []byte("unique unreferenced gc body"), false)
f, err := s.Record(ctx, in, []db.TrafficRef{{TrafficID: "gc"}})
if err != nil {
t.Fatal(err)
}
snap := f.Traffic.Bindings[0].Snapshot
archive := t.TempDir()
if err = s.CopySnapshots(ctx, []db.TrafficEvidenceSnapshot{snap}, archive); err != nil {
t.Fatal(err)
}
if _, err = s.DB.DeleteFinding(f.FindingID); err != nil {
t.Fatal(err)
}
now := time.Now()
if err = s.Collect(ctx, now); err != nil {
t.Fatal(err)
}
path, err := hashPath(s.Dir, snap.RespHash)
if err != nil {
t.Fatal(err)
}
if err = s.Collect(ctx, now.Add(23*time.Hour)); err != nil {
t.Fatal(err)
}
if _, err = os.Stat(path); err != nil {
t.Fatal("GC grace ignored", err)
}
if err = s.Collect(ctx, now.Add(25*time.Hour)); err != nil {
t.Fatal(err)
}
if _, err = os.Stat(path); !os.IsNotExist(err) {
t.Fatal("orphan was not collected", err)
}
entered, release, finished := make(chan struct{}), make(chan struct{}), make(chan error, 1)
go func() {
finished <- s.WithInstalledSnapshots(ctx, []db.TrafficEvidenceSnapshot{snap}, archive, func() error { close(entered); <-release; return nil })
}()
<-entered
short, cancel := context.WithTimeout(ctx, 100*time.Millisecond)
err = s.Collect(short, now.Add(72*time.Hour))
cancel()
close(release)
if err == nil {
t.Fatal("GC entered during active restore")
}
if err = <-finished; err != nil {
t.Fatal(err)
}
if _, err = os.Stat(path); err != nil {
t.Fatal("active restore lost body", err)
}
// A portable package cannot substitute metadata while retaining its old ID.
bad := snap
bad.URL = "http://tampered.local/"
if err = s.InstallSnapshots(ctx, []db.TrafficEvidenceSnapshot{bad}, archive); err == nil {
t.Fatal("accepted tampered snapshot")
}
}
+35
View File
@@ -0,0 +1,35 @@
package evidence
import (
"context"
"os"
"testing"
"github.com/Autumn-27/artex/db"
)
// Initialize an explicitly configured fresh database before taking the same
// suite lock as db, agent and server. Hold it on one pinned connection.
func TestMain(m *testing.M) {
if os.Getenv("ARTEX_PG_DSN") == "" {
os.Exit(m.Run())
}
os.Exit(runEvidenceSuite(m))
}
func runEvidenceSuite(m *testing.M) int {
pg, err := db.Open(os.Getenv("ARTEX_PG_DSN"))
if err != nil {
panic(err)
}
defer pg.Close()
conn, err := pg.Conn(context.Background())
if err != nil {
panic(err)
}
defer conn.Close()
if _, err = conn.ExecContext(context.Background(), `SELECT pg_advisory_lock(7337741002)`); err != nil {
panic(err)
}
defer conn.ExecContext(context.Background(), `SELECT pg_advisory_unlock(7337741002)`)
return m.Run()
}
+36 -1
View File
@@ -84,6 +84,7 @@ func FindingsMarkdown(fs []*db.DBFinding, generatedAt time.Time) string {
b.WriteString(rep)
b.WriteString("\n\n")
}
b.WriteString(findingTrafficMarkdown(f, false))
b.WriteString("---\n\n")
}
return b.String()
@@ -114,6 +115,7 @@ func SingleFindingMarkdown(f *db.DBFinding, generatedAt time.Time) string {
b.WriteString(rep)
b.WriteString("\n")
}
b.WriteString(findingTrafficMarkdown(f, true))
return b.String()
}
@@ -147,7 +149,7 @@ func FindingsCSV(fs []*db.DBFinding) []byte {
var buf bytes.Buffer
buf.WriteString("\xEF\xBB\xBF") // UTF-8 BOM
w := csv.NewWriter(&buf)
_ = w.Write([]string{"ID", "名称", "类别", "严重等级", "状态", "所属任务", "发现时间", "概述"})
_ = w.Write([]string{"ID", "名称", "类别", "严重等级", "状态", "所属任务", "发现时间", "概述", "流量证据数量", "流量证据ID"})
for _, f := range items {
_ = w.Write([]string{
fmt.Sprintf("%d", f.ID),
@@ -158,8 +160,41 @@ func FindingsCSV(fs []*db.DBFinding) []byte {
f.TaskDescription,
f.CreatedAt.Format("2006-01-02 15:04:05"),
strings.TrimSpace(f.Summary),
fmt.Sprint(len(f.TrafficBindings)), findingTrafficIDs(f),
})
}
w.Flush()
return buf.Bytes()
}
func findingTrafficIDs(f *db.DBFinding) string {
ids := make([]string, 0, len(f.TrafficBindings))
for _, b := range f.TrafficBindings {
ids = append(ids, fmt.Sprint(b.ID))
}
return strings.Join(ids, ",")
}
func findingTrafficMarkdown(f *db.DBFinding, attachments bool) string {
stale := f.Report != "" && f.EvidenceVersion != f.ReportEvidenceVersion
if len(f.TrafficBindings) == 0 && !stale {
return ""
}
var out strings.Builder
out.WriteString("\n## 关联流量证据\n\n")
fmt.Fprintf(&out, "证据版本:%d;绑定数量:%d。\n\n", f.EvidenceVersion, len(f.TrafficBindings))
if stale {
out.WriteString("证据已变更,详细报告待更新。\n\n")
}
for i, b := range f.TrafficBindings {
fmt.Fprintf(&out, "%d. **证据 #%d · %s** — `%s %s`,状态码 %d\n", i+1, b.ID, b.Role, b.Snapshot.Method, strings.ReplaceAll(b.Snapshot.URL, "`", "%60"), b.Snapshot.Status)
if b.Note != "" {
fmt.Fprintf(&out, " %s\n", strings.ReplaceAll(b.Note, "\n", "\n "))
}
if attachments {
fmt.Fprintf(&out, " [请求报文](evidence/%d/%d/request.http) · [响应报文](evidence/%d/%d/response.http)\n", f.ID, b.ID, f.ID, b.ID)
}
}
out.WriteString("\n")
return out.String()
}
+1
View File
@@ -210,6 +210,7 @@ func seedPrompts(pg *db.DB) {
// flag and, if kept, wrapped so the model sees the DB-overridden description/schema
// and缺省入参 get injected. MCP/skill/host tools have no row and pass through.
func wireTools(pg *db.DB, domainReg map[string]actool.CoreTool) {
agent.FindingTrafficBindingEnabled = func() bool { return pg.GetBool(settingAgentTrafficBinding, false) }
// Seed the built-in domain tools (first-insert only; DO NOTHING preserves edits).
// No startup prune: rows we didn't seed are left alone so future user-defined
// custom tools (system=false, added via the UI) survive restarts.
+11 -2
View File
@@ -292,6 +292,12 @@ func coverageAssetRefDTO(ref db.AssetRef) CoverageAssetRefDTO {
// ---- Finding (frontend "Finding") ----
type FindingDTO struct {
TrafficCount int `json:"traffic_count"`
EvidenceVersion int64 `json:"evidence_version"`
ReportEvidenceVersion int64 `json:"report_evidence_version"`
ReportStale bool `json:"report_stale"`
TrafficBindings []db.FindingTrafficBinding `json:"traffic_bindings,omitempty"`
ID string `json:"id"`
FindingID string `json:"finding_id,omitempty"` // standalone findings-table id — the handle for status updates
VulnClass string `json:"vulnclass"`
@@ -399,6 +405,7 @@ func findingDTOsForOwner(taskID, description string, in []*db.Node, meta map[int
if m, ok := meta[n.ID]; ok {
d.FindingID = i64s(m.ID)
d.Status = m.Status
d.TrafficCount = m.TrafficCount
d.Assets = findingAssetDTOs(m.AssetIDs, assets)
}
out = append(out, d)
@@ -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,
+412
View File
@@ -0,0 +1,412 @@
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 err := s.agentFindingTrafficAccess(ctx, id, false); err != nil {
return actool.Errorf(err.Error()), nil
}
if len(a.BindingID) == 0 {
list, err := s.m.pg.GetFindingTraffic(ctx, id)
if err != nil {
return actool.Errorf(err.Error()), nil
}
return jsonResult(trafficSummary(list))
}
if a.Side == "" {
a.Side = "response"
}
store := s.evidenceStore()
var result any
err := store.WithBinding(ctx, id, bid, func(b db.FindingTrafficBinding) error {
preview, err := readEvidencePreview(store, b.Snapshot, a.Side, a.Offset, a.Length)
if err != nil {
return err
}
head := b.Snapshot.ReqHead
if a.Side == "response" {
head = b.Snapshot.RespHead
}
result = map[string]any{"binding_id": strconv.FormatInt(bid, 10), "head": head, "body": preview}
return nil
})
if err != nil {
return actool.Errorf(err.Error()), nil
}
return jsonResult(result)
})
}
+121
View File
@@ -0,0 +1,121 @@
package server
import (
"archive/zip"
"encoding/json"
"errors"
"fmt"
"io"
"os"
"strings"
"time"
"github.com/Autumn-27/artex/db"
"github.com/Autumn-27/artex/evidence"
"github.com/Autumn-27/artex/report"
)
func writeFindingsEvidenceZip(out io.Writer, findings []*db.DBFinding, stage string, now time.Time) error {
zw := zip.NewWriter(out)
if err := writeFindingsZipEntries(zw, findings, stage, now); err != nil {
zw.Close()
return err
}
return zw.Close()
}
func writeFindingsZipEntries(zw *zip.Writer, findings []*db.DBFinding, stage string, now time.Time) error {
store := evidence.New(nil, nil, stage) // private export copy; not subject to GC
used := map[string]int{}
for _, f := range findings {
base := report.FindingFilename(f)
name := base
if n := used[base]; n > 0 {
name = fmt.Sprintf("%s-%d.md", strings.TrimSuffix(base, ".md"), n+1)
}
used[base]++
w, err := zw.Create(name)
if err != nil {
return err
}
if _, err = io.WriteString(w, report.SingleFindingMarkdown(f, now)); err != nil {
return err
}
for _, b := range f.TrafficBindings {
prefix := fmt.Sprintf("evidence/%d/%d/", f.ID, b.ID)
w, err := zw.Create(prefix + "manifest.json")
if err != nil {
return err
}
if err = json.NewEncoder(w).Encode(b); err != nil {
return err
}
for _, side := range []string{"request", "response"} {
if err := func() error {
body, _, err := store.OpenBody(b.Snapshot, side)
if err != nil {
return err
}
defer body.Close()
w, err := zw.Create(prefix + side + ".http")
if err != nil {
return err
}
head := b.Snapshot.ReqHead
if side == "response" {
head = b.Snapshot.RespHead
}
if _, err = io.WriteString(w, strings.TrimRight(head, "\r\n")+"\n\n"); err != nil {
return err
}
if _, err = io.Copy(w, body); err != nil {
return err
}
if _, err = body.Seek(0, io.SeekStart); err != nil {
return err
}
w, err = zw.Create(prefix + side + ".bin")
if err != nil {
return err
}
_, err = io.Copy(w, body)
return err
}(); err != nil {
return err
}
}
}
}
return nil
}
func buildFindingsEvidenceZip(path string, findings []*db.DBFinding, stage string, now time.Time) error {
f, err := os.OpenFile(path, os.O_CREATE|os.O_EXCL|os.O_WRONLY, 0o600)
if err != nil {
return err
}
writeErr := writeFindingsEvidenceZip(f, findings, stage, now)
syncErr := f.Sync()
closeErr := f.Close()
if err := errors.Join(writeErr, syncErr, closeErr); err != nil {
return err
}
// Re-read every compressed entry to validate the finished ZIP, including CRCs,
// before the caller starts an HTTP download response.
archive, err := zip.OpenReader(path)
if err != nil {
return err
}
defer archive.Close()
for _, entry := range archive.File {
r, err := entry.Open()
if err != nil {
return err
}
_, readErr := io.Copy(io.Discard, r)
if err := errors.Join(readErr, r.Close()); err != nil {
return err
}
}
return nil
}
+473
View File
@@ -0,0 +1,473 @@
package server
import (
"archive/zip"
"bytes"
"context"
"encoding/json"
"fmt"
"io"
"net/http/httptest"
"os"
"path/filepath"
"strconv"
"strings"
"testing"
"time"
"github.com/Autumn-27/artex/agent"
"github.com/Autumn-27/artex/db"
"github.com/Autumn-27/artex/traffic"
)
func trafficEvidenceServer(t *testing.T) (*Server, *db.RecordedFinding, func(string, string, string) *httptest.ResponseRecorder) {
t.Helper()
m, err := NewManager(t.TempDir(), "")
if err != nil {
t.Fatal(err)
}
t.Cleanup(func() { m.Close() })
m.traffic, err = traffic.Open(filepath.Join(m.dir, "traffic"), ":0")
if err != nil {
t.Fatal(err)
}
task, err := m.CreateTask("traffic evidence API", "local evidence", nil, 0, 0)
if err != nil {
t.Fatal(err)
}
tid, _ := strconv.ParseInt(task.ID, 10, 64)
t.Cleanup(func() { m.pg.Exec(`DELETE FROM task_archives WHERE task_id=$1`, tid); m.pg.DeleteTask(tid) })
ctx, cancel := context.WithCancel(context.Background())
cancel()
s := New(ctx, m, t.TempDir(), t.TempDir(), t.TempDir())
s.archiveWG.Wait()
token, err := signJWT(s.jwtKey)
if err != nil {
t.Fatal(err)
}
request := func(method, path, body string) *httptest.ResponseRecorder {
t.Helper()
r := httptest.NewRequest(method, path, strings.NewReader(body))
r.Header.Set("Authorization", "Bearer "+token)
r.Header.Set("Content-Type", "application/json")
w := httptest.NewRecorder()
s.Handler().ServeHTTP(w, r)
return w
}
input := db.RecordFindingInput{TaskID: tid, ExplorationID: task.ExpID, Worker: "test", VulnClass: "TEST", Summary: "bound evidence", Name: "Evidence API", Severity: "low"}
f, err := s.evidenceStore().Record(context.Background(), input, nil)
if err != nil {
t.Fatal(err)
}
return s, f, request
}
func seedServerEvidenceFlow(t *testing.T, s *Server, id string, body []byte) {
t.Helper()
_, err := s.m.traffic.DB().Exec(`INSERT INTO exchanges(id,ts,host,method,url_template,url,status,content_type,req_len,resp_len,path) VALUES(?,1,'evidence.local','POST','/','http://evidence.local/',200,'application/octet-stream',0,?,'')`, id, len(body))
if err != nil {
t.Fatal(err)
}
_, err = s.m.traffic.DB().Exec(`INSERT INTO exchange_bodies(id,req_head,resp_head,req_body,resp_body) VALUES(?,'POST / HTTP/1.1','HTTP/1.1 200 OK',?,?)`, id, []byte{}, body)
if err != nil {
t.Fatal(err)
}
}
func TestFindingTrafficToolUpgradePreservesCustomization(t *testing.T) {
s, _, _ := trafficEvidenceServer(t)
pg := s.m.pg
worker, err := pg.GetAgentByKey("worker")
if err != nil || worker == nil {
t.Fatal("missing worker", err)
}
oldPrompt, _ := pg.CurrentPrompt(worker.ID)
t.Cleanup(func() { pg.SavePrompt(worker.ID, oldPrompt, "restore test fixture", "test") })
if _, err := pg.SavePrompt(worker.ID, "USER CUSTOM PROMPT", "test", "test"); err != nil {
t.Fatal(err)
}
custom := json.RawMessage(`{"type":"object","properties":{"evidence":{"type":"string","description":"user evidence instructions"}},"required":["evidence"]}`)
for _, key := range []string{"report_finding", "update_finding_report", "get_finding_traffic"} {
old, err := pg.GetTool(key)
if err != nil || old == nil {
t.Fatal("missing tool", key, err)
}
t.Cleanup(func() {
bindings, _ := json.Marshal(old.Agents)
pg.UpdateTool(old.Key, old.Description, old.Schema, bindings, old.Enabled)
})
if err := pg.UpdateTool(key, "user description", custom, json.RawMessage(`["worker"]`), false); err != nil {
t.Fatal(err)
}
}
if err := pg.SetSetting("finding_traffic_tools_v1", "false"); err != nil {
t.Fatal(err)
}
s.seedFindingTrafficTools()
for key, property := range map[string]string{"report_finding": "traffic_refs", "update_finding_report": "evidence_version"} {
tool, err := pg.GetTool(key)
if err != nil || tool == nil {
t.Fatal(err)
}
var schema struct {
Properties map[string]json.RawMessage `json:"properties"`
Required []string `json:"required"`
}
if err := json.Unmarshal(tool.Schema, &schema); err != nil {
t.Fatal(err)
}
if tool.Enabled || tool.Description != "user description" || len(tool.Agents) != 1 || tool.Agents[0] != "worker" || len(schema.Required) != 1 || schema.Required[0] != "evidence" || len(schema.Properties[property]) == 0 || !strings.Contains(string(schema.Properties["evidence"]), "user evidence instructions") {
t.Fatalf("custom configuration overwritten: %+v", tool)
}
}
reader, _ := pg.GetTool("get_finding_traffic")
if reader.Enabled || !contains(reader.Agents, "reporter") {
t.Fatal("reader should be bound without enabling it")
}
if err := pg.RemoveAgentFromTool("reporter", "get_finding_traffic"); err != nil {
t.Fatal(err)
}
s.seedFindingTrafficTools()
reader, _ = pg.GetTool("get_finding_traffic")
if contains(reader.Agents, "reporter") {
t.Fatal("restart overwrote user's unbinding")
}
if got, _ := pg.CurrentPrompt(worker.ID); got != "USER CUSTOM PROMPT" {
t.Fatal("custom prompt replaced")
}
}
func TestFindingTrafficAPIAndExport(t *testing.T) {
s, f, req := trafficEvidenceServer(t)
body := bytes.Repeat([]byte{0, 255, 65, 66}, 5000)
for _, id := range []string{"a", "b", "c"} {
seedServerEvidenceFlow(t, s, id, body)
}
base := fmt.Sprintf("/api/exploration/findings/%d/traffic", f.FindingID)
w := req("POST", base, `{"traffic_refs":[{"traffic_id":"a","role":"baseline"},{"traffic_id":"b","role":"proof","note":"proof note"},{"traffic_id":"c"}]}`)
if w.Code != 200 {
t.Fatalf("bind: %d %s", w.Code, w.Body)
}
var list db.FindingTraffic
if err := json.Unmarshal(w.Body.Bytes(), &list); err != nil {
t.Fatal(err)
}
if len(list.Bindings) != 3 || list.Bindings[0].Snapshot.ReqHead != "" {
t.Fatal("not a bounded summary")
}
first, second, third := list.Bindings[0].ID, list.Bindings[1].ID, list.Bindings[2].ID
if w = req("PATCH", fmt.Sprintf("%s/%d", base, first), `{"version":0,"note":"stale"}`); w.Code != 409 {
t.Fatal(w.Code, w.Body)
}
if w = req("PUT", base+"/order", fmt.Sprintf(`{"version":1,"binding_ids":["%d","%d","%d"]}`, third, first, second)); w.Code != 200 {
t.Fatal(w.Code, w.Body)
}
if w = req("PUT", base+"/order", fmt.Sprintf(`{"version":2,"binding_ids":["%d"]}`, first)); w.Code != 409 {
t.Fatal(w.Code, w.Body)
}
if w = req("GET", fmt.Sprintf("%s/%d", base, first), ""); w.Code != 200 || !strings.Contains(w.Body.String(), `"binary":true`) {
t.Fatal(w.Code, w.Body)
}
if w = req("GET", fmt.Sprintf("%s/%d/body?side=response&offset=20001", base, first), ""); w.Code != 422 {
t.Fatal(w.Code, w.Body)
}
if _, err := s.m.traffic.DeleteHost("evidence.local"); err != nil {
t.Fatal(err)
}
if w = req("GET", fmt.Sprintf("%s/%d/body?side=response&download=1", base, first), ""); w.Code != 200 || !bytes.Equal(w.Body.Bytes(), body) {
t.Fatal("download after original removal failed")
}
// IDs belong to the requested finding; cross-finding detail cannot be read.
finding, err := s.m.pg.GetFinding(f.FindingID)
if err != nil {
t.Fatal(err)
}
task, err := s.m.pg.GetTask(*finding.TaskID)
if err != nil {
t.Fatal(err)
}
other, err := s.evidenceStore().Record(t.Context(), db.RecordFindingInput{TaskID: task.ID, ExplorationID: task.ExplorationID, Summary: "other", Severity: "low"}, nil)
if err != nil {
t.Fatal(err)
}
if w = req("GET", fmt.Sprintf("/api/exploration/findings/%d/traffic/%d", other.FindingID, first), ""); w.Code != 404 {
t.Fatal(w.Code, w.Body)
}
exportURL := fmt.Sprintf("/api/exploration/findings/export?format=md-zip&scope=selected&ids=%d", f.FindingID)
w = req("GET", exportURL, "")
if w.Code != 200 {
t.Fatal(w.Code, w.Body)
}
archive, err := zip.NewReader(bytes.NewReader(w.Body.Bytes()), int64(w.Body.Len()))
if err != nil {
t.Fatal(err)
}
binaries := 0
linked := false
for _, entry := range archive.File {
r, err := entry.Open()
if err != nil {
t.Fatal(err)
}
raw, err := io.ReadAll(r)
r.Close()
if err != nil {
t.Fatal(err)
}
if strings.HasSuffix(entry.Name, "/response.bin") {
binaries++
if !bytes.Equal(raw, body) {
t.Fatal("export changed binary")
}
}
if strings.HasSuffix(entry.Name, ".md") && bytes.Contains(raw, []byte(fmt.Sprintf("evidence/%d/%d/", f.FindingID, third))) {
linked = true
}
}
if binaries != 3 || !linked {
t.Fatalf("missing attachments/links: %d %v", binaries, linked)
}
// Reports can use snapshots with capture stopped and original rows gone.
result, err := s.toolGetFindingTraffic().Call(context.Background(), json.RawMessage(fmt.Sprintf(`{"finding_id":"%d"}`, f.FindingID)), nil)
if err != nil || !strings.Contains(result.Flatten(), "proof note") {
t.Fatal(result, err)
}
// Corrupt body must fail before an attachment response is sent.
snap := list.Bindings[0].Snapshot
path := filepath.Join(s.m.dir, "evidence", "blobs", snap.RespHash[:2], snap.RespHash+".bin")
if err = os.WriteFile(path, []byte("corrupt"), 0600); err != nil {
t.Fatal(err)
}
if w = req("GET", exportURL, ""); w.Code < 400 || strings.Contains(w.Header().Get("Content-Type"), "zip") {
t.Fatal("corrupt export returned an attachment", w.Code)
}
}
func TestFindingTrafficArchiveV3RoundTripAndRetry(t *testing.T) {
s, f, _ := trafficEvidenceServer(t)
ctx := context.Background()
s.ctx = ctx
body := []byte("evidence without original traffic\x00\xff")
seedServerEvidenceFlow(t, s, "original", body)
list, err := s.evidenceStore().Bind(ctx, f.FindingID, []db.TrafficRef{{TrafficID: "original", Role: "proof", Note: "keep note"}})
if err != nil {
t.Fatal(err)
}
finding, err := s.m.pg.GetFinding(f.FindingID)
if err != nil {
t.Fatal(err)
}
tid := *finding.TaskID
sharedTask, err := s.m.CreateTask("shared evidence", "fixture", nil, 0, 0)
if err != nil {
t.Fatal(err)
}
sid, _ := strconv.ParseInt(sharedTask.ID, 10, 64)
defer s.m.pg.DeleteTask(sid)
shared, err := s.evidenceStore().Record(ctx, db.RecordFindingInput{TaskID: sid, ExplorationID: sharedTask.ExpID, Summary: "shared", Severity: "low"}, []db.TrafficRef{{TrafficID: "original"}})
if err != nil {
t.Fatal(err)
}
if _, err = s.m.traffic.DeleteHost("evidence.local"); err != nil {
t.Fatal(err)
}
if err = s.m.SetTaskPaused(strconv.FormatInt(tid, 10), true); err != nil {
t.Fatal(err)
}
job, err := s.m.pg.QueueTaskArchive(tid)
if err != nil {
t.Fatal(err)
}
if err = s.m.pg.EditFindingTraffic(ctx, f.FindingID, list.Bindings[0].ID, 1, nil, nil, true, nil); err == nil {
t.Fatal("write allowed after archive queued")
}
if err = s.runOneTaskArchiveJob(); err != nil {
t.Fatal(err)
}
if got, _ := s.m.pg.GetFinding(f.FindingID); got != nil {
t.Fatal("archive retained hot finding")
}
if err = s.evidenceStore().Collect(ctx, time.Now().Add(48*time.Hour)); err != nil {
t.Fatal(err)
}
if err = s.evidenceStore().WithBinding(ctx, shared.FindingID, shared.Traffic.Bindings[0].ID, func(b db.FindingTrafficBinding) error {
r, _, err := s.evidenceStore().OpenBody(b.Snapshot, "response")
if r != nil {
r.Close()
}
return err
}); err != nil {
t.Fatal("shared body deleted", err)
}
if _, err = s.m.pg.QueueTaskArchiveRestore(job.ID); err != nil {
t.Fatal(err)
}
// Force a database error after validated body installation; then retry normally.
if _, err = s.m.pg.Exec(`CREATE FUNCTION fail_evidence_restore_test() RETURNS trigger LANGUAGE plpgsql AS $$ BEGIN RAISE EXCEPTION 'restore fixture failure'; END $$`); err != nil {
t.Fatal(err)
}
if _, err = s.m.pg.Exec(`CREATE TRIGGER fail_evidence_restore_test BEFORE INSERT ON finding_traffic_bindings FOR EACH ROW EXECUTE FUNCTION fail_evidence_restore_test()`); err != nil {
t.Fatal(err)
}
defer s.m.pg.Exec(`DROP FUNCTION IF EXISTS fail_evidence_restore_test() CASCADE`)
if err = s.runOneTaskArchiveJob(); err == nil {
t.Fatal("expected restore failure")
}
if got, _ := s.m.pg.GetFinding(f.FindingID); got != nil {
t.Fatal("partial restore")
}
if _, err = s.m.pg.Exec(`DROP FUNCTION fail_evidence_restore_test() CASCADE`); err != nil {
t.Fatal(err)
}
if _, err = s.m.pg.QueueTaskArchiveRestore(job.ID); err != nil {
t.Fatal(err)
}
if err = s.runOneTaskArchiveJob(); err != nil {
t.Fatal(err)
}
restored, err := s.m.pg.GetFindingTraffic(ctx, f.FindingID)
if err != nil {
t.Fatal(err)
}
b := restored.Bindings[0]
if restored.Version != list.Version || b.ID != list.Bindings[0].ID || b.Note != "keep note" || b.Role != "proof" || b.Position != 0 {
t.Fatalf("restored: %+v", restored)
}
if err = s.evidenceStore().WithBinding(ctx, f.FindingID, b.ID, func(b db.FindingTrafficBinding) error {
r, _, err := s.evidenceStore().OpenBody(b.Snapshot, "response")
if err != nil {
return err
}
defer r.Close()
got, err := io.ReadAll(r)
if !bytes.Equal(got, body) {
t.Fatal("restore changed body")
}
return err
}); err != nil {
t.Fatal(err)
}
// The read tool enforces task visibility too.
result, err := s.toolGetFindingTraffic().Call(agent.WithRunInfo(ctx, agent.RunInfo{TaskID: sid}), json.RawMessage(fmt.Sprintf(`{"finding_id":"%d"}`, f.FindingID)), nil)
if err != nil || !strings.Contains(result.Flatten(), "不可读取") {
t.Fatal(result, err)
}
}
func TestFindingTrafficFailedReportDoesNotTrigger(t *testing.T) {
s, f, _ := trafficEvidenceServer(t)
finding, err := s.m.pg.GetFinding(f.FindingID)
if err != nil {
t.Fatal(err)
}
task, err := s.m.pg.GetTask(*finding.TaskID)
if err != nil {
t.Fatal(err)
}
var last int64
if err = s.m.pg.QueryRow(`SELECT COALESCE(max(id),0) FROM activity`).Scan(&last); err != nil {
t.Fatal(err)
}
if err = s.m.pg.SetSchedState(schedKeyLastToolCall, strconv.FormatInt(last, 10)); err != nil {
t.Fatal(err)
}
id, err := s.m.pg.Exploration(task.ExplorationID).AppendActivity(db.Activity{Kind: "tool_result", Tool: "report_finding", ToolUseID: "failure-fixture", IsError: true, Detail: "body missing"})
if err != nil {
t.Fatal(err)
}
// A nil server makes any attempted trigger fail this test immediately.
scheduler := &Scheduler{pg: s.m.pg}
scheduler.fireToolCalls([]*db.AgentTrigger{{OnToolCall: true, ToolNames: []string{"report_finding"}, AgentKey: "reporter"}})
if got := scheduler.mustState(schedKeyLastToolCall); got != strconv.FormatInt(id, 10) {
t.Fatal("failed finding must still advance watermark", got)
}
}
func TestFindingTrafficUTF8SegmentsAndInheritedWrites(t *testing.T) {
s, f, req := trafficEvidenceServer(t)
ctx := context.Background()
text := strings.Repeat("中文证据🙂", 2000)
seedServerEvidenceFlow(t, s, "unicode", []byte(text))
list, err := s.evidenceStore().Bind(ctx, f.FindingID, []db.TrafficRef{{TrafficID: "unicode"}})
if err != nil {
t.Fatal(err)
}
var joined strings.Builder
offset := int64(0)
for {
var preview evidencePreview
err = s.evidenceStore().WithBinding(ctx, f.FindingID, list.Bindings[0].ID, func(b db.FindingTrafficBinding) error {
var e error
preview, e = readEvidencePreview(s.evidenceStore(), b.Snapshot, "response", offset, 8192)
return e
})
if err != nil {
t.Fatal(err)
}
if preview.Binary {
t.Fatal("valid UTF-8 classified binary")
}
joined.WriteString(preview.Content)
if !preview.Truncated {
break
}
if preview.NextOffset <= offset {
t.Fatal("preview made no progress")
}
offset = preview.NextOffset
}
if joined.String() != text {
t.Fatal("UTF-8 character split damaged preview")
}
finding, err := s.m.pg.GetFinding(f.FindingID)
if err != nil {
t.Fatal(err)
}
child, err := s.m.CreateTaskWithOptions("inherited evidence API", "fixture", db.TaskCreateOptions{SourceTaskIDs: []int64{*finding.TaskID}})
if err != nil {
t.Fatal(err)
}
cid, _ := strconv.ParseInt(child.ID, 10, 64)
defer s.m.pg.DeleteTask(cid)
base := fmt.Sprintf("/api/exploration/findings/%d/traffic", f.FindingID)
suffix := "?context_task=" + child.ID
if w := req("GET", base+suffix, ""); w.Code != 200 {
t.Fatal(w.Code, w.Body)
}
for _, method := range []string{"POST", "PATCH", "PUT", "DELETE"} {
path := base
if method == "PATCH" || method == "DELETE" {
path += fmt.Sprintf("/%d", list.Bindings[0].ID)
}
if method == "PUT" {
path += "/order"
}
w := req(method, path+suffix, `{"version":1,"traffic_refs":[{"traffic_id":"unicode"}]}`)
if w.Code != 403 {
t.Fatal(method, w.Code, w.Body)
}
}
if _, err = s.m.traffic.DeleteHost("evidence.local"); err != nil {
t.Fatal(err)
}
result, err := s.toolUpdateFindingReport().Call(ctx, json.RawMessage(fmt.Sprintf(`{"finding_id":%d,"evidence_version":1,"report":"## 证据报告\n\n证据 #%d:已验证完整请求响应"}`, f.NodeID, list.Bindings[0].ID)), nil)
if err != nil || result.IsError {
t.Fatal(result, err)
}
updated, err := s.m.pg.GetFinding(f.FindingID)
if err != nil {
t.Fatal(err)
}
if updated.ReportEvidenceVersion != 1 || updated.EvidenceVersion != 1 {
t.Fatal("report did not record the read evidence version")
}
for _, format := range []string{"md-single", "json", "csv"} {
w := req("GET", fmt.Sprintf("/api/exploration/findings/export?scope=selected&ids=%d&format=%s", f.FindingID, format), "")
if w.Code != 200 {
t.Fatal(format, w.Code, w.Body)
}
if !strings.Contains(w.Body.String(), fmt.Sprint(list.Bindings[0].ID)) {
t.Fatal("export omitted evidence ID", format)
}
}
}
+157
View File
@@ -0,0 +1,157 @@
package server
import (
"context"
"encoding/json"
"errors"
"fmt"
"log"
"strconv"
"github.com/Autumn-27/artex/agent"
"github.com/Autumn-27/artex/db"
actool "github.com/Autumn-27/norma/tool"
)
func (s *Server) seedFindingWorkflowTools() {
const flag = "finding_workflow_tools_v2_reporter"
if value, _, _ := s.m.pg.GetSetting(flag); value == "true" {
return
}
for _, key := range []string{"report_finding", "add_hint", "add_task_hint"} {
row, err := s.m.pg.GetTool(key)
if err != nil {
log.Printf("[evidence] load %s: %v", key, err)
return
}
if row == nil || !row.System {
continue
}
var schema map[string]any
if err := json.Unmarshal(row.Schema, &schema); err != nil {
log.Printf("[evidence] invalid schema for %s: %v", key, err)
return
}
if schema == nil {
log.Printf("[evidence] missing object schema for %s", key)
return
}
props := objectProperty(schema, "properties")
if key == "report_finding" {
if _, exists := props["evidence_hint_id"]; !exists {
props["evidence_hint_id"] = map[string]any{"type": "integer", "description": "可选:本任务中对应此漏洞的 hint ID;读取该提示保存的 traffic_refs 一并绑定,无提示时省略"}
}
} else {
if _, exists := props["traffic_refs"]; !exists {
props["traffic_refs"] = agent.HintTrafficSchema()
}
hints := objectProperty(props, "hints")
if _, exists := hints["type"]; !exists {
hints["type"] = "array"
}
items := objectProperty(hints, "items")
if _, ok := items["type"]; !ok {
items["type"] = "object"
}
itemProps := objectProperty(items, "properties")
for name, value := range map[string]any{"text": strParam("提示内容"), "asset_ids": map[string]any{"type": "array", "items": map[string]any{"type": "integer"}}, "traffic_refs": agent.HintTrafficSchema()} {
if _, exists := itemProps[name]; !exists {
itemProps[name] = value
}
}
}
raw, _ := json.Marshal(schema)
result, err := s.m.pg.Exec(`UPDATE tools SET schema=$2::jsonb,updated_at=now() WHERE key=$1 AND system AND schema=$3::jsonb`, key, string(raw), string(row.Schema))
if err != nil {
log.Printf("[evidence] upgrade %s: %v", key, err)
return
}
if n, _ := result.RowsAffected(); n != 1 {
return
} // preserve concurrent user edits
}
// Upgrade only the original default binding. Customized lists and enabled
// flags survive; the one-time flag also preserves future user unbinding.
readers := `["worker","reporter"]`
for _, key := range []string{"traffic_search", "traffic_get", "traffic_blob"} {
if _, err := s.m.pg.Exec(`UPDATE tools SET agents=$2::jsonb WHERE key=$1 AND system AND (agents='["worker"]'::jsonb OR (agents @> '["worker","planner","mainagent","auto","pentest"]'::jsonb AND jsonb_array_length(agents)=5))`, key, readers); err != nil {
return
}
}
if _, err := s.m.pg.Exec(`UPDATE tools SET agents=$1::jsonb WHERE key='get_finding_traffic' AND system AND agents @> '["auto","reporter"]'::jsonb AND jsonb_array_length(agents)=2`, `["auto","reporter","worker","planner","mainagent","pentest"]`); err != nil {
return
}
// Replace the previous code default only; preserve customized binding lists.
if _, err := s.m.pg.Exec(`UPDATE tools SET agents='["reporter"]'::jsonb WHERE key='bind_finding_traffic' AND system AND agents @> '["worker","planner","mainagent","auto","pentest"]'::jsonb AND jsonb_array_length(agents)=5`); err != nil {
return
}
if err := s.m.pg.AddAgentToToolBinding("reporter", []string{"bind_finding_traffic"}); err != nil {
return
}
_ = s.m.pg.SetSetting(flag, "true")
}
func objectProperty(parent map[string]any, key string) map[string]any {
value, ok := parent[key].(map[string]any)
if !ok {
value = map[string]any{}
parent[key] = value
}
return value
}
func (s *Server) agentFindingTrafficAccess(ctx context.Context, id int64, write bool) error {
if id <= 0 {
return errors.New("finding_id 必须为独立漏洞记录 ID;不是探索节点 ID")
}
f, err := s.m.pg.GetFinding(id)
if err != nil {
return err
}
if f == nil {
return fmt.Errorf("%w:finding_id=%d。证据工具使用独立漏洞记录 ID,请从 list_task_findings / get_task_node_detail 的 finding_id 字段读取;不要传 id / finding_node_id", db.ErrFindingNotFound, id)
}
if ri := agent.RunInfoFrom(ctx); ri.TaskID > 0 {
task := s.m.ResolveTask(strconv.FormatInt(ri.TaskID, 10))
if task == nil {
return errors.New("任务不存在")
}
_, inherited, allowed := findingProvenanceInTask(task, f.TaskID)
if !allowed {
return errors.New("当前任务不可读取该漏洞")
}
if write && inherited {
return errors.New("继承漏洞的流量证据只读,请到来源任务修改")
}
}
return nil
}
func (s *Server) toolBindFindingTraffic() actool.CoreTool {
return wrTool("bind_finding_traffic", "为已登记漏洞补绑经核实的真实 HTTP 流量。finding_id 使用独立漏洞记录 ID;不要传探索节点 ID。同批引用全部成功或全部失败,重复引用不覆盖已有说明。补绑会使已有报告标记待更新;不要为补包重新探测或重复创建漏洞。",
objSchema(map[string]any{"finding_id": strParam("独立漏洞记录 ID,从 list_task_findings / get_task_node_detail 的 finding_id 字段读取"), "traffic_refs": agent.HintTrafficSchema()}, "finding_id", "traffic_refs"),
func(ctx context.Context, raw json.RawMessage) (actool.Result, error) {
if !s.m.pg.GetBool(settingAgentTrafficBinding, false) {
return actool.Errorf("Agent 自动绑定流量已关闭;请在系统设置开启,或使用页面人工绑定。"), nil
}
var args struct {
FindingID json.RawMessage `json:"finding_id"`
Refs []db.TrafficRef `json:"traffic_refs"`
}
if err := json.Unmarshal(raw, &args); err != nil {
return actool.Errorf(err.Error()), nil
}
id := parseProfileID(args.FindingID)
if err := s.agentFindingTrafficAccess(ctx, id, true); err != nil {
return actool.Errorf(err.Error()), nil
}
if len(args.Refs) == 0 {
return actool.Errorf("补绑需要至少一条已核实的 traffic_refs;无流量无需调用此工具"), nil
}
list, err := s.evidenceStore().Bind(ctx, id, args.Refs)
if err != nil {
return actool.Errorf(err.Error()), nil
}
return jsonResult(trafficSummary(list))
})
}
+312
View File
@@ -0,0 +1,312 @@
package server
import (
"context"
"encoding/json"
"fmt"
"strconv"
"strings"
"testing"
"github.com/Autumn-27/artex/agent"
"github.com/Autumn-27/artex/db"
actool "github.com/Autumn-27/norma/tool"
)
func workflowCall(t *testing.T, ctx context.Context, tool actool.CoreTool, input any, wantError bool) string {
t.Helper()
raw, err := json.Marshal(input)
if err != nil {
t.Fatal(err)
}
r, err := tool.Call(ctx, raw, nil)
if err != nil || r.IsError != wantError {
t.Fatalf("%s: error=%v result=%s", tool.Name(), err, r.Flatten())
}
return r.Flatten()
}
func workflowTool(t *testing.T, tools []actool.CoreTool, name string) actool.CoreTool {
t.Helper()
for _, tool := range tools {
if tool.Name() == name {
return tool
}
}
t.Fatalf("missing tool: %s", name)
return nil
}
func TestFindingWorkflowAutoHintToPlannerAndSetting(t *testing.T) {
s, initial, request := trafficEvidenceServer(t)
pg := s.m.pg
ctx := context.Background()
if _, err := pg.Exec(`DELETE FROM settings WHERE key=$1`, settingAgentTrafficBinding); err != nil {
t.Fatal(err)
}
t.Cleanup(func() { pg.SetBool(settingAgentTrafficBinding, false) })
if s.settingsPayload()[settingAgentTrafficBinding] != false {
t.Fatal("missing setting must default off")
}
seedServerEvidenceFlow(t, s, "handoff-proof", []byte("verified local proof"))
seedServerEvidenceFlow(t, s, "handoff-baseline", []byte("local baseline"))
seedServerEvidenceFlow(t, s, "handoff-verification", []byte{0, 1, 255})
// Manual binding is independent of the Agent switch.
if r := request("POST", fmt.Sprintf("/api/exploration/findings/%d/traffic", initial.FindingID), `{"traffic_refs":[{"traffic_id":"handoff-proof"}]}`); r.Code != 200 {
t.Fatal(r.Code, r.Body)
}
bind := s.toolBindFindingTraffic()
input := map[string]any{"finding_id": initial.FindingID, "traffic_refs": []db.TrafficRef{{TrafficID: "handoff-baseline"}}}
workflowCall(t, ctx, bind, input, true)
if r := request("PUT", "/api/settings", `{"agent_traffic_binding":true}`); r.Code != 200 {
t.Fatal(r.Code, r.Body)
}
if !pg.GetBool(settingAgentTrafficBinding, false) {
t.Fatal("switch was not persisted")
}
f, _ := pg.GetFinding(initial.FindingID)
task := s.m.ResolveTask(fmt.Sprint(*f.TaskID))
refs := []db.TrafficRef{{TrafficID: "handoff-baseline", Role: "baseline", Note: "normal response"}, {TrafficID: "handoff-proof", Role: "proof", Note: "proves the finding"}, {TrafficID: "handoff-verification", Role: "verification", Note: "binary verification"}}
// Real cross-task Auto hint handler -> persisted graph -> Planner tool.
hintResult := workflowCall(t, ctx, s.toolAddHint(), map[string]any{"task_id": task.ID, "hints": []any{map[string]any{"text": "Report the confirmed local finding with its verified evidence", "traffic_refs": refs}}}, false)
var hints struct {
IDs []int64 `json:"ids"`
}
if err := json.Unmarshal([]byte(hintResult), &hints); err != nil || len(hints.IDs) != 1 || hints.IDs[0] == 0 {
t.Fatal(hintResult, err)
}
graph := workflowCall(t, ctx, s.toolGetTaskGraph(), map[string]any{"task_id": task.ID}, false)
for _, ref := range refs {
if !strings.Contains(graph, ref.TrafficID) {
t.Fatal("handoff lost reference", ref.TrafficID)
}
}
for range 6 {
if _, err := task.Store.AddNode(db.KindFact, map[string]any{"summary": "ID separation"}, 0, "confirmed", "test", nil); err != nil {
t.Fatal(err)
}
}
ts := agent.NewToolSet(task.Store, "planner")
ts.SetTaskID(*f.TaskID)
ts.SetFindingRecorder(s.evidenceStore())
notices := 0
ts.SetNotifyFinding(func(int64, string) { notices++ })
tools, def, cleanup := agent.AugmentTools(ctx, "planner", ts.PlannerTools())
defer cleanup()
if !strings.Contains(def.FindingGuidance, "evidence_hint_id") || !strings.Contains(def.FindingGuidance, "取消 Worker") {
t.Fatal("Planner missed runtime guidance")
}
report := workflowTool(t, tools, "report_finding")
result := workflowCall(t, ctx, report, map[string]any{"vulnclass": "TEST", "severity": "low", "summary": "Planner handoff fixture", "evidence_hint_id": hints.IDs[0]}, false)
var recorded db.RecordedFinding
if err := json.Unmarshal([]byte(strings.SplitN(result, "\n", 2)[1]), &recorded); err != nil {
t.Fatal(err)
}
if recorded.FindingID == recorded.NodeID || len(recorded.Traffic.Bindings) != 3 || notices != 1 {
t.Fatal("bad finding/traffic result", result)
}
for i, b := range recorded.Traffic.Bindings {
if b.Snapshot.SourceTrafficID != refs[i].TrafficID || b.Role != refs[i].Role || b.Note != refs[i].Note {
t.Fatalf("handoff mismatch: %+v", b)
}
}
list := workflowCall(t, ctx, s.toolListTaskFindings(), map[string]any{"task_id": task.ID}, false)
var nodes []struct {
ID int64 `json:"id"`
FindingID int64 `json:"finding_id"`
FindingNodeID int64 `json:"finding_node_id"`
Count int `json:"traffic_count"`
}
if err := json.Unmarshal([]byte(list), &nodes); err != nil {
t.Fatal(err)
}
found := false
for _, n := range nodes {
if n.ID == recorded.NodeID {
found = n.FindingID == recorded.FindingID && n.FindingNodeID == recorded.NodeID && n.Count == 3
}
}
if !found {
t.Fatal("canonical IDs/count missing", list)
}
workflowCall(t, ctx, s.toolGetFindingTraffic(), map[string]any{"finding_id": recorded.FindingID}, false)
wrong := workflowCall(t, ctx, s.toolGetFindingTraffic(), map[string]any{"finding_id": int64(900000000000000000)}, true)
if !strings.Contains(wrong, "独立漏洞记录 ID") {
t.Fatal("ambiguous ID error", wrong)
}
input["finding_id"] = recorded.FindingID
// Duplicate append preserves metadata and version.
workflowCall(t, ctx, bind, input, false)
after, err := pg.GetFindingTraffic(ctx, recorded.FindingID)
if err != nil || len(after.Bindings) != 3 || after.Version != recorded.Traffic.Version || after.Bindings[0].Note != refs[0].Note {
t.Fatal("duplicate changed evidence", after, err)
}
// A malformed handoff must not create a partial finding or notify Planner.
badHint := workflowCall(t, ctx, s.toolAddHint(), map[string]any{"task_id": task.ID, "text": "missing packet", "traffic_refs": []db.TrafficRef{{TrafficID: "missing"}}}, false)
var badID int64
fmt.Sscanf(badHint, "hint added: %d", &badID)
workflowCall(t, ctx, report, map[string]any{"vulnclass": "TEST", "severity": "low", "summary": "must fail", "evidence_hint_id": badID}, true)
if notices != 1 {
t.Fatal("failed handoff notified before commit")
}
child, err := s.m.CreateTaskWithOptions("inherited handoff", "fixture", db.TaskCreateOptions{SourceTaskIDs: []int64{*f.TaskID}})
if err != nil {
t.Fatal(err)
}
childID, _ := strconv.ParseInt(child.ID, 10, 64)
childRow, _ := pg.GetTask(childID)
t.Cleanup(func() { pg.DeleteTask(childRow.ID) })
childCtx := agent.WithRunInfo(ctx, agent.RunInfo{TaskID: childRow.ID})
workflowCall(t, childCtx, bind, input, true)
workflowCall(t, childCtx, s.toolGetFindingTraffic(), map[string]any{"finding_id": recorded.FindingID}, false)
// Existing, already assembled tools must honor a later switch-off.
if r := request("PUT", "/api/settings", `{"agent_traffic_binding":false}`); r.Code != 200 {
t.Fatal(r.Code, r.Body)
}
workflowCall(t, ctx, bind, input, true)
workflowCall(t, ctx, report, map[string]any{"vulnclass": "TEST", "severity": "low", "summary": "off", "evidence_hint_id": hints.IDs[0]}, true)
workflowCall(t, ctx, report, map[string]any{"vulnclass": "TCP", "severity": "low", "summary": "no packet needed"}, false)
if notices != 2 {
t.Fatal("optional no-packet report failed")
}
workflowCall(t, ctx, s.toolGetFindingTraffic(), map[string]any{"finding_id": recorded.FindingID}, false)
tools, def, closeTools := agent.AugmentTools(ctx, "planner", ts.PlannerTools())
defer closeTools()
if def.FindingGuidance != "" {
t.Fatal("disabled feature still injects guidance")
}
for _, tool := range tools {
if tool.Name() == "bind_finding_traffic" {
t.Fatal("disabled binding tool exposed")
}
}
}
func TestFindingWorkflowMigrationPreservesUserConfiguration(t *testing.T) {
s, _, _ := trafficEvidenceServer(t)
pg := s.m.pg
for _, key := range []string{"report_finding", "add_hint", "add_task_hint", "traffic_search", "traffic_get", "traffic_blob", "bind_finding_traffic", "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(key, old.Description, old.Schema, bindings, old.Enabled)
})
}
custom := json.RawMessage(`{"type":"object","properties":{"text":{"type":"string","description":"USER TEXT"},"hints":{"type":"array","items":{"type":"object","properties":{"text":{"type":"string","description":"USER ITEM"}}}}},"required":["text"]}`)
for _, key := range []string{"report_finding", "add_hint", "add_task_hint"} {
if err := pg.UpdateTool(key, "USER DESCRIPTION", custom, json.RawMessage(`["custom-agent"]`), false); err != nil {
t.Fatal(err)
}
}
for _, key := range []string{"traffic_search", "traffic_get"} {
row, _ := pg.GetTool(key)
bindings := json.RawMessage(`["custom-agent"]`)
if key == "traffic_search" {
bindings = json.RawMessage(`["worker"]`)
}
if err := pg.UpdateTool(key, row.Description, row.Schema, bindings, false); err != nil {
t.Fatal(err)
}
}
if err := pg.SetSetting("finding_workflow_tools_v2_reporter", "false"); err != nil {
t.Fatal(err)
}
s.seedFindingWorkflowTools()
for _, key := range []string{"report_finding", "add_hint", "add_task_hint"} {
row, _ := pg.GetTool(key)
if row.Enabled || row.Description != "USER DESCRIPTION" || len(row.Agents) != 1 || row.Agents[0] != "custom-agent" {
t.Fatal("changed user tool configuration", key)
}
if !strings.Contains(string(row.Schema), "USER TEXT") || !strings.Contains(string(row.Schema), "USER ITEM") {
t.Fatal("changed user schema", key)
}
field := "traffic_refs"
if key == "report_finding" {
field = "evidence_hint_id"
}
if !strings.Contains(string(row.Schema), field) {
t.Fatal("missing optional field", key)
}
}
search, _ := pg.GetTool("traffic_search")
get, _ := pg.GetTool("traffic_get")
if search.Enabled || !contains(search.Agents, "reporter") || get.Enabled || len(get.Agents) != 1 || get.Agents[0] != "custom-agent" {
t.Fatal("default/custom reader binding migration incorrect")
}
if err := pg.RemoveAgentFromTool("reporter", "traffic_search"); err != nil {
t.Fatal(err)
}
s.seedFindingWorkflowTools()
search, _ = pg.GetTool("traffic_search")
if contains(search.Agents, "reporter") {
t.Fatal("one-time migration undid later unbinding")
}
}
func TestFindingWorkflowReporterBindsBeforeWritingReport(t *testing.T) {
s, f, request := trafficEvidenceServer(t)
pg := s.m.pg
ctx := context.Background()
if err := pg.SetBool(settingAgentTrafficBinding, true); err != nil {
t.Fatal(err)
}
t.Cleanup(func() { pg.SetBool(settingAgentTrafficBinding, false); s.m.SetTrafficEnabled(false) })
if err := s.m.SetTrafficEnabled(true); err != nil {
t.Fatal(err)
}
seedServerEvidenceFlow(t, s, "reporter-proof", []byte("local proof payload"))
seedServerEvidenceFlow(t, s, "reporter-baseline", []byte("local normal response"))
tools, def, cleanup := agent.AugmentTools(ctx, "reporter", nil)
defer cleanup()
if !strings.Contains(def.FindingGuidance, "报告前自动关联流量") || !strings.Contains(def.FindingGuidance, "绑定成功后重新调用") {
t.Fatal("reporter did not receive binding workflow")
}
for _, name := range []string{"traffic_search", "traffic_get", "get_task_worker_trace", "get_task_node_detail", "bind_finding_traffic", "get_finding_traffic", "update_finding_report"} {
workflowTool(t, tools, name)
}
finding, _ := pg.GetFinding(f.FindingID)
detail := workflowCall(t, ctx, workflowTool(t, tools, "get_task_node_detail"), map[string]any{"task_id": fmt.Sprint(*finding.TaskID), "id": f.NodeID}, false)
if !strings.Contains(detail, `"finding_id"`) || !strings.Contains(detail, `"finding_node_id"`) {
t.Fatal("reporter lacks explicit ID mapping")
}
workflowCall(t, ctx, workflowTool(t, tools, "traffic_search"), map[string]any{"host": "evidence.local"}, false)
for _, id := range []string{"reporter-baseline", "reporter-proof"} {
packet := workflowCall(t, ctx, workflowTool(t, tools, "traffic_get"), map[string]any{"id": id}, false)
if !strings.Contains(packet, "local") {
t.Fatal("reporter cannot inspect packet", packet)
}
}
refs := []db.TrafficRef{{TrafficID: "reporter-baseline", Role: "baseline"}, {TrafficID: "reporter-proof", Role: "proof", Note: "confirmed from recorded response"}}
workflowCall(t, ctx, workflowTool(t, tools, "bind_finding_traffic"), map[string]any{"finding_id": f.FindingID, "traffic_refs": refs}, false)
list := workflowCall(t, ctx, workflowTool(t, tools, "get_finding_traffic"), map[string]any{"finding_id": f.FindingID}, false)
var summary struct {
Version int64 `json:"version"`
}
if err := json.Unmarshal([]byte(list), &summary); err != nil || summary.Version != 1 {
t.Fatal(list, err)
}
workflowCall(t, ctx, workflowTool(t, tools, "update_finding_report"), map[string]any{"finding_id": f.NodeID, "evidence_version": summary.Version, "report": "## Local report\n\nVerified baseline and proof using saved evidence."}, false)
updated, err := pg.GetFinding(f.FindingID)
if err != nil || updated.ReportEvidenceVersion != summary.Version || updated.EvidenceVersion != summary.Version {
t.Fatal("report did not cover post-binding version", updated, err)
}
// Disabling automatic binding still lets Reporter read manually bound snapshots.
if r := request("PUT", "/api/settings", `{"agent_traffic_binding":false}`); r.Code != 200 {
t.Fatal(r.Code, r.Body)
}
off, offDef, closeOff := agent.AugmentTools(ctx, "reporter", nil)
defer closeOff()
if offDef.FindingGuidance != "" {
t.Fatal("off reporter still receives auto-binding guidance")
}
for _, tool := range off {
if tool.Name() == "traffic_search" || tool.Name() == "traffic_get" || tool.Name() == "traffic_blob" || tool.Name() == "bind_finding_traffic" {
t.Fatal("auto-binding tool exposed while off", tool.Name())
}
}
workflowCall(t, ctx, workflowTool(t, off, "get_finding_traffic"), map[string]any{"finding_id": f.FindingID}, false)
workflowTool(t, off, "update_finding_report")
}
+7 -6
View File
@@ -249,12 +249,13 @@ type Manager struct {
// Settings keys the UI toggles at runtime.
const (
settingTrafficCapture = "traffic_capture"
settingWebSearchOn = "web_search_enabled"
settingWebSearchBackend = "web_search_backend"
settingBraveKey = "brave_search_api_key"
settingTavilyKey = "tavily_search_api_key"
settingWebSearchProxy = "web_search_proxy"
settingTrafficCapture = "traffic_capture"
settingAgentTrafficBinding = "agent_traffic_binding"
settingWebSearchOn = "web_search_enabled"
settingWebSearchBackend = "web_search_backend"
settingBraveKey = "brave_search_api_key"
settingTavilyKey = "tavily_search_api_key"
settingWebSearchProxy = "web_search_proxy"
// settingGlobalProxy is the global egress proxy for all target traffic
// (http/https/socks5). Empty = direct. Distinct from web_search_proxy (which
// only routes the search backend) and the per-profile LLM proxy.
+30 -19
View File
@@ -76,6 +76,8 @@ func (s *Server) orchestrationTools() []actool.CoreTool {
s.toolSearchWorkerTraces(),
s.toolGetTaskNodeDetail(),
s.toolUpdateFindingReport(),
s.toolGetFindingTraffic(),
s.toolBindFindingTraffic(),
}
}
@@ -334,10 +336,11 @@ func (s *Server) toolAddHint() actool.CoreTool {
return wrTool("add_task_hint", "给指定任务注入战略提示(该任务的 planner 下轮生成意图时会读到)。\n"+
"★优先批量:多条提示放进 hints 数组一次提交(返回 ids 数组,与 hints 等长同序,失败项 id=0);单条则省略 hints 直接给顶层 text。",
objSchema(map[string]any{
"task_id": strParam("任务 id"),
"hints": map[string]any{"type": "array", "description": "【优先用这个】提示数组,每个元素字段同顶层(text/asset_ids)。", "items": map[string]any{"type": "object"}},
"text": strParam("[单条] 提示内容"),
"asset_ids": map[string]any{"type": "array", "items": map[string]any{"type": "integer"}, "description": "锚定的资产 id(可选,0/1/多个;该任务内的资产 id)"},
"task_id": strParam("任务 id"),
"hints": map[string]any{"type": "array", "description": "【优先用这个】提示数组,每个元素字段同顶层(text/asset_ids/traffic_refs)。", "items": objSchema(map[string]any{"text": strParam("提示内容"), "asset_ids": map[string]any{"type": "array", "items": map[string]any{"type": "integer"}}, "traffic_refs": agent.HintTrafficSchema()})},
"text": strParam("[单条] 提示内容"),
"traffic_refs": agent.HintTrafficSchema(),
"asset_ids": map[string]any{"type": "array", "items": map[string]any{"type": "integer"}, "description": "锚定的资产 id(可选,0/1/多个;该任务内的资产 id)"},
}, "task_id"),
func(ctx context.Context, in json.RawMessage) (actool.Result, error) {
return s.delegateToTask(ctx, in, (*agent.ToolSet).AddHintTool)
@@ -393,20 +396,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
}
@@ -440,7 +445,11 @@ func (s *Server) seedOrchestrationTools() {
autoAgents, _ := json.Marshal([]string{"auto"})
for _, t := range s.orchestrationTools() {
schema, _ := json.Marshal(t.InputSchema())
_ = s.m.PG().SeedTool(t.Name(), t.Description(), schema, autoAgents)
bindings := autoAgents
if t.Name() == "bind_finding_traffic" {
bindings = json.RawMessage(`["reporter"]`)
}
_ = s.m.PG().SeedTool(t.Name(), t.Description(), schema, bindings)
}
for _, t := range s.platformTools() {
schema, _ := json.Marshal(t.InputSchema())
@@ -454,11 +463,13 @@ 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() // 增加可选证据参数及只读证据工具,保留用户配置
s.seedFindingWorkflowTools()
// 注:pentest 的默认工具绑定无需迁移——BuiltinToolSeeds 在全新初始化时就把
// list_assets/insert_assets/report_finding/list_findings/list_companies 连同
// pentest 一起 seed 好了(项目尚无旧库,不做迁移)。
@@ -684,9 +695,9 @@ func (s *Server) seedReporterAgent() {
Enabled: true,
OnToolCall: true,
ToolNames: []string{"report_finding"},
ToolCallMessage: "上面刚有一个漏洞被 report_finding 登记。请从触发上下文里取出 finding_id" +
"(工具返回 \"finding recorded: <id>\" 里的数字)与任务 id,按你的职责撰写该漏洞的详细报告," +
"最后调用 update_finding_report(finding_id, report) 保存。",
ToolCallMessage: "上面刚有一个漏洞被 report_finding 登记。请读取返回 JSON 的 finding_id(独立漏洞记录 ID)与 finding_node_id(探索节点 ID)," +
"若运行指引启用自动绑定,先核实并关联本次漏洞的流量,再用 get_finding_traffic(finding_id) 读取最新证据;节点详情使用 finding_node_id。" +
"最后调用 update_finding_report(finding_id=finding_node_id, report, evidence_version=实际读取版本) 保存。不要混用两种编号。",
}); err != nil {
log.Printf("[reporter] 创建触发器失败: %v", err)
}
+5
View File
@@ -305,6 +305,11 @@ func (sc *Scheduler) fireToolCalls(triggers []*db.AgentTrigger) {
if len(want) == 0 {
continue
}
// A failed finding write has no committed finding to report. Keep other
// tool-error triggers available for user-defined automation.
if e.ToolIsErr && e.Tool == "report_finding" {
continue
}
errTag := ""
if e.ToolIsErr {
errTag = "[error] "
+43 -24
View File
@@ -1,7 +1,6 @@
package server
import (
"archive/zip"
"context"
"encoding/json"
"errors"
@@ -11,6 +10,7 @@ import (
"net"
"net/http"
"net/url"
"os"
"path/filepath"
"regexp"
"sort"
@@ -223,6 +223,7 @@ func New(ctx context.Context, m *Manager, skillDir string, dataDir string, keyDi
if err := s.seedFindingRetester(); err != nil {
log.Printf("[retester] seed: %v", err)
}
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
@@ -363,7 +364,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
@@ -408,6 +409,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")))
@@ -417,6 +419,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
@@ -485,6 +488,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 实时纠偏
@@ -814,6 +818,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)
@@ -2142,6 +2147,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) {
@@ -2154,27 +2169,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))
@@ -2182,6 +2189,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")
@@ -3129,6 +3140,7 @@ func (s *Server) settingsPayload() map[string]any {
}
return map[string]any{
"traffic_capture": s.m.TrafficEnabled(),
"agent_traffic_binding": s.m.pg.GetBool(settingAgentTrafficBinding, false),
"llm_record": s.m.LLMRecordEnabled(),
"web_search_enabled": on,
"web_search_backend": backend,
@@ -3169,8 +3181,9 @@ func (s *Server) pgDetectPython(w http.ResponseWriter, r *http.Request) {
// off, agents get no proxy config, no traffic tools, and no proxy prompt content.
func (s *Server) putSettings(w http.ResponseWriter, r *http.Request) {
var req struct {
TrafficCapture *bool `json:"traffic_capture"`
LLMRecord *bool `json:"llm_record"` // LLM 录制开关(默认关);即时生效,无需重建 agent
TrafficCapture *bool `json:"traffic_capture"`
AgentTrafficBinding *bool `json:"agent_traffic_binding"`
LLMRecord *bool `json:"llm_record"` // LLM 录制开关(默认关);即时生效,无需重建 agent
// Web search. WebSearchEnabled/Backend toggle the tool + backend; BraveKey/TavilyKey
// are optional — omit (null) to leave a stored key untouched, send "" to clear.
WebSearchEnabled *bool `json:"web_search_enabled"`
@@ -3267,6 +3280,12 @@ func (s *Server) putSettings(w http.ResponseWriter, r *http.Request) {
// built under the OLD switch state — drop them so they pick up the new one.
s.invalidateProfileAgents()
}
if req.AgentTrafficBinding != nil {
if err := s.m.pg.SetBool(settingAgentTrafficBinding, *req.AgentTrafficBinding); err != nil {
writeErr(w, 500, err.Error())
return
}
}
if req.TrafficCapture != nil {
if err := s.m.SetTrafficEnabled(*req.TrafficCapture); err != nil {
writeErr(w, 500, err.Error())
+20 -2
View File
@@ -178,6 +178,13 @@ func (s *Server) archiveTask(job *pgdb.TaskArchive) (runErr error) {
snapshot.DataCounts["traffic"] = trafficCount
}
archivePath := taskArchivePath(s.m.dir, job.ID, taskID)
evidenceSnapshots, err := pgdb.ArchiveEvidenceSnapshots(snapshot)
if err != nil {
return err
}
if err = s.evidenceStore().CopySnapshots(s.ctx, evidenceSnapshots, filepath.Join(fileStage.payload, "evidence")); err != nil {
return err
}
_ = s.m.pg.UpdateTaskArchiveProgress(job.ID, "compress", 50)
originalSize, compressedSize, checksum, err := writeTaskArchivePackage(archivePath, fileStage.payload, snapshot)
if err != nil {
@@ -270,6 +277,17 @@ func (s *Server) restoreTaskArchive(job *pgdb.TaskArchive) (runErr error) {
if snapshot.TaskID != job.TaskID || !pgdb.IsTaskArchiveFormatSupported(snapshot.FormatVersion) {
return pgdb.ErrTaskArchiveFormatMismatch
}
evidenceSnapshots, err := pgdb.ArchiveEvidenceSnapshots(&snapshot)
if err != nil {
return err
}
return s.evidenceStore().WithInstalledSnapshots(s.ctx, evidenceSnapshots, filepath.Join(extracted, "evidence"), func() error {
return s.restoreTaskArchivePayload(job, &snapshot, extracted)
})
}
func (s *Server) restoreTaskArchivePayload(job *pgdb.TaskArchive, snapshot *pgdb.TaskArchiveSnapshot, extracted string) (runErr error) {
var err error
var llmRecordsFile *os.File
relative, hasStreamedLLMRecords := snapshot.StreamedTables["llm_records"]
if len(snapshot.StreamedTables) > 0 &&
@@ -309,13 +327,13 @@ func (s *Server) restoreTaskArchive(job *pgdb.TaskArchive) (runErr error) {
var warnings []string
if llmRecordsFile != nil {
warnings, err = s.m.pg.RestoreTaskArchiveWithLLMRecords(
job.ID, &snapshot, job.RemainingTimeoutSeconds, llmRecordsFile,
job.ID, snapshot, job.RemainingTimeoutSeconds, llmRecordsFile,
)
closeErr := llmRecordsFile.Close()
llmRecordsFile = nil
err = errors.Join(err, closeErr)
} else {
warnings, err = s.m.pg.RestoreTaskArchive(job.ID, &snapshot, job.RemainingTimeoutSeconds)
warnings, err = s.m.pg.RestoreTaskArchive(job.ID, snapshot, job.RemainingTimeoutSeconds)
}
if err != nil {
return err
+3
View File
@@ -563,6 +563,7 @@ func (s *Server) agentsForTask(t *Task) *taskAgentBundle {
tx := transcript.NewStore(filepath.Join(s.m.dir, "transcripts"))
window := workerRuntime.CompactionWindow()
wk := agent.NewWorker(workerRuntime, "task-router", s.m.dir, tx, window, s.agentMaxTurns("worker"))
wk.SetFindingRecorder(s.evidenceStore())
wk.SetCompactionWindowResolver(workerRuntime.CompactionWindow)
wk.SetNonStreaming(workerRuntime.nonStreaming) // 按任务当前激活 profile 的流式开关(每轮读)
wk.SetMaxTokens(workerRuntime.maxTokens) // 同上,输出上限也跟随当前激活 profile
@@ -572,6 +573,7 @@ func (s *Server) agentsForTask(t *Task) *taskAgentBundle {
wk.SetWebSearch(s.webSearchFor("worker"))
wk.SetConstraintInject(s.constraintInjectWorker) // 操作约束注入 worker(可配置,默认开)
pl := agent.NewPlanner(plannerRuntime, "task-router", s.m.dir, tx, plannerRuntime.CompactionWindow(), s.agentMaxTurns("planner"))
pl.SetFindingRecorder(s.evidenceStore())
pl.SetCompactionWindowResolver(plannerRuntime.CompactionWindow)
pl.SetNonStreaming(plannerRuntime.nonStreaming)
pl.SetMaxTokens(plannerRuntime.maxTokens)
@@ -581,6 +583,7 @@ func (s *Server) agentsForTask(t *Task) *taskAgentBundle {
pl.SetWebSearch(s.webSearchFor("planner"))
pl.SetConstraintInject(s.constraintInjectPlanner) // 操作约束注入 planner(可配置,默认开)
main := agent.NewMainAgent(mainRuntime, "task-router", s.m.dir, tx, mainRuntime.CompactionWindow(), s.agentMaxTurns("mainagent"))
main.SetFindingRecorder(s.evidenceStore())
main.SetCompactionWindowResolver(mainRuntime.CompactionWindow)
main.SetNonStreaming(mainRuntime.nonStreaming)
main.SetMaxTokens(mainRuntime.maxTokens)
+126
View File
@@ -0,0 +1,126 @@
package traffic
import (
"bufio"
"bytes"
"context"
"database/sql"
"errors"
"fmt"
"io"
"os"
"path/filepath"
"strings"
)
// EvidenceExchange exposes complete captured bodies only during the callback.
// Readers are closed before releasing the writer lock, so host deletion cannot
// remove a blob between its index lookup and its evidence copy.
type EvidenceExchange struct {
ID, URL, Method, ContentType string
TS int64
Status int
ReqHead, RespHead string
ReqLen, RespLen int64
ReqHash, RespHash string
Request, Response io.Reader
}
func (t *Traffic) ReadEvidence(ctx context.Context, ids []string, consume func(EvidenceExchange) error) error {
if t == nil {
return errors.New("流量录制存储不可用")
}
t.wmu.Lock()
defer t.wmu.Unlock()
for _, id := range ids {
if err := ctx.Err(); err != nil {
return err
}
if err := t.readEvidence(id, consume); err != nil {
return fmt.Errorf("流量 %s: %w", id, err)
}
}
return nil
}
func (t *Traffic) readEvidence(id string, consume func(EvidenceExchange) error) error {
e := EvidenceExchange{ID: id}
var legacy string
if err := t.db.QueryRow(`SELECT ts,url,method,status,content_type,req_len,resp_len,path FROM exchanges WHERE id=?`, id).
Scan(&e.TS, &e.URL, &e.Method, &e.Status, &e.ContentType, &e.ReqLen, &e.RespLen, &legacy); err != nil {
return err
}
var req, resp []byte
var reqBlob, respBlob sql.NullString
err := t.db.QueryRow(`SELECT req_head,req_body,req_blob,resp_head,resp_body,resp_blob FROM exchange_bodies WHERE id=?`, id).
Scan(&e.ReqHead, &req, &reqBlob, &e.RespHead, &resp, &respBlob)
if errors.Is(err, sql.ErrNoRows) && legacy != "" {
if !filepath.IsLocal(legacy) {
return errors.New("旧流量路径无效")
}
r, err := os.Open(filepath.Join(t.dir, legacy, "request.http"))
if err != nil {
return err
}
defer r.Close()
p, err := os.Open(filepath.Join(t.dir, legacy, "response.http"))
if err != nil {
return err
}
defer p.Close()
e.ReqHead, e.Request, err = splitLegacyEvidence(r)
if err != nil {
return err
}
e.RespHead, e.Response, err = splitLegacyEvidence(p)
if err != nil {
return err
}
return consume(e)
}
if err != nil {
return err
}
e.Request = bytes.NewReader(req)
e.Response = bytes.NewReader(resp)
if reqBlob.Valid && reqBlob.String != "" {
f, _, err := t.Blob(reqBlob.String)
if err != nil {
return err
}
defer f.Close()
e.Request = f
e.ReqHash = reqBlob.String
}
if respBlob.Valid && respBlob.String != "" {
f, _, err := t.Blob(respBlob.String)
if err != nil {
return err
}
defer f.Close()
e.Response = f
e.RespHash = respBlob.String
}
return consume(e)
}
func splitLegacyEvidence(r io.Reader) (string, io.Reader, error) {
b := bufio.NewReader(r)
var head strings.Builder
for {
line, err := b.ReadString('\n')
if line == "\n" || line == "\r\n" {
return head.String(), b, nil
}
head.WriteString(line)
if head.Len() > 1<<20 {
return "", nil, errors.New("旧流量报文头过大")
}
if errors.Is(err, io.EOF) {
return head.String(), b, nil
}
if err != nil {
return "", nil, err
}
}
}
+17 -12
View File
@@ -19,11 +19,16 @@ export default function LoginPage() {
useEffect(() => {
// 已登录直接进主界面(静态导出下无 middleware 代劳这层跳转)。
if (auth.getToken()) {
router.replace("/function/tasks");
const token = auth.getToken();
if (token) {
// localStorage 可能仍有凭据但 cookie 已丢失。先同步,再发起全新请求,
// 避免服务端守卫或路由缓存把跳转送回仍处于 checking 状态的登录页。
auth.setToken(token);
window.location.replace("/function/tasks");
return;
}
api.authStatus()
api
.authStatus()
.then(({ initialized }) => {
if (!initialized) router.replace("/setup");
})
@@ -38,7 +43,7 @@ export default function LoginPage() {
try {
const { token } = await api.login("ARTEX", password);
auth.setToken(token);
router.replace("/function/tasks");
window.location.replace("/function/tasks");
} catch {
setError("用户名或密码错误");
} finally {
@@ -46,7 +51,13 @@ export default function LoginPage() {
}
}
if (checking) return null;
if (checking) {
return (
<div role="status" className="flex min-h-dvh items-center justify-center text-muted-foreground">
正在检查登录状态…
</div>
);
}
return (
<div className="flex h-dvh">
@@ -57,13 +68,7 @@ export default function LoginPage() {
<div className="absolute size-60 rounded-full border border-primary-foreground/15" />
<div className="absolute size-40 rounded-full border border-primary-foreground/20" />
{/* eslint-disable-next-line @next/next/no-img-element */}
<img
src="/logo.png"
alt="ARTEX"
width={160}
height={160}
className="relative brightness-0 invert"
/>
<img src="/logo.png" alt="ARTEX" width={160} height={160} className="relative brightness-0 invert" />
</div>
</div>
+71 -10
View File
@@ -925,6 +925,7 @@ export default function ChatPage() {
const [agents, setAgents] = React.useState<Agent[]>([]);
const [profiles, setProfiles] = React.useState<LLMProfile[]>([]);
const [convs, setConvs] = React.useState<Conversation[]>([]);
const [agentFilter, setAgentFilter] = React.useState<string | null>(null);
const [selectedId, setSelectedId] = React.useState<number | null>(null);
const [renamingId, setRenamingId] = React.useState<number | null>(null);
const [renameText, setRenameText] = React.useState("");
@@ -1015,17 +1016,44 @@ export default function ChatPage() {
const selected = React.useMemo(() => convs.find((c) => c.id === selectedId) ?? null, [convs, selectedId]);
const agentByKey = React.useMemo(() => new Map(agents.map((agent) => [agent.key, agent])), [agents]);
const filteredConversations = React.useMemo(
() => (agentFilter === null ? convs : convs.filter((conversation) => conversation.agent_key === agentFilter)),
[convs, agentFilter],
);
const visibleConversations = React.useMemo(
() => convs.slice(0, visibleConversationCount),
[convs, visibleConversationCount],
() => filteredConversations.slice(0, visibleConversationCount),
[filteredConversations, visibleConversationCount],
);
// conversation agents: custom agents + conversational built-ins (role=assistant,
// e.g. Auto / 渗透测试). The orchestration built-ins (goals/planner/mainagent/worker)
// are task-specific and stay hidden from the chat page.
const chatAgents = React.useMemo(() => agents.filter((a) => !a.builtin || a.role === "assistant"), [agents]);
const agentFilterOptions = React.useMemo(() => {
const counts = new Map<string, number>();
for (const conversation of convs) {
counts.set(conversation.agent_key, (counts.get(conversation.agent_key) ?? 0) + 1);
}
// Include historical sources even if their Agent has since been disabled or deleted.
const keys = new Set([...chatAgents.map((agent) => agent.key), ...counts.keys()]);
if (agentFilter !== null) keys.add(agentFilter);
return [...keys]
.map((key) => ({ key, name: agentByKey.get(key)?.name || key, count: counts.get(key) ?? 0 }))
.sort((a, b) => a.name.localeCompare(b.name, "zh-CN"));
}, [convs, chatAgents, agentByKey, agentFilter]);
const conversationCountLabel =
agentFilter === null ? `共 ${convs.length} 个` : `${filteredConversations.length} / ${convs.length} 个`;
function changeAgentFilter(key: string | null) {
setAgentFilter(key);
setVisibleConversationCount(CONVERSATION_LIST_PAGE);
setSelectedConversationIds(new Set());
setBulkDeleteOpen(false);
setRenamingId(null);
}
const selectedConversationCount = selectedConversationIds.size;
const allConversationsSelected = convs.length > 0 && selectedConversationCount === convs.length;
const allConversationsSelected =
filteredConversations.length > 0 && selectedConversationCount === filteredConversations.length;
const someConversationsSelected = selectedConversationCount > 0 && !allConversationsSelected;
let conversationHeaderChecked: boolean | "indeterminate" = false;
if (allConversationsSelected) conversationHeaderChecked = true;
@@ -1041,7 +1069,9 @@ export default function ChatPage() {
}, []);
function toggleAllConversations(checked: boolean) {
setSelectedConversationIds(checked ? new Set(convs.map((conversation) => conversation.id)) : new Set());
setSelectedConversationIds(
checked ? new Set(filteredConversations.map((conversation) => conversation.id)) : new Set(),
);
}
function exitSelectionMode() {
@@ -1154,16 +1184,37 @@ export default function ChatPage() {
<Button size="sm" className="w-full" onClick={() => setSelectedId(null)}>
<PlusIcon /> 新建对话
</Button>
<Select
value={agentFilter === null ? "all" : `agent:${agentFilter}`}
onValueChange={(value) => changeAgentFilter(value === "all" ? null : value.slice(6))}
disabled={bulkDeleting}
>
<SelectTrigger size="sm" className="w-full min-w-0" aria-label="按 Agent 筛选对话">
<Bot />
<SelectValue placeholder="全部 Agent" />
</SelectTrigger>
<SelectContent>
<SelectGroup>
<SelectItem value="all">全部 Agent</SelectItem>
{agentFilterOptions.map((agent) => (
<SelectItem key={agent.key} value={`agent:${agent.key}`}>
{agent.name}({agent.count})
</SelectItem>
))}
</SelectGroup>
</SelectContent>
</Select>
{convs.length > 0 &&
(selectionMode ? (
<div className="flex items-center gap-2 px-1">
<Checkbox
checked={conversationHeaderChecked}
onCheckedChange={(checked) => toggleAllConversations(checked === true)}
aria-label="选择全部对话"
aria-label="选择当前筛选的全部对话"
disabled={filteredConversations.length === 0 || bulkDeleting}
/>
<span className="text-muted-foreground min-w-0 flex-1 text-xs tabular-nums">
{selectedConversationCount > 0 ? `已选 ${selectedConversationCount} 个` : `共 ${convs.length} 个`}
{selectedConversationCount > 0 ? `已选 ${selectedConversationCount} 个` : conversationCountLabel}
</span>
{selectedConversationCount > 0 && (
<Button
@@ -1183,13 +1234,14 @@ export default function ChatPage() {
) : (
<div className="flex items-center gap-2 px-1">
<span className="text-muted-foreground min-w-0 flex-1 text-xs tabular-nums">
共 {convs.length} 个
{conversationCountLabel}
</span>
<Button
size="sm"
variant="ghost"
className="text-muted-foreground"
onClick={() => setSelectionMode(true)}
disabled={filteredConversations.length === 0}
>
<ListChecksIcon data-icon="inline-start" />
多选
@@ -1197,9 +1249,17 @@ export default function ChatPage() {
</div>
))}
</div>
<ScrollArea type="auto" className="min-h-0 min-w-0 flex-1 [&_[data-slot=scroll-area-viewport]>div]:block!">
<ScrollArea
key={agentFilter === null ? "all" : `agent:${agentFilter}`}
type="auto"
className="min-h-0 min-w-0 flex-1 [&_[data-slot=scroll-area-viewport]>div]:block!"
>
<div className="flex min-w-0 flex-col gap-0.5 p-2">
{convs.length === 0 && <p className="text-muted-foreground px-2 py-6 text-center text-xs">暂无对话</p>}
{filteredConversations.length === 0 && (
<p className="text-muted-foreground px-2 py-6 text-center text-xs">
{agentFilter === null ? "暂无对话" : "该 Agent 暂无对话"}
</p>
)}
{visibleConversations.map((c) => (
<ConversationItem
key={c.id}
@@ -1220,7 +1280,7 @@ export default function ChatPage() {
onSelectedForDeleteChange={toggleConversationSelected}
/>
))}
{visibleConversationCount < convs.length && (
{visibleConversationCount < filteredConversations.length && (
<Button
type="button"
variant="ghost"
@@ -1252,6 +1312,7 @@ export default function ChatPage() {
agents={chatAgents}
profiles={profiles}
onStarted={(c, pending) => {
if (agentFilter !== null && agentFilter !== c.agent_key) changeAgentFilter(null);
// Insert the new conversation immediately so `selected` resolves to
// it on this render (switching to ChatView right away, before the
// async reloadConvs lands); reloadConvs then reconciles titles etc.
@@ -28,6 +28,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";
@@ -212,6 +213,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>
@@ -1,33 +1,32 @@
"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 { FindingRetestPanel } from "@/components/finding-retest-panel";
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 { FindingRetestPanel } from "@/components/finding-retest-panel";
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",
@@ -46,21 +45,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>
);
}
@@ -125,9 +114,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">
@@ -142,11 +129,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">
@@ -160,16 +143,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>
@@ -190,9 +169,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>
<FindingRetestPanel key={id} findingId={id} readOnly={finding.inherited} onCompleted={load} />
@@ -210,21 +187,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>
@@ -248,14 +233,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">
@@ -276,14 +255,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">
@@ -302,9 +275,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>
)}
@@ -337,9 +308,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";
@@ -65,10 +59,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">
@@ -80,11 +71,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) => (
@@ -96,22 +86,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">
@@ -147,9 +127,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>
@@ -188,25 +166,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)
@@ -244,9 +215,7 @@ export function FindingsTab({ taskId }: { taskId: string }) {
<Row key={f.id} f={f} contextTaskId={taskId} onStatus={onStatus} />
))}
{items.length === 0 && (
<p className="px-4 py-8 text-center text-sm text-muted-foreground">
本任务及直接关联任务暂无确认发现。
</p>
<p className="px-4 py-8 text-center text-sm text-muted-foreground">本任务及直接关联任务暂无确认发现。</p>
)}
</CardContent>
</Card>
+55 -204
View File
@@ -5,19 +5,18 @@ import * as React from "react";
import {
ArrowDownWideNarrowIcon,
ArrowUpNarrowWideIcon,
CheckIcon,
ChevronLeftIcon,
ChevronRightIcon,
CopyIcon,
ListChecksIcon,
Loader2Icon,
RadioTowerIcon,
SearchIcon,
Trash2Icon,
WrapTextIcon,
} from "lucide-react";
import { toast } from "sonner";
import { HttpCodeBlock } from "@/components/http-code-block";
import { LinkTrafficDialog } from "@/components/link-traffic-dialog";
import {
AlertDialog,
AlertDialogAction,
@@ -41,7 +40,7 @@ import { Tabs, TabsContent, TabsList, TabsTrigger } from "@/components/ui/tabs";
import { Tooltip, TooltipContent, TooltipTrigger } from "@/components/ui/tooltip";
import { api } from "@/lib/api";
import type { TrafficDetail, TrafficExchange, TrafficHost, TrafficResp } from "@/lib/types";
import { cn, copyText } from "@/lib/utils";
import { cn } from "@/lib/utils";
function fmtTime(ts: string) {
return new Date(ts).toLocaleString("zh-CN", {
@@ -90,205 +89,6 @@ function requestWithHost(raw: string, exchange: TrafficExchange): string {
return `${raw.slice(0, firstLineEnd + newline.length)}Host: ${host}${newline}${raw.slice(firstLineEnd + newline.length)}`;
}
function StartLine({ line }: { line: string }) {
const response = /^(HTTP\/\S+)(\s+)(\d{3})(.*)$/.exec(line);
if (response) {
return (
<>
<span className="text-muted-foreground">{response[1]}</span>
{response[2]}
<span className={statusTone(Number(response[3]))}>{response[3]}</span>
{response[4]}
</>
);
}
const request = /^([A-Z]+)(\s+)(\S+)(\s+)(HTTP\/\S+)$/.exec(line);
if (request) {
return (
<>
<span className="font-semibold text-primary">{request[1]}</span>
{request[2]}
<span className="text-chart-2">{request[3]}</span>
{request[4]}
<span className="text-muted-foreground">{request[5]}</span>
</>
);
}
return line;
}
function HeaderLine({ line }: { line: string }) {
const separator = line.indexOf(":");
if (separator <= 0) return line;
return (
<>
<span className="text-primary">{line.slice(0, separator)}</span>
<span className="text-muted-foreground">:</span>
<span className="text-chart-2">{line.slice(separator + 1)}</span>
</>
);
}
function JsonBody({ body }: { body: string }) {
const parts: React.ReactNode[] = [];
const tokens = /("(?:\\.|[^"\\])*")(\s*:)?|\b(true|false|null)\b|-?\d+(?:\.\d+)?(?:[eE][+-]?\d+)?/g;
let cursor = 0;
for (const match of body.matchAll(tokens)) {
const index = match.index ?? 0;
if (index > cursor) parts.push(body.slice(cursor, index));
let className = "text-chart-4";
if (match[1]) className = match[2] ? "text-primary" : "text-chart-2";
else if (match[3]) className = "text-chart-3";
parts.push(
<span key={`${index}-${match[0].length}`} className={className}>
{match[0]}
</span>,
);
cursor = index + match[0].length;
}
if (cursor < body.length) parts.push(body.slice(cursor));
return parts;
}
function MarkupBody({ body }: { body: string }) {
const parts: React.ReactNode[] = [];
const tags = /<\/?[A-Za-z][^>]*>|<!--[\s\S]*?-->/g;
let cursor = 0;
for (const match of body.matchAll(tags)) {
const index = match.index ?? 0;
if (index > cursor) parts.push(body.slice(cursor, index));
parts.push(
<span key={`${index}-${match[0].length}`} className="text-primary">
{match[0]}
</span>,
);
cursor = index + match[0].length;
}
if (cursor < body.length) parts.push(body.slice(cursor));
return parts;
}
type BodyFormat = "json" | "markup" | "plain";
function detectBodyFormat(body: string): BodyFormat {
const trimmed = body.trimStart();
if (trimmed.startsWith("{") || trimmed.startsWith("[")) return "json";
if (trimmed.startsWith("<")) return "markup";
return "plain";
}
function HighlightedBody({ body, format }: { body: string; format: BodyFormat }) {
if (format === "json") return <JsonBody body={body} />;
if (format === "markup") return <MarkupBody body={body} />;
return body;
}
function HttpCodeBlock({ raw }: { raw: string }) {
const [wrapLines, setWrapLines] = React.useState(true);
const [copied, setCopied] = React.useState(false);
const value = raw || "(空)";
const lines = value.replaceAll("\r\n", "\n").split("\n");
const separator = lines.indexOf("");
const body = separator >= 0 ? lines.slice(separator + 1).join("\n") : "";
const bodyFormat = detectBodyFormat(body);
React.useEffect(() => {
if (!copied) return;
const timer = window.setTimeout(() => setCopied(false), 1500);
return () => window.clearTimeout(timer);
}, [copied]);
const copyPacket = async () => {
const ok = await copyText(value);
if (ok) {
setCopied(true);
return;
}
toast.error("复制失败,请使用 Ctrl/Cmd+A 后复制");
};
const renderLine = (line: string, index: number) => {
if (index === 0) return <StartLine line={line} />;
if (separator < 0 || index < separator) return <HeaderLine line={line} />;
if (index === separator) return null;
return <HighlightedBody body={line} format={bodyFormat} />;
};
return (
<div className="relative mx-3 my-4 overflow-hidden rounded-md border bg-background shadow-xs">
<div className="absolute top-2 right-2 flex items-center gap-0.5 rounded-md border bg-background/95 p-0.5 shadow-xs">
<Tooltip>
<TooltipTrigger asChild>
<Button
type="button"
variant="ghost"
size="icon-xs"
aria-label={wrapLines ? "关闭自动换行" : "开启自动换行"}
aria-pressed={wrapLines}
onClick={() => setWrapLines((current) => !current)}
>
<WrapTextIcon />
</Button>
</TooltipTrigger>
<TooltipContent side="bottom">{wrapLines ? "关闭自动换行" : "开启自动换行"}</TooltipContent>
</Tooltip>
<Tooltip>
<TooltipTrigger asChild>
<Button
type="button"
variant="ghost"
size="icon-xs"
aria-label={copied ? "已复制报文" : "复制报文"}
onClick={() => void copyPacket()}
>
{copied ? <CheckIcon /> : <CopyIcon />}
</Button>
</TooltipTrigger>
<TooltipContent side="bottom">{copied ? "已复制" : "复制报文"}</TooltipContent>
</Tooltip>
</div>
{/* biome-ignore lint/a11y/useSemanticElements: textarea cannot preserve line numbers and syntax-highlighting markup. */}
<div
role="textbox"
aria-label="HTTP 报文代码"
aria-multiline="true"
aria-readonly="true"
tabIndex={0}
className="max-h-[calc(100vh-15rem)] min-w-0 overflow-auto bg-background py-3 font-mono text-xs outline-none focus-visible:ring-2 focus-visible:ring-inset focus-visible:ring-ring"
onKeyDown={(event) => {
if (!(event.ctrlKey || event.metaKey) || event.key.toLowerCase() !== "a") return;
event.preventDefault();
const selection = window.getSelection();
if (!selection) return;
const range = document.createRange();
range.selectNodeContents(event.currentTarget);
selection.removeAllRanges();
selection.addRange(range);
}}
>
{lines.map((line, index) => (
<div
key={`${index}-${line}`}
data-line={index + 1}
className="grid min-w-0 grid-cols-[1.5rem_minmax(0,1fr)] leading-relaxed before:sticky before:left-0 before:self-stretch before:border-r before:bg-muted/20 before:px-1 before:text-right before:text-muted-foreground/60 before:content-[attr(data-line)]"
>
<code
className={cn(
"min-h-[1lh] min-w-0 pr-16 pl-1.5 [tab-size:4]",
wrapLines ? "break-words whitespace-pre-wrap" : "whitespace-pre",
)}
>
{renderLine(line, index)}
</code>
</div>
))}
</div>
</div>
);
}
// Fixed method set (server filters exact-match); avoids deriving options from a
// single page, which would only ever list the methods on that page.
const METHODS = ["GET", "POST", "PUT", "DELETE", "PATCH", "HEAD", "OPTIONS"];
@@ -296,6 +96,8 @@ const PAGE_SIZES = [25, 50, 100, 200];
type HostCountSortDirection = "asc" | "desc";
export default function TrafficPage() {
const [selectedFlows, setSelectedFlows] = React.useState<Set<string>>(() => new Set());
const [linking, setLinking] = React.useState(false);
const [page, setPage] = React.useState(0);
const [size, setSize] = React.useState(50);
const [host, setHost] = React.useState(""); // raw host input
@@ -623,6 +425,17 @@ export default function TrafficPage() {
</div>
</div>
<div className="flex flex-wrap items-center gap-2">
<span className="text-sm text-muted-foreground">已选 {selectedFlows.size} 条流量</span>
<Button variant="outline" size="sm" disabled={selectedFlows.size === 0} onClick={() => setLinking(true)}>
关联到漏洞
</Button>
{selectedFlows.size > 0 ? (
<Button variant="ghost" size="sm" onClick={() => setSelectedFlows(new Set())}>
清空选择
</Button>
) : null}
</div>
{/* History table */}
<div className="flex h-[calc(100vh-15rem)] min-h-0 flex-col">
<Card className="flex min-h-0 flex-1 flex-col overflow-hidden py-0">
@@ -630,6 +443,22 @@ export default function TrafficPage() {
<Table>
<TableHeader className="sticky top-0 z-10 bg-card">
<TableRow>
<TableHead className="w-10">
<Checkbox
aria-label="选择本页流量"
checked={exchanges.length > 0 && exchanges.every((e) => selectedFlows.has(e.id))}
onCheckedChange={(checked) =>
setSelectedFlows((previous) => {
const next = new Set(previous);
for (const e of exchanges) {
if (checked === true) next.add(e.id);
else next.delete(e.id);
}
return next;
})
}
/>
</TableHead>
<TableHead className="w-36">时间</TableHead>
<TableHead className="w-44">host</TableHead>
<TableHead className="w-20">方法</TableHead>
@@ -646,6 +475,21 @@ export default function TrafficPage() {
className={cn("cursor-pointer", selected?.id === e.id && "bg-accent hover:bg-accent")}
onClick={() => setSelected(e)}
>
<TableCell>
<Checkbox
aria-label={`选择流量 ${e.id}`}
checked={selectedFlows.has(e.id)}
onClick={(event) => event.stopPropagation()}
onCheckedChange={(checked) =>
setSelectedFlows((previous) => {
const next = new Set(previous);
if (checked === true) next.add(e.id);
else next.delete(e.id);
return next;
})
}
/>
</TableCell>
<TableCell className="text-xs text-muted-foreground tabular-nums">{fmtTime(e.ts)}</TableCell>
<TableCell className="font-mono text-xs">{e.host}</TableCell>
<TableCell>
@@ -665,7 +509,7 @@ export default function TrafficPage() {
))}
{exchanges.length === 0 && (
<TableRow>
<TableCell colSpan={7} className="py-12 text-center text-sm text-muted-foreground">
<TableCell colSpan={8} className="py-12 text-center text-sm text-muted-foreground">
{traffic === null ? "加载中…" : "没有匹配的流量。"}
</TableCell>
</TableRow>
@@ -676,6 +520,13 @@ export default function TrafficPage() {
</Card>
</div>
{linking ? (
<LinkTrafficDialog
trafficIds={[...selectedFlows]}
onClose={() => setLinking(false)}
onBound={() => setSelectedFlows(new Set())}
/>
) : null}
<Sheet open={selected !== null} onOpenChange={(open) => !open && setSelected(null)}>
<SheetContent className="w-full! max-w-none! gap-0 p-0 sm:w-[48rem]! sm:max-w-[48rem]!">
{selected && (
+55 -10
View File
@@ -17,6 +17,7 @@ import type { Settings } from "@/lib/types";
export default function SystemSettingsPage() {
const [trafficCapture, setTrafficCapture] = React.useState(false);
const [agentTrafficBinding, setAgentTrafficBinding] = React.useState(false);
const [webSearch, setWebSearch] = React.useState(false);
const [backend, setBackend] = React.useState("ddgs");
const [braveKeySet, setBraveKeySet] = React.useState(false);
@@ -43,13 +44,14 @@ export default function SystemSettingsPage() {
const apply = React.useCallback((s: Settings) => {
setTrafficCapture(!!s.traffic_capture);
setAgentTrafficBinding(!!s.agent_traffic_binding);
setWebSearch(!!s.web_search_enabled);
setBackend(s.web_search_backend || "ddgs");
setBraveKeySet(!!s.brave_key_set);
setTavilyKeySet(!!s.tavily_key_set);
setProxyInput(s.web_search_proxy || "");
setGlobalProxyInput(s.global_proxy || "");
setPyInterp(s.python_interpreter || "");
setProxyInput(s.web_search_proxy ?? "");
setGlobalProxyInput(s.global_proxy ?? "");
setPyInterp(s.python_interpreter ?? "");
setWorkers(String(s.workers ?? 3));
setInjectPlanner(s.constraints_inject_planner !== false);
setInjectWorker(s.constraints_inject_worker !== false);
@@ -118,6 +120,22 @@ export default function SystemSettingsPage() {
.catch(() => setInjectPlanner(!v)); // revert on failure
};
const toggleAgentTrafficBinding = (v: boolean) => {
setAgentTrafficBinding(v);
setSaving(true);
api
.setSettings({ agent_traffic_binding: v })
.then((s) => {
apply(s);
toast.success(v ? "已开启 Agent 自动绑定流量" : "已关闭 Agent 自动绑定流量");
})
.catch((e) => {
setAgentTrafficBinding(!v);
toast.error(`保存失败:${(e as Error).message}`);
})
.finally(() => setSaving(false));
};
const toggleInjectWorker = (v: boolean) => {
setInjectWorker(v); // optimistic
api
@@ -256,6 +274,34 @@ export default function SystemSettingsPage() {
</CardContent>
</Card>
<Card className="mb-4 break-inside-avoid md:mb-6">
<CardHeader>
<CardTitle className="flex items-center gap-2 text-base">
<RadioTowerIcon className="size-4" />
Agent 自动绑定流量
</CardTitle>
<CardDescription id="agent-traffic-binding-description">
默认关闭。开启后,漏洞入库时触发的报告 Agent 会核对已有 HTTP 请求/响应,关联对应流量后再编写报告。
<b>查阅数据包及额外的工具调用会增加 Token 消耗。</b>
<br />
TCP、未抓包或没有匹配流量时仍可正常上报。此开关不影响流量捕获、人工绑定及已保存证据的查看。 对下一轮 Agent
生效;关闭后会立即拒绝新的自动绑定。
</CardDescription>
</CardHeader>
<CardContent className="flex items-center justify-between gap-4">
<Label htmlFor="agent-traffic-binding" className="text-sm font-normal text-muted-foreground">
{agentTrafficBinding ? "已开启 · 会增加 Token 消耗" : "已关闭 · 可继续人工绑定"}
</Label>
<Switch
id="agent-traffic-binding"
aria-describedby="agent-traffic-binding-description"
checked={agentTrafficBinding}
disabled={!loaded || saving}
onCheckedChange={toggleAgentTrafficBinding}
/>
</CardContent>
</Card>
<Card className="mb-4 break-inside-avoid md:mb-6">
<CardHeader>
<CardTitle className="flex items-center gap-2 text-base">
@@ -269,8 +315,9 @@ export default function SystemSettingsPage() {
开启<b>流量捕获</b>时,它作为记录代理的<b>上游</b>(流量仍全量落库,再经此代理出网);关闭捕获时,直接注入
Agent 的 bash / WebFetch 出网。与网络搜索代理、LLM 代理相互独立。
<br />
<b>提示</b>:socks5 在<b>关闭捕获</b>时依赖各命令行工具对 <code>ALL_PROXY</code> 的支持(curl 可用,部分工具可能忽略);
若主要用 socks5,建议开启流量捕获——此路径由 MITM 亲自拨号,工具无感知、稳定生效。
<b>提示</b>:socks5 在<b>关闭捕获</b>时依赖各命令行工具对 <code>ALL_PROXY</code> 的支持(curl
可用,部分工具可能忽略); 若主要用 socks5,建议开启流量捕获——此路径由 MITM
亲自拨号,工具无感知、稳定生效。
</CardDescription>
</CardHeader>
<CardContent className="flex flex-col gap-2">
@@ -291,9 +338,7 @@ export default function SystemSettingsPage() {
</Button>
</div>
<p className="text-muted-foreground text-xs">
{globalProxyInput.trim()
? "已配置 · 所有目标流量经此代理出网"
: "未配置 · 目标流量直连出网"}
{globalProxyInput.trim() ? "已配置 · 所有目标流量经此代理出网" : "未配置 · 目标流量直连出网"}
</p>
</CardContent>
</Card>
@@ -308,8 +353,8 @@ export default function SystemSettingsPage() {
开启后,把每个任务的<b>操作约束</b>(在任务总览「操作约束」里维护的 allow/deny 条目)拼进对应 Agent
的系统提示,用来框定探索边界(如「仅测当前端口」「禁止爆破」)。
<br />
可分别控制注入到 <b>规划者(planner)</b>与 <b>执行者(worker)</b>;默认都开。切换即时生效(下一轮读取),无需重建
Agent。关闭后该 Agent 不再看到约束。
可分别控制注入到 <b>规划者(planner)</b>与 <b>执行者(worker)</b>
;默认都开。切换即时生效(下一轮读取),无需重建 Agent。关闭后该 Agent 不再看到约束。
</CardDescription>
</CardHeader>
<CardContent className="flex flex-col gap-4">
@@ -0,0 +1,285 @@
"use client";
import * as React from "react";
import { ArrowDownIcon, ArrowUpIcon, PlusIcon } from "lucide-react";
import { toast } from "sonner";
import { TrafficEvidenceViewer } from "@/components/traffic-evidence-viewer";
import { TrafficPickerDialog } from "@/components/traffic-picker-dialog";
import { Alert, AlertDescription } from "@/components/ui/alert";
import { Badge } from "@/components/ui/badge";
import { Button } from "@/components/ui/button";
import { Card, CardContent, CardDescription, CardHeader, CardTitle } from "@/components/ui/card";
import {
Dialog,
DialogContent,
DialogDescription,
DialogFooter,
DialogHeader,
DialogTitle,
} from "@/components/ui/dialog";
import { Empty, EmptyDescription, EmptyHeader, EmptyTitle } from "@/components/ui/empty";
import { Field, FieldGroup, FieldLabel } from "@/components/ui/field";
import { Select, SelectContent, SelectGroup, SelectItem, SelectTrigger, SelectValue } from "@/components/ui/select";
import { Skeleton } from "@/components/ui/skeleton";
import { Textarea } from "@/components/ui/textarea";
import { api } from "@/lib/api";
import type { FindingTraffic, FindingTrafficBinding, TrafficEvidenceRole } from "@/lib/types";
const ROLES: Record<TrafficEvidenceRole, string> = {
baseline: "正常对照",
proof: "漏洞证明",
verification: "补充验证",
supporting: "辅助证据",
};
export function FindingTrafficPanel({
findingId,
contextTask,
readOnly,
onChanged,
}: {
findingId: string;
contextTask?: string;
readOnly?: boolean;
onChanged: () => void;
}) {
const [data, setData] = React.useState<FindingTraffic | null>(null);
const [error, setError] = React.useState("");
const [adding, setAdding] = React.useState(false);
const [preview, setPreview] = React.useState<string | null>(null);
const [editing, setEditing] = React.useState<FindingTrafficBinding | null>(null);
const [role, setRole] = React.useState<TrafficEvidenceRole>("supporting");
const [note, setNote] = React.useState("");
const [busy, setBusy] = React.useState(false);
const [reload, setReload] = React.useState(0);
// biome-ignore lint/correctness/useExhaustiveDependencies: retry token intentionally refreshes this request.
React.useEffect(() => {
let active = true;
setData(null);
setError("");
api
.findingTraffic(findingId, contextTask)
.then((value) => {
if (active) setData(value);
})
.catch((e: Error) => {
if (active) setError(e.message);
});
return () => {
active = false;
};
}, [findingId, contextTask, reload]);
async function mutate(action: () => Promise<FindingTraffic>) {
setBusy(true);
try {
setData(await action());
setEditing(null);
onChanged();
toast.success("流量证据已更新");
} catch (e) {
toast.error((e as Error).message);
setReload((n) => n + 1);
} finally {
setBusy(false);
}
}
function move(index: number, delta: number) {
if (!data) return;
const ids = data.bindings.map((b) => b.id);
[ids[index], ids[index + delta]] = [ids[index + delta], ids[index]];
void mutate(() => api.orderFindingTraffic(findingId, ids, data.version, contextTask));
}
return (
<>
<Card>
<CardHeader>
<div className="flex flex-wrap items-center justify-between gap-2">
<CardTitle>关联流量 {data ? `(${data.bindings.length})` : ""}</CardTitle>
{!readOnly ? (
<Button variant="outline" size="sm" disabled={!data || busy} onClick={() => setAdding(true)}>
<PlusIcon data-icon="inline-start" />
绑定流量
</Button>
) : (
<Badge variant="outline">继承证据 · 只读</Badge>
)}
</div>
<CardDescription>按复现顺序组织请求与响应。清理原始流量后,已绑定证据仍然保留。</CardDescription>
</CardHeader>
<CardContent className="flex flex-col gap-3">
{error ? (
<Alert variant="destructive">
<AlertDescription>
{error}
<Button variant="link" onClick={() => setReload((n) => n + 1)}>
重试
</Button>
</AlertDescription>
</Alert>
) : null}
{!data ? (
<Skeleton className="h-28 w-full" />
) : data.bindings.length === 0 ? (
<Empty>
<EmptyHeader>
<EmptyTitle>暂无关联流量</EmptyTitle>
<EmptyDescription>可绑定正常对照、漏洞证明和补充验证请求。</EmptyDescription>
</EmptyHeader>
</Empty>
) : (
data.bindings.map((b, index) => (
<div key={b.id} className="flex min-w-0 flex-col gap-2 rounded-md border p-3">
<div className="flex flex-wrap items-center gap-2">
<Badge variant="outline">
{index + 1}. {ROLES[b.role]}
</Badge>
<Badge variant="secondary">
{b.snapshot.method} · {b.snapshot.status}
</Badge>
<span className="text-xs text-muted-foreground">证据 #{b.id}</span>
</div>
<Button
variant="link"
className="h-auto justify-start p-0 text-left whitespace-normal"
onClick={() => setPreview(b.id)}
>
<span className="break-all font-mono text-xs">{b.snapshot.url}</span>
</Button>
{b.note ? <p className="text-sm whitespace-pre-wrap">{b.note}</p> : null}
<div className="flex flex-wrap items-center justify-between gap-2">
<span className="text-xs text-muted-foreground">
{new Date(b.snapshot.captured_at * 1000).toLocaleString("zh-CN")}
</span>
<div className="flex flex-wrap gap-1">
<Button variant="outline" size="sm" onClick={() => setPreview(b.id)}>
查看报文
</Button>
{!readOnly ? (
<>
<Button
variant="ghost"
size="sm"
disabled={busy}
onClick={() => {
setEditing(b);
setRole(b.role);
setNote(b.note);
}}
>
编辑说明
</Button>
<Button
variant="ghost"
size="icon-sm"
aria-label={`上移证据 ${b.id}`}
disabled={busy || index === 0}
onClick={() => move(index, -1)}
>
<ArrowUpIcon />
</Button>
<Button
variant="ghost"
size="icon-sm"
aria-label={`下移证据 ${b.id}`}
disabled={busy || index === data.bindings.length - 1}
onClick={() => move(index, 1)}
>
<ArrowDownIcon />
</Button>
<Button
variant="ghost"
size="sm"
disabled={busy}
onClick={() =>
void mutate(() => api.removeFindingTraffic(findingId, b.id, data.version, contextTask))
}
>
解除绑定
</Button>
</>
) : null}
</div>
</div>
</div>
))
)}
</CardContent>
</Card>
{adding && data ? (
<TrafficPickerDialog
findingId={findingId}
contextTask={contextTask}
bound={data.bindings}
onClose={() => setAdding(false)}
onBound={() => {
setReload((n) => n + 1);
onChanged();
}}
/>
) : null}
<TrafficEvidenceViewer
findingId={findingId}
bindingId={preview}
contextTask={contextTask}
onClose={() => setPreview(null)}
/>
<Dialog
open={editing !== null}
onOpenChange={(open) => {
if (!open && !busy) setEditing(null);
}}
>
<DialogContent>
<DialogHeader>
<DialogTitle>编辑流量证据</DialogTitle>
<DialogDescription>说明这组请求/响应如何支持漏洞结论。</DialogDescription>
</DialogHeader>
<FieldGroup>
<Field>
<FieldLabel htmlFor="evidence-role">用途</FieldLabel>
<Select value={role} onValueChange={(v) => setRole(v as TrafficEvidenceRole)}>
<SelectTrigger id="evidence-role">
<SelectValue />
</SelectTrigger>
<SelectContent>
<SelectGroup>
{Object.entries(ROLES).map(([value, label]) => (
<SelectItem key={value} value={value}>
{label}
</SelectItem>
))}
</SelectGroup>
</SelectContent>
</Select>
</Field>
<Field>
<FieldLabel htmlFor="evidence-note">证据说明</FieldLabel>
<Textarea id="evidence-note" value={note} onChange={(e) => setNote(e.target.value)} />
</Field>
</FieldGroup>
<DialogFooter>
<Button variant="outline" disabled={busy} onClick={() => setEditing(null)}>
取消
</Button>
<Button
disabled={busy || !editing || !data}
onClick={() => {
if (editing && data)
void mutate(() =>
api.editFindingTraffic(findingId, editing.id, data.version, { role, note }, contextTask),
);
}}
>
保存说明
</Button>
</DialogFooter>
</DialogContent>
</Dialog>
</>
);
}
+216
View File
@@ -0,0 +1,216 @@
"use client";
import * as React from "react";
import { CheckIcon, CopyIcon, WrapTextIcon } from "lucide-react";
import { toast } from "sonner";
import { Button } from "@/components/ui/button";
import { Tooltip, TooltipContent, TooltipTrigger } from "@/components/ui/tooltip";
import { cn, copyText } from "@/lib/utils";
function statusTone(status: number) {
if (status >= 500) return "text-red-500";
if (status >= 400) return "text-amber-500";
if (status >= 300) return "text-blue-500";
if (status >= 200) return "text-emerald-500";
return "text-muted-foreground";
}
function StartLine({ line }: { line: string }) {
const response = /^(HTTP\/\S+)(\s+)(\d{3})(.*)$/.exec(line);
if (response) {
return (
<>
<span className="text-muted-foreground">{response[1]}</span>
{response[2]}
<span className={statusTone(Number(response[3]))}>{response[3]}</span>
{response[4]}
</>
);
}
const request = /^([A-Z]+)(\s+)(\S+)(\s+)(HTTP\/\S+)$/.exec(line);
if (request) {
return (
<>
<span className="font-semibold text-primary">{request[1]}</span>
{request[2]}
<span className="text-chart-2">{request[3]}</span>
{request[4]}
<span className="text-muted-foreground">{request[5]}</span>
</>
);
}
return line;
}
function HeaderLine({ line }: { line: string }) {
const separator = line.indexOf(":");
if (separator <= 0) return line;
return (
<>
<span className="text-primary">{line.slice(0, separator)}</span>
<span className="text-muted-foreground">:</span>
<span className="text-chart-2">{line.slice(separator + 1)}</span>
</>
);
}
function JsonBody({ body }: { body: string }) {
const parts: React.ReactNode[] = [];
const tokens = /("(?:\\.|[^"\\])*")(\s*:)?|\b(true|false|null)\b|-?\d+(?:\.\d+)?(?:[eE][+-]?\d+)?/g;
let cursor = 0;
for (const match of body.matchAll(tokens)) {
const index = match.index ?? 0;
if (index > cursor) parts.push(body.slice(cursor, index));
let className = "text-chart-4";
if (match[1]) className = match[2] ? "text-primary" : "text-chart-2";
else if (match[3]) className = "text-chart-3";
parts.push(
<span key={`${index}-${match[0].length}`} className={className}>
{match[0]}
</span>,
);
cursor = index + match[0].length;
}
if (cursor < body.length) parts.push(body.slice(cursor));
return parts;
}
function MarkupBody({ body }: { body: string }) {
const parts: React.ReactNode[] = [];
const tags = /<\/?[A-Za-z][^>]*>|<!--[\s\S]*?-->/g;
let cursor = 0;
for (const match of body.matchAll(tags)) {
const index = match.index ?? 0;
if (index > cursor) parts.push(body.slice(cursor, index));
parts.push(
<span key={`${index}-${match[0].length}`} className="text-primary">
{match[0]}
</span>,
);
cursor = index + match[0].length;
}
if (cursor < body.length) parts.push(body.slice(cursor));
return parts;
}
type BodyFormat = "json" | "markup" | "plain";
function detectBodyFormat(body: string): BodyFormat {
const trimmed = body.trimStart();
if (trimmed.startsWith("{") || trimmed.startsWith("[")) return "json";
if (trimmed.startsWith("<")) return "markup";
return "plain";
}
function HighlightedBody({ body, format }: { body: string; format: BodyFormat }) {
if (format === "json") return <JsonBody body={body} />;
if (format === "markup") return <MarkupBody body={body} />;
return body;
}
export function HttpCodeBlock({ raw }: { raw: string }) {
const [wrapLines, setWrapLines] = React.useState(true);
const [copied, setCopied] = React.useState(false);
const value = raw || "(空)";
const lines = value.replaceAll("\r\n", "\n").split("\n");
const separator = lines.indexOf("");
const body = separator >= 0 ? lines.slice(separator + 1).join("\n") : "";
const bodyFormat = detectBodyFormat(body);
React.useEffect(() => {
if (!copied) return;
const timer = window.setTimeout(() => setCopied(false), 1500);
return () => window.clearTimeout(timer);
}, [copied]);
const copyPacket = async () => {
const ok = await copyText(value);
if (ok) {
setCopied(true);
return;
}
toast.error("复制失败,请使用 Ctrl/Cmd+A 后复制");
};
const renderLine = (line: string, index: number) => {
if (index === 0) return <StartLine line={line} />;
if (separator < 0 || index < separator) return <HeaderLine line={line} />;
if (index === separator) return null;
return <HighlightedBody body={line} format={bodyFormat} />;
};
return (
<div className="relative mx-3 my-4 overflow-hidden rounded-md border bg-background shadow-xs">
<div className="absolute top-2 right-2 flex items-center gap-0.5 rounded-md border bg-background/95 p-0.5 shadow-xs">
<Tooltip>
<TooltipTrigger asChild>
<Button
type="button"
variant="ghost"
size="icon-xs"
aria-label={wrapLines ? "关闭自动换行" : "开启自动换行"}
aria-pressed={wrapLines}
onClick={() => setWrapLines((current) => !current)}
>
<WrapTextIcon />
</Button>
</TooltipTrigger>
<TooltipContent side="bottom">{wrapLines ? "关闭自动换行" : "开启自动换行"}</TooltipContent>
</Tooltip>
<Tooltip>
<TooltipTrigger asChild>
<Button
type="button"
variant="ghost"
size="icon-xs"
aria-label={copied ? "已复制报文" : "复制报文"}
onClick={() => void copyPacket()}
>
{copied ? <CheckIcon /> : <CopyIcon />}
</Button>
</TooltipTrigger>
<TooltipContent side="bottom">{copied ? "已复制" : "复制报文"}</TooltipContent>
</Tooltip>
</div>
{/* biome-ignore lint/a11y/useSemanticElements: textarea cannot preserve line numbers and syntax-highlighting markup. */}
<div
role="textbox"
aria-label="HTTP 报文代码"
aria-multiline="true"
aria-readonly="true"
tabIndex={0}
className="max-h-[calc(100vh-15rem)] min-w-0 overflow-auto bg-background py-3 font-mono text-xs outline-none focus-visible:ring-2 focus-visible:ring-inset focus-visible:ring-ring"
onKeyDown={(event) => {
if (!(event.ctrlKey || event.metaKey) || event.key.toLowerCase() !== "a") return;
event.preventDefault();
const selection = window.getSelection();
if (!selection) return;
const range = document.createRange();
range.selectNodeContents(event.currentTarget);
selection.removeAllRanges();
selection.addRange(range);
}}
>
{lines.map((line, index) => (
<div
key={`${index}-${line}`}
data-line={index + 1}
className="grid min-w-0 grid-cols-[1.5rem_minmax(0,1fr)] leading-relaxed before:sticky before:left-0 before:self-stretch before:border-r before:bg-muted/20 before:px-1 before:text-right before:text-muted-foreground/60 before:content-[attr(data-line)]"
>
<code
className={cn(
"min-h-[1lh] min-w-0 pr-16 pl-1.5 [tab-size:4]",
wrapLines ? "break-words whitespace-pre-wrap" : "whitespace-pre",
)}
>
{renderLine(line, index)}
</code>
</div>
))}
</div>
</div>
);
}
+163
View File
@@ -0,0 +1,163 @@
"use client";
import * as React from "react";
import { toast } from "sonner";
import { Alert, AlertDescription } from "@/components/ui/alert";
import { Button } from "@/components/ui/button";
import {
Dialog,
DialogContent,
DialogDescription,
DialogFooter,
DialogHeader,
DialogTitle,
} from "@/components/ui/dialog";
import { Field, FieldGroup, FieldLabel } from "@/components/ui/field";
import { Input } from "@/components/ui/input";
import { api } from "@/lib/api";
import type { Finding, FindingsPage } from "@/lib/types";
export function LinkTrafficDialog({
trafficIds,
onClose,
onBound,
}: {
trafficIds: string[];
onClose: () => void;
onBound: () => void;
}) {
const [query, setQuery] = React.useState("");
const [page, setPage] = React.useState(1);
const [data, setData] = React.useState<FindingsPage | null>(null);
const [selected, setSelected] = React.useState<Finding | null>(null);
const [error, setError] = React.useState("");
const [busy, setBusy] = React.useState(false);
const [loading, setLoading] = React.useState(true);
React.useEffect(() => {
let active = true;
setLoading(true);
setError("");
const timer = setTimeout(() => {
api
.findingsPage({ page, pageSize: 20, query })
.then((d) => {
if (active) setData(d);
})
.catch((e: Error) => {
if (active) setError(e.message);
})
.finally(() => {
if (active) setLoading(false);
});
}, 250);
return () => {
active = false;
clearTimeout(timer);
};
}, [query, page]);
async function save() {
if (!selected?.finding_id) return;
setBusy(true);
setError("");
try {
await api.bindFindingTraffic(
selected.finding_id,
trafficIds.map((traffic_id) => ({ traffic_id })),
);
toast.success(`已关联 ${trafficIds.length} 条流量到漏洞 #${selected.finding_id}`);
onBound();
onClose();
} catch (e) {
setError((e as Error).message);
} finally {
setBusy(false);
}
}
return (
<Dialog
open
onOpenChange={(open) => {
if (!open && !busy) onClose();
}}
>
<DialogContent className="max-h-[90vh] overflow-y-auto sm:max-w-2xl">
<DialogHeader>
<DialogTitle>关联到漏洞</DialogTitle>
<DialogDescription>将所选 {trafficIds.length} 条流量保存为已有漏洞的证据。</DialogDescription>
</DialogHeader>
<FieldGroup>
<Field>
<FieldLabel htmlFor="link-finding-query">查找漏洞</FieldLabel>
<Input
id="link-finding-query"
placeholder="名称 / 摘要 / 漏洞类别"
value={query}
onChange={(e) => {
setQuery(e.target.value);
setPage(1);
}}
/>
</Field>
</FieldGroup>
{error ? (
<Alert variant="destructive">
<AlertDescription>{error}</AlertDescription>
</Alert>
) : null}
<div className="flex max-h-[40vh] flex-col gap-2 overflow-auto" aria-busy={loading}>
{(data?.items ?? []).map((f) => (
<Button
key={f.finding_id ?? f.id}
variant={selected?.finding_id === f.finding_id ? "secondary" : "outline"}
className="h-auto justify-start p-3 text-left whitespace-normal"
aria-pressed={selected?.finding_id === f.finding_id}
disabled={busy || !f.finding_id || f.inherited}
onClick={() => setSelected(f)}
>
<span className="flex min-w-0 flex-col gap-1">
<span>
#{f.finding_id ?? f.id} · {f.name || f.vulnclass}
</span>
<span className="line-clamp-2 text-xs text-muted-foreground">{f.summary}</span>
</span>
</Button>
))}
{!data?.items.length ? (
<p className="py-6 text-center text-sm text-muted-foreground">
{loading ? "加载中…" : "没有匹配的漏洞,请先登记漏洞"}
</p>
) : null}
</div>
<div className="flex items-center justify-between gap-2">
<Button variant="outline" size="sm" disabled={loading || page <= 1} onClick={() => setPage((p) => p - 1)}>
上一页
</Button>
<span className="text-xs">
第 {page} 页 · 共 {data?.total ?? 0} 条
</span>
<Button
variant="outline"
size="sm"
disabled={loading || page * 20 >= (data?.total ?? 0)}
onClick={() => setPage((p) => p + 1)}
>
下一页
</Button>
</div>
<p className="text-sm">
{selected ? `已选漏洞:#${selected.finding_id} ${selected.name || selected.vulnclass}` : "请选择一个漏洞"}
</p>
<DialogFooter>
<Button variant="outline" disabled={busy} onClick={onClose}>
取消
</Button>
<Button disabled={busy || !selected} onClick={() => void save()}>
{busy ? "保存中…" : "确认关联"}
</Button>
</DialogFooter>
</DialogContent>
</Dialog>
);
}
@@ -0,0 +1,181 @@
"use client";
import * as React from "react";
import { toast } from "sonner";
import { HttpCodeBlock } from "@/components/http-code-block";
import { Alert, AlertDescription } from "@/components/ui/alert";
import { Button } from "@/components/ui/button";
import { Dialog, DialogContent, DialogDescription, DialogHeader, DialogTitle } from "@/components/ui/dialog";
import { Skeleton } from "@/components/ui/skeleton";
import { Tabs, TabsContent, TabsList, TabsTrigger } from "@/components/ui/tabs";
import { api } from "@/lib/api";
import type { FindingTrafficDetail, TrafficDetail } from "@/lib/types";
export function TrafficEvidenceViewer({
findingId,
bindingId,
contextTask,
onClose,
}: {
findingId: string;
bindingId: string | null;
contextTask?: string;
onClose: () => void;
}) {
const [detail, setDetail] = React.useState<FindingTrafficDetail | null>(null);
const [error, setError] = React.useState("");
const [busy, setBusy] = React.useState(false);
React.useEffect(() => {
if (!bindingId) return;
let active = true;
setDetail(null);
setError("");
api
.findingTrafficDetail(findingId, bindingId, contextTask)
.then((d) => {
if (active) setDetail(d);
})
.catch((e: Error) => {
if (active) setError(e.message);
});
return () => {
active = false;
};
}, [findingId, bindingId, contextTask]);
async function more(side: "request" | "response") {
if (!detail || !bindingId) return;
setBusy(true);
try {
const next = await api.findingTrafficBody(findingId, bindingId, side, detail[side].next_offset, contextTask);
setDetail((current) =>
current?.binding.id === bindingId
? { ...current, [side]: { ...next, content: current[side].content + next.content } }
: current,
);
} catch (e) {
toast.error((e as Error).message);
} finally {
setBusy(false);
}
}
return (
<Dialog
open={bindingId !== null}
onOpenChange={(open) => {
if (!open) onClose();
}}
>
<DialogContent className="max-h-[90vh] overflow-y-auto sm:max-w-4xl">
<DialogHeader>
<DialogTitle>流量证据 #{bindingId}</DialogTitle>
<DialogDescription className="break-all">
{detail?.binding.snapshot.url ?? "查看绑定时保存的请求与响应"}
</DialogDescription>
</DialogHeader>
{error ? (
<Alert variant="destructive">
<AlertDescription>{error}</AlertDescription>
</Alert>
) : detail ? (
<Tabs defaultValue="request">
<TabsList>
<TabsTrigger value="request">请求 Request</TabsTrigger>
<TabsTrigger value="response">响应 Response</TabsTrigger>
</TabsList>
{(["request", "response"] as const).map((side) => (
<TabsContent key={side} value={side}>
<div className="flex flex-wrap items-center justify-between gap-2">
<span className="text-xs text-muted-foreground">
正文 {detail[side].total.toLocaleString()} 字节{detail[side].truncated ? " · 当前为预览" : ""}
</span>
<Button
variant="outline"
size="sm"
onClick={() =>
void api
.downloadFindingTrafficBody(findingId, detail.binding.id, side, contextTask)
.catch((e: Error) => toast.error(e.message))
}
>
下载完整{side === "request" ? "请求" : "响应"}正文
</Button>
</div>
<HttpCodeBlock
raw={`${(side === "request" ? detail.binding.snapshot.req_head : detail.binding.snapshot.resp_head)?.trimEnd() ?? ""}\n\n${detail[side].content}`}
/>
{detail[side].truncated && !detail[side].binary ? (
<Button variant="outline" size="sm" disabled={busy} onClick={() => void more(side)}>
加载更多正文
</Button>
) : null}
</TabsContent>
))}
</Tabs>
) : (
<Skeleton className="h-56 w-full" />
)}
</DialogContent>
</Dialog>
);
}
export function CapturedTrafficViewer({ id, onClose }: { id: string | null; onClose: () => void }) {
const [detail, setDetail] = React.useState<TrafficDetail | null>(null);
const [error, setError] = React.useState("");
React.useEffect(() => {
if (!id) return;
let active = true;
setDetail(null);
setError("");
api
.trafficExchange(id)
.then((d) => {
if (active) setDetail(d);
})
.catch((e: Error) => {
if (active) setError(e.message);
});
return () => {
active = false;
};
}, [id]);
return (
<Dialog
open={id !== null}
onOpenChange={(open) => {
if (!open) onClose();
}}
>
<DialogContent className="max-h-[90vh] overflow-y-auto sm:max-w-4xl">
<DialogHeader>
<DialogTitle>预览流量</DialogTitle>
<DialogDescription>流量 ID:{id}。绑定时会保存完整正文。</DialogDescription>
</DialogHeader>
{error ? (
<Alert variant="destructive">
<AlertDescription>{error}</AlertDescription>
</Alert>
) : detail ? (
<Tabs defaultValue="request">
<TabsList>
<TabsTrigger value="request">请求 Request</TabsTrigger>
<TabsTrigger value="response">响应 Response</TabsTrigger>
</TabsList>
<TabsContent value="request">
<HttpCodeBlock raw={detail.req} />
</TabsContent>
<TabsContent value="response">
<HttpCodeBlock raw={detail.resp} />
</TabsContent>
</Tabs>
) : (
<Skeleton className="h-56 w-full" />
)}
</DialogContent>
</Dialog>
);
}
@@ -0,0 +1,277 @@
"use client";
import * as React from "react";
import { toast } from "sonner";
import { CapturedTrafficViewer } from "@/components/traffic-evidence-viewer";
import { Alert, AlertDescription } from "@/components/ui/alert";
import { Badge } from "@/components/ui/badge";
import { Button } from "@/components/ui/button";
import { Checkbox } from "@/components/ui/checkbox";
import {
Dialog,
DialogContent,
DialogDescription,
DialogFooter,
DialogHeader,
DialogTitle,
} from "@/components/ui/dialog";
import { Field, FieldGroup, FieldLabel } from "@/components/ui/field";
import { Input } from "@/components/ui/input";
import { Select, SelectContent, SelectGroup, SelectItem, SelectTrigger, SelectValue } from "@/components/ui/select";
import { Table, TableBody, TableCell, TableHead, TableHeader, TableRow } from "@/components/ui/table";
import { api } from "@/lib/api";
import type { FindingTrafficBinding, TrafficResp } from "@/lib/types";
const METHODS = ["GET", "POST", "PUT", "DELETE", "PATCH", "HEAD", "OPTIONS"];
export function TrafficPickerDialog({
findingId,
contextTask,
bound,
onClose,
onBound,
}: {
findingId: string;
contextTask?: string;
bound: FindingTrafficBinding[];
onClose: () => void;
onBound: () => void;
}) {
const [host, setHost] = React.useState("");
const [method, setMethod] = React.useState("all");
const [query, setQuery] = React.useState("");
const [page, setPage] = React.useState(0);
const [data, setData] = React.useState<TrafficResp | null>(null);
const [selected, setSelected] = React.useState<Set<string>>(() => new Set());
const [preview, setPreview] = React.useState<string | null>(null);
const [loading, setLoading] = React.useState(true);
const [busy, setBusy] = React.useState(false);
const [error, setError] = React.useState("");
const alreadyBound = new Set(bound.map((b) => b.snapshot.source_traffic_id));
React.useEffect(() => {
let active = true;
setLoading(true);
setError("");
const timer = setTimeout(() => {
api
.traffic(page, 25, host, method === "all" ? "" : method, query)
.then((d) => {
if (active) setData(d);
})
.catch((e: Error) => {
if (active) setError(e.message);
})
.finally(() => {
if (active) setLoading(false);
});
}, 250);
return () => {
active = false;
clearTimeout(timer);
};
}, [page, host, method, query]);
const rows = data?.exchanges ?? [];
const selectable = rows.filter((e) => !alreadyBound.has(e.id));
function toggle(id: string, checked: boolean) {
setSelected((previous) => {
const next = new Set(previous);
if (checked) next.add(id);
else next.delete(id);
return next;
});
}
async function save() {
setBusy(true);
setError("");
try {
await api.bindFindingTraffic(
findingId,
[...selected].map((traffic_id) => ({ traffic_id })),
contextTask,
);
toast.success(`已绑定 ${selected.size} 条流量`);
onBound();
onClose();
} catch (e) {
setError((e as Error).message);
} finally {
setBusy(false);
}
}
return (
<>
<Dialog
open
onOpenChange={(open) => {
if (!open && !busy) onClose();
}}
>
<DialogContent className="max-h-[90vh] overflow-y-auto sm:max-w-5xl">
<DialogHeader>
<DialogTitle>绑定流量</DialogTitle>
<DialogDescription>
筛选并多选请求/响应,已选记录会跨页保留。绑定后可设置用途、说明和顺序。
</DialogDescription>
</DialogHeader>
<FieldGroup className="flex flex-col gap-3 sm:flex-row">
<Field>
<FieldLabel htmlFor="evidence-host">目标 host</FieldLabel>
<Input
id="evidence-host"
value={host}
placeholder="域名或 IP"
onChange={(e) => {
setHost(e.target.value);
setPage(0);
}}
/>
</Field>
<Field>
<FieldLabel htmlFor="evidence-method">请求方法</FieldLabel>
<Select
value={method}
onValueChange={(v) => {
setMethod(v);
setPage(0);
}}
>
<SelectTrigger id="evidence-method">
<SelectValue />
</SelectTrigger>
<SelectContent>
<SelectGroup>
<SelectItem value="all">全部方法</SelectItem>
{METHODS.map((m) => (
<SelectItem key={m} value={m}>
{m}
</SelectItem>
))}
</SelectGroup>
</SelectContent>
</Select>
</Field>
<Field>
<FieldLabel htmlFor="evidence-query">关键词</FieldLabel>
<Input
id="evidence-query"
value={query}
placeholder="URL / 正文关键词"
onChange={(e) => {
setQuery(e.target.value);
setPage(0);
}}
/>
</Field>
</FieldGroup>
{error ? (
<Alert variant="destructive">
<AlertDescription>{error}</AlertDescription>
</Alert>
) : null}
<div className="max-h-[45vh] overflow-auto rounded-md border" aria-busy={loading}>
<Table>
<TableHeader>
<TableRow>
<TableHead>
<Checkbox
aria-label="选择本页未绑定流量"
disabled={loading || busy || selectable.length === 0}
checked={selectable.length > 0 && selectable.every((e) => selected.has(e.id))}
onCheckedChange={(checked) =>
setSelected((previous) => {
const next = new Set(previous);
for (const e of selectable) {
if (checked === true) next.add(e.id);
else next.delete(e.id);
}
return next;
})
}
/>
</TableHead>
<TableHead>方法 / URL</TableHead>
<TableHead>时间</TableHead>
<TableHead>状态码</TableHead>
<TableHead>操作</TableHead>
</TableRow>
</TableHeader>
<TableBody>
{rows.map((e) => (
<TableRow key={e.id}>
<TableCell>
<Checkbox
aria-label={`选择流量 ${e.id}`}
checked={selected.has(e.id) || alreadyBound.has(e.id)}
disabled={busy || loading || alreadyBound.has(e.id)}
onCheckedChange={(checked) => toggle(e.id, checked === true)}
/>
</TableCell>
<TableCell>
<span className="font-mono text-xs">
{e.method} {e.url}
</span>
{alreadyBound.has(e.id) ? <Badge variant="secondary">已绑定</Badge> : null}
</TableCell>
<TableCell className="whitespace-nowrap text-xs">
{new Date(e.ts).toLocaleString("zh-CN")}
</TableCell>
<TableCell>{e.status}</TableCell>
<TableCell>
<Button variant="ghost" size="sm" onClick={() => setPreview(e.id)}>
预览
</Button>
</TableCell>
</TableRow>
))}
{!rows.length ? (
<TableRow>
<TableCell colSpan={5} className="py-8 text-center">
{loading ? "加载中…" : "没有匹配的流量"}
</TableCell>
</TableRow>
) : null}
</TableBody>
</Table>
</div>
<div className="flex flex-wrap items-center justify-between gap-2">
<span className="text-sm">
已选 {selected.size} 条 · 共 {data?.total ?? 0} 条
</span>
<div className="flex items-center gap-2">
<Button
variant="outline"
size="sm"
disabled={loading || page === 0}
onClick={() => setPage((p) => p - 1)}
>
上一页
</Button>
<span className="text-xs">第 {page + 1} 页</span>
<Button
variant="outline"
size="sm"
disabled={loading || (page + 1) * 25 >= (data?.total ?? 0)}
onClick={() => setPage((p) => p + 1)}
>
下一页
</Button>
</div>
</div>
<DialogFooter>
<Button variant="outline" disabled={busy} onClick={onClose}>
取消
</Button>
<Button disabled={busy || selected.size === 0} onClick={() => void save()}>
{busy ? "保存中…" : `绑定 ${selected.size} 条流量`}
</Button>
</DialogFooter>
</DialogContent>
</Dialog>
<CapturedTrafficViewer id={preview} onClose={() => setPreview(null)} />
</>
);
}
+72
View File
@@ -29,6 +29,7 @@ import type {
DeleteTaskOptions,
DeleteTaskResult,
Edge,
EvidenceBodyPreview,
Finding,
FindingAssetTree,
FindingDeepenResponse,
@@ -38,6 +39,8 @@ import type {
FindingStats,
FindingStatus,
FindingsPage,
FindingTraffic,
FindingTrafficDetail,
IntentAsset,
InterceptApprovalRow,
InterceptDetail,
@@ -82,6 +85,8 @@ import type {
Tool,
ToolStat,
TrafficDetail,
TrafficEvidenceRef,
TrafficEvidenceRole,
TrafficHost,
TrafficResp,
UsageStats,
@@ -529,6 +534,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 }),
+1
View File
@@ -3314,6 +3314,7 @@ export const tools: Tool[] = [
// ── Settings ─────────────────────────────────────────────────────────────────
export const settings: Settings = {
traffic_capture: true,
agent_traffic_binding: false,
llm_record: false,
web_search_enabled: true,
web_search_backend: "ddgs",
+56
View File
@@ -463,6 +463,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;
@@ -878,6 +882,7 @@ export interface TrafficHost {
// ---- App settings (runtime toggles) ----
export interface Settings {
traffic_capture: boolean;
agent_traffic_binding: boolean; // Agent 自动绑定流量证据,默认关闭;不影响人工绑定
llm_record: boolean; // LLM 录制开关(默认关);关闭时不记录任何 LLM 调用
// Web search. brave_key_set / tavily_key_set reflect whether a key is stored
// (the values are never returned). On PUT, send the corresponding field to set/clear.
@@ -1299,3 +1304,54 @@ export interface InterceptAudit {
export interface InterceptDetail extends InterceptApprovalRow {
audit: InterceptAudit | null;
}
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;
}