feat: 支持漏洞关联多条流量证据
This commit is contained in:
1 parent
8dae851b9b
commit
ec4ef35fff
41 files changed
+4390
-496
No files matched your search
+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,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)
|
||||
})
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
}
|
||||
+15
-11
@@ -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 好了(项目尚无旧库,不做迁移)。
|
||||
|
||||
@@ -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
@@ -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
@@ -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)
|
||||
|
||||
Reference in new issue
Block a user