feat: 增加 Agent 流量绑定开关并补齐证据交接
This commit is contained in:
1 parent
ec4ef35fff
commit
a3bb2c7959
21 files changed
+810
-62
No files matched your search
@@ -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.
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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))
|
||||
})
|
||||
}
|
||||
@@ -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
@@ -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
@@ -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
@@ -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())
|
||||
|
||||
Reference in new issue
Block a user