feat: 增加 Agent 流量绑定开关并补齐证据交接

This commit is contained in:
ji ruo committed 2026-09-10 14:08:40 +08:00
1 parent ec4ef35fff
commit a3bb2c7959
21 files changed
+810 -62

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(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
+3
View File
@@ -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)
+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(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...))
}
+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")
}
}
+2 -2
View File
@@ -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)
+59 -13
View File
@@ -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 {
+2 -1
View File
@@ -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:
+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"`
+34
View File
@@ -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
+8 -2
View File
@@ -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
+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.
+2 -15
View File
@@ -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)
+160
View File
@@ -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))
})
}
+247
View File
@@ -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")
}
}
+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.
+10 -7
View File
@@ -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>\" 里的数字)与任务 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)
}
+10 -2
View File
@@ -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())
+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">
+1
View File
@@ -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",
+1
View File
@@ -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.