admin/ARTEX
1
0 Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

chore: ARTEX 初始化

Browse Source 
ARTEX — AI 自主渗透测试系统（Go 后端 + Next.js 前端，单二进制内嵌）。

Co-Authored-By: Claude Opus 4.8 (1M context) <[email protected]>
This commit is contained in:
Autumn-27andClaude Opus 4.8 committed 2026-07-26 04:27:45 -04:00
commit 0ef8561411
287 files changed
+70178
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

db/asset_dsl.go
+504
View File
@@ -0,0 +1,504 @@
package db
import (
"fmt"
"strconv"
"strings"
"unicode"
)
// Expr is one leaf DSL clause.
type Expr struct {
Field string // empty = bare-text full-text search
Op string // "=", "==", "!=", ">", ">=", "<", "<="
Value string
}
// astNode is a node in the parsed DSL expression tree.
type astNode struct {
kind string // "and", "or", "leaf"
children []*astNode
expr *Expr // only for "leaf"
}
func andNode(cs []*astNode) *astNode { return &astNode{kind: "and", children: cs} }
func orNode(cs []*astNode) *astNode { return &astNode{kind: "or", children: cs} }
func leafNode(e Expr) *astNode { return &astNode{kind: "leaf", expr: &e} }
// knownStringFields maps DSL field name → SQL column name.
// NOTE: "type" is intentionally excluded — it is a separate parameter, not a DSL field.
var knownStringFields = map[string]string{
"domain": "domain",
"root_domain": "root_domain",
"ip": "ip",
"url": "url",
"page_title": "page_title",
"title": "page_title",
"icp": "icp",
"service_name": "service_name",
"app_name": "app_name",
"method": "method",
"service_type": "service_type",
"record_type": "record_type",
}
// knownArrayFields maps DSL field name → SQL column name (array).
var knownArrayFields = map[string]string{
"technology": "technologies",
"technologies": "technologies",
"tech": "technologies",
}
// knownNumericFields maps DSL field name → SQL column name (integer).
var knownNumericFields = map[string]string{
"port": "port",
"status_code": "status_code",
"status": "status_code",
}
func isKnownField(f string) bool {
f = strings.ToLower(f)
_, s := knownStringFields[f]
_, a := knownArrayFields[f]
_, n := knownNumericFields[f]
return s || a || n || f == "company_id" || f == "task_id"
}
// ── tokeniser ────────────────────────────────────────────────────────────────
const (
tkField = "FIELD"
tkBare = "BARE"
tkAnd = "AND"
tkOr = "OR"
tkLP = "LPAREN"
tkRP = "RPAREN"
tkEOF = "EOF"
)
type tok struct {
kind string
expr *Expr // set for tkField and tkBare
}
func tokenize(s string) ([]tok, error) {
var tokens []tok
i := 0
for i < len(s) {
for i < len(s) && unicode.IsSpace(rune(s[i])) {
i++
}
if i >= len(s) {
break
}
switch s[i] {
case '(':
tokens = append(tokens, tok{kind: tkLP})
i++
case ')':
tokens = append(tokens, tok{kind: tkRP})
i++
default:
if expr, end, ok := tryParseFieldExpr(s, i); ok {
tokens = append(tokens, tok{kind: tkField, expr: &expr})
i = end
continue
}
word, end := readToken(s, i)
if word == "" {
i++
continue
}
switch strings.ToUpper(word) {
case "AND":
tokens = append(tokens, tok{kind: tkAnd})
case "OR":
tokens = append(tokens, tok{kind: tkOr})
default:
e := Expr{Field: "", Op: "=", Value: word}
tokens = append(tokens, tok{kind: tkBare, expr: &e})
}
i = end
}
}
tokens = append(tokens, tok{kind: tkEOF})
return tokens, nil
}
// ── parser ───────────────────────────────────────────────────────────────────
//
// Grammar (AND binds tighter than OR):
// expr = or_expr
// or_expr = and_expr (OR and_expr)*
// and_expr = atom (AND atom)*
// atom = FIELD | BARE | '(' expr ')'
type dslParser struct {
tokens []tok
pos int
}
func (p *dslParser) peek() tok {
if p.pos >= len(p.tokens) {
return tok{kind: tkEOF}
}
return p.tokens[p.pos]
}
func (p *dslParser) consume() tok {
t := p.peek()
p.pos++
return t
}
func (p *dslParser) parseOr() (*astNode, error) {
left, err := p.parseAnd()
if err != nil {
return nil, err
}
children := []*astNode{left}
for p.peek().kind == tkOr {
p.consume()
right, err := p.parseAnd()
if err != nil {
return nil, err
}
children = append(children, right)
}
if len(children) == 1 {
return children[0], nil
}
return orNode(children), nil
}
func (p *dslParser) parseAnd() (*astNode, error) {
left, err := p.parseAtom()
if err != nil {
return nil, err
}
children := []*astNode{left}
for p.peek().kind == tkAnd {
p.consume()
right, err := p.parseAtom()
if err != nil {
return nil, err
}
children = append(children, right)
}
if len(children) == 1 {
return children[0], nil
}
return andNode(children), nil
}
func (p *dslParser) parseAtom() (*astNode, error) {
t := p.peek()
switch t.kind {
case tkField, tkBare:
p.consume()
return leafNode(*t.expr), nil
case tkLP:
p.consume()
node, err := p.parseOr()
if err != nil {
return nil, err
}
if p.peek().kind != tkRP {
return nil, fmt.Errorf("DSL 语法错误：缺少右括号 ')'")
}
p.consume()
return node, nil
case tkEOF:
return nil, fmt.Errorf("DSL 语法错误：表达式不完整")
default:
return nil, fmt.Errorf("DSL 语法错误：意外的 token '%s'", t.kind)
}
}
// ParseDSL parses a DSL query string into an expression tree.
//
// Syntax:
//
// field=value fuzzy match (ILIKE '%value%')
// field==value exact match
// field!=value exclude fuzzy
// port>8080 numeric comparison
// bare word full-text fuzzy across all main text fields
//
// Operators: AND OR (case-insensitive), parentheses for grouping.
// AND binds tighter than OR.
func ParseDSL(s string) (*astNode, error) {
if strings.TrimSpace(s) == "" {
return nil, nil
}
tokens, err := tokenize(s)
if err != nil {
return nil, err
}
p := &dslParser{tokens: tokens}
node, err := p.parseOr()
if err != nil {
return nil, err
}
if p.peek().kind != tkEOF {
return nil, fmt.Errorf("DSL 语法错误：意外的内容 '%s'", p.peek().kind)
}
return node, nil
}
// ── SQL builder ──────────────────────────────────────────────────────────────
// fullTextCols are searched for bare-text tokens.
var fullTextCols = []string{
"domain", "root_domain", "ip", "url", "page_title",
"icp", "service_name", "app_name", "app_description",
}
type whereBuilder struct {
args []any
}
func (b *whereBuilder) next(v any) string {
b.args = append(b.args, v)
return fmt.Sprintf("$%d", len(b.args))
}
func (b *whereBuilder) build(node *astNode) (string, error) {
switch node.kind {
case "and":
parts := make([]string, 0, len(node.children))
for _, child := range node.children {
clause, err := b.build(child)
if err != nil {
return "", err
}
parts = append(parts, "("+clause+")")
}
return strings.Join(parts, " AND "), nil
case "or":
parts := make([]string, 0, len(node.children))
for _, child := range node.children {
clause, err := b.build(child)
if err != nil {
return "", err
}
parts = append(parts, "("+clause+")")
}
return strings.Join(parts, " OR "), nil
case "leaf":
return b.buildLeaf(*node.expr)
}
return "", fmt.Errorf("unknown node kind: %s", node.kind)
}
func (b *whereBuilder) buildLeaf(e Expr) (string, error) {
f := strings.ToLower(e.Field)
// bare-text: OR across all text fields + arrays
if f == "" {
p := b.next("%" + e.Value + "%")
var parts []string
for _, col := range fullTextCols {
parts = append(parts, col+" ILIKE "+p)
}
parts = append(parts,
"EXISTS (SELECT 1 FROM unnest(technologies) t(v) WHERE v ILIKE "+p+")",
"EXISTS (SELECT 1 FROM unnest(bound_domains) t(v) WHERE v ILIKE "+p+")",
)
return "(" + strings.Join(parts, " OR ") + ")", nil
}
// task_id: $N = ANY(task_ids)
if f == "task_id" {
n, err := strconv.ParseInt(e.Value, 10, 64)
if err != nil {
return "", fmt.Errorf("task_id 需要整数值: %s", e.Value)
}
return b.next(n) + " = ANY(task_ids)", nil
}
// company_id: exact integer
if f == "company_id" {
n, err := strconv.ParseInt(e.Value, 10, 64)
if err != nil {
return "", fmt.Errorf("company_id 需要整数值: %s", e.Value)
}
return "company_id = " + b.next(n), nil
}
// numeric fields
if col, ok := knownNumericFields[f]; ok {
n, err := strconv.Atoi(e.Value)
if err != nil {
return "", fmt.Errorf("字段 %s 需要整数值: %s", f, e.Value)
}
op := e.Op
if op == "==" {
op = "="
}
if op != "=" && op != "!=" && op != ">" && op != ">=" && op != "<" && op != "<=" {
return "", fmt.Errorf("字段 %s 不支持运算符 %s", f, e.Op)
}
return fmt.Sprintf("%s %s %s", col, op, b.next(n)), nil
}
// array fields
if col, ok := knownArrayFields[f]; ok {
switch e.Op {
case "==":
return b.next(e.Value) + " = ANY(" + col + ")", nil
case "!=":
return "NOT (" + b.next(e.Value) + " = ANY(" + col + "))", nil
case "=":
p := b.next("%" + e.Value + "%")
return "EXISTS (SELECT 1 FROM unnest(" + col + ") t(v) WHERE v ILIKE " + p + ")", nil
default:
return "", fmt.Errorf("数组字段 %s 不支持运算符 %s", f, e.Op)
}
}
// string fields
if col, ok := knownStringFields[f]; ok {
switch e.Op {
case "=":
return col + " ILIKE " + b.next("%"+e.Value+"%"), nil
case "==":
return col + " = " + b.next(e.Value), nil
case "!=":
return col + " NOT ILIKE " + b.next("%"+e.Value+"%"), nil
default:
return "", fmt.Errorf("字符串字段 %s 不支持运算符 %s", f, e.Op)
}
}
return "", fmt.Errorf("未知字段: %s", f)
}
func buildDSLWhere(node *astNode) (string, []any, error) {
if node == nil {
return "1=1", nil, nil
}
b := &whereBuilder{}
clause, err := b.build(node)
if err != nil {
return "", nil, err
}
return clause, b.args, nil
}
// ── helpers (shared with parser) ─────────────────────────────────────────────
// tryParseFieldExpr tries to parse "field op value" at pos.
func tryParseFieldExpr(s string, pos int) (Expr, int, bool) {
i := pos
if i >= len(s) || !isIdentStart(s[i]) {
return Expr{}, pos, false
}
for i < len(s) && isIdentChar(s[i]) {
i++
}
field := strings.ToLower(s[pos:i])
if !isKnownField(field) {
return Expr{}, pos, false
}
if i >= len(s) {
return Expr{}, pos, false
}
var op string
switch {
case i+1 < len(s) && (s[i] == '=' || s[i] == '!' || s[i] == '>' || s[i] == '<') && s[i+1] == '=':
op = s[i : i+2]
i += 2
case s[i] == '>' || s[i] == '<' || s[i] == '=':
op = string(s[i])
i++
default:
return Expr{}, pos, false
}
value, end := readToken(s, i)
if end == i {
return Expr{}, pos, false
}
return Expr{Field: field, Op: op, Value: value}, end, true
}
// readToken reads a quoted or unquoted token starting at pos.
func readToken(s string, pos int) (string, int) {
if pos >= len(s) {
return "", pos
}
if s[pos] == '"' {
i := pos + 1
for i < len(s) && s[i] != '"' {
i++
}
val := s[pos+1 : i]
if i < len(s) {
i++
}
return val, i
}
i := pos
for i < len(s) && !unicode.IsSpace(rune(s[i])) && s[i] != '(' && s[i] != ')' {
i++
}
return s[pos:i], i
}
func isIdentStart(c byte) bool {
return c == '_' || (c >= 'a' && c <= 'z') || (c >= 'A' && c <= 'Z')
}
func isIdentChar(c byte) bool {
return isIdentStart(c) || (c >= '0' && c <= '9')
}
// ── QueryDSL ─────────────────────────────────────────────────────────────────
// CountDSL returns the total number of assets matching a DSL expression (and optional
// type), for server-side pagination — same WHERE as QueryDSL, without LIMIT/OFFSET.
func (s *AssetStore) CountDSL(dsl, typ string) (int, error) {
node, err := ParseDSL(dsl)
if err != nil {
return 0, err
}
where, args, err := buildDSLWhere(node)
if err != nil {
return 0, err
}
if typ != "" {
args = append(args, typ)
where += fmt.Sprintf(" AND type = $%d", len(args))
}
var n int
err = s.db.QueryRow("SELECT count(*) FROM assets WHERE "+where, args...).Scan(&n)
return n, err
}
// QueryDSL executes a DSL query string against the asset store.
// typ is an optional asset type filter applied independently of the DSL expression.
func (s *AssetStore) QueryDSL(dsl, typ string, limit, offset int) ([]*Asset, error) {
if limit <= 0 {
limit = 50
}
node, err := ParseDSL(dsl)
if err != nil {
return nil, err
}
where, args, err := buildDSLWhere(node)
if err != nil {
return nil, err
}
if typ != "" {
args = append(args, typ)
where += fmt.Sprintf(" AND type = $%d", len(args))
}
args = append(args, limit, offset)
q := assetSelectCols + " WHERE " + where +
fmt.Sprintf(" ORDER BY last_seen DESC LIMIT $%d OFFSET $%d", len(args)-1, len(args))
rows, err := s.db.Query(q, args...)
if err != nil {
return nil, err
}
defer rows.Close()
return scanAssets(rows)
}
db/assets.go
+1103
View File
File diff suppressed because it is too large. Load diff
db/assets_test.go
+705
View File
@@ -0,0 +1,705 @@
package db
import (
"testing"
)
// testSetup opens a DB and returns both stores. Skips if no PG.
func testSetup(t *testing.T) (*DB, *AssetStore, *CompanyStore) {
t.Helper()
d, err := Open(testDSN(t))
if err != nil {
t.Skipf("postgres unavailable (%v)", err)
}
return d, d.Assets(), d.Companies()
}
// deleteAsset removes one v2 asset by id.
func deleteAsset(d *DB, id int64) {
d.Exec(`DELETE FROM assets WHERE id = $1`, id)
}
// =====================================================================
// RootDomain
// =====================================================================
func TestUpsertRootDomain(t *testing.T) {
d, av2, _ := testSetup(t)
defer d.Close()
id1, err := av2.UpsertRootDomain(UpsertRootDomainReq{Domain: "roottest.io", ICP: "A12345", TaskID: 1})
if err != nil {
t.Fatal(err)
}
defer deleteAsset(d, id1)
if id1 == 0 {
t.Fatal("expected non-zero id")
}
// upsert again — same id (dedup), ICP preserved
id2, err := av2.UpsertRootDomain(UpsertRootDomainReq{Domain: "roottest.io", TaskID: 2})
if err != nil {
t.Fatal(err)
}
if id2 != id1 {
t.Errorf("dedup failed: %d != %d", id2, id1)
}
// task_ids should now contain both 1 and 2
var taskIDs []byte
d.QueryRow(`SELECT task_ids FROM assets WHERE id = $1`, id1).Scan(&taskIDs)
// ICP should still be set (COALESCE keeps existing)
var icp *string
d.QueryRow(`SELECT icp FROM assets WHERE id = $1`, id1).Scan(&icp)
if icp == nil || *icp != "A12345" {
t.Errorf("ICP not preserved: %v", icp)
}
}
func TestUpsertRootDomainEmpty(t *testing.T) {
d, av2, _ := testSetup(t)
defer d.Close()
_, err := av2.UpsertRootDomain(UpsertRootDomainReq{Domain: ""})
if err == nil {
t.Error("expected error for empty domain")
}
}
// =====================================================================
// IP
// =====================================================================
func TestUpsertIP(t *testing.T) {
d, av2, _ := testSetup(t)
defer d.Close()
id1, err := av2.UpsertIP(UpsertIPReq{
IP: "192.168.10.5",
OpenPorts: []PortService{{Port: 22, Service: "ssh"}, {Port: 80, Service: "http"}},
TaskID: 1,
})
if err != nil {
t.Fatal(err)
}
defer deleteAsset(d, id1)
// c_segment should be 192.168.10.0/24
var cseg *string
d.QueryRow(`SELECT c_segment::text FROM assets WHERE id = $1`, id1).Scan(&cseg)
if cseg == nil || *cseg != "192.168.10.0/24" {
t.Errorf("c_segment: want 192.168.10.0/24, got %v", cseg)
}
// append another port via UpsertIP (merge)
id2, err := av2.UpsertIP(UpsertIPReq{
IP: "192.168.10.5",
OpenPorts: []PortService{{Port: 443, Service: "https"}},
TaskID: 2,
})
if err != nil {
t.Fatal(err)
}
if id2 != id1 {
t.Errorf("dedup failed: %d != %d", id2, id1)
}
// verify all 3 ports present
var cnt int
d.QueryRow(`SELECT cardinality(open_ports) FROM assets WHERE id = $1`, id1).Scan(&cnt)
if cnt != 3 {
t.Errorf("open_ports merge: want 3 ports, got %d", cnt)
}
}
func TestUpsertIPBoundDomains(t *testing.T) {
d, av2, _ := testSetup(t)
defer d.Close()
id, err := av2.UpsertIP(UpsertIPReq{
IP: "10.1.2.3",
BoundDomains: []string{"a.example.com"},
})
if err != nil {
t.Fatal(err)
}
defer deleteAsset(d, id)
// append a second domain
id2, err := av2.UpsertIP(UpsertIPReq{
IP: "10.1.2.3",
BoundDomains: []string{"b.example.com"},
})
if err != nil {
t.Fatal(err)
}
if id2 != id {
t.Errorf("dedup failed: %d != %d", id2, id)
}
var cnt int
d.QueryRow(`SELECT array_length(bound_domains, 1) FROM assets WHERE id = $1`, id).Scan(&cnt)
if cnt != 2 {
t.Errorf("bound_domains merge: want 2, got %d", cnt)
}
}
func TestAppendIPPort(t *testing.T) {
d, av2, _ := testSetup(t)
defer d.Close()
id, err := av2.UpsertIP(UpsertIPReq{IP: "10.9.8.7"})
if err != nil {
t.Fatal(err)
}
defer deleteAsset(d, id)
if err := av2.AppendIPPort("10.9.8.7", 3306, "mysql"); err != nil {
t.Fatal(err)
}
if err := av2.AppendIPPort("10.9.8.7", 22, "ssh"); err != nil {
t.Fatal(err)
}
var cnt int
d.QueryRow(`SELECT cardinality(open_ports) FROM assets WHERE id = $1`, id).Scan(&cnt)
if cnt != 2 {
t.Errorf("AppendIPPort: want 2 ports, got %d", cnt)
}
}
// =====================================================================
// Subdomain
// =====================================================================
func TestUpsertSubdomain(t *testing.T) {
d, av2, _ := testSetup(t)
defer d.Close()
id, err := av2.UpsertSubdomain(UpsertSubdomainReq{
Domain: "sub.subdtest.com",
RecordType: "A",
RecordValue: []string{"1.2.3.4"},
ICP: "B99",
TaskID: 1,
})
if err != nil {
t.Fatal(err)
}
defer d.Exec(`DELETE FROM assets WHERE domain IN ('sub.subdtest.com', 'subdtest.com') OR ip = '1.2.3.4'`)
if id == 0 {
t.Fatal("expected non-zero id")
}
// verify root_domain auto-populated
var rootDomain string
d.QueryRow(`SELECT COALESCE(root_domain,'') FROM assets WHERE id = $1`, id).Scan(&rootDomain)
if rootDomain != "subdtest.com" {
t.Errorf("root_domain: want subdtest.com, got %q", rootDomain)
}
// side effect: root domain asset should exist
var rootCnt int
d.QueryRow(`SELECT COUNT(*) FROM assets WHERE type = 'root_domain' AND domain = 'subdtest.com'`).Scan(&rootCnt)
if rootCnt != 1 {
t.Error("side-effect root domain not created")
}
// side effect: IP asset with bound_domain
var ipCnt int
d.QueryRow(`SELECT COUNT(*) FROM assets WHERE type = 'ip' AND ip = '1.2.3.4'`).Scan(&ipCnt)
if ipCnt != 1 {
t.Error("side-effect IP asset not created")
}
}
func TestUpsertSubdomainDedup(t *testing.T) {
d, av2, _ := testSetup(t)
defer d.Close()
id1, err := av2.UpsertSubdomain(UpsertSubdomainReq{Domain: "dedup.subdtest2.net", RecordType: "A", RecordValue: []string{"5.5.5.5"}})
if err != nil {
t.Fatal(err)
}
defer d.Exec(`DELETE FROM assets WHERE domain IN ('dedup.subdtest2.net', 'subdtest2.net') OR ip = '5.5.5.5'`)
id2, err := av2.UpsertSubdomain(UpsertSubdomainReq{Domain: "dedup.subdtest2.net", RecordType: "A", RecordValue: []string{"5.5.5.5"}})
if err != nil {
t.Fatal(err)
}
if id2 != id1 {
t.Errorf("subdomain dedup failed: %d vs %d", id2, id1)
}
}
// =====================================================================
// App
// =====================================================================
func TestUpsertAppByBundle(t *testing.T) {
d, av2, _ := testSetup(t)
defer d.Close()
id1, err := av2.UpsertApp(UpsertAppReq{Name: "MyApp", BundleID: "com.example.myapp", Category: "tools"})
if err != nil {
t.Fatal(err)
}
defer deleteAsset(d, id1)
// dedup by bundle
id2, err := av2.UpsertApp(UpsertAppReq{Name: "MyApp Updated", BundleID: "com.example.myapp"})
if err != nil {
t.Fatal(err)
}
if id2 != id1 {
t.Errorf("app bundle dedup failed: %d vs %d", id2, id1)
}
}
func TestUpsertAppByName(t *testing.T) {
d, av2, _ := testSetup(t)
defer d.Close()
id1, err := av2.UpsertApp(UpsertAppReq{Name: "UniqueAppNoBundle", Category: "utility"})
if err != nil {
t.Fatal(err)
}
defer deleteAsset(d, id1)
id2, err := av2.UpsertApp(UpsertAppReq{Name: "UniqueAppNoBundle"})
if err != nil {
t.Fatal(err)
}
if id2 != id1 {
t.Errorf("app name dedup failed: %d vs %d", id2, id1)
}
}
// =====================================================================
// HTTPService
// =====================================================================
func TestUpsertHTTPService(t *testing.T) {
d, av2, _ := testSetup(t)
defer d.Close()
sc := 200
cl := int64(1024)
id1, err := av2.UpsertHTTPService(UpsertHTTPServiceReq{
URL: "https://www.httptest.example.com/",
Technologies: []string{"nginx", "vue"},
StatusCode: &sc,
ContentLength: &cl,
PageTitle: "Test Site",
TaskID: 1,
})
if err != nil {
t.Fatal(err)
}
defer d.Exec(`DELETE FROM assets WHERE url = 'https://www.httptest.example.com' OR domain = 'www.httptest.example.com' OR domain = 'httptest.example.com'`)
if id1 == 0 {
t.Fatal("expected non-zero id")
}
// verify technology was stored
var techCnt int
d.QueryRow(`SELECT array_length(technologies, 1) FROM assets WHERE id = $1`, id1).Scan(&techCnt)
if techCnt != 2 {
t.Errorf("technologies: want 2, got %d", techCnt)
}
// upsert again with additional tech → should merge (append)
id2, err := av2.UpsertHTTPService(UpsertHTTPServiceReq{
URL: "https://www.httptest.example.com/",
Technologies: []string{"react", "webpack"},
TaskID: 2,
})
if err != nil {
t.Fatal(err)
}
if id2 != id1 {
t.Errorf("http service dedup failed: %d vs %d", id2, id1)
}
d.QueryRow(`SELECT array_length(technologies, 1) FROM assets WHERE id = $1`, id1).Scan(&techCnt)
if techCnt != 4 {
t.Errorf("technologies merge: want 4, got %d", techCnt)
}
}
func TestUpsertHTTPServiceAuthAppend(t *testing.T) {
d, av2, _ := testSetup(t)
defer d.Close()
auth1 := []map[string]any{{"type": "basic", "username": "admin", "password": "pass"}}
id1, err := av2.UpsertHTTPService(UpsertHTTPServiceReq{
URL: "http://authtest.example.org/",
Auth: auth1,
TaskID: 1,
})
if err != nil {
t.Fatal(err)
}
defer d.Exec(`DELETE FROM assets WHERE url = 'http://authtest.example.org' OR domain IN ('authtest.example.org', 'example.org')`)
auth2 := []map[string]any{{"type": "bearer", "token": "tok123"}}
id2, err := av2.UpsertHTTPService(UpsertHTTPServiceReq{
URL: "http://authtest.example.org/",
Auth: auth2,
TaskID: 2,
})
if err != nil {
t.Fatal(err)
}
if id2 != id1 {
t.Errorf("dedup failed: %d vs %d", id2, id1)
}
// auth should have 2 items now
var authCnt int
d.QueryRow(`SELECT cardinality(auth) FROM assets WHERE id = $1`, id1).Scan(&authCnt)
if authCnt != 2 {
t.Errorf("auth append: want 2, got %d", authCnt)
}
}
// =====================================================================
// OtherService
// =====================================================================
func TestUpsertOtherService(t *testing.T) {
d, av2, _ := testSetup(t)
defer d.Close()
id1, err := av2.UpsertOtherService(UpsertOtherServiceReq{
IP: "172.16.0.1",
Port: 22,
ServiceName: "ssh",
TaskID: 1,
})
if err != nil {
t.Fatal(err)
}
defer d.Exec(`DELETE FROM assets WHERE (type = 'service' AND service_name = 'ssh' AND ip = '172.16.0.1') OR (type = 'ip' AND ip = '172.16.0.1')`)
if id1 == 0 {
t.Fatal("expected non-zero id")
}
// dedup
id2, err := av2.UpsertOtherService(UpsertOtherServiceReq{
IP: "172.16.0.1",
Port: 22,
ServiceName: "ssh",
})
if err != nil {
t.Fatal(err)
}
if id2 != id1 {
t.Errorf("other service dedup failed: %d vs %d", id2, id1)
}
// side-effect: IP asset should have port 22 in open_ports
var cnt int
d.QueryRow(`SELECT cardinality(open_ports) FROM assets WHERE type='ip' AND ip='172.16.0.1'`).Scan(&cnt)
if cnt == 0 {
t.Error("side-effect: IP open_ports not populated")
}
}
func TestUpsertOtherServiceAuthAppend(t *testing.T) {
d, av2, _ := testSetup(t)
defer d.Close()
auth1 := []map[string]any{{"username": "admin", "password": "secret"}}
id1, err := av2.UpsertOtherService(UpsertOtherServiceReq{
IP: "10.5.5.5",
Port: 3306,
ServiceName: "mysql",
Auth: auth1,
})
if err != nil {
t.Fatal(err)
}
defer d.Exec(`DELETE FROM assets WHERE (type = 'service' AND ip = '10.5.5.5') OR (type = 'ip' AND ip = '10.5.5.5')`)
auth2 := []map[string]any{{"username": "root", "password": "root"}}
id2, err := av2.UpsertOtherService(UpsertOtherServiceReq{
IP: "10.5.5.5",
Port: 3306,
ServiceName: "mysql",
Auth: auth2,
})
if err != nil || id2 != id1 {
t.Fatalf("dedup or error: %v, ids %d vs %d", err, id2, id1)
}
var authCnt int
d.QueryRow(`SELECT cardinality(auth) FROM assets WHERE id = $1`, id1).Scan(&authCnt)
if authCnt != 2 {
t.Errorf("auth append: want 2, got %d", authCnt)
}
}
// =====================================================================
// Endpoint
// =====================================================================
func TestUpsertEndpoint(t *testing.T) {
d, av2, _ := testSetup(t)
defer d.Close()
params1 := []map[string]any{{"location": "query", "name": "id", "value": "1"}}
id1, err := av2.UpsertEndpoint(UpsertEndpointReq{
URL: "https://api.eptest.com/users?id=1",
Method: "GET",
Params: params1,
TaskID: 1,
})
if err != nil {
t.Fatal(err)
}
defer d.Exec(`DELETE FROM assets WHERE url LIKE '%eptest.com%' OR domain IN ('api.eptest.com', 'eptest.com')`)
if id1 == 0 {
t.Fatal("expected non-zero id")
}
// dedup (same URL + method = same endpoint)
id2, err := av2.UpsertEndpoint(UpsertEndpointReq{
URL: "https://api.eptest.com/users?id=1",
Method: "GET",
Params: []map[string]any{{"location": "query", "name": "filter", "value": "active"}},
TaskID: 2,
})
if err != nil {
t.Fatal(err)
}
if id2 != id1 {
t.Errorf("endpoint dedup failed: %d vs %d", id2, id1)
}
// params should be merged (2 distinct params)
var paramCnt int
d.QueryRow(`SELECT cardinality(params) FROM assets WHERE id = $1`, id1).Scan(&paramCnt)
if paramCnt != 2 {
t.Errorf("params merge: want 2, got %d", paramCnt)
}
}
func TestUpsertEndpointRequiredFields(t *testing.T) {
d, av2, _ := testSetup(t)
defer d.Close()
_, err := av2.UpsertEndpoint(UpsertEndpointReq{URL: "", Method: "GET"})
if err == nil {
t.Error("expected error for empty URL")
}
_, err = av2.UpsertEndpoint(UpsertEndpointReq{URL: "https://example.com/", Method: ""})
if err == nil {
t.Error("expected error for empty method")
}
}
// =====================================================================
// Query helpers
// =====================================================================
func TestQueryByType(t *testing.T) {
d, av2, _ := testSetup(t)
defer d.Close()
id, err := av2.UpsertRootDomain(UpsertRootDomainReq{Domain: "querytest.example.net"})
if err != nil {
t.Fatal(err)
}
defer deleteAsset(d, id)
assets, err := av2.QueryByType("root_domain", 10, 0)
if err != nil {
t.Fatal(err)
}
found := false
for _, a := range assets {
if a.ID == id {
found = true
}
}
if !found {
t.Error("QueryByType: inserted asset not found")
}
}
func TestQueryByTask(t *testing.T) {
d, av2, _ := testSetup(t)
defer d.Close()
const taskID = int64(99999)
id, err := av2.UpsertRootDomain(UpsertRootDomainReq{Domain: "taskquery.net", TaskID: taskID})
if err != nil {
t.Fatal(err)
}
defer deleteAsset(d, id)
assets, err := av2.QueryByTask(taskID, "", 10)
if err != nil {
t.Fatal(err)
}
found := false
for _, a := range assets {
if a.ID == id {
found = true
}
}
if !found {
t.Error("QueryByTask: inserted asset not found")
}
}
func TestQueryByCompany(t *testing.T) {
d, av2, cs := testSetup(t)
defer d.Close()
companyID, _, err := cs.UpsertCompany("QueryByCompanyCorp", "")
if err != nil {
t.Fatal(err)
}
defer cleanupCompany(d, companyID)
cs.AddScope(companyID, []string{"qbc-test.io"}, "test")
id, err := av2.UpsertRootDomain(UpsertRootDomainReq{Domain: "qbc-test.io"})
if err != nil {
t.Fatal(err)
}
defer deleteAsset(d, id)
if err := cs.RecomputeAttribution(); err != nil {
t.Fatal(err)
}
assets, err := av2.QueryByCompany(companyID, "root_domain", 10)
if err != nil {
t.Fatal(err)
}
found := false
for _, a := range assets {
if a.ID == id {
found = true
}
}
if !found {
t.Error("QueryByCompany: attributed asset not found")
}
}
func TestCountsByType(t *testing.T) {
d, av2, _ := testSetup(t)
defer d.Close()
id, err := av2.UpsertRootDomain(UpsertRootDomainReq{Domain: "counttest.example.biz"})
if err != nil {
t.Fatal(err)
}
defer deleteAsset(d, id)
counts, err := av2.CountsByType()
if err != nil {
t.Fatal(err)
}
if counts["root_domain"] == 0 {
t.Error("CountsByType: root_domain count should be > 0")
}
}
// =====================================================================
// Company attribution at insert time
// =====================================================================
func TestCompanyAttributionAtInsertTime(t *testing.T) {
d, av2, cs := testSetup(t)
defer d.Close()
companyID, _, err := cs.UpsertCompany("InsertAttrCorp", "")
if err != nil {
t.Fatal(err)
}
defer cleanupCompany(d, companyID)
cs.AddScope(companyID, []string{"insertattr.com"}, "test")
// Now insert a root domain — should be auto-attributed
assetID, err := av2.UpsertRootDomain(UpsertRootDomainReq{Domain: "insertattr.com"})
if err != nil {
t.Fatal(err)
}
defer deleteAsset(d, assetID)
var cid *int64
d.QueryRow(`SELECT company_id FROM assets WHERE id = $1`, assetID).Scan(&cid)
if cid == nil || *cid != companyID {
t.Errorf("attribution at insert time: want company %d, got %v", companyID, cid)
}
}
// =====================================================================
// calcCSegment
// =====================================================================
func TestCalcCSegment(t *testing.T) {
tests := []struct {
ip string
want string
}{
{"192.168.1.5", "192.168.1.0/24"},
{"10.0.0.255", "10.0.0.0/24"},
{"", ""},
{"notanip", ""},
}
for _, tc := range tests {
got := calcCSegment(tc.ip)
if got != tc.want {
t.Errorf("calcCSegment(%q): want %q, got %q", tc.ip, tc.want, got)
}
}
}
// =====================================================================
// marshalStringArray
// =====================================================================
func TestMarshalStringArray(t *testing.T) {
got := marshalStringArray([]string{"a", "b", "c"})
if got != `{"a","b","c"}` {
t.Errorf("unexpected: %q", got)
}
got2 := marshalStringArray(nil)
if got2 != "{}" {
t.Errorf("empty: %q", got2)
}
}
// =====================================================================
// normalizeURL
// =====================================================================
func TestNormalizeURL(t *testing.T) {
tests := []struct {
raw string
want string
}{
{"HTTPS://Example.COM/path", "https://example.com/path"},
{"http://example.com/", "http://example.com"},
{"http://example.com/page", "http://example.com/page"},
}
for _, tc := range tests {
got := normalizeURL(tc.raw)
if got != tc.want {
t.Errorf("normalizeURL(%q): want %q, got %q", tc.raw, tc.want, got)
}
}
}
db/companies.go
+358
View File
@@ -0,0 +1,358 @@
package db
import (
"database/sql"
"fmt"
"net"
"strings"
)
// =====================================================================
// 公司主体层
// =====================================================================
// Company is a row in the companies table.
type Company struct {
ID int64 `json:"id"`
Name string `json:"name"`
NKey string `json:"nkey"`
Logo *string `json:"logo,omitempty"`
CreatedAt string `json:"created_at"`
UpdatedAt string `json:"updated_at"`
}
// CompanyWithScope extends Company with its scope rules and asset count.
type CompanyWithScope struct {
Company
Scope []ScopeRule `json:"scope"`
AssetCount int `json:"asset_count"`
}
// ScopeRule is one company_scope row.
type ScopeRule struct {
ID int64 `json:"id"`
CompanyID int64 `json:"company_id"`
Kind string `json:"kind"`
Domain string `json:"domain,omitempty"`
Net string `json:"net,omitempty"`
Raw string `json:"raw"`
Reason string `json:"reason,omitempty"`
}
// CompanyStore operates on the companies + company_scope tables.
type CompanyStore struct{ db *DB }
// Companies returns the company store.
func (d *DB) Companies() *CompanyStore { return &CompanyStore{db: d} }
// companyNKey normalises a company name: lowercase + trim + collapse whitespace.
func companyNKey(name string) string {
return strings.Join(strings.Fields(strings.ToLower(name)), " ")
}
// UpsertCompany creates or updates a company by name. Returns the id and whether
// a new row was created.
func (s *CompanyStore) UpsertCompany(name, logo string) (id int64, created bool, err error) {
nkey := companyNKey(name)
var logoVal any
if logo != "" {
logoVal = logo
}
err = s.db.QueryRow(`
INSERT INTO companies(name, nkey, logo)
VALUES ($1, $2, $3)
ON CONFLICT (nkey) DO UPDATE SET
name = EXCLUDED.name,
logo = COALESCE(EXCLUDED.logo, companies.logo),
updated_at = now()
RETURNING id, (xmax = 0)`, name, nkey, logoVal).Scan(&id, &created)
return
}
// GetCompany returns one company by id (nil if not found).
func (s *CompanyStore) GetCompany(id int64) (*Company, error) {
c := &Company{}
err := s.db.QueryRow(`
SELECT id, name, nkey, logo, created_at::text, updated_at::text
FROM companies WHERE id = $1`, id).Scan(
&c.ID, &c.Name, &c.NKey, &c.Logo, &c.CreatedAt, &c.UpdatedAt)
if err == sql.ErrNoRows {
return nil, nil
}
return c, err
}
// GetCompanyByName returns one company by normalized name (nil if not found).
func (s *CompanyStore) GetCompanyByName(name string) (*Company, error) {
nkey := companyNKey(name)
c := &Company{}
err := s.db.QueryRow(`
SELECT id, name, nkey, logo, created_at::text, updated_at::text
FROM companies WHERE nkey = $1`, nkey).Scan(
&c.ID, &c.Name, &c.NKey, &c.Logo, &c.CreatedAt, &c.UpdatedAt)
if err == sql.ErrNoRows {
return nil, nil
}
return c, err
}
// UpsertByName creates the company if it doesn't exist, then returns its id.
func (s *CompanyStore) UpsertByName(name string) (int64, error) {
id, _, err := s.UpsertCompany(name, "")
return id, err
}
// DeleteCompany deletes a company (cascades to scope rules; assets.company_id → NULL via FK).
func (s *CompanyStore) DeleteCompany(id int64) error {
_, err := s.db.Exec(`DELETE FROM companies WHERE id = $1`, id)
return err
}
// ListCompanies returns all companies with scope and asset count.
func (s *CompanyStore) ListCompanies() ([]*CompanyWithScope, error) {
rows, err := s.db.Query(`
SELECT c.id, c.name, c.nkey, c.logo, c.created_at::text, c.updated_at::text,
COUNT(DISTINCT a.id) AS asset_count
FROM companies c
LEFT JOIN assets a ON a.company_id = c.id
GROUP BY c.id
ORDER BY c.name`)
if err != nil {
return nil, err
}
defer rows.Close()
var out []*CompanyWithScope
for rows.Next() {
cws := &CompanyWithScope{}
if err := rows.Scan(&cws.ID, &cws.Name, &cws.NKey, &cws.Logo,
&cws.CreatedAt, &cws.UpdatedAt, &cws.AssetCount); err != nil {
return nil, err
}
out = append(out, cws)
}
if err := rows.Err(); err != nil {
return nil, err
}
// fetch scope rules for each company
for _, cws := range out {
cws.Scope, err = s.GetScope(cws.ID)
if err != nil {
return nil, err
}
}
return out, nil
}
// GetScope returns all scope rules for a company.
func (s *CompanyStore) GetScope(companyID int64) ([]ScopeRule, error) {
rows, err := s.db.Query(`
SELECT id, company_id, kind,
COALESCE(domain,''), COALESCE(net::text,''), raw, COALESCE(reason,'')
FROM company_scope
WHERE company_id = $1
ORDER BY id`, companyID)
if err != nil {
return nil, err
}
defer rows.Close()
var out []ScopeRule
for rows.Next() {
var r ScopeRule
if err := rows.Scan(&r.ID, &r.CompanyID, &r.Kind, &r.Domain, &r.Net, &r.Raw, &r.Reason); err != nil {
return nil, err
}
out = append(out, r)
}
return out, rows.Err()
}
// AddScope parses and inserts scope lines for a company, then reattributes assets.
// Returns counts of added, skipped, and invalid lines.
func (s *CompanyStore) AddScope(companyID int64, lines []string, reason string) (added, skipped, invalid int, errors []string) {
for _, line := range lines {
rule, err := ParseScopeLine(line)
if err != nil {
invalid++
errors = append(errors, fmt.Sprintf("%s: %v", line, err))
continue
}
inserted, err := s.insertScopeRule(companyID, rule, reason)
if err != nil {
invalid++
errors = append(errors, fmt.Sprintf("%s: %v", line, err))
continue
}
if inserted {
added++
} else {
skipped++
}
}
// attribute newly-covered assets
if added > 0 {
_ = s.attributeByCompany(companyID)
}
return
}
// insertScopeRule inserts a scope rule. Returns (inserted, error):
// inserted=false means ON CONFLICT DO NOTHING triggered (duplicate), not an error.
func (s *CompanyStore) insertScopeRule(companyID int64, rule ParsedScope, reason string) (inserted bool, err error) {
var res interface{ RowsAffected() (int64, error) }
if rule.Kind == "domain" {
res, err = s.db.Exec(`
INSERT INTO company_scope(company_id, kind, domain, raw, reason)
VALUES ($1, 'domain', $2, $3, $4)
ON CONFLICT ON CONSTRAINT uq_sv2_domain DO NOTHING`,
companyID, rule.Domain, rule.Raw, reason)
} else {
res, err = s.db.Exec(`
INSERT INTO company_scope(company_id, kind, net, raw, reason)
VALUES ($1, $2, $3::cidr, $4, $5)
ON CONFLICT ON CONSTRAINT uq_sv2_net DO NOTHING`,
companyID, rule.Kind, rule.Net, rule.Raw, reason)
}
if err != nil {
return false, err
}
n, _ := res.RowsAffected()
return n > 0, nil
}
// RecomputeAttribution clears all company attribution on assets and rebuilds
// from scope rules. domain rules use root_domain equality; ip/cidr use INET >>= .
func (s *CompanyStore) RecomputeAttribution() error {
tx, err := s.db.Begin()
if err != nil {
return err
}
defer tx.Rollback() //nolint:errcheck
// step 1: clear all attribution
if _, err := tx.Exec(`UPDATE assets SET company_id = NULL`); err != nil {
return err
}
// step 2: domain-based attribution (root_domain exact match)
if _, err := tx.Exec(`
UPDATE assets a SET company_id = (
SELECT cs.company_id FROM company_scope cs
WHERE cs.kind = 'domain'
AND a.root_domain = cs.domain
ORDER BY length(cs.domain) DESC
LIMIT 1
)
WHERE type IN ('root_domain','subdomain','service','endpoint')
AND root_domain IS NOT NULL`); err != nil {
return err
}
// step 3: ip/cidr attribution (only for assets not yet attributed)
if _, err := tx.Exec(`
UPDATE assets a SET company_id = (
SELECT cs.company_id FROM company_scope cs
WHERE cs.kind IN ('ip','cidr')
AND cs.net >>= a.ip::inet
ORDER BY masklen(cs.net) DESC
LIMIT 1
)
WHERE type IN ('ip','subdomain','service','endpoint')
AND ip IS NOT NULL
AND company_id IS NULL`); err != nil {
return err
}
return tx.Commit()
}
// attributeByCompany reattributes only assets matching a specific company's scope.
// Cheaper than full recompute; used when a company's scope is updated.
func (s *CompanyStore) attributeByCompany(companyID int64) error {
tx, err := s.db.Begin()
if err != nil {
return err
}
defer tx.Rollback() //nolint:errcheck
// clear this company's existing claims
if _, err := tx.Exec(`UPDATE assets SET company_id = NULL WHERE company_id = $1`, companyID); err != nil {
return err
}
// domain rules
if _, err := tx.Exec(`
UPDATE assets a SET company_id = $1
FROM company_scope cs
WHERE cs.company_id = $1
AND cs.kind = 'domain'
AND a.root_domain = cs.domain
AND type IN ('root_domain','subdomain','service','endpoint')
AND a.company_id IS NULL`, companyID); err != nil {
return err
}
// ip/cidr rules
if _, err := tx.Exec(`
UPDATE assets a SET company_id = $1
FROM company_scope cs
WHERE cs.company_id = $1
AND cs.kind IN ('ip','cidr')
AND cs.net >>= a.ip::inet
AND type IN ('ip','subdomain','service','endpoint')
AND a.company_id IS NULL`, companyID); err != nil {
return err
}
return tx.Commit()
}
// UpdateScope replaces all scope rules for a company and reattributes.
func (s *CompanyStore) UpdateScope(companyID int64, lines []string, reason string) (added, invalid int, errs []string) {
// delete existing rules
if _, err := s.db.Exec(`DELETE FROM company_scope WHERE company_id = $1`, companyID); err != nil {
errs = append(errs, err.Error())
return
}
added, _, invalid, errs = s.AddScope(companyID, lines, reason)
return
}
// ResolveCompany returns the company_id for a given root_domain and/or ip, or nil
// if no scope rule matches. Mirrors the attribution logic used at asset insert time.
func (s *CompanyStore) ResolveCompany(rootDomain, ipStr string) (*int64, error) {
if rootDomain != "" {
var cid int64
err := s.db.QueryRow(`
SELECT company_id FROM company_scope
WHERE kind = 'domain'
AND domain = $1
ORDER BY length(domain) DESC
LIMIT 1`, rootDomain).Scan(&cid)
if err == nil {
return &cid, nil
}
if err != sql.ErrNoRows {
return nil, err
}
}
if ipStr != "" {
if net.ParseIP(ipStr) != nil {
var cid int64
err := s.db.QueryRow(`
SELECT company_id FROM company_scope
WHERE kind IN ('ip','cidr')
AND net >>= $1::inet
ORDER BY masklen(net) DESC
LIMIT 1`, ipStr).Scan(&cid)
if err == nil {
return &cid, nil
}
if err != sql.ErrNoRows {
return nil, err
}
}
}
return nil, nil
}
db/companies_test.go
+289
View File
@@ -0,0 +1,289 @@
package db
import (
"testing"
)
// cleanup helpers to remove test data
func cleanupCompany(d *DB, id int64) {
d.Exec(`DELETE FROM company_scope WHERE company_id = $1`, id)
d.Exec(`DELETE FROM companies WHERE id = $1`, id)
}
func TestCompanyUpsertAndGet(t *testing.T) {
d, err := Open(testDSN(t))
if err != nil {
t.Skipf("postgres unavailable (%v)", err)
}
defer d.Close()
cs := d.Companies()
id, created, err := cs.UpsertCompany("Test Corp", "https://example.com/logo.png")
if err != nil {
t.Fatal(err)
}
defer cleanupCompany(d, id)
if !created {
t.Error("first upsert should report created=true")
}
// duplicate: same nkey, should not create new
id2, created2, err := cs.UpsertCompany("Test Corp", "")
if err != nil {
t.Fatal(err)
}
if id2 != id {
t.Errorf("dedup failed: %d != %d", id2, id)
}
if created2 {
t.Error("second upsert should report created=false")
}
c, err := cs.GetCompany(id)
if err != nil || c == nil {
t.Fatalf("GetCompany: %v", err)
}
if c.Name != "Test Corp" {
t.Errorf("name: %q", c.Name)
}
// GetCompanyByName
c2, err := cs.GetCompanyByName("test corp") // normalised
if err != nil || c2 == nil {
t.Fatalf("GetCompanyByName: %v", err)
}
if c2.ID != id {
t.Errorf("GetCompanyByName id mismatch: %d vs %d", c2.ID, id)
}
}
func TestCompanyScope(t *testing.T) {
d, err := Open(testDSN(t))
if err != nil {
t.Skipf("postgres unavailable (%v)", err)
}
defer d.Close()
cs := d.Companies()
id, _, err := cs.UpsertCompany("ScopeTestCorp", "")
if err != nil {
t.Fatal(err)
}
defer cleanupCompany(d, id)
lines := []string{"example.com", "192.168.1.0/24", "10.0.0.1"}
added, skipped, invalid, errs := cs.AddScope(id, lines, "test")
if added != 3 {
t.Errorf("want 3 added, got %d (errs: %v)", added, errs)
}
if skipped != 0 || invalid != 0 {
t.Errorf("unexpected skipped=%d invalid=%d", skipped, invalid)
}
// Adding again should skip (duplicate)
added2, skipped2, invalid2, _ := cs.AddScope(id, lines, "test")
if added2 != 0 || skipped2 != 3 {
t.Errorf("want 0 added 3 skipped, got %d added %d skipped %d invalid", added2, skipped2, invalid2)
}
scope, err := cs.GetScope(id)
if err != nil {
t.Fatal(err)
}
if len(scope) != 3 {
t.Errorf("want 3 scope rules, got %d", len(scope))
}
}
func TestCompanyScopeInvalid(t *testing.T) {
d, err := Open(testDSN(t))
if err != nil {
t.Skipf("postgres unavailable (%v)", err)
}
defer d.Close()
cs := d.Companies()
id, _, err := cs.UpsertCompany("InvalidScopeCorp", "")
if err != nil {
t.Fatal(err)
}
defer cleanupCompany(d, id)
// TLD-only domains and overly broad CIDRs should be rejected
lines := []string{"com", "1.2.3.4/8"} // rejected
added, _, invalid, _ := cs.AddScope(id, lines, "test")
if added != 0 {
t.Errorf("want 0 added for invalid lines, got %d", added)
}
if invalid != 2 {
t.Errorf("want 2 invalid, got %d", invalid)
}
}
func TestResolveCompany(t *testing.T) {
d, err := Open(testDSN(t))
if err != nil {
t.Skipf("postgres unavailable (%v)", err)
}
defer d.Close()
cs := d.Companies()
id, _, err := cs.UpsertCompany("ResolveCorp", "")
if err != nil {
t.Fatal(err)
}
defer cleanupCompany(d, id)
cs.AddScope(id, []string{"resolve-test.io", "10.20.0.0/16"}, "test")
// domain match
cid, err := cs.ResolveCompany("resolve-test.io", "")
if err != nil || cid == nil || *cid != id {
t.Errorf("domain resolve: want %d, got %v (err %v)", id, cid, err)
}
// no match
cid2, err := cs.ResolveCompany("notinscope.com", "")
if err != nil || cid2 != nil {
t.Errorf("no-match: want nil, got %v", cid2)
}
// IP/CIDR match
cid3, err := cs.ResolveCompany("", "10.20.5.1")
if err != nil || cid3 == nil || *cid3 != id {
t.Errorf("cidr resolve: want %d, got %v (err %v)", id, cid3, err)
}
// IP outside CIDR
cid4, err := cs.ResolveCompany("", "10.30.0.1")
if err != nil || cid4 != nil {
t.Errorf("cidr no-match: want nil, got %v", cid4)
}
}
func TestUpdateScope(t *testing.T) {
d, err := Open(testDSN(t))
if err != nil {
t.Skipf("postgres unavailable (%v)", err)
}
defer d.Close()
cs := d.Companies()
id, _, err := cs.UpsertCompany("UpdateScopeCorp", "")
if err != nil {
t.Fatal(err)
}
defer cleanupCompany(d, id)
cs.AddScope(id, []string{"old-domain.com"}, "initial")
// UpdateScope replaces
added, invalid, errs := cs.UpdateScope(id, []string{"new-domain.com"}, "replacement")
if added != 1 || invalid != 0 || len(errs) != 0 {
t.Errorf("UpdateScope: added=%d invalid=%d errs=%v", added, invalid, errs)
}
scope, _ := cs.GetScope(id)
if len(scope) != 1 || scope[0].Domain != "new-domain.com" {
t.Errorf("UpdateScope: expected new-domain.com only, got %+v", scope)
}
}
func TestDeleteCompany(t *testing.T) {
d, err := Open(testDSN(t))
if err != nil {
t.Skipf("postgres unavailable (%v)", err)
}
defer d.Close()
cs := d.Companies()
id, _, err := cs.UpsertCompany("DeleteMeCorp", "")
if err != nil {
t.Fatal(err)
}
cs.AddScope(id, []string{"deletetest.com"}, "test")
if err := cs.DeleteCompany(id); err != nil {
t.Fatal(err)
}
c, err := cs.GetCompany(id)
if err != nil || c != nil {
t.Error("expected company to be gone")
}
// scope should be cascade-deleted
scope, _ := cs.GetScope(id)
if len(scope) != 0 {
t.Errorf("expected scope cascade-deleted, got %d rules", len(scope))
}
}
func TestRecomputeAttribution(t *testing.T) {
d, err := Open(testDSN(t))
if err != nil {
t.Skipf("postgres unavailable (%v)", err)
}
defer d.Close()
cs := d.Companies()
as := d.Assets()
id, _, err := cs.UpsertCompany("AttributeTestCorp", "")
if err != nil {
t.Fatal(err)
}
defer cleanupCompany(d, id)
defer d.Exec(`DELETE FROM assets WHERE root_domain = 'attr-test.com'`)
// insert asset before adding scope
assetID, err := as.UpsertRootDomain(UpsertRootDomainReq{Domain: "attr-test.com"})
if err != nil {
t.Fatal(err)
}
defer d.Exec(`DELETE FROM assets WHERE id = $1`, assetID)
// asset should not be attributed yet
var companyID *int64
d.QueryRow(`SELECT company_id FROM assets WHERE id = $1`, assetID).Scan(&companyID)
if companyID != nil {
t.Error("expected no company before scope added")
}
// add scope and recompute
cs.AddScope(id, []string{"attr-test.com"}, "test")
if err := cs.RecomputeAttribution(); err != nil {
t.Fatal(err)
}
d.QueryRow(`SELECT company_id FROM assets WHERE id = $1`, assetID).Scan(&companyID)
if companyID == nil || *companyID != id {
t.Errorf("RecomputeAttribution: expected company %d, got %v", id, companyID)
}
}
func TestListCompanies(t *testing.T) {
d, err := Open(testDSN(t))
if err != nil {
t.Skipf("postgres unavailable (%v)", err)
}
defer d.Close()
cs := d.Companies()
id, _, err := cs.UpsertCompany("ListTestCorp", "")
if err != nil {
t.Fatal(err)
}
defer cleanupCompany(d, id)
companies, err := cs.ListCompanies()
if err != nil {
t.Fatal(err)
}
found := false
for _, c := range companies {
if c.ID == id {
found = true
}
}
if !found {
t.Error("ListCompanies: created company not found")
}
}
db/company_scope.go
+92
View File
@@ -0,0 +1,92 @@
package db
import (
"fmt"
"net"
"strings"
"golang.org/x/net/publicsuffix"
)
// ParsedScope is one parsed asset-scope entry. Exactly one of Domain / Net is set.
// Used internally by ParseScopeLine / ParseScopeLines (passed to CompanyStore).
type ParsedScope struct {
Kind string // "domain" | "ip" | "cidr"
Domain string // normalized registrable/root domain (kind=domain)
Net string // normalized CIDR, single IP as /32 or /128 (kind=ip|cidr)
Raw string // original input line
}
// ParseScopeLine classifies and validates one scope line (root domain / IP /
// CIDR). Guardrails reject bare TLDs and over-broad networks so a rule can never
// swallow the internet. IP ranges must be expressed as CIDR.
func ParseScopeLine(line string) (ParsedScope, error) {
raw := strings.TrimSpace(line)
r := ParsedScope{Raw: raw}
if raw == "" {
return r, fmt.Errorf("空行")
}
s := raw
if i := strings.Index(s, "://"); i >= 0 { // strip scheme if a URL was pasted
s = s[i+3:]
}
// CIDR first — it contains '/', which the path-strip below would remove.
if _, ipnet, err := net.ParseCIDR(s); err == nil {
ones, bits := ipnet.Mask.Size()
if bits == 32 && ones < 16 {
return r, fmt.Errorf("网段过宽(IPv4 需 >= /16): %s", raw)
}
if bits == 128 && ones < 32 {
return r, fmt.Errorf("网段过宽(IPv6 需 >= /32): %s", raw)
}
r.Kind, r.Net = "cidr", ipnet.String()
return r, nil
}
// Strip any path and :port remnants for host-like inputs.
if i := strings.IndexByte(s, '/'); i >= 0 {
s = s[:i]
}
// Single IP.
if ip := net.ParseIP(strings.TrimSpace(s)); ip != nil {
r.Kind = "ip"
if ip.To4() != nil {
r.Net = ip.String() + "/32"
} else {
r.Net = ip.String() + "/128"
}
return r, nil
}
if host, _, ok := strings.Cut(s, ":"); ok { // host:port → host
s = host
}
if strings.Contains(raw, "-") && strings.Count(raw, ".") >= 6 {
return r, fmt.Errorf("IP 段请用 CIDR 表示(如 1.2.3.0/24): %s", raw)
}
// Domain (registrable). Reject bare TLDs / public suffixes.
d := DomainKey(s)
if d == "" || !strings.Contains(d, ".") {
return r, fmt.Errorf("无法识别为根域名/IP/CIDR: %s", raw)
}
if suf, icann := publicsuffix.PublicSuffix(d); icann && suf == d {
return r, fmt.Errorf("不能用裸 TLD 作为范围: %s", raw)
}
r.Kind, r.Domain = "domain", d
return r, nil
}
// ParseScopeLines parses a whole text block (one entry per line), returning the
// valid rules and the invalid lines (with reasons) so the caller can report both.
func ParseScopeLines(text string) (rules []ParsedScope, invalid []string) {
for _, ln := range strings.Split(text, "\n") {
if strings.TrimSpace(ln) == "" {
continue
}
r, err := ParseScopeLine(ln)
if err != nil {
invalid = append(invalid, err.Error())
continue
}
rules = append(rules, r)
}
return rules, invalid
}
db/company_scope_test.go
+142
View File
@@ -0,0 +1,142 @@
package db
import (
"fmt"
"testing"
)
// TestParseScopeLine covers classification + guardrails without a DB.
func TestParseScopeLine(t *testing.T) {
cases := []struct {
in string
kind string
wantErr bool
}{
{"example.com", "domain", false},
{"https://sub.example.com/path", "domain", false},
{"1.2.3.4", "ip", false},
{"10.0.0.0/8", "", true}, // over-broad IPv4 (< /16)
{"198.51.100.0/24", "cidr", false},
{"co.uk", "", true}, // bare public suffix
{"not a host", "", true},
{"1.2.3.1-1.2.3.9", "", true}, // ranges must be CIDR
}
for _, c := range cases {
r, err := ParseScopeLine(c.in)
if c.wantErr {
if err == nil {
t.Errorf("ParseScopeLine(%q) want error, got %+v", c.in, r)
}
continue
}
if err != nil {
t.Errorf("ParseScopeLine(%q) unexpected error: %v", c.in, err)
continue
}
if r.Kind != c.kind {
t.Errorf("ParseScopeLine(%q) kind=%q want %q", c.in, r.Kind, c.kind)
}
}
}
// TestCompanyScopeAttribution exercises the full loop against dev PG: create
// company (unique name), add scope, and verify auto-attribution at insert time,
// backfill of a pre-existing asset, CIDR + domain-suffix matching, and that
// out-of-scope assets stay unattributed.
func TestCompanyScopeAttribution(t *testing.T) {
d, err := Open(testDSN(t))
if err != nil {
t.Skipf("postgres unavailable (%v) — skipping", err)
}
as := d.Assets()
cs := d.Companies()
var startMax int64
if err := d.QueryRow(`SELECT COALESCE(MAX(id),0) FROM assets`).Scan(&startMax); err != nil {
d.Close()
t.Fatalf("startMax: %v", err)
}
t.Cleanup(func() {
_, _ = d.Exec(`DELETE FROM assets WHERE id > $1`, startMax)
d.Close()
})
uniq := startMax + 1
root := fmt.Sprintf("scopetest%d.com", uniq)
sub := "api." + root
ipIn := "198.51.100.9"
ipOut := "203.0.113.9"
outDomain := fmt.Sprintf("other%d.net", uniq)
cid, _, err := cs.UpsertCompany(fmt.Sprintf("ScopeCo %d", uniq), "")
if err != nil {
t.Fatalf("UpsertCompany: %v", err)
}
// a pre-existing asset (inserted BEFORE any scope) — must be back-filled.
preID, err := as.UpsertSubdomain(UpsertSubdomainReq{Domain: sub})
if err != nil {
t.Fatalf("pre upsert: %v", err)
}
var preCompanyID *int64
d.QueryRow(`SELECT company_id FROM assets WHERE id = $1`, preID).Scan(&preCompanyID)
if preCompanyID != nil {
t.Fatalf("pre-scope asset should be unattributed, got %v", *preCompanyID)
}
rules, invalid := ParseScopeLines(root + "\n198.51.100.0/24\nco.uk")
if len(rules) != 2 {
t.Fatalf("want 2 valid rules, got %d (%+v)", len(rules), rules)
}
if len(invalid) != 1 {
t.Fatalf("want 1 invalid line, got %v", invalid)
}
cs.AddScope(cid, []string{root, "198.51.100.0/24"}, "unit test")
mustCid := func(id int64, want int64, label string) {
var cID *int64
d.QueryRow(`SELECT company_id FROM assets WHERE id = $1`, id).Scan(&cID)
if cID == nil {
t.Fatalf("%s company_id = nil, want %d", label, want)
}
if *cID != want {
t.Fatalf("%s company_id = %d, want %d", label, *cID, want)
}
}
mustNil := func(id int64, label string) {
var cID *int64
d.QueryRow(`SELECT company_id FROM assets WHERE id = $1`, id).Scan(&cID)
if cID != nil {
t.Fatalf("%s should be unattributed, got %d", label, *cID)
}
}
// backfill attributed the pre-existing subdomain (domain suffix match).
mustCid(preID, cid, "pre-existing subdomain (backfill)")
// insert-time attribution: ip in CIDR, another subdomain.
ipInID, err := as.UpsertIP(UpsertIPReq{IP: ipIn})
if err != nil {
t.Fatalf("UpsertIP in: %v", err)
}
mustCid(ipInID, cid, "in-CIDR ip (insert-time)")
sub2ID, err := as.UpsertSubdomain(UpsertSubdomainReq{Domain: "www." + root})
if err != nil {
t.Fatalf("UpsertSubdomain: %v", err)
}
mustCid(sub2ID, cid, "new subdomain (insert-time)")
// out of scope stays unattributed.
outID, err := as.UpsertRootDomain(UpsertRootDomainReq{Domain: outDomain})
if err != nil {
t.Fatalf("UpsertRootDomain out: %v", err)
}
mustNil(outID, "out-of-scope domain")
ipOutID, err := as.UpsertIP(UpsertIPReq{IP: ipOut})
if err != nil {
t.Fatalf("UpsertIP out: %v", err)
}
mustNil(ipOutID, "out-of-CIDR ip")
}
db/config.go
+684
View File
@@ -0,0 +1,684 @@
package db
import (
"database/sql"
"encoding/json"
"time"
)
// ---------- LLM profiles ----------
type LLMProfile struct {
ID int64 `json:"id"`
Name string `json:"name"`
Format string `json:"format"`
BaseURL string `json:"base_url,omitempty"`
Proxy string `json:"proxy,omitempty"` // LLM 出站代理(http/https/socks5);空=用环境变量
Model string `json:"model"`
APIKey string `json:"-"` // never serialized to UI
APIKeyHint string `json:"api_key_hint,omitempty"`
RatePerSecond float64 `json:"rate_per_second"`
RatePerMinute float64 `json:"rate_per_minute"`
// ContextWindowK is the model's context window in K tokens, used to size
// compaction thresholds. 0 = use a 200K default; capped at 1000 (1M).
ContextWindowK int `json:"context_window_k"`
// ReasoningEffort is the thinking-mode selector. "" = 默认(不发送思考参数);
// "off" = 显式关闭(thinking.type=disabled); "low"/"medium"/"high"/"max" =
// 开启思考并设强度. Mapped to the provider request in agent.Config.NewProvider.
ReasoningEffort string `json:"reasoning_effort"`
IsDefault bool `json:"is_default"`
}
func (d *DB) ListProfiles() ([]*LLMProfile, error) {
rows, err := d.Query(`SELECT id,name,format,COALESCE(base_url,''),COALESCE(proxy,''),model,COALESCE(api_key_hint,''),rate_per_second,rate_per_minute,context_window_k,COALESCE(reasoning_effort,''),is_default FROM llm_profiles ORDER BY id`)
if err != nil {
return nil, err
}
defer rows.Close()
var out []*LLMProfile
for rows.Next() {
var p LLMProfile
if err := rows.Scan(&p.ID, &p.Name, &p.Format, &p.BaseURL, &p.Proxy, &p.Model, &p.APIKeyHint, &p.RatePerSecond, &p.RatePerMinute, &p.ContextWindowK, &p.ReasoningEffort, &p.IsDefault); err != nil {
return nil, err
}
out = append(out, &p)
}
return out, rows.Err()
}
// ActiveProfile returns the default (active) profile with its api key, or nil.
func (d *DB) ActiveProfile() (*LLMProfile, error) {
var p LLMProfile
err := d.QueryRow(`SELECT id,name,format,COALESCE(base_url,''),COALESCE(proxy,''),model,COALESCE(api_key,''),rate_per_second,rate_per_minute,context_window_k,COALESCE(reasoning_effort,''),is_default FROM llm_profiles WHERE is_default LIMIT 1`).
Scan(&p.ID, &p.Name, &p.Format, &p.BaseURL, &p.Proxy, &p.Model, &p.APIKey, &p.RatePerSecond, &p.RatePerMinute, &p.ContextWindowK, &p.ReasoningEffort, &p.IsDefault)
if err == sql.ErrNoRows {
return nil, nil
}
return &p, err
}
// ProfileByID returns one profile with its api key by id, or nil if not found.
// Used to run a task on a specific (non-default) LLM profile.
func (d *DB) ProfileByID(id int64) (*LLMProfile, error) {
var p LLMProfile
err := d.QueryRow(`SELECT id,name,format,COALESCE(base_url,''),COALESCE(proxy,''),model,COALESCE(api_key,''),rate_per_second,rate_per_minute,context_window_k,COALESCE(reasoning_effort,''),is_default FROM llm_profiles WHERE id=$1`, id).
Scan(&p.ID, &p.Name, &p.Format, &p.BaseURL, &p.Proxy, &p.Model, &p.APIKey, &p.RatePerSecond, &p.RatePerMinute, &p.ContextWindowK, &p.ReasoningEffort, &p.IsDefault)
if err == sql.ErrNoRows {
return nil, nil
}
return &p, err
}
// SaveProfile inserts (id==0) or updates a profile. Empty apiKey on update keeps existing.
func (d *DB) SaveProfile(p *LLMProfile) (int64, error) {
hint := p.APIKeyHint
if len(p.APIKey) >= 4 {
hint = "…" + p.APIKey[len(p.APIKey)-4:]
}
if p.ID == 0 {
var id int64
err := d.QueryRow(`INSERT INTO llm_profiles(name,format,base_url,proxy,model,api_key,api_key_hint,rate_per_second,rate_per_minute,context_window_k,reasoning_effort)
VALUES ($1,$2,NULLIF($3,''),NULLIF($4,''),$5,NULLIF($6,''),NULLIF($7,''),$8,$9,$10,$11) RETURNING id`,
p.Name, p.Format, p.BaseURL, p.Proxy, p.Model, p.APIKey, hint, p.RatePerSecond, p.RatePerMinute, p.ContextWindowK, p.ReasoningEffort).Scan(&id)
return id, err
}
if p.APIKey == "" {
_, err := d.Exec(`UPDATE llm_profiles SET name=$1,format=$2,base_url=NULLIF($3,''),proxy=NULLIF($4,''),model=$5,rate_per_second=$6,rate_per_minute=$7,context_window_k=$8,reasoning_effort=$9 WHERE id=$10`,
p.Name, p.Format, p.BaseURL, p.Proxy, p.Model, p.RatePerSecond, p.RatePerMinute, p.ContextWindowK, p.ReasoningEffort, p.ID)
return p.ID, err
}
_, err := d.Exec(`UPDATE llm_profiles SET name=$1,format=$2,base_url=NULLIF($3,''),proxy=NULLIF($4,''),model=$5,api_key=$6,api_key_hint=$7,rate_per_second=$8,rate_per_minute=$9,context_window_k=$10,reasoning_effort=$11 WHERE id=$12`,
p.Name, p.Format, p.BaseURL, p.Proxy, p.Model, p.APIKey, hint, p.RatePerSecond, p.RatePerMinute, p.ContextWindowK, p.ReasoningEffort, p.ID)
return p.ID, err
}
func (d *DB) DeleteProfile(id int64) error {
_, err := d.Exec(`DELETE FROM llm_profiles WHERE id=$1`, id)
return err
}
// SetActiveProfile makes one profile the global default (single-default invariant).
func (d *DB) SetActiveProfile(id int64) error {
tx, err := d.Begin()
if err != nil {
return err
}
defer tx.Rollback()
if _, err := tx.Exec(`UPDATE llm_profiles SET is_default=false WHERE is_default`); err != nil {
return err
}
if _, err := tx.Exec(`UPDATE llm_profiles SET is_default=true WHERE id=$1`, id); err != nil {
return err
}
return tx.Commit()
}
// ---------- Agents / prompts ----------
type Agent struct {
ID int64 `json:"id"`
Key string `json:"key"`
Name string `json:"name"`
Description string `json:"description"`
Role string `json:"role"`
Builtin bool `json:"builtin"`
Enabled bool `json:"enabled"`
MaxTurns int `json:"max_turns"` // 单次运行最大轮次;0=不限制
RunSecs int `json:"run_seconds"` // worker 单次运行墙钟上限(秒);0=不限制
WebSearch bool `json:"web_search"` // 是否启用网络搜索(受系统全局开关门控)
InteractiveShell bool `json:"interactive_shell"` // 是否启用交互式 shell(持久 PTY 会话工具族)
WrapupPrompt string `json:"wrapup_prompt"` // 收尾提示词(超时/步数耗尽时的 settlement 提示);空=用代码内置默认
WrapupMaxTurns int `json:"wrapup_max_turns"` // 收尾阶段自身的轮数预算;0=用代码内置默认(按 agent)
// 任务级超时收尾词(与 per-run 两套;仅 worker/planner 用);空/0=用代码内置默认。
TaskTimeoutWrapupPrompt string `json:"task_timeout_wrapup_prompt"`
TaskTimeoutWrapupMaxTurns int `json:"task_timeout_wrapup_max_turns"`
}
const agentCols = `id,key,name,COALESCE(description,''),role,builtin,enabled,COALESCE(max_turns,0),COALESCE(run_seconds,600),COALESCE(web_search,false),COALESCE(interactive_shell,false),COALESCE(wrapup_prompt,''),COALESCE(wrapup_max_turns,0),COALESCE(task_timeout_wrapup_prompt,''),COALESCE(task_timeout_wrapup_max_turns,0)`
func scanAgent(sc interface{ Scan(...any) error }) (*Agent, error) {
var a Agent
err := sc.Scan(&a.ID, &a.Key, &a.Name, &a.Description, &a.Role, &a.Builtin, &a.Enabled, &a.MaxTurns, &a.RunSecs, &a.WebSearch, &a.InteractiveShell, &a.WrapupPrompt, &a.WrapupMaxTurns, &a.TaskTimeoutWrapupPrompt, &a.TaskTimeoutWrapupMaxTurns)
return &a, err
}
func (d *DB) ListAgents() ([]*Agent, error) {
rows, err := d.Query(`SELECT ` + agentCols + ` FROM agents ORDER BY id`)
if err != nil {
return nil, err
}
defer rows.Close()
var out []*Agent
for rows.Next() {
a, err := scanAgent(rows)
if err != nil {
return nil, err
}
out = append(out, a)
}
return out, rows.Err()
}
func (d *DB) GetAgentByKey(key string) (*Agent, error) {
a, err := scanAgent(d.QueryRow(`SELECT `+agentCols+` FROM agents WHERE key=$1`, key))
if err == sql.ErrNoRows {
return nil, nil
}
if err != nil {
return nil, err
}
return a, nil
}
// AgentBindingCounts returns per-agent binding counts in a few grouped queries
// (NO N+1): visible MCP servers and skills keyed by agent id, and bound tools
// keyed by agent key (tools.agents is a JSONB array of agent keys). Missing keys
// mean zero. Used to show "MCP N · Skill N · 工具 N" on the agent cards.
func (d *DB) AgentBindingCounts() (mcp map[int64]int, skill map[int64]int, tools map[string]int, err error) {
mcp, skill, tools = map[int64]int{}, map[int64]int{}, map[string]int{}
byID := func(q string, into map[int64]int) error {
rows, e := d.Query(q)
if e != nil {
return e
}
defer rows.Close()
for rows.Next() {
var id int64
var n int
if e := rows.Scan(&id, &n); e != nil {
return e
}
into[id] = n
}
return rows.Err()
}
if err = byID(`SELECT agent_id, count(DISTINCT resource_id) FROM agent_visibility WHERE resource_kind='mcp' AND enabled GROUP BY agent_id`, mcp); err != nil {
return
}
if err = byID(`SELECT agent_id, count(*) FROM agent_skill_visibility WHERE enabled GROUP BY agent_id`, skill); err != nil {
return
}
rows, e := d.Query(`SELECT elem, count(*) FROM tools, jsonb_array_elements_text(agents) AS elem GROUP BY elem`)
if e != nil {
err = e
return
}
defer rows.Close()
for rows.Next() {
var k string
var n int
if e := rows.Scan(&k, &n); e != nil {
err = e
return
}
tools[k] = n
}
err = rows.Err()
return
}
// CreateAgent inserts a custom (builtin=false) conversational agent with role
// 'assistant'. Returns the new row. Callers validate key/name upstream; the DB
// enforces key charset + uniqueness and the role check constraint.
func (d *DB) CreateAgent(key, name, description string) (*Agent, error) {
a, err := scanAgent(d.QueryRow(`
INSERT INTO agents(key, name, description, role, builtin, enabled)
VALUES ($1, $2, NULLIF($3,''), 'assistant', false, true)
RETURNING `+agentCols, key, name, description))
if err != nil {
return nil, err
}
return a, nil
}
// UpdateAgentMeta updates a custom agent's display name + description. Built-in
// agents are left untouched (guarded by the caller / the builtin flag).
func (d *DB) UpdateAgentMeta(key, name, description string) error {
_, err := d.Exec(`UPDATE agents SET name=$2, description=NULLIF($3,'') WHERE key=$1 AND builtin=false`, key, name, description)
return err
}
// DeleteAgent removes a custom agent. Built-in agents are protected by the
// builtin=false guard. agent_prompts / agent_prompt_vars / visibility rows cascade
// via FK; tools.agents bindings for the key are cleaned by the caller.
func (d *DB) DeleteAgent(key string) error {
_, err := d.Exec(`DELETE FROM agents WHERE key=$1 AND builtin=false`, key)
return err
}
// SetAgentMaxTurns updates an agent's max_turns (0 = unlimited).
func (d *DB) SetAgentMaxTurns(key string, maxTurns int) error {
if maxTurns < 0 {
maxTurns = 0
}
_, err := d.Exec(`UPDATE agents SET max_turns=$1 WHERE key=$2`, maxTurns, key)
return err
}
// SetAgentWebSearch toggles whether an agent uses network search (still gated by
// the global web-search master switch + backend/key config).
func (d *DB) SetAgentWebSearch(key string, on bool) error {
_, err := d.Exec(`UPDATE agents SET web_search=$1 WHERE key=$2`, on, key)
return err
}
// SetAgentInteractiveShell toggles whether an agent gets the interactive shell
// (持久 PTY 会话) tool family + Bash 提示词联动(见 docs/交互式shell设计.md §14.2).
func (d *DB) SetAgentInteractiveShell(key string, on bool) error {
_, err := d.Exec(`UPDATE agents SET interactive_shell=$1 WHERE key=$2`, on, key)
return err
}
// SetAgentWrapupPrompt stores an agent's wrap-up (settlement) prompt. Empty string
// means "use the code built-in default" — resolved at runtime by resolveWrapup.
func (d *DB) SetAgentWrapupPrompt(key, prompt string) error {
_, err := d.Exec(`UPDATE agents SET wrapup_prompt=$1 WHERE key=$2`, prompt, key)
return err
}
// SetAgentWrapupMaxTurns stores the wrap-up phase's own turn budget. 0 means "use
// the code built-in default" — resolved at runtime by resolveWrapupTurns.
func (d *DB) SetAgentWrapupMaxTurns(key string, n int) error {
_, err := d.Exec(`UPDATE agents SET wrapup_max_turns=$1 WHERE key=$2`, n, key)
return err
}
// SetAgentTaskTimeoutWrapup stores an agent's task-timeout wrap-up prompt (empty =
// use code built-in default; only worker/planner have one) and its turn budget
// (0 = default). Resolved at runtime by resolveTaskTimeoutWrapup / …Turns.
func (d *DB) SetAgentTaskTimeoutWrapup(key, prompt string, maxTurns int) error {
_, err := d.Exec(`UPDATE agents SET task_timeout_wrapup_prompt=$1, task_timeout_wrapup_max_turns=$2 WHERE key=$3`, prompt, maxTurns, key)
return err
}
// SetAgentRunSeconds updates an agent's run_seconds wall-clock budget (0 = unlimited).
func (d *DB) SetAgentRunSeconds(key string, runSecs int) error {
if runSecs < 0 {
runSecs = 0
}
_, err := d.Exec(`UPDATE agents SET run_seconds=$1 WHERE key=$2`, runSecs, key)
return err
}
type PromptVar struct {
Name string `json:"name"`
Description string `json:"description"`
Example string `json:"example"`
Source string `json:"source"`
}
func (d *DB) PromptVars(agentID int64) ([]PromptVar, error) {
rows, err := d.Query(`SELECT var_name,COALESCE(description,''),COALESCE(example,''),source FROM agent_prompt_vars WHERE agent_id=$1 ORDER BY var_name`, agentID)
if err != nil {
return nil, err
}
defer rows.Close()
out := []PromptVar{}
for rows.Next() {
var v PromptVar
if err := rows.Scan(&v.Name, &v.Description, &v.Example, &v.Source); err != nil {
return nil, err
}
out = append(out, v)
}
return out, rows.Err()
}
// CurrentPrompt returns the agent's active template text ("" if none set yet).
func (d *DB) CurrentPrompt(agentID int64) (string, error) {
var tmpl sql.NullString
err := d.QueryRow(`SELECT p.template_text FROM agents a JOIN agent_prompts p ON p.id=a.current_prompt_id WHERE a.id=$1`, agentID).Scan(&tmpl)
if err == sql.ErrNoRows {
return "", nil
}
return tmpl.String, err
}
// SeedPromptIfEmpty writes the code-default template as the agent's first prompt
// version ONLY when it has none yet (current_prompt_id IS NULL). Mirrors
// SeedTool's first-insert-only philosophy: a user's edited prompt is never
// clobbered on restart. Idempotent — a no-op once any version exists.
func (d *DB) SeedPromptIfEmpty(agentID int64, tmpl string) error {
var cur sql.NullInt64
if err := d.QueryRow(`SELECT current_prompt_id FROM agents WHERE id=$1`, agentID).Scan(&cur); err != nil {
return err
}
if cur.Valid {
return nil // already seeded or user-edited → leave it
}
_, err := d.SavePrompt(agentID, tmpl, "内置默认", "system")
return err
}
// ResetPromptToDefault appends the code-default template as a new version and
// points current at it — the explicit "恢复为内置默认" action.
func (d *DB) ResetPromptToDefault(agentID int64, tmpl string) (int, error) {
return d.SavePrompt(agentID, tmpl, "恢复为内置默认", "system")
}
// SavePrompt appends a new version and points current_prompt_id at it.
func (d *DB) SavePrompt(agentID int64, template, note, by string) (int, error) {
tx, err := d.Begin()
if err != nil {
return 0, err
}
defer tx.Rollback()
var ver int
if err := tx.QueryRow(`SELECT COALESCE(max(version),0)+1 FROM agent_prompts WHERE agent_id=$1`, agentID).Scan(&ver); err != nil {
return 0, err
}
var pid int64
if err := tx.QueryRow(`INSERT INTO agent_prompts(agent_id,version,template_text,note,updated_by) VALUES ($1,$2,$3,NULLIF($4,''),NULLIF($5,'')) RETURNING id`,
agentID, ver, template, note, by).Scan(&pid); err != nil {
return 0, err
}
if _, err := tx.Exec(`UPDATE agents SET current_prompt_id=$1 WHERE id=$2`, pid, agentID); err != nil {
return 0, err
}
return ver, tx.Commit()
}
type PromptVersion struct {
Version int `json:"version"`
Template string `json:"template_text"`
Note string `json:"note,omitempty"`
CreatedAt time.Time `json:"ts"`
}
func (d *DB) ListPromptVersions(agentID int64) ([]PromptVersion, error) {
rows, err := d.Query(`SELECT version,template_text,COALESCE(note,''),created_at FROM agent_prompts WHERE agent_id=$1 ORDER BY version DESC`, agentID)
if err != nil {
return nil, err
}
defer rows.Close()
out := []PromptVersion{}
for rows.Next() {
var v PromptVersion
if err := rows.Scan(&v.Version, &v.Template, &v.Note, &v.CreatedAt); err != nil {
return nil, err
}
out = append(out, v)
}
return out, rows.Err()
}
// ---------- MCP servers ----------
type MCPServer struct {
ID int64 `json:"id"`
Name string `json:"name"`
Transport string `json:"transport"`
Command string `json:"command,omitempty"`
Args json.RawMessage `json:"args"`
Env json.RawMessage `json:"env"`
URL string `json:"url,omitempty"`
Enabled bool `json:"enabled"`
Tools []string `json:"tools,omitempty"` // cached tool names (mcp_tools_cache)
}
func (d *DB) ListMCP() ([]*MCPServer, error) {
rows, err := d.Query(`SELECT id,name,transport,COALESCE(command,''),args,env,COALESCE(url,''),enabled FROM mcp_servers ORDER BY id`)
if err != nil {
return nil, err
}
var out []*MCPServer
for rows.Next() {
var m MCPServer
var args, env []byte
if err := rows.Scan(&m.ID, &m.Name, &m.Transport, &m.Command, &args, &env, &m.URL, &m.Enabled); err != nil {
rows.Close()
return nil, err
}
m.Args, m.Env = json.RawMessage(args), json.RawMessage(env)
out = append(out, &m)
}
if err := rows.Err(); err != nil {
rows.Close()
return nil, err
}
rows.Close() // free the connection before the per-server tool-cache queries below
// Attach each server's cached tool names (best-effort; empty until discovered).
for _, m := range out {
m.Tools, _ = d.MCPToolNames(m.ID)
}
return out, nil
}
// MCPToolNames returns the cached tool names for a server (empty until discovered).
func (d *DB) MCPToolNames(serverID int64) ([]string, error) {
rows, err := d.Query(`SELECT tool_name FROM mcp_tools_cache WHERE server_id=$1 ORDER BY tool_name`, serverID)
if err != nil {
return nil, err
}
defer rows.Close()
var out []string
for rows.Next() {
var n string
if err := rows.Scan(&n); err != nil {
return nil, err
}
out = append(out, n)
}
return out, rows.Err()
}
// MCPTool is one cached tool of an MCP server (name + description).
type MCPTool struct {
Name string `json:"name"`
Description string `json:"description"`
}
// MCPToolsDetailed returns the cached tools (name + description) for a server.
func (d *DB) MCPToolsDetailed(serverID int64) ([]MCPTool, error) {
rows, err := d.Query(`SELECT tool_name, COALESCE(description,'') FROM mcp_tools_cache WHERE server_id=$1 ORDER BY tool_name`, serverID)
if err != nil {
return nil, err
}
defer rows.Close()
var out []MCPTool
for rows.Next() {
var t MCPTool
if err := rows.Scan(&t.Name, &t.Description); err != nil {
return nil, err
}
out = append(out, t)
}
return out, rows.Err()
}
// SaveMCPTools replaces the cached tool list for a server (called after discovery).
func (d *DB) SaveMCPTools(serverID int64, tools []MCPTool) error {
tx, err := d.Begin()
if err != nil {
return err
}
defer tx.Rollback()
if _, err := tx.Exec(`DELETE FROM mcp_tools_cache WHERE server_id=$1`, serverID); err != nil {
return err
}
for _, t := range tools {
if _, err := tx.Exec(`INSERT INTO mcp_tools_cache(server_id, tool_name, description) VALUES ($1,$2,$3)
ON CONFLICT (server_id, tool_name) DO UPDATE SET description=EXCLUDED.description`, serverID, t.Name, t.Description); err != nil {
return err
}
}
return tx.Commit()
}
func (d *DB) SaveMCP(m *MCPServer) (int64, error) {
args, env := string(m.Args), string(m.Env)
if args == "" {
args = "[]"
}
if env == "" {
env = "{}"
}
if m.ID == 0 {
var id int64
err := d.QueryRow(`INSERT INTO mcp_servers(name,transport,command,args,env,url,enabled) VALUES ($1,$2,NULLIF($3,''),$4,$5,NULLIF($6,''),$7) RETURNING id`,
m.Name, m.Transport, m.Command, args, env, m.URL, m.Enabled).Scan(&id)
return id, err
}
_, err := d.Exec(`UPDATE mcp_servers SET name=$1,transport=$2,command=NULLIF($3,''),args=$4,env=$5,url=NULLIF($6,''),enabled=$7 WHERE id=$8`,
m.Name, m.Transport, m.Command, args, env, m.URL, m.Enabled, m.ID)
return m.ID, err
}
func (d *DB) DeleteMCP(id int64) error {
tx, err := d.Begin()
if err != nil {
return err
}
defer tx.Rollback()
if _, err := tx.Exec(`DELETE FROM agent_visibility WHERE resource_kind='mcp' AND resource_id=$1`, id); err != nil {
return err
}
if _, err := tx.Exec(`DELETE FROM mcp_servers WHERE id=$1`, id); err != nil {
return err
}
return tx.Commit()
}
// ---------- Skill visibility (agent × skill_name) ----------
// AgentSkillNames returns the skill directory names visible to an agent.
func (d *DB) AgentSkillNames(agentID int64) ([]string, error) {
rows, err := d.Query(`SELECT skill_name FROM agent_skill_visibility WHERE agent_id=$1 AND enabled ORDER BY skill_name`, agentID)
if err != nil {
return nil, err
}
defer rows.Close()
var out []string
for rows.Next() {
var name string
if err := rows.Scan(&name); err != nil {
return nil, err
}
out = append(out, name)
}
return out, rows.Err()
}
// SkillAgents returns the agent IDs that can see a skill.
func (d *DB) SkillAgents(skillName string) ([]int64, error) {
rows, err := d.Query(`SELECT agent_id FROM agent_skill_visibility WHERE skill_name=$1 AND enabled`, skillName)
if err != nil {
return nil, err
}
defer rows.Close()
var out []int64
for rows.Next() {
var id int64
if err := rows.Scan(&id); err != nil {
return nil, err
}
out = append(out, id)
}
return out, rows.Err()
}
// SetAgentSkillVisibility replaces all skill visibility for an agent.
func (d *DB) SetAgentSkillVisibility(agentID int64, names []string) error {
tx, err := d.Begin()
if err != nil {
return err
}
defer tx.Rollback()
if _, err := tx.Exec(`DELETE FROM agent_skill_visibility WHERE agent_id=$1`, agentID); err != nil {
return err
}
for _, name := range names {
if _, err := tx.Exec(`INSERT INTO agent_skill_visibility(agent_id,skill_name,enabled) VALUES ($1,$2,true)
ON CONFLICT (agent_id,skill_name) DO UPDATE SET enabled=true`, agentID, name); err != nil {
return err
}
}
return tx.Commit()
}
// ToggleSkillVisibility sets one (agent, skill_name) visibility on/off.
func (d *DB) ToggleSkillVisibility(agentID int64, skillName string, on bool) error {
if on {
_, err := d.Exec(`INSERT INTO agent_skill_visibility(agent_id,skill_name,enabled) VALUES ($1,$2,true)
ON CONFLICT (agent_id,skill_name) DO UPDATE SET enabled=true`, agentID, skillName)
return err
}
_, err := d.Exec(`DELETE FROM agent_skill_visibility WHERE agent_id=$1 AND skill_name=$2`, agentID, skillName)
return err
}
// DeleteSkillVisibility removes all visibility rows for a skill (called on skill delete).
func (d *DB) DeleteSkillVisibility(skillName string) error {
_, err := d.Exec(`DELETE FROM agent_skill_visibility WHERE skill_name=$1`, skillName)
return err
}
// ---------- Visibility (agent × mcp) ----------
// AgentVisible returns the resource ids of a kind visible to an agent.
func (d *DB) AgentVisible(agentID int64, kind string) ([]int64, error) {
rows, err := d.Query(`SELECT resource_id FROM agent_visibility WHERE agent_id=$1 AND resource_kind=$2 AND enabled`, agentID, kind)
if err != nil {
return nil, err
}
defer rows.Close()
var out []int64
for rows.Next() {
var id int64
if err := rows.Scan(&id); err != nil {
return nil, err
}
out = append(out, id)
}
return out, rows.Err()
}
// ResourceAgents returns the agent ids that can see a resource.
func (d *DB) ResourceAgents(kind string, resourceID int64) ([]int64, error) {
rows, err := d.Query(`SELECT agent_id FROM agent_visibility WHERE resource_kind=$1 AND resource_id=$2 AND enabled`, kind, resourceID)
if err != nil {
return nil, err
}
defer rows.Close()
var out []int64
for rows.Next() {
var id int64
if err := rows.Scan(&id); err != nil {
return nil, err
}
out = append(out, id)
}
return out, rows.Err()
}
// SetAgentVisibilityKind replaces the full set of visible resources of a kind for
// an agent (agent-side bulk write). Bidirectional with the resource-side view —
// both read/write the same agent_visibility rows.
func (d *DB) SetAgentVisibilityKind(agentID int64, kind string, resourceIDs []int64) error {
tx, err := d.Begin()
if err != nil {
return err
}
defer tx.Rollback()
if _, err := tx.Exec(`DELETE FROM agent_visibility WHERE agent_id=$1 AND resource_kind=$2`, agentID, kind); err != nil {
return err
}
for _, rid := range resourceIDs {
if _, err := tx.Exec(`INSERT INTO agent_visibility(agent_id,resource_kind,resource_id,enabled) VALUES ($1,$2,$3,true)
ON CONFLICT (agent_id,resource_kind,resource_id,mcp_tool_name) DO UPDATE SET enabled=true`, agentID, kind, rid); err != nil {
return err
}
}
return tx.Commit()
}
// ToggleVisibility sets one (agent, kind, resource) visibility on/off (idempotent).
func (d *DB) ToggleVisibility(agentID int64, kind string, resourceID int64, on bool) error {
if on {
_, err := d.Exec(`INSERT INTO agent_visibility(agent_id,resource_kind,resource_id,enabled) VALUES ($1,$2,$3,true)
ON CONFLICT (agent_id,resource_kind,resource_id,mcp_tool_name) DO UPDATE SET enabled=true`, agentID, kind, resourceID)
return err
}
_, err := d.Exec(`DELETE FROM agent_visibility WHERE agent_id=$1 AND resource_kind=$2 AND resource_id=$3`, agentID, kind, resourceID)
return err
}
db/config_test.go
+134
View File
@@ -0,0 +1,134 @@
package db
import (
"encoding/json"
"testing"
)
func TestConfigStores(t *testing.T) {
d, err := Open(testDSN(t))
if err != nil {
t.Skipf("postgres unavailable (%v) — skipping", err)
}
defer d.Close()
// LLM profile: save, active, key never serialized
pid, err := d.SaveProfile(&LLMProfile{Name: "t-default", Format: "openai", Model: "gpt-x", APIKey: "secret123"})
if err != nil {
t.Fatal(err)
}
defer d.Exec(`DELETE FROM llm_profiles WHERE id=$1`, pid)
if err := d.SetActiveProfile(pid); err != nil {
t.Fatal(err)
}
act, err := d.ActiveProfile()
if err != nil || act == nil || act.APIKey != "secret123" {
t.Fatalf("active profile/key: %+v err=%v", act, err)
}
// list must hide the key, expose hint
list, _ := d.ListProfiles()
for _, p := range list {
if p.ID == pid {
b, _ := json.Marshal(p)
if string(b) == "" || contains(string(b), "secret123") {
t.Fatalf("api key leaked in list json: %s", b)
}
if p.APIKeyHint != "…t123" {
t.Fatalf("hint want …t123, got %q", p.APIKeyHint)
}
}
}
// agents seeded; prompt versioning
ag, err := d.GetAgentByKey("planner")
if err != nil || ag == nil {
t.Fatalf("planner agent: %v", err)
}
v1, err := d.SavePrompt(ag.ID, "你是规划者 {{.Goal}}", "init", "test")
if err != nil {
t.Fatal(err)
}
v2, _ := d.SavePrompt(ag.ID, "你是规划者 v2 {{.Goal}} {{.Scope}}", "edit", "test")
if v2 != v1+1 {
t.Fatalf("version should increment: %d -> %d", v1, v2)
}
cur, _ := d.CurrentPrompt(ag.ID)
if cur != "你是规划者 v2 {{.Goal}} {{.Scope}}" {
t.Fatalf("current prompt wrong: %q", cur)
}
vers, _ := d.ListPromptVersions(ag.ID)
if len(vers) < 2 {
t.Fatalf("want >=2 versions, got %d", len(vers))
}
pv, _ := d.PromptVars(ag.ID)
// planner must have at least the seeded catalog vars (Goal, AssetSummary)
if len(pv) < 2 {
t.Fatalf("planner catalog want >=2 vars, got %d", len(pv))
}
d.Exec(`DELETE FROM agent_prompts WHERE agent_id=$1`, ag.ID)
d.Exec(`UPDATE agents SET current_prompt_id=NULL WHERE id=$1`, ag.ID)
// mcp + skill + visibility (bidirectional via one join)
// Clean up any leftover MCP from prior runs to keep this test idempotent.
d.Exec(`DELETE FROM mcp_servers WHERE name = 't-gh'`)
mid, err := d.SaveMCP(&MCPServer{Name: "t-gh", Transport: "stdio", Command: "npx", Args: json.RawMessage(`["server-github"]`), Env: json.RawMessage(`{"GITHUB_TOKEN":"x"}`), Enabled: true})
if err != nil {
t.Fatal(err)
}
// agent-side write. MCP is id-keyed (generic visibility join); skills are now
// filesystem-based, so their visibility is keyed by skill (directory) name in a
// dedicated table. Clear any pre-existing MCP visibility rows so the assertions
// below isolate on exactly what this test sets.
d.Exec(`DELETE FROM agent_visibility WHERE agent_id = $1 AND resource_kind = 'mcp'`, ag.ID)
if err := d.ToggleVisibility(ag.ID, "mcp", mid, true); err != nil {
t.Fatal(err)
}
if err := d.SetAgentSkillVisibility(ag.ID, nil); err != nil {
t.Fatal(err)
}
if err := d.ToggleSkillVisibility(ag.ID, "t-skill", true); err != nil {
t.Fatal(err)
}
// agent-side read
vm, _ := d.AgentVisible(ag.ID, "mcp")
if len(vm) != 1 || vm[0] != mid {
t.Fatalf("agent visible mcp: %+v", vm)
}
// resource-side read (same join row) → bidirectional
ra, _ := d.ResourceAgents("mcp", mid)
if len(ra) != 1 || ra[0] != ag.ID {
t.Fatalf("resource agents: %+v", ra)
}
// toggle off
d.ToggleVisibility(ag.ID, "mcp", mid, false)
vm2, _ := d.AgentVisible(ag.ID, "mcp")
if len(vm2) != 0 {
t.Fatalf("after toggle off: %+v", vm2)
}
// skill visibility is name-keyed: verify the read, then deleting the skill's
// visibility rows (called when a skill is removed) clears it.
names, _ := d.AgentSkillNames(ag.ID)
if len(names) != 1 || names[0] != "t-skill" {
t.Fatalf("agent visible skills: %+v", names)
}
if err := d.DeleteSkillVisibility("t-skill"); err != nil {
t.Fatal(err)
}
names2, _ := d.AgentSkillNames(ag.ID)
if len(names2) != 0 {
t.Fatalf("skill visibility should be cleared on delete: %+v", names2)
}
d.DeleteMCP(mid)
}
func contains(s, sub string) bool {
return len(s) >= len(sub) && (indexOf(s, sub) >= 0)
}
func indexOf(s, sub string) int {
for i := 0; i+len(sub) <= len(s); i++ {
if s[i:i+len(sub)] == sub {
return i
}
}
return -1
}
db/constants.go
+25
View File
@@ -0,0 +1,25 @@
package db
// Exploration node kinds (exploration_nodes.kind).
const (
KindBegin = "begin" // DEPRECATED: legacy task root; new tasks seed an origin fact (KindFact + StateOrigin) instead
KindGoal = "goal" // a task objective
KindIntent = "intent" // an exploration direction (planner-generated)
KindFact = "fact" // a worker's exploration result/conclusion (incl. negative results), tied to its intent
KindFinding = "finding" // a confirmed vulnerability (report_finding), distinct from a fact
KindHint = "hint"
)
// StateOrigin marks the task root fact (KindFact) seeded at task creation — the
// exploration graph's origin. Every intent traces back to it, so "an intent must
// connect to a fact" holds uniformly from the very first intent. Worker-produced
// facts use state 'confirmed', so this never collides.
const StateOrigin = "origin"
// Exploration edge relations (exploration_edges.rel).
const (
RelSpawns = "spawns"
RelDerivedFrom = "derived_from"
RelYields = "yields"
RelProves = "proves"
)
db/conversation.go
+185
View File
@@ -0,0 +1,185 @@
package db
import (
"database/sql"
"time"
)
// ConvTokenSummary is one conversation's token total (sum of its kind='result'
// rows) with its profile + created_at, used to merge conversation usage into the
// dashboard's per-profile / daily token stats (which otherwise cover only tasks).
type ConvTokenSummary struct {
LLMProfileID *int64 `json:"llm_profile_id"`
CreatedAt string `json:"created_at"`
InputTokens int `json:"input_tokens"`
OutputTokens int `json:"output_tokens"`
CacheReadTokens int `json:"cache_read_tokens"`
CacheWriteTokens int `json:"cache_write_tokens"`
}
// ConversationTokenSummaries returns one row per conversation with its summed
// result-row token usage (0 for conversations with no completed run yet).
func (d *DB) ConversationTokenSummaries() ([]ConvTokenSummary, error) {
rows, err := d.Query(`
SELECT c.llm_profile_id, c.created_at::text,
COALESCE(sum(ca.input_tokens),0), COALESCE(sum(ca.output_tokens),0),
COALESCE(sum(ca.cache_read_tokens),0), COALESCE(sum(ca.cache_write_tokens),0)
FROM conversations c
LEFT JOIN conversation_activities ca ON ca.conversation_id = c.id AND ca.kind = 'result'
GROUP BY c.id`)
if err != nil {
return nil, err
}
defer rows.Close()
out := []ConvTokenSummary{}
for rows.Next() {
var s ConvTokenSummary
var pid sql.NullInt64
if err := rows.Scan(&pid, &s.CreatedAt, &s.InputTokens, &s.OutputTokens, &s.CacheReadTokens, &s.CacheWriteTokens); err != nil {
return nil, err
}
if pid.Valid {
v := pid.Int64
s.LLMProfileID = &v
}
out = append(out, s)
}
return out, rows.Err()
}
// Conversation is one ChatGPT-style chat thread bound to an agent key. It lives
// independent of the pentest exploration graph — see schema.sql §I.
type Conversation struct {
ID int64 `json:"id"`
AgentKey string `json:"agent_key"`
Title string `json:"title"`
LLMProfileID *int64 `json:"llm_profile_id,omitempty"`
CreatedAt time.Time `json:"created_at"`
UpdatedAt time.Time `json:"updated_at"`
}
const convCols = `id, agent_key, title, llm_profile_id, created_at, updated_at`
func scanConv(row interface{ Scan(...any) error }) (Conversation, error) {
var c Conversation
err := row.Scan(&c.ID, &c.AgentKey, &c.Title, &c.LLMProfileID, &c.CreatedAt, &c.UpdatedAt)
return c, err
}
// CreateConversation opens a new chat thread for agentKey with an initial title.
// llmProfileID may be nil to use the globally active profile.
func (d *DB) CreateConversation(agentKey, title string, llmProfileID *int64) (*Conversation, error) {
c, err := scanConv(d.QueryRow(`
INSERT INTO conversations(agent_key, title, llm_profile_id) VALUES ($1, $2, $3)
RETURNING `+convCols, agentKey, title, llmProfileID))
if err != nil {
return nil, err
}
return &c, nil
}
// UpdateConversationProfile sets (or clears) the LLM profile override for a conversation.
func (d *DB) UpdateConversationProfile(id int64, llmProfileID *int64) error {
_, err := d.Exec(`UPDATE conversations SET llm_profile_id=$2 WHERE id=$1`, id, llmProfileID)
return err
}
// ListConversations returns all threads, most-recently-updated first.
func (d *DB) ListConversations() ([]*Conversation, error) {
rows, err := d.Query(`SELECT ` + convCols + ` FROM conversations ORDER BY updated_at DESC`)
if err != nil {
return nil, err
}
defer rows.Close()
out := []*Conversation{}
for rows.Next() {
c, err := scanConv(rows)
if err != nil {
return nil, err
}
out = append(out, &c)
}
return out, rows.Err()
}
// GetConversation returns one thread (nil, nil if absent).
func (d *DB) GetConversation(id int64) (*Conversation, error) {
c, err := scanConv(d.QueryRow(`SELECT `+convCols+` FROM conversations WHERE id=$1`, id))
if err == sql.ErrNoRows {
return nil, nil
}
if err != nil {
return nil, err
}
return &c, nil
}
// RenameConversation sets a thread's title.
func (d *DB) RenameConversation(id int64, title string) error {
_, err := d.Exec(`UPDATE conversations SET title=$2 WHERE id=$1`, id, title)
return err
}
// TouchConversation bumps updated_at so the thread floats to the top of the list.
func (d *DB) TouchConversation(id int64) error {
_, err := d.Exec(`UPDATE conversations SET updated_at=now() WHERE id=$1`, id)
return err
}
// DeleteConversation removes a thread; its activities cascade via FK.
func (d *DB) DeleteConversation(id int64) error {
_, err := d.Exec(`DELETE FROM conversations WHERE id=$1`, id)
return err
}
// AppendConvActivity records one step of a conversation (human message or an agent
// execution step) and returns its id. Mirrors ExplorationStore.AppendActivity but
// keyed by conversation_id. Reuses the Activity struct (NodeID is ignored here).
func (d *DB) AppendConvActivity(convID int64, a Activity) (int64, error) {
var id int64
err := d.QueryRow(`
INSERT INTO conversation_activities(conversation_id, worker, kind, tool, tool_use_id, is_error, summary, detail, input_tokens, output_tokens, cache_read_tokens, cache_write_tokens)
VALUES ($1,NULLIF($2,''),NULLIF($3,''),NULLIF($4,''),NULLIF($5,''),$6,NULLIF($7,''),NULLIF($8,''),$9,$10,$11,$12)
RETURNING id`, convID, utf8Clean(a.Worker), utf8Clean(a.Kind), utf8Clean(a.Tool), utf8Clean(a.ToolUseID), a.IsError,
utf8Clean(a.Summary), utf8Clean(a.Detail), a.InputTokens, a.OutputTokens, a.CacheReadTokens, a.CacheWriteTokens).Scan(&id)
return id, err
}
// ConvActivityList returns a conversation's steps after sinceID (exclusive) with
// the summary-only column set (detail is lazy-loaded via ConvActivityDetail).
func (d *DB) ConvActivityList(convID, sinceID int64, limit int) ([]Activity, int64, error) {
if limit <= 0 {
limit = 500
}
const cols = `id, COALESCE(worker,''), COALESCE(kind,''), COALESCE(tool,''), COALESCE(tool_use_id,''), is_error, COALESCE(summary,''), created_at, input_tokens, output_tokens, cache_read_tokens, cache_write_tokens`
rows, err := d.Query(`SELECT `+cols+`
FROM conversation_activities WHERE conversation_id=$1 AND id>$2 ORDER BY id LIMIT $3`, convID, sinceID, limit)
if err != nil {
return nil, sinceID, err
}
defer rows.Close()
out := []Activity{}
cursor := sinceID
for rows.Next() {
var a Activity
if err := rows.Scan(&a.ID, &a.Worker, &a.Kind, &a.Tool, &a.ToolUseID, &a.IsError, &a.Summary, &a.CreatedAt,
&a.InputTokens, &a.OutputTokens, &a.CacheReadTokens, &a.CacheWriteTokens); err != nil {
return nil, sinceID, err
}
if a.ID > cursor {
cursor = a.ID
}
out = append(out, a)
}
return out, cursor, rows.Err()
}
// ConvActivityDetail lazily returns the full detail blob for one step.
func (d *DB) ConvActivityDetail(convID, id int64) (string, error) {
var s sql.NullString
err := d.QueryRow(`SELECT detail FROM conversation_activities WHERE id=$1 AND conversation_id=$2`, id, convID).Scan(&s)
if err == sql.ErrNoRows {
return "", nil
}
return s.String, err
}
db/customtool_test.go
+82
View File
@@ -0,0 +1,82 @@
package db
import (
"encoding/json"
"testing"
)
// TestCustomToolCRUD exercises create/list/update/delete of a user-defined tool
// (system=false) with kind/exec/deferred. Skips when no Postgres is configured.
func TestCustomToolCRUD(t *testing.T) {
d, err := Open(testDSN(t))
if err != nil {
t.Skipf("postgres unavailable (%v) — skipping", err)
}
defer d.Close()
key := "ct_test_tool"
_ = d.DeleteCustomTool(key) // clean slate
in := &Tool{
Key: key,
Description: "test",
Schema: json.RawMessage(`{"type":"object","properties":{"x":{"type":"string"}}}`),
Agents: []string{"worker"},
Enabled: true,
Kind: "script",
Exec: json.RawMessage(`{"code":"print(1)"}`),
Deferred: true,
}
if err := d.CreateCustomTool(in); err != nil {
t.Fatalf("create: %v", err)
}
t.Cleanup(func() { _ = d.DeleteCustomTool(key) })
got, err := d.GetTool(key)
if err != nil || got == nil {
t.Fatalf("get after create: %v (nil=%v)", err, got == nil)
}
if got.System || got.Kind != "script" || !got.Deferred || len(got.Agents) != 1 {
t.Fatalf("unexpected row: %+v", got)
}
// custom tools appear in ListCustomTools, not among system-only.
customs, err := d.ListCustomTools()
if err != nil {
t.Fatal(err)
}
found := false
for _, c := range customs {
if c.Key == key {
found = true
if c.System {
t.Fatal("custom tool should have system=false")
}
}
}
if !found {
t.Fatal("created tool not in ListCustomTools")
}
// update: change kind + disable + rebind.
in.Kind = "command"
in.Exec = json.RawMessage(`{"command":"echo hi"}`)
in.Enabled = false
in.Agents = []string{"worker", "planner"}
in.Deferred = false
if err := d.UpdateCustomTool(in); err != nil {
t.Fatalf("update: %v", err)
}
got2, _ := d.GetTool(key)
if got2.Kind != "command" || got2.Enabled || got2.Deferred || len(got2.Agents) != 2 {
t.Fatalf("update not applied: %+v", got2)
}
// delete only removes non-system rows.
if err := d.DeleteCustomTool(key); err != nil {
t.Fatalf("delete: %v", err)
}
if g, _ := d.GetTool(key); g != nil {
t.Fatal("tool still present after delete")
}
}
db/db.go
+419
View File
@@ -0,0 +1,419 @@
// Package db is the PostgreSQL data source for ARTEX (取代旧 graph 单文件 SQLite)。
// 它打开连接、应用 schema、并 seed 内置 agent 与变量目录。
package db
import (
"database/sql"
_ "embed"
"fmt"
"net/url"
"strings"
"github.com/Autumn-27/artex/config"
_ "github.com/jackc/pgx/v5/stdlib" // pgx database/sql driver ("pgx")
)
//go:embed schema.sql
var schemaSQL string
// DSN resolves the PostgreSQL connection string and reports where it came from.
// Precedence: env ARTEX_PG_DSN > config file (config.json). There is no
// built-in default — it errors if neither source is configured.
func DSN() (dsn, source string, err error) {
return config.PostgresDSN()
}
// DB wraps the shared *sql.DB. PG handles its own connection pool + concurrency
// (MVCC), so unlike the old SQLite store there is no process-wide write mutex.
type DB struct{ *sql.DB }
// ensureDatabase connects to the postgres system database and creates the target
// database if it does not exist. dsn must be a postgres:// URL.
func ensureDatabase(dsn string) error {
u, err := url.Parse(dsn)
if err != nil {
return nil // unparseable DSN — let the normal Open fail with a clear error
}
dbName := strings.TrimPrefix(u.Path, "/")
if dbName == "" || dbName == "postgres" {
return nil
}
// connect to the postgres maintenance database instead
adminDSN := *u
adminDSN.Path = "/postgres"
admin, err := sql.Open("pgx", adminDSN.String())
if err != nil {
return nil // best-effort; let Open surface the real error
}
defer admin.Close()
if err := admin.Ping(); err != nil {
return nil
}
var exists bool
_ = admin.QueryRow(`SELECT true FROM pg_database WHERE datname=$1`, dbName).Scan(&exists)
if !exists {
if _, err := admin.Exec(`CREATE DATABASE "` + dbName + `"`); err != nil {
return fmt.Errorf("create database %q: %w", dbName, err)
}
}
return nil
}
// Open connects, applies the schema (idempotent), and seeds builtin rows.
func Open(dsn string) (*DB, error) {
if err := ensureDatabase(dsn); err != nil {
return nil, err
}
sqlDB, err := sql.Open("pgx", dsn)
if err != nil {
return nil, err
}
if err := sqlDB.Ping(); err != nil {
sqlDB.Close()
return nil, fmt.Errorf("ping postgres (%s): %w", config.Redact(dsn), err)
}
// Serialize schema migration + seeding across concurrent openers (e.g. parallel
// test packages) using a session-level advisory lock. The lock is released
// automatically when the connection is returned to the pool or closed.
if _, err := sqlDB.Exec(`SELECT pg_advisory_lock(7337741001)`); err != nil {
sqlDB.Close()
return nil, fmt.Errorf("advisory lock: %w", err)
}
// pgx runs multi-statement Exec via the simple protocol when there are no args.
if _, err := sqlDB.Exec(schemaSQL); err != nil {
sqlDB.Exec(`SELECT pg_advisory_unlock(7337741001)`) //nolint:errcheck
sqlDB.Close()
return nil, fmt.Errorf("apply schema: %w", err)
}
d := &DB{sqlDB}
if err := d.seedBuiltins(); err != nil {
sqlDB.Exec(`SELECT pg_advisory_unlock(7337741001)`) //nolint:errcheck
sqlDB.Close()
return nil, fmt.Errorf("seed builtins: %w", err)
}
if _, err := sqlDB.Exec(`SELECT pg_advisory_unlock(7337741001)`); err != nil {
sqlDB.Close()
return nil, fmt.Errorf("advisory unlock: %w", err)
}
return d, nil
}
// builtinAgent describes one of the fixed agents and its prompt-variable catalog.
type builtinAgent struct {
key, name, role, desc string
vars []promptVar
interactiveShell bool // 建行时的默认交互式 shell 开关；ON CONFLICT 不覆盖用户后续手动开关
runSeconds *int // 建行时的单次 run 墙钟上限(秒)；nil=用表默认(600)，0=不限时
}
type promptVar struct{ name, desc, example, source string }
// intp 返回 v 的指针，用于给 builtinAgent 可选字段(如 runSeconds)显式取值。
func intp(v int) *int { return &v }
// builtinAgents mirrors docs §5(a). 内置工具不入库；这里只 seed agent + 变量目录。
// 注：planner/worker/mainagent/auto 的交互式 shell 默认由下方 interactive_shell_default_v1
// 块统一置 true（尊重后续 toggle）；这里的 interactiveShell 只给需要「建行即默认开」的新 agent。
var builtinAgents = []builtinAgent{
{"goals", "目标拆解", "goals", "把渗透任务目标拆解成若干独立、可验证的子目标。", []promptVar{
{"EngagementDescription", "任务描述（测试对象/背景）", "测试 example.com 站点", "exploration"},
{"Now", "服务端当前时间(RFC3339)", "2026-06-25T00:00:00Z", "runtime"},
}, false, nil},
{"planner", "规划", "planner", "读取态势、判定目标，只在确有未覆盖的新方向时补充探索意图（每任务一个规划循环）。", []promptVar{
{"Goal", "任务总目标", "拿下 example.com 的管理员权限", "exploration"},
{"AssetSummary", "资产计数/类型分布摘要(可选)", "domain:3 ip:5 site:2", "distilled"},
}, false, nil},
{"mainagent", "主", "main", "人机接口：观察进展，把人的意图落成 hint 或高优先级意图。", []promptVar{
{"Goal", "当前任务目标", "拿下 example.com 的管理员权限", "exploration"},
{"AssetSummary", "开局态势摘要(可选)", "domain:3 ip:5", "distilled"},
{"FindingsSummary", "已确认漏洞摘要(可选)", "high:1 medium:2", "distilled"},
}, false, nil},
{"worker", "执行", "worker", "领取一条意图执行，把发现的事实/漏洞写回知识图谱后停止。", []promptVar{
{"ProxyAddr", "记录代理地址(驱动 if 双文案)", "127.0.0.1:8080", "runtime"},
{"WorkerName", "worker 自我标识(可选)", "worker-1", "runtime"},
}, false, nil},
// Auto:内置「平台操作」agent。不参与渗透编排循环,经对话页驱动,用工具操作平台。
{"auto", "Auto", "assistant", "平台操作助手：用工具管理任务(建/看/暂停/给提示)与资产，并可创建/修改 skill、自定义工具、MCP。", nil, false, nil},
// 渗透测试:内置「独立渗透」agent。经对话页驱动,一人从侦察到收尾走完整条渗透链,自己规划自己执行自己验证。默认开启交互式 shell。
{"pentest", "渗透测试", "assistant", "独立渗透 agent：一人从侦察→找攻击面→深入利用→验证→收尾走完整条链，自己规划、自己执行、自己对抗式验证。", nil, true, intp(0)},
}
// seedBuiltins inserts the fixed built-in agents and their variable catalog (idempotent).
func (d *DB) seedBuiltins() error {
for _, a := range builtinAgents {
var agentID int64
err := d.QueryRow(`
INSERT INTO agents(key, name, description, role, builtin, enabled, interactive_shell, run_seconds)
VALUES ($1, $2, NULLIF($3,''), $4, true, true, $5, COALESCE($6, 600))
ON CONFLICT (key) DO UPDATE SET name = EXCLUDED.name, description = EXCLUDED.description
RETURNING id`, a.key, a.name, a.desc, a.role, a.interactiveShell, a.runSeconds).Scan(&agentID)
if err != nil {
return fmt.Errorf("agent %s: %w", a.key, err)
}
for _, v := range a.vars {
if _, err := d.Exec(`
INSERT INTO agent_prompt_vars(agent_id, var_name, description, example, source)
VALUES ($1, $2, $3, $4, $5)
ON CONFLICT (agent_id, var_name) DO UPDATE
SET description = EXCLUDED.description, example = EXCLUDED.example, source = EXCLUDED.source`,
agentID, v.name, v.desc, v.example, v.source); err != nil {
return fmt.Errorf("agent %s var %s: %w", a.key, v.name, err)
}
}
}
// Drop catalog entries for variables that were renamed, so the white-list no
// longer advertises a name templates can't resolve (EngagementTitle→Description).
if _, err := d.Exec(`DELETE FROM agent_prompt_vars WHERE var_name IN ('EngagementTitle', 'CoverageGaps')`); err != nil {
return fmt.Errorf("cleanup renamed vars: %w", err)
}
// Default-on interactive_shell for the runtime agents (planner/worker/mainagent/auto)
// ONCE — respects a later user toggle-off (guarded by a settings flag). goals(one-shot
// decomposer) stays off. Runs after the column exists (schema applied before seed).
if v, _, _ := d.GetSetting("interactive_shell_default_v1"); v != "true" {
if _, err := d.Exec(`UPDATE agents SET interactive_shell=true WHERE key IN ('planner','worker','mainagent','auto')`); err != nil {
return fmt.Errorf("seed interactive_shell defaults: %w", err)
}
_ = d.SetSetting("interactive_shell_default_v1", "true")
}
// Seed the built-in browser (Playwright) MCP once — DISABLED by default (用户
// 需要时自行启用), no proxy by default. The traffic-capture toggle injects/strips
// the recording proxy + CA at runtime (server.Manager.syncBrowserMCPProxy).
// Insert only if absent so we never clobber user edits (args/env/enabled/
// visibility) on restart.
if _, err := d.Exec(`
INSERT INTO mcp_servers(name, transport, command, args, env, enabled)
VALUES ('browser', 'stdio', 'npx', $1, '{}', false)
ON CONFLICT (name) DO NOTHING`,
`["-y","@playwright/mcp@latest","--headless"]`); err != nil {
return fmt.Errorf("seed browser mcp: %w", err)
}
// NOTE: the placeholder ScopeSentry data-source MCP (empty URL + empty X-API-Key,
// disabled) is seeded directly in schema.sql §F so a raw `psql < schema.sql` init
// also gets it. schema.sql is Exec'd on every startup, so it stays idempotent.
if err := d.seedBuiltinSkillVisibility(); err != nil {
return fmt.Errorf("seed skill visibility: %w", err)
}
if err := d.seedDefaultInterceptRules(); err != nil {
return fmt.Errorf("seed intercept rules: %w", err)
}
return nil
}
// builtinSkillVisibility maps a shipped skill's directory name → the built-in
// agent keys that should see it by default. The skill FILES themselves live on the
// filesystem (SkillDir, loaded by norma at runtime); DB only carries this visibility
// binding. Skills omitted here (e.g. playwright-cli, scopesentry) ship invisible by
// default — the user turns them on per-agent when needed. scopesentry additionally
// declares `mcps: ScopeSentry`, which only takes effect once it's made visible and
// that MCP is enabled/configured.
var builtinSkillVisibility = map[string][]string{
"api-recon": {"auto", "pentest", "worker"},
}
// seedBuiltinSkillVisibility binds the shipped built-in skills to their default
// agents. Insert-if-absent (ON CONFLICT DO NOTHING) so a user's later toggle-off is
// never resurrected on restart — matches the browser-MCP / intercept-rule seed policy.
func (d *DB) seedBuiltinSkillVisibility() error {
for skillName, agentKeys := range builtinSkillVisibility {
for _, key := range agentKeys {
if _, err := d.Exec(`
INSERT INTO agent_skill_visibility(agent_id, skill_name, enabled)
SELECT id, $2, true FROM agents WHERE key=$1
ON CONFLICT (agent_id, skill_name) DO NOTHING`, key, skillName); err != nil {
return fmt.Errorf("skill %s → agent %s: %w", skillName, key, err)
}
}
}
return nil
}
// seedDefaultInterceptRules inserts built-in safety intercept rules once on
// first startup. The seed is gated by a settings flag so user edits (disable,
// delete, re-order) are never overwritten on subsequent restarts.
func (d *DB) seedDefaultInterceptRules() error {
if v, _, _ := d.GetSetting("intercept_default_rules_v1"); v == "done" {
return nil
}
type rule struct {
name string
target string // tool_name | tool_input
typ string // string | regex
pattern string
action string
message string
priority int
}
rules := []rule{
// ── 系统破坏性命令 (priority 100) ──────────────────────────────────
{
name: "[内置] 递归强制删除 rm -rf",
target: "tool_input",
typ: "regex",
pattern: `(?i)\brm\b.{0,80}(?:-[a-z]*r[a-z]*f[a-z]*|-[a-z]*f[a-z]*r[a-z]*|--recursive|--no-preserve-root)`,
action: "deny",
message: "禁止执行递归强制删除（rm -rf / rm --recursive），可能永久损坏系统或靶机环境",
priority: 100,
},
{
name: "[内置] 删除系统关键目录",
target: "tool_input",
typ: "regex",
pattern: `\brm\b[^"'\n]{0,60}["'\s](/|/etc|/bin|/usr|/boot|/var|/lib|/sys|/proc|/dev|/sbin|/root)`,
action: "deny",
message: "禁止删除系统关键路径",
priority: 100,
},
{
name: "[内置] 磁盘格式化 mkfs",
target: "tool_input",
typ: "regex",
pattern: `\bmkfs\b`,
action: "deny",
message: "禁止格式化磁盘（mkfs）",
priority: 100,
},
{
name: "[内置] 覆写磁盘设备 dd",
target: "tool_input",
typ: "regex",
pattern: `\bdd\b[^|\n]{0,100}\bof=\s*/dev/[a-zA-Z]`,
action: "deny",
message: "禁止使用 dd 覆写磁盘设备",
priority: 100,
},
{
name: "[内置] Fork 炸弹",
target: "tool_input",
typ: "regex",
pattern: `:\(\)\s*\{[^}]*:\|:`,
action: "deny",
message: "禁止执行 Fork 炸弹",
priority: 100,
},
{
name: "[内置] 关机 / 重启",
target: "tool_input",
typ: "regex",
pattern: `\b(?:shutdown|reboot|halt|poweroff|init\s+[06])\b`,
action: "deny",
message: "禁止执行关机或重启命令",
priority: 100,
},
{
name: "[内置] 杀死全部进程",
target: "tool_input",
typ: "regex",
pattern: `\bkill\s+-9\s+-1\b|\bkillall\s+-9\b`,
action: "deny",
message: "禁止 kill -9 -1 或 killall -9（杀死所有进程）",
priority: 100,
},
{
name: "[内置] 磁盘擦除 shred / wipe",
target: "tool_input",
typ: "regex",
pattern: `\b(?:shred|wipe)\b[^|\n]{0,80}/dev/[a-zA-Z]`,
action: "deny",
message: "禁止对磁盘设备执行 shred/wipe 擦除",
priority: 100,
},
{
name: "[内置] 清空防火墙规则",
target: "tool_input",
typ: "regex",
pattern: `\biptables\s+(?:-F|--flush)\b|\bnft\s+flush\s+ruleset\b`,
action: "deny",
message: "禁止清空防火墙规则（iptables -F / nft flush）",
priority: 100,
},
// ── 数据库破坏性操作 (priority 90) ─────────────────────────────────
{
name: "[内置] SQL DROP DATABASE / TABLE / SCHEMA",
target: "tool_input",
typ: "regex",
pattern: `(?i)\bDROP\s+(?:DATABASE|TABLE|SCHEMA|INDEX|VIEW|TABLESPACE|USER|ROLE)\b`,
action: "deny",
message: "禁止执行 DROP 操作，可能不可逆地销毁数据库对象",
priority: 90,
},
{
name: "[内置] SQL TRUNCATE",
target: "tool_input",
typ: "regex",
pattern: `(?i)\bTRUNCATE\s+(?:TABLE\s+)?\w`,
action: "deny",
message: "禁止执行 TRUNCATE，可能清空数据表所有数据",
priority: 90,
},
{
name: "[内置] MongoDB drop / dropDatabase",
target: "tool_input",
typ: "regex",
pattern: `(?i)\.(?:dropDatabase|dropCollection|drop)\s*\(`,
action: "deny",
message: "禁止执行 MongoDB drop 操作",
priority: 90,
},
{
name: "[内置] Redis FLUSHALL / FLUSHDB",
target: "tool_input",
typ: "regex",
pattern: `(?i)\b(?:FLUSHALL|FLUSHDB)\b`,
action: "deny",
message: "禁止执行 Redis FLUSHALL / FLUSHDB，可能清空全部缓存数据",
priority: 90,
},
// ── HTTP 破坏性请求 (priority 80) ──────────────────────────────────
// Agent 发送 DELETE 请求的三种常见方式：
// 1. curl -X DELETE / --request DELETE（Bash 工具直接执行或写入脚本）
// 2. Python HTTP 客户端 .delete() 方法
// 3. JS/通用脚本里的 method: 'DELETE' / method="DELETE"
{
name: "[内置] curl / wget 发送 DELETE 请求",
target: "tool_input",
typ: "regex",
pattern: `(?i)\bcurl\b[^|\n&;"]{0,300}(?:-X\s*DELETE|--request\s+DELETE|-XDELETE)|\bwget\b[^|\n&;"]{0,300}--method[=\s]+DELETE`,
action: "deny",
message: "禁止通过 curl/wget 发送 HTTP DELETE 请求，可能删除目标系统数据",
priority: 80,
},
{
name: "[内置] Python HTTP 客户端 DELETE（requests/httpx/aiohttp）",
target: "tool_input",
typ: "regex",
pattern: `(?i)\b(?:requests|httpx|aiohttp|urllib\.request)\.delete\s*\(|session\.delete\s*\(|client\.delete\s*\(`,
action: "deny",
message: "禁止使用 Python HTTP 客户端发送 DELETE 请求",
priority: 80,
},
{
name: "[内置] 脚本中声明 HTTP DELETE 方法（JS/通用）",
target: "tool_input",
typ: "regex",
pattern: `(?i)axios\.delete\s*\(|method\s*[:=]\s*['"]DELETE['"]`,
action: "deny",
message: "禁止在脚本中声明并发送 HTTP DELETE 请求",
priority: 80,
},
{
name: "[内置] 批量清空 / 清除接口路径",
target: "tool_input",
typ: "regex",
pattern: `(?i)/(?:clear|wipe|flush|purge|truncate|drop|destroy|factory[-_]reset|reset[-_]all)(?:[/?#"'\s]|$)`,
action: "deny",
message: "禁止调用批量清空或销毁类接口（/clear /wipe /flush /purge 等）",
priority: 80,
},
}
for _, r := range rules {
if _, err := d.Exec(`
INSERT INTO intercept_rules(name, enabled, priority, match_target, match_type, pattern, action, message, timeout_enabled, timeout_seconds, timeout_action)
VALUES ($1, true, $2, $3, $4, $5, $6, $7, false, 60, 'deny')
ON CONFLICT DO NOTHING`,
r.name, r.priority, r.target, r.typ, r.pattern, r.action, r.message,
); err != nil {
return fmt.Errorf("rule %q: %w", r.name, err)
}
}
return d.SetSetting("intercept_default_rules_v1", "done")
}
db/db_test.go
+67
View File
@@ -0,0 +1,67 @@
package db
import (
"testing"
)
// testDSN returns the configured DSN, skipping the test when neither the env var
// nor a config file supplies one (DSN no longer has a built-in default).
func testDSN(t *testing.T) string {
t.Helper()
dsn, _, err := DSN()
if err != nil {
t.Skipf("no database config (%v) — skipping", err)
}
return dsn
}
// TestOpenSeed opens the live dev PG, applies schema, seeds, and verifies the
// builtin agents + their variable catalog exist. Skips if PG is unreachable.
func TestOpenSeed(t *testing.T) {
d, err := Open(testDSN(t))
if err != nil {
t.Skipf("postgres unavailable (%v) — skipping", err)
}
defer d.Close()
wantAgents := len(builtinAgents)
var agents int
if err := d.QueryRow(`SELECT count(*) FROM agents WHERE builtin`).Scan(&agents); err != nil {
t.Fatal(err)
}
if agents != wantAgents {
t.Fatalf("want %d builtin agents, got %d", wantAgents, agents)
}
// planner should have at least its seeded catalog vars
wantPlannerVars := 0
for _, a := range builtinAgents {
if a.key == "planner" {
wantPlannerVars = len(a.vars)
break
}
}
var plannerVars int
if err := d.QueryRow(`
SELECT count(*) FROM agent_prompt_vars v
JOIN agents a ON a.id = v.agent_id
WHERE a.key = 'planner'`).Scan(&plannerVars); err != nil {
t.Fatal(err)
}
if plannerVars < wantPlannerVars {
t.Fatalf("want at least %d planner vars, got %d", wantPlannerVars, plannerVars)
}
// re-open must be idempotent (no duplicate agents)
d2, err := Open(testDSN(t))
if err != nil {
t.Fatal(err)
}
defer d2.Close()
if err := d2.QueryRow(`SELECT count(*) FROM agents WHERE builtin`).Scan(&agents); err != nil {
t.Fatal(err)
}
if agents != wantAgents {
t.Fatalf("after re-open want %d agents, got %d", wantAgents, agents)
}
}
db/exploration.go
+654
View File
@@ -0,0 +1,654 @@
package db
import (
"database/sql"
"encoding/json"
"fmt"
"strings"
"time"
)
// utf8Clean makes a string safe for a PostgreSQL text column. It (1) replaces
// invalid UTF-8 byte sequences with U+FFFD and (2) strips NUL (0x00) bytes.
// Tool output (nmap/curl/… stdout) can carry raw/truncated bytes that PostgreSQL's
// UTF8 encoding rejects ("invalid byte sequence for encoding UTF8"); without this
// the INSERT fails and the activity record is silently lost. NUL is *valid* UTF-8
// (U+0000) so ToValidUTF8 leaves it in place, yet PostgreSQL text still rejects it
// (SQLSTATE 22021) — so it must be removed separately. JSONB payloads are fine
// (json.Marshal already sanitizes), so only the plain text columns need it.
func utf8Clean(s string) string {
if strings.IndexByte(s, 0) >= 0 {
s = strings.ReplaceAll(s, "\x00", "")
}
return strings.ToValidUTF8(s, "�")
}
// Node is a typed reasoning node (= old task_nodes). kind ∈ goal|intent|finding|hint.
type Node struct {
ID int64 `json:"id"`
Kind string `json:"kind"`
Payload json.RawMessage `json:"payload"`
Priority int `json:"priority"`
State string `json:"state"`
Origin string `json:"origin,omitempty"`
Owner string `json:"owner,omitempty"`
Anchors []int64 `json:"anchors,omitempty"`
CreatedAt time.Time `json:"created_at"`
}
// Activity is one worker execution step (= old task_activity).
type Activity struct {
ID int64 `json:"id"`
NodeID *int64 `json:"node_id,omitempty"`
Worker string `json:"worker,omitempty"`
Kind string `json:"kind,omitempty"`
Tool string `json:"tool,omitempty"`
ToolUseID string `json:"tool_use_id,omitempty"`
IsError bool `json:"is_error"`
Summary string `json:"summary,omitempty"`
Detail string `json:"-"` // full blob; lazy-loaded, omitted from list payloads
CreatedAt time.Time `json:"created_at"`
// Token usage — set only on the terminal kind='result' record; nil elsewhere.
InputTokens *int `json:"input_tokens,omitempty"`
OutputTokens *int `json:"output_tokens,omitempty"`
CacheReadTokens *int `json:"cache_read_tokens,omitempty"`
CacheWriteTokens *int `json:"cache_write_tokens,omitempty"`
}
// TokenUsage is a per-worker token aggregate (TokenStatsByWorker).
type TokenUsage struct {
Worker string `json:"worker"`
InputTokens int `json:"input_tokens"`
OutputTokens int `json:"output_tokens"`
CacheReadTokens int `json:"cache_read_tokens"`
CacheWriteTokens int `json:"cache_write_tokens"`
}
// DailyTokenBucket is one day's global token aggregate across all tasks.
type DailyTokenBucket struct {
Day string `json:"day"` // "YYYY-MM-DD"
InputTokens int `json:"input_tokens"`
OutputTokens int `json:"output_tokens"`
CacheReadTokens int `json:"cache_read_tokens"`
}
// TokenDailyAll aggregates token consumption by calendar day (UTC) for the
// past `days` days across all explorations. Only kind='result' rows carry token
// counts (the terminal per-run summary), so there is no double-counting.
func (d *DB) TokenDailyAll(days int) ([]DailyTokenBucket, error) {
if days <= 0 {
days = 30
}
rows, err := d.Query(`
SELECT TO_CHAR(DATE_TRUNC('day', created_at AT TIME ZONE 'UTC'), 'YYYY-MM-DD') AS day,
COALESCE(SUM(input_tokens), 0),
COALESCE(SUM(output_tokens), 0),
COALESCE(SUM(cache_read_tokens), 0)
FROM activity
WHERE kind = 'result'
AND created_at >= NOW() - ($1 * INTERVAL '1 day')
GROUP BY day
ORDER BY day`, days)
if err != nil {
return nil, err
}
defer rows.Close()
var out []DailyTokenBucket
for rows.Next() {
var b DailyTokenBucket
if err := rows.Scan(&b.Day, &b.InputTokens, &b.OutputTokens, &b.CacheReadTokens); err != nil {
return nil, err
}
out = append(out, b)
}
return out, rows.Err()
}
// ExplorationStore is a handle onto one exploration's reasoning graph + activity.
type ExplorationStore struct {
db *DB
expID int64
}
// CreateExploration creates a new exploration root and returns its id.
func (d *DB) CreateExploration(description, goal string) (int64, error) {
var id int64
err := d.QueryRow(`INSERT INTO explorations(description, goal) VALUES ($1, $2) RETURNING id`, description, goal).Scan(&id)
return id, err
}
// Exploration returns a handle bound to an exploration id.
func (d *DB) Exploration(id int64) *ExplorationStore { return &ExplorationStore{db: d, expID: id} }
func (s *ExplorationStore) ID() int64 { return s.expID }
// Root returns the exploration's description and goal.
func (s *ExplorationStore) Root() (description, goal string, err error) {
var d sql.NullString
err = s.db.QueryRow(`SELECT description, goal FROM explorations WHERE id=$1`, s.expID).Scan(&d, &goal)
return d.String, goal, err
}
// AddNode writes a typed reasoning node with optional anchors to asset ids.
func (s *ExplorationStore) AddNode(kind string, payload map[string]any, priority int, state, origin string, anchors []int64) (int64, error) {
raw, _ := json.Marshal(payload)
tx, err := s.db.Begin()
if err != nil {
return 0, err
}
defer tx.Rollback()
var id int64
if err := tx.QueryRow(`
INSERT INTO exploration_nodes(exploration_id, kind, payload, priority, state, origin)
VALUES ($1, $2, $3, $4, $5, $6) RETURNING id`,
s.expID, kind, string(raw), priority, state, origin).Scan(&id); err != nil {
return 0, err
}
for _, a := range anchors {
if _, err := tx.Exec(`INSERT INTO exploration_anchors(node_id, asset_id) VALUES ($1,$2) ON CONFLICT DO NOTHING`, id, a); err != nil {
return 0, err
}
}
return id, tx.Commit()
}
// AddIntent is a convenience: an open intent.
func (s *ExplorationStore) AddIntent(payload map[string]any, priority int, anchors []int64, origin string) (int64, error) {
if origin == "" {
origin = "planner"
}
return s.AddNode("intent", payload, priority, "open", origin, anchors)
}
// AddGoal writes a goal node (state open).
func (s *ExplorationStore) AddGoal(payload map[string]any, origin string) (int64, error) {
return s.AddNode("goal", payload, 0, "open", origin, nil)
}
// OriginFactID returns this exploration's root fact id — the KindFact node with
// state='origin' seeded at task creation (the task root). Every intent traces
// back to it. Returns 0 (no error) when none exists (legacy explorations created
// under the old 'begin' root, or none yet).
func (s *ExplorationStore) OriginFactID() (int64, error) {
var id int64
err := s.db.QueryRow(`SELECT id FROM exploration_nodes WHERE exploration_id=$1 AND kind='fact' AND state='origin' ORDER BY id LIMIT 1`, s.expID).Scan(&id)
if err == sql.ErrNoRows {
return 0, nil
}
return id, err
}
// Anchor records a node→asset reference edge (many-to-many): it captures which
// global assets an exploration node touched, as lineage/provenance. The asset
// graph itself is global and shared — anchoring no longer gates visibility (any
// task may read any asset via search); it only preserves the reasoning trail.
// Idempotent.
func (s *ExplorationStore) Anchor(nodeID, assetID int64) error {
if nodeID <= 0 || assetID <= 0 {
return nil
}
_, err := s.db.Exec(`INSERT INTO exploration_anchors(node_id, asset_id) VALUES ($1,$2) ON CONFLICT DO NOTHING`, nodeID, assetID)
return err
}
// Link adds a typed exploration edge (idempotent).
func (s *ExplorationStore) Link(from int64, rel string, to int64) error {
_, err := s.db.Exec(`
INSERT INTO exploration_edges(exploration_id, src_id, rel, dst_id) VALUES ($1,$2,$3,$4)
ON CONFLICT (exploration_id, src_id, rel, dst_id) DO NOTHING`, s.expID, from, rel, to)
return err
}
// SetNodeState updates any node's state (never deletes).
func (s *ExplorationStore) SetNodeState(id int64, state string) error {
_, err := s.db.Exec(`UPDATE exploration_nodes SET state=$1 WHERE id=$2 AND exploration_id=$3`, state, id, s.expID)
return err
}
// SetIntentState updates an intent node's state (done/blocked/exhausted/open).
// ResetRunningIntents returns any intent left in 'running' back to 'open' so it
// is re-claimed. Called on startup: a 'running' intent with no live worker (a
// backend restart or crashed worker goroutine left it stuck) would otherwise spin
// forever in the UI. Workers re-run an intent from scratch, so reopening is safe.
func (s *ExplorationStore) ResetRunningIntents() (int64, error) {
res, err := s.db.Exec(`UPDATE exploration_nodes SET state='open', completed_at=NULL
WHERE exploration_id=$1 AND kind='intent' AND state='running'`, s.expID)
if err != nil {
return 0, err
}
n, _ := res.RowsAffected()
return n, nil
}
func (s *ExplorationStore) SetIntentState(id int64, state string) error {
// terminal states stamp completed_at; reopening (back to open/running) clears it.
terminal := state == "done" || state == "blocked" || state == "exhausted" || state == "stopped"
_, err := s.db.Exec(`UPDATE exploration_nodes
SET state=$1, completed_at = CASE WHEN $4 THEN now() ELSE NULL END
WHERE id=$2 AND exploration_id=$3 AND kind='intent'`, state, id, s.expID, terminal)
return err
}
const nodeCols = `id, kind, payload, priority, state, COALESCE(origin,''), COALESCE(owner,''), created_at`
func scanNode(sc interface{ Scan(...any) error }) (*Node, error) {
var n Node
var payload []byte
if err := sc.Scan(&n.ID, &n.Kind, &payload, &n.Priority, &n.State, &n.Origin, &n.Owner, &n.CreatedAt); err != nil {
return nil, err
}
n.Payload = json.RawMessage(payload)
return &n, nil
}
func scanNodes(rows *sql.Rows) ([]*Node, error) {
var out []*Node
for rows.Next() {
n, err := scanNode(rows)
if err != nil {
return nil, err
}
out = append(out, n)
}
return out, rows.Err()
}
// ListByKind lists nodes of a kind (newest first).
func (s *ExplorationStore) ListByKind(kind string, limit int) ([]*Node, error) {
if limit <= 0 {
limit = 500
}
rows, err := s.db.Query(`SELECT `+nodeCols+` FROM exploration_nodes
WHERE exploration_id=$1 AND kind=$2 ORDER BY id DESC LIMIT $3`, s.expID, kind, limit)
if err != nil {
return nil, err
}
defer rows.Close()
return scanNodes(rows)
}
// GetNode returns one node of this exploration by id (nil, nil if not found).
func (s *ExplorationStore) GetNode(id int64) (*Node, error) {
n, err := scanNode(s.db.QueryRow(`SELECT `+nodeCols+` FROM exploration_nodes WHERE id=$1 AND exploration_id=$2`, id, s.expID))
if err == sql.ErrNoRows {
return nil, nil
}
return n, err
}
// Nodes lists all nodes (for graph viz).
func (s *ExplorationStore) Nodes(limit int) ([]*Node, error) {
if limit <= 0 {
limit = 2000
}
rows, err := s.db.Query(`SELECT `+nodeCols+` FROM exploration_nodes WHERE exploration_id=$1 ORDER BY id LIMIT $2`, s.expID, limit)
if err != nil {
return nil, err
}
defer rows.Close()
return scanNodes(rows)
}
// Edge is one row from exploration_edges.
type Edge struct {
From int64
Rel string
To int64
}
// Edges lists exploration edges (for graph viz).
func (s *ExplorationStore) Edges(limit int) ([]Edge, error) {
if limit <= 0 {
limit = 5000
}
rows, err := s.db.Query(`SELECT src_id, rel, dst_id FROM exploration_edges WHERE exploration_id=$1 LIMIT $2`, s.expID, limit)
if err != nil {
return nil, err
}
defer rows.Close()
var out []Edge
for rows.Next() {
var e Edge
if err := rows.Scan(&e.From, &e.Rel, &e.To); err != nil {
return nil, err
}
out = append(out, e)
}
return out, rows.Err()
}
// FindingIntents maps each finding id to the intent that produced it (the
// intent --yields--> finding edge; report_finding links it). Findings with no
// producing intent are absent. Precise (JOIN, no edge-scan limit).
func (s *ExplorationStore) FindingIntents() (map[int64]int64, error) {
rows, err := s.db.Query(`SELECT e.dst_id, e.src_id
FROM exploration_edges e
JOIN exploration_nodes n ON n.id = e.dst_id AND n.exploration_id = e.exploration_id
WHERE e.exploration_id=$1 AND e.rel=$2 AND n.kind='finding'`, s.expID, RelYields)
if err != nil {
return nil, err
}
defer rows.Close()
out := map[int64]int64{}
for rows.Next() {
var findingID, intentID int64
if err := rows.Scan(&findingID, &intentID); err != nil {
return nil, err
}
out[findingID] = intentID
}
return out, rows.Err()
}
// Frontier returns open intents ordered by priority desc then id (FIFO tiebreak).
func (s *ExplorationStore) Frontier(limit int) ([]*Node, error) {
if limit <= 0 {
limit = 50
}
rows, err := s.db.Query(`SELECT `+nodeCols+` FROM exploration_nodes
WHERE exploration_id=$1 AND kind='intent' AND state='open'
ORDER BY priority DESC, id ASC LIMIT $2`, s.expID, limit)
if err != nil {
return nil, err
}
defer rows.Close()
return scanNodes(rows)
}
// ClaimIntent atomically moves an open intent to running. Returns true if claimed.
func (s *ExplorationStore) ClaimIntent(id int64, owner string) (bool, error) {
res, err := s.db.Exec(`UPDATE exploration_nodes SET state='running', owner=$1
WHERE id=$2 AND exploration_id=$3 AND kind='intent' AND state='open'`, owner, id, s.expID)
if err != nil {
return false, err
}
n, _ := res.RowsAffected()
return n == 1, nil
}
// Stats returns node counts grouped by kind (for dashboard).
func (s *ExplorationStore) Stats() (map[string]int, error) {
rows, err := s.db.Query(`SELECT kind, count(*) FROM exploration_nodes WHERE exploration_id=$1 GROUP BY kind`, s.expID)
if err != nil {
return nil, err
}
defer rows.Close()
out := map[string]int{}
for rows.Next() {
var k string
var c int
if err := rows.Scan(&k, &c); err != nil {
return nil, err
}
out[k] = c
}
return out, rows.Err()
}
// --- activity (poll by global id cursor) ---
// AppendActivity records one worker step and returns its id.
func (s *ExplorationStore) AppendActivity(a Activity) (int64, error) {
var id int64
err := s.db.QueryRow(`
INSERT INTO activity(exploration_id, node_id, worker, kind, tool, tool_use_id, is_error, summary, detail, input_tokens, output_tokens, cache_read_tokens, cache_write_tokens)
VALUES ($1,$2,NULLIF($3,''),NULLIF($4,''),NULLIF($5,''),NULLIF($6,''),$7,NULLIF($8,''),NULLIF($9,''),$10,$11,$12,$13)
RETURNING id`, s.expID, a.NodeID, utf8Clean(a.Worker), utf8Clean(a.Kind), utf8Clean(a.Tool), utf8Clean(a.ToolUseID), a.IsError, utf8Clean(a.Summary), utf8Clean(a.Detail),
a.InputTokens, a.OutputTokens, a.CacheReadTokens, a.CacheWriteTokens).Scan(&id)
return id, err
}
// TokenTotal sums token usage across ALL workers for this exploration (whole-task
// total). Same source as TokenStatsByWorker (kind='result' rows), just ungrouped.
func (s *ExplorationStore) TokenTotal() (TokenUsage, error) {
var u TokenUsage
err := s.db.QueryRow(`SELECT
COALESCE(SUM(input_tokens),0), COALESCE(SUM(output_tokens),0),
COALESCE(SUM(cache_read_tokens),0), COALESCE(SUM(cache_write_tokens),0)
FROM activity WHERE exploration_id=$1 AND kind='result'`, s.expID).
Scan(&u.InputTokens, &u.OutputTokens, &u.CacheReadTokens, &u.CacheWriteTokens)
return u, err
}
// TokenTotalsAll returns the whole-task token total for every exploration in one
// query (exploration_id → total), so the task list can show per-task consumption
// without a query per row.
func (d *DB) TokenTotalsAll() (map[int64]TokenUsage, error) {
rows, err := d.Query(`SELECT exploration_id,
COALESCE(SUM(input_tokens),0), COALESCE(SUM(output_tokens),0),
COALESCE(SUM(cache_read_tokens),0), COALESCE(SUM(cache_write_tokens),0)
FROM activity WHERE kind='result' GROUP BY exploration_id`)
if err != nil {
return nil, err
}
defer rows.Close()
out := map[int64]TokenUsage{}
for rows.Next() {
var eid int64
var u TokenUsage
if err := rows.Scan(&eid, &u.InputTokens, &u.OutputTokens, &u.CacheReadTokens, &u.CacheWriteTokens); err != nil {
return nil, err
}
out[eid] = u
}
return out, rows.Err()
}
// LastActivityAll returns the unix time of the most recent activity per
// exploration (exploration_id → max created_at epoch), one query for all tasks.
// Persisted (unlike Engine.LastActivity's in-memory map), so it survives restarts
// and gives终态任务 a stable "ran until" time for computing run duration.
func (d *DB) LastActivityAll() (map[int64]int64, error) {
rows, err := d.Query(`SELECT exploration_id, EXTRACT(EPOCH FROM MAX(created_at))::bigint FROM activity GROUP BY exploration_id`)
if err != nil {
return nil, err
}
defer rows.Close()
out := map[int64]int64{}
for rows.Next() {
var eid, ts int64
if err := rows.Scan(&eid, &ts); err != nil {
return nil, err
}
out[eid] = ts
}
return out, rows.Err()
}
// GoalCounts is the goal summary for one exploration.
type GoalCounts struct{ Total, Met int }
// GoalCountsAll returns goal totals (total / met) per exploration id, one query
// for all tasks — used by listTasks so the task list shows progress for every task.
func (d *DB) GoalCountsAll() (map[int64]GoalCounts, error) {
rows, err := d.Query(`
SELECT exploration_id,
COUNT(*) AS total,
COUNT(*) FILTER (WHERE state = 'met') AS met
FROM exploration_nodes
WHERE kind = 'goal'
GROUP BY exploration_id`)
if err != nil {
return nil, err
}
defer rows.Close()
out := map[int64]GoalCounts{}
for rows.Next() {
var eid int64
var gc GoalCounts
if err := rows.Scan(&eid, &gc.Total, &gc.Met); err != nil {
return nil, err
}
out[eid] = gc
}
return out, rows.Err()
}
// TokenStatsByWorker sums token usage per worker for this exploration (from the
// kind='result' records that carry usage). Used by the per-agent token display.
func (s *ExplorationStore) TokenStatsByWorker() ([]TokenUsage, error) {
rows, err := s.db.Query(`SELECT COALESCE(worker,''),
COALESCE(SUM(input_tokens),0), COALESCE(SUM(output_tokens),0),
COALESCE(SUM(cache_read_tokens),0), COALESCE(SUM(cache_write_tokens),0)
FROM activity WHERE exploration_id=$1 AND kind='result'
GROUP BY worker ORDER BY worker`, s.expID)
if err != nil {
return nil, err
}
defer rows.Close()
out := []TokenUsage{}
for rows.Next() {
var u TokenUsage
if err := rows.Scan(&u.Worker, &u.InputTokens, &u.OutputTokens, &u.CacheReadTokens, &u.CacheWriteTokens); err != nil {
return nil, err
}
out = append(out, u)
}
return out, rows.Err()
}
// ActivityList returns steps after sinceID (exclusive), optionally filtered by node.
// Returns items and the new cursor (max id seen).
func (s *ExplorationStore) ActivityList(nodeID *int64, sinceID int64, limit int) ([]Activity, int64, error) {
if limit <= 0 {
limit = 300
}
var rows *sql.Rows
var err error
const cols = `id, node_id, COALESCE(worker,''), COALESCE(kind,''), COALESCE(tool,''), COALESCE(tool_use_id,''), is_error, COALESCE(summary,''), created_at, input_tokens, output_tokens, cache_read_tokens, cache_write_tokens`
if nodeID != nil {
rows, err = s.db.Query(`SELECT `+cols+`
FROM activity WHERE exploration_id=$1 AND node_id=$2 AND id>$3 ORDER BY id LIMIT $4`, s.expID, *nodeID, sinceID, limit)
} else {
rows, err = s.db.Query(`SELECT `+cols+`
FROM activity WHERE exploration_id=$1 AND id>$2 ORDER BY id LIMIT $3`, s.expID, sinceID, limit)
}
if err != nil {
return nil, sinceID, err
}
defer rows.Close()
out := []Activity{}
cursor := sinceID
for rows.Next() {
var a Activity
if err := rows.Scan(&a.ID, &a.NodeID, &a.Worker, &a.Kind, &a.Tool, &a.ToolUseID, &a.IsError, &a.Summary, &a.CreatedAt,
&a.InputTokens, &a.OutputTokens, &a.CacheReadTokens, &a.CacheWriteTokens); err != nil {
return nil, sinceID, err
}
if a.ID > cursor {
cursor = a.ID
}
out = append(out, a)
}
return out, cursor, rows.Err()
}
// ActivityDetail lazily returns the full detail blob for one step.
func (s *ExplorationStore) ActivityDetail(id int64) (string, error) {
var d sql.NullString
err := s.db.QueryRow(`SELECT detail FROM activity WHERE id=$1 AND exploration_id=$2`, id, s.expID).Scan(&d)
if err == sql.ErrNoRows {
return "", nil
}
return d.String, err
}
// scanTrace reads the summary-only column set shared by the worker-trace tools
// (ActivityTrace / ActivityTraceSearch). Detail is never selected here — it is
// pulled separately, on demand, via ActivityByIDs.
func scanTrace(rows *sql.Rows) ([]Activity, error) {
out := []Activity{}
for rows.Next() {
var a Activity
if err := rows.Scan(&a.ID, &a.NodeID, &a.Worker, &a.Kind, &a.Tool, &a.IsError, &a.Summary); err != nil {
return nil, err
}
out = append(out, a)
}
return out, rows.Err()
}
const traceCols = `id, node_id, COALESCE(worker,''), COALESCE(kind,''), COALESCE(tool,''), is_error, COALESCE(summary,'')`
// ActivityTrace lists one work's steps (by intent/node id) for the trace tools:
// summaries only (detail is lazy — fetch via ActivityByIDs). thinking/usage rows
// are dropped so the caller sees the work's actions + observations, not the
// model's internal reasoning or token-accounting noise.
func (s *ExplorationStore) ActivityTrace(nodeID int64, limit int) ([]Activity, error) {
if limit <= 0 {
limit = 500
}
rows, err := s.db.Query(`SELECT `+traceCols+`
FROM activity WHERE exploration_id=$1 AND node_id=$2 AND kind NOT IN ('thinking','usage')
ORDER BY id LIMIT $3`, s.expID, nodeID, limit)
if err != nil {
return nil, err
}
defer rows.Close()
return scanTrace(rows)
}
// ActivityTraceSearch finds steps whose summary OR detail matches q
// (case-insensitive), returning summaries only. nodeID != nil scopes to one work;
// nil searches across ALL works (node_id IS NOT NULL — planner/main steps have no
// node_id, so this cleanly means "worker traces"). thinking/usage rows dropped.
func (s *ExplorationStore) ActivityTraceSearch(nodeID *int64, q string, limit int) ([]Activity, error) {
if limit <= 0 {
limit = 100
}
like := "%" + q + "%"
var rows *sql.Rows
var err error
if nodeID != nil {
rows, err = s.db.Query(`SELECT `+traceCols+`
FROM activity WHERE exploration_id=$1 AND node_id=$2 AND kind NOT IN ('thinking','usage')
AND (summary ILIKE $3 OR detail ILIKE $3) ORDER BY id LIMIT $4`, s.expID, *nodeID, like, limit)
} else {
rows, err = s.db.Query(`SELECT `+traceCols+`
FROM activity WHERE exploration_id=$1 AND node_id IS NOT NULL AND kind NOT IN ('thinking','usage')
AND (summary ILIKE $2 OR detail ILIKE $2) ORDER BY id LIMIT $3`, s.expID, like, limit)
}
if err != nil {
return nil, err
}
defer rows.Close()
return scanTrace(rows)
}
// ActivityByIDs loads full detail for specific step ids (the trace drill-down),
// scoped to this exploration. thinking rows are skipped (never returned, even if
// their id is asked for). Order is ascending id, not the input order.
func (s *ExplorationStore) ActivityByIDs(ids []int64) ([]Activity, error) {
if len(ids) == 0 {
return nil, nil
}
ph := make([]string, len(ids))
args := make([]any, len(ids)+1)
args[0] = s.expID
for i, id := range ids {
ph[i] = fmt.Sprintf("$%d", i+2)
args[i+1] = id
}
rows, err := s.db.Query(`SELECT id, node_id, COALESCE(kind,''), COALESCE(tool,''), is_error, COALESCE(detail,'')
FROM activity WHERE exploration_id=$1 AND kind<>'thinking' AND id IN (`+strings.Join(ph, ",")+`) ORDER BY id`, args...)
if err != nil {
return nil, err
}
defer rows.Close()
out := []Activity{}
for rows.Next() {
var a Activity
if err := rows.Scan(&a.ID, &a.NodeID, &a.Kind, &a.Tool, &a.IsError, &a.Detail); err != nil {
return nil, err
}
out = append(out, a)
}
return out, rows.Err()
}
// AddStandaloneFinding writes a finding to the standalone findings table, which
// persists across task deletion (task_id / node_id become NULL when the task or
// exploration node is deleted). taskID and nodeID may be 0 (stored as NULL).
func (s *ExplorationStore) AddStandaloneFinding(taskID, nodeID int64, vulnclass, severity, summary, evidence, worker string, assetIDs []int64) (int64, error) {
return s.db.AddFinding(taskID, nodeID, vulnclass, severity, summary, evidence, worker, assetIDs)
}
db/exploration_test.go
+93
View File
@@ -0,0 +1,93 @@
package db
import "testing"
func TestExplorationFlow(t *testing.T) {
d, err := Open(testDSN(t))
if err != nil {
t.Skipf("postgres unavailable (%v) — skipping", err)
}
defer d.Close()
expID, err := d.CreateExploration("test", "拿下测试目标")
if err != nil {
t.Fatal(err)
}
defer d.Exec(`DELETE FROM explorations WHERE id=$1`, expID) // cascades nodes/edges/activity
es := d.Exploration(expID)
// goal node + two intents
goal, err := es.AddGoal(map[string]any{"text": "getadmin", "vulnclass": "authz"}, "human")
if err != nil {
t.Fatal(err)
}
if _, err := es.AddIntent(map[string]any{"summary": "enumerate endpoints"}, 5, nil, "planner"); err != nil {
t.Fatal(err)
}
i2, err := es.AddIntent(map[string]any{"summary": "test idor"}, 8, nil, "planner")
if err != nil {
t.Fatal(err)
}
// frontier ordered by priority desc → i2(8) before i1(5)
fr, err := es.Frontier(10)
if err != nil {
t.Fatal(err)
}
if len(fr) != 2 || fr[0].ID != i2 {
t.Fatalf("frontier order wrong: %+v", fr)
}
// atomic claim: first wins, second on same id fails
ok, err := es.ClaimIntent(i2, "worker-1")
if err != nil || !ok {
t.Fatalf("claim i2: ok=%v err=%v", ok, err)
}
ok2, _ := es.ClaimIntent(i2, "worker-2")
if ok2 {
t.Fatalf("double-claim should fail")
}
// finding yields from intent, proves goal
find, err := es.AddNode("finding", map[string]any{"vulnclass": "idor", "severity": "high"}, 9, "confirmed", "worker-1", nil)
if err != nil {
t.Fatal(err)
}
if err := es.Link(i2, "yields", find); err != nil {
t.Fatal(err)
}
if err := es.Link(find, "proves", goal); err != nil {
t.Fatal(err)
}
if err := es.SetNodeState(goal, "met"); err != nil {
t.Fatal(err)
}
// activity poll by id cursor
id1, err := es.AppendActivity(Activity{Worker: "worker-1", Kind: "tool_use", Tool: "Bash", Summary: "ran curl", Detail: "full output"})
if err != nil {
t.Fatal(err)
}
items, cursor, err := es.ActivityList(nil, 0, 100)
if err != nil {
t.Fatal(err)
}
if len(items) != 1 || cursor != id1 {
t.Fatalf("activity list: items=%d cursor=%d", len(items), cursor)
}
det, _ := es.ActivityDetail(id1)
if det != "full output" {
t.Fatalf("detail want 'full output', got %q", det)
}
// incremental: nothing new after cursor
items2, _, _ := es.ActivityList(nil, cursor, 100)
if len(items2) != 0 {
t.Fatalf("incremental poll should be empty, got %d", len(items2))
}
// stats
st, _ := es.Stats()
if st["intent"] != 2 || st["goal"] != 1 || st["finding"] != 1 {
t.Fatalf("stats: %+v", st)
}
}
db/findings.go
+75
View File
@@ -0,0 +1,75 @@
package db
import (
"encoding/json"
"time"
)
// DBFinding is a row in the standalone findings table (persists across task deletion).
type DBFinding struct {
ID int64
TaskID *int64
NodeID *int64
VulnClass string
Severity string
Summary string
Evidence string
Worker string
AssetIDs []int64
CreatedAt time.Time
TaskDescription string // populated via LEFT JOIN on tasks
}
// AddFinding inserts a finding into the standalone findings table. taskID and
// nodeID may be 0 (stored as NULL). Returns the new finding id.
func (d *DB) AddFinding(taskID, nodeID int64, vulnclass, severity, summary, evidence, worker string, assetIDs []int64) (int64, error) {
aidsJSON, _ := json.Marshal(assetIDs)
if assetIDs == nil {
aidsJSON = []byte("[]")
}
var tid, nid *int64
if taskID > 0 {
tid = &taskID
}
if nodeID > 0 {
nid = &nodeID
}
var id int64
err := d.QueryRow(
`INSERT INTO findings (task_id, node_id, vulnclass, severity, summary, evidence, worker, asset_ids)
VALUES ($1, $2, $3, $4, $5, $6, $7, $8) RETURNING id`,
tid, nid, vulnclass, severity, summary, evidence, worker, string(aidsJSON),
).Scan(&id)
return id, err
}
// ListFindings returns all findings (newest first), joined with task description.
func (d *DB) ListFindings(limit int) ([]*DBFinding, error) {
if limit <= 0 {
limit = 500
}
rows, err := d.Query(`
SELECT f.id, f.task_id, f.node_id, f.vulnclass, f.severity, f.summary,
f.evidence, f.worker, f.asset_ids, f.created_at,
COALESCE(t.description, '') AS task_description
FROM findings f
LEFT JOIN tasks t ON f.task_id = t.id
ORDER BY f.created_at DESC
LIMIT $1`, limit)
if err != nil {
return nil, err
}
defer rows.Close()
var out []*DBFinding
for rows.Next() {
f := &DBFinding{}
var aidsJSON string
if err := rows.Scan(&f.ID, &f.TaskID, &f.NodeID, &f.VulnClass, &f.Severity,
&f.Summary, &f.Evidence, &f.Worker, &aidsJSON, &f.CreatedAt, &f.TaskDescription); err != nil {
return nil, err
}
_ = json.Unmarshal([]byte(aidsJSON), &f.AssetIDs)
out = append(out, f)
}
return out, rows.Err()
}
db/intercept.go
+234
View File
@@ -0,0 +1,234 @@
package db
import (
"database/sql"
"encoding/json"
"time"
)
// InterceptRule is one row of intercept_rules.
type InterceptRule struct {
ID int64 `json:"id"`
Name string `json:"name"`
Enabled bool `json:"enabled"`
Priority int `json:"priority"`
MatchTarget string `json:"match_target"`
MatchType string `json:"match_type"`
Pattern string `json:"pattern"`
Action string `json:"action"`
Message string `json:"message"`
TimeoutEnabled bool `json:"timeout_enabled"`
TimeoutSeconds int `json:"timeout_seconds"`
TimeoutAction string `json:"timeout_action"`
CreatedAt time.Time `json:"created_at"`
UpdatedAt time.Time `json:"updated_at"`
}
// InterceptPending is one row of intercept_pending.
type InterceptPending struct {
ID int64 `json:"id"`
RuleID *int64 `json:"rule_id"`
ConversationID *int64 `json:"conversation_id"`
TaskID *string `json:"task_id"`
AgentName string `json:"agent_name"`
ToolName string `json:"tool_name"`
ToolInput json.RawMessage `json:"tool_input"`
Status string `json:"status"`
DecidedAt *time.Time `json:"decided_at"`
CreatedAt time.Time `json:"created_at"`
}
const interceptRuleCols = `id, name, enabled, priority, match_target, match_type, pattern, action, message, timeout_enabled, timeout_seconds, timeout_action, created_at, updated_at`
func scanInterceptRule(row interface{ Scan(...any) error }) (InterceptRule, error) {
var r InterceptRule
err := row.Scan(&r.ID, &r.Name, &r.Enabled, &r.Priority,
&r.MatchTarget, &r.MatchType, &r.Pattern, &r.Action, &r.Message,
&r.TimeoutEnabled, &r.TimeoutSeconds, &r.TimeoutAction,
&r.CreatedAt, &r.UpdatedAt)
return r, err
}
// ListInterceptRules returns all rules ordered by priority DESC then id.
func (d *DB) ListInterceptRules() ([]InterceptRule, error) {
rows, err := d.Query(`SELECT ` + interceptRuleCols + ` FROM intercept_rules ORDER BY priority DESC, id`)
if err != nil {
return nil, err
}
defer rows.Close()
var out []InterceptRule
for rows.Next() {
r, err := scanInterceptRule(rows)
if err != nil {
return nil, err
}
out = append(out, r)
}
return out, rows.Err()
}
// CreateInterceptRule inserts a new rule.
func (d *DB) CreateInterceptRule(name, matchTarget, matchType, pattern, action, message string, priority int, enabled bool, timeoutEnabled bool, timeoutSeconds int, timeoutAction string) (InterceptRule, error) {
row := d.QueryRow(`
INSERT INTO intercept_rules(name, enabled, priority, match_target, match_type, pattern, action, message, timeout_enabled, timeout_seconds, timeout_action)
VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11)
RETURNING `+interceptRuleCols,
name, enabled, priority, matchTarget, matchType, pattern, action, message, timeoutEnabled, timeoutSeconds, timeoutAction)
return scanInterceptRule(row)
}
// UpdateInterceptRule replaces all editable fields of an existing rule.
func (d *DB) UpdateInterceptRule(id int64, name, matchTarget, matchType, pattern, action, message string, priority int, enabled bool, timeoutEnabled bool, timeoutSeconds int, timeoutAction string) (InterceptRule, error) {
row := d.QueryRow(`
UPDATE intercept_rules
SET name=$2, enabled=$3, priority=$4, match_target=$5,
match_type=$6, pattern=$7, action=$8, message=$9,
timeout_enabled=$10, timeout_seconds=$11, timeout_action=$12
WHERE id=$1
RETURNING `+interceptRuleCols,
id, name, enabled, priority, matchTarget, matchType, pattern, action, message, timeoutEnabled, timeoutSeconds, timeoutAction)
return scanInterceptRule(row)
}
// DeleteInterceptRule removes a rule.
func (d *DB) DeleteInterceptRule(id int64) error {
_, err := d.Exec(`DELETE FROM intercept_rules WHERE id=$1`, id)
return err
}
// ToggleInterceptRule flips the enabled state of a rule.
func (d *DB) ToggleInterceptRule(id int64, enabled bool) error {
_, err := d.Exec(`UPDATE intercept_rules SET enabled=$2 WHERE id=$1`, id, enabled)
return err
}
// CreateInterceptPending inserts a pending approval record and returns its ID.
// convID == 0 → conversation_id stored as NULL (background task).
// taskID == "" → task_id stored as NULL.
func (d *DB) CreateInterceptPending(ruleID, convID int64, taskID, agentName, toolName string, input []byte) (int64, error) {
raw := json.RawMessage(input)
if len(raw) == 0 {
raw = json.RawMessage("{}")
}
var convIDPtr *int64
if convID != 0 {
convIDPtr = &convID
}
var taskIDPtr *string
if taskID != "" {
taskIDPtr = &taskID
}
var id int64
err := d.QueryRow(`
INSERT INTO intercept_pending(rule_id, conversation_id, task_id, agent_name, tool_name, tool_input)
VALUES ($1, $2, $3, $4, $5, $6) RETURNING id`,
ruleID, convIDPtr, taskIDPtr, agentName, toolName, raw).Scan(&id)
return id, err
}
// DecideInterceptPending updates a pending record's status (allowed/denied/timeout).
func (d *DB) DecideInterceptPending(id int64, status string) error {
_, err := d.Exec(`UPDATE intercept_pending SET status=$2, decided_at=NOW() WHERE id=$1`, id, status)
return err
}
const interceptPendingCols = `id, rule_id, conversation_id, task_id, agent_name, tool_name, tool_input, status, decided_at, created_at`
func scanInterceptPending(s interface{ Scan(...any) error }, p *InterceptPending) error {
return s.Scan(&p.ID, &p.RuleID, &p.ConversationID, &p.TaskID, &p.AgentName,
&p.ToolName, &p.ToolInput, &p.Status, &p.DecidedAt, &p.CreatedAt)
}
// ListPendingIntercepts returns all unresolved approval requests, newest first.
func (d *DB) ListPendingIntercepts() ([]InterceptPending, error) {
rows, err := d.Query(`SELECT ` + interceptPendingCols + ` FROM intercept_pending WHERE status='pending' ORDER BY created_at DESC`)
if err != nil {
return nil, err
}
defer rows.Close()
var out []InterceptPending
for rows.Next() {
var p InterceptPending
if err := scanInterceptPending(rows, &p); err != nil {
return nil, err
}
out = append(out, p)
}
return out, rows.Err()
}
// GetInterceptPending returns one pending record (nil if absent).
func (d *DB) GetInterceptPending(id int64) (*InterceptPending, error) {
var p InterceptPending
err := scanInterceptPending(
d.QueryRow(`SELECT `+interceptPendingCols+` FROM intercept_pending WHERE id=$1`, id),
&p,
)
if err == sql.ErrNoRows {
return nil, nil
}
return &p, err
}
// InterceptApprovalRow is intercept_pending enriched with conversation and rule info.
type InterceptApprovalRow struct {
InterceptPending
ConvTitle string `json:"conv_title"`
ConvAgentKey string `json:"conv_agent_key"`
RuleName string `json:"rule_name"`
}
func scanInterceptApprovalRow(rows interface{ Scan(...any) error }, r *InterceptApprovalRow) error {
return rows.Scan(
&r.ID, &r.RuleID, &r.ConversationID, &r.TaskID, &r.AgentName,
&r.ToolName, &r.ToolInput, &r.Status, &r.DecidedAt, &r.CreatedAt,
&r.ConvTitle, &r.ConvAgentKey, &r.RuleName,
)
}
const approvalRowSelect = `
SELECT ip.id, ip.rule_id, ip.conversation_id, ip.task_id, ip.agent_name,
ip.tool_name, ip.tool_input, ip.status, ip.decided_at, ip.created_at,
COALESCE(c.title,'') AS conv_title,
COALESCE(c.agent_key,'') AS conv_agent_key,
COALESCE(ir.name,'') AS rule_name
FROM intercept_pending ip
LEFT JOIN conversations c ON c.id = ip.conversation_id
LEFT JOIN intercept_rules ir ON ir.id = ip.rule_id`
// ListAllIntercepts returns up to limit intercept_pending rows (newest first)
// joined with conversation and rule info.
func (d *DB) ListAllIntercepts(limit int) ([]InterceptApprovalRow, error) {
rows, err := d.Query(approvalRowSelect+` ORDER BY ip.created_at DESC LIMIT $1`, limit)
if err != nil {
return nil, err
}
defer rows.Close()
var out []InterceptApprovalRow
for rows.Next() {
var r InterceptApprovalRow
if err := scanInterceptApprovalRow(rows, &r); err != nil {
return nil, err
}
out = append(out, r)
}
return out, rows.Err()
}
// ListTaskIntercepts returns all intercept_pending rows for a specific task (newest first).
func (d *DB) ListTaskIntercepts(taskID string) ([]InterceptApprovalRow, error) {
rows, err := d.Query(approvalRowSelect+` WHERE ip.task_id=$1 ORDER BY ip.created_at DESC`, taskID)
if err != nil {
return nil, err
}
defer rows.Close()
var out []InterceptApprovalRow
for rows.Next() {
var r InterceptApprovalRow
if err := scanInterceptApprovalRow(rows, &r); err != nil {
return nil, err
}
out = append(out, r)
}
return out, rows.Err()
}
db/logs.go
+78
View File
@@ -0,0 +1,78 @@
package db
import (
"context"
"time"
)
// DBLog is one persisted backend log row from the server_logs table.
type DBLog struct {
ID int64
CreatedAt time.Time
Level string
Tag string
Text string
}
// InsertLog appends one log line and returns its auto-assigned id.
func (d *DB) InsertLog(level, tag, text string) (int64, error) {
var id int64
err := d.QueryRowContext(context.Background(),
"INSERT INTO server_logs(level,tag,text) VALUES($1,$2,$3) RETURNING id",
level, tag, text,
).Scan(&id)
return id, err
}
// RecentLogs returns the most recent limit rows, oldest-first.
func (d *DB) RecentLogs(limit int) ([]*DBLog, error) {
rows, err := d.QueryContext(context.Background(),
`SELECT id, created_at, level, tag, text
FROM server_logs
ORDER BY id DESC
LIMIT $1`, limit)
if err != nil {
return nil, err
}
defer rows.Close()
var out []*DBLog
for rows.Next() {
l := &DBLog{}
if err := rows.Scan(&l.ID, &l.CreatedAt, &l.Level, &l.Tag, &l.Text); err != nil {
return nil, err
}
out = append(out, l)
}
// Reverse to oldest-first
for i, j := 0, len(out)-1; i < j; i, j = i+1, j-1 {
out[i], out[j] = out[j], out[i]
}
return out, nil
}
// ListLogsBefore returns up to limit rows with id < beforeID, oldest-first.
func (d *DB) ListLogsBefore(beforeID int64, limit int) ([]*DBLog, error) {
rows, err := d.QueryContext(context.Background(),
`SELECT id, created_at, level, tag, text
FROM server_logs
WHERE id < $1
ORDER BY id DESC
LIMIT $2`, beforeID, limit)
if err != nil {
return nil, err
}
defer rows.Close()
var out []*DBLog
for rows.Next() {
l := &DBLog{}
if err := rows.Scan(&l.ID, &l.CreatedAt, &l.Level, &l.Tag, &l.Text); err != nil {
return nil, err
}
out = append(out, l)
}
// Reverse to oldest-first
for i, j := 0, len(out)-1; i < j; i, j = i+1, j-1 {
out[i], out[j] = out[j], out[i]
}
return out, nil
}
db/nkey.go
+134
View File
@@ -0,0 +1,134 @@
package db
import (
"net"
"net/url"
"regexp"
"strconv"
"strings"
"golang.org/x/net/publicsuffix"
)
// 归一化自然键 (nkey)：移植自旧 graph/id.go，去掉 StableID 哈希（PG 用 BIGSERIAL 主键 +
// UNIQUE(type, nkey) 去重）。子资产的 nkey 内嵌父资产的 int64 id，把层级编码进键。
func DomainKey(fqdn string) string {
return strings.TrimSuffix(strings.ToLower(strings.TrimSpace(fqdn)), ".")
}
func IPKey(ip string) string { return strings.TrimSpace(ip) }
// RootDomain returns the registrable domain (eTLD+1) for a host and whether the
// host itself IS that apex (§3.1). Edge cases (§3.1 边界处理): an IP literal or a
// host publicsuffix can't classify (localhost / internal / non-ICANN TLD) is
// returned unchanged as its own root with isApex=true — best-effort, never treated
// as a subdomain.
func RootDomain(host string) (root string, isApex bool) {
h := DomainKey(host)
if h == "" || net.ParseIP(h) != nil {
return h, true
}
etld1, err := publicsuffix.EffectiveTLDPlusOne(h)
if err != nil || etld1 == "" {
return h, true
}
return etld1, h == etld1
}
func PortKey(ipID int64, proto string, port int) string {
return itoa(ipID) + "|" + strings.ToLower(proto) + "|" + strconv.Itoa(port)
}
func ServiceKey(portID int64, svcName string) string {
return itoa(portID) + "|" + strings.ToLower(svcName)
}
func SiteKey(scheme, host string, port int) string {
return strings.ToLower(scheme) + "|" + strings.ToLower(host) + "|" + strconv.Itoa(port)
}
func EndpointKey(siteID int64, method, urlTemplate string) string {
return itoa(siteID) + "|" + strings.ToUpper(method) + "|" + urlTemplate
}
func ParameterKey(endpointID int64, location, name string) string {
return itoa(endpointID) + "|" + strings.ToLower(location) + "|" + name
}
// NormalizeParamName 归一化参数名（endpoint.params 元素的「相同引用」判定）。
// 规则：lower + trim，不做同义词合并(userId/user_id/uid 视为不同)。写入与查询共享此实现，
// 保证「按参数名查同公司接口」可复现。
func NormalizeParamName(name string) string {
return strings.ToLower(strings.TrimSpace(name))
}
func TechKey(name, version string) string {
return strings.ToLower(name) + "|" + version
}
func itoa(n int64) string { return strconv.FormatInt(n, 10) }
var (
reNumeric = regexp.MustCompile(`^\d+$`)
reUUID = regexp.MustCompile(`^[0-9a-fA-F]{8}-[0-9a-fA-F]{4}-[0-9a-fA-F]{4}-[0-9a-fA-F]{4}-[0-9a-fA-F]{12}$`)
reHex = regexp.MustCompile(`^[0-9a-fA-F]{16,}$`)
reLong = regexp.MustCompile(`^[A-Za-z0-9_-]{24,}$`)
)
// TemplatePath collapses high-cardinality path segments into placeholders so the
// graph is not flooded by instances: /user/123 -> /user/{id}.
func TemplatePath(path string) string {
if path == "" {
return "/"
}
segs := strings.Split(path, "/")
for i, s := range segs {
switch {
case s == "":
continue
case reNumeric.MatchString(s):
segs[i] = "{id}"
case reUUID.MatchString(s):
segs[i] = "{uuid}"
case reHex.MatchString(s):
segs[i] = "{hex}"
case reLong.MatchString(s):
segs[i] = "{token}"
}
}
return strings.Join(segs, "/")
}
// SplitURL parses a raw URL into scheme/host/port/urlTemplate/params for building
// site/endpoint/param keys.
func SplitURL(raw, method string) (scheme, host string, port int, urlTemplate string, params []string, err error) {
u, err := url.Parse(raw)
if err != nil {
return "", "", 0, "", nil, err
}
scheme = strings.ToLower(u.Scheme)
host = strings.ToLower(u.Hostname())
port = defaultPort(scheme, u.Port())
urlTemplate = TemplatePath(u.EscapedPath())
for k := range u.Query() {
params = append(params, k)
}
return scheme, host, port, urlTemplate, params, nil
}
func defaultPort(scheme, p string) int {
if p != "" {
if n, err := strconv.Atoi(p); err == nil {
return n
}
}
switch scheme {
case "https":
return 443
case "http":
return 80
default:
return 0
}
}
db/schema.sql
+513
View File
@@ -0,0 +1,513 @@
-- ARTEX PostgreSQL schema (单一数据源)
-- 幂等：可重复执行（IF NOT EXISTS / OR REPLACE / DROP TRIGGER IF EXISTS）。
-- =====================================================================
-- 0. 通用：updated_at 触发器
-- =====================================================================
CREATE OR REPLACE FUNCTION set_updated_at() RETURNS trigger AS $$
BEGIN NEW.updated_at = now(); RETURN NEW; END;
$$ LANGUAGE plpgsql;
-- =====================================================================
-- A. 资产层：companies / assets / company_scope
-- =====================================================================
CREATE TABLE IF NOT EXISTS companies (
id BIGSERIAL PRIMARY KEY,
name TEXT NOT NULL,
nkey TEXT NOT NULL UNIQUE,
logo TEXT,
created_at TIMESTAMPTZ NOT NULL DEFAULT now(),
updated_at TIMESTAMPTZ NOT NULL DEFAULT now()
);
CREATE INDEX IF NOT EXISTS idx_companies_nkey ON companies(nkey);
DROP TRIGGER IF EXISTS trg_companies_upd ON companies;
CREATE TRIGGER trg_companies_upd BEFORE UPDATE ON companies
FOR EACH ROW EXECUTE FUNCTION set_updated_at();
CREATE TABLE IF NOT EXISTS assets (
id BIGSERIAL PRIMARY KEY,
type TEXT NOT NULL CHECK (type IN (
'root_domain','ip','subdomain','app','service','endpoint'
)),
company_id BIGINT REFERENCES companies(id) ON DELETE SET NULL,
task_ids BIGINT[] NOT NULL DEFAULT '{}',
domain TEXT,
root_domain TEXT,
ip TEXT,
c_segment CIDR,
port INTEGER CHECK (port BETWEEN 1 AND 65535),
icp TEXT,
bound_domains TEXT[] NOT NULL DEFAULT '{}',
open_ports JSONB[] NOT NULL DEFAULT '{}',
record_type TEXT,
record_value TEXT[],
bundle_id TEXT,
app_name TEXT,
category TEXT,
app_description TEXT,
app_icp TEXT,
url TEXT,
service_type TEXT CHECK (service_type IN ('http','other')),
service_name TEXT,
favicon_mmh3 TEXT,
status_code INTEGER,
content_length BIGINT,
page_title TEXT,
technologies TEXT[] NOT NULL DEFAULT '{}',
auth JSONB[] NOT NULL DEFAULT '{}',
method TEXT,
params JSONB[] NOT NULL DEFAULT '{}',
extra JSONB NOT NULL DEFAULT '{}',
created_at TIMESTAMPTZ NOT NULL DEFAULT now(),
last_seen TIMESTAMPTZ NOT NULL DEFAULT now(),
updated_at TIMESTAMPTZ NOT NULL DEFAULT now()
);
CREATE UNIQUE INDEX IF NOT EXISTS uq_av2_root_domain ON assets(domain) WHERE type = 'root_domain';
CREATE UNIQUE INDEX IF NOT EXISTS uq_av2_ip ON assets(ip) WHERE type = 'ip';
CREATE UNIQUE INDEX IF NOT EXISTS uq_av2_subdomain ON assets(domain, COALESCE(record_type,'')) WHERE type = 'subdomain';
CREATE UNIQUE INDEX IF NOT EXISTS uq_av2_app_bundle ON assets(bundle_id) WHERE type = 'app' AND bundle_id IS NOT NULL;
CREATE UNIQUE INDEX IF NOT EXISTS uq_av2_app_name ON assets(app_name) WHERE type = 'app' AND bundle_id IS NULL;
CREATE UNIQUE INDEX IF NOT EXISTS uq_av2_service_http ON assets(url) WHERE type = 'service' AND service_type = 'http';
CREATE UNIQUE INDEX IF NOT EXISTS uq_av2_service_other
ON assets(COALESCE(domain,''), COALESCE(ip,''), port, service_name) WHERE type = 'service' AND service_type = 'other';
CREATE UNIQUE INDEX IF NOT EXISTS uq_av2_endpoint ON assets(url, method) WHERE type = 'endpoint';
CREATE INDEX IF NOT EXISTS idx_av2_company ON assets(company_id) WHERE company_id IS NOT NULL;
CREATE INDEX IF NOT EXISTS idx_av2_company_type ON assets(company_id, type) WHERE company_id IS NOT NULL;
CREATE INDEX IF NOT EXISTS idx_av2_task_ids ON assets USING GIN(task_ids);
CREATE INDEX IF NOT EXISTS idx_av2_domain ON assets(domain) WHERE domain IS NOT NULL;
CREATE INDEX IF NOT EXISTS idx_av2_root_domain ON assets(root_domain) WHERE root_domain IS NOT NULL;
CREATE INDEX IF NOT EXISTS idx_av2_ip ON assets(ip) WHERE ip IS NOT NULL;
CREATE INDEX IF NOT EXISTS idx_av2_c_segment ON assets USING GIST(c_segment inet_ops) WHERE c_segment IS NOT NULL;
CREATE INDEX IF NOT EXISTS idx_av2_technologies ON assets USING GIN(technologies) WHERE type = 'service';
CREATE INDEX IF NOT EXISTS idx_av2_bound_domains ON assets USING GIN(bound_domains) WHERE type = 'ip';
CREATE INDEX IF NOT EXISTS idx_av2_open_ports ON assets USING GIN(open_ports) WHERE type = 'ip';
CREATE INDEX IF NOT EXISTS idx_av2_last_seen ON assets(last_seen DESC);
CREATE INDEX IF NOT EXISTS idx_av2_type_seen ON assets(type, last_seen DESC);
DROP TRIGGER IF EXISTS trg_av2_upd ON assets;
CREATE TRIGGER trg_av2_upd BEFORE UPDATE ON assets
FOR EACH ROW EXECUTE FUNCTION set_updated_at();
CREATE TABLE IF NOT EXISTS company_scope (
id BIGSERIAL PRIMARY KEY,
company_id BIGINT NOT NULL REFERENCES companies(id) ON DELETE CASCADE,
kind TEXT NOT NULL CHECK (kind IN ('domain','ip','cidr')),
domain TEXT,
net CIDR,
raw TEXT NOT NULL,
reason TEXT,
created_at TIMESTAMPTZ NOT NULL DEFAULT now(),
CONSTRAINT uq_sv2_domain UNIQUE (company_id, domain),
CONSTRAINT uq_sv2_net UNIQUE (company_id, net)
);
CREATE INDEX IF NOT EXISTS idx_sv2_domain ON company_scope(domain) WHERE kind = 'domain';
CREATE INDEX IF NOT EXISTS idx_sv2_net ON company_scope USING GIST(net inet_ops) WHERE kind IN ('ip','cidr');
CREATE INDEX IF NOT EXISTS idx_sv2_company ON company_scope(company_id);
-- =====================================================================
-- B. 推理探索层
-- =====================================================================
CREATE TABLE IF NOT EXISTS explorations (
id BIGSERIAL PRIMARY KEY,
description TEXT,
goal TEXT NOT NULL,
status TEXT NOT NULL DEFAULT 'open'
CHECK (status IN ('open','achieved','failed')),
created_at TIMESTAMPTZ NOT NULL DEFAULT now(),
updated_at TIMESTAMPTZ NOT NULL DEFAULT now()
);
DROP TRIGGER IF EXISTS trg_exp_upd ON explorations;
CREATE TRIGGER trg_exp_upd BEFORE UPDATE ON explorations
FOR EACH ROW EXECUTE FUNCTION set_updated_at();
CREATE TABLE IF NOT EXISTS exploration_nodes (
id BIGSERIAL PRIMARY KEY,
exploration_id BIGINT NOT NULL REFERENCES explorations(id) ON DELETE CASCADE,
kind TEXT NOT NULL,
payload JSONB NOT NULL DEFAULT '{}',
priority INT NOT NULL DEFAULT 0,
state TEXT NOT NULL DEFAULT 'open',
origin TEXT,
owner TEXT,
created_at TIMESTAMPTZ NOT NULL DEFAULT now(),
updated_at TIMESTAMPTZ NOT NULL DEFAULT now(),
completed_at TIMESTAMPTZ,
CONSTRAINT ck_node_kind CHECK (kind IN ('begin','goal','intent','fact','finding','hint')),
CONSTRAINT ck_node_state CHECK (
(kind='begin' AND state IN ('open')) OR
(kind='intent' AND state IN ('open','running','done','blocked','exhausted','stopped')) OR
(kind='goal' AND state IN ('open','met','abandoned')) OR
(kind='fact' AND state IN ('confirmed','dismissed','origin')) OR
(kind='finding' AND state IN ('confirmed','dismissed')) OR
(kind='hint' AND state IN ('active','consumed'))
)
);
CREATE INDEX IF NOT EXISTS idx_expnodes_part ON exploration_nodes(exploration_id, kind);
CREATE INDEX IF NOT EXISTS idx_expnodes_frontier ON exploration_nodes(exploration_id, priority DESC)
WHERE kind='intent' AND state='open';
DROP TRIGGER IF EXISTS trg_expnodes_upd ON exploration_nodes;
CREATE TRIGGER trg_expnodes_upd BEFORE UPDATE ON exploration_nodes
FOR EACH ROW EXECUTE FUNCTION set_updated_at();
CREATE TABLE IF NOT EXISTS exploration_edges (
exploration_id BIGINT NOT NULL REFERENCES explorations(id) ON DELETE CASCADE,
src_id BIGINT NOT NULL REFERENCES exploration_nodes(id) ON DELETE CASCADE,
dst_id BIGINT NOT NULL REFERENCES exploration_nodes(id) ON DELETE CASCADE,
rel TEXT NOT NULL CHECK (rel IN ('spawns','derived_from','yields','proves')),
created_at TIMESTAMPTZ NOT NULL DEFAULT now(),
PRIMARY KEY (exploration_id, src_id, rel, dst_id),
CONSTRAINT ck_edge_noself CHECK (src_id <> dst_id)
);
CREATE INDEX IF NOT EXISTS idx_expedges_src ON exploration_edges(src_id, rel);
CREATE INDEX IF NOT EXISTS idx_expedges_dst ON exploration_edges(dst_id, rel);
CREATE TABLE IF NOT EXISTS exploration_anchors (
node_id BIGINT NOT NULL REFERENCES exploration_nodes(id) ON DELETE CASCADE,
asset_id BIGINT NOT NULL REFERENCES assets(id) ON DELETE CASCADE,
PRIMARY KEY (node_id, asset_id)
);
CREATE INDEX IF NOT EXISTS idx_anchor_asset ON exploration_anchors(asset_id);
CREATE TABLE IF NOT EXISTS activity (
id BIGSERIAL PRIMARY KEY,
exploration_id BIGINT NOT NULL REFERENCES explorations(id) ON DELETE CASCADE,
node_id BIGINT REFERENCES exploration_nodes(id) ON DELETE SET NULL,
worker TEXT,
kind TEXT,
tool TEXT,
tool_use_id TEXT,
is_error BOOLEAN NOT NULL DEFAULT false,
summary TEXT,
detail TEXT,
input_tokens INTEGER,
output_tokens INTEGER,
cache_read_tokens INTEGER,
cache_write_tokens INTEGER,
created_at TIMESTAMPTZ NOT NULL DEFAULT now()
);
CREATE INDEX IF NOT EXISTS idx_act_node ON activity(exploration_id, node_id, id);
CREATE INDEX IF NOT EXISTS idx_act_since ON activity(exploration_id, id);
-- =====================================================================
-- C. LLM profiles
-- =====================================================================
CREATE TABLE IF NOT EXISTS settings (
key TEXT PRIMARY KEY,
value TEXT NOT NULL,
updated_at TIMESTAMPTZ NOT NULL DEFAULT now()
);
CREATE TABLE IF NOT EXISTS llm_profiles (
id BIGSERIAL PRIMARY KEY,
name TEXT NOT NULL UNIQUE,
format TEXT NOT NULL CHECK (format IN ('openai','anthropic')),
base_url TEXT,
proxy TEXT,
model TEXT NOT NULL,
api_key TEXT,
api_key_hint TEXT,
rate_per_second DOUBLE PRECISION NOT NULL DEFAULT 0,
rate_per_minute DOUBLE PRECISION NOT NULL DEFAULT 0,
context_window_k INTEGER NOT NULL DEFAULT 0,
reasoning_effort TEXT NOT NULL DEFAULT '',
is_default BOOLEAN NOT NULL DEFAULT false,
created_at TIMESTAMPTZ NOT NULL DEFAULT now(),
updated_at TIMESTAMPTZ NOT NULL DEFAULT now()
);
CREATE UNIQUE INDEX IF NOT EXISTS uq_llm_one_default ON llm_profiles(is_default) WHERE is_default;
DROP TRIGGER IF EXISTS trg_llm_upd ON llm_profiles;
CREATE TRIGGER trg_llm_upd BEFORE UPDATE ON llm_profiles
FOR EACH ROW EXECUTE FUNCTION set_updated_at();
-- =====================================================================
-- D. 任务层
-- =====================================================================
CREATE TABLE IF NOT EXISTS tasks (
id BIGSERIAL PRIMARY KEY,
description TEXT NOT NULL,
goal TEXT NOT NULL,
exploration_id BIGINT NOT NULL UNIQUE
REFERENCES explorations(id) ON DELETE RESTRICT,
status TEXT NOT NULL DEFAULT 'created'
CHECK (status IN ('created','running','paused','done','failed','timeout')),
paused BOOLEAN NOT NULL DEFAULT false,
llm_profile_id BIGINT REFERENCES llm_profiles(id) ON DELETE SET NULL,
company_id BIGINT REFERENCES companies(id) ON DELETE SET NULL,
parent_ref TEXT,
timeout_seconds INTEGER NOT NULL DEFAULT 0,
first_run_at TIMESTAMPTZ,
deadline_at TIMESTAMPTZ,
deleted_at TIMESTAMPTZ,
completed_at TIMESTAMPTZ,
created_at TIMESTAMPTZ NOT NULL DEFAULT now(),
updated_at TIMESTAMPTZ NOT NULL DEFAULT now()
);
CREATE INDEX IF NOT EXISTS idx_tasks_alive ON tasks(created_at DESC) WHERE deleted_at IS NULL;
CREATE INDEX IF NOT EXISTS idx_tasks_status ON tasks(status) WHERE deleted_at IS NULL;
DROP TRIGGER IF EXISTS trg_tasks_upd ON tasks;
CREATE TRIGGER trg_tasks_upd BEFORE UPDATE ON tasks
FOR EACH ROW EXECUTE FUNCTION set_updated_at();
-- =====================================================================
-- E. Agents / 提示词模板 / 变量目录
-- =====================================================================
CREATE TABLE IF NOT EXISTS agents (
id BIGSERIAL PRIMARY KEY,
key TEXT NOT NULL UNIQUE CHECK (key ~ '^[a-z][a-z0-9_]*$'),
name TEXT NOT NULL,
description TEXT,
role TEXT NOT NULL,
builtin BOOLEAN NOT NULL DEFAULT true,
enabled BOOLEAN NOT NULL DEFAULT true,
llm_profile_id BIGINT REFERENCES llm_profiles(id) ON DELETE SET NULL,
current_prompt_id BIGINT,
max_turns INTEGER NOT NULL DEFAULT 0,
run_seconds INTEGER NOT NULL DEFAULT 600,
web_search BOOLEAN NOT NULL DEFAULT false,
interactive_shell BOOLEAN NOT NULL DEFAULT false,
wrapup_prompt TEXT NOT NULL DEFAULT '',
wrapup_max_turns INTEGER NOT NULL DEFAULT 0,
task_timeout_wrapup_prompt TEXT NOT NULL DEFAULT '',
task_timeout_wrapup_max_turns INTEGER NOT NULL DEFAULT 0,
created_at TIMESTAMPTZ NOT NULL DEFAULT now(),
updated_at TIMESTAMPTZ NOT NULL DEFAULT now(),
CONSTRAINT agents_role_ck CHECK (role IN ('goals','main','planner','worker','assistant'))
);
DROP TRIGGER IF EXISTS trg_agents_upd ON agents;
CREATE TRIGGER trg_agents_upd BEFORE UPDATE ON agents
FOR EACH ROW EXECUTE FUNCTION set_updated_at();
CREATE TABLE IF NOT EXISTS agent_prompts (
id BIGSERIAL PRIMARY KEY,
agent_id BIGINT NOT NULL REFERENCES agents(id) ON DELETE CASCADE,
version INT NOT NULL,
template_text TEXT NOT NULL,
note TEXT,
updated_by TEXT,
created_at TIMESTAMPTZ NOT NULL DEFAULT now(),
UNIQUE (agent_id, version)
);
-- 循环外键：agents.current_prompt_id → agent_prompts.id（需在两表创建后加）
DO $$ BEGIN
ALTER TABLE agents ADD CONSTRAINT fk_agents_curprompt
FOREIGN KEY (current_prompt_id) REFERENCES agent_prompts(id) ON DELETE SET NULL;
EXCEPTION WHEN duplicate_object THEN NULL; END $$;
CREATE TABLE IF NOT EXISTS agent_prompt_vars (
id BIGSERIAL PRIMARY KEY,
agent_id BIGINT NOT NULL REFERENCES agents(id) ON DELETE CASCADE,
var_name TEXT NOT NULL,
description TEXT,
example TEXT,
source TEXT NOT NULL CHECK (source IN ('exploration','runtime','distilled')),
UNIQUE (agent_id, var_name)
);
-- =====================================================================
-- F. MCP 服务
-- =====================================================================
CREATE TABLE IF NOT EXISTS mcp_servers (
id BIGSERIAL PRIMARY KEY,
name TEXT NOT NULL UNIQUE,
transport TEXT NOT NULL CHECK (transport IN ('stdio','http')),
command TEXT,
args JSONB NOT NULL DEFAULT '[]',
env JSONB NOT NULL DEFAULT '{}',
url TEXT,
enabled BOOLEAN NOT NULL DEFAULT true,
created_at TIMESTAMPTZ NOT NULL DEFAULT now(),
updated_at TIMESTAMPTZ NOT NULL DEFAULT now()
);
DROP TRIGGER IF EXISTS trg_mcp_upd ON mcp_servers;
CREATE TRIGGER trg_mcp_upd BEFORE UPDATE ON mcp_servers
FOR EACH ROW EXECUTE FUNCTION set_updated_at();
-- 默认数据源占位：ScopeSentry 资产同步 MCP（地址与认证均留空、未启用）。
-- 供「资产同步」页检测数据源是否已配置；用户在页面填入 url 与 X-API-Key 后再启用。
-- 仅在缺失时插入，绝不覆盖用户已配置/已启用的服务器（schema.sql 每次启动都会 Exec）。
INSERT INTO mcp_servers (name, transport, url, env, enabled)
VALUES ('ScopeSentry', 'http', NULL, '{"X-API-Key":""}', false)
ON CONFLICT (name) DO NOTHING;
CREATE TABLE IF NOT EXISTS mcp_tools_cache (
id BIGSERIAL PRIMARY KEY,
server_id BIGINT NOT NULL REFERENCES mcp_servers(id) ON DELETE CASCADE,
tool_name TEXT NOT NULL,
description TEXT,
schema JSONB,
discovered_at TIMESTAMPTZ NOT NULL DEFAULT now(),
UNIQUE (server_id, tool_name)
);
-- =====================================================================
-- G. 可见性：agent × mcp / skill
-- =====================================================================
CREATE TABLE IF NOT EXISTS agent_visibility (
agent_id BIGINT NOT NULL REFERENCES agents(id) ON DELETE CASCADE,
resource_kind TEXT NOT NULL CHECK (resource_kind IN ('mcp')),
resource_id BIGINT NOT NULL,
mcp_tool_name TEXT NOT NULL DEFAULT '',
enabled BOOLEAN NOT NULL DEFAULT true,
created_at TIMESTAMPTZ NOT NULL DEFAULT now(),
PRIMARY KEY (agent_id, resource_kind, resource_id, mcp_tool_name)
);
CREATE INDEX IF NOT EXISTS idx_vis_resource ON agent_visibility(resource_kind, resource_id);
CREATE TABLE IF NOT EXISTS agent_skill_visibility (
agent_id BIGINT NOT NULL REFERENCES agents(id) ON DELETE CASCADE,
skill_name TEXT NOT NULL,
enabled BOOLEAN NOT NULL DEFAULT true,
created_at TIMESTAMPTZ NOT NULL DEFAULT now(),
PRIMARY KEY (agent_id, skill_name)
);
CREATE INDEX IF NOT EXISTS idx_askv_skill ON agent_skill_visibility(skill_name);
-- =====================================================================
-- H. 内置工具目录
-- =====================================================================
CREATE TABLE IF NOT EXISTS tools (
key TEXT PRIMARY KEY,
system BOOLEAN NOT NULL DEFAULT true,
description TEXT NOT NULL DEFAULT '',
schema JSONB NOT NULL DEFAULT '{}',
agents JSONB NOT NULL DEFAULT '[]',
enabled BOOLEAN NOT NULL DEFAULT true,
kind TEXT NOT NULL DEFAULT 'builtin',
exec JSONB NOT NULL DEFAULT '{}',
deferred BOOLEAN NOT NULL DEFAULT false,
updated_at TIMESTAMPTZ NOT NULL DEFAULT now()
);
DROP TRIGGER IF EXISTS trg_tools_upd ON tools;
CREATE TRIGGER trg_tools_upd BEFORE UPDATE ON tools
FOR EACH ROW EXECUTE FUNCTION set_updated_at();
-- =====================================================================
-- I. 会话（对话页）
-- =====================================================================
CREATE TABLE IF NOT EXISTS conversations (
id BIGSERIAL PRIMARY KEY,
agent_key TEXT NOT NULL,
title TEXT NOT NULL DEFAULT '',
llm_profile_id BIGINT REFERENCES llm_profiles(id) ON DELETE SET NULL,
created_at TIMESTAMPTZ NOT NULL DEFAULT now(),
updated_at TIMESTAMPTZ NOT NULL DEFAULT now()
);
DROP TRIGGER IF EXISTS trg_conversations_upd ON conversations;
CREATE TRIGGER trg_conversations_upd BEFORE UPDATE ON conversations
FOR EACH ROW EXECUTE FUNCTION set_updated_at();
CREATE TABLE IF NOT EXISTS conversation_activities (
id BIGSERIAL PRIMARY KEY,
conversation_id BIGINT NOT NULL REFERENCES conversations(id) ON DELETE CASCADE,
worker TEXT,
kind TEXT,
tool TEXT,
tool_use_id TEXT,
is_error BOOLEAN NOT NULL DEFAULT false,
summary TEXT,
detail TEXT,
input_tokens INTEGER,
output_tokens INTEGER,
cache_read_tokens INTEGER,
cache_write_tokens INTEGER,
created_at TIMESTAMPTZ NOT NULL DEFAULT now()
);
CREATE INDEX IF NOT EXISTS idx_conv_act ON conversation_activities(conversation_id, id);
-- =====================================================================
-- J. Agent 触发器
-- =====================================================================
CREATE TABLE IF NOT EXISTS agent_triggers (
id BIGSERIAL PRIMARY KEY,
agent_key TEXT NOT NULL,
enabled BOOLEAN NOT NULL DEFAULT true,
interval_sec INTEGER NOT NULL DEFAULT 0,
on_finding BOOLEAN NOT NULL DEFAULT false,
on_goal_met BOOLEAN NOT NULL DEFAULT false,
on_task_timeout BOOLEAN NOT NULL DEFAULT false,
interval_message TEXT NOT NULL DEFAULT '',
finding_message TEXT NOT NULL DEFAULT '',
goal_message TEXT NOT NULL DEFAULT '',
task_timeout_message TEXT NOT NULL DEFAULT '',
last_fire TIMESTAMPTZ,
created_at TIMESTAMPTZ NOT NULL DEFAULT now(),
updated_at TIMESTAMPTZ NOT NULL DEFAULT now()
);
CREATE INDEX IF NOT EXISTS idx_agent_triggers_agent ON agent_triggers(agent_key);
DROP TRIGGER IF EXISTS trg_agent_triggers_upd ON agent_triggers;
CREATE TRIGGER trg_agent_triggers_upd BEFORE UPDATE ON agent_triggers
FOR EACH ROW EXECUTE FUNCTION set_updated_at();
CREATE TABLE IF NOT EXISTS scheduler_state (
key TEXT PRIMARY KEY,
value TEXT NOT NULL DEFAULT ''
);
-- =====================================================================
-- K. 拦截规则
-- =====================================================================
CREATE TABLE IF NOT EXISTS intercept_rules (
id BIGSERIAL PRIMARY KEY,
name TEXT NOT NULL,
enabled BOOLEAN NOT NULL DEFAULT true,
priority INTEGER NOT NULL DEFAULT 0,
match_target TEXT NOT NULL CHECK (match_target IN ('tool_name', 'tool_input')),
match_type TEXT NOT NULL CHECK (match_type IN ('string', 'regex')),
pattern TEXT NOT NULL,
action TEXT NOT NULL CHECK (action IN ('allow', 'deny', 'ask')),
message TEXT NOT NULL DEFAULT '',
timeout_enabled BOOLEAN NOT NULL DEFAULT true,
timeout_seconds INTEGER NOT NULL DEFAULT 60,
timeout_action TEXT NOT NULL DEFAULT 'deny',
created_at TIMESTAMPTZ NOT NULL DEFAULT NOW(),
updated_at TIMESTAMPTZ NOT NULL DEFAULT NOW()
);
DROP TRIGGER IF EXISTS trg_intercept_rules_upd ON intercept_rules;
CREATE TRIGGER trg_intercept_rules_upd BEFORE UPDATE ON intercept_rules
FOR EACH ROW EXECUTE FUNCTION set_updated_at();
CREATE TABLE IF NOT EXISTS intercept_pending (
id BIGSERIAL PRIMARY KEY,
rule_id BIGINT REFERENCES intercept_rules(id) ON DELETE SET NULL,
conversation_id BIGINT REFERENCES conversations(id) ON DELETE CASCADE,
task_id TEXT,
agent_name TEXT NOT NULL DEFAULT '',
tool_name TEXT NOT NULL,
tool_input JSONB NOT NULL DEFAULT '{}',
status TEXT NOT NULL DEFAULT 'pending'
CHECK (status IN ('pending', 'allowed', 'denied', 'timeout')),
decided_at TIMESTAMPTZ,
created_at TIMESTAMPTZ NOT NULL DEFAULT NOW()
);
CREATE INDEX IF NOT EXISTS idx_intercept_pending_status ON intercept_pending(status, created_at DESC);
CREATE INDEX IF NOT EXISTS idx_intercept_pending_task ON intercept_pending(task_id, created_at DESC);
-- =====================================================================
-- L. 漏洞发现持久化
-- =====================================================================
CREATE TABLE IF NOT EXISTS findings (
id BIGSERIAL PRIMARY KEY,
task_id BIGINT REFERENCES tasks(id) ON DELETE SET NULL,
node_id BIGINT REFERENCES exploration_nodes(id) ON DELETE SET NULL,
vulnclass TEXT NOT NULL DEFAULT '',
severity TEXT NOT NULL DEFAULT '',
summary TEXT NOT NULL DEFAULT '',
evidence TEXT NOT NULL DEFAULT '',
worker TEXT NOT NULL DEFAULT '',
asset_ids JSONB NOT NULL DEFAULT '[]',
created_at TIMESTAMPTZ NOT NULL DEFAULT now()
);
CREATE INDEX IF NOT EXISTS idx_findings_task ON findings(task_id, created_at DESC);
CREATE INDEX IF NOT EXISTS idx_findings_time ON findings(created_at DESC);
-- =====================================================================
-- M. 后端日志持久化
-- =====================================================================
CREATE TABLE IF NOT EXISTS server_logs (
id BIGSERIAL PRIMARY KEY,
created_at TIMESTAMPTZ NOT NULL DEFAULT now(),
level TEXT NOT NULL DEFAULT 'info',
tag TEXT NOT NULL DEFAULT '',
text TEXT NOT NULL DEFAULT ''
);
CREATE INDEX IF NOT EXISTS idx_server_logs_id ON server_logs(id DESC);
db/settings.go
+43
View File
@@ -0,0 +1,43 @@
package db
import "database/sql"
// Settings is a tiny key-value store for global app config the UI toggles at
// runtime (e.g. traffic_capture). Missing keys fall back to caller defaults.
// GetSetting returns the stored value and ok=false when the key is unset.
func (d *DB) GetSetting(key string) (value string, ok bool, err error) {
err = d.QueryRow(`SELECT value FROM settings WHERE key=$1`, key).Scan(&value)
if err == sql.ErrNoRows {
return "", false, nil
}
if err != nil {
return "", false, err
}
return value, true, nil
}
// SetSetting upserts a setting value.
func (d *DB) SetSetting(key, value string) error {
_, err := d.Exec(`
INSERT INTO settings(key, value) VALUES ($1, $2)
ON CONFLICT (key) DO UPDATE SET value = EXCLUDED.value, updated_at = now()`, key, value)
return err
}
// GetBool returns the boolean setting, or def when unset/unparseable.
func (d *DB) GetBool(key string, def bool) bool {
v, ok, err := d.GetSetting(key)
if err != nil || !ok {
return def
}
return v == "true" || v == "1"
}
// SetBool stores a boolean setting as "true"/"false".
func (d *DB) SetBool(key string, val bool) error {
if val {
return d.SetSetting(key, "true")
}
return d.SetSetting(key, "false")
}
db/tasks.go
+196
View File
@@ -0,0 +1,196 @@
package db
import (
"encoding/json"
"time"
)
// Task is a row in the task registry (1:1 with an exploration).
type Task struct {
ID int64 `json:"id"`
Description string `json:"description"`
Goal string `json:"goal"`
ExplorationID int64 `json:"exploration_id"`
Status string `json:"status"`
Paused bool `json:"paused"`
LLMProfileID *int64 `json:"llm_profile_id,omitempty"`
ParentRef string `json:"parent_ref,omitempty"` // 父任务 id(编排 spawn 记录;空=顶层)
CreatedAt time.Time `json:"created_at"`
CompletedAt *time.Time `json:"completed_at,omitempty"` // 进入终态(done/failed/timeout)的时刻;非终态为 nil
// 任务级超时(见 docs/任务级超时与收尾设计.md)。
TimeoutSeconds int `json:"timeout_seconds"` // 0=不限时
FirstRunAt *time.Time `json:"first_run_at,omitempty"` // 首次真正开始运行的时刻(非 created_at);nil=尚未运行
DeadlineAt *time.Time `json:"deadline_at,omitempty"` // = first_run_at + timeout_seconds;nil=不限或未运行
}
// IsTerminal reports whether a task status is a terminal (finished) state.
// 单一真源，替换散落各处的 done/failed 硬编码判定。
func IsTerminal(status string) bool {
return status == "done" || status == "failed" || status == "timeout"
}
// CreateTask creates an exploration + task in one transaction and returns the task.
// timeoutSeconds is the task-level wall-clock budget (0 = 不限时); deadline_at is
// stamped later at first real run (see engine), not here.
func (d *DB) CreateTask(description, goal string, llmProfileID *int64, timeoutSeconds int) (*Task, error) {
tx, err := d.Begin()
if err != nil {
return nil, err
}
defer tx.Rollback()
var expID int64
if err := tx.QueryRow(`INSERT INTO explorations(description, goal) VALUES ($1,$2) RETURNING id`, description, goal).Scan(&expID); err != nil {
return nil, err
}
// origin fact: the exploration graph's root, a KindFact node (state='origin')
// holding the original task description. Goals/intents/findings descend from
// it, and the seeded target asset is anchored to it as lineage (the asset graph
// is global and shared, not isolated per task). Being a fact (not a special 'begin' kind) lets every
// intent uniformly connect to a fact node, including the first ones.
originPayload, _ := json.Marshal(map[string]any{
"summary": "任务起点：" + description + "；目标：" + goal,
"description": description,
"goal": goal,
})
if _, err := tx.Exec(`
INSERT INTO exploration_nodes(exploration_id, kind, payload, priority, state, origin)
VALUES ($1, 'fact', $2, 0, 'origin', 'system')`, expID, string(originPayload)); err != nil {
return nil, err
}
if timeoutSeconds < 0 {
timeoutSeconds = 0
}
t := &Task{Description: description, Goal: goal, ExplorationID: expID, LLMProfileID: llmProfileID, TimeoutSeconds: timeoutSeconds}
if err := tx.QueryRow(`
INSERT INTO tasks(description, goal, exploration_id, llm_profile_id, timeout_seconds) VALUES ($1,$2,$3,$4,$5)
RETURNING id, status, paused, created_at`, description, goal, expID, llmProfileID, timeoutSeconds).Scan(&t.ID, &t.Status, &t.Paused, &t.CreatedAt); err != nil {
return nil, err
}
return t, tx.Commit()
}
const taskCols = `id, description, goal, exploration_id, status, paused, llm_profile_id, COALESCE(parent_ref,''), created_at, completed_at, COALESCE(timeout_seconds,0), first_run_at, deadline_at`
func scanTask(sc interface{ Scan(...any) error }) (*Task, error) {
var t Task
if err := sc.Scan(&t.ID, &t.Description, &t.Goal, &t.ExplorationID, &t.Status, &t.Paused, &t.LLMProfileID, &t.ParentRef, &t.CreatedAt, &t.CompletedAt, &t.TimeoutSeconds, &t.FirstRunAt, &t.DeadlineAt); err != nil {
return nil, err
}
return &t, nil
}
// SetParentRef records a task's parent task id (编排 agent spawn_task 关联).
func (d *DB) SetParentRef(id int64, parentRef string) error {
_, err := d.Exec(`UPDATE tasks SET parent_ref=NULLIF($2,'') WHERE id=$1`, id, parentRef)
return err
}
// ListTasks returns alive tasks, newest first.
func (d *DB) ListTasks() ([]*Task, error) {
rows, err := d.Query(`SELECT ` + taskCols + ` FROM tasks WHERE deleted_at IS NULL ORDER BY created_at DESC`)
if err != nil {
return nil, err
}
defer rows.Close()
var out []*Task
for rows.Next() {
t, err := scanTask(rows)
if err != nil {
return nil, err
}
out = append(out, t)
}
return out, rows.Err()
}
// GetTask returns one alive task (nil if not found/deleted).
func (d *DB) GetTask(id int64) (*Task, error) {
t, err := scanTask(d.QueryRow(`SELECT `+taskCols+` FROM tasks WHERE id=$1 AND deleted_at IS NULL`, id))
if err != nil {
if err.Error() == "sql: no rows in result set" {
return nil, nil
}
return nil, err
}
return t, nil
}
// SetPaused persists a task's paused flag.
func (d *DB) SetPaused(id int64, paused bool) error {
_, err := d.Exec(`UPDATE tasks SET paused=$1 WHERE id=$2`, paused, id)
return err
}
// SetStatus updates a task's lifecycle status. Entering a terminal state
// (done/failed/timeout) stamps completed_at once (COALESCE keeps the first stamp
// stable); moving back to a non-terminal state clears it, so a re-run has no stale
// finish time.
func (d *DB) SetStatus(id int64, status string) error {
_, err := d.Exec(`
UPDATE tasks
SET status = $1,
completed_at = CASE WHEN $1 IN ('done','failed','timeout') THEN COALESCE(completed_at, now()) ELSE NULL END
WHERE id = $2`, status, id)
return err
}
// SetTerminalStatusGuarded sets a terminal status only when the task is NOT already
// terminal, so a completed↔timeout race resolves to the first writer (won=true).
// Returns won=false (no error) when another terminal status already stuck — the
// caller then leaves it alone. Non-terminal transitions / re-run still use SetStatus.
func (d *DB) SetTerminalStatusGuarded(id int64, status string) (won bool, err error) {
res, err := d.Exec(`
UPDATE tasks
SET status = $1,
completed_at = COALESCE(completed_at, now())
WHERE id = $2 AND status NOT IN ('done','failed','timeout')`, status, id)
if err != nil {
return false, err
}
n, _ := res.RowsAffected()
return n > 0, nil
}
// StampFirstRun records a task's first-real-run moment and computes its absolute
// deadline (= now + timeoutSeconds). Idempotent: only stamps when first_run_at is
// still NULL, so restarts / re-entries keep the original clock. timeoutSeconds<=0
// leaves deadline_at NULL (不限时). Returns the resulting deadline (nil = 不限/未变).
func (d *DB) StampFirstRun(id int64, timeoutSeconds int) (*time.Time, error) {
var deadline *time.Time
err := d.QueryRow(`
UPDATE tasks
SET first_run_at = COALESCE(first_run_at, now()),
deadline_at = CASE
WHEN first_run_at IS NOT NULL THEN deadline_at -- 已盖过章：不动
WHEN $2 > 0 THEN now() + make_interval(secs => $2)
ELSE NULL END
WHERE id = $1
RETURNING deadline_at`, id, timeoutSeconds).Scan(&deadline)
return deadline, err
}
// DeleteTask hard-deletes a task and its exploration subgraph (nodes/edges/
// activity via ON DELETE CASCADE). The task row is removed first so the
// ON DELETE RESTRICT on tasks.exploration_id does not block the exploration delete.
// Global assets are NOT touched.
func (d *DB) DeleteTask(id int64) error {
tx, err := d.Begin()
if err != nil {
return err
}
defer tx.Rollback()
var expID int64
if err := tx.QueryRow(`SELECT exploration_id FROM tasks WHERE id=$1`, id).Scan(&expID); err != nil {
if err.Error() == "sql: no rows in result set" {
return nil // already gone
}
return err
}
if _, err := tx.Exec(`DELETE FROM tasks WHERE id=$1`, id); err != nil {
return err
}
if _, err := tx.Exec(`DELETE FROM explorations WHERE id=$1`, expID); err != nil {
return err
}
return tx.Commit()
}
db/tasks_test.go
+60
View File
@@ -0,0 +1,60 @@
package db
import "testing"
func TestTaskLifecycleAndDeleteCascade(t *testing.T) {
d, err := Open(testDSN(t))
if err != nil {
t.Skipf("postgres unavailable (%v) — skipping", err)
}
defer d.Close()
tk, err := d.CreateTask("迁移测试", "目标X", nil, 0)
if err != nil {
t.Fatal(err)
}
// populate the exploration subgraph
es := d.Exploration(tk.ExplorationID)
if _, err := es.AddIntent(map[string]any{"summary": "x"}, 5, nil, "planner"); err != nil {
t.Fatal(err)
}
// pause + status
if err := d.SetPaused(tk.ID, true); err != nil {
t.Fatal(err)
}
got, err := d.GetTask(tk.ID)
if err != nil || got == nil || !got.Paused {
t.Fatalf("paused not persisted: %+v err=%v", got, err)
}
// list contains it
list, _ := d.ListTasks()
found := false
for _, x := range list {
if x.ID == tk.ID {
found = true
}
}
if !found {
t.Fatalf("task not in list")
}
// delete cascades exploration subgraph
if err := d.DeleteTask(tk.ID); err != nil {
t.Fatal(err)
}
if g, _ := d.GetTask(tk.ID); g != nil {
t.Fatalf("task should be gone")
}
var nodes int
d.QueryRow(`SELECT count(*) FROM exploration_nodes WHERE exploration_id=$1`, tk.ExplorationID).Scan(&nodes)
if nodes != 0 {
t.Fatalf("exploration nodes should be cascade-deleted, got %d", nodes)
}
var exps int
d.QueryRow(`SELECT count(*) FROM explorations WHERE id=$1`, tk.ExplorationID).Scan(&exps)
if exps != 0 {
t.Fatalf("exploration should be deleted, got %d", exps)
}
}
db/testmain_test.go
+32
View File
@@ -0,0 +1,32 @@
package db
import (
"database/sql"
"os"
"testing"
_ "github.com/jackc/pgx/v5/stdlib"
)
// TestMain acquires a PostgreSQL advisory lock (7337741002) for the entire
// db test suite. Packages agent and server hold the same lock, so parallel
// `go test ./...` runs serialize across packages on the shared dev DB and
// avoid cross-package cleanup races (e.g. DELETE FROM assets WHERE id > X
// from one package deleting assets created by another).
func TestMain(m *testing.M) {
dsn, _, err := DSN()
if err != nil {
// No DB configured — tests that need PG will skip themselves.
os.Exit(m.Run())
}
conn, err := sql.Open("pgx", dsn)
if err != nil || conn.Ping() != nil {
os.Exit(m.Run())
}
defer conn.Close()
if _, err := conn.Exec(`SELECT pg_advisory_lock(7337741002)`); err != nil {
os.Exit(m.Run())
}
defer conn.Exec(`SELECT pg_advisory_unlock(7337741002)`) //nolint:errcheck
os.Exit(m.Run())
}
db/tools.go
+223
View File
@@ -0,0 +1,223 @@
package db
import (
"database/sql"
"encoding/json"
)
// Tool is one row of the built-in tool catalog. key + handler live in code; this
// row carries only the page-editable surface: description, parameter schema
// (structure read-only, per-param description/default editable), agent binding,
// and the on/off switch. See schema.sql §H and agent/toolcatalog.go.
type Tool struct {
Key string `json:"key"`
System bool `json:"system"`
Description string `json:"description"`
Schema json.RawMessage `json:"schema"`
Agents []string `json:"agents"`
Enabled bool `json:"enabled"`
Kind string `json:"kind"` // builtin | command | script | http
Exec json.RawMessage `json:"exec"` // 自定义工具执行规格(kind!=builtin)
Deferred bool `json:"deferred"` // schema 延迟(走 SearchExtraTools/ExecuteExtraTool)
}
const toolCols = `key, system, description, schema, agents, enabled, kind, exec, deferred`
// SeedTool inserts a built-in tool's code-defined defaults ONCE. ON CONFLICT DO
// NOTHING: an existing row (possibly edited in the UI) is never overwritten on
// startup — that's what keeps page edits from being wiped every restart. Use
// UpsertToolForce for an explicit "reset to code default".
func (d *DB) SeedTool(key, desc string, schema, agents json.RawMessage) error {
if len(schema) == 0 {
schema = json.RawMessage("{}")
}
if len(agents) == 0 {
agents = json.RawMessage("[]")
}
_, err := d.Exec(`
INSERT INTO tools(key, system, description, schema, agents, enabled)
VALUES ($1, true, $2, $3, $4, true)
ON CONFLICT (key) DO NOTHING`, key, desc, schema, agents)
return err
}
// AddAgentToToolBinding adds an agent key to the given tools' `agents` arrays if
// not already present (idempotent). Used to give the built-in Auto agent its
// default toolset on existing DBs.
func (d *DB) AddAgentToToolBinding(agentKey string, keys []string) error {
for _, k := range keys {
if _, err := d.Exec(`UPDATE tools SET agents = agents || to_jsonb($1::text) WHERE key=$2 AND NOT (agents ? $1)`, agentKey, k); err != nil {
return err
}
}
return nil
}
// RemoveAgentFromToolBindings strips an agent key from every tool's `agents`
// JSONB array — called when a custom agent is deleted so no tool keeps a dangling
// binding. Uses jsonb `-` (remove array element) guarded by `?` (membership).
func (d *DB) RemoveAgentFromToolBindings(agentKey string) error {
_, err := d.Exec(`UPDATE tools SET agents = agents - $1 WHERE agents ? $1`, agentKey)
return err
}
// UpsertToolForce overwrites a tool row with the given code-default values (used by
// the per-tool "reset" action). It resets description/schema/agents and re-enables
// the tool, but keeps system=true.
func (d *DB) UpsertToolForce(key, desc string, schema, agents json.RawMessage) error {
if len(schema) == 0 {
schema = json.RawMessage("{}")
}
if len(agents) == 0 {
agents = json.RawMessage("[]")
}
_, err := d.Exec(`
INSERT INTO tools(key, system, description, schema, agents, enabled)
VALUES ($1, true, $2, $3, $4, true)
ON CONFLICT (key) DO UPDATE
SET description = EXCLUDED.description,
schema = EXCLUDED.schema,
agents = EXCLUDED.agents,
enabled = true`, key, desc, schema, agents)
return err
}
// RefreshToolDefaults updates a system tool's model-facing description + schema to the
// code defaults, PRESERVING the user's agent binding + enabled flag. Used by one-time
// migrations to propagate a code schema change (SeedTool is first-insert-only, so a new
// parameter added in code otherwise never reaches an already-seeded row). No-op for
// custom tools or unknown keys.
func (d *DB) RefreshToolDefaults(key, desc string, schema json.RawMessage) error {
if len(schema) == 0 {
schema = json.RawMessage("{}")
}
_, err := d.Exec(`UPDATE tools SET description=$2, schema=$3, updated_at=now() WHERE key=$1 AND system`, key, desc, schema)
return err
}
func scanTool(rows interface{ Scan(...any) error }) (*Tool, error) {
var t Tool
var agents []byte
if err := rows.Scan(&t.Key, &t.System, &t.Description, &t.Schema, &agents, &t.Enabled, &t.Kind, &t.Exec, &t.Deferred); err != nil {
return nil, err
}
if len(agents) > 0 {
_ = json.Unmarshal(agents, &t.Agents)
}
if t.Agents == nil {
t.Agents = []string{}
}
return &t, nil
}
// ListTools returns the whole tool catalog, ordered by key.
func (d *DB) ListTools() ([]*Tool, error) {
rows, err := d.Query(`SELECT ` + toolCols + ` FROM tools ORDER BY key`)
if err != nil {
return nil, err
}
defer rows.Close()
var out []*Tool
for rows.Next() {
t, err := scanTool(rows)
if err != nil {
return nil, err
}
out = append(out, t)
}
return out, rows.Err()
}
// GetTool fetches one tool row (nil if absent).
func (d *DB) GetTool(key string) (*Tool, error) {
row := d.QueryRow(`SELECT `+toolCols+` FROM tools WHERE key=$1`, key)
t, err := scanTool(row)
if err == sql.ErrNoRows {
return nil, nil
}
return t, err
}
// CreateCustomTool inserts a user-defined tool (system=false) with an execution
// spec. Fails if the key already exists.
func (d *DB) CreateCustomTool(t *Tool) error {
schema := t.Schema
if len(schema) == 0 {
schema = json.RawMessage("{}")
}
exec := t.Exec
if len(exec) == 0 {
exec = json.RawMessage("{}")
}
agents, _ := json.Marshal(t.Agents)
if len(agents) == 0 {
agents = json.RawMessage("[]")
}
_, err := d.Exec(`
INSERT INTO tools(key, system, description, schema, agents, enabled, kind, exec, deferred)
VALUES ($1, false, $2, $3, $4, $5, $6, $7, $8)`,
t.Key, t.Description, schema, agents, t.Enabled, t.Kind, exec, t.Deferred)
return err
}
// UpdateCustomTool updates a custom tool's editable fields (kind/exec/deferred +
// desc/schema/agents/enabled). Only touches system=false rows.
func (d *DB) UpdateCustomTool(t *Tool) error {
schema := t.Schema
if len(schema) == 0 {
schema = json.RawMessage("{}")
}
exec := t.Exec
if len(exec) == 0 {
exec = json.RawMessage("{}")
}
agents, _ := json.Marshal(t.Agents)
if len(agents) == 0 {
agents = json.RawMessage("[]")
}
_, err := d.Exec(`
UPDATE tools SET description=$2, schema=$3, agents=$4, enabled=$5, kind=$6, exec=$7, deferred=$8
WHERE key=$1 AND system=false`,
t.Key, t.Description, schema, agents, t.Enabled, t.Kind, exec, t.Deferred)
return err
}
// DeleteCustomTool removes a custom tool (system=false only; built-ins protected).
func (d *DB) DeleteCustomTool(key string) error {
_, err := d.Exec(`DELETE FROM tools WHERE key=$1 AND system=false`, key)
return err
}
// ListCustomTools returns only the user-defined (system=false) tools.
func (d *DB) ListCustomTools() ([]*Tool, error) {
rows, err := d.Query(`SELECT ` + toolCols + ` FROM tools WHERE system=false ORDER BY key`)
if err != nil {
return nil, err
}
defer rows.Close()
var out []*Tool
for rows.Next() {
t, err := scanTool(rows)
if err != nil {
return nil, err
}
out = append(out, t)
}
return out, rows.Err()
}
// UpdateTool saves the page-editable fields. key is never changed (it is welded to
// the Go handler). system tools: the caller must keep the schema structure — only
// per-param description/default and the agent binding are meant to move.
func (d *DB) UpdateTool(key, desc string, schema, agents json.RawMessage, enabled bool) error {
if len(schema) == 0 {
schema = json.RawMessage("{}")
}
if len(agents) == 0 {
agents = json.RawMessage("[]")
}
_, err := d.Exec(`
UPDATE tools SET description=$2, schema=$3, agents=$4, enabled=$5 WHERE key=$1`,
key, desc, schema, agents, enabled)
return err
}
db/triggers.go
+215
View File
@@ -0,0 +1,215 @@
package db
import (
"database/sql"
"encoding/json"
"time"
)
// AgentTrigger is one P3 trigger attached to a custom agent. Four trigger
// conditions can be on at once: interval / on_finding / on_goal_met / on_task_timeout.
type AgentTrigger struct {
ID int64 `json:"id"`
AgentKey string `json:"agent_key"`
Enabled bool `json:"enabled"`
IntervalSec int `json:"interval_sec"`
OnFinding bool `json:"on_finding"`
OnGoalMet bool `json:"on_goal_met"`
OnTaskTimeout bool `json:"on_task_timeout"`
IntervalMessage string `json:"interval_message"` // 各触发条件的独立用户消息
FindingMessage string `json:"finding_message"`
GoalMessage string `json:"goal_message"`
TaskTimeoutMessage string `json:"task_timeout_message"`
LastFire *time.Time `json:"last_fire,omitempty"`
}
const triggerCols = `id, agent_key, enabled, interval_sec, on_finding, on_goal_met, on_task_timeout, interval_message, finding_message, goal_message, task_timeout_message, last_fire`
func scanTrigger(sc interface{ Scan(...any) error }) (*AgentTrigger, error) {
var t AgentTrigger
var lf sql.NullTime
if err := sc.Scan(&t.ID, &t.AgentKey, &t.Enabled, &t.IntervalSec, &t.OnFinding, &t.OnGoalMet, &t.OnTaskTimeout,
&t.IntervalMessage, &t.FindingMessage, &t.GoalMessage, &t.TaskTimeoutMessage, &lf); err != nil {
return nil, err
}
if lf.Valid {
t.LastFire = &lf.Time
}
return &t, nil
}
// CreateTrigger inserts a trigger for agentKey and returns it.
func (d *DB) CreateTrigger(t *AgentTrigger) (*AgentTrigger, error) {
row := d.QueryRow(`
INSERT INTO agent_triggers(agent_key, enabled, interval_sec, on_finding, on_goal_met, on_task_timeout, interval_message, finding_message, goal_message, task_timeout_message)
VALUES ($1,$2,$3,$4,$5,$6,$7,$8,$9,$10) RETURNING `+triggerCols,
t.AgentKey, t.Enabled, t.IntervalSec, t.OnFinding, t.OnGoalMet, t.OnTaskTimeout,
t.IntervalMessage, t.FindingMessage, t.GoalMessage, t.TaskTimeoutMessage)
return scanTrigger(row)
}
// UpdateTrigger updates a trigger's fields (not last_fire).
func (d *DB) UpdateTrigger(t *AgentTrigger) error {
_, err := d.Exec(`UPDATE agent_triggers SET enabled=$2, interval_sec=$3, on_finding=$4, on_goal_met=$5, on_task_timeout=$6, interval_message=$7, finding_message=$8, goal_message=$9, task_timeout_message=$10 WHERE id=$1`,
t.ID, t.Enabled, t.IntervalSec, t.OnFinding, t.OnGoalMet, t.OnTaskTimeout,
t.IntervalMessage, t.FindingMessage, t.GoalMessage, t.TaskTimeoutMessage)
return err
}
// DeleteTrigger removes a trigger.
func (d *DB) DeleteTrigger(id int64) error {
_, err := d.Exec(`DELETE FROM agent_triggers WHERE id=$1`, id)
return err
}
// ListTriggersFor returns an agent's triggers.
func (d *DB) ListTriggersFor(agentKey string) ([]*AgentTrigger, error) {
return d.queryTriggers(`SELECT `+triggerCols+` FROM agent_triggers WHERE agent_key=$1 ORDER BY id`, agentKey)
}
// ListEnabledTriggers returns all enabled triggers (for the scheduler).
func (d *DB) ListEnabledTriggers() ([]*AgentTrigger, error) {
return d.queryTriggers(`SELECT ` + triggerCols + ` FROM agent_triggers WHERE enabled ORDER BY id`)
}
func (d *DB) queryTriggers(q string, args ...any) ([]*AgentTrigger, error) {
rows, err := d.Query(q, args...)
if err != nil {
return nil, err
}
defer rows.Close()
out := []*AgentTrigger{}
for rows.Next() {
t, err := scanTrigger(rows)
if err != nil {
return nil, err
}
out = append(out, t)
}
return out, rows.Err()
}
// TouchTriggerFire records an interval trigger's fire time (now).
func (d *DB) TouchTriggerFire(id int64) error {
_, err := d.Exec(`UPDATE agent_triggers SET last_fire=now() WHERE id=$1`, id)
return err
}
// DeleteTriggersForAgent removes all triggers of an agent (custom agent delete).
func (d *DB) DeleteTriggersForAgent(agentKey string) error {
_, err := d.Exec(`DELETE FROM agent_triggers WHERE agent_key=$1`, agentKey)
return err
}
// ---------- scheduler_state (kv watermarks) ----------
func (d *DB) GetSchedState(key string) (string, error) {
var v string
err := d.QueryRow(`SELECT value FROM scheduler_state WHERE key=$1`, key).Scan(&v)
if err == sql.ErrNoRows {
return "", nil
}
return v, err
}
func (d *DB) SetSchedState(key, value string) error {
_, err := d.Exec(`INSERT INTO scheduler_state(key,value) VALUES ($1,$2)
ON CONFLICT (key) DO UPDATE SET value=EXCLUDED.value`, key, value)
return err
}
// ---------- event queries (cross-exploration, for the scheduler) ----------
// TaskEvent is a finding/goal event carrying the owning task's info, used to
// compose the trigger message context.
type TaskEvent struct {
NodeID int64 `json:"node_id"`
TaskID int64 `json:"task_id"`
TaskDesc string `json:"task_description"`
TaskGoal string `json:"task_goal"`
Summary string `json:"summary"` // finding summary / goal text
VulnClass string `json:"vulnclass"` // finding only
Severity string `json:"severity"` // finding only
}
// NewFindingsSince returns findings with node id > lastID across all live tasks,
// ordered by id (monotonic watermark → no double-fire).
func (d *DB) NewFindingsSince(lastID int64) ([]TaskEvent, error) {
rows, err := d.Query(`
SELECT n.id, t.id, t.description, t.goal, n.payload
FROM exploration_nodes n JOIN tasks t ON t.exploration_id = n.exploration_id
WHERE n.kind='finding' AND n.id > $1 AND t.deleted_at IS NULL
ORDER BY n.id`, lastID)
if err != nil {
return nil, err
}
defer rows.Close()
out := []TaskEvent{}
for rows.Next() {
var e TaskEvent
var payload []byte
if err := rows.Scan(&e.NodeID, &e.TaskID, &e.TaskDesc, &e.TaskGoal, &payload); err != nil {
return nil, err
}
var p struct{ Summary, Vulnclass, Severity string }
_ = json.Unmarshal(payload, &p)
e.Summary, e.VulnClass, e.Severity = p.Summary, p.Vulnclass, p.Severity
out = append(out, e)
}
return out, rows.Err()
}
// TimedOutTasksSince returns tasks that reached status='timeout' with id > lastID,
// ordered by id (monotonic watermark → no double-fire across restarts).
func (d *DB) TimedOutTasksSince(lastID int64) ([]TaskEvent, error) {
rows, err := d.Query(`
SELECT id, description, goal FROM tasks
WHERE status='timeout' AND deleted_at IS NULL AND id > $1
ORDER BY id`, lastID)
if err != nil {
return nil, err
}
defer rows.Close()
out := []TaskEvent{}
for rows.Next() {
var e TaskEvent
if err := rows.Scan(&e.NodeID, &e.TaskDesc, &e.TaskGoal); err != nil {
return nil, err
}
e.TaskID = e.NodeID // task id doubles as NodeID for the watermark
e.Summary = e.TaskGoal
out = append(out, e)
}
return out, rows.Err()
}
// MetGoals returns all met goals across live tasks (the scheduler filters out the
// ones it already fired for via the persisted fired-set).
func (d *DB) MetGoals() ([]TaskEvent, error) {
rows, err := d.Query(`
SELECT n.id, t.id, t.description, t.goal, n.payload
FROM exploration_nodes n JOIN tasks t ON t.exploration_id = n.exploration_id
WHERE n.kind='goal' AND n.state='met' AND t.deleted_at IS NULL
ORDER BY n.id`)
if err != nil {
return nil, err
}
defer rows.Close()
out := []TaskEvent{}
for rows.Next() {
var e TaskEvent
var payload []byte
if err := rows.Scan(&e.NodeID, &e.TaskID, &e.TaskDesc, &e.TaskGoal, &payload); err != nil {
return nil, err
}
var p struct{ Text, Summary string }
_ = json.Unmarshal(payload, &p)
if p.Text != "" {
e.Summary = p.Text
} else {
e.Summary = p.Summary
}
out = append(out, e)
}
return out, rows.Err()
}