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

2026-09-10 14:08:40 +08:00
1 parent ec4ef35fff
commit a3bb2c7959
@@ -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))
+		})
+}