405 lines
13 KiB
Go
405 lines
13 KiB
Go
package db
|
|
|
|
import (
|
|
"context"
|
|
"crypto/sha256"
|
|
"database/sql"
|
|
"encoding/hex"
|
|
"encoding/json"
|
|
"errors"
|
|
"fmt"
|
|
"strings"
|
|
"time"
|
|
)
|
|
|
|
var (
|
|
ErrEvidenceConflict = errors.New("流量证据已变更，请刷新后重试")
|
|
ErrFindingNotFound = errors.New("漏洞不存在")
|
|
ErrEvidenceNotFound = errors.New("流量证据不存在")
|
|
)
|
|
|
|
// This lock covers the evidence filesystem as well as its SQL references. All
|
|
// processes sharing the database use it, including readers, exports and GC.
|
|
const findingEvidenceLockKey int64 = 7337741004
|
|
|
|
func (d *DB) WithEvidenceTx(ctx context.Context, fn func(*sql.Tx) error) error {
|
|
tx, err := d.BeginTx(ctx, nil)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
defer tx.Rollback()
|
|
if _, err = tx.ExecContext(ctx, `SELECT pg_advisory_xact_lock($1)`, findingEvidenceLockKey); err != nil {
|
|
return err
|
|
}
|
|
if err = fn(tx); err != nil {
|
|
return err
|
|
}
|
|
return tx.Commit()
|
|
}
|
|
|
|
type TrafficRef struct {
|
|
TrafficID string `json:"traffic_id"`
|
|
Role string `json:"role"`
|
|
Note string `json:"note"`
|
|
}
|
|
|
|
func NormalizeTrafficRefs(refs []TrafficRef) ([]TrafficRef, error) {
|
|
out := make([]TrafficRef, 0, len(refs))
|
|
seen := map[string]bool{}
|
|
for _, ref := range refs {
|
|
ref.TrafficID = strings.TrimSpace(ref.TrafficID)
|
|
if ref.TrafficID == "" {
|
|
return nil, errors.New("traffic_id 不能为空")
|
|
}
|
|
if ref.Role == "" {
|
|
ref.Role = "supporting"
|
|
}
|
|
if !ValidTrafficRole(ref.Role) {
|
|
return nil, fmt.Errorf("无效的流量用途 %q", ref.Role)
|
|
}
|
|
if !seen[ref.TrafficID] {
|
|
out = append(out, ref)
|
|
seen[ref.TrafficID] = true
|
|
}
|
|
}
|
|
return out, nil
|
|
}
|
|
|
|
func ValidTrafficRole(role string) bool {
|
|
return role == "baseline" || role == "proof" || role == "verification" || role == "supporting"
|
|
}
|
|
|
|
type TrafficEvidenceSnapshot struct {
|
|
ID string `json:"id"`
|
|
SourceTrafficID string `json:"source_traffic_id"`
|
|
CapturedAt int64 `json:"captured_at"`
|
|
URL string `json:"url"`
|
|
Method string `json:"method"`
|
|
Status int `json:"status"`
|
|
ContentType string `json:"content_type"`
|
|
ReqHead string `json:"req_head,omitempty"`
|
|
RespHead string `json:"resp_head,omitempty"`
|
|
ReqHash string `json:"req_hash"`
|
|
RespHash string `json:"resp_hash"`
|
|
ReqLen int64 `json:"req_len"`
|
|
RespLen int64 `json:"resp_len"`
|
|
CreatedAt time.Time `json:"created_at"`
|
|
}
|
|
|
|
type FindingTrafficBinding struct {
|
|
ID int64 `json:"id,string"`
|
|
FindingID int64 `json:"finding_id,string"`
|
|
SnapshotID string `json:"snapshot_id"`
|
|
Role string `json:"role"`
|
|
Note string `json:"note"`
|
|
Position int `json:"position"`
|
|
CreatedAt time.Time `json:"created_at"`
|
|
Snapshot TrafficEvidenceSnapshot `json:"snapshot"`
|
|
}
|
|
|
|
type FindingTraffic struct {
|
|
FindingID int64 `json:"finding_id,string"`
|
|
Version int64 `json:"version"`
|
|
ReportVersion int64 `json:"report_version"`
|
|
Bindings []FindingTrafficBinding `json:"bindings"`
|
|
}
|
|
|
|
type PreparedTrafficEvidence struct {
|
|
Ref TrafficRef
|
|
Snapshot TrafficEvidenceSnapshot
|
|
}
|
|
|
|
func TrafficSnapshotID(snapshot TrafficEvidenceSnapshot) string {
|
|
snapshot.ID = ""
|
|
snapshot.CreatedAt = time.Time{}
|
|
raw, _ := json.Marshal(snapshot)
|
|
sum := sha256.Sum256(raw)
|
|
return hex.EncodeToString(sum[:])
|
|
}
|
|
|
|
// LockTaskEvidenceTx serializes writes with archive queueing (which locks the
|
|
// same task row). Once queued, its snapshot must not acquire new evidence.
|
|
func LockTaskEvidenceTx(tx *sql.Tx, taskID int64) error {
|
|
if taskID == 0 {
|
|
return nil
|
|
}
|
|
var deleted sql.NullTime
|
|
if err := tx.QueryRow(`SELECT deleted_at FROM tasks WHERE id=$1 FOR UPDATE`, taskID).Scan(&deleted); err != nil {
|
|
return err
|
|
}
|
|
if deleted.Valid {
|
|
return ErrTaskArchiveState
|
|
}
|
|
var state string
|
|
err := tx.QueryRow(`SELECT state FROM task_archives WHERE task_id=$1`, taskID).Scan(&state)
|
|
if err != nil && !errors.Is(err, sql.ErrNoRows) {
|
|
return err
|
|
}
|
|
if err == nil && state != ArchiveFailed {
|
|
return ErrTaskArchiveState
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func LockFindingEvidenceTx(tx *sql.Tx, findingID int64, version *int64) error {
|
|
var taskID sql.NullInt64
|
|
if err := tx.QueryRow(`SELECT task_id FROM findings WHERE id=$1`, findingID).Scan(&taskID); err != nil {
|
|
if errors.Is(err, sql.ErrNoRows) {
|
|
return ErrFindingNotFound
|
|
}
|
|
return err
|
|
}
|
|
if err := LockTaskEvidenceTx(tx, taskID.Int64); err != nil {
|
|
return err
|
|
}
|
|
var current int64
|
|
if err := tx.QueryRow(`SELECT evidence_version FROM findings WHERE id=$1 FOR UPDATE`, findingID).Scan(¤t); err != nil {
|
|
if errors.Is(err, sql.ErrNoRows) {
|
|
return ErrFindingNotFound
|
|
}
|
|
return err
|
|
}
|
|
if version != nil && *version != current {
|
|
return ErrEvidenceConflict
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func InsertEvidenceSnapshotTx(tx *sql.Tx, s TrafficEvidenceSnapshot) error {
|
|
if s.ID != TrafficSnapshotID(s) {
|
|
return errors.New("证据快照元数据哈希不匹配")
|
|
}
|
|
if s.CreatedAt.IsZero() {
|
|
s.CreatedAt = time.Now().UTC()
|
|
}
|
|
_, err := tx.Exec(`INSERT INTO traffic_evidence_snapshots
|
|
(id,source_traffic_id,captured_at,url,method,status,content_type,req_head,resp_head,req_hash,resp_hash,req_len,resp_len,created_at)
|
|
VALUES($1,$2,$3,$4,$5,$6,$7,$8,$9,$10,$11,$12,$13,$14) ON CONFLICT(id) DO NOTHING`,
|
|
s.ID, s.SourceTrafficID, s.CapturedAt, s.URL, s.Method, s.Status, s.ContentType, s.ReqHead, s.RespHead, s.ReqHash, s.RespHash, s.ReqLen, s.RespLen, s.CreatedAt)
|
|
return err
|
|
}
|
|
|
|
func AddFindingTrafficTx(tx *sql.Tx, findingID int64, prepared []PreparedTrafficEvidence) error {
|
|
var position int
|
|
if err := tx.QueryRow(`SELECT COALESCE(MAX(position)+1,0) FROM finding_traffic_bindings WHERE finding_id=$1`, findingID).Scan(&position); err != nil {
|
|
return err
|
|
}
|
|
changed := false
|
|
for _, item := range prepared {
|
|
if err := InsertEvidenceSnapshotTx(tx, item.Snapshot); err != nil {
|
|
return err
|
|
}
|
|
if _, err := tx.Exec(`UPDATE traffic_evidence_snapshots SET unreferenced_at=NULL WHERE id=$1`, item.Snapshot.ID); err != nil {
|
|
return err
|
|
}
|
|
res, err := tx.Exec(`INSERT INTO finding_traffic_bindings(finding_id,snapshot_id,role,note,position)
|
|
VALUES($1,$2,$3,$4,$5) ON CONFLICT(finding_id,snapshot_id) DO NOTHING`, findingID, item.Snapshot.ID, item.Ref.Role, item.Ref.Note, position)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
n, err := res.RowsAffected()
|
|
if err != nil {
|
|
return err
|
|
}
|
|
if n > 0 {
|
|
changed = true
|
|
position++
|
|
}
|
|
}
|
|
if changed {
|
|
return bumpEvidenceVersionTx(tx, findingID)
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func bumpEvidenceVersionTx(tx *sql.Tx, findingID int64) error {
|
|
_, err := tx.Exec(`UPDATE findings SET evidence_version=evidence_version+1 WHERE id=$1`, findingID)
|
|
return err
|
|
}
|
|
|
|
func FindingTrafficTx(tx *sql.Tx, findingID int64) (*FindingTraffic, error) {
|
|
out := &FindingTraffic{FindingID: findingID, Bindings: []FindingTrafficBinding{}}
|
|
if err := tx.QueryRow(`SELECT evidence_version,report_evidence_version FROM findings WHERE id=$1`, findingID).Scan(&out.Version, &out.ReportVersion); err != nil {
|
|
if errors.Is(err, sql.ErrNoRows) {
|
|
return nil, ErrFindingNotFound
|
|
}
|
|
return nil, err
|
|
}
|
|
rows, err := tx.Query(`SELECT b.id,b.finding_id,b.snapshot_id,b.role,b.note,b.position,b.created_at,to_jsonb(s)
|
|
FROM finding_traffic_bindings b JOIN traffic_evidence_snapshots s ON s.id=b.snapshot_id WHERE b.finding_id=$1 ORDER BY b.position,b.id`, findingID)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
defer rows.Close()
|
|
for rows.Next() {
|
|
var b FindingTrafficBinding
|
|
var raw []byte
|
|
if err := rows.Scan(&b.ID, &b.FindingID, &b.SnapshotID, &b.Role, &b.Note, &b.Position, &b.CreatedAt, &raw); err != nil {
|
|
return nil, err
|
|
}
|
|
if err := json.Unmarshal(raw, &b.Snapshot); err != nil {
|
|
return nil, err
|
|
}
|
|
out.Bindings = append(out.Bindings, b)
|
|
}
|
|
return out, rows.Err()
|
|
}
|
|
|
|
func (d *DB) GetFindingTraffic(ctx context.Context, findingID int64) (out *FindingTraffic, err error) {
|
|
err = d.WithEvidenceTx(ctx, func(tx *sql.Tx) error { var e error; out, e = FindingTrafficTx(tx, findingID); return e })
|
|
return
|
|
}
|
|
|
|
func (d *DB) EditFindingTraffic(ctx context.Context, findingID, bindingID, version int64, role, note *string, remove bool, order []int64) error {
|
|
if role != nil && !ValidTrafficRole(*role) {
|
|
return errors.New("无效的流量用途")
|
|
}
|
|
return d.WithEvidenceTx(ctx, func(tx *sql.Tx) error {
|
|
if err := LockFindingEvidenceTx(tx, findingID, &version); err != nil {
|
|
return err
|
|
}
|
|
if order != nil {
|
|
current, err := FindingTrafficTx(tx, findingID)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
if len(order) != len(current.Bindings) {
|
|
return ErrEvidenceConflict
|
|
}
|
|
ids := map[int64]bool{}
|
|
for _, b := range current.Bindings {
|
|
ids[b.ID] = true
|
|
}
|
|
for position, id := range order {
|
|
if !ids[id] {
|
|
return ErrEvidenceConflict
|
|
}
|
|
delete(ids, id)
|
|
if _, err := tx.Exec(`UPDATE finding_traffic_bindings SET position=$1 WHERE finding_id=$2 AND id=$3`, position, findingID, id); err != nil {
|
|
return err
|
|
}
|
|
}
|
|
} else {
|
|
var res sql.Result
|
|
var err error
|
|
if remove {
|
|
res, err = tx.Exec(`DELETE FROM finding_traffic_bindings WHERE finding_id=$1 AND id=$2`, findingID, bindingID)
|
|
} else {
|
|
res, err = tx.Exec(`UPDATE finding_traffic_bindings SET role=COALESCE($3,role),note=COALESCE($4,note) WHERE finding_id=$1 AND id=$2`, findingID, bindingID, role, note)
|
|
}
|
|
if err != nil {
|
|
return err
|
|
}
|
|
n, err := res.RowsAffected()
|
|
if err != nil {
|
|
return err
|
|
}
|
|
if n == 0 {
|
|
return ErrEvidenceNotFound
|
|
}
|
|
}
|
|
return bumpEvidenceVersionTx(tx, findingID)
|
|
})
|
|
}
|
|
|
|
type RecordFindingInput struct {
|
|
TaskID, ExplorationID, IntentID int64
|
|
VulnClass, Name, Severity, Summary, Evidence, Worker string
|
|
AssetIDs []int64
|
|
}
|
|
|
|
type RecordedFinding struct {
|
|
FindingID int64 `json:"finding_id,string"`
|
|
NodeID int64 `json:"finding_node_id,string"`
|
|
Traffic *FindingTraffic `json:"traffic"`
|
|
}
|
|
|
|
func RecordFindingTx(tx *sql.Tx, in RecordFindingInput, prepared []PreparedTrafficEvidence) (*RecordedFinding, error) {
|
|
if err := LockTaskEvidenceTx(tx, in.TaskID); err != nil {
|
|
return nil, err
|
|
}
|
|
if in.TaskID > 0 {
|
|
var expID int64
|
|
if err := tx.QueryRow(`SELECT exploration_id FROM tasks WHERE id=$1`, in.TaskID).Scan(&expID); err != nil {
|
|
return nil, err
|
|
}
|
|
if expID != in.ExplorationID {
|
|
return nil, errors.New("漏洞所属任务与探索记录不匹配")
|
|
}
|
|
}
|
|
if in.IntentID > 0 {
|
|
var ok bool
|
|
if err := tx.QueryRow(`SELECT EXISTS(SELECT 1 FROM exploration_nodes WHERE id=$1 AND exploration_id=$2 AND kind='intent')`, in.IntentID, in.ExplorationID).Scan(&ok); err != nil {
|
|
return nil, err
|
|
}
|
|
if !ok {
|
|
return nil, errors.New("intent_id 必须是本任务的意图（关联任务意图只读）")
|
|
}
|
|
}
|
|
payload, _ := json.Marshal(map[string]any{"vulnclass": in.VulnClass, "name": in.Name, "severity": in.Severity, "summary": in.Summary, "evidence": map[string]any{"by": in.Worker, "poc": in.Evidence}})
|
|
out := &RecordedFinding{}
|
|
if err := tx.QueryRow(`INSERT INTO exploration_nodes(exploration_id,kind,payload,priority,state,origin)
|
|
VALUES($1,'finding',$2,9,'confirmed',$3) RETURNING id`, in.ExplorationID, string(payload), in.Worker).Scan(&out.NodeID); err != nil {
|
|
return nil, err
|
|
}
|
|
for _, asset := range in.AssetIDs {
|
|
if _, err := tx.Exec(`INSERT INTO exploration_anchors(node_id,asset_id) VALUES($1,$2) ON CONFLICT DO NOTHING`, out.NodeID, asset); err != nil {
|
|
return nil, err
|
|
}
|
|
}
|
|
if in.IntentID > 0 {
|
|
if _, err := tx.Exec(`INSERT INTO exploration_edges(exploration_id,src_id,rel,dst_id) VALUES($1,$2,$3,$4)`, in.ExplorationID, in.IntentID, RelYields, out.NodeID); err != nil {
|
|
return nil, err
|
|
}
|
|
}
|
|
assets := in.AssetIDs
|
|
if assets == nil {
|
|
assets = []int64{}
|
|
}
|
|
raw, _ := json.Marshal(assets)
|
|
if err := tx.QueryRow(`INSERT INTO findings(task_id,node_id,vulnclass,name,severity,summary,evidence,worker,asset_ids)
|
|
VALUES(NULLIF($1,0),$2,$3,$4,$5,$6,$7,$8,$9) RETURNING id`, in.TaskID, out.NodeID, in.VulnClass, in.Name, in.Severity, in.Summary, in.Evidence, in.Worker, string(raw)).Scan(&out.FindingID); err != nil {
|
|
return nil, err
|
|
}
|
|
if err := AddFindingTrafficTx(tx, out.FindingID, prepared); err != nil {
|
|
return nil, err
|
|
}
|
|
var err error
|
|
out.Traffic, err = FindingTrafficTx(tx, out.FindingID)
|
|
return out, err
|
|
}
|
|
|
|
// RecordFinding is the atomic legacy/no-recorder path used by tools and tests.
|
|
func (s *ExplorationStore) RecordFinding(ctx context.Context, in RecordFindingInput) (out *RecordedFinding, err error) {
|
|
in.ExplorationID = s.expID
|
|
err = s.db.WithEvidenceTx(ctx, func(tx *sql.Tx) error { var e error; out, e = RecordFindingTx(tx, in, nil); return e })
|
|
return
|
|
}
|
|
|
|
func (d *DB) FindingIDByNodeID(nodeID int64) (id int64, err error) {
|
|
err = d.QueryRow(`SELECT id FROM findings WHERE node_id=$1`, nodeID).Scan(&id)
|
|
return
|
|
}
|
|
|
|
func (d *DB) SetFindingReportVersionByNodeID(ctx context.Context, nodeID int64, report string, version *int64) (n int64, err error) {
|
|
err = d.WithEvidenceTx(ctx, func(tx *sql.Tx) error {
|
|
var id int64
|
|
if err := tx.QueryRow(`SELECT id FROM findings WHERE node_id=$1`, nodeID).Scan(&id); err != nil {
|
|
if errors.Is(err, sql.ErrNoRows) {
|
|
return nil
|
|
}
|
|
return err
|
|
}
|
|
if err := LockFindingEvidenceTx(tx, id, version); err != nil {
|
|
return err
|
|
}
|
|
res, err := tx.Exec(`UPDATE findings SET report=$2,report_evidence_version=COALESCE($3::bigint,CASE WHEN evidence_version=0 THEN 0 ELSE -1 END) WHERE id=$1`, id, report, version)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
n, err = res.RowsAffected()
|
|
return err
|
|
})
|
|
return
|
|
}
|