Autumn-27andClaude Opus 4.8
3f798ec74e
镜像内全局预装 @playwright/mcp、@playwright/cli、playwright，并 `playwright install --with-deps chromium` 预置浏览器及系统依赖，运行时 不再联网下载。browser MCP 启动参数去掉 -y 与 @latest（改用已装好的本地版本）。 Co-Authored-By: Claude Opus 4.8 (1M context) <[email protected]>
420 lines
17 KiB
Go
420 lines
17 KiB
Go
// 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`,
|
||
`["@playwright/mcp","--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")
|
||
}
|