Merge PR #102: feat: 支持漏洞关联多条流量证据
冲突解决: - db/schema.sql: 复测表(#103)与流量证据表(#102)均为追加,两者保留 - server/server.go: 启动序列同时保留 seedFindingRetester 与 evidence GC - web/src/lib/types.ts: InterceptAudit/InterceptDetail(#101) 与流量证据类型(#102)并存 - findings/detail/page.tsx: 去掉 #103 带入的重复 import,仅保留 FindingRetestPanel - chat/page.tsx: 保留 #103 的 Spinner import
This commit is contained in:
commit
e8efce605c
54 files changed
+5328
-552
No files matched your search
+6
-4
@@ -12,10 +12,11 @@ import (
|
||||
// UnlockSkill unlocks a named skill's MCPs — hosts call it to rebuild the unlock set
|
||||
// from history on a resumed session (design doc C2).
|
||||
type DeferredInfo struct {
|
||||
Deferred []string // all MCP tool names (schema withheld)
|
||||
GlobalNames []string // MCP names to list in the system-prompt block
|
||||
Unlock *actool.UnlockSet // shared call-gate; nil when no MCP tools
|
||||
UnlockSkill func(skillName string)
|
||||
FindingGuidance string // derived from the final permitted tools, including DB overrides
|
||||
Deferred []string // all MCP tool names (schema withheld)
|
||||
GlobalNames []string // MCP names to list in the system-prompt block
|
||||
Unlock *actool.UnlockSet // shared call-gate; nil when no MCP tools
|
||||
UnlockSkill func(skillName string)
|
||||
}
|
||||
|
||||
// ToolAugment, if set, returns the EXTRA tools an agent should see beyond its
|
||||
@@ -55,5 +56,6 @@ func AugmentTools(ctx context.Context, agentKey string, base []actool.CoreTool)
|
||||
if ToolResolve != nil {
|
||||
out = ToolResolve(ctx, agentKey, out)
|
||||
}
|
||||
out, def.FindingGuidance = findingWorkflowTools(agentKey, out)
|
||||
return out, def, cleanup
|
||||
}
|
||||
@@ -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
|
||||
|
||||
@@ -0,0 +1,22 @@
|
||||
package agent
|
||||
|
||||
import (
|
||||
"context"
|
||||
|
||||
"github.com/Autumn-27/artex/db"
|
||||
)
|
||||
|
||||
// FindingRecorder is injected by the host; agents never synthesize or copy
|
||||
// evidence bodies themselves. Its implementation owns the atomic write.
|
||||
type FindingRecorder interface {
|
||||
Record(context.Context, db.RecordFindingInput, []db.TrafficRef) (*db.RecordedFinding, error)
|
||||
}
|
||||
|
||||
// Tool-use guidance is appended without replacing the user's editable prompt.
|
||||
// It does not require capture or claim that unavailable traffic tools exist.
|
||||
const findingTrafficGuidance = "\n\n**漏洞流量证据(可选)**:调用 report_finding 上报漏洞时,如有已查看并确认支持漏洞结论的 HTTP 请求/响应,可用 traffic_refs 按复现顺序绑定真实 ID;域名和时间只作候选筛选,不推定关联。TCP 等非 HTTP 漏洞、未采集或无确切匹配时省略或传 [],在 evidence 保留命令输出、日志等其他可验证证据,建议说明未绑定原因。不要猜测 ID,也不要仅为补包重复探测。"
|
||||
|
||||
func (t *ToolSet) SetFindingRecorder(r FindingRecorder) { t.findingRecorder = r }
|
||||
func (w *Worker) SetFindingRecorder(r FindingRecorder) { w.findingRecorder = r }
|
||||
func (p *Planner) SetFindingRecorder(r FindingRecorder) { p.findingRecorder = r }
|
||||
func (m *MainAgent) SetFindingRecorder(r FindingRecorder) { m.findingRecorder = r }
|
||||
@@ -0,0 +1,108 @@
|
||||
package agent
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"github.com/Autumn-27/artex/db"
|
||||
"github.com/Autumn-27/artex/evidence"
|
||||
)
|
||||
|
||||
type failingFindingRecorder struct{}
|
||||
|
||||
func (failingFindingRecorder) Record(context.Context, db.RecordFindingInput, []db.TrafficRef) (*db.RecordedFinding, error) {
|
||||
return nil, errors.New("fixture persistence failed")
|
||||
}
|
||||
|
||||
func TestReportFindingOptionalTrafficWithoutCapture(t *testing.T) {
|
||||
d := testDB(t)
|
||||
defer d.Close()
|
||||
task, err := d.CreateTask("TCP evidence", "fixture", nil, 0, 0)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer d.DeleteTask(task.ID)
|
||||
ts := NewToolSet(nil, "")
|
||||
ts.ts, ts.taskID = d.Exploration(task.ExplorationID), task.ID
|
||||
ts.SetFindingRecorder(evidence.New(d, nil, t.TempDir()))
|
||||
notices := 0
|
||||
ts.notifyFinding = func(int64, string) { notices++ }
|
||||
for _, refs := range []string{"", `,"traffic_refs":[]`, `,"traffic_refs":null`} {
|
||||
res, err := ts.addFinding().Call(t.Context(), json.RawMessage(`{"vulnclass":"TCP","severity":"low","summary":"verified TCP fixture","evidence":"command output proves the finding"`+refs+`}`), nil)
|
||||
if err != nil || !strings.HasPrefix(res.Flatten(), "finding recorded: ") {
|
||||
t.Fatalf("optional traffic rejected: %v %s", err, res.Flatten())
|
||||
}
|
||||
var out db.RecordedFinding
|
||||
if err := json.Unmarshal([]byte(strings.SplitN(res.Flatten(), "\n", 2)[1]), &out); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
f, err := d.GetFinding(out.FindingID)
|
||||
if err != nil || f == nil || len(out.Traffic.Bindings) != 0 || f.Evidence != "command output proves the finding" {
|
||||
t.Fatalf("lost non-HTTP evidence: %+v %v", f, err)
|
||||
}
|
||||
}
|
||||
if notices != 3 {
|
||||
t.Fatalf("successful reports notified %d times", notices)
|
||||
}
|
||||
}
|
||||
|
||||
func TestReportFindingAtomicContract(t *testing.T) {
|
||||
old := FindingTrafficBindingEnabled
|
||||
FindingTrafficBindingEnabled = func() bool { return true }
|
||||
t.Cleanup(func() { FindingTrafficBindingEnabled = old })
|
||||
d := testDB(t)
|
||||
defer d.Close()
|
||||
task, err := d.CreateTask("report contract", "fixture", nil, 0, 0)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer d.DeleteTask(task.ID)
|
||||
ts := NewToolSet(nil, "")
|
||||
ts.ts = d.Exploration(task.ExplorationID)
|
||||
ts.taskID = task.ID
|
||||
ts.worker = "fixture"
|
||||
notices := 0
|
||||
ts.notifyFinding = func(int64, string) { notices++ }
|
||||
call := func(body string) string {
|
||||
t.Helper()
|
||||
res, err := ts.addFinding().Call(t.Context(), json.RawMessage(body), nil)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
return res.Flatten()
|
||||
}
|
||||
ts.SetFindingRecorder(failingFindingRecorder{})
|
||||
if got := call(`{"vulnclass":"TEST","summary":"fail","severity":"low","traffic_refs":[{"traffic_id":"x"}]}`); !strings.Contains(got, "persistence failed") {
|
||||
t.Fatal(got)
|
||||
}
|
||||
if notices != 0 || ts.writes.Findings != 0 {
|
||||
t.Fatal("notified before commit")
|
||||
}
|
||||
ts.SetFindingRecorder(nil)
|
||||
got := call(`{"vulnclass":"TEST","summary":"legacy report","severity":"low"}`)
|
||||
lines := strings.SplitN(got, "\n", 2)
|
||||
if len(lines) != 2 {
|
||||
t.Fatal(got)
|
||||
}
|
||||
var out db.RecordedFinding
|
||||
if err = json.Unmarshal([]byte(lines[1]), &out); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if lines[0] != fmt.Sprintf("finding recorded: %d", out.NodeID) || out.FindingID <= 0 || notices != 1 {
|
||||
t.Fatal(got)
|
||||
}
|
||||
f, err := d.GetFinding(out.FindingID)
|
||||
if err != nil || f == nil || f.NodeID == nil || *f.NodeID != out.NodeID {
|
||||
t.Fatalf("wrong finding/node mapping: %+v %v", f, err)
|
||||
}
|
||||
if got = call(`{"vulnclass":"TEST","summary":"no storage","severity":"low","traffic_refs":[{"traffic_id":"x"}]}`); !strings.Contains(got, "未登记") {
|
||||
t.Fatal(got)
|
||||
}
|
||||
if notices != 1 {
|
||||
t.Fatal("failed tool triggered reporter")
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,130 @@
|
||||
package agent
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
|
||||
"github.com/Autumn-27/artex/db"
|
||||
actool "github.com/Autumn-27/norma/tool"
|
||||
)
|
||||
|
||||
const findingIDGuidance = "\n\n**漏洞编号约定**:finding_id 是独立漏洞记录 ID;finding_node_id 是探索节点 ID。list_findings / list_task_findings / node_detail / get_task_node_detail 的 id 保留为探索节点 ID,应从同一返回的 finding_id 读取独立编号。get_finding_traffic / bind_finding_traffic 用独立 finding_id。旧 update_finding_report 的 finding_id 参数仍传 finding_node_id。不要把 report_finding 第一行的数字用于证据工具,也不要遇到编号错误后猜测其他数字。"
|
||||
|
||||
// The server supplies the persisted setting. A missing setting/host is off.
|
||||
// Consulted at assembly and again on writes so an already-running session
|
||||
// cannot keep binding after the user switches the feature off.
|
||||
var FindingTrafficBindingEnabled func() bool
|
||||
|
||||
func findingTrafficBindingEnabled() bool {
|
||||
return FindingTrafficBindingEnabled != nil && FindingTrafficBindingEnabled()
|
||||
}
|
||||
|
||||
const findingTrafficDisabled = "Agent 自动绑定流量已关闭;请在系统设置开启后重试,或使用页面人工绑定。未登记本次携带流量的漏洞。无流量上报可省略 traffic_refs / evidence_hint_id。"
|
||||
|
||||
// Applied after ToolResolve: user descriptions and prompts remain intact, while
|
||||
// all actual reporters (including Planner and custom chat agents) see the same
|
||||
// API contract. Disabled/unbound tools are never reintroduced here.
|
||||
func findingWorkflowTools(agentKey string, tools []actool.CoreTool) ([]actool.CoreTool, string) {
|
||||
if !findingTrafficBindingEnabled() {
|
||||
out := make([]actool.CoreTool, 0, len(tools))
|
||||
for _, tool := range tools {
|
||||
if tool.Name() == "bind_finding_traffic" {
|
||||
continue
|
||||
}
|
||||
if agentKey == "reporter" && (tool.Name() == "traffic_search" || tool.Name() == "traffic_get" || tool.Name() == "traffic_blob") {
|
||||
continue
|
||||
}
|
||||
switch tool.Name() {
|
||||
case "report_finding", "add_hint", "add_task_hint":
|
||||
// Work on a copy: toggling back on must restore the original schema.
|
||||
raw, _ := json.Marshal(tool.InputSchema())
|
||||
var schema map[string]any
|
||||
if json.Unmarshal(raw, &schema) == nil {
|
||||
stripTrafficParameters(schema)
|
||||
tool = DecorateTool(tool, tool.Description(), schema)
|
||||
}
|
||||
}
|
||||
out = append(out, tool)
|
||||
}
|
||||
return out, ""
|
||||
}
|
||||
out := append([]actool.CoreTool(nil), tools...)
|
||||
has := map[string]bool{}
|
||||
for i, tool := range out {
|
||||
has[tool.Name()] = true
|
||||
note := ""
|
||||
switch tool.Name() {
|
||||
case "report_finding":
|
||||
note = "\n默认由报告 Agent 在编写报告前核对并绑定流量。上报者在 evidence 中保留验证命令、关键输出、已有的真实流量 ID 及其用途,供报告 Agent 对照执行记录核实;无需为绑定额外查包。兼容显式即时绑定:traffic_refs 或 evidence_hint_id 可提交已核实的引用,后者读取本任务指定 hint 的结构化引用;任一无效则本次上报全部失败。TCP/无包不需要这些可选参数。返回 finding_id 与 finding_node_id 分别表示独立记录和探索节点。"
|
||||
case "add_hint", "add_task_hint":
|
||||
note = "\n交接已确认漏洞时,在对应提示的 traffic_refs 中保留已核实流量的 ID、用途、说明和顺序(单条放顶层,批量放对应 hints 元素),并在 text 中说明它证明的具体漏洞。调用方不能只交接文字而丢弃已有流量引用。未核实的候选不能作为证据传递。"
|
||||
case "get_finding_traffic", "bind_finding_traffic", "list_findings", "list_task_findings", "node_detail", "get_task_node_detail", "update_finding_report":
|
||||
note = findingIDGuidance
|
||||
}
|
||||
if note != "" {
|
||||
out[i] = DecorateTool(tool, tool.Description()+note, tool.InputSchema())
|
||||
}
|
||||
}
|
||||
guidance := ""
|
||||
if has["report_finding"] || has["add_task_hint"] || has["add_hint"] {
|
||||
guidance = "\n\n**流量证据交接(可选)**:自动绑定默认由报告 Agent 在漏洞入库后、编写报告前完成。上报者应在 evidence 保留验证命令、关键输出、已有真实流量 ID 及其用途,任务中带 intent_id,便于报告 Agent 追溯;不必为了绑定额外查包。Auto / Planner 代为上报时不要丢弃执行者已有的引用。add_hint / add_task_hint 可用 traffic_refs 交接;显式即时绑定仍兼容 report_finding 的 traffic_refs / evidence_hint_id。TCP 或无包时正常登记,不能猜测 ID,也不能仅为补包重复探测。"
|
||||
if has["add_task_hint"] && !has["add_hint"] {
|
||||
guidance += "\n平台对话没有任务上下文时,不直接调用 report_finding;通过 add_task_hint 向已有对应任务交接,由任务 Agent 登记,并用 list_task_findings 核对结果。"
|
||||
}
|
||||
if has["prove_goal"] || has["goal_met"] {
|
||||
guidance += "\n判定目标完成前,先完成本次已有证据的上报/交接。不要在证据交接尚未完成时仅因文字漏洞已登记就结束任务、取消 Worker;无包不要求等待或强行抓包。"
|
||||
}
|
||||
}
|
||||
if has["update_finding_report"] && has["bind_finding_traffic"] && has["get_finding_traffic"] {
|
||||
guidance += "\n\n**报告前自动关联流量(已开启)**:你负责为本次触发的漏洞核对并绑定流量,再撰写报告。先从 report_finding 返回 JSON 或 get_task_node_detail / list_task_findings 取得明确的 finding_id 与 finding_node_id。读取漏洞详情、对应意图的执行记录及已有证据清单,优先使用上报者交接的真实 ID。若本次验证为 HTTP 且流量工具可用,用 traffic_search 筛选候选,再用 traffic_get 逐条核实请求/响应确实支持该漏洞;域名和时间只用于筛选,不证明归属。将确认的证据按复现顺序用 bind_finding_traffic(finding_id, traffic_refs) 关联,选择 baseline / proof / verification / supporting 并说明用途。只能操作本次漏洞,不重复创建漏洞或重新探测目标。绑定成功后重新调用 get_finding_traffic 获取最新 version,读取所需正文,再将实际读取的 version 作为 evidence_version 传给 update_finding_report(其 finding_id 参数仍用 finding_node_id)。已有绑定不必重复追加。TCP、未采集、工具不可用或没有确切匹配时,跳过自动绑定,依据文字/命令证据正常写报告并说明原因,不得为凑齐流量而猜测。绑定失败不宣称成功;保留已有证据并在报告说明未绑定原因。"
|
||||
}
|
||||
if guidance != "" || has["get_finding_traffic"] || has["update_finding_report"] {
|
||||
guidance += findingIDGuidance
|
||||
}
|
||||
return out, guidance
|
||||
}
|
||||
|
||||
func stripTrafficParameters(schema map[string]any) {
|
||||
props, _ := schema["properties"].(map[string]any)
|
||||
delete(props, "traffic_refs")
|
||||
delete(props, "evidence_hint_id")
|
||||
if required, ok := schema["required"].([]any); ok {
|
||||
kept := required[:0]
|
||||
for _, key := range required {
|
||||
if key != "traffic_refs" && key != "evidence_hint_id" {
|
||||
kept = append(kept, key)
|
||||
}
|
||||
}
|
||||
schema["required"] = kept
|
||||
}
|
||||
if hints, ok := props["hints"].(map[string]any); ok {
|
||||
if items, ok := hints["items"].(map[string]any); ok {
|
||||
stripTrafficParameters(items)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// HintTrafficSchema is shared by the task-local and cross-task hint tools.
|
||||
func HintTrafficSchema() map[string]any {
|
||||
return map[string]any{"type": "array", "description": "可选:已核实且对应本提示中具体漏洞的流量引用,保留顺序;交接后 report_finding 可传 evidence_hint_id 携带这些引用。", "items": obj(map[string]any{"traffic_id": str("真实流量 ID"), "role": str("baseline / proof / verification / supporting"), "note": str("该流量支持什么结论")}, "traffic_id")}
|
||||
}
|
||||
|
||||
func (t *ToolSet) findingRefsFromHint(hintID int64, explicit []db.TrafficRef) ([]db.TrafficRef, error) {
|
||||
if hintID <= 0 {
|
||||
return db.NormalizeTrafficRefs(explicit)
|
||||
}
|
||||
n, err := t.ts.GetNode(hintID) // local store only: inherited hints cannot supply evidence
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if n == nil || n.Kind != db.KindHint {
|
||||
return nil, fmt.Errorf("evidence_hint_id=%d 必须是本任务的提示节点(继承提示不可直接用于绑定)", hintID)
|
||||
}
|
||||
var payload struct {
|
||||
Refs []db.TrafficRef `json:"traffic_refs"`
|
||||
}
|
||||
if err := json.Unmarshal(n.Payload, &payload); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return db.NormalizeTrafficRefs(append(append([]db.TrafficRef{}, explicit...), payload.Refs...))
|
||||
}
|
||||
@@ -0,0 +1,68 @@
|
||||
package agent
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
actool "github.com/Autumn-27/norma/tool"
|
||||
)
|
||||
|
||||
func TestFindingWorkflowSharedAssemblyAndSwitch(t *testing.T) {
|
||||
oldPolicy, oldResolve, oldAugment := FindingTrafficBindingEnabled, ToolResolve, ToolAugment
|
||||
t.Cleanup(func() { FindingTrafficBindingEnabled, ToolResolve, ToolAugment = oldPolicy, oldResolve, oldAugment })
|
||||
ToolAugment = nil
|
||||
on := false
|
||||
FindingTrafficBindingEnabled = func() bool { return on }
|
||||
ts := NewToolSet(nil, "fixture")
|
||||
ToolResolve = func(_ context.Context, _ string, base []actool.CoreTool) []actool.CoreTool {
|
||||
out := make([]actool.CoreTool, len(base))
|
||||
for i, tool := range base {
|
||||
out[i] = DecorateTool(tool, "CUSTOM DESCRIPTION", tool.InputSchema())
|
||||
}
|
||||
return out
|
||||
}
|
||||
for _, role := range []string{"worker", "planner", "mainagent", "auto", "pentest", "custom-agent"} {
|
||||
for _, enabled := range []bool{false, true, false} {
|
||||
on = enabled
|
||||
out, def, cleanup := AugmentTools(t.Context(), role, []actool.CoreTool{ts.addFinding(), ts.addHint()})
|
||||
cleanup()
|
||||
system, _ := deferredSystem("USER CUSTOM PROMPT", def)
|
||||
if !strings.HasPrefix(system[0], "USER CUSTOM PROMPT") {
|
||||
t.Fatal("custom prompt replaced")
|
||||
}
|
||||
if strings.Contains(system[0], "traffic_refs") != on {
|
||||
t.Fatalf("%s: guidance ignored switch: %v", role, on)
|
||||
}
|
||||
if on && (!strings.Contains(system[0], "TCP") || !strings.Contains(system[0], "evidence_hint_id")) {
|
||||
t.Fatal("missing optional/handoff contract")
|
||||
}
|
||||
for _, tool := range out {
|
||||
if !strings.HasPrefix(tool.Description(), "CUSTOM DESCRIPTION") {
|
||||
t.Fatal("custom description replaced")
|
||||
}
|
||||
raw, _ := json.Marshal(tool.InputSchema())
|
||||
var schema map[string]any
|
||||
json.Unmarshal(raw, &schema)
|
||||
props := schema["properties"].(map[string]any)
|
||||
if (props["traffic_refs"] != nil) != on {
|
||||
t.Fatalf("%s: binding schema ignored switch", role)
|
||||
}
|
||||
if tool.Name() == "add_hint" {
|
||||
nested := props["hints"].(map[string]any)["items"].(map[string]any)["properties"].(map[string]any)
|
||||
if (nested["traffic_refs"] != nil) != on {
|
||||
t.Fatal("nested hint schema ignored switch")
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
on = true
|
||||
ToolResolve = func(context.Context, string, []actool.CoreTool) []actool.CoreTool { return nil }
|
||||
out, def, cleanup := AugmentTools(t.Context(), "planner", []actool.CoreTool{ts.addFinding()})
|
||||
defer cleanup()
|
||||
if len(out) != 0 || def.FindingGuidance != "" {
|
||||
t.Fatal("reintroduced disabled tool or its guidance")
|
||||
}
|
||||
}
|
||||
+16
-14
@@ -17,19 +17,20 @@ import (
|
||||
// (→planner) or direct high-priority intents (→frontier). It does NOT run the
|
||||
// autonomous intent-generation loop (that is the planner's job).
|
||||
type MainAgent struct {
|
||||
prov llm.Provider
|
||||
model string
|
||||
tx *transcript.Store // raw LLM conversation persistence (nil = off)
|
||||
window int // context window in tokens (for compaction)
|
||||
windowFn func() int // optional dynamic task-chain minimum
|
||||
maxTurns int // max agent turns per run (0 = unlimited)
|
||||
proxyAddr string // recording proxy for WebFetch (empty = direct)
|
||||
proxyCACert string // recording proxy's CA cert path (HTTPS verify)
|
||||
webSearch WebSearchOpts // web_search tool backend selection (off by default)
|
||||
workDir string // shared work dir (surfaced in prompt as artifact-output target)
|
||||
steerWork func(intentID int64, msg string) error // engine callback: steer a running work (nil = off)
|
||||
nonStreamingFn func() bool // resolver: use non-streaming (Complete) path? (nil = streaming)
|
||||
maxTokensFn func() int // resolver: per-reply output cap (nil/0 = send no cap)
|
||||
findingRecorder FindingRecorder
|
||||
prov llm.Provider
|
||||
model string
|
||||
tx *transcript.Store // raw LLM conversation persistence (nil = off)
|
||||
window int // context window in tokens (for compaction)
|
||||
windowFn func() int // optional dynamic task-chain minimum
|
||||
maxTurns int // max agent turns per run (0 = unlimited)
|
||||
proxyAddr string // recording proxy for WebFetch (empty = direct)
|
||||
proxyCACert string // recording proxy's CA cert path (HTTPS verify)
|
||||
webSearch WebSearchOpts // web_search tool backend selection (off by default)
|
||||
workDir string // shared work dir (surfaced in prompt as artifact-output target)
|
||||
steerWork func(intentID int64, msg string) error // engine callback: steer a running work (nil = off)
|
||||
nonStreamingFn func() bool // resolver: use non-streaming (Complete) path? (nil = streaming)
|
||||
maxTokensFn func() int // resolver: per-reply output cap (nil/0 = send no cap)
|
||||
}
|
||||
|
||||
// SetNonStreaming wires a resolver deciding whether runs use the non-streaming
|
||||
@@ -103,6 +104,7 @@ func mainAgentSystem(goal, dataDir, workDir string) string {
|
||||
// worker/planner sessions — not just the final answer.
|
||||
func (m *MainAgent) Chat(ctx context.Context, taskID int64, as *db.AssetStore, ts *db.ExplorationStore, goal, message string, emit func(db.Activity), notify, resume func(), notifyGoal func([]string)) (string, error) {
|
||||
tsx := NewToolSet(ts, "human")
|
||||
tsx.SetFindingRecorder(m.findingRecorder)
|
||||
if as != nil {
|
||||
tsx.SetAssetStore(as, as.Companies())
|
||||
}
|
||||
@@ -149,7 +151,7 @@ func (m *MainAgent) Chat(ctx context.Context, taskID int64, as *db.AssetStore, t
|
||||
// 命中预算(步数)→ SDK 跑收尾:向用户输出一句进展总结。Prompt 与收尾轮数可后台编辑(默认 10 轮)。
|
||||
Settlement: wrapupSettlement("mainagent", nil),
|
||||
NonStreaming: m.nonStreaming(), // 该 profile 选非流式时走 Provider.Complete
|
||||
MaxTokens: m.maxTokens(), // 0 = 不发上限,由服务端默认值决定
|
||||
MaxTokens: m.maxTokens(), // 0 = 不发上限,由服务端默认值决定
|
||||
}
|
||||
if m.tx != nil { // persist raw human↔AI conversation; one accumulating file per task
|
||||
opts.Transcript = m.tx
|
||||
|
||||
@@ -20,6 +20,7 @@ import (
|
||||
// assets, judges whether the task goal is met, and emits 0..N exploration intents
|
||||
// into the frontier. It is the sole intent generator.
|
||||
type Planner struct {
|
||||
findingRecorder FindingRecorder
|
||||
prov llm.Provider
|
||||
model string
|
||||
tx *transcript.Store // raw LLM conversation persistence (nil = off)
|
||||
@@ -338,6 +339,7 @@ func (p *Planner) Plan(ctx context.Context, taskID int64, as *db.AssetStore, ts
|
||||
// runs in a detached goroutine so it never adds latency to this round.
|
||||
p.compactor.OnPlannerRound(ctx, ts)
|
||||
tsx := NewToolSet(ts, "planner")
|
||||
tsx.SetFindingRecorder(p.findingRecorder)
|
||||
if as != nil {
|
||||
tsx.SetAssetStore(as, as.Companies())
|
||||
}
|
||||
|
||||
@@ -57,6 +57,9 @@ func TestRenderSystemOverrideAndFallback(t *testing.T) {
|
||||
if !strings.Contains(recording, "traffic_search") {
|
||||
t.Fatalf("worker while recording should inject trafficTool: %q", recording)
|
||||
}
|
||||
if strings.Contains(recording, "traffic_refs") {
|
||||
t.Fatalf("worker bypassed shared optional evidence policy: %q", recording)
|
||||
}
|
||||
if !strings.Contains(recording, "中间产物输出规约") {
|
||||
t.Fatalf("worker missing artifact tail: %q", recording)
|
||||
}
|
||||
|
||||
+10
-7
@@ -52,7 +52,9 @@ const pentestDefaultTmpl = `你是一个授权渗透测试系统的"独立渗透
|
||||
- 每得出一个结果**立刻**落地,别攒到最后(会话步数耗尽就全丢;记下来的才算数,活在脑子里的不算)。这些记录也是你抗 compaction 的长期记忆。
|
||||
- **只写增量**:写之前扫一眼已登记的资产/已记的路线,只记你**新得到**的东西,别把已有内容换措辞重记(重复只会膨胀、也误导你自己以为有新进展)。只是印证已有结论而无新增,就不必再记。
|
||||
- **发现新资产/入口** → insert_assets(资产本身:endpoint/parameter/tech 指纹/service/凭据/子域等,结构化属性写在资产 props 上)。回看已登记资产用 list_assets,避免重复登记。
|
||||
- **确认漏洞** → report_finding(含可复现 PoC)。**只有你在本次运行里真实触发过、拿到可复现证据(请求/响应或命令输出)才用它**;回看已报漏洞用 list_findings。严禁把仅凭版本/CVE 匹配、"看起来可注入"、外部漏洞库/更新日志/代码 diff 推断的东西当已确认漏洞上报。**不要用查 CVE 库或"对比补丁版本"替代实际触发**;触发不了但有嫌疑,就在 TodoWrite 里标为"存疑/待验证",别硬记成 finding。
|
||||
- **确认漏洞** → report_finding(含可复现 PoC)。**只有你在本次运行里真实触发过、拿到可复现证据(请求/响应或命令输出)才用它**;回看已报漏洞用 list_findings。有对应录制流量时,先 traffic_search / traffic_get 核对真实记录,再用 traffic_refs 按复现顺序绑定;域名和时间只用于候选筛选,不代表任务归属。严禁把仅凭版本/CVE 匹配、"看起来可注入"、外部漏洞库/更新日志/代码 diff 推断的东西当已确认漏洞上报。**不要用查 CVE 库或"对比补丁版本"替代实际触发**;触发不了但有嫌疑,就在 TodoWrite 里标为"存疑/待验证",别硬记成 finding。
|
||||
|
||||
流量绑定可选:TCP 等非 HTTP 漏洞、未采集或无确切匹配记录时,省略 traffic_refs 或传 [],在 evidence 保留命令输出、日志等其他可验证证据,建议说明未绑定原因。不要猜测 ID,也不要仅为补包重复探测。
|
||||
|
||||
━━ 判定与收尾 ━━
|
||||
- 随时对照任务目标:已被你**验证过**的成果满足了目标,就据此判定达成并说明依据。判"达成"的前提是心法 4 的自检已通过——没独立复现过的战果不算达成依据。
|
||||
@@ -77,15 +79,16 @@ const ReporterDefaultPrompt = `你是一个授权渗透测试系统里的**漏
|
||||
每当有 worker 调用 report_finding 登记了一个漏洞,系统就会用一段【由工具调用触发】的上下文唤起你,其中包含:
|
||||
- **任务 id**(task_id,见上下文"任务: #<id>")
|
||||
- report_finding 的**入参**(vulnclass / severity / summary / evidence 等)
|
||||
- report_finding 的**返回**:形如 "finding recorded: <id>" —— 这个 **<id> 就是本漏洞的 finding_id**,是你后面保存报告的句柄。
|
||||
- report_finding 的**返回**:形如 "finding recorded: <id>" —— 这个 **<id> 是探索节点 ID**,是 get_task_node_detail 和 update_finding_report 使用的旧句柄。返回 JSON 中的 finding_id 则是独立漏洞记录 ID,get_finding_traffic 使用它。
|
||||
|
||||
先从上下文里**准确抽取 finding_id 和 task_id**。抽取不到 finding_id 就不要瞎写,说明情况即可。
|
||||
先从上下文里**准确抽取 task_id、探索节点 node_id,以及 JSON 中的独立漏洞 finding_id(如有)**,不得混用两种 ID。抽取不到 node_id 就不要瞎写,说明情况即可。
|
||||
|
||||
━━ 工作步骤 ━━
|
||||
1. **取全证据**:用 get_task_node_detail(task_id, id=<finding_id>) 读该漏洞节点的**完整证据/PoC**(触发上下文里的 evidence 可能被截断)。
|
||||
2. **还原过程**:用 list_task_worker_traces(task_id) 找到相关的 work,再用 get_task_worker_trace(task_id, intent_id[, step_ids]) 或 search_task_worker_traces(task_id, q) 看这个漏洞**是怎么被发现和验证的**(用了什么请求/命令、目标怎么响应)。必要时 get_task_graph(task_id) 看整体态势、list_task_findings(task_id) 看是否有关联漏洞。
|
||||
3. **写报告**:综合以上,写一份结构化 Markdown 报告(见下方模板)。
|
||||
4. **保存**:调用 **update_finding_report(finding_id=<那个 id>, report=<Markdown 全文>)** 保存。这是你的最终产物——不写进去等于没做。
|
||||
1. **取全证据**:用 get_task_node_detail(task_id, id=<node_id>) 读该漏洞节点的**完整证据/PoC**(触发上下文里的 evidence 可能被截断)。
|
||||
2. **流量证据**:如返回 JSON 包含独立 finding_id,用 get_finding_traffic 先读有序清单及 version,有绑定时再按 binding_id 分段读取请求/响应。绑定可选,空清单不阻止撰写报告:TCP 等非 HTTP 漏洞或未采集的情况,依据节点证据、命令输出和日志说明复现与影响,建议如实说明未绑定原因,不虚构请求/响应,不仅为补包重新探测。报告引用稳定证据编号及用途;仅按真实内容描述。保存报告时传入所读 version 作为 evidence_version;如版本冲突,重新读取并生成,不得直接换版本重试。
|
||||
3. **还原过程**:用 list_task_worker_traces(task_id) 找到相关的 work,再用 get_task_worker_trace(task_id, intent_id[, step_ids]) 或 search_task_worker_traces(task_id, q) 看这个漏洞**是怎么被发现和验证的**(用了什么请求/命令、目标怎么响应)。必要时 get_task_graph(task_id) 看整体态势、list_task_findings(task_id) 看是否有关联漏洞。
|
||||
4. **写报告**:综合以上,写一份结构化 Markdown 报告(见下方模板)。
|
||||
5. **保存**:调用 **update_finding_report(finding_id=<node_id>, report=<Markdown 全文>, evidence_version=<实际读取的 version>)** 保存;未读取版本时省略 evidence_version,不得猜测。这是你的最终产物——不写进去等于没做。
|
||||
|
||||
━━ 报告结构(Markdown,按需裁剪,但证据/复现/修复必须有)━━
|
||||
- ` + "`## 概述`" + `:一句话说清是什么漏洞、在哪、能造成什么。
|
||||
|
||||
+111
-72
@@ -73,11 +73,12 @@ func compactIntents(ns []*db.Node, parentsOf, yieldsOf map[int64][]int64) []map[
|
||||
// ToolSet exposes the PG-backed dual graph (asset + exploration) to an LLM agent.
|
||||
// One ToolSet is created per planner/worker run; per-run signals live here.
|
||||
type ToolSet struct {
|
||||
as *db.AssetStore // asset store (optional; nil = asset tools not available)
|
||||
cs *db.CompanyStore // company store (optional)
|
||||
ts *db.ExplorationStore
|
||||
worker string
|
||||
taskID int64 // PG tasks.id; 0 when unknown (tests / orchestrator cross-task reads)
|
||||
findingRecorder FindingRecorder
|
||||
as *db.AssetStore // asset store (optional; nil = asset tools not available)
|
||||
cs *db.CompanyStore // company store (optional)
|
||||
ts *db.ExplorationStore
|
||||
worker string
|
||||
taskID int64 // PG tasks.id; 0 when unknown (tests / orchestrator cross-task reads)
|
||||
// coverageDisabled mirrors tasks.coverage_enabled=false. Stored inverted so the
|
||||
// zero value (all existing ToolSet constructions) means ENABLED — matching the
|
||||
// DB default (true). When true: graphOverviewData drops the coverage block, the
|
||||
@@ -398,7 +399,11 @@ func (t *ToolSet) graphOverviewData() map[string]any {
|
||||
for _, h := range hints {
|
||||
var p map[string]any
|
||||
_ = json.Unmarshal(h.Payload, &p)
|
||||
hsum = append(hsum, map[string]any{"id": h.ID, "state": h.State, "text": p["text"]})
|
||||
hint := map[string]any{"id": h.ID, "state": h.State, "text": p["text"]}
|
||||
if findingTrafficBindingEnabled() && p["traffic_refs"] != nil {
|
||||
hint["traffic_refs"] = p["traffic_refs"]
|
||||
}
|
||||
hsum = append(hsum, hint)
|
||||
}
|
||||
out["hints"] = hsum
|
||||
// lineage from the exploration edges: an intent's parents (what it
|
||||
@@ -1068,6 +1073,9 @@ func (t *ToolSet) listFindings() actool.CoreTool {
|
||||
obj(map[string]any{}),
|
||||
func(context.Context, json.RawMessage) (actool.Result, error) {
|
||||
f, _ := t.ts.ListByKindWithSources(db.KindFinding, 500)
|
||||
if err := t.ts.PopulateFindingTrafficIDs(f); err != nil {
|
||||
return actool.Errorf(err.Error()), nil
|
||||
}
|
||||
intentOf, _ := t.ts.FindingIntentsWithSources() // finding id -> 产生它的 intent id
|
||||
taskID := t.taskID
|
||||
if taskID <= 0 {
|
||||
@@ -1076,6 +1084,9 @@ func (t *ToolSet) listFindings() actool.CoreTool {
|
||||
out := make([]map[string]any, 0, len(f))
|
||||
for _, n := range f {
|
||||
m := compactFinding(n)
|
||||
if n.FindingID > 0 {
|
||||
m["finding_id"], m["finding_node_id"], m["traffic_count"] = n.FindingID, n.ID, n.TrafficCount
|
||||
}
|
||||
if n.Inherited {
|
||||
m["task_id"] = n.SourceTaskID
|
||||
} else {
|
||||
@@ -1167,7 +1178,10 @@ func (t *ToolSet) nodeDetail() actool.CoreTool {
|
||||
if n == nil {
|
||||
return actool.Errorf(fmt.Sprintf("未找到探索节点 %d。若你想查的是资产,请用 list_assets / asset_neighbors(资产与探索节点是不同的 id 空间,资产 id 不能传给 node_detail)。", id)), nil
|
||||
}
|
||||
return jsonResult(n) // full payload incl. detail / evidence
|
||||
if err := t.ts.PopulateFindingTrafficIDs([]*db.Node{n}); err != nil {
|
||||
return actool.Errorf(err.Error()), nil
|
||||
}
|
||||
return jsonResult(n) // full payload incl. detail / evidence, plus explicit finding IDs
|
||||
})
|
||||
}
|
||||
|
||||
@@ -1358,67 +1372,79 @@ func (t *ToolSet) goalMet() actool.CoreTool {
|
||||
// --- worker write tools ---
|
||||
|
||||
func (t *ToolSet) addFinding() actool.CoreTool {
|
||||
return writeTool("report_finding", "[重要]发现漏洞时必须调用该工具进行记录!记录一个确认的漏洞发现。在任务上下文中 intent_id 必填(当前正在执行的意图 id);在会话上下文中 intent_id 可不填。",
|
||||
obj(map[string]any{
|
||||
"vulnclass": str("漏洞类(分类,如 SQL Injection / IDOR / XSS)"),
|
||||
"name": str("漏洞名称(具体可读的标题,如『用户中心订单接口存在越权访问』;建议填写,留空时前端回退展示 vulnclass)"),
|
||||
"severity": str("critical|high|medium|low(严重/高/中/低)"),
|
||||
"summary": str("发现摘要"),
|
||||
"intent_id": idp("产生本发现的意图 id(任务上下文必填;会话上下文可不填)"),
|
||||
"asset_ids": map[string]any{"type": "array", "items": map[string]any{"type": "integer"}, "description": "【存在时尽量填写,否则在摘要中必须要写清楚漏洞位置】受影响资产 id(可选,0/1/多个):参数/端点/站点等。一个漏洞影响多处可全填,纯观察可不填。"},
|
||||
"evidence": str("证据/PoC 文本"),
|
||||
}, "vulnclass", "severity", "summary"),
|
||||
func(_ context.Context, in json.RawMessage) (actool.Result, error) {
|
||||
var a struct {
|
||||
VulnClass, Name, Severity, Summary, Evidence string
|
||||
IntentID json.RawMessage `json:"intent_id"`
|
||||
AssetIDs []json.RawMessage `json:"asset_ids"`
|
||||
return writeTool("report_finding", "记录确认的漏洞,用 evidence 提供命令输出、日志等可验证证据。任务上下文传当前 intent_id。返回的 finding_id 是独立漏洞记录 ID,finding_node_id 是探索节点 ID(第一行保留该节点编号)。", obj(map[string]any{
|
||||
"vulnclass": str("漏洞类别"), "name": str("漏洞名称"), "severity": str("critical|high|medium|low"), "summary": str("发现摘要"),
|
||||
"intent_id": idp("当前任务的意图 id"), "asset_ids": map[string]any{"type": "array", "items": map[string]any{"type": "integer"}, "description": "受影响资产 id"},
|
||||
"evidence": str("证据/PoC 文本"),
|
||||
"evidence_hint_id": idp("可选:本任务中对应此漏洞的提示节点 ID,自动携带其结构化 traffic_refs;不能引用继承提示或其他漏洞的提示"),
|
||||
"traffic_refs": map[string]any{"type": "array", "description": "可选;HTTP/HTTPS 漏洞先检索并逐条核实请求/响应确实支持漏洞结论,再按复现顺序填写真实 ID。TCP 等非 HTTP 漏洞、未采集或找不到确切记录时省略或传 [],不阻止上报;可在 evidence 说明原因并提供其他可验证证据。不要猜测 ID、按域名/时间推定关联或仅为补包重复探测。用途 baseline 正常对照 / proof 漏洞证明 / verification 补充验证 / supporting 辅助证据。",
|
||||
"items": obj(map[string]any{"traffic_id": str("traffic_search 返回的真实流量 ID"), "role": map[string]any{"type": "string", "enum": []string{"baseline", "proof", "verification", "supporting"}}, "note": str("该流量如何支持漏洞结论")}, "traffic_id")},
|
||||
}, "vulnclass", "severity", "summary"), func(ctx context.Context, in json.RawMessage) (actool.Result, error) {
|
||||
var a struct {
|
||||
VulnClass, Name, Severity, Summary, Evidence string
|
||||
IntentID json.RawMessage `json:"intent_id"`
|
||||
AssetIDs []json.RawMessage `json:"asset_ids"`
|
||||
TrafficRefs []db.TrafficRef `json:"traffic_refs"`
|
||||
EvidenceHintID json.RawMessage `json:"evidence_hint_id"`
|
||||
}
|
||||
if err := json.Unmarshal(in, &a); err != nil {
|
||||
return actool.Errorf(err.Error()), nil
|
||||
}
|
||||
if t.ts == nil {
|
||||
return actool.Errorf("report_finding 需要任务上下文;平台对话请通过 add_task_hint 向对应任务交接漏洞,并在提示中携带已有的 traffic_refs,由任务 Agent 登记。已登记漏洞可用 bind_finding_traffic 补绑。"), nil
|
||||
}
|
||||
if !findingTrafficBindingEnabled() && (len(a.TrafficRefs) > 0 || pid(a.EvidenceHintID) > 0) {
|
||||
return actool.Errorf(findingTrafficDisabled), nil
|
||||
}
|
||||
if len(a.EvidenceHintID) > 0 && pid(a.EvidenceHintID) <= 0 {
|
||||
return actool.Errorf("evidence_hint_id 必须为有效的提示节点 ID;无交接提示时省略"), nil
|
||||
}
|
||||
refs, err := t.findingRefsFromHint(pid(a.EvidenceHintID), a.TrafficRefs)
|
||||
if err != nil {
|
||||
return actool.Errorf(err.Error()), nil
|
||||
}
|
||||
input := db.RecordFindingInput{TaskID: t.taskID, ExplorationID: t.ts.ID(), IntentID: pid(a.IntentID), VulnClass: a.VulnClass, Name: a.Name, Severity: a.Severity, Summary: a.Summary, Evidence: a.Evidence, Worker: t.worker, AssetIDs: pidList(a.AssetIDs)}
|
||||
var recorded *db.RecordedFinding
|
||||
if t.findingRecorder != nil {
|
||||
recorded, err = t.findingRecorder.Record(ctx, input, refs)
|
||||
} else if len(refs) > 0 {
|
||||
return actool.Errorf("流量证据存储不可用;未登记漏洞"), nil
|
||||
} else {
|
||||
recorded, err = t.ts.RecordFinding(ctx, input)
|
||||
}
|
||||
if err != nil {
|
||||
return actool.Errorf(err.Error()), nil
|
||||
}
|
||||
if t.notifyFinding != nil {
|
||||
iid := input.IntentID
|
||||
if iid <= 0 {
|
||||
iid = t.ownerNode
|
||||
}
|
||||
_ = json.Unmarshal(in, &a)
|
||||
payload := map[string]any{"vulnclass": a.VulnClass, "name": a.Name, "severity": a.Severity, "summary": a.Summary,
|
||||
"evidence": map[string]any{"by": t.worker, "poc": a.Evidence}}
|
||||
var anchors []int64
|
||||
for _, raw := range a.AssetIDs {
|
||||
if p := pid(raw); p > 0 {
|
||||
anchors = append(anchors, p)
|
||||
}
|
||||
t.notifyFinding(iid, a.Summary)
|
||||
} else if t.notify != nil {
|
||||
t.notify()
|
||||
}
|
||||
t.writes.Findings++
|
||||
// Keep the first line's node-ID contract for existing reporter triggers.
|
||||
for i := range recorded.Traffic.Bindings {
|
||||
recorded.Traffic.Bindings[i].Snapshot.ReqHead = ""
|
||||
recorded.Traffic.Bindings[i].Snapshot.RespHead = ""
|
||||
}
|
||||
result := struct {
|
||||
*db.RecordedFinding
|
||||
EvidenceStatus string `json:"evidence_status"`
|
||||
EvidenceNote string `json:"evidence_note,omitempty"`
|
||||
}{RecordedFinding: recorded, EvidenceStatus: "bound"}
|
||||
if len(recorded.Traffic.Bindings) == 0 {
|
||||
result.EvidenceStatus = "not_bound"
|
||||
result.EvidenceNote = "漏洞已保存,未绑定流量。TCP/无包情形可正常继续;若已有核实的 HTTP 流量,请用可用的 bind_finding_traffic 或漏洞页面补绑,再完成证据交接。不要重复创建漏洞。"
|
||||
if !findingTrafficBindingEnabled() {
|
||||
result.EvidenceNote = "漏洞已保存。Agent 自动绑定流量已关闭,可在页面人工关联流量。"
|
||||
}
|
||||
var id int64
|
||||
if t.ts != nil {
|
||||
intent := pid(a.IntentID)
|
||||
if intent > 0 {
|
||||
node, err := t.ts.GetNode(intent)
|
||||
if err != nil || node == nil || node.Kind != db.KindIntent {
|
||||
return actool.Errorf("intent_id 必须是本任务的意图(关联任务意图只读)"), nil
|
||||
}
|
||||
}
|
||||
var err error
|
||||
id, err = t.ts.AddNode(db.KindFinding, payload, 9, "confirmed", t.worker, anchors)
|
||||
if err != nil {
|
||||
return actool.Errorf(err.Error()), nil
|
||||
}
|
||||
if intent > 0 {
|
||||
_ = t.ts.Link(intent, db.RelYields, id) // chain: intent -> finding
|
||||
}
|
||||
_, _ = t.ts.AddStandaloneFinding(t.taskID, id, a.VulnClass, a.Name, a.Severity, a.Summary, a.Evidence, t.worker, anchors)
|
||||
// 确证漏洞落库 → 当场唤醒本任务 planner(不等 worker 收工,debounce 合并)。
|
||||
// 优先带上下文(哪个意图+finding摘要);intent 用工具参数,缺省回退到 owner 意图。
|
||||
if t.notifyFinding != nil {
|
||||
iid := pid(a.IntentID)
|
||||
if iid <= 0 {
|
||||
iid = t.ownerNode
|
||||
}
|
||||
t.notifyFinding(iid, a.Summary)
|
||||
} else if t.notify != nil {
|
||||
t.notify()
|
||||
}
|
||||
} else {
|
||||
// conversation context: no exploration store available, cannot record finding
|
||||
return actool.Errorf("report_finding 需要任务上下文(exploration store 未初始化)"), nil
|
||||
}
|
||||
t.writes.Findings++
|
||||
return actool.Text(fmt.Sprintf("finding recorded: %d", id)), nil
|
||||
})
|
||||
}
|
||||
raw, _ := json.Marshal(result)
|
||||
return actool.Text(fmt.Sprintf("finding recorded: %d\n%s", recorded.NodeID, raw)), nil
|
||||
})
|
||||
}
|
||||
|
||||
// recordFact writes a general exploration RESULT/conclusion (not a vuln, not a
|
||||
@@ -1530,12 +1556,16 @@ func (t *ToolSet) recordFact() actool.CoreTool {
|
||||
}
|
||||
|
||||
type hintItem struct {
|
||||
Text string `json:"text"`
|
||||
AssetIDs []json.RawMessage `json:"asset_ids"`
|
||||
Text string `json:"text"`
|
||||
AssetIDs []json.RawMessage `json:"asset_ids"`
|
||||
TrafficRefs []db.TrafficRef `json:"traffic_refs"`
|
||||
}
|
||||
|
||||
// addOneHint 挂一条 hint 节点(active/human)到探索图,可锚定资产,返回 id。
|
||||
func (t *ToolSet) addOneHint(it hintItem) (int64, error) {
|
||||
if len(it.TrafficRefs) > 0 && !findingTrafficBindingEnabled() {
|
||||
return 0, fmt.Errorf("Agent 自动绑定流量已关闭,未保存携带 traffic_refs 的提示;可在系统设置开启,或仅交接文字")
|
||||
}
|
||||
if strings.TrimSpace(it.Text) == "" {
|
||||
return 0, fmt.Errorf("text 不能为空")
|
||||
}
|
||||
@@ -1545,7 +1575,15 @@ func (t *ToolSet) addOneHint(it hintItem) (int64, error) {
|
||||
anchors = append(anchors, tid)
|
||||
}
|
||||
}
|
||||
id, err := t.ts.AddNode(db.KindHint, map[string]any{"text": it.Text}, 0, "active", "human", anchors)
|
||||
refs, err := db.NormalizeTrafficRefs(it.TrafficRefs)
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
payload := map[string]any{"text": it.Text}
|
||||
if len(refs) > 0 {
|
||||
payload["traffic_refs"] = refs
|
||||
}
|
||||
id, err := t.ts.AddNode(db.KindHint, payload, 0, "active", "human", anchors)
|
||||
if err == nil && t.notify != nil {
|
||||
t.notify() // wake the planner so the new hint is read promptly (debounced)
|
||||
}
|
||||
@@ -1731,9 +1769,10 @@ func (t *ToolSet) addHint() actool.CoreTool {
|
||||
return writeTool("add_hint", "把人类/主 agent 的战略提示挂到探索图,规划者下次生成意图时会读到它。\n"+
|
||||
"★优先批量:多条提示放进 hints 数组一次提交(比逐条调用省往返)。返回 ids 数组,与 hints 等长同序(失败项 id=0,详情见 errors)。单条则省略 hints 直接给顶层 text。",
|
||||
obj(map[string]any{
|
||||
"hints": map[string]any{"type": "array", "description": "【优先用这个】要新增的提示数组,按顺序处理。每个元素字段同下方顶层字段(text/asset_ids)。返回 ids 与本数组等长、同序。", "items": map[string]any{"type": "object"}},
|
||||
"text": str("[单条] 提示内容,如'重点挖认证后接口'"),
|
||||
"asset_ids": map[string]any{"type": "array", "items": map[string]any{"type": "integer"}, "description": "锚定的资产 id(可选,0/1/多个)"},
|
||||
"hints": map[string]any{"type": "array", "description": "【优先用这个】要新增的提示数组,按顺序处理。每个元素字段同下方顶层字段(text/asset_ids/traffic_refs)。返回 ids 与本数组等长、同序。", "items": obj(map[string]any{"text": str("提示内容"), "asset_ids": map[string]any{"type": "array", "items": map[string]any{"type": "integer"}}, "traffic_refs": HintTrafficSchema()})},
|
||||
"text": str("[单条] 提示内容,如'重点挖认证后接口'"),
|
||||
"traffic_refs": HintTrafficSchema(),
|
||||
"asset_ids": map[string]any{"type": "array", "items": map[string]any{"type": "integer"}, "description": "锚定的资产 id(可选,0/1/多个)"},
|
||||
}),
|
||||
func(_ context.Context, in json.RawMessage) (actool.Result, error) {
|
||||
var a struct {
|
||||
|
||||
+17
-12
@@ -41,17 +41,18 @@ type WebSearchOpts struct {
|
||||
}
|
||||
|
||||
type Worker struct {
|
||||
prov llm.Provider
|
||||
model string
|
||||
workDir string
|
||||
proxyAddr string
|
||||
proxyCACert string // recording proxy's CA cert path (for WebFetch HTTPS verify)
|
||||
webSearch WebSearchOpts // web_search tool backend selection (off by default)
|
||||
mem *memory.Store // cross-engagement tradecraft memory (G4)
|
||||
tx *transcript.Store // raw LLM conversation persistence (nil = off)
|
||||
window int // context window in tokens (for compaction)
|
||||
windowFn func() int // optional dynamic task-chain minimum
|
||||
maxTurns int // max agent turns per run (0 = unlimited)
|
||||
findingRecorder FindingRecorder
|
||||
prov llm.Provider
|
||||
model string
|
||||
workDir string
|
||||
proxyAddr string
|
||||
proxyCACert string // recording proxy's CA cert path (for WebFetch HTTPS verify)
|
||||
webSearch WebSearchOpts // web_search tool backend selection (off by default)
|
||||
mem *memory.Store // cross-engagement tradecraft memory (G4)
|
||||
tx *transcript.Store // raw LLM conversation persistence (nil = off)
|
||||
window int // context window in tokens (for compaction)
|
||||
windowFn func() int // optional dynamic task-chain minimum
|
||||
maxTurns int // max agent turns per run (0 = unlimited)
|
||||
// runTimeout is the wall-clock budget for the main exploration of one intent
|
||||
// (0 = unlimited). When it fires, the run is cut and a settlement round is
|
||||
// forced so already-identified facts get written back instead of being lost.
|
||||
@@ -204,7 +205,9 @@ const workerDefaultTmpl = `你是一个 ARTEX 平台授权渗透测试系统的"
|
||||
**边发现边写回**(写进图才算数,脑子/文字里的不算;每得一个结果立刻写,别攒到最后被步数耗尽丢掉)。三种写回,别串图:
|
||||
- **新资产/资源 → insert_assets(资产图)**:子域 / service / endpoint / 指纹 / 凭据 等一切资产【本身】。**这里只登记资产;探索结论/判断不写这里,用 record_fact。**
|
||||
- **探索结论/事实 → record_fact(探索图,传 intent_id)**:都用它。**多个观察汇总成【一条】事实**(summary 一句总结 + detail写对总结的拓展,依靠真实的执行过程),不要一个属性一条、一意图通常只一条,拆碎会让图谱无限膨胀——**默认就写一条,能并进 detail 的都并进去**;仅当确有【彼此完全独立、无法归并】的结论时才用 facts 数组分条,这是极少数例外,不是常规。**只写增量**:只记这次【新得到】的,别把已有事实换措辞重记(只印证已有、无新增就不必记)。**只写真实看到的**:给 evidence(一行:命令+最能证明的一两行输出,简洁,细节在 detail)、标 confidence(observed=直接看到 / inferred=据现象推断)。
|
||||
- **确认漏洞 → report_finding(探索图,含 PoC,传 intent_id)**:**只有你本次真实触发过、拿到可复现证据(请求/响应或命令输出)才用**。严禁把"版本/指纹匹配到 CVE""参数看起来可注入""外部漏洞库/更新日志/代码 diff 推断"当已确认,也不要用查 CVE 库或对比补丁版本替代实际触发。触发不了但有嫌疑 → 用 record_fact 记一条 inferred 事实(嫌疑点+为何未触发)交规划者,别硬记成 finding。
|
||||
- **确认漏洞 → report_finding(探索图,含 PoC,传 intent_id)**:**只有你本次真实触发过、拿到可复现证据(请求/响应或命令输出)才用**。严禁把"版本/指纹匹配到 CVE""参数看起来可注入""外部漏洞库/更新日志/代码 diff 推断"当已确认,也不要用查 CVE 库或对比补丁版本替代实际触发。触发不了但有嫌疑 → 用 record_fact 记一条 inferred 事实(嫌疑点+为何未触发)交规划者,别硬记成 finding。有对应录制流量时,先用 traffic_search / traffic_get 核对真实 ID,再用 traffic_refs 按复现顺序绑定多条请求/响应;域名和时间只用于候选筛选,不推定任务归属。
|
||||
|
||||
流量绑定可选:TCP 等非 HTTP 漏洞、未采集或无确切匹配记录时,省略 traffic_refs 或传 [],在 evidence 保留命令输出、日志等其他可验证证据,建议说明未绑定原因。不要猜测 ID,也不要仅为补包重复探测。
|
||||
|
||||
完成本意图后用一句话总结你做了什么、写回了哪些事实。务实、克制、聚焦这一条意图。`
|
||||
|
||||
@@ -255,6 +258,7 @@ func workerSystem(proxyAddr, caCert, dataDir, runDir string) string {
|
||||
body := renderSystem("worker", workerDefaultTmpl, WorkerVars{ProxyAddr: proxyAddr, DataDir: dataDir, Now: nowStr()})
|
||||
// caCert is present only when the recording MITM is on, which is exactly when
|
||||
// the traffic_* tools are registered — so it gates the traffic-tool note.
|
||||
// Optional finding guidance is added for every role after tool resolution.
|
||||
return body + workerTrafficBlock(caCert != "") + workerArtifactSpec(runDir)
|
||||
}
|
||||
|
||||
@@ -325,6 +329,7 @@ func (w *Worker) ExecuteWithMessage(ctx context.Context, name string, taskID int
|
||||
|
||||
func (w *Worker) execute(ctx context.Context, name string, taskID int64, as *db.AssetStore, ts *db.ExplorationStore, intent *db.Node, hooks harness.HookRunner, emit func(db.Activity), enr EnrichTrigger, notifyFinding func(int64, string), requestID, message string) (harness.TerminalReason, WriteCounts, error) {
|
||||
tsx := NewToolSet(ts, name)
|
||||
tsx.SetFindingRecorder(w.findingRecorder)
|
||||
tsx.SetTaskID(taskID)
|
||||
coverageEnabled := as == nil || as.CoverageEnabled(taskID)
|
||||
tsx.SetCoverageEnabled(coverageEnabled)
|
||||
|
||||
@@ -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"`
|
||||
|
||||
@@ -0,0 +1,438 @@
|
||||
package db
|
||||
|
||||
import (
|
||||
"context"
|
||||
"crypto/sha256"
|
||||
"database/sql"
|
||||
"encoding/hex"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"strings"
|
||||
"time"
|
||||
)
|
||||
|
||||
var (
|
||||
ErrEvidenceConflict = errors.New("流量证据已变更,请刷新后重试")
|
||||
ErrFindingNotFound = errors.New("漏洞不存在")
|
||||
ErrEvidenceNotFound = errors.New("流量证据不存在")
|
||||
)
|
||||
|
||||
// This lock covers the evidence filesystem as well as its SQL references. All
|
||||
// processes sharing the database use it, including readers, exports and GC.
|
||||
const findingEvidenceLockKey int64 = 7337741004
|
||||
|
||||
func (d *DB) WithEvidenceTx(ctx context.Context, fn func(*sql.Tx) error) error {
|
||||
tx, err := d.BeginTx(ctx, nil)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer tx.Rollback()
|
||||
if _, err = tx.ExecContext(ctx, `SELECT pg_advisory_xact_lock($1)`, findingEvidenceLockKey); err != nil {
|
||||
return err
|
||||
}
|
||||
if err = fn(tx); err != nil {
|
||||
return err
|
||||
}
|
||||
return tx.Commit()
|
||||
}
|
||||
|
||||
type TrafficRef struct {
|
||||
TrafficID string `json:"traffic_id"`
|
||||
Role string `json:"role"`
|
||||
Note string `json:"note"`
|
||||
}
|
||||
|
||||
func NormalizeTrafficRefs(refs []TrafficRef) ([]TrafficRef, error) {
|
||||
out := make([]TrafficRef, 0, len(refs))
|
||||
seen := map[string]bool{}
|
||||
for _, ref := range refs {
|
||||
ref.TrafficID = strings.TrimSpace(ref.TrafficID)
|
||||
if ref.TrafficID == "" {
|
||||
return nil, errors.New("traffic_id 不能为空")
|
||||
}
|
||||
if ref.Role == "" {
|
||||
ref.Role = "supporting"
|
||||
}
|
||||
if !ValidTrafficRole(ref.Role) {
|
||||
return nil, fmt.Errorf("无效的流量用途 %q", ref.Role)
|
||||
}
|
||||
if !seen[ref.TrafficID] {
|
||||
out = append(out, ref)
|
||||
seen[ref.TrafficID] = true
|
||||
}
|
||||
}
|
||||
return out, nil
|
||||
}
|
||||
|
||||
func ValidTrafficRole(role string) bool {
|
||||
return role == "baseline" || role == "proof" || role == "verification" || role == "supporting"
|
||||
}
|
||||
|
||||
type TrafficEvidenceSnapshot struct {
|
||||
ID string `json:"id"`
|
||||
SourceTrafficID string `json:"source_traffic_id"`
|
||||
CapturedAt int64 `json:"captured_at"`
|
||||
URL string `json:"url"`
|
||||
Method string `json:"method"`
|
||||
Status int `json:"status"`
|
||||
ContentType string `json:"content_type"`
|
||||
ReqHead string `json:"req_head,omitempty"`
|
||||
RespHead string `json:"resp_head,omitempty"`
|
||||
ReqHash string `json:"req_hash"`
|
||||
RespHash string `json:"resp_hash"`
|
||||
ReqLen int64 `json:"req_len"`
|
||||
RespLen int64 `json:"resp_len"`
|
||||
CreatedAt time.Time `json:"created_at"`
|
||||
}
|
||||
|
||||
type FindingTrafficBinding struct {
|
||||
ID int64 `json:"id,string"`
|
||||
FindingID int64 `json:"finding_id,string"`
|
||||
SnapshotID string `json:"snapshot_id"`
|
||||
Role string `json:"role"`
|
||||
Note string `json:"note"`
|
||||
Position int `json:"position"`
|
||||
CreatedAt time.Time `json:"created_at"`
|
||||
Snapshot TrafficEvidenceSnapshot `json:"snapshot"`
|
||||
}
|
||||
|
||||
type FindingTraffic struct {
|
||||
FindingID int64 `json:"finding_id,string"`
|
||||
Version int64 `json:"version"`
|
||||
ReportVersion int64 `json:"report_version"`
|
||||
Bindings []FindingTrafficBinding `json:"bindings"`
|
||||
}
|
||||
|
||||
type PreparedTrafficEvidence struct {
|
||||
Ref TrafficRef
|
||||
Snapshot TrafficEvidenceSnapshot
|
||||
}
|
||||
|
||||
func TrafficSnapshotID(snapshot TrafficEvidenceSnapshot) string {
|
||||
snapshot.ID = ""
|
||||
snapshot.CreatedAt = time.Time{}
|
||||
raw, _ := json.Marshal(snapshot)
|
||||
sum := sha256.Sum256(raw)
|
||||
return hex.EncodeToString(sum[:])
|
||||
}
|
||||
|
||||
// LockTaskEvidenceTx serializes writes with archive queueing (which locks the
|
||||
// same task row). Once queued, its snapshot must not acquire new evidence.
|
||||
func LockTaskEvidenceTx(tx *sql.Tx, taskID int64) error {
|
||||
if taskID == 0 {
|
||||
return nil
|
||||
}
|
||||
var deleted sql.NullTime
|
||||
if err := tx.QueryRow(`SELECT deleted_at FROM tasks WHERE id=$1 FOR UPDATE`, taskID).Scan(&deleted); err != nil {
|
||||
return err
|
||||
}
|
||||
if deleted.Valid {
|
||||
return ErrTaskArchiveState
|
||||
}
|
||||
var state string
|
||||
err := tx.QueryRow(`SELECT state FROM task_archives WHERE task_id=$1`, taskID).Scan(&state)
|
||||
if err != nil && !errors.Is(err, sql.ErrNoRows) {
|
||||
return err
|
||||
}
|
||||
if err == nil && state != ArchiveFailed {
|
||||
return ErrTaskArchiveState
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func LockFindingEvidenceTx(tx *sql.Tx, findingID int64, version *int64) error {
|
||||
var taskID sql.NullInt64
|
||||
if err := tx.QueryRow(`SELECT task_id FROM findings WHERE id=$1`, findingID).Scan(&taskID); err != nil {
|
||||
if errors.Is(err, sql.ErrNoRows) {
|
||||
return ErrFindingNotFound
|
||||
}
|
||||
return err
|
||||
}
|
||||
if err := LockTaskEvidenceTx(tx, taskID.Int64); err != nil {
|
||||
return err
|
||||
}
|
||||
var current int64
|
||||
if err := tx.QueryRow(`SELECT evidence_version FROM findings WHERE id=$1 FOR UPDATE`, findingID).Scan(¤t); err != nil {
|
||||
if errors.Is(err, sql.ErrNoRows) {
|
||||
return ErrFindingNotFound
|
||||
}
|
||||
return err
|
||||
}
|
||||
if version != nil && *version != current {
|
||||
return ErrEvidenceConflict
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func InsertEvidenceSnapshotTx(tx *sql.Tx, s TrafficEvidenceSnapshot) error {
|
||||
if s.ID != TrafficSnapshotID(s) {
|
||||
return errors.New("证据快照元数据哈希不匹配")
|
||||
}
|
||||
if s.CreatedAt.IsZero() {
|
||||
s.CreatedAt = time.Now().UTC()
|
||||
}
|
||||
_, err := tx.Exec(`INSERT INTO traffic_evidence_snapshots
|
||||
(id,source_traffic_id,captured_at,url,method,status,content_type,req_head,resp_head,req_hash,resp_hash,req_len,resp_len,created_at)
|
||||
VALUES($1,$2,$3,$4,$5,$6,$7,$8,$9,$10,$11,$12,$13,$14) ON CONFLICT(id) DO NOTHING`,
|
||||
s.ID, s.SourceTrafficID, s.CapturedAt, s.URL, s.Method, s.Status, s.ContentType, s.ReqHead, s.RespHead, s.ReqHash, s.RespHash, s.ReqLen, s.RespLen, s.CreatedAt)
|
||||
return err
|
||||
}
|
||||
|
||||
func AddFindingTrafficTx(tx *sql.Tx, findingID int64, prepared []PreparedTrafficEvidence) error {
|
||||
var position int
|
||||
if err := tx.QueryRow(`SELECT COALESCE(MAX(position)+1,0) FROM finding_traffic_bindings WHERE finding_id=$1`, findingID).Scan(&position); err != nil {
|
||||
return err
|
||||
}
|
||||
changed := false
|
||||
for _, item := range prepared {
|
||||
if err := InsertEvidenceSnapshotTx(tx, item.Snapshot); err != nil {
|
||||
return err
|
||||
}
|
||||
if _, err := tx.Exec(`UPDATE traffic_evidence_snapshots SET unreferenced_at=NULL WHERE id=$1`, item.Snapshot.ID); err != nil {
|
||||
return err
|
||||
}
|
||||
res, err := tx.Exec(`INSERT INTO finding_traffic_bindings(finding_id,snapshot_id,role,note,position)
|
||||
VALUES($1,$2,$3,$4,$5) ON CONFLICT(finding_id,snapshot_id) DO NOTHING`, findingID, item.Snapshot.ID, item.Ref.Role, item.Ref.Note, position)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
n, err := res.RowsAffected()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if n > 0 {
|
||||
changed = true
|
||||
position++
|
||||
}
|
||||
}
|
||||
if changed {
|
||||
return bumpEvidenceVersionTx(tx, findingID)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func bumpEvidenceVersionTx(tx *sql.Tx, findingID int64) error {
|
||||
_, err := tx.Exec(`UPDATE findings SET evidence_version=evidence_version+1 WHERE id=$1`, findingID)
|
||||
return err
|
||||
}
|
||||
|
||||
func FindingTrafficTx(tx *sql.Tx, findingID int64) (*FindingTraffic, error) {
|
||||
out := &FindingTraffic{FindingID: findingID, Bindings: []FindingTrafficBinding{}}
|
||||
if err := tx.QueryRow(`SELECT evidence_version,report_evidence_version FROM findings WHERE id=$1`, findingID).Scan(&out.Version, &out.ReportVersion); err != nil {
|
||||
if errors.Is(err, sql.ErrNoRows) {
|
||||
return nil, ErrFindingNotFound
|
||||
}
|
||||
return nil, err
|
||||
}
|
||||
rows, err := tx.Query(`SELECT b.id,b.finding_id,b.snapshot_id,b.role,b.note,b.position,b.created_at,to_jsonb(s)
|
||||
FROM finding_traffic_bindings b JOIN traffic_evidence_snapshots s ON s.id=b.snapshot_id WHERE b.finding_id=$1 ORDER BY b.position,b.id`, findingID)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer rows.Close()
|
||||
for rows.Next() {
|
||||
var b FindingTrafficBinding
|
||||
var raw []byte
|
||||
if err := rows.Scan(&b.ID, &b.FindingID, &b.SnapshotID, &b.Role, &b.Note, &b.Position, &b.CreatedAt, &raw); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if err := json.Unmarshal(raw, &b.Snapshot); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
out.Bindings = append(out.Bindings, b)
|
||||
}
|
||||
return out, rows.Err()
|
||||
}
|
||||
|
||||
func (d *DB) GetFindingTraffic(ctx context.Context, findingID int64) (out *FindingTraffic, err error) {
|
||||
err = d.WithEvidenceTx(ctx, func(tx *sql.Tx) error { var e error; out, e = FindingTrafficTx(tx, findingID); return e })
|
||||
return
|
||||
}
|
||||
|
||||
func (d *DB) EditFindingTraffic(ctx context.Context, findingID, bindingID, version int64, role, note *string, remove bool, order []int64) error {
|
||||
if role != nil && !ValidTrafficRole(*role) {
|
||||
return errors.New("无效的流量用途")
|
||||
}
|
||||
return d.WithEvidenceTx(ctx, func(tx *sql.Tx) error {
|
||||
if err := LockFindingEvidenceTx(tx, findingID, &version); err != nil {
|
||||
return err
|
||||
}
|
||||
if order != nil {
|
||||
current, err := FindingTrafficTx(tx, findingID)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if len(order) != len(current.Bindings) {
|
||||
return ErrEvidenceConflict
|
||||
}
|
||||
ids := map[int64]bool{}
|
||||
for _, b := range current.Bindings {
|
||||
ids[b.ID] = true
|
||||
}
|
||||
for position, id := range order {
|
||||
if !ids[id] {
|
||||
return ErrEvidenceConflict
|
||||
}
|
||||
delete(ids, id)
|
||||
if _, err := tx.Exec(`UPDATE finding_traffic_bindings SET position=$1 WHERE finding_id=$2 AND id=$3`, position, findingID, id); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
} else {
|
||||
var res sql.Result
|
||||
var err error
|
||||
if remove {
|
||||
res, err = tx.Exec(`DELETE FROM finding_traffic_bindings WHERE finding_id=$1 AND id=$2`, findingID, bindingID)
|
||||
} else {
|
||||
res, err = tx.Exec(`UPDATE finding_traffic_bindings SET role=COALESCE($3,role),note=COALESCE($4,note) WHERE finding_id=$1 AND id=$2`, findingID, bindingID, role, note)
|
||||
}
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
n, err := res.RowsAffected()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if n == 0 {
|
||||
return ErrEvidenceNotFound
|
||||
}
|
||||
}
|
||||
return bumpEvidenceVersionTx(tx, findingID)
|
||||
})
|
||||
}
|
||||
|
||||
type RecordFindingInput struct {
|
||||
TaskID, ExplorationID, IntentID int64
|
||||
VulnClass, Name, Severity, Summary, Evidence, Worker string
|
||||
AssetIDs []int64
|
||||
}
|
||||
|
||||
type RecordedFinding struct {
|
||||
FindingID int64 `json:"finding_id,string"`
|
||||
NodeID int64 `json:"finding_node_id,string"`
|
||||
Traffic *FindingTraffic `json:"traffic"`
|
||||
}
|
||||
|
||||
func RecordFindingTx(tx *sql.Tx, in RecordFindingInput, prepared []PreparedTrafficEvidence) (*RecordedFinding, error) {
|
||||
if err := LockTaskEvidenceTx(tx, in.TaskID); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if in.TaskID > 0 {
|
||||
var expID int64
|
||||
if err := tx.QueryRow(`SELECT exploration_id FROM tasks WHERE id=$1`, in.TaskID).Scan(&expID); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if expID != in.ExplorationID {
|
||||
return nil, errors.New("漏洞所属任务与探索记录不匹配")
|
||||
}
|
||||
}
|
||||
if in.IntentID > 0 {
|
||||
var ok bool
|
||||
if err := tx.QueryRow(`SELECT EXISTS(SELECT 1 FROM exploration_nodes WHERE id=$1 AND exploration_id=$2 AND kind='intent')`, in.IntentID, in.ExplorationID).Scan(&ok); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if !ok {
|
||||
return nil, errors.New("intent_id 必须是本任务的意图(关联任务意图只读)")
|
||||
}
|
||||
}
|
||||
payload, _ := json.Marshal(map[string]any{"vulnclass": in.VulnClass, "name": in.Name, "severity": in.Severity, "summary": in.Summary, "evidence": map[string]any{"by": in.Worker, "poc": in.Evidence}})
|
||||
out := &RecordedFinding{}
|
||||
if err := tx.QueryRow(`INSERT INTO exploration_nodes(exploration_id,kind,payload,priority,state,origin)
|
||||
VALUES($1,'finding',$2,9,'confirmed',$3) RETURNING id`, in.ExplorationID, string(payload), in.Worker).Scan(&out.NodeID); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
for _, asset := range in.AssetIDs {
|
||||
if _, err := tx.Exec(`INSERT INTO exploration_anchors(node_id,asset_id) VALUES($1,$2) ON CONFLICT DO NOTHING`, out.NodeID, asset); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
}
|
||||
if in.IntentID > 0 {
|
||||
if _, err := tx.Exec(`INSERT INTO exploration_edges(exploration_id,src_id,rel,dst_id) VALUES($1,$2,$3,$4)`, in.ExplorationID, in.IntentID, RelYields, out.NodeID); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
}
|
||||
assets := in.AssetIDs
|
||||
if assets == nil {
|
||||
assets = []int64{}
|
||||
}
|
||||
raw, _ := json.Marshal(assets)
|
||||
if err := tx.QueryRow(`INSERT INTO findings(task_id,node_id,vulnclass,name,severity,summary,evidence,worker,asset_ids)
|
||||
VALUES(NULLIF($1,0),$2,$3,$4,$5,$6,$7,$8,$9) RETURNING id`, in.TaskID, out.NodeID, in.VulnClass, in.Name, in.Severity, in.Summary, in.Evidence, in.Worker, string(raw)).Scan(&out.FindingID); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if err := AddFindingTrafficTx(tx, out.FindingID, prepared); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
var err error
|
||||
out.Traffic, err = FindingTrafficTx(tx, out.FindingID)
|
||||
return out, err
|
||||
}
|
||||
|
||||
// RecordFinding is the atomic legacy/no-recorder path used by tools and tests.
|
||||
func (s *ExplorationStore) RecordFinding(ctx context.Context, in RecordFindingInput) (out *RecordedFinding, err error) {
|
||||
in.ExplorationID = s.expID
|
||||
err = s.db.WithEvidenceTx(ctx, func(tx *sql.Tx) error { var e error; out, e = RecordFindingTx(tx, in, nil); return e })
|
||||
return
|
||||
}
|
||||
|
||||
func (d *DB) FindingIDByNodeID(nodeID int64) (id int64, err error) {
|
||||
err = d.QueryRow(`SELECT id FROM findings WHERE node_id=$1`, nodeID).Scan(&id)
|
||||
return
|
||||
}
|
||||
|
||||
// PopulateFindingTrafficIDs only enriches already-visible nodes. It performs no
|
||||
// discovery or ID guessing, and leaves the legacy node ID unchanged.
|
||||
func (s *ExplorationStore) PopulateFindingTrafficIDs(nodes []*Node) error {
|
||||
byID := map[int64]*Node{}
|
||||
var args []any
|
||||
var placeholders []string
|
||||
for _, n := range nodes {
|
||||
if n == nil || n.Kind != KindFinding || byID[n.ID] != nil {
|
||||
continue
|
||||
}
|
||||
byID[n.ID] = n
|
||||
args = append(args, n.ID)
|
||||
placeholders = append(placeholders, fmt.Sprintf("$%d", len(args)))
|
||||
}
|
||||
if len(args) == 0 {
|
||||
return nil
|
||||
}
|
||||
rows, err := s.db.Query(`SELECT f.id,f.node_id,(SELECT count(*) FROM finding_traffic_bindings b WHERE b.finding_id=f.id) FROM findings f WHERE f.node_id IN (`+strings.Join(placeholders, ",")+`)`, args...)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer rows.Close()
|
||||
for rows.Next() {
|
||||
var findingID, nodeID int64
|
||||
var count int
|
||||
if err := rows.Scan(&findingID, &nodeID, &count); err != nil {
|
||||
return err
|
||||
}
|
||||
n := byID[nodeID]
|
||||
n.FindingID, n.FindingNodeID, n.TrafficCount = findingID, nodeID, count
|
||||
}
|
||||
return rows.Err()
|
||||
}
|
||||
|
||||
func (d *DB) SetFindingReportVersionByNodeID(ctx context.Context, nodeID int64, report string, version *int64) (n int64, err error) {
|
||||
err = d.WithEvidenceTx(ctx, func(tx *sql.Tx) error {
|
||||
var id int64
|
||||
if err := tx.QueryRow(`SELECT id FROM findings WHERE node_id=$1`, nodeID).Scan(&id); err != nil {
|
||||
if errors.Is(err, sql.ErrNoRows) {
|
||||
return nil
|
||||
}
|
||||
return err
|
||||
}
|
||||
if err := LockFindingEvidenceTx(tx, id, version); err != nil {
|
||||
return err
|
||||
}
|
||||
res, err := tx.Exec(`UPDATE findings SET report=$2,report_evidence_version=COALESCE($3::bigint,CASE WHEN evidence_version=0 THEN 0 ELSE -1 END) WHERE id=$1`, id, report, version)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
n, err = res.RowsAffected()
|
||||
return err
|
||||
})
|
||||
return
|
||||
}
|
||||
@@ -0,0 +1,85 @@
|
||||
package db
|
||||
|
||||
import (
|
||||
"database/sql"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
)
|
||||
|
||||
func ArchiveEvidenceSnapshots(snapshot *TaskArchiveSnapshot) ([]TrafficEvidenceSnapshot, error) {
|
||||
var out []TrafficEvidenceSnapshot
|
||||
if rawRowCount(snapshot.Tables["traffic_evidence_snapshots"]) == 0 {
|
||||
return out, nil
|
||||
}
|
||||
if snapshot.FormatVersion < 3 {
|
||||
return nil, ErrTaskArchiveFormatMismatch
|
||||
}
|
||||
err := json.Unmarshal(snapshot.Tables["traffic_evidence_snapshots"], &out)
|
||||
return out, err
|
||||
}
|
||||
|
||||
func restoreFindingTrafficTx(tx *sql.Tx, snapshot *TaskArchiveSnapshot) error {
|
||||
snapshots, err := ArchiveEvidenceSnapshots(snapshot)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
allowed := map[string]bool{}
|
||||
for _, v := range snapshots {
|
||||
if allowed[v.ID] {
|
||||
return errors.New("duplicate archived evidence snapshot")
|
||||
}
|
||||
allowed[v.ID] = true
|
||||
if err = InsertEvidenceSnapshotTx(tx, v); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
rows, err := decodeArchiveRows(snapshot.Tables["finding_traffic_bindings"])
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if len(rows) > 0 && snapshot.FormatVersion < 3 {
|
||||
return ErrTaskArchiveFormatMismatch
|
||||
}
|
||||
for _, row := range rows {
|
||||
fid, ok := jsonInt64(row["finding_id"])
|
||||
if !ok {
|
||||
return errors.New("invalid archived evidence finding id")
|
||||
}
|
||||
sid, _ := row["snapshot_id"].(string)
|
||||
role, _ := row["role"].(string)
|
||||
if !allowed[sid] || !ValidTrafficRole(role) {
|
||||
return errors.New("invalid archived evidence binding")
|
||||
}
|
||||
var owned bool
|
||||
if err = tx.QueryRow(`SELECT EXISTS(SELECT 1 FROM findings WHERE id=$1 AND task_id=$2)`, fid, snapshot.TaskID).Scan(&owned); err != nil {
|
||||
return err
|
||||
}
|
||||
if !owned {
|
||||
return fmt.Errorf("evidence references finding outside archived task: %d", fid)
|
||||
}
|
||||
}
|
||||
if len(rows) > 0 {
|
||||
raw, _ := json.Marshal(rows)
|
||||
if _, err = tx.Exec(`INSERT INTO finding_traffic_bindings SELECT * FROM json_populate_recordset(NULL::finding_traffic_bindings,$1::json)`, string(raw)); err != nil {
|
||||
return err
|
||||
}
|
||||
_, err = tx.Exec(`UPDATE traffic_evidence_snapshots s SET unreferenced_at=NULL WHERE EXISTS(SELECT 1 FROM finding_traffic_bindings b WHERE b.snapshot_id=s.id)`)
|
||||
}
|
||||
return err
|
||||
}
|
||||
|
||||
func normalizeArchivedFindingVersions(raw json.RawMessage) (json.RawMessage, error) {
|
||||
rows, err := decodeArchiveRows(raw)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
for _, row := range rows {
|
||||
for _, key := range []string{"evidence_version", "report_evidence_version"} {
|
||||
if row[key] == nil {
|
||||
row[key] = 0
|
||||
}
|
||||
}
|
||||
}
|
||||
return json.Marshal(rows)
|
||||
}
|
||||
@@ -0,0 +1,87 @@
|
||||
package db
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestFindingEvidenceLockNamespace(t *testing.T) {
|
||||
for _, reserved := range []int64{7337741001, 7337741002, 7337741003} {
|
||||
if findingEvidenceLockKey == reserved {
|
||||
t.Fatal("evidence lock collides with migration/test/company lock")
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestFindingTrafficLegacyArchiveDefaults(t *testing.T) {
|
||||
d, err := Open(testDSN(t))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer d.Close()
|
||||
if err = d.EnsureLLMRecordsTable(); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err = d.EnsureLLMUsageTable(); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
for _, version := range []int{1, 2} {
|
||||
t.Run(fmt.Sprint(version), func(t *testing.T) {
|
||||
task, err := d.CreateTask("legacy archive evidence defaults", "fixture", nil, 0, 0)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer func() { d.Exec(`DELETE FROM task_archives WHERE task_id=$1`, task.ID); d.DeleteTask(task.ID) }()
|
||||
f, err := d.Exploration(task.ExplorationID).RecordFinding(t.Context(), RecordFindingInput{TaskID: task.ID, ExplorationID: task.ExplorationID, Summary: "legacy", Severity: "low"})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err = d.SetPaused(task.ID, true); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
archive, err := d.QueueTaskArchive(task.ID)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if _, err = d.ClaimTaskArchiveJob(t.Context()); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
snapshot, err := d.SnapshotTaskArchive(task.ID)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err = d.CompleteTaskArchive(archive.ID, snapshot, "/tmp/legacy-test.tar.zst", "fixture", 1, 1); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
snapshot.FormatVersion = version
|
||||
delete(snapshot.Tables, "finding_traffic_bindings")
|
||||
delete(snapshot.Tables, "traffic_evidence_snapshots")
|
||||
var rows []map[string]any
|
||||
if err = json.Unmarshal(snapshot.Tables["findings"], &rows); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
for _, row := range rows {
|
||||
delete(row, "evidence_version")
|
||||
delete(row, "report_evidence_version")
|
||||
}
|
||||
snapshot.Tables["findings"], _ = json.Marshal(rows)
|
||||
if _, err = d.QueueTaskArchiveRestore(archive.ID); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if _, err = d.ClaimTaskArchiveJob(t.Context()); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if _, err = d.RestoreTaskArchive(archive.ID, snapshot, 0); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
got, err := d.GetFindingTraffic(t.Context(), f.FindingID)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if got.Version != 0 || got.ReportVersion != 0 || len(got.Bindings) != 0 {
|
||||
t.Fatal(got)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
+37
-25
@@ -1,6 +1,7 @@
|
||||
package db
|
||||
|
||||
import (
|
||||
"context"
|
||||
"database/sql"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
@@ -14,6 +15,11 @@ import (
|
||||
// DBFinding is a row in the standalone findings table. It persists across task
|
||||
// deletion unless the caller explicitly requests related finding cleanup.
|
||||
type DBFinding struct {
|
||||
TrafficCount int
|
||||
EvidenceVersion int64
|
||||
ReportEvidenceVersion int64
|
||||
TrafficBindings []FindingTrafficBinding // populated only for export
|
||||
|
||||
ID int64
|
||||
TaskID *int64
|
||||
NodeID *int64
|
||||
@@ -98,7 +104,8 @@ func (d *DB) AddFinding(taskID, nodeID int64, vulnclass, name, severity, summary
|
||||
// list query selects, so scanFinding stays in sync across callers.
|
||||
const findingSelectCols = `f.id, f.task_id, f.node_id, f.vulnclass, COALESCE(f.name, ''), f.severity, f.summary,
|
||||
f.evidence, f.worker, f.asset_ids, COALESCE(f.status, 'pending'), f.created_at,
|
||||
COALESCE(t.description, '') AS task_description`
|
||||
COALESCE(t.description, '') AS task_description, f.evidence_version, f.report_evidence_version,
|
||||
(SELECT count(*) FROM finding_traffic_bindings b WHERE b.finding_id=f.id)`
|
||||
|
||||
// scanFindings materializes rows selected via findingSelectCols.
|
||||
func scanFindings(rows interface {
|
||||
@@ -111,7 +118,7 @@ func scanFindings(rows interface {
|
||||
f := &DBFinding{}
|
||||
var aidsJSON string
|
||||
if err := rows.Scan(&f.ID, &f.TaskID, &f.NodeID, &f.VulnClass, &f.Name, &f.Severity,
|
||||
&f.Summary, &f.Evidence, &f.Worker, &aidsJSON, &f.Status, &f.CreatedAt, &f.TaskDescription); err != nil {
|
||||
&f.Summary, &f.Evidence, &f.Worker, &aidsJSON, &f.Status, &f.CreatedAt, &f.TaskDescription, &f.EvidenceVersion, &f.ReportEvidenceVersion, &f.TrafficCount); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
_ = json.Unmarshal([]byte(aidsJSON), &f.AssetIDs)
|
||||
@@ -504,7 +511,7 @@ func (d *DB) ListFindingsForExport(f FindingFilter, ids []int64) ([]*DBFinding,
|
||||
var aidsJSON string
|
||||
if err := rows.Scan(&f.ID, &f.TaskID, &f.NodeID, &f.VulnClass, &f.Name, &f.Severity,
|
||||
&f.Summary, &f.Evidence, &f.Worker, &aidsJSON, &f.Status, &f.CreatedAt,
|
||||
&f.TaskDescription, &f.Report); err != nil {
|
||||
&f.TaskDescription, &f.EvidenceVersion, &f.ReportEvidenceVersion, &f.TrafficCount, &f.Report); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
_ = json.Unmarshal([]byte(aidsJSON), &f.AssetIDs)
|
||||
@@ -632,7 +639,7 @@ func (d *DB) GetFinding(id int64) (*DBFinding, error) {
|
||||
WHERE f.id = $1`, id).Scan(
|
||||
&f.ID, &f.TaskID, &f.NodeID, &f.VulnClass, &f.Name, &f.Severity,
|
||||
&f.Summary, &f.Evidence, &f.Worker, &aidsJSON, &f.Status, &f.CreatedAt,
|
||||
&f.TaskDescription, &f.Report)
|
||||
&f.TaskDescription, &f.EvidenceVersion, &f.ReportEvidenceVersion, &f.TrafficCount, &f.Report)
|
||||
if err == sql.ErrNoRows {
|
||||
return nil, nil
|
||||
}
|
||||
@@ -648,20 +655,27 @@ func (d *DB) GetFinding(id int64) (*DBFinding, error) {
|
||||
// list, the per-task 发现 Tab, and the exploration graph alike. Deleting the node
|
||||
// cascades its edges + node_assets and nulls any activity referencing it. Returns
|
||||
// rows affected (0 = no finding with that id).
|
||||
func (d *DB) DeleteFinding(id int64) (int64, error) {
|
||||
var nodeID *int64
|
||||
err := d.QueryRow(`DELETE FROM findings WHERE id=$1 RETURNING node_id`, id).Scan(&nodeID)
|
||||
if err == sql.ErrNoRows {
|
||||
return 0, nil
|
||||
}
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
if nodeID != nil {
|
||||
// best-effort: the finding is already gone; a stray node must not fail the op.
|
||||
_, _ = d.Exec(`DELETE FROM exploration_nodes WHERE id=$1 AND kind='finding'`, *nodeID)
|
||||
}
|
||||
return 1, nil
|
||||
func (d *DB) DeleteFinding(id int64) (n int64, err error) {
|
||||
err = d.WithEvidenceTx(context.Background(), func(tx *sql.Tx) error {
|
||||
if err := LockFindingEvidenceTx(tx, id, nil); err != nil {
|
||||
if errors.Is(err, ErrFindingNotFound) {
|
||||
return nil
|
||||
}
|
||||
return err
|
||||
}
|
||||
var nodeID sql.NullInt64
|
||||
if err := tx.QueryRow(`DELETE FROM findings WHERE id=$1 RETURNING node_id`, id).Scan(&nodeID); err != nil {
|
||||
return err
|
||||
}
|
||||
if nodeID.Valid {
|
||||
if _, err := tx.Exec(`DELETE FROM exploration_nodes WHERE id=$1 AND kind='finding'`, nodeID.Int64); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
n = 1
|
||||
return nil
|
||||
})
|
||||
return
|
||||
}
|
||||
|
||||
// DeleteFindingsByTask removes all findings rows of a task. The originating
|
||||
@@ -688,11 +702,7 @@ func (d *DB) SetFindingStatus(id int64, status string) (int64, error) {
|
||||
// whose node_id matches — report_finding returns that node id, so an agent tool
|
||||
// can address the finding it just created. Returns rows affected (0 when no row).
|
||||
func (d *DB) SetFindingReportByNodeID(nodeID int64, report string) (int64, error) {
|
||||
res, err := d.Exec(`UPDATE findings SET report=$1 WHERE node_id=$2`, report, nodeID)
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
return res.RowsAffected()
|
||||
return d.SetFindingReportVersionByNodeID(context.Background(), nodeID, report, nil)
|
||||
}
|
||||
|
||||
// setFindingCol updates one text column on the standalone finding row AND mirrors
|
||||
@@ -738,6 +748,8 @@ func (d *DB) SetFindingVulnClass(id int64, vulnclass string) (int64, error) {
|
||||
// FindingMeta is the standalone-row data (id, triage state, anchored assets) the
|
||||
// per-task view grafts onto its exploration-node findings.
|
||||
type FindingMeta struct {
|
||||
TrafficCount int
|
||||
|
||||
ID int64
|
||||
Status string
|
||||
AssetIDs []int64
|
||||
@@ -758,7 +770,7 @@ func (d *DB) FindingMetaByNodeID(taskID int64) (map[int64]FindingMeta, error) {
|
||||
if taskID <= 0 {
|
||||
return out, nil
|
||||
}
|
||||
rows, err := d.Query(`SELECT node_id, id, COALESCE(status,'pending'), asset_ids FROM findings
|
||||
rows, err := d.Query(`SELECT node_id, id, COALESCE(status,'pending'), asset_ids, (SELECT count(*) FROM finding_traffic_bindings b WHERE b.finding_id=findings.id) FROM findings
|
||||
WHERE task_id=$1 AND node_id IS NOT NULL`, taskID)
|
||||
if err != nil {
|
||||
return out, err
|
||||
@@ -768,7 +780,7 @@ func (d *DB) FindingMetaByNodeID(taskID int64) (map[int64]FindingMeta, error) {
|
||||
var nid int64
|
||||
var m FindingMeta
|
||||
var aidsJSON string
|
||||
if err := rows.Scan(&nid, &m.ID, &m.Status, &aidsJSON); err != nil {
|
||||
if err := rows.Scan(&nid, &m.ID, &m.Status, &aidsJSON, &m.TrafficCount); err != nil {
|
||||
return out, err
|
||||
}
|
||||
_ = json.Unmarshal([]byte(aidsJSON), &m.AssetIDs)
|
||||
|
||||
@@ -1095,6 +1095,39 @@ DROP TRIGGER IF EXISTS trg_conversation_retest_delete ON conversations;
|
||||
CREATE TRIGGER trg_conversation_retest_delete BEFORE DELETE ON conversations
|
||||
FOR EACH ROW EXECUTE FUNCTION stop_deleted_conversation_retest();
|
||||
|
||||
ALTER TABLE findings ADD COLUMN IF NOT EXISTS evidence_version BIGINT NOT NULL DEFAULT 0;
|
||||
ALTER TABLE findings ADD COLUMN IF NOT EXISTS report_evidence_version BIGINT NOT NULL DEFAULT 0;
|
||||
|
||||
CREATE TABLE IF NOT EXISTS traffic_evidence_snapshots (
|
||||
id TEXT PRIMARY KEY,
|
||||
source_traffic_id TEXT NOT NULL,
|
||||
captured_at BIGINT NOT NULL,
|
||||
url TEXT NOT NULL,
|
||||
method TEXT NOT NULL,
|
||||
status INTEGER NOT NULL,
|
||||
content_type TEXT NOT NULL DEFAULT '',
|
||||
req_head TEXT NOT NULL,
|
||||
resp_head TEXT NOT NULL,
|
||||
req_hash TEXT NOT NULL,
|
||||
resp_hash TEXT NOT NULL,
|
||||
req_len BIGINT NOT NULL,
|
||||
resp_len BIGINT NOT NULL,
|
||||
unreferenced_at TIMESTAMPTZ,
|
||||
created_at TIMESTAMPTZ NOT NULL DEFAULT now()
|
||||
);
|
||||
CREATE TABLE IF NOT EXISTS finding_traffic_bindings (
|
||||
id BIGSERIAL PRIMARY KEY,
|
||||
finding_id BIGINT NOT NULL REFERENCES findings(id) ON DELETE CASCADE,
|
||||
snapshot_id TEXT NOT NULL REFERENCES traffic_evidence_snapshots(id),
|
||||
role TEXT NOT NULL DEFAULT 'supporting',
|
||||
note TEXT NOT NULL DEFAULT '',
|
||||
position INTEGER NOT NULL,
|
||||
created_at TIMESTAMPTZ NOT NULL DEFAULT now(),
|
||||
UNIQUE(finding_id, snapshot_id)
|
||||
);
|
||||
CREATE INDEX IF NOT EXISTS idx_finding_traffic_order ON finding_traffic_bindings(finding_id, position, id);
|
||||
CREATE INDEX IF NOT EXISTS idx_finding_traffic_snapshot ON finding_traffic_bindings(snapshot_id);
|
||||
|
||||
-- =====================================================================
|
||||
-- M. 后端日志持久化
|
||||
-- =====================================================================
|
||||
|
||||
+3
-1
@@ -16,7 +16,7 @@ import (
|
||||
)
|
||||
|
||||
const (
|
||||
TaskArchiveFormatVersion = 2
|
||||
TaskArchiveFormatVersion = 3
|
||||
TaskArchiveLegacyFormatVersion = 1
|
||||
TaskArchiveLLMRecordsPath = "database/llm_records.ndjson"
|
||||
)
|
||||
@@ -587,6 +587,8 @@ func (d *DB) snapshotTaskArchive(taskID int64, llmRecords io.Writer) (*TaskArchi
|
||||
{"task_llm_profiles", `SELECT * FROM task_llm_profiles WHERE task_id=$1 ORDER BY position`, []any{taskID}},
|
||||
{"task_scope", `SELECT * FROM task_scope WHERE task_id=$1 ORDER BY id`, []any{taskID}},
|
||||
{"findings", `SELECT * FROM findings WHERE task_id=$1 ORDER BY id`, []any{taskID}},
|
||||
{"finding_traffic_bindings", `SELECT b.* FROM finding_traffic_bindings b JOIN findings f ON f.id=b.finding_id WHERE f.task_id=$1 ORDER BY b.finding_id,b.position,b.id`, []any{taskID}},
|
||||
{"traffic_evidence_snapshots", `SELECT s.* FROM traffic_evidence_snapshots s WHERE EXISTS(SELECT 1 FROM finding_traffic_bindings b JOIN findings f ON f.id=b.finding_id WHERE b.snapshot_id=s.id AND f.task_id=$1) ORDER BY s.id`, []any{taskID}},
|
||||
{"llm_records", `SELECT * FROM llm_records WHERE COALESCE(task_id,'')=$1 ORDER BY id`, []any{strconv.FormatInt(taskID, 10)}},
|
||||
{"llm_usage", `SELECT * FROM llm_usage WHERE COALESCE(task_id,'')=$1 OR exploration_id=$2 ORDER BY id`, []any{strconv.FormatInt(taskID, 10), expID}},
|
||||
{"skill_usage", `SELECT * FROM skill_usage WHERE task_id=$1 OR exploration_id=$2 ORDER BY id`, []any{taskID, expID}},
|
||||
|
||||
@@ -233,6 +233,10 @@ WHERE archive.id=$1 FOR UPDATE OF archive,task`, archiveID).Scan(&taskID, &expID
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
remappedTables["findings"], err = normalizeArchivedFindingVersions(remappedTables["findings"])
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
// Insert graph rows in foreign-key order. The archived stub has no graph rows,
|
||||
// so an ID conflict signals external corruption and must stop the restore.
|
||||
for _, table := range []string{"exploration_nodes", "exploration_edges", "exploration_anchors", "task_constraints", "activity"} {
|
||||
@@ -260,6 +264,9 @@ WHERE archive.id=$1 FOR UPDATE OF archive,task`, archiveID).Scan(&taskID, &expID
|
||||
return nil, fmt.Errorf("restore %s: %w", table, err)
|
||||
}
|
||||
}
|
||||
if err := restoreFindingTrafficTx(tx, snapshot); err != nil {
|
||||
return nil, fmt.Errorf("restore finding traffic: %w", err)
|
||||
}
|
||||
if streamedLLMRecords != "" {
|
||||
count, err := insertArchiveJSONSequenceRows(tx, "llm_records", llmRecords)
|
||||
if err != nil {
|
||||
|
||||
@@ -0,0 +1,99 @@
|
||||
# 漏洞多流量证据
|
||||
|
||||
漏洞详情的「关联流量」支持跨页多选、用途和说明、排序及解除绑定。流量页也可多选记录,一次关联到一个已有漏洞。任务继承的漏洞证据只读,修改需进入来源任务。
|
||||
|
||||
系统设置中的「Agent 自动绑定流量」默认关闭,读取和修改接口为 `/api/settings` 的 `agent_traffic_binding`。开启会增加查阅请求/响应、工具调用及提示词带来的 Token 消耗;下一轮 Agent 使用新设置。关闭时隐藏自动绑定参数和补绑工具,不注入自动绑定指引,并拒绝已运行会话新提交的自动绑定。人工绑定、流量捕获、已保存证据读取和导出不受影响。
|
||||
|
||||
开启后,默认流程是「发现入库 → 自动触发报告 Agent → 核对并绑定流量 → 按最新证据版本写报告」。上报者在 `evidence` 保留验证命令、关键输出、已有真实流量 ID 及用途;报告 Agent 结合漏洞详情和执行记录,用 `traffic_search` / `traffic_get` 核实后调用 `bind_finding_traffic`,再读取最新 `version` 保存报告。关闭开关时,报告 Agent 不额外获得原始流量检索/读取工具,也不自动绑定;仍能读取人工绑定的快照并生成报告。
|
||||
|
||||
兼容已有调用者:`report_finding` 的 `traffic_refs` / `evidence_hint_id` 仍可显式即时绑定。绑定可选:TCP 等非 HTTP 漏洞、未采集或找不到确切记录时,省略引用仍可正常上报和撰写报告;保留命令输出、日志等其他可验证证据,建议说明未绑定原因,不新增必填字段。提交的 ID 必须全部有效且正文完整。任何一条失败都会回滚本次绑定操作;显式随上报绑定失败时整次上报回滚。重复追加同一快照不增加绑定,也不覆盖说明。
|
||||
|
||||
## Agent 编号与报告版本
|
||||
|
||||
`report_finding` 新增可选参数,数组顺序即初始证据顺序:
|
||||
|
||||
```json
|
||||
{
|
||||
"traffic_refs": [
|
||||
{"traffic_id": "真实流量ID", "role": "baseline", "note": "正常账户请求"},
|
||||
{"traffic_id": "另一个真实流量ID", "role": "proof", "note": "复现请求"}
|
||||
]
|
||||
}
|
||||
```
|
||||
|
||||
用途为 `baseline`(正常对照)、`proof`(漏洞证明)、`verification`(补充验证)、`supporting`(辅助证据,默认)。先用 `traffic_search` / `traffic_get` 核对真实记录;域名和时间仅用于候选筛选,不推定任务归属。
|
||||
|
||||
提示词参考 [CyberStrikeAI 的漏洞上报工具指引](https://github.com/RuoJi6/CyberStrikeAI/blob/54d56774b8bd285817d16d48b70a4a5e6e0963f7/internal/app/vulnerability_tools.go),结合本项目的可选绑定约定,不引入无包原因必填校验。不得猜测 ID,也不应仅为补包重复探测。
|
||||
|
||||
返回第一行仍为 `finding recorded: <探索节点 ID>`;随后 JSON 提供独立漏洞记录 `finding_id`、探索节点 `finding_node_id` 和绑定摘要。
|
||||
|
||||
- `get_finding_traffic(finding_id)` 使用**独立漏洞记录 ID**,返回有序清单及 `version`;传 `binding_id`、`side=request|response`、`offset`、`length` 可分段读取,每段最多 8192 字节。
|
||||
- `update_finding_report` 的 `finding_id` **继续使用探索节点 ID**。新增 `evidence_version` 填实际读取的版本;生成期间证据变化会拒绝旧版本写入,必须重新读取并生成。
|
||||
- 旧报告调用未传版本时,不宣称覆盖已有流量证据。绑定、说明、用途或排序变化后,已有报告提示待更新。
|
||||
|
||||
开启自动绑定后,`add_hint` / `add_task_hint` 支持在单条提示或批量 `hints` 的每个元素中保存 `traffic_refs`。Planner 代为上报可传 `evidence_hint_id`,明确选取本任务对应提示中的引用,不能引用继承提示。系统不会按域名、时间或浏览记录猜测绑定。提交失败不生成部分漏洞或提前触发报告。
|
||||
|
||||
已有漏洞漏绑时可用 `bind_finding_traffic(finding_id, traffic_refs)` 补绑,无需重复登记。`list_findings` / `list_task_findings` / `node_detail` / `get_task_node_detail` 返回明确的 `finding_id` 和 `finding_node_id`;旧 `id` 保持探索节点语义。
|
||||
|
||||
启动时仅为旧工具 schema 增加可选属性,原始默认流量工具绑定扩展至报告 Agent;补绑工具默认交给报告 Agent。自定义绑定列表、提示词、描述及启用状态保留。指引在最终工具装配后统一加入:上报角色负责交接已有证据,报告 Agent 负责核对、绑定和写报告。平台对话缺少任务上下文时应通过结构化提示交接给任务 Agent,不直接上报。判定任务完成前应先交接已有证据;无包不强制等待。失败的 `report_finding` 不触发报告 Agent。
|
||||
|
||||
## API
|
||||
|
||||
基础路径:`/api/exploration/findings/{finding_id}/traffic`,使用独立漏洞 ID。沿用认证;`context_task` 校验任务可见性及继承只读。
|
||||
|
||||
| 方法 / 相对路径 | 请求 / 返回 |
|
||||
| --- | --- |
|
||||
| `GET` | 有序摘要、证据版本、报告采用版本 |
|
||||
| `POST` | `{"traffic_refs":[...]}` 整批追加 |
|
||||
| `PATCH /{binding_id}` | `{"version":1,"role":"proof","note":"说明"}` |
|
||||
| `DELETE /{binding_id}` | `{"version":1}` |
|
||||
| `PUT /order` | `{"version":1,"binding_ids":["2","1"]}`,必须是完整列表 |
|
||||
| `GET /{binding_id}` | 快照元数据和有界正文预览 |
|
||||
| `GET /{binding_id}/body` | `side`、`offset`、`length`;`download=1` 下载完整原字节 |
|
||||
|
||||
版本/排序集合冲突、归档中写入返回 `409`;继承写入 `403`;不存在或不属于漏洞的绑定 `404`;流量/附件读取及校验失败明确返回错误。
|
||||
|
||||
## 存储与迁移
|
||||
|
||||
启动幂等迁移 PostgreSQL:新增 `traffic_evidence_snapshots`、`finding_traffic_bindings`,以及 `findings.evidence_version` / `report_evidence_version`(默认 0)。不根据历史文字猜测补绑。
|
||||
|
||||
快照保存原流量 ID、采集时间、URL、方法、状态、请求/响应头、正文长度和 SHA-256。正文按哈希存放在 `<data>/evidence/blobs/<前两位>/<hash>.bin`,与可清理的 `data/traffic` 独立,多个漏洞可共享快照/正文。快照不提供内容更新接口,校验不一致时读取和导出失败。
|
||||
|
||||
在原流量写锁下读取完整正文,包括大正文 blob 和旧目录记录。先持久化并校验文件,再用单个 PostgreSQL 事务写入探索节点、意图关系、漏洞、快照和绑定;提交后才通知规划者。失败可能遗留无引用文件,但不产生部分业务记录。
|
||||
|
||||
PostgreSQL advisory lock `7337741004` 协调证据文件和 SQL 引用;任务行锁禁止归档排队后的证据修改。恢复从正文安装到元数据提交全程保持证据锁。删除漏洞级联移除绑定。
|
||||
|
||||
清理器每小时运行,只回收无引用且非进行中操作的内容,至少延迟 24 小时。普通流量清理不触及证据目录。备份热数据时同时备份 PostgreSQL 和 `data/evidence`。
|
||||
|
||||
## 导出与归档
|
||||
|
||||
Markdown 包含有序证据清单及版本;JSON 包含元数据;CSV 增加数量和绑定 ID。`md-zip` 保留漏洞 Markdown,并提供:
|
||||
|
||||
```text
|
||||
evidence/<finding_id>/<binding_id>/
|
||||
manifest.json
|
||||
request.http
|
||||
response.http
|
||||
request.bin
|
||||
response.bin
|
||||
```
|
||||
|
||||
Markdown 通过相对链接引用报文。发送下载前完成附件复制、哈希校验、压缩、磁盘同步及全部 ZIP 条目的 CRC 读取校验;缺失/损坏使整个下载失败。完整附件保留二进制原字节。
|
||||
|
||||
归档 v3 按漏洞绑定关系收集快照和正文,不依赖原流量或域名。包校验完成后才清理热数据;共享证据继续保留。恢复先校验安装正文,再事务恢复元数据和绑定,支持失败重试。v1/v2 继续可恢复,缺失的新字段显式补为 0。
|
||||
|
||||
## 验证与边界
|
||||
|
||||
为每个测试包设置独立、新建的 PostgreSQL 测试库,通过 `ARTEX_PG_DSN` 指定,避免残留任务/模型夹具触发后台运行。运行相关包完整测试,并确认没有配置缺失导致的跳过:
|
||||
|
||||
```sh
|
||||
# 每个包运行前将 ARTEX_PG_DSN 设为对应的独立测试库;显式配置失败必须报错。
|
||||
go test ./<package> -count=1
|
||||
go test -race -p 1 ./evidence ./db ./agent ./server -run 'TestEvidence|TestFindingTraffic|TestFindingEvidence|TestReportFindingAtomicContract|TestTaskArchive'
|
||||
```
|
||||
|
||||
前端验证包含 `npx tsc --noEmit`、受影响文件的 Biome 检查、Webpack 构建和 `NEXT_EXPORT=1` 静态导出。使用独立缓存目录,避免覆盖运行中的开发服务。
|
||||
|
||||
本机端到端验收使用独立端口、受控 HTTP / 域名 HTTPS 目标和临时数据目录,覆盖两个绑定入口、跨页选择、排序/说明、错误提示、继承只读、下载,以及删除原流量后导出、归档、回收热正文、恢复并校验哈希。
|
||||
|
||||
首版采用全局证据协调锁;大批量绑定/导出或慢速附件下载期间,其他证据操作可能等待。缺少录制记录或完整正文时不能补造证据。本功能不改变捕获开关,也不处理 HTTPS 直接访问 IP 的证书问题。
|
||||
@@ -0,0 +1,401 @@
|
||||
// Package evidence preserves finding evidence independently of disposable traffic.
|
||||
package evidence
|
||||
|
||||
import (
|
||||
"context"
|
||||
"crypto/sha256"
|
||||
"database/sql"
|
||||
"encoding/hex"
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
"io/fs"
|
||||
"log"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/Autumn-27/artex/db"
|
||||
"github.com/Autumn-27/artex/traffic"
|
||||
)
|
||||
|
||||
type Store struct {
|
||||
DB *db.DB
|
||||
Traffic *traffic.Traffic
|
||||
Dir string
|
||||
}
|
||||
|
||||
func New(pg *db.DB, tr *traffic.Traffic, dir string) *Store {
|
||||
return &Store{DB: pg, Traffic: tr, Dir: dir}
|
||||
}
|
||||
|
||||
func hashPath(dir, hash string) (string, error) {
|
||||
if len(hash) != 64 {
|
||||
return "", errors.New("invalid evidence hash")
|
||||
}
|
||||
if _, err := hex.DecodeString(hash); err != nil || strings.ToLower(hash) != hash {
|
||||
return "", errors.New("invalid evidence hash")
|
||||
}
|
||||
return filepath.Join(dir, "blobs", hash[:2], hash+".bin"), nil
|
||||
}
|
||||
|
||||
func verifyFile(path, hash string, length int64) error {
|
||||
f, err := os.Open(path)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer f.Close()
|
||||
h := sha256.New()
|
||||
n, err := io.Copy(h, f)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if n != length || hex.EncodeToString(h.Sum(nil)) != hash {
|
||||
return fmt.Errorf("证据正文校验失败: %s", hash)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// A new body becomes visible only after a durable write. Failed SQL commits may
|
||||
// leave unreferenced files; GC reaps those after a full day's grace period.
|
||||
func (s *Store) writeBody(r io.Reader, expectedLength int64, expectedHash string) (hash string, err error) {
|
||||
stage := filepath.Join(s.Dir, ".staging")
|
||||
if err = os.MkdirAll(stage, 0o700); err != nil {
|
||||
return
|
||||
}
|
||||
f, err := os.CreateTemp(stage, "body-")
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
defer func() { f.Close(); os.Remove(f.Name()) }()
|
||||
h := sha256.New()
|
||||
n, err := io.Copy(io.MultiWriter(f, h), r)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
hash = hex.EncodeToString(h.Sum(nil))
|
||||
if n != expectedLength {
|
||||
return "", fmt.Errorf("正文不完整: 预期 %d 字节,读取 %d 字节", expectedLength, n)
|
||||
}
|
||||
if expectedHash != "" && expectedHash != hash {
|
||||
return "", errors.New("原始流量正文哈希不匹配")
|
||||
}
|
||||
if err = f.Sync(); err != nil {
|
||||
return "", err
|
||||
}
|
||||
if err = f.Close(); err != nil {
|
||||
return "", err
|
||||
}
|
||||
path, err := hashPath(s.Dir, hash)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
if _, err = os.Stat(path); err == nil {
|
||||
if err = verifyFile(path, hash, n); err != nil {
|
||||
return "", err
|
||||
}
|
||||
// Refresh the grace period for a restored but not-yet-committed body.
|
||||
now := time.Now()
|
||||
return hash, os.Chtimes(path, now, now)
|
||||
} else if !os.IsNotExist(err) {
|
||||
return "", err
|
||||
}
|
||||
if err = os.MkdirAll(filepath.Dir(path), 0o700); err != nil {
|
||||
return "", err
|
||||
}
|
||||
if err = os.Rename(f.Name(), path); err != nil {
|
||||
return "", err
|
||||
}
|
||||
d, err := os.Open(filepath.Dir(path))
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
defer d.Close()
|
||||
return hash, d.Sync()
|
||||
}
|
||||
|
||||
// StageFindingsExport freezes bindings/report versions and makes private body
|
||||
// copies before the HTTP response is started. The caller owns and removes dest.
|
||||
func (s *Store) StageFindingsExport(ctx context.Context, findings []*db.DBFinding, dest string, copyBodies bool) error {
|
||||
return s.DB.WithEvidenceTx(ctx, func(tx *sql.Tx) error {
|
||||
var snapshots []db.TrafficEvidenceSnapshot
|
||||
for _, f := range findings {
|
||||
list, err := db.FindingTrafficTx(tx, f.ID)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
f.TrafficBindings = list.Bindings
|
||||
f.TrafficCount = len(list.Bindings)
|
||||
f.EvidenceVersion = list.Version
|
||||
f.ReportEvidenceVersion = list.ReportVersion
|
||||
if err = tx.QueryRow(`SELECT report FROM findings WHERE id=$1`, f.ID).Scan(&f.Report); err != nil {
|
||||
return err
|
||||
}
|
||||
for _, b := range list.Bindings {
|
||||
snapshots = append(snapshots, b.Snapshot)
|
||||
}
|
||||
}
|
||||
if copyBodies {
|
||||
return s.copySnapshots(snapshots, dest)
|
||||
}
|
||||
return nil
|
||||
})
|
||||
}
|
||||
|
||||
func (s *Store) prepare(ctx context.Context, refs []db.TrafficRef) ([]db.PreparedTrafficEvidence, error) {
|
||||
refs, err := db.NormalizeTrafficRefs(refs)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
out := make([]db.PreparedTrafficEvidence, 0, len(refs))
|
||||
if len(refs) == 0 {
|
||||
return out, nil
|
||||
}
|
||||
ids := make([]string, len(refs))
|
||||
byID := map[string]db.TrafficRef{}
|
||||
for i, ref := range refs {
|
||||
ids[i] = ref.TrafficID
|
||||
byID[ref.TrafficID] = ref
|
||||
}
|
||||
err = s.Traffic.ReadEvidence(ctx, ids, func(e traffic.EvidenceExchange) error {
|
||||
rh, err := s.writeBody(e.Request, e.ReqLen, e.ReqHash)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
ph, err := s.writeBody(e.Response, e.RespLen, e.RespHash)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
v := db.TrafficEvidenceSnapshot{SourceTrafficID: e.ID, CapturedAt: e.TS, URL: e.URL, Method: e.Method, Status: e.Status, ContentType: e.ContentType,
|
||||
ReqHead: e.ReqHead, RespHead: e.RespHead, ReqHash: rh, RespHash: ph, ReqLen: e.ReqLen, RespLen: e.RespLen}
|
||||
v.ID = db.TrafficSnapshotID(v)
|
||||
out = append(out, db.PreparedTrafficEvidence{Ref: byID[e.ID], Snapshot: v})
|
||||
return nil
|
||||
})
|
||||
return out, err
|
||||
}
|
||||
|
||||
func (s *Store) Record(ctx context.Context, in db.RecordFindingInput, refs []db.TrafficRef) (out *db.RecordedFinding, err error) {
|
||||
err = s.DB.WithEvidenceTx(ctx, func(tx *sql.Tx) error {
|
||||
if err := db.LockTaskEvidenceTx(tx, in.TaskID); err != nil {
|
||||
return err
|
||||
}
|
||||
prepared, err := s.prepare(ctx, refs)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
out, err = db.RecordFindingTx(tx, in, prepared)
|
||||
return err
|
||||
})
|
||||
return
|
||||
}
|
||||
|
||||
func (s *Store) Bind(ctx context.Context, findingID int64, refs []db.TrafficRef) (out *db.FindingTraffic, err error) {
|
||||
err = s.DB.WithEvidenceTx(ctx, func(tx *sql.Tx) error {
|
||||
if err := db.LockFindingEvidenceTx(tx, findingID, nil); err != nil {
|
||||
return err
|
||||
}
|
||||
prepared, err := s.prepare(ctx, refs)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if err = db.AddFindingTrafficTx(tx, findingID, prepared); err != nil {
|
||||
return err
|
||||
}
|
||||
out, err = db.FindingTrafficTx(tx, findingID)
|
||||
return err
|
||||
})
|
||||
return
|
||||
}
|
||||
|
||||
func (s *Store) WithBinding(ctx context.Context, findingID, bindingID int64, fn func(db.FindingTrafficBinding) error) error {
|
||||
return s.DB.WithEvidenceTx(ctx, func(tx *sql.Tx) error {
|
||||
list, err := db.FindingTrafficTx(tx, findingID)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
for _, b := range list.Bindings {
|
||||
if b.ID == bindingID {
|
||||
return fn(b)
|
||||
}
|
||||
}
|
||||
return db.ErrEvidenceNotFound
|
||||
})
|
||||
}
|
||||
|
||||
// OpenBody must be called while holding WithEvidenceTx/WithBinding, and the
|
||||
// returned file closed before that callback returns.
|
||||
func (s *Store) OpenBody(snapshot db.TrafficEvidenceSnapshot, side string) (*os.File, int64, error) {
|
||||
hash, length := snapshot.ReqHash, snapshot.ReqLen
|
||||
if side == "response" {
|
||||
hash, length = snapshot.RespHash, snapshot.RespLen
|
||||
} else if side != "request" {
|
||||
return nil, 0, errors.New("side 必须为 request 或 response")
|
||||
}
|
||||
path, err := hashPath(s.Dir, hash)
|
||||
if err != nil {
|
||||
return nil, 0, err
|
||||
}
|
||||
if err = verifyFile(path, hash, length); err != nil {
|
||||
return nil, 0, err
|
||||
}
|
||||
f, err := os.Open(path)
|
||||
return f, length, err
|
||||
}
|
||||
|
||||
// CopySnapshots is used by both report downloads and portable task archives.
|
||||
// The destination owns real copies, never links into either disposable store.
|
||||
func (s *Store) CopySnapshots(ctx context.Context, snapshots []db.TrafficEvidenceSnapshot, dest string) error {
|
||||
return s.DB.WithEvidenceTx(ctx, func(*sql.Tx) error { return s.copySnapshots(snapshots, dest) })
|
||||
}
|
||||
|
||||
func (s *Store) copySnapshots(snapshots []db.TrafficEvidenceSnapshot, dest string) error {
|
||||
for _, v := range snapshots {
|
||||
for _, side := range []string{"request", "response"} {
|
||||
if err := func() error {
|
||||
f, length, err := s.OpenBody(v, side)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer f.Close()
|
||||
hash := v.ReqHash
|
||||
if side == "response" {
|
||||
hash = v.RespHash
|
||||
}
|
||||
target, err := hashPath(dest, hash)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if err = os.MkdirAll(filepath.Dir(target), 0o700); err != nil {
|
||||
return err
|
||||
}
|
||||
out, err := os.OpenFile(target, os.O_CREATE|os.O_WRONLY|os.O_EXCL, 0o600)
|
||||
if os.IsExist(err) {
|
||||
return verifyFile(target, hash, length)
|
||||
}
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
_, copyErr := io.Copy(out, f)
|
||||
syncErr := out.Sync()
|
||||
closeErr := out.Close()
|
||||
return errors.Join(copyErr, syncErr, closeErr)
|
||||
}(); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (s *Store) InstallSnapshots(ctx context.Context, snapshots []db.TrafficEvidenceSnapshot, source string) error {
|
||||
return s.WithInstalledSnapshots(ctx, snapshots, source, func() error { return nil })
|
||||
}
|
||||
|
||||
// WithInstalledSnapshots pins installed bodies until the metadata restore finishes.
|
||||
// The callback must not acquire another evidence advisory lock.
|
||||
func (s *Store) WithInstalledSnapshots(ctx context.Context, snapshots []db.TrafficEvidenceSnapshot, source string, restore func() error) error {
|
||||
return s.DB.WithEvidenceTx(ctx, func(*sql.Tx) error {
|
||||
for _, v := range snapshots {
|
||||
if v.ID != db.TrafficSnapshotID(v) {
|
||||
return errors.New("归档证据快照元数据哈希不匹配")
|
||||
}
|
||||
for _, body := range []struct {
|
||||
hash string
|
||||
length int64
|
||||
}{{v.ReqHash, v.ReqLen}, {v.RespHash, v.RespLen}} {
|
||||
if err := func() error {
|
||||
path, err := hashPath(source, body.hash)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
f, err := os.Open(path)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer f.Close()
|
||||
_, err = s.writeBody(f, body.length, body.hash)
|
||||
return err
|
||||
}(); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
}
|
||||
return restore()
|
||||
})
|
||||
}
|
||||
|
||||
func (s *Store) Collect(ctx context.Context, now time.Time) error {
|
||||
return s.DB.WithEvidenceTx(ctx, func(tx *sql.Tx) error {
|
||||
if _, err := tx.Exec(`UPDATE traffic_evidence_snapshots s SET unreferenced_at=$1
|
||||
WHERE unreferenced_at IS NULL AND NOT EXISTS(SELECT 1 FROM finding_traffic_bindings b WHERE b.snapshot_id=s.id)`, now); err != nil {
|
||||
return err
|
||||
}
|
||||
if _, err := tx.Exec(`DELETE FROM traffic_evidence_snapshots s WHERE unreferenced_at<$1
|
||||
AND NOT EXISTS(SELECT 1 FROM finding_traffic_bindings b WHERE b.snapshot_id=s.id)`, now.Add(-24*time.Hour)); err != nil {
|
||||
return err
|
||||
}
|
||||
rows, err := tx.Query(`SELECT req_hash FROM traffic_evidence_snapshots UNION SELECT resp_hash FROM traffic_evidence_snapshots`)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
refs := map[string]bool{}
|
||||
for rows.Next() {
|
||||
var hash string
|
||||
if err := rows.Scan(&hash); err != nil {
|
||||
rows.Close()
|
||||
return err
|
||||
}
|
||||
refs[hash] = true
|
||||
}
|
||||
err = rows.Err()
|
||||
rows.Close()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
for _, root := range []string{filepath.Join(s.Dir, "blobs"), filepath.Join(s.Dir, ".staging")} {
|
||||
err := filepath.WalkDir(root, func(path string, entry fs.DirEntry, err error) error {
|
||||
if os.IsNotExist(err) {
|
||||
return nil
|
||||
}
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if entry.IsDir() {
|
||||
return nil
|
||||
}
|
||||
if refs[strings.TrimSuffix(entry.Name(), ".bin")] {
|
||||
return nil
|
||||
}
|
||||
info, err := entry.Info()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if now.Sub(info.ModTime()) < 24*time.Hour {
|
||||
return nil
|
||||
}
|
||||
return os.Remove(path)
|
||||
})
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
return nil
|
||||
})
|
||||
}
|
||||
|
||||
func (s *Store) RunGC(ctx context.Context) {
|
||||
timer := time.NewTicker(time.Hour)
|
||||
defer timer.Stop()
|
||||
for {
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return
|
||||
case now := <-timer.C:
|
||||
if err := s.Collect(ctx, now); err != nil && !errors.Is(err, context.Canceled) {
|
||||
log.Printf("[evidence] cleanup failed: %v", err)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,333 @@
|
||||
package evidence
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"crypto/sha256"
|
||||
"encoding/hex"
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"sync"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/Autumn-27/artex/db"
|
||||
"github.com/Autumn-27/artex/traffic"
|
||||
)
|
||||
|
||||
func evidenceFixture(t *testing.T) (*Store, db.RecordFindingInput, string) {
|
||||
t.Helper()
|
||||
dsn := os.Getenv("ARTEX_PG_DSN")
|
||||
if dsn == "" {
|
||||
t.Skip("ARTEX_PG_DSN is required for evidence integration tests")
|
||||
}
|
||||
pg, err := db.Open(dsn)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
t.Cleanup(func() { pg.Close() })
|
||||
task, err := pg.CreateTask("evidence integration "+t.Name(), "local fixtures", nil, 0, 0)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
t.Cleanup(func() { pg.DeleteFindingsByTask(task.ID); pg.DeleteTask(task.ID) })
|
||||
dir := t.TempDir()
|
||||
tr, err := traffic.Open(filepath.Join(dir, "traffic"), ":0")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
t.Cleanup(func() { tr.Close() })
|
||||
store := New(pg, tr, filepath.Join(dir, "evidence"))
|
||||
in := db.RecordFindingInput{TaskID: task.ID, ExplorationID: task.ExplorationID, Worker: "test", VulnClass: "TEST", Name: "Evidence fixture", Severity: "low", Summary: "local test"}
|
||||
return store, in, dir
|
||||
}
|
||||
|
||||
func seedExchange(t *testing.T, s *Store, id string, body []byte, spill bool) {
|
||||
t.Helper()
|
||||
var blob any
|
||||
inline := body
|
||||
if spill {
|
||||
sum := sha256.Sum256(body)
|
||||
hash := hex.EncodeToString(sum[:])
|
||||
blob = hash
|
||||
path := filepath.Join(filepath.Dir(s.Dir), "traffic", "_blobs", "sha256", hash[:2], hash+".bin")
|
||||
if err := os.MkdirAll(filepath.Dir(path), 0o700); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := os.WriteFile(path, body, 0o600); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
inline = []byte("TRUNCATED PREVIEW")
|
||||
}
|
||||
_, err := s.Traffic.DB().Exec(`INSERT INTO exchanges(id,ts,host,method,url_template,url,status,content_type,req_len,resp_len,path) VALUES(?,?,'fixture.local','POST','/test',?,200,'application/octet-stream',2,?,'')`, id, time.Now().Unix(), "https://fixture.local/"+id, len(body))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
_, err = s.Traffic.DB().Exec(`INSERT INTO exchange_bodies(id,req_head,req_body,resp_head,resp_body,resp_blob) VALUES(?,?,?, ?,?,?)`, id, "POST /test HTTP/1.1\nHost: fixture.local\n", []byte("{}"), "HTTP 200\nContent-Type: application/octet-stream\n", inline, blob)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestEvidenceBindingLifecycle(t *testing.T) {
|
||||
s, in, _ := evidenceFixture(t)
|
||||
ctx := context.Background()
|
||||
large := bytes.Repeat([]byte{0, 1, 2, 255, 'a', 'b'}, 180000)
|
||||
for i := 0; i < 3; i++ {
|
||||
seedExchange(t, s, fmt.Sprint(i), large, i == 2)
|
||||
}
|
||||
refs := []db.TrafficRef{{TrafficID: "0", Role: "baseline", Note: "normal"}, {TrafficID: "1", Role: "proof", Note: "proof"}, {TrafficID: "2", Role: "verification", Note: "binary"}}
|
||||
r, err := s.Record(ctx, in, refs)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(r.Traffic.Bindings) != 3 || r.Traffic.Version != 1 {
|
||||
t.Fatalf("record=%+v", r)
|
||||
}
|
||||
again, err := s.Bind(ctx, r.FindingID, []db.TrafficRef{{TrafficID: "1", Note: "must not overwrite"}})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(again.Bindings) != 3 || again.Version != 1 || again.Bindings[1].Note != "proof" {
|
||||
t.Fatalf("duplicate changed bindings: %+v", again)
|
||||
}
|
||||
other, err := s.Record(ctx, in, []db.TrafficRef{refs[2]})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if other.Traffic.Bindings[0].SnapshotID != r.Traffic.Bindings[2].SnapshotID {
|
||||
t.Fatal("snapshot was not shared")
|
||||
}
|
||||
if _, err = s.Traffic.DeleteHost("fixture.local"); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
for _, binding := range r.Traffic.Bindings {
|
||||
if err = s.WithBinding(ctx, r.FindingID, binding.ID, func(b db.FindingTrafficBinding) error {
|
||||
f, _, err := s.OpenBody(b.Snapshot, "response")
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer f.Close()
|
||||
got, err := io.ReadAll(f)
|
||||
if !bytes.Equal(got, large) {
|
||||
t.Fatal("body was truncated/changed")
|
||||
}
|
||||
return err
|
||||
}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
v := again.Version
|
||||
if _, err = s.DB.SetFindingReportVersionByNodeID(ctx, r.NodeID, "report", &v); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
note := "updated"
|
||||
if err = s.DB.EditFindingTraffic(ctx, r.FindingID, again.Bindings[0].ID, v, nil, ¬e, false, nil); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if _, err = s.DB.SetFindingReportVersionByNodeID(ctx, r.NodeID, "stale report", &v); !errors.Is(err, db.ErrEvidenceConflict) {
|
||||
t.Fatalf("stale report: %v", err)
|
||||
}
|
||||
f, err := s.DB.GetFinding(r.FindingID)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if f.Report != "report" || f.EvidenceVersion == f.ReportEvidenceVersion || f.TrafficCount != 3 {
|
||||
t.Fatalf("finding versions: %+v", f)
|
||||
}
|
||||
if err = s.DB.EditFindingTraffic(ctx, r.FindingID, again.Bindings[0].ID, v, nil, nil, true, nil); !errors.Is(err, db.ErrEvidenceConflict) {
|
||||
t.Fatalf("stale delete: %v", err)
|
||||
}
|
||||
if _, err = s.DB.SetFindingReportByNodeID(r.NodeID, "legacy report"); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
f, err = s.DB.GetFinding(r.FindingID)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if f.ReportEvidenceVersion != -1 {
|
||||
t.Fatal("legacy write claimed current evidence")
|
||||
}
|
||||
if _, err = s.DB.DeleteFinding(r.FindingID); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err = s.Collect(ctx, time.Now().Add(48*time.Hour)); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err = s.WithBinding(ctx, other.FindingID, other.Traffic.Bindings[0].ID, func(b db.FindingTrafficBinding) error {
|
||||
f, _, err := s.OpenBody(b.Snapshot, "response")
|
||||
if f != nil {
|
||||
f.Close()
|
||||
}
|
||||
return err
|
||||
}); err != nil {
|
||||
t.Fatal("GC removed shared body:", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestEvidenceAtomicFailuresAndLegacy(t *testing.T) {
|
||||
s, in, dir := evidenceFixture(t)
|
||||
ctx := context.Background()
|
||||
seedExchange(t, s, "valid", []byte("OK"), false)
|
||||
counts := func() (findings, nodes int) {
|
||||
s.DB.QueryRow(`SELECT count(*) FROM findings WHERE task_id=$1`, in.TaskID).Scan(&findings)
|
||||
s.DB.QueryRow(`SELECT count(*) FROM exploration_nodes WHERE exploration_id=$1 AND kind='finding'`, in.ExplorationID).Scan(&nodes)
|
||||
return
|
||||
}
|
||||
for _, kind := range []string{"missing_id", "missing_body", "disk_failure", "database_failure"} {
|
||||
t.Run(kind, func(t *testing.T) {
|
||||
input := in
|
||||
refs := []db.TrafficRef{{TrafficID: "valid"}, {TrafficID: "absent"}}
|
||||
local := *s
|
||||
switch kind {
|
||||
case "missing_body":
|
||||
seedExchange(t, s, "broken", []byte("more than preview"), false)
|
||||
s.Traffic.DB().Exec(`UPDATE exchange_bodies SET resp_body=? WHERE id='broken'`, []byte("x"))
|
||||
refs = []db.TrafficRef{{TrafficID: "valid"}, {TrafficID: "broken"}}
|
||||
case "disk_failure":
|
||||
local.Dir = filepath.Join(dir, "file-not-dir")
|
||||
if err := os.WriteFile(local.Dir, []byte("x"), 0o600); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
refs = refs[:1]
|
||||
case "database_failure":
|
||||
input.AssetIDs = []int64{9223372036854775807}
|
||||
refs = refs[:1]
|
||||
}
|
||||
if _, err := local.Record(ctx, input, refs); err == nil {
|
||||
t.Fatal("expected failure")
|
||||
}
|
||||
if f, n := counts(); f != 0 || n != 0 {
|
||||
t.Fatalf("partial record: findings=%d nodes=%d", f, n)
|
||||
}
|
||||
})
|
||||
}
|
||||
legacy := filepath.Join(dir, "traffic", "legacy")
|
||||
os.MkdirAll(legacy, 0o700)
|
||||
os.WriteFile(filepath.Join(legacy, "request.http"), []byte("GET / HTTP/1.1\nHost: old.local\n"), 0o600)
|
||||
os.WriteFile(filepath.Join(legacy, "response.http"), []byte("HTTP 200\n\nlegacy body"), 0o600)
|
||||
_, err := s.Traffic.DB().Exec(`INSERT INTO exchanges(id,ts,host,method,url_template,url,status,content_type,req_len,resp_len,path) VALUES('old',1,'old.local','GET','/','http://old.local/',200,'text/plain',0,11,'legacy')`)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if _, err = s.Record(ctx, in, []db.TrafficRef{{TrafficID: "old"}}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if _, err = s.Record(ctx, in, nil); err != nil {
|
||||
t.Fatal("legacy no-traffic report failed:", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestEvidenceConcurrentBindAndDelete(t *testing.T) {
|
||||
s, in, _ := evidenceFixture(t)
|
||||
ctx := context.Background()
|
||||
seedExchange(t, s, "race", bytes.Repeat([]byte("large"), 90000), true)
|
||||
r, err := s.Record(ctx, in, nil)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
var wg sync.WaitGroup
|
||||
errs := make(chan error, 10)
|
||||
for i := 0; i < 8; i++ {
|
||||
wg.Add(1)
|
||||
go func() {
|
||||
defer wg.Done()
|
||||
_, err := s.Bind(ctx, r.FindingID, []db.TrafficRef{{TrafficID: "race"}})
|
||||
errs <- err
|
||||
}()
|
||||
}
|
||||
wg.Wait()
|
||||
close(errs)
|
||||
for err := range errs {
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
list, err := s.DB.GetFindingTraffic(ctx, r.FindingID)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(list.Bindings) != 1 || list.Version != 1 {
|
||||
t.Fatalf("concurrent duplicates: %+v", list)
|
||||
}
|
||||
// Race a new capture copy with source deletion: either the full copy commits
|
||||
// or nothing does. An already-bound snapshot remains readable in both cases.
|
||||
wg.Add(2)
|
||||
go func() { defer wg.Done(); s.Bind(ctx, r.FindingID, []db.TrafficRef{{TrafficID: "race"}}) }()
|
||||
go func() { defer wg.Done(); s.Traffic.DeleteHost("fixture.local") }()
|
||||
wg.Wait()
|
||||
if err = s.WithBinding(ctx, r.FindingID, list.Bindings[0].ID, func(b db.FindingTrafficBinding) error {
|
||||
f, _, err := s.OpenBody(b.Snapshot, "response")
|
||||
if f != nil {
|
||||
f.Close()
|
||||
}
|
||||
return err
|
||||
}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestEvidenceGCGraceAndActiveRestore(t *testing.T) {
|
||||
s, in, _ := evidenceFixture(t)
|
||||
ctx := context.Background()
|
||||
seedExchange(t, s, "gc", []byte("unique unreferenced gc body"), false)
|
||||
f, err := s.Record(ctx, in, []db.TrafficRef{{TrafficID: "gc"}})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
snap := f.Traffic.Bindings[0].Snapshot
|
||||
archive := t.TempDir()
|
||||
if err = s.CopySnapshots(ctx, []db.TrafficEvidenceSnapshot{snap}, archive); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if _, err = s.DB.DeleteFinding(f.FindingID); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
now := time.Now()
|
||||
if err = s.Collect(ctx, now); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
path, err := hashPath(s.Dir, snap.RespHash)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err = s.Collect(ctx, now.Add(23*time.Hour)); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if _, err = os.Stat(path); err != nil {
|
||||
t.Fatal("GC grace ignored", err)
|
||||
}
|
||||
if err = s.Collect(ctx, now.Add(25*time.Hour)); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if _, err = os.Stat(path); !os.IsNotExist(err) {
|
||||
t.Fatal("orphan was not collected", err)
|
||||
}
|
||||
entered, release, finished := make(chan struct{}), make(chan struct{}), make(chan error, 1)
|
||||
go func() {
|
||||
finished <- s.WithInstalledSnapshots(ctx, []db.TrafficEvidenceSnapshot{snap}, archive, func() error { close(entered); <-release; return nil })
|
||||
}()
|
||||
<-entered
|
||||
short, cancel := context.WithTimeout(ctx, 100*time.Millisecond)
|
||||
err = s.Collect(short, now.Add(72*time.Hour))
|
||||
cancel()
|
||||
close(release)
|
||||
if err == nil {
|
||||
t.Fatal("GC entered during active restore")
|
||||
}
|
||||
if err = <-finished; err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if _, err = os.Stat(path); err != nil {
|
||||
t.Fatal("active restore lost body", err)
|
||||
}
|
||||
// A portable package cannot substitute metadata while retaining its old ID.
|
||||
bad := snap
|
||||
bad.URL = "http://tampered.local/"
|
||||
if err = s.InstallSnapshots(ctx, []db.TrafficEvidenceSnapshot{bad}, archive); err == nil {
|
||||
t.Fatal("accepted tampered snapshot")
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,35 @@
|
||||
package evidence
|
||||
|
||||
import (
|
||||
"context"
|
||||
"os"
|
||||
"testing"
|
||||
|
||||
"github.com/Autumn-27/artex/db"
|
||||
)
|
||||
|
||||
// Initialize an explicitly configured fresh database before taking the same
|
||||
// suite lock as db, agent and server. Hold it on one pinned connection.
|
||||
func TestMain(m *testing.M) {
|
||||
if os.Getenv("ARTEX_PG_DSN") == "" {
|
||||
os.Exit(m.Run())
|
||||
}
|
||||
os.Exit(runEvidenceSuite(m))
|
||||
}
|
||||
func runEvidenceSuite(m *testing.M) int {
|
||||
pg, err := db.Open(os.Getenv("ARTEX_PG_DSN"))
|
||||
if err != nil {
|
||||
panic(err)
|
||||
}
|
||||
defer pg.Close()
|
||||
conn, err := pg.Conn(context.Background())
|
||||
if err != nil {
|
||||
panic(err)
|
||||
}
|
||||
defer conn.Close()
|
||||
if _, err = conn.ExecContext(context.Background(), `SELECT pg_advisory_lock(7337741002)`); err != nil {
|
||||
panic(err)
|
||||
}
|
||||
defer conn.ExecContext(context.Background(), `SELECT pg_advisory_unlock(7337741002)`)
|
||||
return m.Run()
|
||||
}
|
||||
+36
-1
@@ -84,6 +84,7 @@ func FindingsMarkdown(fs []*db.DBFinding, generatedAt time.Time) string {
|
||||
b.WriteString(rep)
|
||||
b.WriteString("\n\n")
|
||||
}
|
||||
b.WriteString(findingTrafficMarkdown(f, false))
|
||||
b.WriteString("---\n\n")
|
||||
}
|
||||
return b.String()
|
||||
@@ -114,6 +115,7 @@ func SingleFindingMarkdown(f *db.DBFinding, generatedAt time.Time) string {
|
||||
b.WriteString(rep)
|
||||
b.WriteString("\n")
|
||||
}
|
||||
b.WriteString(findingTrafficMarkdown(f, true))
|
||||
return b.String()
|
||||
}
|
||||
|
||||
@@ -147,7 +149,7 @@ func FindingsCSV(fs []*db.DBFinding) []byte {
|
||||
var buf bytes.Buffer
|
||||
buf.WriteString("\xEF\xBB\xBF") // UTF-8 BOM
|
||||
w := csv.NewWriter(&buf)
|
||||
_ = w.Write([]string{"ID", "名称", "类别", "严重等级", "状态", "所属任务", "发现时间", "概述"})
|
||||
_ = w.Write([]string{"ID", "名称", "类别", "严重等级", "状态", "所属任务", "发现时间", "概述", "流量证据数量", "流量证据ID"})
|
||||
for _, f := range items {
|
||||
_ = w.Write([]string{
|
||||
fmt.Sprintf("%d", f.ID),
|
||||
@@ -158,8 +160,41 @@ func FindingsCSV(fs []*db.DBFinding) []byte {
|
||||
f.TaskDescription,
|
||||
f.CreatedAt.Format("2006-01-02 15:04:05"),
|
||||
strings.TrimSpace(f.Summary),
|
||||
fmt.Sprint(len(f.TrafficBindings)), findingTrafficIDs(f),
|
||||
})
|
||||
}
|
||||
w.Flush()
|
||||
return buf.Bytes()
|
||||
}
|
||||
|
||||
func findingTrafficIDs(f *db.DBFinding) string {
|
||||
ids := make([]string, 0, len(f.TrafficBindings))
|
||||
for _, b := range f.TrafficBindings {
|
||||
ids = append(ids, fmt.Sprint(b.ID))
|
||||
}
|
||||
return strings.Join(ids, ",")
|
||||
}
|
||||
|
||||
func findingTrafficMarkdown(f *db.DBFinding, attachments bool) string {
|
||||
stale := f.Report != "" && f.EvidenceVersion != f.ReportEvidenceVersion
|
||||
if len(f.TrafficBindings) == 0 && !stale {
|
||||
return ""
|
||||
}
|
||||
var out strings.Builder
|
||||
out.WriteString("\n## 关联流量证据\n\n")
|
||||
fmt.Fprintf(&out, "证据版本:%d;绑定数量:%d。\n\n", f.EvidenceVersion, len(f.TrafficBindings))
|
||||
if stale {
|
||||
out.WriteString("证据已变更,详细报告待更新。\n\n")
|
||||
}
|
||||
for i, b := range f.TrafficBindings {
|
||||
fmt.Fprintf(&out, "%d. **证据 #%d · %s** — `%s %s`,状态码 %d\n", i+1, b.ID, b.Role, b.Snapshot.Method, strings.ReplaceAll(b.Snapshot.URL, "`", "%60"), b.Snapshot.Status)
|
||||
if b.Note != "" {
|
||||
fmt.Fprintf(&out, " %s\n", strings.ReplaceAll(b.Note, "\n", "\n "))
|
||||
}
|
||||
if attachments {
|
||||
fmt.Fprintf(&out, " [请求报文](evidence/%d/%d/request.http) · [响应报文](evidence/%d/%d/response.http)\n", f.ID, b.ID, f.ID, b.ID)
|
||||
}
|
||||
}
|
||||
out.WriteString("\n")
|
||||
return out.String()
|
||||
}
|
||||
@@ -210,6 +210,7 @@ func seedPrompts(pg *db.DB) {
|
||||
// flag and, if kept, wrapped so the model sees the DB-overridden description/schema
|
||||
// and缺省入参 get injected. MCP/skill/host tools have no row and pass through.
|
||||
func wireTools(pg *db.DB, domainReg map[string]actool.CoreTool) {
|
||||
agent.FindingTrafficBindingEnabled = func() bool { return pg.GetBool(settingAgentTrafficBinding, false) }
|
||||
// Seed the built-in domain tools (first-insert only; DO NOTHING preserves edits).
|
||||
// No startup prune: rows we didn't seed are left alone so future user-defined
|
||||
// custom tools (system=false, added via the UI) survive restarts.
|
||||
|
||||
+11
-2
@@ -292,6 +292,12 @@ func coverageAssetRefDTO(ref db.AssetRef) CoverageAssetRefDTO {
|
||||
// ---- Finding (frontend "Finding") ----
|
||||
|
||||
type FindingDTO struct {
|
||||
TrafficCount int `json:"traffic_count"`
|
||||
EvidenceVersion int64 `json:"evidence_version"`
|
||||
ReportEvidenceVersion int64 `json:"report_evidence_version"`
|
||||
ReportStale bool `json:"report_stale"`
|
||||
TrafficBindings []db.FindingTrafficBinding `json:"traffic_bindings,omitempty"`
|
||||
|
||||
ID string `json:"id"`
|
||||
FindingID string `json:"finding_id,omitempty"` // standalone findings-table id — the handle for status updates
|
||||
VulnClass string `json:"vulnclass"`
|
||||
@@ -399,6 +405,7 @@ func findingDTOsForOwner(taskID, description string, in []*db.Node, meta map[int
|
||||
if m, ok := meta[n.ID]; ok {
|
||||
d.FindingID = i64s(m.ID)
|
||||
d.Status = m.Status
|
||||
d.TrafficCount = m.TrafficCount
|
||||
d.Assets = findingAssetDTOs(m.AssetIDs, assets)
|
||||
}
|
||||
out = append(out, d)
|
||||
@@ -426,8 +433,10 @@ func findingFromDB(f *db.DBFinding, assets map[int64]*db.Asset) FindingDTO {
|
||||
status = db.FindingPending
|
||||
}
|
||||
d := FindingDTO{
|
||||
ID: i64s(f.ID),
|
||||
FindingID: i64s(f.ID),
|
||||
ID: i64s(f.ID),
|
||||
FindingID: i64s(f.ID),
|
||||
TrafficCount: f.TrafficCount, EvidenceVersion: f.EvidenceVersion, ReportEvidenceVersion: f.ReportEvidenceVersion,
|
||||
ReportStale: f.Report != "" && f.EvidenceVersion != f.ReportEvidenceVersion, TrafficBindings: f.TrafficBindings,
|
||||
VulnClass: f.VulnClass,
|
||||
Name: f.Name,
|
||||
Severity: f.Severity,
|
||||
|
||||
@@ -0,0 +1,412 @@
|
||||
package server
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
"log"
|
||||
"net/http"
|
||||
"path/filepath"
|
||||
"strconv"
|
||||
"unicode/utf8"
|
||||
|
||||
"github.com/Autumn-27/artex/agent"
|
||||
"github.com/Autumn-27/artex/db"
|
||||
"github.com/Autumn-27/artex/evidence"
|
||||
actool "github.com/Autumn-27/norma/tool"
|
||||
)
|
||||
|
||||
func (s *Server) evidenceStore() *evidence.Store {
|
||||
return evidence.New(s.m.pg, s.m.traffic, filepath.Join(s.m.dir, "evidence"))
|
||||
}
|
||||
|
||||
// Add only new optional properties; preserve edited descriptions, existing
|
||||
// properties, agent bindings and disabled flags. The one-time flag also keeps
|
||||
// subsequent user unbinding of the evidence reader intact.
|
||||
func (s *Server) seedFindingTrafficTools() {
|
||||
const flag = "finding_traffic_tools_v1"
|
||||
if v, _, _ := s.m.pg.GetSetting(flag); v == "true" {
|
||||
return
|
||||
}
|
||||
for _, key := range []string{"report_finding", "update_finding_report"} {
|
||||
var schema any
|
||||
if key == "update_finding_report" {
|
||||
schema = s.toolUpdateFindingReport().InputSchema()
|
||||
} else {
|
||||
for _, seed := range agent.BuiltinToolSeeds() {
|
||||
if seed.Key == key {
|
||||
schema = seed.Schema
|
||||
break
|
||||
}
|
||||
}
|
||||
}
|
||||
raw, err := json.Marshal(schema)
|
||||
if err != nil {
|
||||
log.Printf("[evidence] tool schema: %v", err)
|
||||
return
|
||||
}
|
||||
var obj map[string]json.RawMessage
|
||||
if err = json.Unmarshal(raw, &obj); err != nil {
|
||||
return
|
||||
}
|
||||
// BuiltinToolSeeds stores Schema as RawMessage; both forms marshal as JSON.
|
||||
var properties map[string]json.RawMessage
|
||||
if err = json.Unmarshal(obj["properties"], &properties); err != nil {
|
||||
return
|
||||
}
|
||||
name := "traffic_refs"
|
||||
if key == "update_finding_report" {
|
||||
name = "evidence_version"
|
||||
}
|
||||
if len(properties[name]) == 0 {
|
||||
return
|
||||
}
|
||||
_, err = s.m.pg.Exec(`UPDATE tools SET schema=jsonb_set(schema,ARRAY['properties',$2::text],$3::jsonb,true),updated_at=now()
|
||||
WHERE key=$1 AND system AND NOT(COALESCE(schema->'properties','{}'::jsonb) ? $2)`, key, name, string(properties[name]))
|
||||
if err != nil {
|
||||
log.Printf("[evidence] upgrade tool %s: %v", key, err)
|
||||
return
|
||||
}
|
||||
}
|
||||
if reporter, _ := s.m.pg.GetAgentByKey("reporter"); reporter != nil {
|
||||
if err := s.m.pg.AddAgentToToolBinding("reporter", []string{"get_finding_traffic"}); err != nil {
|
||||
return
|
||||
}
|
||||
}
|
||||
_ = s.m.pg.SetSetting(flag, "true")
|
||||
}
|
||||
|
||||
func (s *Server) registerFindingTraffic(mux *http.ServeMux) {
|
||||
base := "/api/exploration/findings/{id}/traffic"
|
||||
mux.HandleFunc("GET "+base, s.getFindingTraffic)
|
||||
mux.HandleFunc("POST "+base, s.bindFindingTraffic)
|
||||
mux.HandleFunc("PATCH "+base+"/{binding_id}", s.editFindingTraffic)
|
||||
mux.HandleFunc("DELETE "+base+"/{binding_id}", s.editFindingTraffic)
|
||||
mux.HandleFunc("PUT "+base+"/order", s.editFindingTraffic)
|
||||
mux.HandleFunc("GET "+base+"/{binding_id}", s.getFindingTrafficDetail)
|
||||
mux.HandleFunc("GET "+base+"/{binding_id}/body", s.getFindingTrafficBody)
|
||||
}
|
||||
|
||||
func evidenceError(w http.ResponseWriter, err error) {
|
||||
status := http.StatusUnprocessableEntity
|
||||
if errors.Is(err, db.ErrEvidenceConflict) || errors.Is(err, db.ErrTaskArchiveState) {
|
||||
status = http.StatusConflict
|
||||
}
|
||||
if errors.Is(err, db.ErrFindingNotFound) || errors.Is(err, db.ErrEvidenceNotFound) {
|
||||
status = http.StatusNotFound
|
||||
}
|
||||
writeErr(w, status, err.Error())
|
||||
}
|
||||
|
||||
func (s *Server) findingTrafficAccess(w http.ResponseWriter, r *http.Request, write bool) (int64, bool) {
|
||||
id, err := strconv.ParseInt(r.PathValue("id"), 10, 64)
|
||||
if err != nil || id <= 0 {
|
||||
writeErr(w, 400, "invalid finding id")
|
||||
return 0, false
|
||||
}
|
||||
f, err := s.m.pg.GetFinding(id)
|
||||
if err != nil {
|
||||
writeErr(w, 500, err.Error())
|
||||
return 0, false
|
||||
}
|
||||
if f == nil {
|
||||
writeErr(w, 404, "finding not found")
|
||||
return 0, false
|
||||
}
|
||||
if taskID := r.URL.Query().Get("context_task"); taskID != "" {
|
||||
task := s.m.ResolveTask(taskID)
|
||||
if task == nil {
|
||||
writeErr(w, 404, "context task not found")
|
||||
return 0, false
|
||||
}
|
||||
_, inherited, allowed := findingProvenanceInTask(task, f.TaskID)
|
||||
if !allowed {
|
||||
writeErr(w, 404, "finding not available in task context")
|
||||
return 0, false
|
||||
}
|
||||
if write && inherited {
|
||||
writeErr(w, 403, "继承漏洞只读,请在来源任务中修改")
|
||||
return 0, false
|
||||
}
|
||||
}
|
||||
return id, true
|
||||
}
|
||||
|
||||
func trafficSummary(in *db.FindingTraffic) *db.FindingTraffic {
|
||||
out := *in
|
||||
out.Bindings = append([]db.FindingTrafficBinding{}, in.Bindings...)
|
||||
for i := range out.Bindings {
|
||||
out.Bindings[i].Snapshot.ReqHead = ""
|
||||
out.Bindings[i].Snapshot.RespHead = ""
|
||||
}
|
||||
return &out
|
||||
}
|
||||
|
||||
func (s *Server) getFindingTraffic(w http.ResponseWriter, r *http.Request) {
|
||||
id, ok := s.findingTrafficAccess(w, r, false)
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
out, err := s.m.pg.GetFindingTraffic(r.Context(), id)
|
||||
if err != nil {
|
||||
evidenceError(w, err)
|
||||
return
|
||||
}
|
||||
writeJSON(w, 200, trafficSummary(out))
|
||||
}
|
||||
|
||||
func (s *Server) bindFindingTraffic(w http.ResponseWriter, r *http.Request) {
|
||||
id, ok := s.findingTrafficAccess(w, r, true)
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
var body struct {
|
||||
Refs []db.TrafficRef `json:"traffic_refs"`
|
||||
}
|
||||
if err := json.NewDecoder(http.MaxBytesReader(w, r.Body, 1<<20)).Decode(&body); err != nil {
|
||||
writeErr(w, 400, "invalid body")
|
||||
return
|
||||
}
|
||||
if len(body.Refs) == 0 {
|
||||
writeErr(w, 400, "请选择流量")
|
||||
return
|
||||
}
|
||||
out, err := s.evidenceStore().Bind(r.Context(), id, body.Refs)
|
||||
if err != nil {
|
||||
evidenceError(w, err)
|
||||
return
|
||||
}
|
||||
writeJSON(w, 200, trafficSummary(out))
|
||||
}
|
||||
|
||||
func (s *Server) editFindingTraffic(w http.ResponseWriter, r *http.Request) {
|
||||
id, ok := s.findingTrafficAccess(w, r, true)
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
var body struct {
|
||||
Version *int64 `json:"version"`
|
||||
Role *string `json:"role"`
|
||||
Note *string `json:"note"`
|
||||
Order []string `json:"binding_ids"`
|
||||
}
|
||||
if err := json.NewDecoder(http.MaxBytesReader(w, r.Body, 1<<20)).Decode(&body); err != nil || body.Version == nil {
|
||||
writeErr(w, 400, "version 和有效请求体必填")
|
||||
return
|
||||
}
|
||||
var order []int64
|
||||
bindingID := int64(0)
|
||||
if r.Method == http.MethodPut {
|
||||
if body.Order == nil {
|
||||
writeErr(w, 400, "binding_ids 必填")
|
||||
return
|
||||
}
|
||||
order = []int64{}
|
||||
for _, raw := range body.Order {
|
||||
v, err := strconv.ParseInt(raw, 10, 64)
|
||||
if err != nil || v <= 0 {
|
||||
writeErr(w, 400, "invalid binding id")
|
||||
return
|
||||
}
|
||||
order = append(order, v)
|
||||
}
|
||||
} else {
|
||||
var err error
|
||||
bindingID, err = strconv.ParseInt(r.PathValue("binding_id"), 10, 64)
|
||||
if err != nil || bindingID <= 0 {
|
||||
writeErr(w, 400, "invalid binding id")
|
||||
return
|
||||
}
|
||||
}
|
||||
err := s.m.pg.EditFindingTraffic(r.Context(), id, bindingID, *body.Version, body.Role, body.Note, r.Method == http.MethodDelete, order)
|
||||
if err != nil {
|
||||
evidenceError(w, err)
|
||||
return
|
||||
}
|
||||
s.getFindingTraffic(w, r)
|
||||
}
|
||||
|
||||
type evidencePreview struct {
|
||||
Content string `json:"content"`
|
||||
Offset int64 `json:"offset"`
|
||||
Total int64 `json:"total"`
|
||||
NextOffset int64 `json:"next_offset"`
|
||||
Truncated bool `json:"truncated"`
|
||||
Binary bool `json:"binary"`
|
||||
}
|
||||
|
||||
func readEvidencePreview(store *evidence.Store, snapshot db.TrafficEvidenceSnapshot, side string, offset, length int64) (out evidencePreview, err error) {
|
||||
if offset < 0 || length < 0 {
|
||||
return out, errors.New("offset / length 不能为负数")
|
||||
}
|
||||
if length == 0 || length > 8192 {
|
||||
length = 8192
|
||||
}
|
||||
f, total, err := store.OpenBody(snapshot, side)
|
||||
if err != nil {
|
||||
return out, err
|
||||
}
|
||||
defer f.Close()
|
||||
if offset > total {
|
||||
return out, errors.New("offset 超出正文长度")
|
||||
}
|
||||
if _, err = f.Seek(offset, io.SeekStart); err != nil {
|
||||
return out, err
|
||||
}
|
||||
raw, err := io.ReadAll(io.LimitReader(f, length))
|
||||
if err != nil {
|
||||
return out, err
|
||||
}
|
||||
// Leave an incomplete trailing UTF-8 rune for the next page. Explicit byte
|
||||
// offsets are still accepted; complete downloads always retain original bytes.
|
||||
if offset+int64(len(raw)) < total && len(raw) >= utf8.UTFMax {
|
||||
start := len(raw) - 1
|
||||
for start > 0 && raw[start]&0xc0 == 0x80 {
|
||||
start--
|
||||
}
|
||||
if !utf8.FullRune(raw[start:]) {
|
||||
raw = raw[:start]
|
||||
}
|
||||
}
|
||||
out = evidencePreview{Offset: offset, Total: total, NextOffset: offset + int64(len(raw)), Truncated: offset+int64(len(raw)) < total,
|
||||
Binary: bytes.IndexByte(raw, 0) >= 0 || (offset == 0 && !utf8.Valid(raw))}
|
||||
if out.Binary {
|
||||
out.Content = fmt.Sprintf("[二进制正文,%d 字节;请下载查看]", total)
|
||||
} else {
|
||||
out.Content = string(bytes.ToValidUTF8(raw, []byte("�")))
|
||||
}
|
||||
return out, nil
|
||||
}
|
||||
|
||||
func (s *Server) getFindingTrafficDetail(w http.ResponseWriter, r *http.Request) {
|
||||
id, ok := s.findingTrafficAccess(w, r, false)
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
bid, err := strconv.ParseInt(r.PathValue("binding_id"), 10, 64)
|
||||
if err != nil || bid <= 0 {
|
||||
writeErr(w, 400, "invalid binding id")
|
||||
return
|
||||
}
|
||||
store := s.evidenceStore()
|
||||
var result any
|
||||
err = store.WithBinding(r.Context(), id, bid, func(b db.FindingTrafficBinding) error {
|
||||
req, err := readEvidencePreview(store, b.Snapshot, "request", 0, 8192)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
resp, err := readEvidencePreview(store, b.Snapshot, "response", 0, 8192)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
result = map[string]any{"binding": b, "request": req, "response": resp}
|
||||
return nil
|
||||
})
|
||||
if err != nil {
|
||||
evidenceError(w, err)
|
||||
return
|
||||
}
|
||||
writeJSON(w, 200, result)
|
||||
}
|
||||
|
||||
func (s *Server) getFindingTrafficBody(w http.ResponseWriter, r *http.Request) {
|
||||
id, ok := s.findingTrafficAccess(w, r, false)
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
bid, err := strconv.ParseInt(r.PathValue("binding_id"), 10, 64)
|
||||
if err != nil || bid <= 0 {
|
||||
writeErr(w, 400, "invalid binding id")
|
||||
return
|
||||
}
|
||||
side := r.URL.Query().Get("side")
|
||||
store := s.evidenceStore()
|
||||
started := false
|
||||
err = store.WithBinding(r.Context(), id, bid, func(b db.FindingTrafficBinding) error {
|
||||
if r.URL.Query().Get("download") == "1" {
|
||||
f, length, err := store.OpenBody(b.Snapshot, side)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer f.Close()
|
||||
w.Header().Set("Content-Type", "application/octet-stream")
|
||||
w.Header().Set("X-Content-Type-Options", "nosniff")
|
||||
w.Header().Set("Content-Disposition", fmt.Sprintf("attachment; filename=\"evidence-%d-%s.bin\"", bid, side))
|
||||
w.Header().Set("Content-Length", strconv.FormatInt(length, 10))
|
||||
started = true
|
||||
_, err = io.Copy(w, f)
|
||||
return err
|
||||
}
|
||||
offset, length := int64(0), int64(8192)
|
||||
for name, dst := range map[string]*int64{"offset": &offset, "length": &length} {
|
||||
if raw := r.URL.Query().Get(name); raw != "" {
|
||||
v, err := strconv.ParseInt(raw, 10, 64)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
*dst = v
|
||||
}
|
||||
}
|
||||
preview, err := readEvidencePreview(store, b.Snapshot, side, offset, length)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
started = true
|
||||
writeJSON(w, 200, preview)
|
||||
return nil
|
||||
})
|
||||
if err != nil && !started {
|
||||
evidenceError(w, err)
|
||||
}
|
||||
}
|
||||
|
||||
func (s *Server) toolGetFindingTraffic() actool.CoreTool {
|
||||
return roTool("get_finding_traffic", "读取漏洞已绑定的真实流量证据,不依赖捕获开关。finding_id 使用 report_finding JSON 返回的独立漏洞记录 ID(不是第一行的探索节点 ID)。先不传 binding_id 获取清单及 version;空清单是正常情况,TCP 等非 HTTP 漏洞或未采集时仍可依据文字/命令证据编写报告,不强制绑定。有绑定时按 binding_id、side(request/response)、offset 分段读取正文。写报告时将读取的 version 作为 evidence_version 传给 update_finding_report,后者 finding_id 仍使用探索节点 ID。",
|
||||
objSchema(map[string]any{"finding_id": strParam("独立漏洞记录 ID"), "binding_id": strParam("清单里的绑定 ID,省略则返回清单"), "side": strParam("request 或 response,默认 response"), "offset": map[string]any{"type": "integer"}, "length": map[string]any{"type": "integer"}}, "finding_id"),
|
||||
func(ctx context.Context, in json.RawMessage) (actool.Result, error) {
|
||||
var a struct {
|
||||
FindingID json.RawMessage `json:"finding_id"`
|
||||
BindingID json.RawMessage `json:"binding_id"`
|
||||
Side string `json:"side"`
|
||||
Offset, Length int64
|
||||
}
|
||||
if err := json.Unmarshal(in, &a); err != nil {
|
||||
return actool.Errorf(err.Error()), nil
|
||||
}
|
||||
id, bid := parseProfileID(a.FindingID), parseProfileID(a.BindingID)
|
||||
if err := s.agentFindingTrafficAccess(ctx, id, false); err != nil {
|
||||
return actool.Errorf(err.Error()), nil
|
||||
}
|
||||
if len(a.BindingID) == 0 {
|
||||
list, err := s.m.pg.GetFindingTraffic(ctx, id)
|
||||
if err != nil {
|
||||
return actool.Errorf(err.Error()), nil
|
||||
}
|
||||
return jsonResult(trafficSummary(list))
|
||||
}
|
||||
if a.Side == "" {
|
||||
a.Side = "response"
|
||||
}
|
||||
store := s.evidenceStore()
|
||||
var result any
|
||||
err := store.WithBinding(ctx, id, bid, func(b db.FindingTrafficBinding) error {
|
||||
preview, err := readEvidencePreview(store, b.Snapshot, a.Side, a.Offset, a.Length)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
head := b.Snapshot.ReqHead
|
||||
if a.Side == "response" {
|
||||
head = b.Snapshot.RespHead
|
||||
}
|
||||
result = map[string]any{"binding_id": strconv.FormatInt(bid, 10), "head": head, "body": preview}
|
||||
return nil
|
||||
})
|
||||
if err != nil {
|
||||
return actool.Errorf(err.Error()), nil
|
||||
}
|
||||
return jsonResult(result)
|
||||
})
|
||||
}
|
||||
@@ -0,0 +1,121 @@
|
||||
package server
|
||||
|
||||
import (
|
||||
"archive/zip"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
"os"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/Autumn-27/artex/db"
|
||||
"github.com/Autumn-27/artex/evidence"
|
||||
"github.com/Autumn-27/artex/report"
|
||||
)
|
||||
|
||||
func writeFindingsEvidenceZip(out io.Writer, findings []*db.DBFinding, stage string, now time.Time) error {
|
||||
zw := zip.NewWriter(out)
|
||||
if err := writeFindingsZipEntries(zw, findings, stage, now); err != nil {
|
||||
zw.Close()
|
||||
return err
|
||||
}
|
||||
return zw.Close()
|
||||
}
|
||||
|
||||
func writeFindingsZipEntries(zw *zip.Writer, findings []*db.DBFinding, stage string, now time.Time) error {
|
||||
store := evidence.New(nil, nil, stage) // private export copy; not subject to GC
|
||||
used := map[string]int{}
|
||||
for _, f := range findings {
|
||||
base := report.FindingFilename(f)
|
||||
name := base
|
||||
if n := used[base]; n > 0 {
|
||||
name = fmt.Sprintf("%s-%d.md", strings.TrimSuffix(base, ".md"), n+1)
|
||||
}
|
||||
used[base]++
|
||||
w, err := zw.Create(name)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if _, err = io.WriteString(w, report.SingleFindingMarkdown(f, now)); err != nil {
|
||||
return err
|
||||
}
|
||||
for _, b := range f.TrafficBindings {
|
||||
prefix := fmt.Sprintf("evidence/%d/%d/", f.ID, b.ID)
|
||||
w, err := zw.Create(prefix + "manifest.json")
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if err = json.NewEncoder(w).Encode(b); err != nil {
|
||||
return err
|
||||
}
|
||||
for _, side := range []string{"request", "response"} {
|
||||
if err := func() error {
|
||||
body, _, err := store.OpenBody(b.Snapshot, side)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer body.Close()
|
||||
w, err := zw.Create(prefix + side + ".http")
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
head := b.Snapshot.ReqHead
|
||||
if side == "response" {
|
||||
head = b.Snapshot.RespHead
|
||||
}
|
||||
if _, err = io.WriteString(w, strings.TrimRight(head, "\r\n")+"\n\n"); err != nil {
|
||||
return err
|
||||
}
|
||||
if _, err = io.Copy(w, body); err != nil {
|
||||
return err
|
||||
}
|
||||
if _, err = body.Seek(0, io.SeekStart); err != nil {
|
||||
return err
|
||||
}
|
||||
w, err = zw.Create(prefix + side + ".bin")
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
_, err = io.Copy(w, body)
|
||||
return err
|
||||
}(); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func buildFindingsEvidenceZip(path string, findings []*db.DBFinding, stage string, now time.Time) error {
|
||||
f, err := os.OpenFile(path, os.O_CREATE|os.O_EXCL|os.O_WRONLY, 0o600)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
writeErr := writeFindingsEvidenceZip(f, findings, stage, now)
|
||||
syncErr := f.Sync()
|
||||
closeErr := f.Close()
|
||||
if err := errors.Join(writeErr, syncErr, closeErr); err != nil {
|
||||
return err
|
||||
}
|
||||
// Re-read every compressed entry to validate the finished ZIP, including CRCs,
|
||||
// before the caller starts an HTTP download response.
|
||||
archive, err := zip.OpenReader(path)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer archive.Close()
|
||||
for _, entry := range archive.File {
|
||||
r, err := entry.Open()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
_, readErr := io.Copy(io.Discard, r)
|
||||
if err := errors.Join(readErr, r.Close()); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
@@ -0,0 +1,473 @@
|
||||
package server
|
||||
|
||||
import (
|
||||
"archive/zip"
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http/httptest"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strconv"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/Autumn-27/artex/agent"
|
||||
"github.com/Autumn-27/artex/db"
|
||||
"github.com/Autumn-27/artex/traffic"
|
||||
)
|
||||
|
||||
func trafficEvidenceServer(t *testing.T) (*Server, *db.RecordedFinding, func(string, string, string) *httptest.ResponseRecorder) {
|
||||
t.Helper()
|
||||
m, err := NewManager(t.TempDir(), "")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
t.Cleanup(func() { m.Close() })
|
||||
m.traffic, err = traffic.Open(filepath.Join(m.dir, "traffic"), ":0")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
task, err := m.CreateTask("traffic evidence API", "local evidence", nil, 0, 0)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
tid, _ := strconv.ParseInt(task.ID, 10, 64)
|
||||
t.Cleanup(func() { m.pg.Exec(`DELETE FROM task_archives WHERE task_id=$1`, tid); m.pg.DeleteTask(tid) })
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
cancel()
|
||||
s := New(ctx, m, t.TempDir(), t.TempDir(), t.TempDir())
|
||||
s.archiveWG.Wait()
|
||||
token, err := signJWT(s.jwtKey)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
request := func(method, path, body string) *httptest.ResponseRecorder {
|
||||
t.Helper()
|
||||
r := httptest.NewRequest(method, path, strings.NewReader(body))
|
||||
r.Header.Set("Authorization", "Bearer "+token)
|
||||
r.Header.Set("Content-Type", "application/json")
|
||||
w := httptest.NewRecorder()
|
||||
s.Handler().ServeHTTP(w, r)
|
||||
return w
|
||||
}
|
||||
input := db.RecordFindingInput{TaskID: tid, ExplorationID: task.ExpID, Worker: "test", VulnClass: "TEST", Summary: "bound evidence", Name: "Evidence API", Severity: "low"}
|
||||
f, err := s.evidenceStore().Record(context.Background(), input, nil)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
return s, f, request
|
||||
}
|
||||
|
||||
func seedServerEvidenceFlow(t *testing.T, s *Server, id string, body []byte) {
|
||||
t.Helper()
|
||||
_, err := s.m.traffic.DB().Exec(`INSERT INTO exchanges(id,ts,host,method,url_template,url,status,content_type,req_len,resp_len,path) VALUES(?,1,'evidence.local','POST','/','http://evidence.local/',200,'application/octet-stream',0,?,'')`, id, len(body))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
_, err = s.m.traffic.DB().Exec(`INSERT INTO exchange_bodies(id,req_head,resp_head,req_body,resp_body) VALUES(?,'POST / HTTP/1.1','HTTP/1.1 200 OK',?,?)`, id, []byte{}, body)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestFindingTrafficToolUpgradePreservesCustomization(t *testing.T) {
|
||||
s, _, _ := trafficEvidenceServer(t)
|
||||
pg := s.m.pg
|
||||
worker, err := pg.GetAgentByKey("worker")
|
||||
if err != nil || worker == nil {
|
||||
t.Fatal("missing worker", err)
|
||||
}
|
||||
oldPrompt, _ := pg.CurrentPrompt(worker.ID)
|
||||
t.Cleanup(func() { pg.SavePrompt(worker.ID, oldPrompt, "restore test fixture", "test") })
|
||||
if _, err := pg.SavePrompt(worker.ID, "USER CUSTOM PROMPT", "test", "test"); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
custom := json.RawMessage(`{"type":"object","properties":{"evidence":{"type":"string","description":"user evidence instructions"}},"required":["evidence"]}`)
|
||||
for _, key := range []string{"report_finding", "update_finding_report", "get_finding_traffic"} {
|
||||
old, err := pg.GetTool(key)
|
||||
if err != nil || old == nil {
|
||||
t.Fatal("missing tool", key, err)
|
||||
}
|
||||
t.Cleanup(func() {
|
||||
bindings, _ := json.Marshal(old.Agents)
|
||||
pg.UpdateTool(old.Key, old.Description, old.Schema, bindings, old.Enabled)
|
||||
})
|
||||
if err := pg.UpdateTool(key, "user description", custom, json.RawMessage(`["worker"]`), false); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
if err := pg.SetSetting("finding_traffic_tools_v1", "false"); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
s.seedFindingTrafficTools()
|
||||
for key, property := range map[string]string{"report_finding": "traffic_refs", "update_finding_report": "evidence_version"} {
|
||||
tool, err := pg.GetTool(key)
|
||||
if err != nil || tool == nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
var schema struct {
|
||||
Properties map[string]json.RawMessage `json:"properties"`
|
||||
Required []string `json:"required"`
|
||||
}
|
||||
if err := json.Unmarshal(tool.Schema, &schema); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if tool.Enabled || tool.Description != "user description" || len(tool.Agents) != 1 || tool.Agents[0] != "worker" || len(schema.Required) != 1 || schema.Required[0] != "evidence" || len(schema.Properties[property]) == 0 || !strings.Contains(string(schema.Properties["evidence"]), "user evidence instructions") {
|
||||
t.Fatalf("custom configuration overwritten: %+v", tool)
|
||||
}
|
||||
}
|
||||
reader, _ := pg.GetTool("get_finding_traffic")
|
||||
if reader.Enabled || !contains(reader.Agents, "reporter") {
|
||||
t.Fatal("reader should be bound without enabling it")
|
||||
}
|
||||
if err := pg.RemoveAgentFromTool("reporter", "get_finding_traffic"); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
s.seedFindingTrafficTools()
|
||||
reader, _ = pg.GetTool("get_finding_traffic")
|
||||
if contains(reader.Agents, "reporter") {
|
||||
t.Fatal("restart overwrote user's unbinding")
|
||||
}
|
||||
if got, _ := pg.CurrentPrompt(worker.ID); got != "USER CUSTOM PROMPT" {
|
||||
t.Fatal("custom prompt replaced")
|
||||
}
|
||||
}
|
||||
|
||||
func TestFindingTrafficAPIAndExport(t *testing.T) {
|
||||
s, f, req := trafficEvidenceServer(t)
|
||||
body := bytes.Repeat([]byte{0, 255, 65, 66}, 5000)
|
||||
for _, id := range []string{"a", "b", "c"} {
|
||||
seedServerEvidenceFlow(t, s, id, body)
|
||||
}
|
||||
base := fmt.Sprintf("/api/exploration/findings/%d/traffic", f.FindingID)
|
||||
w := req("POST", base, `{"traffic_refs":[{"traffic_id":"a","role":"baseline"},{"traffic_id":"b","role":"proof","note":"proof note"},{"traffic_id":"c"}]}`)
|
||||
if w.Code != 200 {
|
||||
t.Fatalf("bind: %d %s", w.Code, w.Body)
|
||||
}
|
||||
var list db.FindingTraffic
|
||||
if err := json.Unmarshal(w.Body.Bytes(), &list); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(list.Bindings) != 3 || list.Bindings[0].Snapshot.ReqHead != "" {
|
||||
t.Fatal("not a bounded summary")
|
||||
}
|
||||
first, second, third := list.Bindings[0].ID, list.Bindings[1].ID, list.Bindings[2].ID
|
||||
if w = req("PATCH", fmt.Sprintf("%s/%d", base, first), `{"version":0,"note":"stale"}`); w.Code != 409 {
|
||||
t.Fatal(w.Code, w.Body)
|
||||
}
|
||||
if w = req("PUT", base+"/order", fmt.Sprintf(`{"version":1,"binding_ids":["%d","%d","%d"]}`, third, first, second)); w.Code != 200 {
|
||||
t.Fatal(w.Code, w.Body)
|
||||
}
|
||||
if w = req("PUT", base+"/order", fmt.Sprintf(`{"version":2,"binding_ids":["%d"]}`, first)); w.Code != 409 {
|
||||
t.Fatal(w.Code, w.Body)
|
||||
}
|
||||
if w = req("GET", fmt.Sprintf("%s/%d", base, first), ""); w.Code != 200 || !strings.Contains(w.Body.String(), `"binary":true`) {
|
||||
t.Fatal(w.Code, w.Body)
|
||||
}
|
||||
if w = req("GET", fmt.Sprintf("%s/%d/body?side=response&offset=20001", base, first), ""); w.Code != 422 {
|
||||
t.Fatal(w.Code, w.Body)
|
||||
}
|
||||
if _, err := s.m.traffic.DeleteHost("evidence.local"); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if w = req("GET", fmt.Sprintf("%s/%d/body?side=response&download=1", base, first), ""); w.Code != 200 || !bytes.Equal(w.Body.Bytes(), body) {
|
||||
t.Fatal("download after original removal failed")
|
||||
}
|
||||
// IDs belong to the requested finding; cross-finding detail cannot be read.
|
||||
finding, err := s.m.pg.GetFinding(f.FindingID)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
task, err := s.m.pg.GetTask(*finding.TaskID)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
other, err := s.evidenceStore().Record(t.Context(), db.RecordFindingInput{TaskID: task.ID, ExplorationID: task.ExplorationID, Summary: "other", Severity: "low"}, nil)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if w = req("GET", fmt.Sprintf("/api/exploration/findings/%d/traffic/%d", other.FindingID, first), ""); w.Code != 404 {
|
||||
t.Fatal(w.Code, w.Body)
|
||||
}
|
||||
exportURL := fmt.Sprintf("/api/exploration/findings/export?format=md-zip&scope=selected&ids=%d", f.FindingID)
|
||||
w = req("GET", exportURL, "")
|
||||
if w.Code != 200 {
|
||||
t.Fatal(w.Code, w.Body)
|
||||
}
|
||||
archive, err := zip.NewReader(bytes.NewReader(w.Body.Bytes()), int64(w.Body.Len()))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
binaries := 0
|
||||
linked := false
|
||||
for _, entry := range archive.File {
|
||||
r, err := entry.Open()
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
raw, err := io.ReadAll(r)
|
||||
r.Close()
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if strings.HasSuffix(entry.Name, "/response.bin") {
|
||||
binaries++
|
||||
if !bytes.Equal(raw, body) {
|
||||
t.Fatal("export changed binary")
|
||||
}
|
||||
}
|
||||
if strings.HasSuffix(entry.Name, ".md") && bytes.Contains(raw, []byte(fmt.Sprintf("evidence/%d/%d/", f.FindingID, third))) {
|
||||
linked = true
|
||||
}
|
||||
}
|
||||
if binaries != 3 || !linked {
|
||||
t.Fatalf("missing attachments/links: %d %v", binaries, linked)
|
||||
}
|
||||
// Reports can use snapshots with capture stopped and original rows gone.
|
||||
result, err := s.toolGetFindingTraffic().Call(context.Background(), json.RawMessage(fmt.Sprintf(`{"finding_id":"%d"}`, f.FindingID)), nil)
|
||||
if err != nil || !strings.Contains(result.Flatten(), "proof note") {
|
||||
t.Fatal(result, err)
|
||||
}
|
||||
// Corrupt body must fail before an attachment response is sent.
|
||||
snap := list.Bindings[0].Snapshot
|
||||
path := filepath.Join(s.m.dir, "evidence", "blobs", snap.RespHash[:2], snap.RespHash+".bin")
|
||||
if err = os.WriteFile(path, []byte("corrupt"), 0600); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if w = req("GET", exportURL, ""); w.Code < 400 || strings.Contains(w.Header().Get("Content-Type"), "zip") {
|
||||
t.Fatal("corrupt export returned an attachment", w.Code)
|
||||
}
|
||||
}
|
||||
|
||||
func TestFindingTrafficArchiveV3RoundTripAndRetry(t *testing.T) {
|
||||
s, f, _ := trafficEvidenceServer(t)
|
||||
ctx := context.Background()
|
||||
s.ctx = ctx
|
||||
body := []byte("evidence without original traffic\x00\xff")
|
||||
seedServerEvidenceFlow(t, s, "original", body)
|
||||
list, err := s.evidenceStore().Bind(ctx, f.FindingID, []db.TrafficRef{{TrafficID: "original", Role: "proof", Note: "keep note"}})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
finding, err := s.m.pg.GetFinding(f.FindingID)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
tid := *finding.TaskID
|
||||
sharedTask, err := s.m.CreateTask("shared evidence", "fixture", nil, 0, 0)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
sid, _ := strconv.ParseInt(sharedTask.ID, 10, 64)
|
||||
defer s.m.pg.DeleteTask(sid)
|
||||
shared, err := s.evidenceStore().Record(ctx, db.RecordFindingInput{TaskID: sid, ExplorationID: sharedTask.ExpID, Summary: "shared", Severity: "low"}, []db.TrafficRef{{TrafficID: "original"}})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if _, err = s.m.traffic.DeleteHost("evidence.local"); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err = s.m.SetTaskPaused(strconv.FormatInt(tid, 10), true); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
job, err := s.m.pg.QueueTaskArchive(tid)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err = s.m.pg.EditFindingTraffic(ctx, f.FindingID, list.Bindings[0].ID, 1, nil, nil, true, nil); err == nil {
|
||||
t.Fatal("write allowed after archive queued")
|
||||
}
|
||||
if err = s.runOneTaskArchiveJob(); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if got, _ := s.m.pg.GetFinding(f.FindingID); got != nil {
|
||||
t.Fatal("archive retained hot finding")
|
||||
}
|
||||
if err = s.evidenceStore().Collect(ctx, time.Now().Add(48*time.Hour)); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err = s.evidenceStore().WithBinding(ctx, shared.FindingID, shared.Traffic.Bindings[0].ID, func(b db.FindingTrafficBinding) error {
|
||||
r, _, err := s.evidenceStore().OpenBody(b.Snapshot, "response")
|
||||
if r != nil {
|
||||
r.Close()
|
||||
}
|
||||
return err
|
||||
}); err != nil {
|
||||
t.Fatal("shared body deleted", err)
|
||||
}
|
||||
if _, err = s.m.pg.QueueTaskArchiveRestore(job.ID); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
// Force a database error after validated body installation; then retry normally.
|
||||
if _, err = s.m.pg.Exec(`CREATE FUNCTION fail_evidence_restore_test() RETURNS trigger LANGUAGE plpgsql AS $$ BEGIN RAISE EXCEPTION 'restore fixture failure'; END $$`); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if _, err = s.m.pg.Exec(`CREATE TRIGGER fail_evidence_restore_test BEFORE INSERT ON finding_traffic_bindings FOR EACH ROW EXECUTE FUNCTION fail_evidence_restore_test()`); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer s.m.pg.Exec(`DROP FUNCTION IF EXISTS fail_evidence_restore_test() CASCADE`)
|
||||
if err = s.runOneTaskArchiveJob(); err == nil {
|
||||
t.Fatal("expected restore failure")
|
||||
}
|
||||
if got, _ := s.m.pg.GetFinding(f.FindingID); got != nil {
|
||||
t.Fatal("partial restore")
|
||||
}
|
||||
if _, err = s.m.pg.Exec(`DROP FUNCTION fail_evidence_restore_test() CASCADE`); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if _, err = s.m.pg.QueueTaskArchiveRestore(job.ID); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err = s.runOneTaskArchiveJob(); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
restored, err := s.m.pg.GetFindingTraffic(ctx, f.FindingID)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
b := restored.Bindings[0]
|
||||
if restored.Version != list.Version || b.ID != list.Bindings[0].ID || b.Note != "keep note" || b.Role != "proof" || b.Position != 0 {
|
||||
t.Fatalf("restored: %+v", restored)
|
||||
}
|
||||
if err = s.evidenceStore().WithBinding(ctx, f.FindingID, b.ID, func(b db.FindingTrafficBinding) error {
|
||||
r, _, err := s.evidenceStore().OpenBody(b.Snapshot, "response")
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer r.Close()
|
||||
got, err := io.ReadAll(r)
|
||||
if !bytes.Equal(got, body) {
|
||||
t.Fatal("restore changed body")
|
||||
}
|
||||
return err
|
||||
}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
// The read tool enforces task visibility too.
|
||||
result, err := s.toolGetFindingTraffic().Call(agent.WithRunInfo(ctx, agent.RunInfo{TaskID: sid}), json.RawMessage(fmt.Sprintf(`{"finding_id":"%d"}`, f.FindingID)), nil)
|
||||
if err != nil || !strings.Contains(result.Flatten(), "不可读取") {
|
||||
t.Fatal(result, err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestFindingTrafficFailedReportDoesNotTrigger(t *testing.T) {
|
||||
s, f, _ := trafficEvidenceServer(t)
|
||||
finding, err := s.m.pg.GetFinding(f.FindingID)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
task, err := s.m.pg.GetTask(*finding.TaskID)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
var last int64
|
||||
if err = s.m.pg.QueryRow(`SELECT COALESCE(max(id),0) FROM activity`).Scan(&last); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err = s.m.pg.SetSchedState(schedKeyLastToolCall, strconv.FormatInt(last, 10)); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
id, err := s.m.pg.Exploration(task.ExplorationID).AppendActivity(db.Activity{Kind: "tool_result", Tool: "report_finding", ToolUseID: "failure-fixture", IsError: true, Detail: "body missing"})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
// A nil server makes any attempted trigger fail this test immediately.
|
||||
scheduler := &Scheduler{pg: s.m.pg}
|
||||
scheduler.fireToolCalls([]*db.AgentTrigger{{OnToolCall: true, ToolNames: []string{"report_finding"}, AgentKey: "reporter"}})
|
||||
if got := scheduler.mustState(schedKeyLastToolCall); got != strconv.FormatInt(id, 10) {
|
||||
t.Fatal("failed finding must still advance watermark", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestFindingTrafficUTF8SegmentsAndInheritedWrites(t *testing.T) {
|
||||
s, f, req := trafficEvidenceServer(t)
|
||||
ctx := context.Background()
|
||||
text := strings.Repeat("中文证据🙂", 2000)
|
||||
seedServerEvidenceFlow(t, s, "unicode", []byte(text))
|
||||
list, err := s.evidenceStore().Bind(ctx, f.FindingID, []db.TrafficRef{{TrafficID: "unicode"}})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
var joined strings.Builder
|
||||
offset := int64(0)
|
||||
for {
|
||||
var preview evidencePreview
|
||||
err = s.evidenceStore().WithBinding(ctx, f.FindingID, list.Bindings[0].ID, func(b db.FindingTrafficBinding) error {
|
||||
var e error
|
||||
preview, e = readEvidencePreview(s.evidenceStore(), b.Snapshot, "response", offset, 8192)
|
||||
return e
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if preview.Binary {
|
||||
t.Fatal("valid UTF-8 classified binary")
|
||||
}
|
||||
joined.WriteString(preview.Content)
|
||||
if !preview.Truncated {
|
||||
break
|
||||
}
|
||||
if preview.NextOffset <= offset {
|
||||
t.Fatal("preview made no progress")
|
||||
}
|
||||
offset = preview.NextOffset
|
||||
}
|
||||
if joined.String() != text {
|
||||
t.Fatal("UTF-8 character split damaged preview")
|
||||
}
|
||||
finding, err := s.m.pg.GetFinding(f.FindingID)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
child, err := s.m.CreateTaskWithOptions("inherited evidence API", "fixture", db.TaskCreateOptions{SourceTaskIDs: []int64{*finding.TaskID}})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
cid, _ := strconv.ParseInt(child.ID, 10, 64)
|
||||
defer s.m.pg.DeleteTask(cid)
|
||||
base := fmt.Sprintf("/api/exploration/findings/%d/traffic", f.FindingID)
|
||||
suffix := "?context_task=" + child.ID
|
||||
if w := req("GET", base+suffix, ""); w.Code != 200 {
|
||||
t.Fatal(w.Code, w.Body)
|
||||
}
|
||||
for _, method := range []string{"POST", "PATCH", "PUT", "DELETE"} {
|
||||
path := base
|
||||
if method == "PATCH" || method == "DELETE" {
|
||||
path += fmt.Sprintf("/%d", list.Bindings[0].ID)
|
||||
}
|
||||
if method == "PUT" {
|
||||
path += "/order"
|
||||
}
|
||||
w := req(method, path+suffix, `{"version":1,"traffic_refs":[{"traffic_id":"unicode"}]}`)
|
||||
if w.Code != 403 {
|
||||
t.Fatal(method, w.Code, w.Body)
|
||||
}
|
||||
}
|
||||
if _, err = s.m.traffic.DeleteHost("evidence.local"); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
result, err := s.toolUpdateFindingReport().Call(ctx, json.RawMessage(fmt.Sprintf(`{"finding_id":%d,"evidence_version":1,"report":"## 证据报告\n\n证据 #%d:已验证完整请求响应"}`, f.NodeID, list.Bindings[0].ID)), nil)
|
||||
if err != nil || result.IsError {
|
||||
t.Fatal(result, err)
|
||||
}
|
||||
updated, err := s.m.pg.GetFinding(f.FindingID)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if updated.ReportEvidenceVersion != 1 || updated.EvidenceVersion != 1 {
|
||||
t.Fatal("report did not record the read evidence version")
|
||||
}
|
||||
for _, format := range []string{"md-single", "json", "csv"} {
|
||||
w := req("GET", fmt.Sprintf("/api/exploration/findings/export?scope=selected&ids=%d&format=%s", f.FindingID, format), "")
|
||||
if w.Code != 200 {
|
||||
t.Fatal(format, w.Code, w.Body)
|
||||
}
|
||||
if !strings.Contains(w.Body.String(), fmt.Sprint(list.Bindings[0].ID)) {
|
||||
t.Fatal("export omitted evidence ID", format)
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,157 @@
|
||||
package server
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"log"
|
||||
"strconv"
|
||||
|
||||
"github.com/Autumn-27/artex/agent"
|
||||
"github.com/Autumn-27/artex/db"
|
||||
actool "github.com/Autumn-27/norma/tool"
|
||||
)
|
||||
|
||||
func (s *Server) seedFindingWorkflowTools() {
|
||||
const flag = "finding_workflow_tools_v2_reporter"
|
||||
if value, _, _ := s.m.pg.GetSetting(flag); value == "true" {
|
||||
return
|
||||
}
|
||||
for _, key := range []string{"report_finding", "add_hint", "add_task_hint"} {
|
||||
row, err := s.m.pg.GetTool(key)
|
||||
if err != nil {
|
||||
log.Printf("[evidence] load %s: %v", key, err)
|
||||
return
|
||||
}
|
||||
if row == nil || !row.System {
|
||||
continue
|
||||
}
|
||||
var schema map[string]any
|
||||
if err := json.Unmarshal(row.Schema, &schema); err != nil {
|
||||
log.Printf("[evidence] invalid schema for %s: %v", key, err)
|
||||
return
|
||||
}
|
||||
if schema == nil {
|
||||
log.Printf("[evidence] missing object schema for %s", key)
|
||||
return
|
||||
}
|
||||
props := objectProperty(schema, "properties")
|
||||
if key == "report_finding" {
|
||||
if _, exists := props["evidence_hint_id"]; !exists {
|
||||
props["evidence_hint_id"] = map[string]any{"type": "integer", "description": "可选:本任务中对应此漏洞的 hint ID;读取该提示保存的 traffic_refs 一并绑定,无提示时省略"}
|
||||
}
|
||||
} else {
|
||||
if _, exists := props["traffic_refs"]; !exists {
|
||||
props["traffic_refs"] = agent.HintTrafficSchema()
|
||||
}
|
||||
hints := objectProperty(props, "hints")
|
||||
if _, exists := hints["type"]; !exists {
|
||||
hints["type"] = "array"
|
||||
}
|
||||
items := objectProperty(hints, "items")
|
||||
if _, ok := items["type"]; !ok {
|
||||
items["type"] = "object"
|
||||
}
|
||||
itemProps := objectProperty(items, "properties")
|
||||
for name, value := range map[string]any{"text": strParam("提示内容"), "asset_ids": map[string]any{"type": "array", "items": map[string]any{"type": "integer"}}, "traffic_refs": agent.HintTrafficSchema()} {
|
||||
if _, exists := itemProps[name]; !exists {
|
||||
itemProps[name] = value
|
||||
}
|
||||
}
|
||||
}
|
||||
raw, _ := json.Marshal(schema)
|
||||
result, err := s.m.pg.Exec(`UPDATE tools SET schema=$2::jsonb,updated_at=now() WHERE key=$1 AND system AND schema=$3::jsonb`, key, string(raw), string(row.Schema))
|
||||
if err != nil {
|
||||
log.Printf("[evidence] upgrade %s: %v", key, err)
|
||||
return
|
||||
}
|
||||
if n, _ := result.RowsAffected(); n != 1 {
|
||||
return
|
||||
} // preserve concurrent user edits
|
||||
}
|
||||
// Upgrade only the original default binding. Customized lists and enabled
|
||||
// flags survive; the one-time flag also preserves future user unbinding.
|
||||
readers := `["worker","reporter"]`
|
||||
for _, key := range []string{"traffic_search", "traffic_get", "traffic_blob"} {
|
||||
if _, err := s.m.pg.Exec(`UPDATE tools SET agents=$2::jsonb WHERE key=$1 AND system AND (agents='["worker"]'::jsonb OR (agents @> '["worker","planner","mainagent","auto","pentest"]'::jsonb AND jsonb_array_length(agents)=5))`, key, readers); err != nil {
|
||||
return
|
||||
}
|
||||
}
|
||||
if _, err := s.m.pg.Exec(`UPDATE tools SET agents=$1::jsonb WHERE key='get_finding_traffic' AND system AND agents @> '["auto","reporter"]'::jsonb AND jsonb_array_length(agents)=2`, `["auto","reporter","worker","planner","mainagent","pentest"]`); err != nil {
|
||||
return
|
||||
}
|
||||
// Replace the previous code default only; preserve customized binding lists.
|
||||
if _, err := s.m.pg.Exec(`UPDATE tools SET agents='["reporter"]'::jsonb WHERE key='bind_finding_traffic' AND system AND agents @> '["worker","planner","mainagent","auto","pentest"]'::jsonb AND jsonb_array_length(agents)=5`); err != nil {
|
||||
return
|
||||
}
|
||||
if err := s.m.pg.AddAgentToToolBinding("reporter", []string{"bind_finding_traffic"}); err != nil {
|
||||
return
|
||||
}
|
||||
_ = s.m.pg.SetSetting(flag, "true")
|
||||
}
|
||||
|
||||
func objectProperty(parent map[string]any, key string) map[string]any {
|
||||
value, ok := parent[key].(map[string]any)
|
||||
if !ok {
|
||||
value = map[string]any{}
|
||||
parent[key] = value
|
||||
}
|
||||
return value
|
||||
}
|
||||
|
||||
func (s *Server) agentFindingTrafficAccess(ctx context.Context, id int64, write bool) error {
|
||||
if id <= 0 {
|
||||
return errors.New("finding_id 必须为独立漏洞记录 ID;不是探索节点 ID")
|
||||
}
|
||||
f, err := s.m.pg.GetFinding(id)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if f == nil {
|
||||
return fmt.Errorf("%w:finding_id=%d。证据工具使用独立漏洞记录 ID,请从 list_task_findings / get_task_node_detail 的 finding_id 字段读取;不要传 id / finding_node_id", db.ErrFindingNotFound, id)
|
||||
}
|
||||
if ri := agent.RunInfoFrom(ctx); ri.TaskID > 0 {
|
||||
task := s.m.ResolveTask(strconv.FormatInt(ri.TaskID, 10))
|
||||
if task == nil {
|
||||
return errors.New("任务不存在")
|
||||
}
|
||||
_, inherited, allowed := findingProvenanceInTask(task, f.TaskID)
|
||||
if !allowed {
|
||||
return errors.New("当前任务不可读取该漏洞")
|
||||
}
|
||||
if write && inherited {
|
||||
return errors.New("继承漏洞的流量证据只读,请到来源任务修改")
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (s *Server) toolBindFindingTraffic() actool.CoreTool {
|
||||
return wrTool("bind_finding_traffic", "为已登记漏洞补绑经核实的真实 HTTP 流量。finding_id 使用独立漏洞记录 ID;不要传探索节点 ID。同批引用全部成功或全部失败,重复引用不覆盖已有说明。补绑会使已有报告标记待更新;不要为补包重新探测或重复创建漏洞。",
|
||||
objSchema(map[string]any{"finding_id": strParam("独立漏洞记录 ID,从 list_task_findings / get_task_node_detail 的 finding_id 字段读取"), "traffic_refs": agent.HintTrafficSchema()}, "finding_id", "traffic_refs"),
|
||||
func(ctx context.Context, raw json.RawMessage) (actool.Result, error) {
|
||||
if !s.m.pg.GetBool(settingAgentTrafficBinding, false) {
|
||||
return actool.Errorf("Agent 自动绑定流量已关闭;请在系统设置开启,或使用页面人工绑定。"), nil
|
||||
}
|
||||
var args struct {
|
||||
FindingID json.RawMessage `json:"finding_id"`
|
||||
Refs []db.TrafficRef `json:"traffic_refs"`
|
||||
}
|
||||
if err := json.Unmarshal(raw, &args); err != nil {
|
||||
return actool.Errorf(err.Error()), nil
|
||||
}
|
||||
id := parseProfileID(args.FindingID)
|
||||
if err := s.agentFindingTrafficAccess(ctx, id, true); err != nil {
|
||||
return actool.Errorf(err.Error()), nil
|
||||
}
|
||||
if len(args.Refs) == 0 {
|
||||
return actool.Errorf("补绑需要至少一条已核实的 traffic_refs;无流量无需调用此工具"), nil
|
||||
}
|
||||
list, err := s.evidenceStore().Bind(ctx, id, args.Refs)
|
||||
if err != nil {
|
||||
return actool.Errorf(err.Error()), nil
|
||||
}
|
||||
return jsonResult(trafficSummary(list))
|
||||
})
|
||||
}
|
||||
@@ -0,0 +1,312 @@
|
||||
package server
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"strconv"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"github.com/Autumn-27/artex/agent"
|
||||
"github.com/Autumn-27/artex/db"
|
||||
actool "github.com/Autumn-27/norma/tool"
|
||||
)
|
||||
|
||||
func workflowCall(t *testing.T, ctx context.Context, tool actool.CoreTool, input any, wantError bool) string {
|
||||
t.Helper()
|
||||
raw, err := json.Marshal(input)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
r, err := tool.Call(ctx, raw, nil)
|
||||
if err != nil || r.IsError != wantError {
|
||||
t.Fatalf("%s: error=%v result=%s", tool.Name(), err, r.Flatten())
|
||||
}
|
||||
return r.Flatten()
|
||||
}
|
||||
|
||||
func workflowTool(t *testing.T, tools []actool.CoreTool, name string) actool.CoreTool {
|
||||
t.Helper()
|
||||
for _, tool := range tools {
|
||||
if tool.Name() == name {
|
||||
return tool
|
||||
}
|
||||
}
|
||||
t.Fatalf("missing tool: %s", name)
|
||||
return nil
|
||||
}
|
||||
|
||||
func TestFindingWorkflowAutoHintToPlannerAndSetting(t *testing.T) {
|
||||
s, initial, request := trafficEvidenceServer(t)
|
||||
pg := s.m.pg
|
||||
ctx := context.Background()
|
||||
if _, err := pg.Exec(`DELETE FROM settings WHERE key=$1`, settingAgentTrafficBinding); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
t.Cleanup(func() { pg.SetBool(settingAgentTrafficBinding, false) })
|
||||
if s.settingsPayload()[settingAgentTrafficBinding] != false {
|
||||
t.Fatal("missing setting must default off")
|
||||
}
|
||||
seedServerEvidenceFlow(t, s, "handoff-proof", []byte("verified local proof"))
|
||||
seedServerEvidenceFlow(t, s, "handoff-baseline", []byte("local baseline"))
|
||||
seedServerEvidenceFlow(t, s, "handoff-verification", []byte{0, 1, 255})
|
||||
// Manual binding is independent of the Agent switch.
|
||||
if r := request("POST", fmt.Sprintf("/api/exploration/findings/%d/traffic", initial.FindingID), `{"traffic_refs":[{"traffic_id":"handoff-proof"}]}`); r.Code != 200 {
|
||||
t.Fatal(r.Code, r.Body)
|
||||
}
|
||||
bind := s.toolBindFindingTraffic()
|
||||
input := map[string]any{"finding_id": initial.FindingID, "traffic_refs": []db.TrafficRef{{TrafficID: "handoff-baseline"}}}
|
||||
workflowCall(t, ctx, bind, input, true)
|
||||
if r := request("PUT", "/api/settings", `{"agent_traffic_binding":true}`); r.Code != 200 {
|
||||
t.Fatal(r.Code, r.Body)
|
||||
}
|
||||
if !pg.GetBool(settingAgentTrafficBinding, false) {
|
||||
t.Fatal("switch was not persisted")
|
||||
}
|
||||
f, _ := pg.GetFinding(initial.FindingID)
|
||||
task := s.m.ResolveTask(fmt.Sprint(*f.TaskID))
|
||||
refs := []db.TrafficRef{{TrafficID: "handoff-baseline", Role: "baseline", Note: "normal response"}, {TrafficID: "handoff-proof", Role: "proof", Note: "proves the finding"}, {TrafficID: "handoff-verification", Role: "verification", Note: "binary verification"}}
|
||||
// Real cross-task Auto hint handler -> persisted graph -> Planner tool.
|
||||
hintResult := workflowCall(t, ctx, s.toolAddHint(), map[string]any{"task_id": task.ID, "hints": []any{map[string]any{"text": "Report the confirmed local finding with its verified evidence", "traffic_refs": refs}}}, false)
|
||||
var hints struct {
|
||||
IDs []int64 `json:"ids"`
|
||||
}
|
||||
if err := json.Unmarshal([]byte(hintResult), &hints); err != nil || len(hints.IDs) != 1 || hints.IDs[0] == 0 {
|
||||
t.Fatal(hintResult, err)
|
||||
}
|
||||
graph := workflowCall(t, ctx, s.toolGetTaskGraph(), map[string]any{"task_id": task.ID}, false)
|
||||
for _, ref := range refs {
|
||||
if !strings.Contains(graph, ref.TrafficID) {
|
||||
t.Fatal("handoff lost reference", ref.TrafficID)
|
||||
}
|
||||
}
|
||||
for range 6 {
|
||||
if _, err := task.Store.AddNode(db.KindFact, map[string]any{"summary": "ID separation"}, 0, "confirmed", "test", nil); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
ts := agent.NewToolSet(task.Store, "planner")
|
||||
ts.SetTaskID(*f.TaskID)
|
||||
ts.SetFindingRecorder(s.evidenceStore())
|
||||
notices := 0
|
||||
ts.SetNotifyFinding(func(int64, string) { notices++ })
|
||||
tools, def, cleanup := agent.AugmentTools(ctx, "planner", ts.PlannerTools())
|
||||
defer cleanup()
|
||||
if !strings.Contains(def.FindingGuidance, "evidence_hint_id") || !strings.Contains(def.FindingGuidance, "取消 Worker") {
|
||||
t.Fatal("Planner missed runtime guidance")
|
||||
}
|
||||
report := workflowTool(t, tools, "report_finding")
|
||||
result := workflowCall(t, ctx, report, map[string]any{"vulnclass": "TEST", "severity": "low", "summary": "Planner handoff fixture", "evidence_hint_id": hints.IDs[0]}, false)
|
||||
var recorded db.RecordedFinding
|
||||
if err := json.Unmarshal([]byte(strings.SplitN(result, "\n", 2)[1]), &recorded); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if recorded.FindingID == recorded.NodeID || len(recorded.Traffic.Bindings) != 3 || notices != 1 {
|
||||
t.Fatal("bad finding/traffic result", result)
|
||||
}
|
||||
for i, b := range recorded.Traffic.Bindings {
|
||||
if b.Snapshot.SourceTrafficID != refs[i].TrafficID || b.Role != refs[i].Role || b.Note != refs[i].Note {
|
||||
t.Fatalf("handoff mismatch: %+v", b)
|
||||
}
|
||||
}
|
||||
list := workflowCall(t, ctx, s.toolListTaskFindings(), map[string]any{"task_id": task.ID}, false)
|
||||
var nodes []struct {
|
||||
ID int64 `json:"id"`
|
||||
FindingID int64 `json:"finding_id"`
|
||||
FindingNodeID int64 `json:"finding_node_id"`
|
||||
Count int `json:"traffic_count"`
|
||||
}
|
||||
if err := json.Unmarshal([]byte(list), &nodes); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
found := false
|
||||
for _, n := range nodes {
|
||||
if n.ID == recorded.NodeID {
|
||||
found = n.FindingID == recorded.FindingID && n.FindingNodeID == recorded.NodeID && n.Count == 3
|
||||
}
|
||||
}
|
||||
if !found {
|
||||
t.Fatal("canonical IDs/count missing", list)
|
||||
}
|
||||
workflowCall(t, ctx, s.toolGetFindingTraffic(), map[string]any{"finding_id": recorded.FindingID}, false)
|
||||
wrong := workflowCall(t, ctx, s.toolGetFindingTraffic(), map[string]any{"finding_id": int64(900000000000000000)}, true)
|
||||
if !strings.Contains(wrong, "独立漏洞记录 ID") {
|
||||
t.Fatal("ambiguous ID error", wrong)
|
||||
}
|
||||
input["finding_id"] = recorded.FindingID
|
||||
// Duplicate append preserves metadata and version.
|
||||
workflowCall(t, ctx, bind, input, false)
|
||||
after, err := pg.GetFindingTraffic(ctx, recorded.FindingID)
|
||||
if err != nil || len(after.Bindings) != 3 || after.Version != recorded.Traffic.Version || after.Bindings[0].Note != refs[0].Note {
|
||||
t.Fatal("duplicate changed evidence", after, err)
|
||||
}
|
||||
// A malformed handoff must not create a partial finding or notify Planner.
|
||||
badHint := workflowCall(t, ctx, s.toolAddHint(), map[string]any{"task_id": task.ID, "text": "missing packet", "traffic_refs": []db.TrafficRef{{TrafficID: "missing"}}}, false)
|
||||
var badID int64
|
||||
fmt.Sscanf(badHint, "hint added: %d", &badID)
|
||||
workflowCall(t, ctx, report, map[string]any{"vulnclass": "TEST", "severity": "low", "summary": "must fail", "evidence_hint_id": badID}, true)
|
||||
if notices != 1 {
|
||||
t.Fatal("failed handoff notified before commit")
|
||||
}
|
||||
child, err := s.m.CreateTaskWithOptions("inherited handoff", "fixture", db.TaskCreateOptions{SourceTaskIDs: []int64{*f.TaskID}})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
childID, _ := strconv.ParseInt(child.ID, 10, 64)
|
||||
childRow, _ := pg.GetTask(childID)
|
||||
t.Cleanup(func() { pg.DeleteTask(childRow.ID) })
|
||||
childCtx := agent.WithRunInfo(ctx, agent.RunInfo{TaskID: childRow.ID})
|
||||
workflowCall(t, childCtx, bind, input, true)
|
||||
workflowCall(t, childCtx, s.toolGetFindingTraffic(), map[string]any{"finding_id": recorded.FindingID}, false)
|
||||
// Existing, already assembled tools must honor a later switch-off.
|
||||
if r := request("PUT", "/api/settings", `{"agent_traffic_binding":false}`); r.Code != 200 {
|
||||
t.Fatal(r.Code, r.Body)
|
||||
}
|
||||
workflowCall(t, ctx, bind, input, true)
|
||||
workflowCall(t, ctx, report, map[string]any{"vulnclass": "TEST", "severity": "low", "summary": "off", "evidence_hint_id": hints.IDs[0]}, true)
|
||||
workflowCall(t, ctx, report, map[string]any{"vulnclass": "TCP", "severity": "low", "summary": "no packet needed"}, false)
|
||||
if notices != 2 {
|
||||
t.Fatal("optional no-packet report failed")
|
||||
}
|
||||
workflowCall(t, ctx, s.toolGetFindingTraffic(), map[string]any{"finding_id": recorded.FindingID}, false)
|
||||
tools, def, closeTools := agent.AugmentTools(ctx, "planner", ts.PlannerTools())
|
||||
defer closeTools()
|
||||
if def.FindingGuidance != "" {
|
||||
t.Fatal("disabled feature still injects guidance")
|
||||
}
|
||||
for _, tool := range tools {
|
||||
if tool.Name() == "bind_finding_traffic" {
|
||||
t.Fatal("disabled binding tool exposed")
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestFindingWorkflowMigrationPreservesUserConfiguration(t *testing.T) {
|
||||
s, _, _ := trafficEvidenceServer(t)
|
||||
pg := s.m.pg
|
||||
for _, key := range []string{"report_finding", "add_hint", "add_task_hint", "traffic_search", "traffic_get", "traffic_blob", "bind_finding_traffic", "get_finding_traffic"} {
|
||||
old, err := pg.GetTool(key)
|
||||
if err != nil || old == nil {
|
||||
t.Fatal("missing tool", key, err)
|
||||
}
|
||||
t.Cleanup(func() {
|
||||
bindings, _ := json.Marshal(old.Agents)
|
||||
pg.UpdateTool(key, old.Description, old.Schema, bindings, old.Enabled)
|
||||
})
|
||||
}
|
||||
custom := json.RawMessage(`{"type":"object","properties":{"text":{"type":"string","description":"USER TEXT"},"hints":{"type":"array","items":{"type":"object","properties":{"text":{"type":"string","description":"USER ITEM"}}}}},"required":["text"]}`)
|
||||
for _, key := range []string{"report_finding", "add_hint", "add_task_hint"} {
|
||||
if err := pg.UpdateTool(key, "USER DESCRIPTION", custom, json.RawMessage(`["custom-agent"]`), false); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
for _, key := range []string{"traffic_search", "traffic_get"} {
|
||||
row, _ := pg.GetTool(key)
|
||||
bindings := json.RawMessage(`["custom-agent"]`)
|
||||
if key == "traffic_search" {
|
||||
bindings = json.RawMessage(`["worker"]`)
|
||||
}
|
||||
if err := pg.UpdateTool(key, row.Description, row.Schema, bindings, false); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
if err := pg.SetSetting("finding_workflow_tools_v2_reporter", "false"); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
s.seedFindingWorkflowTools()
|
||||
for _, key := range []string{"report_finding", "add_hint", "add_task_hint"} {
|
||||
row, _ := pg.GetTool(key)
|
||||
if row.Enabled || row.Description != "USER DESCRIPTION" || len(row.Agents) != 1 || row.Agents[0] != "custom-agent" {
|
||||
t.Fatal("changed user tool configuration", key)
|
||||
}
|
||||
if !strings.Contains(string(row.Schema), "USER TEXT") || !strings.Contains(string(row.Schema), "USER ITEM") {
|
||||
t.Fatal("changed user schema", key)
|
||||
}
|
||||
field := "traffic_refs"
|
||||
if key == "report_finding" {
|
||||
field = "evidence_hint_id"
|
||||
}
|
||||
if !strings.Contains(string(row.Schema), field) {
|
||||
t.Fatal("missing optional field", key)
|
||||
}
|
||||
}
|
||||
search, _ := pg.GetTool("traffic_search")
|
||||
get, _ := pg.GetTool("traffic_get")
|
||||
if search.Enabled || !contains(search.Agents, "reporter") || get.Enabled || len(get.Agents) != 1 || get.Agents[0] != "custom-agent" {
|
||||
t.Fatal("default/custom reader binding migration incorrect")
|
||||
}
|
||||
if err := pg.RemoveAgentFromTool("reporter", "traffic_search"); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
s.seedFindingWorkflowTools()
|
||||
search, _ = pg.GetTool("traffic_search")
|
||||
if contains(search.Agents, "reporter") {
|
||||
t.Fatal("one-time migration undid later unbinding")
|
||||
}
|
||||
}
|
||||
|
||||
func TestFindingWorkflowReporterBindsBeforeWritingReport(t *testing.T) {
|
||||
s, f, request := trafficEvidenceServer(t)
|
||||
pg := s.m.pg
|
||||
ctx := context.Background()
|
||||
if err := pg.SetBool(settingAgentTrafficBinding, true); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
t.Cleanup(func() { pg.SetBool(settingAgentTrafficBinding, false); s.m.SetTrafficEnabled(false) })
|
||||
if err := s.m.SetTrafficEnabled(true); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
seedServerEvidenceFlow(t, s, "reporter-proof", []byte("local proof payload"))
|
||||
seedServerEvidenceFlow(t, s, "reporter-baseline", []byte("local normal response"))
|
||||
tools, def, cleanup := agent.AugmentTools(ctx, "reporter", nil)
|
||||
defer cleanup()
|
||||
if !strings.Contains(def.FindingGuidance, "报告前自动关联流量") || !strings.Contains(def.FindingGuidance, "绑定成功后重新调用") {
|
||||
t.Fatal("reporter did not receive binding workflow")
|
||||
}
|
||||
for _, name := range []string{"traffic_search", "traffic_get", "get_task_worker_trace", "get_task_node_detail", "bind_finding_traffic", "get_finding_traffic", "update_finding_report"} {
|
||||
workflowTool(t, tools, name)
|
||||
}
|
||||
finding, _ := pg.GetFinding(f.FindingID)
|
||||
detail := workflowCall(t, ctx, workflowTool(t, tools, "get_task_node_detail"), map[string]any{"task_id": fmt.Sprint(*finding.TaskID), "id": f.NodeID}, false)
|
||||
if !strings.Contains(detail, `"finding_id"`) || !strings.Contains(detail, `"finding_node_id"`) {
|
||||
t.Fatal("reporter lacks explicit ID mapping")
|
||||
}
|
||||
workflowCall(t, ctx, workflowTool(t, tools, "traffic_search"), map[string]any{"host": "evidence.local"}, false)
|
||||
for _, id := range []string{"reporter-baseline", "reporter-proof"} {
|
||||
packet := workflowCall(t, ctx, workflowTool(t, tools, "traffic_get"), map[string]any{"id": id}, false)
|
||||
if !strings.Contains(packet, "local") {
|
||||
t.Fatal("reporter cannot inspect packet", packet)
|
||||
}
|
||||
}
|
||||
refs := []db.TrafficRef{{TrafficID: "reporter-baseline", Role: "baseline"}, {TrafficID: "reporter-proof", Role: "proof", Note: "confirmed from recorded response"}}
|
||||
workflowCall(t, ctx, workflowTool(t, tools, "bind_finding_traffic"), map[string]any{"finding_id": f.FindingID, "traffic_refs": refs}, false)
|
||||
list := workflowCall(t, ctx, workflowTool(t, tools, "get_finding_traffic"), map[string]any{"finding_id": f.FindingID}, false)
|
||||
var summary struct {
|
||||
Version int64 `json:"version"`
|
||||
}
|
||||
if err := json.Unmarshal([]byte(list), &summary); err != nil || summary.Version != 1 {
|
||||
t.Fatal(list, err)
|
||||
}
|
||||
workflowCall(t, ctx, workflowTool(t, tools, "update_finding_report"), map[string]any{"finding_id": f.NodeID, "evidence_version": summary.Version, "report": "## Local report\n\nVerified baseline and proof using saved evidence."}, false)
|
||||
updated, err := pg.GetFinding(f.FindingID)
|
||||
if err != nil || updated.ReportEvidenceVersion != summary.Version || updated.EvidenceVersion != summary.Version {
|
||||
t.Fatal("report did not cover post-binding version", updated, err)
|
||||
}
|
||||
// Disabling automatic binding still lets Reporter read manually bound snapshots.
|
||||
if r := request("PUT", "/api/settings", `{"agent_traffic_binding":false}`); r.Code != 200 {
|
||||
t.Fatal(r.Code, r.Body)
|
||||
}
|
||||
off, offDef, closeOff := agent.AugmentTools(ctx, "reporter", nil)
|
||||
defer closeOff()
|
||||
if offDef.FindingGuidance != "" {
|
||||
t.Fatal("off reporter still receives auto-binding guidance")
|
||||
}
|
||||
for _, tool := range off {
|
||||
if tool.Name() == "traffic_search" || tool.Name() == "traffic_get" || tool.Name() == "traffic_blob" || tool.Name() == "bind_finding_traffic" {
|
||||
t.Fatal("auto-binding tool exposed while off", tool.Name())
|
||||
}
|
||||
}
|
||||
workflowCall(t, ctx, workflowTool(t, off, "get_finding_traffic"), map[string]any{"finding_id": f.FindingID}, false)
|
||||
workflowTool(t, off, "update_finding_report")
|
||||
}
|
||||
+7
-6
@@ -249,12 +249,13 @@ type Manager struct {
|
||||
|
||||
// Settings keys the UI toggles at runtime.
|
||||
const (
|
||||
settingTrafficCapture = "traffic_capture"
|
||||
settingWebSearchOn = "web_search_enabled"
|
||||
settingWebSearchBackend = "web_search_backend"
|
||||
settingBraveKey = "brave_search_api_key"
|
||||
settingTavilyKey = "tavily_search_api_key"
|
||||
settingWebSearchProxy = "web_search_proxy"
|
||||
settingTrafficCapture = "traffic_capture"
|
||||
settingAgentTrafficBinding = "agent_traffic_binding"
|
||||
settingWebSearchOn = "web_search_enabled"
|
||||
settingWebSearchBackend = "web_search_backend"
|
||||
settingBraveKey = "brave_search_api_key"
|
||||
settingTavilyKey = "tavily_search_api_key"
|
||||
settingWebSearchProxy = "web_search_proxy"
|
||||
// settingGlobalProxy is the global egress proxy for all target traffic
|
||||
// (http/https/socks5). Empty = direct. Distinct from web_search_proxy (which
|
||||
// only routes the search backend) and the per-profile LLM proxy.
|
||||
|
||||
+30
-19
@@ -76,6 +76,8 @@ func (s *Server) orchestrationTools() []actool.CoreTool {
|
||||
s.toolSearchWorkerTraces(),
|
||||
s.toolGetTaskNodeDetail(),
|
||||
s.toolUpdateFindingReport(),
|
||||
s.toolGetFindingTraffic(),
|
||||
s.toolBindFindingTraffic(),
|
||||
}
|
||||
}
|
||||
|
||||
@@ -334,10 +336,11 @@ func (s *Server) toolAddHint() actool.CoreTool {
|
||||
return wrTool("add_task_hint", "给指定任务注入战略提示(该任务的 planner 下轮生成意图时会读到)。\n"+
|
||||
"★优先批量:多条提示放进 hints 数组一次提交(返回 ids 数组,与 hints 等长同序,失败项 id=0);单条则省略 hints 直接给顶层 text。",
|
||||
objSchema(map[string]any{
|
||||
"task_id": strParam("任务 id"),
|
||||
"hints": map[string]any{"type": "array", "description": "【优先用这个】提示数组,每个元素字段同顶层(text/asset_ids)。", "items": map[string]any{"type": "object"}},
|
||||
"text": strParam("[单条] 提示内容"),
|
||||
"asset_ids": map[string]any{"type": "array", "items": map[string]any{"type": "integer"}, "description": "锚定的资产 id(可选,0/1/多个;该任务内的资产 id)"},
|
||||
"task_id": strParam("任务 id"),
|
||||
"hints": map[string]any{"type": "array", "description": "【优先用这个】提示数组,每个元素字段同顶层(text/asset_ids/traffic_refs)。", "items": objSchema(map[string]any{"text": strParam("提示内容"), "asset_ids": map[string]any{"type": "array", "items": map[string]any{"type": "integer"}}, "traffic_refs": agent.HintTrafficSchema()})},
|
||||
"text": strParam("[单条] 提示内容"),
|
||||
"traffic_refs": agent.HintTrafficSchema(),
|
||||
"asset_ids": map[string]any{"type": "array", "items": map[string]any{"type": "integer"}, "description": "锚定的资产 id(可选,0/1/多个;该任务内的资产 id)"},
|
||||
}, "task_id"),
|
||||
func(ctx context.Context, in json.RawMessage) (actool.Result, error) {
|
||||
return s.delegateToTask(ctx, in, (*agent.ToolSet).AddHintTool)
|
||||
@@ -393,20 +396,22 @@ func (s *Server) toolUpdateFindingReport() actool.CoreTool {
|
||||
return wrTool("update_finding_report",
|
||||
"为已登记的漏洞写入/更新【详细报告】(Markdown 全文,整段覆盖旧内容)。finding_id 传 report_finding 返回的那个 id(\"finding recorded: <id>\" 里的数字)。报告建议包含:漏洞概述、影响与危害、复现步骤、证据/PoC、修复建议。",
|
||||
objSchema(map[string]any{
|
||||
"finding_id": map[string]any{"type": "integer", "description": "目标漏洞 id(report_finding 返回的 id)"},
|
||||
"report": strParam("详细报告全文,Markdown 格式"),
|
||||
"finding_id": map[string]any{"type": "integer", "description": "目标漏洞 id(report_finding 返回的 id)"},
|
||||
"report": strParam("详细报告全文,Markdown 格式"),
|
||||
"evidence_version": map[string]any{"type": "integer", "description": "get_finding_traffic 返回的证据 version;用于防止报告覆盖新的证据变更"},
|
||||
}, "finding_id", "report"),
|
||||
func(_ context.Context, in json.RawMessage) (actool.Result, error) {
|
||||
func(ctx context.Context, in json.RawMessage) (actool.Result, error) {
|
||||
var a struct {
|
||||
FindingID json.RawMessage `json:"finding_id"`
|
||||
Report string `json:"report"`
|
||||
EvidenceVersion *int64 `json:"evidence_version"`
|
||||
FindingID json.RawMessage `json:"finding_id"`
|
||||
Report string `json:"report"`
|
||||
}
|
||||
_ = json.Unmarshal(in, &a)
|
||||
nodeID := parseProfileID(a.FindingID) // 复用「数字或数字字符串」解析
|
||||
if nodeID <= 0 {
|
||||
return actool.Errorf("finding_id 无效"), nil
|
||||
}
|
||||
n, err := s.m.pg.SetFindingReportByNodeID(nodeID, a.Report)
|
||||
n, err := s.m.pg.SetFindingReportVersionByNodeID(ctx, nodeID, a.Report, a.EvidenceVersion)
|
||||
if err != nil {
|
||||
return actool.Errorf(err.Error()), nil
|
||||
}
|
||||
@@ -440,7 +445,11 @@ func (s *Server) seedOrchestrationTools() {
|
||||
autoAgents, _ := json.Marshal([]string{"auto"})
|
||||
for _, t := range s.orchestrationTools() {
|
||||
schema, _ := json.Marshal(t.InputSchema())
|
||||
_ = s.m.PG().SeedTool(t.Name(), t.Description(), schema, autoAgents)
|
||||
bindings := autoAgents
|
||||
if t.Name() == "bind_finding_traffic" {
|
||||
bindings = json.RawMessage(`["reporter"]`)
|
||||
}
|
||||
_ = s.m.PG().SeedTool(t.Name(), t.Description(), schema, bindings)
|
||||
}
|
||||
for _, t := range s.platformTools() {
|
||||
schema, _ := json.Marshal(t.InputSchema())
|
||||
@@ -454,11 +463,13 @@ func (s *Server) seedOrchestrationTools() {
|
||||
s.seedWorkerReadToolsUnbind() // list_facts/node_detail/list_companies/跨 work 检索从 worker 默认解绑(一次性)
|
||||
s.seedAutoReportFindingBinding()
|
||||
s.unbindGoalMetDefault()
|
||||
s.reseedGoalsPrompt() // goals 提示词加入「抽操作约束」步 → 旧库追加一版新默认(一次性)
|
||||
s.reseedMainAgentPrompt() // mainagent 提示词加入「目标达成后 add_intent 反问是否建目标」(一次性)
|
||||
s.reseedPlannerPrompt() // planner 提示词:重写「0 意图」正当理由 + 加量化验收核对(一次性)
|
||||
s.reseedWorkerPrompt() // worker 提示词:加否定结论证据门槛(一次性)
|
||||
s.seedReporterAgent() // 预置「报告撰写」agent + 工具绑定 + finding 触发器(一次性)
|
||||
s.reseedGoalsPrompt() // goals 提示词加入「抽操作约束」步 → 旧库追加一版新默认(一次性)
|
||||
s.reseedMainAgentPrompt() // mainagent 提示词加入「目标达成后 add_intent 反问是否建目标」(一次性)
|
||||
s.reseedPlannerPrompt() // planner 提示词:重写「0 意图」正当理由 + 加量化验收核对(一次性)
|
||||
s.reseedWorkerPrompt() // worker 提示词:加否定结论证据门槛(一次性)
|
||||
s.seedReporterAgent() // 预置「报告撰写」agent + 工具绑定 + finding 触发器(一次性)
|
||||
s.seedFindingTrafficTools() // 增加可选证据参数及只读证据工具,保留用户配置
|
||||
s.seedFindingWorkflowTools()
|
||||
// 注:pentest 的默认工具绑定无需迁移——BuiltinToolSeeds 在全新初始化时就把
|
||||
// list_assets/insert_assets/report_finding/list_findings/list_companies 连同
|
||||
// pentest 一起 seed 好了(项目尚无旧库,不做迁移)。
|
||||
@@ -684,9 +695,9 @@ func (s *Server) seedReporterAgent() {
|
||||
Enabled: true,
|
||||
OnToolCall: true,
|
||||
ToolNames: []string{"report_finding"},
|
||||
ToolCallMessage: "上面刚有一个漏洞被 report_finding 登记。请从触发上下文里取出 finding_id" +
|
||||
"(工具返回 \"finding recorded: <id>\" 里的数字)与任务 id,按你的职责撰写该漏洞的详细报告," +
|
||||
"最后调用 update_finding_report(finding_id, report) 保存。",
|
||||
ToolCallMessage: "上面刚有一个漏洞被 report_finding 登记。请读取返回 JSON 的 finding_id(独立漏洞记录 ID)与 finding_node_id(探索节点 ID)," +
|
||||
"若运行指引启用自动绑定,先核实并关联本次漏洞的流量,再用 get_finding_traffic(finding_id) 读取最新证据;节点详情使用 finding_node_id。" +
|
||||
"最后调用 update_finding_report(finding_id=finding_node_id, report, evidence_version=实际读取版本) 保存。不要混用两种编号。",
|
||||
}); err != nil {
|
||||
log.Printf("[reporter] 创建触发器失败: %v", err)
|
||||
}
|
||||
|
||||
@@ -305,6 +305,11 @@ func (sc *Scheduler) fireToolCalls(triggers []*db.AgentTrigger) {
|
||||
if len(want) == 0 {
|
||||
continue
|
||||
}
|
||||
// A failed finding write has no committed finding to report. Keep other
|
||||
// tool-error triggers available for user-defined automation.
|
||||
if e.ToolIsErr && e.Tool == "report_finding" {
|
||||
continue
|
||||
}
|
||||
errTag := ""
|
||||
if e.ToolIsErr {
|
||||
errTag = "[error] "
|
||||
|
||||
+43
-24
@@ -1,7 +1,6 @@
|
||||
package server
|
||||
|
||||
import (
|
||||
"archive/zip"
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
@@ -11,6 +10,7 @@ import (
|
||||
"net"
|
||||
"net/http"
|
||||
"net/url"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"regexp"
|
||||
"sort"
|
||||
@@ -223,6 +223,7 @@ func New(ctx context.Context, m *Manager, skillDir string, dataDir string, keyDi
|
||||
if err := s.seedFindingRetester(); err != nil {
|
||||
log.Printf("[retester] seed: %v", err)
|
||||
}
|
||||
go s.evidenceStore().RunGC(s.ctx)
|
||||
s.seedPythonInterpreter() // 自定义脚本工具:开机检测 python 解释器入库(仅空时)
|
||||
go newScheduler(s).Run(s.ctx) // P3 触发器调度(定时/finding/目标事件),仅自定义 agent
|
||||
// Fill the tool cache for any enabled MCP that has none yet (notably the
|
||||
@@ -363,7 +364,7 @@ func (s *Server) saveLLMConfig(cfg agent.Config) error {
|
||||
ContextWindowK: cfg.ContextWindowK, ThinkingType: cfg.ThinkingType, ReasoningEffort: cfg.ReasoningEffort, IsDefault: true,
|
||||
Priority: priority, PoolExclude: poolExclude, Streaming: streaming,
|
||||
MaxTokens: maxTokens, MaxTokensField: maxTokensField, SessionHeaderKey: sessionHeaderKey,
|
||||
Retry: retry,
|
||||
Retry: retry,
|
||||
})
|
||||
if err != nil {
|
||||
return err
|
||||
@@ -408,6 +409,7 @@ func (s *Server) buildPlannerWorker(pinID *int64, gProv llm.Provider, gCfg agent
|
||||
// the tools-table binding (default = worker), so worker behavior is unchanged.
|
||||
wProv, wCfg := s.providerForAgent("worker", pinID, gProv, gCfg)
|
||||
wk := agent.NewWorker(wProv, wCfg.Model, s.m.dir, tx, wCfg.CompactionWindow(), s.agentMaxTurns("worker"))
|
||||
wk.SetFindingRecorder(s.evidenceStore())
|
||||
wk.SetRunTimeout(time.Duration(s.agentRunSeconds("worker")) * time.Second)
|
||||
wk.SetProxy(s.m.ProxyAddr(), s.m.ProxyCACert())
|
||||
wk.SetMemory(memory.NewStore(filepath.Join(s.m.dir, "memory")))
|
||||
@@ -417,6 +419,7 @@ func (s *Server) buildPlannerWorker(pinID *int64, gProv llm.Provider, gCfg agent
|
||||
wk.SetMaxTokens(maxTokensResolver(wCfg)) // 单次回复输出上限(0 = 不发)
|
||||
pProv, pCfg := s.providerForAgent("planner", pinID, gProv, gCfg)
|
||||
pl := agent.NewPlanner(pProv, pCfg.Model, s.m.dir, tx, pCfg.CompactionWindow(), s.agentMaxTurns("planner"))
|
||||
pl.SetFindingRecorder(s.evidenceStore())
|
||||
pl.SetKillWork(s.engine.KillWork) // planner kill_work → terminate a running work
|
||||
pl.SetSteerWork(s.engine.SteerWork) // planner steer_work → inject mid-run course-correction
|
||||
pl.SetProxy(s.m.ProxyAddr(), s.m.ProxyCACert()) // WebFetch through the recording proxy
|
||||
@@ -485,6 +488,7 @@ func (s *Server) applyLLM(cfg agent.Config) error {
|
||||
mProv, mCfg := s.providerForAgent("mainagent", nil, prov, cfg)
|
||||
s.cfgMu.Lock()
|
||||
s.mainAgent = agent.NewMainAgent(mProv, mCfg.Model, s.m.dir, tx, mCfg.CompactionWindow(), s.agentMaxTurns("mainagent"))
|
||||
s.mainAgent.SetFindingRecorder(s.evidenceStore())
|
||||
s.mainAgent.SetProxy(s.m.ProxyAddr(), s.m.ProxyCACert()) // WebFetch through the recording proxy
|
||||
s.mainAgent.SetWebSearch(s.webSearchFor("mainagent"))
|
||||
s.mainAgent.SetSteerWork(s.engine.SteerWork) // steer_work:人对运行中 work 实时纠偏
|
||||
@@ -814,6 +818,7 @@ func (s *Server) Handler() http.Handler {
|
||||
mux.HandleFunc("GET /api/exploration/findings/asset-tree", s.findingAssetTree)
|
||||
mux.HandleFunc("GET /api/exploration/findings/stats", s.findingStats)
|
||||
mux.HandleFunc("GET /api/exploration/findings/export", s.findingsExport)
|
||||
s.registerFindingTraffic(mux)
|
||||
mux.HandleFunc("GET /api/exploration/findings/{id}", s.getFinding)
|
||||
mux.HandleFunc("GET /api/exploration/findings/{id}/lineage", s.findingLineage)
|
||||
mux.HandleFunc("POST /api/exploration/findings/{id}/deepen", s.deepenFinding)
|
||||
@@ -2142,6 +2147,16 @@ func (s *Server) findingsExport(w http.ResponseWriter, r *http.Request) {
|
||||
return
|
||||
}
|
||||
|
||||
stage, err := os.MkdirTemp("", "artex-finding-export-")
|
||||
if err != nil {
|
||||
writeErr(w, 500, err.Error())
|
||||
return
|
||||
}
|
||||
defer os.RemoveAll(stage)
|
||||
if err = s.evidenceStore().StageFindingsExport(r.Context(), fs, stage, format == "md-zip"); err != nil {
|
||||
evidenceError(w, err)
|
||||
return
|
||||
}
|
||||
now := time.Now()
|
||||
stamp := now.Format("20060102-150405")
|
||||
setDownload := func(contentType, filename string) {
|
||||
@@ -2154,27 +2169,19 @@ func (s *Server) findingsExport(w http.ResponseWriter, r *http.Request) {
|
||||
setDownload("text/markdown; charset=utf-8", "findings-"+stamp+".md")
|
||||
_, _ = w.Write([]byte(report.FindingsMarkdown(fs, now)))
|
||||
case "md-zip":
|
||||
path := filepath.Join(stage, "findings.zip")
|
||||
if err := buildFindingsEvidenceZip(path, fs, stage, now); err != nil {
|
||||
writeErr(w, 500, err.Error())
|
||||
return
|
||||
}
|
||||
file, err := os.Open(path)
|
||||
if err != nil {
|
||||
writeErr(w, 500, err.Error())
|
||||
return
|
||||
}
|
||||
defer file.Close()
|
||||
setDownload("application/zip", "findings-"+stamp+".zip")
|
||||
zw := zip.NewWriter(w)
|
||||
used := map[string]int{}
|
||||
for _, f := range fs {
|
||||
base := report.FindingFilename(f)
|
||||
name := base
|
||||
// 去重:同名文件追加 -2、-3……
|
||||
if n := used[base]; n > 0 {
|
||||
name = fmt.Sprintf("%s-%d.md", strings.TrimSuffix(base, ".md"), n+1)
|
||||
}
|
||||
used[base]++
|
||||
fw, werr := zw.Create(name)
|
||||
if werr != nil {
|
||||
log.Printf("[findings-export] zip create %s: %v", name, werr)
|
||||
continue
|
||||
}
|
||||
_, _ = fw.Write([]byte(report.SingleFindingMarkdown(f, now)))
|
||||
}
|
||||
if cerr := zw.Close(); cerr != nil {
|
||||
log.Printf("[findings-export] zip close: %v", cerr)
|
||||
}
|
||||
http.ServeContent(w, r, "findings.zip", now, file)
|
||||
case "csv":
|
||||
setDownload("text/csv; charset=utf-8", "findings-"+stamp+".csv")
|
||||
_, _ = w.Write(report.FindingsCSV(fs))
|
||||
@@ -2182,6 +2189,10 @@ func (s *Server) findingsExport(w http.ResponseWriter, r *http.Request) {
|
||||
assets := s.resolveFindingAssets(fs)
|
||||
out := make([]FindingDTO, 0, len(fs))
|
||||
for _, f := range fs {
|
||||
for i := range f.TrafficBindings {
|
||||
f.TrafficBindings[i].Snapshot.ReqHead = ""
|
||||
f.TrafficBindings[i].Snapshot.RespHead = ""
|
||||
}
|
||||
out = append(out, findingFromDB(f, assets))
|
||||
}
|
||||
setDownload("application/json; charset=utf-8", "findings-"+stamp+".json")
|
||||
@@ -3129,6 +3140,7 @@ func (s *Server) settingsPayload() map[string]any {
|
||||
}
|
||||
return map[string]any{
|
||||
"traffic_capture": s.m.TrafficEnabled(),
|
||||
"agent_traffic_binding": s.m.pg.GetBool(settingAgentTrafficBinding, false),
|
||||
"llm_record": s.m.LLMRecordEnabled(),
|
||||
"web_search_enabled": on,
|
||||
"web_search_backend": backend,
|
||||
@@ -3169,8 +3181,9 @@ func (s *Server) pgDetectPython(w http.ResponseWriter, r *http.Request) {
|
||||
// off, agents get no proxy config, no traffic tools, and no proxy prompt content.
|
||||
func (s *Server) putSettings(w http.ResponseWriter, r *http.Request) {
|
||||
var req struct {
|
||||
TrafficCapture *bool `json:"traffic_capture"`
|
||||
LLMRecord *bool `json:"llm_record"` // LLM 录制开关(默认关);即时生效,无需重建 agent
|
||||
TrafficCapture *bool `json:"traffic_capture"`
|
||||
AgentTrafficBinding *bool `json:"agent_traffic_binding"`
|
||||
LLMRecord *bool `json:"llm_record"` // LLM 录制开关(默认关);即时生效,无需重建 agent
|
||||
// Web search. WebSearchEnabled/Backend toggle the tool + backend; BraveKey/TavilyKey
|
||||
// are optional — omit (null) to leave a stored key untouched, send "" to clear.
|
||||
WebSearchEnabled *bool `json:"web_search_enabled"`
|
||||
@@ -3267,6 +3280,12 @@ func (s *Server) putSettings(w http.ResponseWriter, r *http.Request) {
|
||||
// built under the OLD switch state — drop them so they pick up the new one.
|
||||
s.invalidateProfileAgents()
|
||||
}
|
||||
if req.AgentTrafficBinding != nil {
|
||||
if err := s.m.pg.SetBool(settingAgentTrafficBinding, *req.AgentTrafficBinding); err != nil {
|
||||
writeErr(w, 500, err.Error())
|
||||
return
|
||||
}
|
||||
}
|
||||
if req.TrafficCapture != nil {
|
||||
if err := s.m.SetTrafficEnabled(*req.TrafficCapture); err != nil {
|
||||
writeErr(w, 500, err.Error())
|
||||
|
||||
+20
-2
@@ -178,6 +178,13 @@ func (s *Server) archiveTask(job *pgdb.TaskArchive) (runErr error) {
|
||||
snapshot.DataCounts["traffic"] = trafficCount
|
||||
}
|
||||
archivePath := taskArchivePath(s.m.dir, job.ID, taskID)
|
||||
evidenceSnapshots, err := pgdb.ArchiveEvidenceSnapshots(snapshot)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if err = s.evidenceStore().CopySnapshots(s.ctx, evidenceSnapshots, filepath.Join(fileStage.payload, "evidence")); err != nil {
|
||||
return err
|
||||
}
|
||||
_ = s.m.pg.UpdateTaskArchiveProgress(job.ID, "compress", 50)
|
||||
originalSize, compressedSize, checksum, err := writeTaskArchivePackage(archivePath, fileStage.payload, snapshot)
|
||||
if err != nil {
|
||||
@@ -270,6 +277,17 @@ func (s *Server) restoreTaskArchive(job *pgdb.TaskArchive) (runErr error) {
|
||||
if snapshot.TaskID != job.TaskID || !pgdb.IsTaskArchiveFormatSupported(snapshot.FormatVersion) {
|
||||
return pgdb.ErrTaskArchiveFormatMismatch
|
||||
}
|
||||
evidenceSnapshots, err := pgdb.ArchiveEvidenceSnapshots(&snapshot)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
return s.evidenceStore().WithInstalledSnapshots(s.ctx, evidenceSnapshots, filepath.Join(extracted, "evidence"), func() error {
|
||||
return s.restoreTaskArchivePayload(job, &snapshot, extracted)
|
||||
})
|
||||
}
|
||||
|
||||
func (s *Server) restoreTaskArchivePayload(job *pgdb.TaskArchive, snapshot *pgdb.TaskArchiveSnapshot, extracted string) (runErr error) {
|
||||
var err error
|
||||
var llmRecordsFile *os.File
|
||||
relative, hasStreamedLLMRecords := snapshot.StreamedTables["llm_records"]
|
||||
if len(snapshot.StreamedTables) > 0 &&
|
||||
@@ -309,13 +327,13 @@ func (s *Server) restoreTaskArchive(job *pgdb.TaskArchive) (runErr error) {
|
||||
var warnings []string
|
||||
if llmRecordsFile != nil {
|
||||
warnings, err = s.m.pg.RestoreTaskArchiveWithLLMRecords(
|
||||
job.ID, &snapshot, job.RemainingTimeoutSeconds, llmRecordsFile,
|
||||
job.ID, snapshot, job.RemainingTimeoutSeconds, llmRecordsFile,
|
||||
)
|
||||
closeErr := llmRecordsFile.Close()
|
||||
llmRecordsFile = nil
|
||||
err = errors.Join(err, closeErr)
|
||||
} else {
|
||||
warnings, err = s.m.pg.RestoreTaskArchive(job.ID, &snapshot, job.RemainingTimeoutSeconds)
|
||||
warnings, err = s.m.pg.RestoreTaskArchive(job.ID, snapshot, job.RemainingTimeoutSeconds)
|
||||
}
|
||||
if err != nil {
|
||||
return err
|
||||
|
||||
@@ -563,6 +563,7 @@ func (s *Server) agentsForTask(t *Task) *taskAgentBundle {
|
||||
tx := transcript.NewStore(filepath.Join(s.m.dir, "transcripts"))
|
||||
window := workerRuntime.CompactionWindow()
|
||||
wk := agent.NewWorker(workerRuntime, "task-router", s.m.dir, tx, window, s.agentMaxTurns("worker"))
|
||||
wk.SetFindingRecorder(s.evidenceStore())
|
||||
wk.SetCompactionWindowResolver(workerRuntime.CompactionWindow)
|
||||
wk.SetNonStreaming(workerRuntime.nonStreaming) // 按任务当前激活 profile 的流式开关(每轮读)
|
||||
wk.SetMaxTokens(workerRuntime.maxTokens) // 同上,输出上限也跟随当前激活 profile
|
||||
@@ -572,6 +573,7 @@ func (s *Server) agentsForTask(t *Task) *taskAgentBundle {
|
||||
wk.SetWebSearch(s.webSearchFor("worker"))
|
||||
wk.SetConstraintInject(s.constraintInjectWorker) // 操作约束注入 worker(可配置,默认开)
|
||||
pl := agent.NewPlanner(plannerRuntime, "task-router", s.m.dir, tx, plannerRuntime.CompactionWindow(), s.agentMaxTurns("planner"))
|
||||
pl.SetFindingRecorder(s.evidenceStore())
|
||||
pl.SetCompactionWindowResolver(plannerRuntime.CompactionWindow)
|
||||
pl.SetNonStreaming(plannerRuntime.nonStreaming)
|
||||
pl.SetMaxTokens(plannerRuntime.maxTokens)
|
||||
@@ -581,6 +583,7 @@ func (s *Server) agentsForTask(t *Task) *taskAgentBundle {
|
||||
pl.SetWebSearch(s.webSearchFor("planner"))
|
||||
pl.SetConstraintInject(s.constraintInjectPlanner) // 操作约束注入 planner(可配置,默认开)
|
||||
main := agent.NewMainAgent(mainRuntime, "task-router", s.m.dir, tx, mainRuntime.CompactionWindow(), s.agentMaxTurns("mainagent"))
|
||||
main.SetFindingRecorder(s.evidenceStore())
|
||||
main.SetCompactionWindowResolver(mainRuntime.CompactionWindow)
|
||||
main.SetNonStreaming(mainRuntime.nonStreaming)
|
||||
main.SetMaxTokens(mainRuntime.maxTokens)
|
||||
|
||||
@@ -0,0 +1,126 @@
|
||||
package traffic
|
||||
|
||||
import (
|
||||
"bufio"
|
||||
"bytes"
|
||||
"context"
|
||||
"database/sql"
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
)
|
||||
|
||||
// EvidenceExchange exposes complete captured bodies only during the callback.
|
||||
// Readers are closed before releasing the writer lock, so host deletion cannot
|
||||
// remove a blob between its index lookup and its evidence copy.
|
||||
type EvidenceExchange struct {
|
||||
ID, URL, Method, ContentType string
|
||||
TS int64
|
||||
Status int
|
||||
ReqHead, RespHead string
|
||||
ReqLen, RespLen int64
|
||||
ReqHash, RespHash string
|
||||
Request, Response io.Reader
|
||||
}
|
||||
|
||||
func (t *Traffic) ReadEvidence(ctx context.Context, ids []string, consume func(EvidenceExchange) error) error {
|
||||
if t == nil {
|
||||
return errors.New("流量录制存储不可用")
|
||||
}
|
||||
t.wmu.Lock()
|
||||
defer t.wmu.Unlock()
|
||||
for _, id := range ids {
|
||||
if err := ctx.Err(); err != nil {
|
||||
return err
|
||||
}
|
||||
if err := t.readEvidence(id, consume); err != nil {
|
||||
return fmt.Errorf("流量 %s: %w", id, err)
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (t *Traffic) readEvidence(id string, consume func(EvidenceExchange) error) error {
|
||||
e := EvidenceExchange{ID: id}
|
||||
var legacy string
|
||||
if err := t.db.QueryRow(`SELECT ts,url,method,status,content_type,req_len,resp_len,path FROM exchanges WHERE id=?`, id).
|
||||
Scan(&e.TS, &e.URL, &e.Method, &e.Status, &e.ContentType, &e.ReqLen, &e.RespLen, &legacy); err != nil {
|
||||
return err
|
||||
}
|
||||
var req, resp []byte
|
||||
var reqBlob, respBlob sql.NullString
|
||||
err := t.db.QueryRow(`SELECT req_head,req_body,req_blob,resp_head,resp_body,resp_blob FROM exchange_bodies WHERE id=?`, id).
|
||||
Scan(&e.ReqHead, &req, &reqBlob, &e.RespHead, &resp, &respBlob)
|
||||
if errors.Is(err, sql.ErrNoRows) && legacy != "" {
|
||||
if !filepath.IsLocal(legacy) {
|
||||
return errors.New("旧流量路径无效")
|
||||
}
|
||||
r, err := os.Open(filepath.Join(t.dir, legacy, "request.http"))
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer r.Close()
|
||||
p, err := os.Open(filepath.Join(t.dir, legacy, "response.http"))
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer p.Close()
|
||||
e.ReqHead, e.Request, err = splitLegacyEvidence(r)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
e.RespHead, e.Response, err = splitLegacyEvidence(p)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
return consume(e)
|
||||
}
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
e.Request = bytes.NewReader(req)
|
||||
e.Response = bytes.NewReader(resp)
|
||||
if reqBlob.Valid && reqBlob.String != "" {
|
||||
f, _, err := t.Blob(reqBlob.String)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer f.Close()
|
||||
e.Request = f
|
||||
e.ReqHash = reqBlob.String
|
||||
}
|
||||
if respBlob.Valid && respBlob.String != "" {
|
||||
f, _, err := t.Blob(respBlob.String)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer f.Close()
|
||||
e.Response = f
|
||||
e.RespHash = respBlob.String
|
||||
}
|
||||
return consume(e)
|
||||
}
|
||||
|
||||
func splitLegacyEvidence(r io.Reader) (string, io.Reader, error) {
|
||||
b := bufio.NewReader(r)
|
||||
var head strings.Builder
|
||||
for {
|
||||
line, err := b.ReadString('\n')
|
||||
if line == "\n" || line == "\r\n" {
|
||||
return head.String(), b, nil
|
||||
}
|
||||
head.WriteString(line)
|
||||
if head.Len() > 1<<20 {
|
||||
return "", nil, errors.New("旧流量报文头过大")
|
||||
}
|
||||
if errors.Is(err, io.EOF) {
|
||||
return head.String(), b, nil
|
||||
}
|
||||
if err != nil {
|
||||
return "", nil, err
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -19,11 +19,16 @@ export default function LoginPage() {
|
||||
|
||||
useEffect(() => {
|
||||
// 已登录直接进主界面(静态导出下无 middleware 代劳这层跳转)。
|
||||
if (auth.getToken()) {
|
||||
router.replace("/function/tasks");
|
||||
const token = auth.getToken();
|
||||
if (token) {
|
||||
// localStorage 可能仍有凭据但 cookie 已丢失。先同步,再发起全新请求,
|
||||
// 避免服务端守卫或路由缓存把跳转送回仍处于 checking 状态的登录页。
|
||||
auth.setToken(token);
|
||||
window.location.replace("/function/tasks");
|
||||
return;
|
||||
}
|
||||
api.authStatus()
|
||||
api
|
||||
.authStatus()
|
||||
.then(({ initialized }) => {
|
||||
if (!initialized) router.replace("/setup");
|
||||
})
|
||||
@@ -38,7 +43,7 @@ export default function LoginPage() {
|
||||
try {
|
||||
const { token } = await api.login("ARTEX", password);
|
||||
auth.setToken(token);
|
||||
router.replace("/function/tasks");
|
||||
window.location.replace("/function/tasks");
|
||||
} catch {
|
||||
setError("用户名或密码错误");
|
||||
} finally {
|
||||
@@ -46,7 +51,13 @@ export default function LoginPage() {
|
||||
}
|
||||
}
|
||||
|
||||
if (checking) return null;
|
||||
if (checking) {
|
||||
return (
|
||||
<div role="status" className="flex min-h-dvh items-center justify-center text-muted-foreground">
|
||||
正在检查登录状态…
|
||||
</div>
|
||||
);
|
||||
}
|
||||
|
||||
return (
|
||||
<div className="flex h-dvh">
|
||||
@@ -57,13 +68,7 @@ export default function LoginPage() {
|
||||
<div className="absolute size-60 rounded-full border border-primary-foreground/15" />
|
||||
<div className="absolute size-40 rounded-full border border-primary-foreground/20" />
|
||||
{/* eslint-disable-next-line @next/next/no-img-element */}
|
||||
<img
|
||||
src="/logo.png"
|
||||
alt="ARTEX"
|
||||
width={160}
|
||||
height={160}
|
||||
className="relative brightness-0 invert"
|
||||
/>
|
||||
<img src="/logo.png" alt="ARTEX" width={160} height={160} className="relative brightness-0 invert" />
|
||||
</div>
|
||||
</div>
|
||||
|
||||
|
||||
@@ -925,6 +925,7 @@ export default function ChatPage() {
|
||||
const [agents, setAgents] = React.useState<Agent[]>([]);
|
||||
const [profiles, setProfiles] = React.useState<LLMProfile[]>([]);
|
||||
const [convs, setConvs] = React.useState<Conversation[]>([]);
|
||||
const [agentFilter, setAgentFilter] = React.useState<string | null>(null);
|
||||
const [selectedId, setSelectedId] = React.useState<number | null>(null);
|
||||
const [renamingId, setRenamingId] = React.useState<number | null>(null);
|
||||
const [renameText, setRenameText] = React.useState("");
|
||||
@@ -1015,17 +1016,44 @@ export default function ChatPage() {
|
||||
|
||||
const selected = React.useMemo(() => convs.find((c) => c.id === selectedId) ?? null, [convs, selectedId]);
|
||||
const agentByKey = React.useMemo(() => new Map(agents.map((agent) => [agent.key, agent])), [agents]);
|
||||
const filteredConversations = React.useMemo(
|
||||
() => (agentFilter === null ? convs : convs.filter((conversation) => conversation.agent_key === agentFilter)),
|
||||
[convs, agentFilter],
|
||||
);
|
||||
const visibleConversations = React.useMemo(
|
||||
() => convs.slice(0, visibleConversationCount),
|
||||
[convs, visibleConversationCount],
|
||||
() => filteredConversations.slice(0, visibleConversationCount),
|
||||
[filteredConversations, visibleConversationCount],
|
||||
);
|
||||
// conversation agents: custom agents + conversational built-ins (role=assistant,
|
||||
// e.g. Auto / 渗透测试). The orchestration built-ins (goals/planner/mainagent/worker)
|
||||
// are task-specific and stay hidden from the chat page.
|
||||
const chatAgents = React.useMemo(() => agents.filter((a) => !a.builtin || a.role === "assistant"), [agents]);
|
||||
const agentFilterOptions = React.useMemo(() => {
|
||||
const counts = new Map<string, number>();
|
||||
for (const conversation of convs) {
|
||||
counts.set(conversation.agent_key, (counts.get(conversation.agent_key) ?? 0) + 1);
|
||||
}
|
||||
// Include historical sources even if their Agent has since been disabled or deleted.
|
||||
const keys = new Set([...chatAgents.map((agent) => agent.key), ...counts.keys()]);
|
||||
if (agentFilter !== null) keys.add(agentFilter);
|
||||
return [...keys]
|
||||
.map((key) => ({ key, name: agentByKey.get(key)?.name || key, count: counts.get(key) ?? 0 }))
|
||||
.sort((a, b) => a.name.localeCompare(b.name, "zh-CN"));
|
||||
}, [convs, chatAgents, agentByKey, agentFilter]);
|
||||
const conversationCountLabel =
|
||||
agentFilter === null ? `共 ${convs.length} 个` : `${filteredConversations.length} / ${convs.length} 个`;
|
||||
|
||||
function changeAgentFilter(key: string | null) {
|
||||
setAgentFilter(key);
|
||||
setVisibleConversationCount(CONVERSATION_LIST_PAGE);
|
||||
setSelectedConversationIds(new Set());
|
||||
setBulkDeleteOpen(false);
|
||||
setRenamingId(null);
|
||||
}
|
||||
|
||||
const selectedConversationCount = selectedConversationIds.size;
|
||||
const allConversationsSelected = convs.length > 0 && selectedConversationCount === convs.length;
|
||||
const allConversationsSelected =
|
||||
filteredConversations.length > 0 && selectedConversationCount === filteredConversations.length;
|
||||
const someConversationsSelected = selectedConversationCount > 0 && !allConversationsSelected;
|
||||
let conversationHeaderChecked: boolean | "indeterminate" = false;
|
||||
if (allConversationsSelected) conversationHeaderChecked = true;
|
||||
@@ -1041,7 +1069,9 @@ export default function ChatPage() {
|
||||
}, []);
|
||||
|
||||
function toggleAllConversations(checked: boolean) {
|
||||
setSelectedConversationIds(checked ? new Set(convs.map((conversation) => conversation.id)) : new Set());
|
||||
setSelectedConversationIds(
|
||||
checked ? new Set(filteredConversations.map((conversation) => conversation.id)) : new Set(),
|
||||
);
|
||||
}
|
||||
|
||||
function exitSelectionMode() {
|
||||
@@ -1154,16 +1184,37 @@ export default function ChatPage() {
|
||||
<Button size="sm" className="w-full" onClick={() => setSelectedId(null)}>
|
||||
<PlusIcon /> 新建对话
|
||||
</Button>
|
||||
<Select
|
||||
value={agentFilter === null ? "all" : `agent:${agentFilter}`}
|
||||
onValueChange={(value) => changeAgentFilter(value === "all" ? null : value.slice(6))}
|
||||
disabled={bulkDeleting}
|
||||
>
|
||||
<SelectTrigger size="sm" className="w-full min-w-0" aria-label="按 Agent 筛选对话">
|
||||
<Bot />
|
||||
<SelectValue placeholder="全部 Agent" />
|
||||
</SelectTrigger>
|
||||
<SelectContent>
|
||||
<SelectGroup>
|
||||
<SelectItem value="all">全部 Agent</SelectItem>
|
||||
{agentFilterOptions.map((agent) => (
|
||||
<SelectItem key={agent.key} value={`agent:${agent.key}`}>
|
||||
{agent.name}({agent.count})
|
||||
</SelectItem>
|
||||
))}
|
||||
</SelectGroup>
|
||||
</SelectContent>
|
||||
</Select>
|
||||
{convs.length > 0 &&
|
||||
(selectionMode ? (
|
||||
<div className="flex items-center gap-2 px-1">
|
||||
<Checkbox
|
||||
checked={conversationHeaderChecked}
|
||||
onCheckedChange={(checked) => toggleAllConversations(checked === true)}
|
||||
aria-label="选择全部对话"
|
||||
aria-label="选择当前筛选的全部对话"
|
||||
disabled={filteredConversations.length === 0 || bulkDeleting}
|
||||
/>
|
||||
<span className="text-muted-foreground min-w-0 flex-1 text-xs tabular-nums">
|
||||
{selectedConversationCount > 0 ? `已选 ${selectedConversationCount} 个` : `共 ${convs.length} 个`}
|
||||
{selectedConversationCount > 0 ? `已选 ${selectedConversationCount} 个` : conversationCountLabel}
|
||||
</span>
|
||||
{selectedConversationCount > 0 && (
|
||||
<Button
|
||||
@@ -1183,13 +1234,14 @@ export default function ChatPage() {
|
||||
) : (
|
||||
<div className="flex items-center gap-2 px-1">
|
||||
<span className="text-muted-foreground min-w-0 flex-1 text-xs tabular-nums">
|
||||
共 {convs.length} 个
|
||||
{conversationCountLabel}
|
||||
</span>
|
||||
<Button
|
||||
size="sm"
|
||||
variant="ghost"
|
||||
className="text-muted-foreground"
|
||||
onClick={() => setSelectionMode(true)}
|
||||
disabled={filteredConversations.length === 0}
|
||||
>
|
||||
<ListChecksIcon data-icon="inline-start" />
|
||||
多选
|
||||
@@ -1197,9 +1249,17 @@ export default function ChatPage() {
|
||||
</div>
|
||||
))}
|
||||
</div>
|
||||
<ScrollArea type="auto" className="min-h-0 min-w-0 flex-1 [&_[data-slot=scroll-area-viewport]>div]:block!">
|
||||
<ScrollArea
|
||||
key={agentFilter === null ? "all" : `agent:${agentFilter}`}
|
||||
type="auto"
|
||||
className="min-h-0 min-w-0 flex-1 [&_[data-slot=scroll-area-viewport]>div]:block!"
|
||||
>
|
||||
<div className="flex min-w-0 flex-col gap-0.5 p-2">
|
||||
{convs.length === 0 && <p className="text-muted-foreground px-2 py-6 text-center text-xs">暂无对话</p>}
|
||||
{filteredConversations.length === 0 && (
|
||||
<p className="text-muted-foreground px-2 py-6 text-center text-xs">
|
||||
{agentFilter === null ? "暂无对话" : "该 Agent 暂无对话"}
|
||||
</p>
|
||||
)}
|
||||
{visibleConversations.map((c) => (
|
||||
<ConversationItem
|
||||
key={c.id}
|
||||
@@ -1220,7 +1280,7 @@ export default function ChatPage() {
|
||||
onSelectedForDeleteChange={toggleConversationSelected}
|
||||
/>
|
||||
))}
|
||||
{visibleConversationCount < convs.length && (
|
||||
{visibleConversationCount < filteredConversations.length && (
|
||||
<Button
|
||||
type="button"
|
||||
variant="ghost"
|
||||
@@ -1252,6 +1312,7 @@ export default function ChatPage() {
|
||||
agents={chatAgents}
|
||||
profiles={profiles}
|
||||
onStarted={(c, pending) => {
|
||||
if (agentFilter !== null && agentFilter !== c.agent_key) changeAgentFilter(null);
|
||||
// Insert the new conversation immediately so `selected` resolves to
|
||||
// it on this render (switching to ChatView right away, before the
|
||||
// async reloadConvs lands); reloadConvs then reconciles titles etc.
|
||||
|
||||
@@ -28,6 +28,7 @@ import {
|
||||
AlertDialogTitle,
|
||||
AlertDialogTrigger,
|
||||
} from "@/components/ui/alert-dialog";
|
||||
import { Badge } from "@/components/ui/badge";
|
||||
import { Button } from "@/components/ui/button";
|
||||
import { Checkbox } from "@/components/ui/checkbox";
|
||||
import { Input } from "@/components/ui/input";
|
||||
@@ -212,6 +213,7 @@ export function FindingsTable({
|
||||
<span className="truncate font-medium">{f.name || f.vulnclass || "未分类"}</span>
|
||||
)}
|
||||
<span className="truncate text-xs text-muted-foreground">{f.summary}</span>
|
||||
<Badge variant="outline">流量证据 {f.traffic_count ?? 0} 条</Badge>
|
||||
</div>
|
||||
</TableCell>
|
||||
<TableCell>
|
||||
|
||||
@@ -1,33 +1,32 @@
|
||||
"use client";
|
||||
|
||||
import * as React from "react";
|
||||
|
||||
import Link from "next/link";
|
||||
import { useSearchParams } from "next/navigation";
|
||||
import { toast } from "sonner";
|
||||
import { ArrowLeftIcon, ArrowUpRightIcon, ShieldAlertIcon } from "lucide-react";
|
||||
|
||||
import { SidebarTrigger } from "@/components/ui/sidebar";
|
||||
import { Separator } from "@/components/ui/separator";
|
||||
import { ArrowLeftIcon, ArrowUpRightIcon, ShieldAlertIcon } from "lucide-react";
|
||||
import { toast } from "sonner";
|
||||
|
||||
import { CopyButton } from "@/components/copy-button";
|
||||
import { FindingRetestPanel } from "@/components/finding-retest-panel";
|
||||
import { FindingTrafficPanel } from "@/components/finding-traffic-panel";
|
||||
import { Markdown } from "@/components/markdown";
|
||||
import { StatusBadge } from "@/components/status-badge";
|
||||
import { Alert, AlertDescription } from "@/components/ui/alert";
|
||||
import { Badge } from "@/components/ui/badge";
|
||||
import { Button } from "@/components/ui/button";
|
||||
import { Card, CardContent, CardHeader, CardTitle } from "@/components/ui/card";
|
||||
import { Select, SelectContent, SelectGroup, SelectItem, SelectTrigger } from "@/components/ui/select";
|
||||
import { Separator } from "@/components/ui/separator";
|
||||
import { SidebarTrigger } from "@/components/ui/sidebar";
|
||||
import { Tabs, TabsContent, TabsList, TabsTrigger } from "@/components/ui/tabs";
|
||||
import {
|
||||
Select,
|
||||
SelectContent,
|
||||
SelectGroup,
|
||||
SelectItem,
|
||||
SelectTrigger,
|
||||
} from "@/components/ui/select";
|
||||
import { StatusBadge } from "@/components/status-badge";
|
||||
import { CopyButton } from "@/components/copy-button";
|
||||
import { Markdown } from "@/components/markdown";
|
||||
import { FindingRetestPanel } from "@/components/finding-retest-panel";
|
||||
import { FindingLineageView } from "./lineage";
|
||||
import { statusMeta } from "@/lib/status";
|
||||
import { api } from "@/lib/api";
|
||||
import { statusMeta } from "@/lib/status";
|
||||
import type { Finding, FindingStatus, Severity } from "@/lib/types";
|
||||
|
||||
import { FindingLineageView } from "./lineage";
|
||||
|
||||
const SEVERITIES: Severity[] = ["critical", "high", "medium", "low"];
|
||||
const FINDING_STATUSES: FindingStatus[] = [
|
||||
"pending",
|
||||
@@ -46,21 +45,11 @@ function fmtTime(ts: string) {
|
||||
}
|
||||
|
||||
// FieldRow is one label/value line in the right-hand status panel.
|
||||
function FieldRow({
|
||||
label,
|
||||
children,
|
||||
}: {
|
||||
label: string;
|
||||
children: React.ReactNode;
|
||||
}) {
|
||||
function FieldRow({ label, children }: { label: string; children: React.ReactNode }) {
|
||||
return (
|
||||
<div className="flex items-start justify-between gap-3 py-2.5">
|
||||
<span className="shrink-0 pt-0.5 text-xs text-muted-foreground">
|
||||
{label}
|
||||
</span>
|
||||
<div className="flex min-w-0 flex-col items-end gap-1 text-right text-sm">
|
||||
{children}
|
||||
</div>
|
||||
<span className="shrink-0 pt-0.5 text-xs text-muted-foreground">{label}</span>
|
||||
<div className="flex min-w-0 flex-col items-end gap-1 text-right text-sm">{children}</div>
|
||||
</div>
|
||||
);
|
||||
}
|
||||
@@ -125,9 +114,7 @@ function FindingDetailInner() {
|
||||
if (!finding) {
|
||||
return (
|
||||
<div className="flex flex-1 flex-col items-center justify-center gap-3 p-10 text-center">
|
||||
<p className="text-muted-foreground">
|
||||
{loaded ? `未找到发现 ${id}` : "加载中…"}
|
||||
</p>
|
||||
<p className="text-muted-foreground">{loaded ? `未找到发现 ${id}` : "加载中…"}</p>
|
||||
{loaded && (
|
||||
<Button asChild variant="outline">
|
||||
<Link href="/function/findings">
|
||||
@@ -142,11 +129,7 @@ function FindingDetailInner() {
|
||||
const title = finding.name || finding.vulnclass || "未分类";
|
||||
|
||||
return (
|
||||
<Tabs
|
||||
value={tab}
|
||||
onValueChange={setTab}
|
||||
className="flex flex-1 flex-col gap-0"
|
||||
>
|
||||
<Tabs value={tab} onValueChange={setTab} className="flex flex-1 flex-col gap-0">
|
||||
{/* Sticky header */}
|
||||
<header className="sticky top-0 z-10 flex flex-col gap-2 border-b bg-background/95 px-4 py-2.5 backdrop-blur lg:px-6">
|
||||
<div className="flex flex-wrap items-center gap-2">
|
||||
@@ -160,16 +143,12 @@ function FindingDetailInner() {
|
||||
<h1 className="max-w-md truncate text-sm font-semibold" title={title}>
|
||||
{title}
|
||||
</h1>
|
||||
<code className="rounded bg-muted px-1.5 py-0.5 font-mono text-xs text-muted-foreground">
|
||||
#{finding.id}
|
||||
</code>
|
||||
<code className="rounded bg-muted px-1.5 py-0.5 font-mono text-xs text-muted-foreground">#{finding.id}</code>
|
||||
<Separator orientation="vertical" className="mx-1 h-4" />
|
||||
<StatusBadge domain="severity" value={finding.severity} dot />
|
||||
<StatusBadge domain="finding" value={finding.status} dot />
|
||||
{finding.inherited && finding.source_task_id && (
|
||||
<Badge variant="outline">
|
||||
来源任务 #{finding.source_task_id} · 只读
|
||||
</Badge>
|
||||
<Badge variant="outline">来源任务 #{finding.source_task_id} · 只读</Badge>
|
||||
)}
|
||||
</div>
|
||||
<TabsList>
|
||||
@@ -190,9 +169,7 @@ function FindingDetailInner() {
|
||||
<CardTitle className="text-sm">摘要</CardTitle>
|
||||
</CardHeader>
|
||||
<CardContent>
|
||||
<p className="text-sm leading-relaxed whitespace-pre-wrap">
|
||||
{finding.summary || "(无摘要)"}
|
||||
</p>
|
||||
<p className="text-sm leading-relaxed whitespace-pre-wrap">{finding.summary || "(无摘要)"}</p>
|
||||
</CardContent>
|
||||
</Card>
|
||||
<FindingRetestPanel key={id} findingId={id} readOnly={finding.inherited} onCompleted={load} />
|
||||
@@ -210,21 +187,29 @@ function FindingDetailInner() {
|
||||
)}
|
||||
</CardContent>
|
||||
</Card>
|
||||
<FindingTrafficPanel
|
||||
key={id}
|
||||
findingId={id}
|
||||
contextTask={contextTaskId || undefined}
|
||||
readOnly={finding.inherited}
|
||||
onChanged={load}
|
||||
/>
|
||||
{/* 证据下方:详细报告(Markdown 渲染) */}
|
||||
<Card>
|
||||
<CardHeader className="flex-row items-center justify-between">
|
||||
<CardTitle className="text-sm">详细报告</CardTitle>
|
||||
{finding.report && (
|
||||
<CopyButton text={finding.report} successMessage="已复制详细报告" />
|
||||
)}
|
||||
{finding.report && <CopyButton text={finding.report} successMessage="已复制详细报告" />}
|
||||
</CardHeader>
|
||||
<CardContent>
|
||||
{finding.report_stale ? (
|
||||
<Alert>
|
||||
<AlertDescription>流量证据已变更,详细报告待更新。</AlertDescription>
|
||||
</Alert>
|
||||
) : null}
|
||||
{finding.report ? (
|
||||
<Markdown text={finding.report} />
|
||||
) : (
|
||||
<p className="text-sm text-muted-foreground">
|
||||
暂无详细报告。
|
||||
</p>
|
||||
<p className="text-sm text-muted-foreground">暂无详细报告。</p>
|
||||
)}
|
||||
</CardContent>
|
||||
</Card>
|
||||
@@ -248,14 +233,8 @@ function FindingDetailInner() {
|
||||
{finding.inherited ? (
|
||||
<StatusBadge domain="severity" value={finding.severity} dot />
|
||||
) : (
|
||||
<Select
|
||||
value={finding.severity}
|
||||
onValueChange={(v) => changeSeverity(v as Severity)}
|
||||
>
|
||||
<SelectTrigger
|
||||
size="sm"
|
||||
className="h-7 w-auto border-none px-1 shadow-none focus-visible:ring-0"
|
||||
>
|
||||
<Select value={finding.severity} onValueChange={(v) => changeSeverity(v as Severity)}>
|
||||
<SelectTrigger size="sm" className="h-7 w-auto border-none px-1 shadow-none focus-visible:ring-0">
|
||||
<StatusBadge domain="severity" value={finding.severity} dot />
|
||||
</SelectTrigger>
|
||||
<SelectContent position="popper" align="end">
|
||||
@@ -276,14 +255,8 @@ function FindingDetailInner() {
|
||||
{finding.inherited ? (
|
||||
<StatusBadge domain="finding" value={finding.status} dot />
|
||||
) : (
|
||||
<Select
|
||||
value={finding.status}
|
||||
onValueChange={(v) => changeStatus(v as FindingStatus)}
|
||||
>
|
||||
<SelectTrigger
|
||||
size="sm"
|
||||
className="h-7 w-auto border-none px-1 shadow-none focus-visible:ring-0"
|
||||
>
|
||||
<Select value={finding.status} onValueChange={(v) => changeStatus(v as FindingStatus)}>
|
||||
<SelectTrigger size="sm" className="h-7 w-auto border-none px-1 shadow-none focus-visible:ring-0">
|
||||
<StatusBadge domain="finding" value={finding.status} dot />
|
||||
</SelectTrigger>
|
||||
<SelectContent position="popper" align="end">
|
||||
@@ -302,9 +275,7 @@ function FindingDetailInner() {
|
||||
{/* 漏洞类型 */}
|
||||
<FieldRow label="漏洞类型">
|
||||
{finding.vulnclass ? (
|
||||
<code className="rounded bg-muted px-1.5 py-0.5 font-mono text-xs">
|
||||
{finding.vulnclass}
|
||||
</code>
|
||||
<code className="rounded bg-muted px-1.5 py-0.5 font-mono text-xs">{finding.vulnclass}</code>
|
||||
) : (
|
||||
<span className="text-muted-foreground">—</span>
|
||||
)}
|
||||
@@ -337,9 +308,7 @@ function FindingDetailInner() {
|
||||
className="inline-flex max-w-[16rem] items-center gap-1 text-primary hover:underline"
|
||||
title={finding.task_description}
|
||||
>
|
||||
<span className="truncate">
|
||||
{finding.task_description || `#${finding.task_id}`}
|
||||
</span>
|
||||
<span className="truncate">{finding.task_description || `#${finding.task_id}`}</span>
|
||||
<ArrowUpRightIcon className="size-3 shrink-0" />
|
||||
</Link>
|
||||
) : (
|
||||
|
||||
@@ -10,13 +10,7 @@ import { toast } from "sonner";
|
||||
import { StatusBadge } from "@/components/status-badge";
|
||||
import { Badge } from "@/components/ui/badge";
|
||||
import { Card, CardContent } from "@/components/ui/card";
|
||||
import {
|
||||
Select,
|
||||
SelectContent,
|
||||
SelectGroup,
|
||||
SelectItem,
|
||||
SelectTrigger,
|
||||
} from "@/components/ui/select";
|
||||
import { Select, SelectContent, SelectGroup, SelectItem, SelectTrigger } from "@/components/ui/select";
|
||||
import { api } from "@/lib/api";
|
||||
import { useStoredSortPreference } from "@/lib/sort-preference";
|
||||
import { statusMeta } from "@/lib/status";
|
||||
@@ -65,10 +59,7 @@ function Row({
|
||||
className="flex min-w-0 flex-1 items-center gap-3 text-left"
|
||||
>
|
||||
<ChevronRightIcon
|
||||
className={cn(
|
||||
"size-4 shrink-0 text-muted-foreground transition-transform",
|
||||
open && "rotate-90",
|
||||
)}
|
||||
className={cn("size-4 shrink-0 text-muted-foreground transition-transform", open && "rotate-90")}
|
||||
/>
|
||||
<StatusBadge domain="severity" value={f.severity} dot />
|
||||
<div className="flex min-w-0 flex-1 flex-col gap-1">
|
||||
@@ -80,11 +71,10 @@ function Row({
|
||||
</Badge>
|
||||
)}
|
||||
</div>
|
||||
<span className="truncate text-xs text-muted-foreground">
|
||||
{f.summary}
|
||||
</span>
|
||||
<span className="truncate text-xs text-muted-foreground">{f.summary}</span>
|
||||
</div>
|
||||
</button>
|
||||
<Badge variant="outline">流量证据 {f.traffic_count ?? 0} 条</Badge>
|
||||
{f.assets && f.assets.length > 0 && (
|
||||
<div className="hidden shrink-0 flex-wrap justify-end gap-1 sm:flex">
|
||||
{f.assets.slice(0, 2).map((a) => (
|
||||
@@ -96,22 +86,12 @@ function Row({
|
||||
{a.label}
|
||||
</code>
|
||||
))}
|
||||
{f.assets.length > 2 && (
|
||||
<span className="text-xs text-muted-foreground">
|
||||
+{f.assets.length - 2}
|
||||
</span>
|
||||
)}
|
||||
{f.assets.length > 2 && <span className="text-xs text-muted-foreground">+{f.assets.length - 2}</span>}
|
||||
</div>
|
||||
)}
|
||||
{f.finding_id && !f.inherited ? (
|
||||
<Select
|
||||
value={f.status}
|
||||
onValueChange={(v) => onStatus(f, v as FindingStatus)}
|
||||
>
|
||||
<SelectTrigger
|
||||
size="sm"
|
||||
className="h-7 w-28 shrink-0 border-none px-1 shadow-none focus-visible:ring-0"
|
||||
>
|
||||
<Select value={f.status} onValueChange={(v) => onStatus(f, v as FindingStatus)}>
|
||||
<SelectTrigger size="sm" className="h-7 w-28 shrink-0 border-none px-1 shadow-none focus-visible:ring-0">
|
||||
<StatusBadge domain="finding" value={f.status} dot />
|
||||
</SelectTrigger>
|
||||
<SelectContent position="popper" align="end">
|
||||
@@ -147,9 +127,7 @@ function Row({
|
||||
</div>
|
||||
{open && (
|
||||
<div className="bg-muted/30 px-4 pb-4 pl-11">
|
||||
<div className="mb-1 text-xs font-medium text-muted-foreground">
|
||||
证据 / PoC
|
||||
</div>
|
||||
<div className="mb-1 text-xs font-medium text-muted-foreground">证据 / PoC</div>
|
||||
<pre className="overflow-auto rounded-md border bg-background p-3 font-mono text-xs whitespace-pre-wrap">
|
||||
{f.evidence}
|
||||
</pre>
|
||||
@@ -188,25 +166,18 @@ export function FindingsTab({ taskId }: { taskId: string }) {
|
||||
};
|
||||
}, [taskId]);
|
||||
|
||||
const onStatus = React.useCallback(
|
||||
async (f: Finding, next: FindingStatus) => {
|
||||
if (f.inherited || !f.finding_id || next === f.status) return;
|
||||
const prev = f.status;
|
||||
setFindings((cur) =>
|
||||
cur.map((x) => (x.id === f.id ? { ...x, status: next } : x)),
|
||||
);
|
||||
try {
|
||||
await api.setFindingStatus(f.finding_id, next);
|
||||
toast.success(`已标记为「${statusMeta("finding", next).label}」`);
|
||||
} catch (e) {
|
||||
setFindings((cur) =>
|
||||
cur.map((x) => (x.id === f.id ? { ...x, status: prev } : x)),
|
||||
);
|
||||
toast.error("更新失败:" + (e as Error).message);
|
||||
}
|
||||
},
|
||||
[],
|
||||
);
|
||||
const onStatus = React.useCallback(async (f: Finding, next: FindingStatus) => {
|
||||
if (f.inherited || !f.finding_id || next === f.status) return;
|
||||
const prev = f.status;
|
||||
setFindings((cur) => cur.map((x) => (x.id === f.id ? { ...x, status: next } : x)));
|
||||
try {
|
||||
await api.setFindingStatus(f.finding_id, next);
|
||||
toast.success(`已标记为「${statusMeta("finding", next).label}」`);
|
||||
} catch (e) {
|
||||
setFindings((cur) => cur.map((x) => (x.id === f.id ? { ...x, status: prev } : x)));
|
||||
toast.error("更新失败:" + (e as Error).message);
|
||||
}
|
||||
}, []);
|
||||
|
||||
const items = findings
|
||||
.filter((f) => f.task_id === taskId || f.inherited)
|
||||
@@ -244,9 +215,7 @@ export function FindingsTab({ taskId }: { taskId: string }) {
|
||||
<Row key={f.id} f={f} contextTaskId={taskId} onStatus={onStatus} />
|
||||
))}
|
||||
{items.length === 0 && (
|
||||
<p className="px-4 py-8 text-center text-sm text-muted-foreground">
|
||||
本任务及直接关联任务暂无确认发现。
|
||||
</p>
|
||||
<p className="px-4 py-8 text-center text-sm text-muted-foreground">本任务及直接关联任务暂无确认发现。</p>
|
||||
)}
|
||||
</CardContent>
|
||||
</Card>
|
||||
|
||||
@@ -5,19 +5,18 @@ import * as React from "react";
|
||||
import {
|
||||
ArrowDownWideNarrowIcon,
|
||||
ArrowUpNarrowWideIcon,
|
||||
CheckIcon,
|
||||
ChevronLeftIcon,
|
||||
ChevronRightIcon,
|
||||
CopyIcon,
|
||||
ListChecksIcon,
|
||||
Loader2Icon,
|
||||
RadioTowerIcon,
|
||||
SearchIcon,
|
||||
Trash2Icon,
|
||||
WrapTextIcon,
|
||||
} from "lucide-react";
|
||||
import { toast } from "sonner";
|
||||
|
||||
import { HttpCodeBlock } from "@/components/http-code-block";
|
||||
import { LinkTrafficDialog } from "@/components/link-traffic-dialog";
|
||||
import {
|
||||
AlertDialog,
|
||||
AlertDialogAction,
|
||||
@@ -41,7 +40,7 @@ import { Tabs, TabsContent, TabsList, TabsTrigger } from "@/components/ui/tabs";
|
||||
import { Tooltip, TooltipContent, TooltipTrigger } from "@/components/ui/tooltip";
|
||||
import { api } from "@/lib/api";
|
||||
import type { TrafficDetail, TrafficExchange, TrafficHost, TrafficResp } from "@/lib/types";
|
||||
import { cn, copyText } from "@/lib/utils";
|
||||
import { cn } from "@/lib/utils";
|
||||
|
||||
function fmtTime(ts: string) {
|
||||
return new Date(ts).toLocaleString("zh-CN", {
|
||||
@@ -90,205 +89,6 @@ function requestWithHost(raw: string, exchange: TrafficExchange): string {
|
||||
return `${raw.slice(0, firstLineEnd + newline.length)}Host: ${host}${newline}${raw.slice(firstLineEnd + newline.length)}`;
|
||||
}
|
||||
|
||||
function StartLine({ line }: { line: string }) {
|
||||
const response = /^(HTTP\/\S+)(\s+)(\d{3})(.*)$/.exec(line);
|
||||
if (response) {
|
||||
return (
|
||||
<>
|
||||
<span className="text-muted-foreground">{response[1]}</span>
|
||||
{response[2]}
|
||||
<span className={statusTone(Number(response[3]))}>{response[3]}</span>
|
||||
{response[4]}
|
||||
</>
|
||||
);
|
||||
}
|
||||
|
||||
const request = /^([A-Z]+)(\s+)(\S+)(\s+)(HTTP\/\S+)$/.exec(line);
|
||||
if (request) {
|
||||
return (
|
||||
<>
|
||||
<span className="font-semibold text-primary">{request[1]}</span>
|
||||
{request[2]}
|
||||
<span className="text-chart-2">{request[3]}</span>
|
||||
{request[4]}
|
||||
<span className="text-muted-foreground">{request[5]}</span>
|
||||
</>
|
||||
);
|
||||
}
|
||||
|
||||
return line;
|
||||
}
|
||||
|
||||
function HeaderLine({ line }: { line: string }) {
|
||||
const separator = line.indexOf(":");
|
||||
if (separator <= 0) return line;
|
||||
return (
|
||||
<>
|
||||
<span className="text-primary">{line.slice(0, separator)}</span>
|
||||
<span className="text-muted-foreground">:</span>
|
||||
<span className="text-chart-2">{line.slice(separator + 1)}</span>
|
||||
</>
|
||||
);
|
||||
}
|
||||
|
||||
function JsonBody({ body }: { body: string }) {
|
||||
const parts: React.ReactNode[] = [];
|
||||
const tokens = /("(?:\\.|[^"\\])*")(\s*:)?|\b(true|false|null)\b|-?\d+(?:\.\d+)?(?:[eE][+-]?\d+)?/g;
|
||||
let cursor = 0;
|
||||
for (const match of body.matchAll(tokens)) {
|
||||
const index = match.index ?? 0;
|
||||
if (index > cursor) parts.push(body.slice(cursor, index));
|
||||
let className = "text-chart-4";
|
||||
if (match[1]) className = match[2] ? "text-primary" : "text-chart-2";
|
||||
else if (match[3]) className = "text-chart-3";
|
||||
parts.push(
|
||||
<span key={`${index}-${match[0].length}`} className={className}>
|
||||
{match[0]}
|
||||
</span>,
|
||||
);
|
||||
cursor = index + match[0].length;
|
||||
}
|
||||
if (cursor < body.length) parts.push(body.slice(cursor));
|
||||
return parts;
|
||||
}
|
||||
|
||||
function MarkupBody({ body }: { body: string }) {
|
||||
const parts: React.ReactNode[] = [];
|
||||
const tags = /<\/?[A-Za-z][^>]*>|<!--[\s\S]*?-->/g;
|
||||
let cursor = 0;
|
||||
for (const match of body.matchAll(tags)) {
|
||||
const index = match.index ?? 0;
|
||||
if (index > cursor) parts.push(body.slice(cursor, index));
|
||||
parts.push(
|
||||
<span key={`${index}-${match[0].length}`} className="text-primary">
|
||||
{match[0]}
|
||||
</span>,
|
||||
);
|
||||
cursor = index + match[0].length;
|
||||
}
|
||||
if (cursor < body.length) parts.push(body.slice(cursor));
|
||||
return parts;
|
||||
}
|
||||
|
||||
type BodyFormat = "json" | "markup" | "plain";
|
||||
|
||||
function detectBodyFormat(body: string): BodyFormat {
|
||||
const trimmed = body.trimStart();
|
||||
if (trimmed.startsWith("{") || trimmed.startsWith("[")) return "json";
|
||||
if (trimmed.startsWith("<")) return "markup";
|
||||
return "plain";
|
||||
}
|
||||
|
||||
function HighlightedBody({ body, format }: { body: string; format: BodyFormat }) {
|
||||
if (format === "json") return <JsonBody body={body} />;
|
||||
if (format === "markup") return <MarkupBody body={body} />;
|
||||
return body;
|
||||
}
|
||||
|
||||
function HttpCodeBlock({ raw }: { raw: string }) {
|
||||
const [wrapLines, setWrapLines] = React.useState(true);
|
||||
const [copied, setCopied] = React.useState(false);
|
||||
const value = raw || "(空)";
|
||||
const lines = value.replaceAll("\r\n", "\n").split("\n");
|
||||
const separator = lines.indexOf("");
|
||||
const body = separator >= 0 ? lines.slice(separator + 1).join("\n") : "";
|
||||
const bodyFormat = detectBodyFormat(body);
|
||||
|
||||
React.useEffect(() => {
|
||||
if (!copied) return;
|
||||
const timer = window.setTimeout(() => setCopied(false), 1500);
|
||||
return () => window.clearTimeout(timer);
|
||||
}, [copied]);
|
||||
|
||||
const copyPacket = async () => {
|
||||
const ok = await copyText(value);
|
||||
if (ok) {
|
||||
setCopied(true);
|
||||
return;
|
||||
}
|
||||
toast.error("复制失败,请使用 Ctrl/Cmd+A 后复制");
|
||||
};
|
||||
|
||||
const renderLine = (line: string, index: number) => {
|
||||
if (index === 0) return <StartLine line={line} />;
|
||||
if (separator < 0 || index < separator) return <HeaderLine line={line} />;
|
||||
if (index === separator) return null;
|
||||
return <HighlightedBody body={line} format={bodyFormat} />;
|
||||
};
|
||||
|
||||
return (
|
||||
<div className="relative mx-3 my-4 overflow-hidden rounded-md border bg-background shadow-xs">
|
||||
<div className="absolute top-2 right-2 flex items-center gap-0.5 rounded-md border bg-background/95 p-0.5 shadow-xs">
|
||||
<Tooltip>
|
||||
<TooltipTrigger asChild>
|
||||
<Button
|
||||
type="button"
|
||||
variant="ghost"
|
||||
size="icon-xs"
|
||||
aria-label={wrapLines ? "关闭自动换行" : "开启自动换行"}
|
||||
aria-pressed={wrapLines}
|
||||
onClick={() => setWrapLines((current) => !current)}
|
||||
>
|
||||
<WrapTextIcon />
|
||||
</Button>
|
||||
</TooltipTrigger>
|
||||
<TooltipContent side="bottom">{wrapLines ? "关闭自动换行" : "开启自动换行"}</TooltipContent>
|
||||
</Tooltip>
|
||||
<Tooltip>
|
||||
<TooltipTrigger asChild>
|
||||
<Button
|
||||
type="button"
|
||||
variant="ghost"
|
||||
size="icon-xs"
|
||||
aria-label={copied ? "已复制报文" : "复制报文"}
|
||||
onClick={() => void copyPacket()}
|
||||
>
|
||||
{copied ? <CheckIcon /> : <CopyIcon />}
|
||||
</Button>
|
||||
</TooltipTrigger>
|
||||
<TooltipContent side="bottom">{copied ? "已复制" : "复制报文"}</TooltipContent>
|
||||
</Tooltip>
|
||||
</div>
|
||||
{/* biome-ignore lint/a11y/useSemanticElements: textarea cannot preserve line numbers and syntax-highlighting markup. */}
|
||||
<div
|
||||
role="textbox"
|
||||
aria-label="HTTP 报文代码"
|
||||
aria-multiline="true"
|
||||
aria-readonly="true"
|
||||
tabIndex={0}
|
||||
className="max-h-[calc(100vh-15rem)] min-w-0 overflow-auto bg-background py-3 font-mono text-xs outline-none focus-visible:ring-2 focus-visible:ring-inset focus-visible:ring-ring"
|
||||
onKeyDown={(event) => {
|
||||
if (!(event.ctrlKey || event.metaKey) || event.key.toLowerCase() !== "a") return;
|
||||
event.preventDefault();
|
||||
const selection = window.getSelection();
|
||||
if (!selection) return;
|
||||
const range = document.createRange();
|
||||
range.selectNodeContents(event.currentTarget);
|
||||
selection.removeAllRanges();
|
||||
selection.addRange(range);
|
||||
}}
|
||||
>
|
||||
{lines.map((line, index) => (
|
||||
<div
|
||||
key={`${index}-${line}`}
|
||||
data-line={index + 1}
|
||||
className="grid min-w-0 grid-cols-[1.5rem_minmax(0,1fr)] leading-relaxed before:sticky before:left-0 before:self-stretch before:border-r before:bg-muted/20 before:px-1 before:text-right before:text-muted-foreground/60 before:content-[attr(data-line)]"
|
||||
>
|
||||
<code
|
||||
className={cn(
|
||||
"min-h-[1lh] min-w-0 pr-16 pl-1.5 [tab-size:4]",
|
||||
wrapLines ? "break-words whitespace-pre-wrap" : "whitespace-pre",
|
||||
)}
|
||||
>
|
||||
{renderLine(line, index)}
|
||||
</code>
|
||||
</div>
|
||||
))}
|
||||
</div>
|
||||
</div>
|
||||
);
|
||||
}
|
||||
|
||||
// Fixed method set (server filters exact-match); avoids deriving options from a
|
||||
// single page, which would only ever list the methods on that page.
|
||||
const METHODS = ["GET", "POST", "PUT", "DELETE", "PATCH", "HEAD", "OPTIONS"];
|
||||
@@ -296,6 +96,8 @@ const PAGE_SIZES = [25, 50, 100, 200];
|
||||
type HostCountSortDirection = "asc" | "desc";
|
||||
|
||||
export default function TrafficPage() {
|
||||
const [selectedFlows, setSelectedFlows] = React.useState<Set<string>>(() => new Set());
|
||||
const [linking, setLinking] = React.useState(false);
|
||||
const [page, setPage] = React.useState(0);
|
||||
const [size, setSize] = React.useState(50);
|
||||
const [host, setHost] = React.useState(""); // raw host input
|
||||
@@ -623,6 +425,17 @@ export default function TrafficPage() {
|
||||
</div>
|
||||
</div>
|
||||
|
||||
<div className="flex flex-wrap items-center gap-2">
|
||||
<span className="text-sm text-muted-foreground">已选 {selectedFlows.size} 条流量</span>
|
||||
<Button variant="outline" size="sm" disabled={selectedFlows.size === 0} onClick={() => setLinking(true)}>
|
||||
关联到漏洞
|
||||
</Button>
|
||||
{selectedFlows.size > 0 ? (
|
||||
<Button variant="ghost" size="sm" onClick={() => setSelectedFlows(new Set())}>
|
||||
清空选择
|
||||
</Button>
|
||||
) : null}
|
||||
</div>
|
||||
{/* History table */}
|
||||
<div className="flex h-[calc(100vh-15rem)] min-h-0 flex-col">
|
||||
<Card className="flex min-h-0 flex-1 flex-col overflow-hidden py-0">
|
||||
@@ -630,6 +443,22 @@ export default function TrafficPage() {
|
||||
<Table>
|
||||
<TableHeader className="sticky top-0 z-10 bg-card">
|
||||
<TableRow>
|
||||
<TableHead className="w-10">
|
||||
<Checkbox
|
||||
aria-label="选择本页流量"
|
||||
checked={exchanges.length > 0 && exchanges.every((e) => selectedFlows.has(e.id))}
|
||||
onCheckedChange={(checked) =>
|
||||
setSelectedFlows((previous) => {
|
||||
const next = new Set(previous);
|
||||
for (const e of exchanges) {
|
||||
if (checked === true) next.add(e.id);
|
||||
else next.delete(e.id);
|
||||
}
|
||||
return next;
|
||||
})
|
||||
}
|
||||
/>
|
||||
</TableHead>
|
||||
<TableHead className="w-36">时间</TableHead>
|
||||
<TableHead className="w-44">host</TableHead>
|
||||
<TableHead className="w-20">方法</TableHead>
|
||||
@@ -646,6 +475,21 @@ export default function TrafficPage() {
|
||||
className={cn("cursor-pointer", selected?.id === e.id && "bg-accent hover:bg-accent")}
|
||||
onClick={() => setSelected(e)}
|
||||
>
|
||||
<TableCell>
|
||||
<Checkbox
|
||||
aria-label={`选择流量 ${e.id}`}
|
||||
checked={selectedFlows.has(e.id)}
|
||||
onClick={(event) => event.stopPropagation()}
|
||||
onCheckedChange={(checked) =>
|
||||
setSelectedFlows((previous) => {
|
||||
const next = new Set(previous);
|
||||
if (checked === true) next.add(e.id);
|
||||
else next.delete(e.id);
|
||||
return next;
|
||||
})
|
||||
}
|
||||
/>
|
||||
</TableCell>
|
||||
<TableCell className="text-xs text-muted-foreground tabular-nums">{fmtTime(e.ts)}</TableCell>
|
||||
<TableCell className="font-mono text-xs">{e.host}</TableCell>
|
||||
<TableCell>
|
||||
@@ -665,7 +509,7 @@ export default function TrafficPage() {
|
||||
))}
|
||||
{exchanges.length === 0 && (
|
||||
<TableRow>
|
||||
<TableCell colSpan={7} className="py-12 text-center text-sm text-muted-foreground">
|
||||
<TableCell colSpan={8} className="py-12 text-center text-sm text-muted-foreground">
|
||||
{traffic === null ? "加载中…" : "没有匹配的流量。"}
|
||||
</TableCell>
|
||||
</TableRow>
|
||||
@@ -676,6 +520,13 @@ export default function TrafficPage() {
|
||||
</Card>
|
||||
</div>
|
||||
|
||||
{linking ? (
|
||||
<LinkTrafficDialog
|
||||
trafficIds={[...selectedFlows]}
|
||||
onClose={() => setLinking(false)}
|
||||
onBound={() => setSelectedFlows(new Set())}
|
||||
/>
|
||||
) : null}
|
||||
<Sheet open={selected !== null} onOpenChange={(open) => !open && setSelected(null)}>
|
||||
<SheetContent className="w-full! max-w-none! gap-0 p-0 sm:w-[48rem]! sm:max-w-[48rem]!">
|
||||
{selected && (
|
||||
|
||||
@@ -17,6 +17,7 @@ import type { Settings } from "@/lib/types";
|
||||
|
||||
export default function SystemSettingsPage() {
|
||||
const [trafficCapture, setTrafficCapture] = React.useState(false);
|
||||
const [agentTrafficBinding, setAgentTrafficBinding] = React.useState(false);
|
||||
const [webSearch, setWebSearch] = React.useState(false);
|
||||
const [backend, setBackend] = React.useState("ddgs");
|
||||
const [braveKeySet, setBraveKeySet] = React.useState(false);
|
||||
@@ -43,13 +44,14 @@ export default function SystemSettingsPage() {
|
||||
|
||||
const apply = React.useCallback((s: Settings) => {
|
||||
setTrafficCapture(!!s.traffic_capture);
|
||||
setAgentTrafficBinding(!!s.agent_traffic_binding);
|
||||
setWebSearch(!!s.web_search_enabled);
|
||||
setBackend(s.web_search_backend || "ddgs");
|
||||
setBraveKeySet(!!s.brave_key_set);
|
||||
setTavilyKeySet(!!s.tavily_key_set);
|
||||
setProxyInput(s.web_search_proxy || "");
|
||||
setGlobalProxyInput(s.global_proxy || "");
|
||||
setPyInterp(s.python_interpreter || "");
|
||||
setProxyInput(s.web_search_proxy ?? "");
|
||||
setGlobalProxyInput(s.global_proxy ?? "");
|
||||
setPyInterp(s.python_interpreter ?? "");
|
||||
setWorkers(String(s.workers ?? 3));
|
||||
setInjectPlanner(s.constraints_inject_planner !== false);
|
||||
setInjectWorker(s.constraints_inject_worker !== false);
|
||||
@@ -118,6 +120,22 @@ export default function SystemSettingsPage() {
|
||||
.catch(() => setInjectPlanner(!v)); // revert on failure
|
||||
};
|
||||
|
||||
const toggleAgentTrafficBinding = (v: boolean) => {
|
||||
setAgentTrafficBinding(v);
|
||||
setSaving(true);
|
||||
api
|
||||
.setSettings({ agent_traffic_binding: v })
|
||||
.then((s) => {
|
||||
apply(s);
|
||||
toast.success(v ? "已开启 Agent 自动绑定流量" : "已关闭 Agent 自动绑定流量");
|
||||
})
|
||||
.catch((e) => {
|
||||
setAgentTrafficBinding(!v);
|
||||
toast.error(`保存失败:${(e as Error).message}`);
|
||||
})
|
||||
.finally(() => setSaving(false));
|
||||
};
|
||||
|
||||
const toggleInjectWorker = (v: boolean) => {
|
||||
setInjectWorker(v); // optimistic
|
||||
api
|
||||
@@ -256,6 +274,34 @@ export default function SystemSettingsPage() {
|
||||
</CardContent>
|
||||
</Card>
|
||||
|
||||
<Card className="mb-4 break-inside-avoid md:mb-6">
|
||||
<CardHeader>
|
||||
<CardTitle className="flex items-center gap-2 text-base">
|
||||
<RadioTowerIcon className="size-4" />
|
||||
Agent 自动绑定流量
|
||||
</CardTitle>
|
||||
<CardDescription id="agent-traffic-binding-description">
|
||||
默认关闭。开启后,漏洞入库时触发的报告 Agent 会核对已有 HTTP 请求/响应,关联对应流量后再编写报告。
|
||||
<b>查阅数据包及额外的工具调用会增加 Token 消耗。</b>
|
||||
<br />
|
||||
TCP、未抓包或没有匹配流量时仍可正常上报。此开关不影响流量捕获、人工绑定及已保存证据的查看。 对下一轮 Agent
|
||||
生效;关闭后会立即拒绝新的自动绑定。
|
||||
</CardDescription>
|
||||
</CardHeader>
|
||||
<CardContent className="flex items-center justify-between gap-4">
|
||||
<Label htmlFor="agent-traffic-binding" className="text-sm font-normal text-muted-foreground">
|
||||
{agentTrafficBinding ? "已开启 · 会增加 Token 消耗" : "已关闭 · 可继续人工绑定"}
|
||||
</Label>
|
||||
<Switch
|
||||
id="agent-traffic-binding"
|
||||
aria-describedby="agent-traffic-binding-description"
|
||||
checked={agentTrafficBinding}
|
||||
disabled={!loaded || saving}
|
||||
onCheckedChange={toggleAgentTrafficBinding}
|
||||
/>
|
||||
</CardContent>
|
||||
</Card>
|
||||
|
||||
<Card className="mb-4 break-inside-avoid md:mb-6">
|
||||
<CardHeader>
|
||||
<CardTitle className="flex items-center gap-2 text-base">
|
||||
@@ -269,8 +315,9 @@ export default function SystemSettingsPage() {
|
||||
开启<b>流量捕获</b>时,它作为记录代理的<b>上游</b>(流量仍全量落库,再经此代理出网);关闭捕获时,直接注入
|
||||
Agent 的 bash / WebFetch 出网。与网络搜索代理、LLM 代理相互独立。
|
||||
<br />
|
||||
<b>提示</b>:socks5 在<b>关闭捕获</b>时依赖各命令行工具对 <code>ALL_PROXY</code> 的支持(curl 可用,部分工具可能忽略);
|
||||
若主要用 socks5,建议开启流量捕获——此路径由 MITM 亲自拨号,工具无感知、稳定生效。
|
||||
<b>提示</b>:socks5 在<b>关闭捕获</b>时依赖各命令行工具对 <code>ALL_PROXY</code> 的支持(curl
|
||||
可用,部分工具可能忽略); 若主要用 socks5,建议开启流量捕获——此路径由 MITM
|
||||
亲自拨号,工具无感知、稳定生效。
|
||||
</CardDescription>
|
||||
</CardHeader>
|
||||
<CardContent className="flex flex-col gap-2">
|
||||
@@ -291,9 +338,7 @@ export default function SystemSettingsPage() {
|
||||
</Button>
|
||||
</div>
|
||||
<p className="text-muted-foreground text-xs">
|
||||
{globalProxyInput.trim()
|
||||
? "已配置 · 所有目标流量经此代理出网"
|
||||
: "未配置 · 目标流量直连出网"}
|
||||
{globalProxyInput.trim() ? "已配置 · 所有目标流量经此代理出网" : "未配置 · 目标流量直连出网"}
|
||||
</p>
|
||||
</CardContent>
|
||||
</Card>
|
||||
@@ -308,8 +353,8 @@ export default function SystemSettingsPage() {
|
||||
开启后,把每个任务的<b>操作约束</b>(在任务总览「操作约束」里维护的 allow/deny 条目)拼进对应 Agent
|
||||
的系统提示,用来框定探索边界(如「仅测当前端口」「禁止爆破」)。
|
||||
<br />
|
||||
可分别控制注入到 <b>规划者(planner)</b>与 <b>执行者(worker)</b>;默认都开。切换即时生效(下一轮读取),无需重建
|
||||
Agent。关闭后该 Agent 不再看到约束。
|
||||
可分别控制注入到 <b>规划者(planner)</b>与 <b>执行者(worker)</b>
|
||||
;默认都开。切换即时生效(下一轮读取),无需重建 Agent。关闭后该 Agent 不再看到约束。
|
||||
</CardDescription>
|
||||
</CardHeader>
|
||||
<CardContent className="flex flex-col gap-4">
|
||||
|
||||
@@ -0,0 +1,285 @@
|
||||
"use client";
|
||||
|
||||
import * as React from "react";
|
||||
|
||||
import { ArrowDownIcon, ArrowUpIcon, PlusIcon } from "lucide-react";
|
||||
import { toast } from "sonner";
|
||||
|
||||
import { TrafficEvidenceViewer } from "@/components/traffic-evidence-viewer";
|
||||
import { TrafficPickerDialog } from "@/components/traffic-picker-dialog";
|
||||
import { Alert, AlertDescription } from "@/components/ui/alert";
|
||||
import { Badge } from "@/components/ui/badge";
|
||||
import { Button } from "@/components/ui/button";
|
||||
import { Card, CardContent, CardDescription, CardHeader, CardTitle } from "@/components/ui/card";
|
||||
import {
|
||||
Dialog,
|
||||
DialogContent,
|
||||
DialogDescription,
|
||||
DialogFooter,
|
||||
DialogHeader,
|
||||
DialogTitle,
|
||||
} from "@/components/ui/dialog";
|
||||
import { Empty, EmptyDescription, EmptyHeader, EmptyTitle } from "@/components/ui/empty";
|
||||
import { Field, FieldGroup, FieldLabel } from "@/components/ui/field";
|
||||
import { Select, SelectContent, SelectGroup, SelectItem, SelectTrigger, SelectValue } from "@/components/ui/select";
|
||||
import { Skeleton } from "@/components/ui/skeleton";
|
||||
import { Textarea } from "@/components/ui/textarea";
|
||||
import { api } from "@/lib/api";
|
||||
import type { FindingTraffic, FindingTrafficBinding, TrafficEvidenceRole } from "@/lib/types";
|
||||
|
||||
const ROLES: Record<TrafficEvidenceRole, string> = {
|
||||
baseline: "正常对照",
|
||||
proof: "漏洞证明",
|
||||
verification: "补充验证",
|
||||
supporting: "辅助证据",
|
||||
};
|
||||
|
||||
export function FindingTrafficPanel({
|
||||
findingId,
|
||||
contextTask,
|
||||
readOnly,
|
||||
onChanged,
|
||||
}: {
|
||||
findingId: string;
|
||||
contextTask?: string;
|
||||
readOnly?: boolean;
|
||||
onChanged: () => void;
|
||||
}) {
|
||||
const [data, setData] = React.useState<FindingTraffic | null>(null);
|
||||
const [error, setError] = React.useState("");
|
||||
const [adding, setAdding] = React.useState(false);
|
||||
const [preview, setPreview] = React.useState<string | null>(null);
|
||||
const [editing, setEditing] = React.useState<FindingTrafficBinding | null>(null);
|
||||
const [role, setRole] = React.useState<TrafficEvidenceRole>("supporting");
|
||||
const [note, setNote] = React.useState("");
|
||||
const [busy, setBusy] = React.useState(false);
|
||||
const [reload, setReload] = React.useState(0);
|
||||
// biome-ignore lint/correctness/useExhaustiveDependencies: retry token intentionally refreshes this request.
|
||||
React.useEffect(() => {
|
||||
let active = true;
|
||||
setData(null);
|
||||
setError("");
|
||||
api
|
||||
.findingTraffic(findingId, contextTask)
|
||||
.then((value) => {
|
||||
if (active) setData(value);
|
||||
})
|
||||
.catch((e: Error) => {
|
||||
if (active) setError(e.message);
|
||||
});
|
||||
return () => {
|
||||
active = false;
|
||||
};
|
||||
}, [findingId, contextTask, reload]);
|
||||
|
||||
async function mutate(action: () => Promise<FindingTraffic>) {
|
||||
setBusy(true);
|
||||
try {
|
||||
setData(await action());
|
||||
setEditing(null);
|
||||
onChanged();
|
||||
toast.success("流量证据已更新");
|
||||
} catch (e) {
|
||||
toast.error((e as Error).message);
|
||||
setReload((n) => n + 1);
|
||||
} finally {
|
||||
setBusy(false);
|
||||
}
|
||||
}
|
||||
|
||||
function move(index: number, delta: number) {
|
||||
if (!data) return;
|
||||
const ids = data.bindings.map((b) => b.id);
|
||||
[ids[index], ids[index + delta]] = [ids[index + delta], ids[index]];
|
||||
void mutate(() => api.orderFindingTraffic(findingId, ids, data.version, contextTask));
|
||||
}
|
||||
|
||||
return (
|
||||
<>
|
||||
<Card>
|
||||
<CardHeader>
|
||||
<div className="flex flex-wrap items-center justify-between gap-2">
|
||||
<CardTitle>关联流量 {data ? `(${data.bindings.length})` : ""}</CardTitle>
|
||||
{!readOnly ? (
|
||||
<Button variant="outline" size="sm" disabled={!data || busy} onClick={() => setAdding(true)}>
|
||||
<PlusIcon data-icon="inline-start" />
|
||||
绑定流量
|
||||
</Button>
|
||||
) : (
|
||||
<Badge variant="outline">继承证据 · 只读</Badge>
|
||||
)}
|
||||
</div>
|
||||
<CardDescription>按复现顺序组织请求与响应。清理原始流量后,已绑定证据仍然保留。</CardDescription>
|
||||
</CardHeader>
|
||||
<CardContent className="flex flex-col gap-3">
|
||||
{error ? (
|
||||
<Alert variant="destructive">
|
||||
<AlertDescription>
|
||||
{error}
|
||||
<Button variant="link" onClick={() => setReload((n) => n + 1)}>
|
||||
重试
|
||||
</Button>
|
||||
</AlertDescription>
|
||||
</Alert>
|
||||
) : null}
|
||||
{!data ? (
|
||||
<Skeleton className="h-28 w-full" />
|
||||
) : data.bindings.length === 0 ? (
|
||||
<Empty>
|
||||
<EmptyHeader>
|
||||
<EmptyTitle>暂无关联流量</EmptyTitle>
|
||||
<EmptyDescription>可绑定正常对照、漏洞证明和补充验证请求。</EmptyDescription>
|
||||
</EmptyHeader>
|
||||
</Empty>
|
||||
) : (
|
||||
data.bindings.map((b, index) => (
|
||||
<div key={b.id} className="flex min-w-0 flex-col gap-2 rounded-md border p-3">
|
||||
<div className="flex flex-wrap items-center gap-2">
|
||||
<Badge variant="outline">
|
||||
{index + 1}. {ROLES[b.role]}
|
||||
</Badge>
|
||||
<Badge variant="secondary">
|
||||
{b.snapshot.method} · {b.snapshot.status}
|
||||
</Badge>
|
||||
<span className="text-xs text-muted-foreground">证据 #{b.id}</span>
|
||||
</div>
|
||||
<Button
|
||||
variant="link"
|
||||
className="h-auto justify-start p-0 text-left whitespace-normal"
|
||||
onClick={() => setPreview(b.id)}
|
||||
>
|
||||
<span className="break-all font-mono text-xs">{b.snapshot.url}</span>
|
||||
</Button>
|
||||
{b.note ? <p className="text-sm whitespace-pre-wrap">{b.note}</p> : null}
|
||||
<div className="flex flex-wrap items-center justify-between gap-2">
|
||||
<span className="text-xs text-muted-foreground">
|
||||
{new Date(b.snapshot.captured_at * 1000).toLocaleString("zh-CN")}
|
||||
</span>
|
||||
<div className="flex flex-wrap gap-1">
|
||||
<Button variant="outline" size="sm" onClick={() => setPreview(b.id)}>
|
||||
查看报文
|
||||
</Button>
|
||||
{!readOnly ? (
|
||||
<>
|
||||
<Button
|
||||
variant="ghost"
|
||||
size="sm"
|
||||
disabled={busy}
|
||||
onClick={() => {
|
||||
setEditing(b);
|
||||
setRole(b.role);
|
||||
setNote(b.note);
|
||||
}}
|
||||
>
|
||||
编辑说明
|
||||
</Button>
|
||||
<Button
|
||||
variant="ghost"
|
||||
size="icon-sm"
|
||||
aria-label={`上移证据 ${b.id}`}
|
||||
disabled={busy || index === 0}
|
||||
onClick={() => move(index, -1)}
|
||||
>
|
||||
<ArrowUpIcon />
|
||||
</Button>
|
||||
<Button
|
||||
variant="ghost"
|
||||
size="icon-sm"
|
||||
aria-label={`下移证据 ${b.id}`}
|
||||
disabled={busy || index === data.bindings.length - 1}
|
||||
onClick={() => move(index, 1)}
|
||||
>
|
||||
<ArrowDownIcon />
|
||||
</Button>
|
||||
<Button
|
||||
variant="ghost"
|
||||
size="sm"
|
||||
disabled={busy}
|
||||
onClick={() =>
|
||||
void mutate(() => api.removeFindingTraffic(findingId, b.id, data.version, contextTask))
|
||||
}
|
||||
>
|
||||
解除绑定
|
||||
</Button>
|
||||
</>
|
||||
) : null}
|
||||
</div>
|
||||
</div>
|
||||
</div>
|
||||
))
|
||||
)}
|
||||
</CardContent>
|
||||
</Card>
|
||||
{adding && data ? (
|
||||
<TrafficPickerDialog
|
||||
findingId={findingId}
|
||||
contextTask={contextTask}
|
||||
bound={data.bindings}
|
||||
onClose={() => setAdding(false)}
|
||||
onBound={() => {
|
||||
setReload((n) => n + 1);
|
||||
onChanged();
|
||||
}}
|
||||
/>
|
||||
) : null}
|
||||
<TrafficEvidenceViewer
|
||||
findingId={findingId}
|
||||
bindingId={preview}
|
||||
contextTask={contextTask}
|
||||
onClose={() => setPreview(null)}
|
||||
/>
|
||||
<Dialog
|
||||
open={editing !== null}
|
||||
onOpenChange={(open) => {
|
||||
if (!open && !busy) setEditing(null);
|
||||
}}
|
||||
>
|
||||
<DialogContent>
|
||||
<DialogHeader>
|
||||
<DialogTitle>编辑流量证据</DialogTitle>
|
||||
<DialogDescription>说明这组请求/响应如何支持漏洞结论。</DialogDescription>
|
||||
</DialogHeader>
|
||||
<FieldGroup>
|
||||
<Field>
|
||||
<FieldLabel htmlFor="evidence-role">用途</FieldLabel>
|
||||
<Select value={role} onValueChange={(v) => setRole(v as TrafficEvidenceRole)}>
|
||||
<SelectTrigger id="evidence-role">
|
||||
<SelectValue />
|
||||
</SelectTrigger>
|
||||
<SelectContent>
|
||||
<SelectGroup>
|
||||
{Object.entries(ROLES).map(([value, label]) => (
|
||||
<SelectItem key={value} value={value}>
|
||||
{label}
|
||||
</SelectItem>
|
||||
))}
|
||||
</SelectGroup>
|
||||
</SelectContent>
|
||||
</Select>
|
||||
</Field>
|
||||
<Field>
|
||||
<FieldLabel htmlFor="evidence-note">证据说明</FieldLabel>
|
||||
<Textarea id="evidence-note" value={note} onChange={(e) => setNote(e.target.value)} />
|
||||
</Field>
|
||||
</FieldGroup>
|
||||
<DialogFooter>
|
||||
<Button variant="outline" disabled={busy} onClick={() => setEditing(null)}>
|
||||
取消
|
||||
</Button>
|
||||
<Button
|
||||
disabled={busy || !editing || !data}
|
||||
onClick={() => {
|
||||
if (editing && data)
|
||||
void mutate(() =>
|
||||
api.editFindingTraffic(findingId, editing.id, data.version, { role, note }, contextTask),
|
||||
);
|
||||
}}
|
||||
>
|
||||
保存说明
|
||||
</Button>
|
||||
</DialogFooter>
|
||||
</DialogContent>
|
||||
</Dialog>
|
||||
</>
|
||||
);
|
||||
}
|
||||
@@ -0,0 +1,216 @@
|
||||
"use client";
|
||||
import * as React from "react";
|
||||
|
||||
import { CheckIcon, CopyIcon, WrapTextIcon } from "lucide-react";
|
||||
import { toast } from "sonner";
|
||||
|
||||
import { Button } from "@/components/ui/button";
|
||||
import { Tooltip, TooltipContent, TooltipTrigger } from "@/components/ui/tooltip";
|
||||
import { cn, copyText } from "@/lib/utils";
|
||||
|
||||
function statusTone(status: number) {
|
||||
if (status >= 500) return "text-red-500";
|
||||
if (status >= 400) return "text-amber-500";
|
||||
if (status >= 300) return "text-blue-500";
|
||||
if (status >= 200) return "text-emerald-500";
|
||||
return "text-muted-foreground";
|
||||
}
|
||||
|
||||
function StartLine({ line }: { line: string }) {
|
||||
const response = /^(HTTP\/\S+)(\s+)(\d{3})(.*)$/.exec(line);
|
||||
if (response) {
|
||||
return (
|
||||
<>
|
||||
<span className="text-muted-foreground">{response[1]}</span>
|
||||
{response[2]}
|
||||
<span className={statusTone(Number(response[3]))}>{response[3]}</span>
|
||||
{response[4]}
|
||||
</>
|
||||
);
|
||||
}
|
||||
|
||||
const request = /^([A-Z]+)(\s+)(\S+)(\s+)(HTTP\/\S+)$/.exec(line);
|
||||
if (request) {
|
||||
return (
|
||||
<>
|
||||
<span className="font-semibold text-primary">{request[1]}</span>
|
||||
{request[2]}
|
||||
<span className="text-chart-2">{request[3]}</span>
|
||||
{request[4]}
|
||||
<span className="text-muted-foreground">{request[5]}</span>
|
||||
</>
|
||||
);
|
||||
}
|
||||
|
||||
return line;
|
||||
}
|
||||
|
||||
function HeaderLine({ line }: { line: string }) {
|
||||
const separator = line.indexOf(":");
|
||||
if (separator <= 0) return line;
|
||||
return (
|
||||
<>
|
||||
<span className="text-primary">{line.slice(0, separator)}</span>
|
||||
<span className="text-muted-foreground">:</span>
|
||||
<span className="text-chart-2">{line.slice(separator + 1)}</span>
|
||||
</>
|
||||
);
|
||||
}
|
||||
|
||||
function JsonBody({ body }: { body: string }) {
|
||||
const parts: React.ReactNode[] = [];
|
||||
const tokens = /("(?:\\.|[^"\\])*")(\s*:)?|\b(true|false|null)\b|-?\d+(?:\.\d+)?(?:[eE][+-]?\d+)?/g;
|
||||
let cursor = 0;
|
||||
for (const match of body.matchAll(tokens)) {
|
||||
const index = match.index ?? 0;
|
||||
if (index > cursor) parts.push(body.slice(cursor, index));
|
||||
let className = "text-chart-4";
|
||||
if (match[1]) className = match[2] ? "text-primary" : "text-chart-2";
|
||||
else if (match[3]) className = "text-chart-3";
|
||||
parts.push(
|
||||
<span key={`${index}-${match[0].length}`} className={className}>
|
||||
{match[0]}
|
||||
</span>,
|
||||
);
|
||||
cursor = index + match[0].length;
|
||||
}
|
||||
if (cursor < body.length) parts.push(body.slice(cursor));
|
||||
return parts;
|
||||
}
|
||||
|
||||
function MarkupBody({ body }: { body: string }) {
|
||||
const parts: React.ReactNode[] = [];
|
||||
const tags = /<\/?[A-Za-z][^>]*>|<!--[\s\S]*?-->/g;
|
||||
let cursor = 0;
|
||||
for (const match of body.matchAll(tags)) {
|
||||
const index = match.index ?? 0;
|
||||
if (index > cursor) parts.push(body.slice(cursor, index));
|
||||
parts.push(
|
||||
<span key={`${index}-${match[0].length}`} className="text-primary">
|
||||
{match[0]}
|
||||
</span>,
|
||||
);
|
||||
cursor = index + match[0].length;
|
||||
}
|
||||
if (cursor < body.length) parts.push(body.slice(cursor));
|
||||
return parts;
|
||||
}
|
||||
|
||||
type BodyFormat = "json" | "markup" | "plain";
|
||||
|
||||
function detectBodyFormat(body: string): BodyFormat {
|
||||
const trimmed = body.trimStart();
|
||||
if (trimmed.startsWith("{") || trimmed.startsWith("[")) return "json";
|
||||
if (trimmed.startsWith("<")) return "markup";
|
||||
return "plain";
|
||||
}
|
||||
|
||||
function HighlightedBody({ body, format }: { body: string; format: BodyFormat }) {
|
||||
if (format === "json") return <JsonBody body={body} />;
|
||||
if (format === "markup") return <MarkupBody body={body} />;
|
||||
return body;
|
||||
}
|
||||
|
||||
export function HttpCodeBlock({ raw }: { raw: string }) {
|
||||
const [wrapLines, setWrapLines] = React.useState(true);
|
||||
const [copied, setCopied] = React.useState(false);
|
||||
const value = raw || "(空)";
|
||||
const lines = value.replaceAll("\r\n", "\n").split("\n");
|
||||
const separator = lines.indexOf("");
|
||||
const body = separator >= 0 ? lines.slice(separator + 1).join("\n") : "";
|
||||
const bodyFormat = detectBodyFormat(body);
|
||||
|
||||
React.useEffect(() => {
|
||||
if (!copied) return;
|
||||
const timer = window.setTimeout(() => setCopied(false), 1500);
|
||||
return () => window.clearTimeout(timer);
|
||||
}, [copied]);
|
||||
|
||||
const copyPacket = async () => {
|
||||
const ok = await copyText(value);
|
||||
if (ok) {
|
||||
setCopied(true);
|
||||
return;
|
||||
}
|
||||
toast.error("复制失败,请使用 Ctrl/Cmd+A 后复制");
|
||||
};
|
||||
|
||||
const renderLine = (line: string, index: number) => {
|
||||
if (index === 0) return <StartLine line={line} />;
|
||||
if (separator < 0 || index < separator) return <HeaderLine line={line} />;
|
||||
if (index === separator) return null;
|
||||
return <HighlightedBody body={line} format={bodyFormat} />;
|
||||
};
|
||||
|
||||
return (
|
||||
<div className="relative mx-3 my-4 overflow-hidden rounded-md border bg-background shadow-xs">
|
||||
<div className="absolute top-2 right-2 flex items-center gap-0.5 rounded-md border bg-background/95 p-0.5 shadow-xs">
|
||||
<Tooltip>
|
||||
<TooltipTrigger asChild>
|
||||
<Button
|
||||
type="button"
|
||||
variant="ghost"
|
||||
size="icon-xs"
|
||||
aria-label={wrapLines ? "关闭自动换行" : "开启自动换行"}
|
||||
aria-pressed={wrapLines}
|
||||
onClick={() => setWrapLines((current) => !current)}
|
||||
>
|
||||
<WrapTextIcon />
|
||||
</Button>
|
||||
</TooltipTrigger>
|
||||
<TooltipContent side="bottom">{wrapLines ? "关闭自动换行" : "开启自动换行"}</TooltipContent>
|
||||
</Tooltip>
|
||||
<Tooltip>
|
||||
<TooltipTrigger asChild>
|
||||
<Button
|
||||
type="button"
|
||||
variant="ghost"
|
||||
size="icon-xs"
|
||||
aria-label={copied ? "已复制报文" : "复制报文"}
|
||||
onClick={() => void copyPacket()}
|
||||
>
|
||||
{copied ? <CheckIcon /> : <CopyIcon />}
|
||||
</Button>
|
||||
</TooltipTrigger>
|
||||
<TooltipContent side="bottom">{copied ? "已复制" : "复制报文"}</TooltipContent>
|
||||
</Tooltip>
|
||||
</div>
|
||||
{/* biome-ignore lint/a11y/useSemanticElements: textarea cannot preserve line numbers and syntax-highlighting markup. */}
|
||||
<div
|
||||
role="textbox"
|
||||
aria-label="HTTP 报文代码"
|
||||
aria-multiline="true"
|
||||
aria-readonly="true"
|
||||
tabIndex={0}
|
||||
className="max-h-[calc(100vh-15rem)] min-w-0 overflow-auto bg-background py-3 font-mono text-xs outline-none focus-visible:ring-2 focus-visible:ring-inset focus-visible:ring-ring"
|
||||
onKeyDown={(event) => {
|
||||
if (!(event.ctrlKey || event.metaKey) || event.key.toLowerCase() !== "a") return;
|
||||
event.preventDefault();
|
||||
const selection = window.getSelection();
|
||||
if (!selection) return;
|
||||
const range = document.createRange();
|
||||
range.selectNodeContents(event.currentTarget);
|
||||
selection.removeAllRanges();
|
||||
selection.addRange(range);
|
||||
}}
|
||||
>
|
||||
{lines.map((line, index) => (
|
||||
<div
|
||||
key={`${index}-${line}`}
|
||||
data-line={index + 1}
|
||||
className="grid min-w-0 grid-cols-[1.5rem_minmax(0,1fr)] leading-relaxed before:sticky before:left-0 before:self-stretch before:border-r before:bg-muted/20 before:px-1 before:text-right before:text-muted-foreground/60 before:content-[attr(data-line)]"
|
||||
>
|
||||
<code
|
||||
className={cn(
|
||||
"min-h-[1lh] min-w-0 pr-16 pl-1.5 [tab-size:4]",
|
||||
wrapLines ? "break-words whitespace-pre-wrap" : "whitespace-pre",
|
||||
)}
|
||||
>
|
||||
{renderLine(line, index)}
|
||||
</code>
|
||||
</div>
|
||||
))}
|
||||
</div>
|
||||
</div>
|
||||
);
|
||||
}
|
||||
@@ -0,0 +1,163 @@
|
||||
"use client";
|
||||
|
||||
import * as React from "react";
|
||||
|
||||
import { toast } from "sonner";
|
||||
|
||||
import { Alert, AlertDescription } from "@/components/ui/alert";
|
||||
import { Button } from "@/components/ui/button";
|
||||
import {
|
||||
Dialog,
|
||||
DialogContent,
|
||||
DialogDescription,
|
||||
DialogFooter,
|
||||
DialogHeader,
|
||||
DialogTitle,
|
||||
} from "@/components/ui/dialog";
|
||||
import { Field, FieldGroup, FieldLabel } from "@/components/ui/field";
|
||||
import { Input } from "@/components/ui/input";
|
||||
import { api } from "@/lib/api";
|
||||
import type { Finding, FindingsPage } from "@/lib/types";
|
||||
|
||||
export function LinkTrafficDialog({
|
||||
trafficIds,
|
||||
onClose,
|
||||
onBound,
|
||||
}: {
|
||||
trafficIds: string[];
|
||||
onClose: () => void;
|
||||
onBound: () => void;
|
||||
}) {
|
||||
const [query, setQuery] = React.useState("");
|
||||
const [page, setPage] = React.useState(1);
|
||||
const [data, setData] = React.useState<FindingsPage | null>(null);
|
||||
const [selected, setSelected] = React.useState<Finding | null>(null);
|
||||
const [error, setError] = React.useState("");
|
||||
const [busy, setBusy] = React.useState(false);
|
||||
const [loading, setLoading] = React.useState(true);
|
||||
React.useEffect(() => {
|
||||
let active = true;
|
||||
setLoading(true);
|
||||
setError("");
|
||||
const timer = setTimeout(() => {
|
||||
api
|
||||
.findingsPage({ page, pageSize: 20, query })
|
||||
.then((d) => {
|
||||
if (active) setData(d);
|
||||
})
|
||||
.catch((e: Error) => {
|
||||
if (active) setError(e.message);
|
||||
})
|
||||
.finally(() => {
|
||||
if (active) setLoading(false);
|
||||
});
|
||||
}, 250);
|
||||
return () => {
|
||||
active = false;
|
||||
clearTimeout(timer);
|
||||
};
|
||||
}, [query, page]);
|
||||
async function save() {
|
||||
if (!selected?.finding_id) return;
|
||||
setBusy(true);
|
||||
setError("");
|
||||
try {
|
||||
await api.bindFindingTraffic(
|
||||
selected.finding_id,
|
||||
trafficIds.map((traffic_id) => ({ traffic_id })),
|
||||
);
|
||||
toast.success(`已关联 ${trafficIds.length} 条流量到漏洞 #${selected.finding_id}`);
|
||||
onBound();
|
||||
onClose();
|
||||
} catch (e) {
|
||||
setError((e as Error).message);
|
||||
} finally {
|
||||
setBusy(false);
|
||||
}
|
||||
}
|
||||
return (
|
||||
<Dialog
|
||||
open
|
||||
onOpenChange={(open) => {
|
||||
if (!open && !busy) onClose();
|
||||
}}
|
||||
>
|
||||
<DialogContent className="max-h-[90vh] overflow-y-auto sm:max-w-2xl">
|
||||
<DialogHeader>
|
||||
<DialogTitle>关联到漏洞</DialogTitle>
|
||||
<DialogDescription>将所选 {trafficIds.length} 条流量保存为已有漏洞的证据。</DialogDescription>
|
||||
</DialogHeader>
|
||||
<FieldGroup>
|
||||
<Field>
|
||||
<FieldLabel htmlFor="link-finding-query">查找漏洞</FieldLabel>
|
||||
<Input
|
||||
id="link-finding-query"
|
||||
placeholder="名称 / 摘要 / 漏洞类别"
|
||||
value={query}
|
||||
onChange={(e) => {
|
||||
setQuery(e.target.value);
|
||||
setPage(1);
|
||||
}}
|
||||
/>
|
||||
</Field>
|
||||
</FieldGroup>
|
||||
{error ? (
|
||||
<Alert variant="destructive">
|
||||
<AlertDescription>{error}</AlertDescription>
|
||||
</Alert>
|
||||
) : null}
|
||||
<div className="flex max-h-[40vh] flex-col gap-2 overflow-auto" aria-busy={loading}>
|
||||
{(data?.items ?? []).map((f) => (
|
||||
<Button
|
||||
key={f.finding_id ?? f.id}
|
||||
variant={selected?.finding_id === f.finding_id ? "secondary" : "outline"}
|
||||
className="h-auto justify-start p-3 text-left whitespace-normal"
|
||||
aria-pressed={selected?.finding_id === f.finding_id}
|
||||
disabled={busy || !f.finding_id || f.inherited}
|
||||
onClick={() => setSelected(f)}
|
||||
>
|
||||
<span className="flex min-w-0 flex-col gap-1">
|
||||
<span>
|
||||
#{f.finding_id ?? f.id} · {f.name || f.vulnclass}
|
||||
</span>
|
||||
<span className="line-clamp-2 text-xs text-muted-foreground">{f.summary}</span>
|
||||
</span>
|
||||
</Button>
|
||||
))}
|
||||
{!data?.items.length ? (
|
||||
<p className="py-6 text-center text-sm text-muted-foreground">
|
||||
{loading ? "加载中…" : "没有匹配的漏洞,请先登记漏洞"}
|
||||
</p>
|
||||
) : null}
|
||||
</div>
|
||||
<div className="flex items-center justify-between gap-2">
|
||||
<Button variant="outline" size="sm" disabled={loading || page <= 1} onClick={() => setPage((p) => p - 1)}>
|
||||
上一页
|
||||
</Button>
|
||||
<span className="text-xs">
|
||||
第 {page} 页 · 共 {data?.total ?? 0} 条
|
||||
</span>
|
||||
<Button
|
||||
variant="outline"
|
||||
size="sm"
|
||||
disabled={loading || page * 20 >= (data?.total ?? 0)}
|
||||
onClick={() => setPage((p) => p + 1)}
|
||||
>
|
||||
下一页
|
||||
</Button>
|
||||
</div>
|
||||
<p className="text-sm">
|
||||
{selected ? `已选漏洞:#${selected.finding_id} ${selected.name || selected.vulnclass}` : "请选择一个漏洞"}
|
||||
</p>
|
||||
<DialogFooter>
|
||||
<Button variant="outline" disabled={busy} onClick={onClose}>
|
||||
取消
|
||||
</Button>
|
||||
<Button disabled={busy || !selected} onClick={() => void save()}>
|
||||
{busy ? "保存中…" : "确认关联"}
|
||||
</Button>
|
||||
</DialogFooter>
|
||||
</DialogContent>
|
||||
</Dialog>
|
||||
);
|
||||
}
|
||||
@@ -0,0 +1,181 @@
|
||||
"use client";
|
||||
|
||||
import * as React from "react";
|
||||
|
||||
import { toast } from "sonner";
|
||||
|
||||
import { HttpCodeBlock } from "@/components/http-code-block";
|
||||
import { Alert, AlertDescription } from "@/components/ui/alert";
|
||||
import { Button } from "@/components/ui/button";
|
||||
import { Dialog, DialogContent, DialogDescription, DialogHeader, DialogTitle } from "@/components/ui/dialog";
|
||||
import { Skeleton } from "@/components/ui/skeleton";
|
||||
import { Tabs, TabsContent, TabsList, TabsTrigger } from "@/components/ui/tabs";
|
||||
import { api } from "@/lib/api";
|
||||
import type { FindingTrafficDetail, TrafficDetail } from "@/lib/types";
|
||||
|
||||
export function TrafficEvidenceViewer({
|
||||
findingId,
|
||||
bindingId,
|
||||
contextTask,
|
||||
onClose,
|
||||
}: {
|
||||
findingId: string;
|
||||
bindingId: string | null;
|
||||
contextTask?: string;
|
||||
onClose: () => void;
|
||||
}) {
|
||||
const [detail, setDetail] = React.useState<FindingTrafficDetail | null>(null);
|
||||
const [error, setError] = React.useState("");
|
||||
const [busy, setBusy] = React.useState(false);
|
||||
React.useEffect(() => {
|
||||
if (!bindingId) return;
|
||||
let active = true;
|
||||
setDetail(null);
|
||||
setError("");
|
||||
api
|
||||
.findingTrafficDetail(findingId, bindingId, contextTask)
|
||||
.then((d) => {
|
||||
if (active) setDetail(d);
|
||||
})
|
||||
.catch((e: Error) => {
|
||||
if (active) setError(e.message);
|
||||
});
|
||||
return () => {
|
||||
active = false;
|
||||
};
|
||||
}, [findingId, bindingId, contextTask]);
|
||||
|
||||
async function more(side: "request" | "response") {
|
||||
if (!detail || !bindingId) return;
|
||||
setBusy(true);
|
||||
try {
|
||||
const next = await api.findingTrafficBody(findingId, bindingId, side, detail[side].next_offset, contextTask);
|
||||
setDetail((current) =>
|
||||
current?.binding.id === bindingId
|
||||
? { ...current, [side]: { ...next, content: current[side].content + next.content } }
|
||||
: current,
|
||||
);
|
||||
} catch (e) {
|
||||
toast.error((e as Error).message);
|
||||
} finally {
|
||||
setBusy(false);
|
||||
}
|
||||
}
|
||||
|
||||
return (
|
||||
<Dialog
|
||||
open={bindingId !== null}
|
||||
onOpenChange={(open) => {
|
||||
if (!open) onClose();
|
||||
}}
|
||||
>
|
||||
<DialogContent className="max-h-[90vh] overflow-y-auto sm:max-w-4xl">
|
||||
<DialogHeader>
|
||||
<DialogTitle>流量证据 #{bindingId}</DialogTitle>
|
||||
<DialogDescription className="break-all">
|
||||
{detail?.binding.snapshot.url ?? "查看绑定时保存的请求与响应"}
|
||||
</DialogDescription>
|
||||
</DialogHeader>
|
||||
{error ? (
|
||||
<Alert variant="destructive">
|
||||
<AlertDescription>{error}</AlertDescription>
|
||||
</Alert>
|
||||
) : detail ? (
|
||||
<Tabs defaultValue="request">
|
||||
<TabsList>
|
||||
<TabsTrigger value="request">请求 Request</TabsTrigger>
|
||||
<TabsTrigger value="response">响应 Response</TabsTrigger>
|
||||
</TabsList>
|
||||
{(["request", "response"] as const).map((side) => (
|
||||
<TabsContent key={side} value={side}>
|
||||
<div className="flex flex-wrap items-center justify-between gap-2">
|
||||
<span className="text-xs text-muted-foreground">
|
||||
正文 {detail[side].total.toLocaleString()} 字节{detail[side].truncated ? " · 当前为预览" : ""}
|
||||
</span>
|
||||
<Button
|
||||
variant="outline"
|
||||
size="sm"
|
||||
onClick={() =>
|
||||
void api
|
||||
.downloadFindingTrafficBody(findingId, detail.binding.id, side, contextTask)
|
||||
.catch((e: Error) => toast.error(e.message))
|
||||
}
|
||||
>
|
||||
下载完整{side === "request" ? "请求" : "响应"}正文
|
||||
</Button>
|
||||
</div>
|
||||
<HttpCodeBlock
|
||||
raw={`${(side === "request" ? detail.binding.snapshot.req_head : detail.binding.snapshot.resp_head)?.trimEnd() ?? ""}\n\n${detail[side].content}`}
|
||||
/>
|
||||
{detail[side].truncated && !detail[side].binary ? (
|
||||
<Button variant="outline" size="sm" disabled={busy} onClick={() => void more(side)}>
|
||||
加载更多正文
|
||||
</Button>
|
||||
) : null}
|
||||
</TabsContent>
|
||||
))}
|
||||
</Tabs>
|
||||
) : (
|
||||
<Skeleton className="h-56 w-full" />
|
||||
)}
|
||||
</DialogContent>
|
||||
</Dialog>
|
||||
);
|
||||
}
|
||||
|
||||
export function CapturedTrafficViewer({ id, onClose }: { id: string | null; onClose: () => void }) {
|
||||
const [detail, setDetail] = React.useState<TrafficDetail | null>(null);
|
||||
const [error, setError] = React.useState("");
|
||||
React.useEffect(() => {
|
||||
if (!id) return;
|
||||
let active = true;
|
||||
setDetail(null);
|
||||
setError("");
|
||||
api
|
||||
.trafficExchange(id)
|
||||
.then((d) => {
|
||||
if (active) setDetail(d);
|
||||
})
|
||||
.catch((e: Error) => {
|
||||
if (active) setError(e.message);
|
||||
});
|
||||
return () => {
|
||||
active = false;
|
||||
};
|
||||
}, [id]);
|
||||
return (
|
||||
<Dialog
|
||||
open={id !== null}
|
||||
onOpenChange={(open) => {
|
||||
if (!open) onClose();
|
||||
}}
|
||||
>
|
||||
<DialogContent className="max-h-[90vh] overflow-y-auto sm:max-w-4xl">
|
||||
<DialogHeader>
|
||||
<DialogTitle>预览流量</DialogTitle>
|
||||
<DialogDescription>流量 ID:{id}。绑定时会保存完整正文。</DialogDescription>
|
||||
</DialogHeader>
|
||||
{error ? (
|
||||
<Alert variant="destructive">
|
||||
<AlertDescription>{error}</AlertDescription>
|
||||
</Alert>
|
||||
) : detail ? (
|
||||
<Tabs defaultValue="request">
|
||||
<TabsList>
|
||||
<TabsTrigger value="request">请求 Request</TabsTrigger>
|
||||
<TabsTrigger value="response">响应 Response</TabsTrigger>
|
||||
</TabsList>
|
||||
<TabsContent value="request">
|
||||
<HttpCodeBlock raw={detail.req} />
|
||||
</TabsContent>
|
||||
<TabsContent value="response">
|
||||
<HttpCodeBlock raw={detail.resp} />
|
||||
</TabsContent>
|
||||
</Tabs>
|
||||
) : (
|
||||
<Skeleton className="h-56 w-full" />
|
||||
)}
|
||||
</DialogContent>
|
||||
</Dialog>
|
||||
);
|
||||
}
|
||||
@@ -0,0 +1,277 @@
|
||||
"use client";
|
||||
|
||||
import * as React from "react";
|
||||
|
||||
import { toast } from "sonner";
|
||||
|
||||
import { CapturedTrafficViewer } from "@/components/traffic-evidence-viewer";
|
||||
import { Alert, AlertDescription } from "@/components/ui/alert";
|
||||
import { Badge } from "@/components/ui/badge";
|
||||
import { Button } from "@/components/ui/button";
|
||||
import { Checkbox } from "@/components/ui/checkbox";
|
||||
import {
|
||||
Dialog,
|
||||
DialogContent,
|
||||
DialogDescription,
|
||||
DialogFooter,
|
||||
DialogHeader,
|
||||
DialogTitle,
|
||||
} from "@/components/ui/dialog";
|
||||
import { Field, FieldGroup, FieldLabel } from "@/components/ui/field";
|
||||
import { Input } from "@/components/ui/input";
|
||||
import { Select, SelectContent, SelectGroup, SelectItem, SelectTrigger, SelectValue } from "@/components/ui/select";
|
||||
import { Table, TableBody, TableCell, TableHead, TableHeader, TableRow } from "@/components/ui/table";
|
||||
import { api } from "@/lib/api";
|
||||
import type { FindingTrafficBinding, TrafficResp } from "@/lib/types";
|
||||
|
||||
const METHODS = ["GET", "POST", "PUT", "DELETE", "PATCH", "HEAD", "OPTIONS"];
|
||||
|
||||
export function TrafficPickerDialog({
|
||||
findingId,
|
||||
contextTask,
|
||||
bound,
|
||||
onClose,
|
||||
onBound,
|
||||
}: {
|
||||
findingId: string;
|
||||
contextTask?: string;
|
||||
bound: FindingTrafficBinding[];
|
||||
onClose: () => void;
|
||||
onBound: () => void;
|
||||
}) {
|
||||
const [host, setHost] = React.useState("");
|
||||
const [method, setMethod] = React.useState("all");
|
||||
const [query, setQuery] = React.useState("");
|
||||
const [page, setPage] = React.useState(0);
|
||||
const [data, setData] = React.useState<TrafficResp | null>(null);
|
||||
const [selected, setSelected] = React.useState<Set<string>>(() => new Set());
|
||||
const [preview, setPreview] = React.useState<string | null>(null);
|
||||
const [loading, setLoading] = React.useState(true);
|
||||
const [busy, setBusy] = React.useState(false);
|
||||
const [error, setError] = React.useState("");
|
||||
const alreadyBound = new Set(bound.map((b) => b.snapshot.source_traffic_id));
|
||||
React.useEffect(() => {
|
||||
let active = true;
|
||||
setLoading(true);
|
||||
setError("");
|
||||
const timer = setTimeout(() => {
|
||||
api
|
||||
.traffic(page, 25, host, method === "all" ? "" : method, query)
|
||||
.then((d) => {
|
||||
if (active) setData(d);
|
||||
})
|
||||
.catch((e: Error) => {
|
||||
if (active) setError(e.message);
|
||||
})
|
||||
.finally(() => {
|
||||
if (active) setLoading(false);
|
||||
});
|
||||
}, 250);
|
||||
return () => {
|
||||
active = false;
|
||||
clearTimeout(timer);
|
||||
};
|
||||
}, [page, host, method, query]);
|
||||
|
||||
const rows = data?.exchanges ?? [];
|
||||
const selectable = rows.filter((e) => !alreadyBound.has(e.id));
|
||||
function toggle(id: string, checked: boolean) {
|
||||
setSelected((previous) => {
|
||||
const next = new Set(previous);
|
||||
if (checked) next.add(id);
|
||||
else next.delete(id);
|
||||
return next;
|
||||
});
|
||||
}
|
||||
async function save() {
|
||||
setBusy(true);
|
||||
setError("");
|
||||
try {
|
||||
await api.bindFindingTraffic(
|
||||
findingId,
|
||||
[...selected].map((traffic_id) => ({ traffic_id })),
|
||||
contextTask,
|
||||
);
|
||||
toast.success(`已绑定 ${selected.size} 条流量`);
|
||||
onBound();
|
||||
onClose();
|
||||
} catch (e) {
|
||||
setError((e as Error).message);
|
||||
} finally {
|
||||
setBusy(false);
|
||||
}
|
||||
}
|
||||
|
||||
return (
|
||||
<>
|
||||
<Dialog
|
||||
open
|
||||
onOpenChange={(open) => {
|
||||
if (!open && !busy) onClose();
|
||||
}}
|
||||
>
|
||||
<DialogContent className="max-h-[90vh] overflow-y-auto sm:max-w-5xl">
|
||||
<DialogHeader>
|
||||
<DialogTitle>绑定流量</DialogTitle>
|
||||
<DialogDescription>
|
||||
筛选并多选请求/响应,已选记录会跨页保留。绑定后可设置用途、说明和顺序。
|
||||
</DialogDescription>
|
||||
</DialogHeader>
|
||||
<FieldGroup className="flex flex-col gap-3 sm:flex-row">
|
||||
<Field>
|
||||
<FieldLabel htmlFor="evidence-host">目标 host</FieldLabel>
|
||||
<Input
|
||||
id="evidence-host"
|
||||
value={host}
|
||||
placeholder="域名或 IP"
|
||||
onChange={(e) => {
|
||||
setHost(e.target.value);
|
||||
setPage(0);
|
||||
}}
|
||||
/>
|
||||
</Field>
|
||||
<Field>
|
||||
<FieldLabel htmlFor="evidence-method">请求方法</FieldLabel>
|
||||
<Select
|
||||
value={method}
|
||||
onValueChange={(v) => {
|
||||
setMethod(v);
|
||||
setPage(0);
|
||||
}}
|
||||
>
|
||||
<SelectTrigger id="evidence-method">
|
||||
<SelectValue />
|
||||
</SelectTrigger>
|
||||
<SelectContent>
|
||||
<SelectGroup>
|
||||
<SelectItem value="all">全部方法</SelectItem>
|
||||
{METHODS.map((m) => (
|
||||
<SelectItem key={m} value={m}>
|
||||
{m}
|
||||
</SelectItem>
|
||||
))}
|
||||
</SelectGroup>
|
||||
</SelectContent>
|
||||
</Select>
|
||||
</Field>
|
||||
<Field>
|
||||
<FieldLabel htmlFor="evidence-query">关键词</FieldLabel>
|
||||
<Input
|
||||
id="evidence-query"
|
||||
value={query}
|
||||
placeholder="URL / 正文关键词"
|
||||
onChange={(e) => {
|
||||
setQuery(e.target.value);
|
||||
setPage(0);
|
||||
}}
|
||||
/>
|
||||
</Field>
|
||||
</FieldGroup>
|
||||
{error ? (
|
||||
<Alert variant="destructive">
|
||||
<AlertDescription>{error}</AlertDescription>
|
||||
</Alert>
|
||||
) : null}
|
||||
<div className="max-h-[45vh] overflow-auto rounded-md border" aria-busy={loading}>
|
||||
<Table>
|
||||
<TableHeader>
|
||||
<TableRow>
|
||||
<TableHead>
|
||||
<Checkbox
|
||||
aria-label="选择本页未绑定流量"
|
||||
disabled={loading || busy || selectable.length === 0}
|
||||
checked={selectable.length > 0 && selectable.every((e) => selected.has(e.id))}
|
||||
onCheckedChange={(checked) =>
|
||||
setSelected((previous) => {
|
||||
const next = new Set(previous);
|
||||
for (const e of selectable) {
|
||||
if (checked === true) next.add(e.id);
|
||||
else next.delete(e.id);
|
||||
}
|
||||
return next;
|
||||
})
|
||||
}
|
||||
/>
|
||||
</TableHead>
|
||||
<TableHead>方法 / URL</TableHead>
|
||||
<TableHead>时间</TableHead>
|
||||
<TableHead>状态码</TableHead>
|
||||
<TableHead>操作</TableHead>
|
||||
</TableRow>
|
||||
</TableHeader>
|
||||
<TableBody>
|
||||
{rows.map((e) => (
|
||||
<TableRow key={e.id}>
|
||||
<TableCell>
|
||||
<Checkbox
|
||||
aria-label={`选择流量 ${e.id}`}
|
||||
checked={selected.has(e.id) || alreadyBound.has(e.id)}
|
||||
disabled={busy || loading || alreadyBound.has(e.id)}
|
||||
onCheckedChange={(checked) => toggle(e.id, checked === true)}
|
||||
/>
|
||||
</TableCell>
|
||||
<TableCell>
|
||||
<span className="font-mono text-xs">
|
||||
{e.method} {e.url}
|
||||
</span>
|
||||
{alreadyBound.has(e.id) ? <Badge variant="secondary">已绑定</Badge> : null}
|
||||
</TableCell>
|
||||
<TableCell className="whitespace-nowrap text-xs">
|
||||
{new Date(e.ts).toLocaleString("zh-CN")}
|
||||
</TableCell>
|
||||
<TableCell>{e.status}</TableCell>
|
||||
<TableCell>
|
||||
<Button variant="ghost" size="sm" onClick={() => setPreview(e.id)}>
|
||||
预览
|
||||
</Button>
|
||||
</TableCell>
|
||||
</TableRow>
|
||||
))}
|
||||
{!rows.length ? (
|
||||
<TableRow>
|
||||
<TableCell colSpan={5} className="py-8 text-center">
|
||||
{loading ? "加载中…" : "没有匹配的流量"}
|
||||
</TableCell>
|
||||
</TableRow>
|
||||
) : null}
|
||||
</TableBody>
|
||||
</Table>
|
||||
</div>
|
||||
<div className="flex flex-wrap items-center justify-between gap-2">
|
||||
<span className="text-sm">
|
||||
已选 {selected.size} 条 · 共 {data?.total ?? 0} 条
|
||||
</span>
|
||||
<div className="flex items-center gap-2">
|
||||
<Button
|
||||
variant="outline"
|
||||
size="sm"
|
||||
disabled={loading || page === 0}
|
||||
onClick={() => setPage((p) => p - 1)}
|
||||
>
|
||||
上一页
|
||||
</Button>
|
||||
<span className="text-xs">第 {page + 1} 页</span>
|
||||
<Button
|
||||
variant="outline"
|
||||
size="sm"
|
||||
disabled={loading || (page + 1) * 25 >= (data?.total ?? 0)}
|
||||
onClick={() => setPage((p) => p + 1)}
|
||||
>
|
||||
下一页
|
||||
</Button>
|
||||
</div>
|
||||
</div>
|
||||
<DialogFooter>
|
||||
<Button variant="outline" disabled={busy} onClick={onClose}>
|
||||
取消
|
||||
</Button>
|
||||
<Button disabled={busy || selected.size === 0} onClick={() => void save()}>
|
||||
{busy ? "保存中…" : `绑定 ${selected.size} 条流量`}
|
||||
</Button>
|
||||
</DialogFooter>
|
||||
</DialogContent>
|
||||
</Dialog>
|
||||
<CapturedTrafficViewer id={preview} onClose={() => setPreview(null)} />
|
||||
</>
|
||||
);
|
||||
}
|
||||
@@ -29,6 +29,7 @@ import type {
|
||||
DeleteTaskOptions,
|
||||
DeleteTaskResult,
|
||||
Edge,
|
||||
EvidenceBodyPreview,
|
||||
Finding,
|
||||
FindingAssetTree,
|
||||
FindingDeepenResponse,
|
||||
@@ -38,6 +39,8 @@ import type {
|
||||
FindingStats,
|
||||
FindingStatus,
|
||||
FindingsPage,
|
||||
FindingTraffic,
|
||||
FindingTrafficDetail,
|
||||
IntentAsset,
|
||||
InterceptApprovalRow,
|
||||
InterceptDetail,
|
||||
@@ -82,6 +85,8 @@ import type {
|
||||
Tool,
|
||||
ToolStat,
|
||||
TrafficDetail,
|
||||
TrafficEvidenceRef,
|
||||
TrafficEvidenceRole,
|
||||
TrafficHost,
|
||||
TrafficResp,
|
||||
UsageStats,
|
||||
@@ -529,6 +534,73 @@ export const api = {
|
||||
get<Finding>(
|
||||
`/exploration/findings/${id}${contextTaskId ? `?context_task=${encodeURIComponent(contextTaskId)}` : ""}`,
|
||||
),
|
||||
findingTraffic: (id: string, contextTask?: string) =>
|
||||
get<FindingTraffic>(
|
||||
`/exploration/findings/${id}/traffic${contextTask ? `?context_task=${encodeURIComponent(contextTask)}` : ""}`,
|
||||
),
|
||||
bindFindingTraffic: (id: string, traffic_refs: TrafficEvidenceRef[], contextTask?: string) =>
|
||||
post<FindingTraffic>(
|
||||
`/exploration/findings/${id}/traffic${contextTask ? `?context_task=${encodeURIComponent(contextTask)}` : ""}`,
|
||||
{ traffic_refs },
|
||||
),
|
||||
editFindingTraffic: (
|
||||
id: string,
|
||||
bindingId: string,
|
||||
version: number,
|
||||
fields: { role: TrafficEvidenceRole; note: string },
|
||||
contextTask?: string,
|
||||
) =>
|
||||
patch<FindingTraffic>(
|
||||
`/exploration/findings/${id}/traffic/${bindingId}${contextTask ? `?context_task=${encodeURIComponent(contextTask)}` : ""}`,
|
||||
{ version, ...fields },
|
||||
),
|
||||
removeFindingTraffic: (id: string, bindingId: string, version: number, contextTask?: string) =>
|
||||
del<FindingTraffic>(
|
||||
`/exploration/findings/${id}/traffic/${bindingId}${contextTask ? `?context_task=${encodeURIComponent(contextTask)}` : ""}`,
|
||||
{ version },
|
||||
),
|
||||
orderFindingTraffic: (id: string, binding_ids: string[], version: number, contextTask?: string) =>
|
||||
http<FindingTraffic>(
|
||||
`/exploration/findings/${id}/traffic/order${contextTask ? `?context_task=${encodeURIComponent(contextTask)}` : ""}`,
|
||||
{ method: "PUT", body: JSON.stringify({ binding_ids, version }) },
|
||||
),
|
||||
findingTrafficDetail: (id: string, bindingId: string, contextTask?: string) =>
|
||||
get<FindingTrafficDetail>(
|
||||
`/exploration/findings/${id}/traffic/${bindingId}${contextTask ? `?context_task=${encodeURIComponent(contextTask)}` : ""}`,
|
||||
),
|
||||
findingTrafficBody: (
|
||||
id: string,
|
||||
bindingId: string,
|
||||
side: "request" | "response",
|
||||
offset: number,
|
||||
contextTask?: string,
|
||||
) =>
|
||||
get<EvidenceBodyPreview>(
|
||||
`/exploration/findings/${id}/traffic/${bindingId}/body?side=${side}&offset=${offset}${contextTask ? `&context_task=${encodeURIComponent(contextTask)}` : ""}`,
|
||||
),
|
||||
downloadFindingTrafficBody: async (
|
||||
id: string,
|
||||
bindingId: string,
|
||||
side: "request" | "response",
|
||||
contextTask?: string,
|
||||
) => {
|
||||
const token = getToken();
|
||||
const response = await fetch(
|
||||
`/api/exploration/findings/${id}/traffic/${bindingId}/body?side=${side}&download=1${contextTask ? `&context_task=${encodeURIComponent(contextTask)}` : ""}`,
|
||||
{ headers: token ? { Authorization: `Bearer ${token}` } : {} },
|
||||
);
|
||||
if (!response.ok) {
|
||||
const error = await response.json().catch(() => ({ error: "下载失败" }));
|
||||
throw new Error(error.error ?? "下载失败");
|
||||
}
|
||||
const blob = await response.blob();
|
||||
const url = URL.createObjectURL(blob);
|
||||
const a = document.createElement("a");
|
||||
a.href = url;
|
||||
a.download = `evidence-${bindingId}-${side}.bin`;
|
||||
a.click();
|
||||
setTimeout(() => URL.revokeObjectURL(url), 1000);
|
||||
},
|
||||
// 漏洞链路:该漏洞节点回溯到任务初始节点的子图(节点 + 关系)。
|
||||
findingLineage: (id: string) => get<{ nodes: TaskNode[]; edges: Edge[] }>(`/exploration/findings/${id}/lineage`),
|
||||
setFindingStatus: (id: string, status: FindingStatus) => patch<Finding>(`/exploration/findings/${id}`, { status }),
|
||||
|
||||
@@ -3314,6 +3314,7 @@ export const tools: Tool[] = [
|
||||
// ── Settings ─────────────────────────────────────────────────────────────────
|
||||
export const settings: Settings = {
|
||||
traffic_capture: true,
|
||||
agent_traffic_binding: false,
|
||||
llm_record: false,
|
||||
web_search_enabled: true,
|
||||
web_search_backend: "ddgs",
|
||||
|
||||
@@ -463,6 +463,10 @@ export interface FindingAsset {
|
||||
}
|
||||
|
||||
export interface Finding {
|
||||
traffic_count?: number;
|
||||
evidence_version?: number;
|
||||
report_evidence_version?: number;
|
||||
report_stale?: boolean;
|
||||
id: string;
|
||||
finding_id?: string; // 独立 findings 表的行 id,状态更新的句柄(任务内旧节点可能缺失)
|
||||
vulnclass: string;
|
||||
@@ -878,6 +882,7 @@ export interface TrafficHost {
|
||||
// ---- App settings (runtime toggles) ----
|
||||
export interface Settings {
|
||||
traffic_capture: boolean;
|
||||
agent_traffic_binding: boolean; // Agent 自动绑定流量证据,默认关闭;不影响人工绑定
|
||||
llm_record: boolean; // LLM 录制开关(默认关);关闭时不记录任何 LLM 调用
|
||||
// Web search. brave_key_set / tavily_key_set reflect whether a key is stored
|
||||
// (the values are never returned). On PUT, send the corresponding field to set/clear.
|
||||
@@ -1299,3 +1304,54 @@ export interface InterceptAudit {
|
||||
export interface InterceptDetail extends InterceptApprovalRow {
|
||||
audit: InterceptAudit | null;
|
||||
}
|
||||
|
||||
export type TrafficEvidenceRole = "baseline" | "proof" | "verification" | "supporting";
|
||||
export interface TrafficEvidenceRef {
|
||||
traffic_id: string;
|
||||
role?: TrafficEvidenceRole;
|
||||
note?: string;
|
||||
}
|
||||
export interface TrafficEvidenceSnapshot {
|
||||
id: string;
|
||||
source_traffic_id: string;
|
||||
captured_at: number;
|
||||
url: string;
|
||||
method: string;
|
||||
status: number;
|
||||
content_type: string;
|
||||
req_head?: string;
|
||||
resp_head?: string;
|
||||
req_hash: string;
|
||||
resp_hash: string;
|
||||
req_len: number;
|
||||
resp_len: number;
|
||||
}
|
||||
export interface FindingTrafficBinding {
|
||||
id: string;
|
||||
finding_id: string;
|
||||
snapshot_id: string;
|
||||
role: TrafficEvidenceRole;
|
||||
note: string;
|
||||
position: number;
|
||||
created_at: string;
|
||||
snapshot: TrafficEvidenceSnapshot;
|
||||
}
|
||||
export interface FindingTraffic {
|
||||
finding_id: string;
|
||||
version: number;
|
||||
report_version: number;
|
||||
bindings: FindingTrafficBinding[];
|
||||
}
|
||||
export interface EvidenceBodyPreview {
|
||||
content: string;
|
||||
offset: number;
|
||||
total: number;
|
||||
next_offset: number;
|
||||
truncated: boolean;
|
||||
binary: boolean;
|
||||
}
|
||||
export interface FindingTrafficDetail {
|
||||
binding: FindingTrafficBinding;
|
||||
request: EvidenceBodyPreview;
|
||||
response: EvidenceBodyPreview;
|
||||
}
|
||||
Reference in new issue
Block a user