From a3bb2c79592c4d79a41479da1f5d52657bfad708 Mon Sep 17 00:00:00 2001 From: ji ruo <79234113+RuoJi6@users.noreply.github.com> Date: Thu, 10 Sep 2026 14:08:40 +0800 Subject: [PATCH] =?UTF-8?q?feat:=20=E5=A2=9E=E5=8A=A0=20Agent=20=E6=B5=81?= =?UTF-8?q?=E9=87=8F=E7=BB=91=E5=AE=9A=E5=BC=80=E5=85=B3=E5=B9=B6=E8=A1=A5?= =?UTF-8?q?=E9=BD=90=E8=AF=81=E6=8D=AE=E4=BA=A4=E6=8E=A5?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- agent/assembly.go | 10 +- agent/deferred.go | 1 + agent/finding_recorder_test.go | 3 + agent/finding_workflow.go | 130 +++++++++++ agent/finding_workflow_test.go | 68 ++++++ agent/prompt_test.go | 4 +- agent/tools.go | 72 ++++-- agent/worker.go | 3 +- db/exploration.go | 3 + db/finding_traffic.go | 34 +++ docs/漏洞流量证据.md | 10 +- server/assembly.go | 1 + server/finding_traffic.go | 17 +- server/finding_workflow.go | 160 +++++++++++++ server/finding_workflow_test.go | 247 ++++++++++++++++++++ server/manager.go | 13 +- server/orchestration.go | 17 +- server/server.go | 12 +- web/src/app/(main)/system/settings/page.tsx | 65 +++++- web/src/lib/mock/data.ts | 1 + web/src/lib/types.ts | 1 + 21 files changed, 810 insertions(+), 62 deletions(-) create mode 100644 agent/finding_workflow.go create mode 100644 agent/finding_workflow_test.go create mode 100644 server/finding_workflow.go create mode 100644 server/finding_workflow_test.go diff --git a/agent/assembly.go b/agent/assembly.go index 98cf40a..28ce000 100644 --- a/agent/assembly.go +++ b/agent/assembly.go @@ -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(out) return out, def, cleanup } diff --git a/agent/deferred.go b/agent/deferred.go index f718e1e..3789680 100644 --- a/agent/deferred.go +++ b/agent/deferred.go @@ -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 diff --git a/agent/finding_recorder_test.go b/agent/finding_recorder_test.go index 57d8055..db7d90a 100644 --- a/agent/finding_recorder_test.go +++ b/agent/finding_recorder_test.go @@ -51,6 +51,9 @@ func TestReportFindingOptionalTrafficWithoutCapture(t *testing.T) { } 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) diff --git a/agent/finding_workflow.go b/agent/finding_workflow.go new file mode 100644 index 0000000..b82ac28 --- /dev/null +++ b/agent/finding_workflow.go @@ -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(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 + } + 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 = findingTrafficGuidance + "\n已有对应的、经核实的 HTTP 流量时，随本次上报提交 traffic_refs。代为上报时保留交接中的引用，或传 evidence_hint_id 读取本任务指定 hint 的结构化 traffic_refs；不要只复制文字结论。无数据包仍可上报。返回 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 = findingTrafficGuidance + "\n**代上报与交接**：有已核实 HTTP 证据时，报告保存应同时携带 traffic_refs；用 add_hint / add_task_hint 交接时保留结构化引用。收到交接提示后用 evidence_hint_id 明确选择本任务对应 hint，不能混用其他漏洞的证据。若只收到文字，先查已有执行记录或请原执行者交接证据，不要仅为补包重新探测。只有拿到真实 ID 才绑定，TCP 或无包时正常上报。" + if has["traffic_search"] && has["traffic_get"] { + guidance += "\n可用 traffic_search 筛选，再用 traffic_get 逐条核实请求/响应及真实 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["bind_finding_traffic"] { + guidance += "\n若漏洞已登记而遗漏了已有的真实流量，用 bind_finding_traffic 补绑，不重复创建漏洞。" + } + } + 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...)) +} diff --git a/agent/finding_workflow_test.go b/agent/finding_workflow_test.go new file mode 100644 index 0000000..fd2df65 --- /dev/null +++ b/agent/finding_workflow_test.go @@ -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") + } +} diff --git a/agent/prompt_test.go b/agent/prompt_test.go index 7dc5335..822e933 100644 --- a/agent/prompt_test.go +++ b/agent/prompt_test.go @@ -57,8 +57,8 @@ func TestRenderSystemOverrideAndFallback(t *testing.T) { if !strings.Contains(recording, "traffic_search") { t.Fatalf("worker while recording should inject trafficTool: %q", recording) } - if !strings.Contains(recording, "traffic_refs") || !strings.Contains(recording, "TCP") { - t.Fatalf("custom worker prompt lost optional finding evidence guidance: %q", recording) + if strings.Contains(recording, "traffic_refs") { + t.Fatalf("worker bypassed shared optional evidence policy: %q", recording) } if !strings.Contains(recording, "中间产物输出规约") { t.Fatalf("worker missing artifact tail: %q", recording) diff --git a/agent/tools.go b/agent/tools.go index ad47d69..8ae47c7 100644 --- a/agent/tools.go +++ b/agent/tools.go @@ -399,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 @@ -1069,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 { @@ -1077,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 { @@ -1168,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 }) } @@ -1359,10 +1372,11 @@ func (t *ToolSet) goalMet() actool.CoreTool { // --- worker write tools --- func (t *ToolSet) addFinding() actool.CoreTool { - return writeTool("report_finding", "记录确认的漏洞。可选 traffic_refs 绑定多条真实 HTTP 请求/响应：先用 traffic_search / traffic_get 核对 ID；提交的流量必须全部绑定成功。TCP 等非 HTTP 漏洞、未采集或无确切匹配记录时可省略 traffic_refs 或传空数组，仍用 evidence 提供命令输出、日志等可验证证据，建议说明未绑定原因。任务上下文传当前 intent_id。", obj(map[string]any{ + 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": 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) { @@ -1371,14 +1385,21 @@ func (t *ToolSet) addFinding() actool.CoreTool { 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 需要任务上下文（exploration store 未初始化）"), nil + return actool.Errorf("report_finding 需要任务上下文；平台对话请通过 add_task_hint 向对应任务交接漏洞，并在提示中携带已有的 traffic_refs，由任务 Agent 登记。已登记漏洞可用 bind_finding_traffic 补绑。"), nil } - refs, err := db.NormalizeTrafficRefs(a.TrafficRefs) + 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 } @@ -1409,7 +1430,19 @@ func (t *ToolSet) addFinding() actool.CoreTool { recorded.Traffic.Bindings[i].Snapshot.ReqHead = "" recorded.Traffic.Bindings[i].Snapshot.RespHead = "" } - raw, _ := json.Marshal(recorded) + 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 自动绑定流量已关闭，可在页面人工关联流量。" + } + } + raw, _ := json.Marshal(result) return actool.Text(fmt.Sprintf("finding recorded: %d\n%s", recorded.NodeID, raw)), nil }) } @@ -1523,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 不能为空") } @@ -1538,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) } @@ -1724,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 { diff --git a/agent/worker.go b/agent/worker.go index a8781b5..b8d861d 100644 --- a/agent/worker.go +++ b/agent/worker.go @@ -258,7 +258,8 @@ func workerSystem(proxyAddr, caCert, dataDir, runDir string) string { body := renderSystem("worker", workerDefaultTmpl, WorkerVars{ProxyAddr: proxyAddr, DataDir: dataDir, Now: nowStr()}) // caCert is present only when the recording MITM is on, which is exactly when // the traffic_* tools are registered — so it gates the traffic-tool note. - return body + workerTrafficBlock(caCert != "") + findingTrafficGuidance + workerArtifactSpec(runDir) + // Optional finding guidance is added for every role after tool resolution. + return body + workerTrafficBlock(caCert != "") + workerArtifactSpec(runDir) } // renderIntentTask formats the claimed intent for the worker's launch USER message: diff --git a/db/exploration.go b/db/exploration.go index ad85ee5..cd844e4 100644 --- a/db/exploration.go +++ b/db/exploration.go @@ -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"` diff --git a/db/finding_traffic.go b/db/finding_traffic.go index a1ffb63..f898303 100644 --- a/db/finding_traffic.go +++ b/db/finding_traffic.go @@ -381,6 +381,40 @@ func (d *DB) FindingIDByNodeID(nodeID int64) (id int64, err error) { 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 diff --git a/docs/漏洞流量证据.md b/docs/漏洞流量证据.md index bcbf204..7566261 100644 --- a/docs/漏洞流量证据.md +++ b/docs/漏洞流量证据.md @@ -2,7 +2,9 @@ 漏洞详情的「关联流量」支持跨页多选、用途和说明、排序及解除绑定。流量页也可多选记录，一次关联到一个已有漏洞。任务继承的漏洞证据只读，修改需进入来源任务。 -Agent 在发现漏洞、调用 `report_finding` 时可一并绑定多条流量，报告 Agent 随后读取这些快照撰写报告。绑定可选：TCP 等非 HTTP 漏洞、未采集或找不到确切记录时，省略 `traffic_refs` 或传空数组仍可正常上报；在 `evidence` 保留命令输出、日志等其他可验证证据，建议说明未绑定原因，不新增必填字段。提交的 ID 必须全部有效且正文完整。任何一条失败都会回滚本次业务记录。重复追加同一快照不增加绑定，也不覆盖说明。 +系统设置中的「Agent 自动绑定流量」默认关闭，读取和修改接口为 `/api/settings` 的 `agent_traffic_binding`。开启会增加查阅请求/响应、工具调用及提示词带来的 Token 消耗；下一轮 Agent 使用新设置。关闭时隐藏自动绑定参数和补绑工具，不注入自动绑定指引，并拒绝已运行会话新提交的自动绑定。人工绑定、流量捕获、已保存证据读取和导出不受影响。 + +开启后，Agent 在发现漏洞、调用 `report_finding` 时可一并绑定多条流量，报告 Agent 随后读取这些快照撰写报告。绑定可选：TCP 等非 HTTP 漏洞、未采集或找不到确切记录时，省略 `traffic_refs` 或传空数组仍可正常上报；在 `evidence` 保留命令输出、日志等其他可验证证据，建议说明未绑定原因，不新增必填字段。提交的 ID 必须全部有效且正文完整。任何一条失败都会回滚本次业务记录。重复追加同一快照不增加绑定，也不覆盖说明。 ## Agent 编号与报告版本 @@ -27,7 +29,11 @@ Agent 在发现漏洞、调用 `report_finding` 时可一并绑定多条流量 - `update_finding_report` 的 `finding_id` **继续使用探索节点 ID**。新增 `evidence_version` 填实际读取的版本；生成期间证据变化会拒绝旧版本写入，必须重新读取并生成。 - 旧报告调用未传版本时，不宣称覆盖已有流量证据。绑定、说明、用途或排序变化后，已有报告提示待更新。 -启动时仅为旧工具 schema 增加可选属性，并首次给 reporter 绑定只读工具；保留已有提示词、描述、启用状态及其他配置。Worker 运行时附加可选绑定的工具使用指引，不替换用户编辑的提示词。失败的 `report_finding` 不触发报告 Agent。 +开启自动绑定后，`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 增加可选属性，原始默认流量工具绑定扩展至内置上报角色；自定义绑定列表、提示词、描述及启用状态保留。自动绑定指引在最终工具装配后统一加入，覆盖 Worker、Planner、MainAgent、Auto 及获授权的自定义 Agent。平台对话缺少任务上下文时应通过结构化提示交接给任务 Agent，不直接上报。判定任务完成前应先交接已有证据；无包不强制等待。失败的 `report_finding` 不触发报告 Agent。 ## API diff --git a/server/assembly.go b/server/assembly.go index a8b2d41..dbb89bd 100644 --- a/server/assembly.go +++ b/server/assembly.go @@ -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. diff --git a/server/finding_traffic.go b/server/finding_traffic.go index dfc511b..92eea8d 100644 --- a/server/finding_traffic.go +++ b/server/finding_traffic.go @@ -377,21 +377,8 @@ func (s *Server) toolGetFindingTraffic() actool.CoreTool { return actool.Errorf(err.Error()), nil } id, bid := parseProfileID(a.FindingID), parseProfileID(a.BindingID) - if id <= 0 { - return actool.Errorf("finding_id 无效"), nil - } - if ri := agent.RunInfoFrom(ctx); ri.TaskID > 0 { - f, err := s.m.pg.GetFinding(id) - if err != nil || f == nil { - return actool.Errorf("漏洞不存在"), nil - } - task := s.m.ResolveTask(strconv.FormatInt(ri.TaskID, 10)) - if task == nil { - return actool.Errorf("任务不存在"), nil - } - if _, _, allowed := findingProvenanceInTask(task, f.TaskID); !allowed { - return actool.Errorf("当前任务不可读取该漏洞"), nil - } + if 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) diff --git a/server/finding_workflow.go b/server/finding_workflow.go new file mode 100644 index 0000000..458d1b3 --- /dev/null +++ b/server/finding_workflow.go @@ -0,0 +1,160 @@ +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_v1" + 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","planner","mainagent","auto","pentest"]` + 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`, 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 + } + if report, err := s.m.pg.GetTool("report_finding"); err == nil && report != nil { + binding, _ := s.m.pg.GetTool("bind_finding_traffic") + if binding != nil && binding.System { + for _, key := range report.Agents { + if err := s.m.pg.AddAgentToToolBinding(key, []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)) + }) +} diff --git a/server/finding_workflow_test.go b/server/finding_workflow_test.go new file mode 100644 index 0000000..39911ab --- /dev/null +++ b/server/finding_workflow_test.go @@ -0,0 +1,247 @@ +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_v1", "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, "planner") || get.Enabled || len(get.Agents) != 1 || get.Agents[0] != "custom-agent" { + t.Fatal("default/custom reader binding migration incorrect") + } + if err := pg.RemoveAgentFromTool("planner", "traffic_search"); err != nil { + t.Fatal(err) + } + s.seedFindingWorkflowTools() + search, _ = pg.GetTool("traffic_search") + if contains(search.Agents, "planner") { + t.Fatal("one-time migration undid later unbinding") + } +} diff --git a/server/manager.go b/server/manager.go index df46d7e..cfbb3fa 100644 --- a/server/manager.go +++ b/server/manager.go @@ -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. diff --git a/server/orchestration.go b/server/orchestration.go index 50c2a47..b532e9f 100644 --- a/server/orchestration.go +++ b/server/orchestration.go @@ -76,6 +76,7 @@ func (s *Server) orchestrationTools() []actool.CoreTool { s.toolGetTaskNodeDetail(), s.toolUpdateFindingReport(), s.toolGetFindingTraffic(), + s.toolBindFindingTraffic(), } } @@ -334,10 +335,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) @@ -462,6 +464,7 @@ func (s *Server) seedOrchestrationTools() { 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 好了（项目尚无旧库，不做迁移）。 @@ -687,9 +690,9 @@ func (s *Server) seedReporterAgent() { Enabled: true, OnToolCall: true, ToolNames: []string{"report_finding"}, - ToolCallMessage: "上面刚有一个漏洞被 report_finding 登记。请从触发上下文里取出 finding_id" + - "（工具返回 \"finding recorded: \" 里的数字）与任务 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) } diff --git a/server/server.go b/server/server.go index fe2e712..e45d3c1 100644 --- a/server/server.go +++ b/server/server.go @@ -3133,6 +3133,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, @@ -3173,8 +3174,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"` @@ -3271,6 +3273,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()) diff --git a/web/src/app/(main)/system/settings/page.tsx b/web/src/app/(main)/system/settings/page.tsx index 5a94837..a1d41a2 100644 --- a/web/src/app/(main)/system/settings/page.tsx +++ b/web/src/app/(main)/system/settings/page.tsx @@ -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() { + + + + + Agent 自动绑定流量 + + + 默认关闭。开启后，Agent 会核对已有 HTTP 请求/响应，并在上报漏洞或交接任务时关联对应流量。 + 查阅数据包及额外的工具调用会增加 Token 消耗。 +
+ TCP、未抓包或没有匹配流量时仍可正常上报。此开关不影响流量捕获、人工绑定及已保存证据的查看。 对下一轮 Agent + 生效；关闭后会立即拒绝新的自动绑定。 + + + + + + + + @@ -269,8 +315,9 @@ export default function SystemSettingsPage() { 开启流量捕获时，它作为记录代理的上游（流量仍全量落库，再经此代理出网）；关闭捕获时，直接注入 Agent 的 bash / WebFetch 出网。与网络搜索代理、LLM 代理相互独立。
- 提示：socks5 在关闭捕获时依赖各命令行工具对 ALL_PROXY 的支持（curl 可用，部分工具可能忽略）； - 若主要用 socks5，建议开启流量捕获——此路径由 MITM 亲自拨号，工具无感知、稳定生效。 + 提示：socks5 在关闭捕获时依赖各命令行工具对 ALL_PROXY 的支持（curl + 可用，部分工具可能忽略）； 若主要用 socks5，建议开启流量捕获——此路径由 MITM + 亲自拨号，工具无感知、稳定生效。 @@ -291,9 +338,7 @@ export default function SystemSettingsPage() {

- {globalProxyInput.trim() - ? "已配置 · 所有目标流量经此代理出网" - : "未配置 · 目标流量直连出网"} + {globalProxyInput.trim() ? "已配置 · 所有目标流量经此代理出网" : "未配置 · 目标流量直连出网"}

@@ -308,8 +353,8 @@ export default function SystemSettingsPage() { 开启后，把每个任务的操作约束（在任务总览「操作约束」里维护的 allow/deny 条目）拼进对应 Agent 的系统提示，用来框定探索边界（如「仅测当前端口」「禁止爆破」）。
- 可分别控制注入到 规划者（planner）与 执行者（worker）；默认都开。切换即时生效（下一轮读取），无需重建 - Agent。关闭后该 Agent 不再看到约束。 + 可分别控制注入到 规划者（planner）与 执行者（worker） + ；默认都开。切换即时生效（下一轮读取），无需重建 Agent。关闭后该 Agent 不再看到约束。 diff --git a/web/src/lib/mock/data.ts b/web/src/lib/mock/data.ts index be0e186..ff5c667 100644 --- a/web/src/lib/mock/data.ts +++ b/web/src/lib/mock/data.ts @@ -3309,6 +3309,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", diff --git a/web/src/lib/types.ts b/web/src/lib/types.ts index ea757a0..339c2b7 100644 --- a/web/src/lib/types.ts +++ b/web/src/lib/types.ts @@ -858,6 +858,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.