feat: 支持漏洞关联多条流量证据

This commit is contained in:
ji ruo committed 2026-09-10 13:03:20 +08:00
1 parent 8dae851b9b
commit ec4ef35fff
41 files changed
+4390 -496

No files matched your search

+11 -2
View File
@@ -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,
+425
View File
@@ -0,0 +1,425 @@
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 id <= 0 {
return actool.Errorf("finding_id 无效"), nil
}
if ri := agent.RunInfoFrom(ctx); ri.TaskID > 0 {
f, err := s.m.pg.GetFinding(id)
if err != nil || f == nil {
return actool.Errorf("漏洞不存在"), nil
}
task := s.m.ResolveTask(strconv.FormatInt(ri.TaskID, 10))
if task == nil {
return actool.Errorf("任务不存在"), nil
}
if _, _, allowed := findingProvenanceInTask(task, f.TaskID); !allowed {
return actool.Errorf("当前任务不可读取该漏洞"), nil
}
}
if 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)
})
}
+121
View File
@@ -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
}
+473
View File
@@ -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)
}
}
}
+15 -11
View File
@@ -75,6 +75,7 @@ func (s *Server) orchestrationTools() []actool.CoreTool {
s.toolSearchWorkerTraces(),
s.toolGetTaskNodeDetail(),
s.toolUpdateFindingReport(),
s.toolGetFindingTraffic(),
}
}
@@ -392,20 +393,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
}
@@ -453,11 +456,12 @@ 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() // 增加可选证据参数及只读证据工具,保留用户配置
// 注:pentest 的默认工具绑定无需迁移——BuiltinToolSeeds 在全新初始化时就把
// list_assets/insert_assets/report_finding/list_findings/list_companies 连同
// pentest 一起 seed 好了(项目尚无旧库,不做迁移)。
+5
View File
@@ -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] "
+36 -25
View File
@@ -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"
@@ -217,9 +217,10 @@ func New(ctx context.Context, m *Manager, skillDir string, dataDir string, keyDi
}
wireAgentAugment(m.pg, s.skillDir, s.hostTools) // 可见 skills/MCP + 流量/编排 host 工具装配进 agent 工具集
domainReg := buildDomainReg(m.Assets())
wireTools(m.pg, domainReg) // 内置工具表:按 agent 过滤 + 覆盖描述/schema + 注入默认值
seedPrompts(m.pg) // 内置 agent 默认提示词正文播种进 agent_prompts(仅空时)
s.seedOrchestrationTools() // P2 跨任务编排工具 seed 进 tools 表(可按 agent 绑定)
wireTools(m.pg, domainReg) // 内置工具表:按 agent 过滤 + 覆盖描述/schema + 注入默认值
seedPrompts(m.pg) // 内置 agent 默认提示词正文播种进 agent_prompts(仅空时)
s.seedOrchestrationTools() // P2 跨任务编排工具 seed 进 tools 表(可按 agent 绑定)
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
@@ -360,7 +361,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
@@ -405,6 +406,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")))
@@ -414,6 +416,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
@@ -482,6 +485,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 实时纠偏
@@ -811,6 +815,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)
@@ -2135,6 +2140,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) {
@@ -2147,27 +2162,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))
@@ -2175,6 +2182,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")
+20 -2
View File
@@ -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
+3
View File
@@ -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)