chore: ARTEX 初始化
ARTEX — AI 自主渗透测试系统(Go 后端 + Next.js 前端,单二进制内嵌)。 Co-Authored-By: Claude Opus 4.8 (1M context) <[email protected]>
This commit is contained in:
commit
0ef8561411
287 files changed
+70178
No files matched your search
@@ -0,0 +1,59 @@
|
||||
package agent
|
||||
|
||||
import (
|
||||
"context"
|
||||
|
||||
actool "github.com/Autumn-27/norma/tool"
|
||||
)
|
||||
|
||||
// DeferredInfo carries the deferred-tools wiring an agent needs to build its
|
||||
// Options: the MCP tool names whose schemas are withheld, the subset listed in the
|
||||
// global system-prompt block (non-skill-gated), and the shared session unlock set.
|
||||
// UnlockSkill unlocks a named skill's MCPs — hosts call it to rebuild the unlock set
|
||||
// from history on a resumed session (design doc C2).
|
||||
type DeferredInfo struct {
|
||||
Deferred []string // all MCP tool names (schema withheld)
|
||||
GlobalNames []string // MCP names to list in the system-prompt block
|
||||
Unlock *actool.UnlockSet // shared call-gate; nil when no MCP tools
|
||||
UnlockSkill func(skillName string)
|
||||
}
|
||||
|
||||
// ToolAugment, if set, returns the EXTRA tools an agent should see beyond its
|
||||
// built-in base set — the agent's visible skills (packed into one Skill meta-tool)
|
||||
// and visible MCP servers (expanded to mcp__server__tool). It also returns the
|
||||
// DeferredInfo describing how those MCP tools are deferred/gated. The server wires
|
||||
// it to the PG agent_visibility table. cleanup releases any spawned MCP clients.
|
||||
//
|
||||
// When nil, agents run with only their built-in tools — behavior is unchanged
|
||||
// until a user assigns a skill/MCP to the agent in the UI.
|
||||
var ToolAugment func(ctx context.Context, agentKey string) (extra []actool.CoreTool, def DeferredInfo, cleanup func())
|
||||
|
||||
// AugmentTools returns base plus the agent's visible skill/MCP tools, the
|
||||
// DeferredInfo, and a cleanup func the caller must defer (closes MCP clients).
|
||||
// Built-in base tools are kept as-is — never filtered (内置工具留代码层,不做可见性过滤).
|
||||
func AugmentTools(ctx context.Context, agentKey string, base []actool.CoreTool) ([]actool.CoreTool, DeferredInfo, func()) {
|
||||
var (
|
||||
def DeferredInfo
|
||||
cleanup = func() {}
|
||||
out = base
|
||||
)
|
||||
if ToolAugment != nil {
|
||||
var extra []actool.CoreTool
|
||||
var cl func()
|
||||
extra, def, cl = ToolAugment(ctx, agentKey)
|
||||
if cl != nil {
|
||||
cleanup = cl
|
||||
}
|
||||
if len(extra) > 0 {
|
||||
out = append(append([]actool.CoreTool{}, base...), extra...)
|
||||
}
|
||||
}
|
||||
// DB tools table has the final say on the built-in tools: drop the ones this
|
||||
// agent isn't bound to (or that are disabled) and swap in overridden
|
||||
// descriptions/schemas + default injection. MCP/skill/host tools have no row
|
||||
// and pass through untouched, so deferred/unlock wiring stays consistent.
|
||||
if ToolResolve != nil {
|
||||
out = ToolResolve(ctx, agentKey, out)
|
||||
}
|
||||
return out, def, cleanup
|
||||
}
|
||||
@@ -0,0 +1,178 @@
|
||||
package agent
|
||||
|
||||
import (
|
||||
"context"
|
||||
"strings"
|
||||
|
||||
"github.com/Autumn-27/norma/agentcore"
|
||||
"github.com/Autumn-27/norma/harness"
|
||||
"github.com/Autumn-27/norma/llm"
|
||||
"github.com/Autumn-27/artex/db"
|
||||
)
|
||||
|
||||
// captureRun drives one agent turn-to-completion over Session.Prompt and emits a
|
||||
// coalesced ActivityRecord per execution step (tool_use / tool_result / text /
|
||||
// thinking / result). It is shared by every LLM agent in the system (worker,
|
||||
// planner, …) so their execution is visible instead of a black box — the old
|
||||
// agentcore.Run discarded every event. The emitted records carry only
|
||||
// Kind/Tool/ToolUseID/IsError/Summary/Detail; the caller's emit fills in
|
||||
// IntentID/Worker. Returns the final assistant text + terminal error.
|
||||
//
|
||||
// KindText/KindThinking arrive as streaming deltas (one event per fragment); a
|
||||
// contiguous run is coalesced into a single record so the trace shows whole
|
||||
// messages, not dozens of fragments.
|
||||
func captureRun(ctx context.Context, opts agentcore.Options, input string, emit func(db.Activity)) (string, harness.TerminalReason, error) {
|
||||
s := agentcore.NewSession(opts)
|
||||
defer s.Close() // release the session's background-task manager (temp dir + processes)
|
||||
return captureRunSession(ctx, s, input, emit)
|
||||
}
|
||||
|
||||
// captureRunSession is captureRun over an existing session, so a caller can run
|
||||
// multiple prompts on the SAME conversation (e.g. a settlement round that reuses
|
||||
// the worker's accumulated context after the main run hit max_turns).
|
||||
func captureRunSession(ctx context.Context, s *agentcore.Session, input string, emit func(db.Activity)) (string, harness.TerminalReason, error) {
|
||||
var reason harness.TerminalReason
|
||||
rec := func(r db.Activity) {
|
||||
if emit != nil {
|
||||
emit(r)
|
||||
}
|
||||
}
|
||||
toolNames := map[string]string{} // tool_use id -> name, to label results
|
||||
|
||||
var tbuf strings.Builder
|
||||
var tkind string
|
||||
flush := func() {
|
||||
if tbuf.Len() == 0 {
|
||||
return
|
||||
}
|
||||
s := strings.TrimSpace(tbuf.String())
|
||||
k := tkind
|
||||
tbuf.Reset()
|
||||
tkind = ""
|
||||
if s != "" {
|
||||
rec(db.Activity{Kind: k, Summary: firstLine(s, 200), Detail: s})
|
||||
}
|
||||
}
|
||||
addDelta := func(kind, text string) {
|
||||
if text == "" {
|
||||
return
|
||||
}
|
||||
if tkind != "" && tkind != kind {
|
||||
flush()
|
||||
}
|
||||
tkind = kind
|
||||
tbuf.WriteString(text)
|
||||
}
|
||||
|
||||
var finalText string
|
||||
var rerr error
|
||||
for ev, err := range s.Prompt(ctx, input) {
|
||||
if err != nil {
|
||||
flush()
|
||||
if ctx.Err() != nil { // ctx cancelled = manual stop, not a failure
|
||||
rec(db.Activity{Kind: "result", Summary: "(已手动停止本次运行)"})
|
||||
return finalText, reason, ctx.Err()
|
||||
}
|
||||
rec(db.Activity{Kind: "result", IsError: true, Summary: "执行出错: " + err.Error(), Detail: err.Error()})
|
||||
return finalText, reason, err
|
||||
}
|
||||
switch ev.Kind {
|
||||
case harness.KindToolUse:
|
||||
if ev.ToolUse == nil {
|
||||
continue
|
||||
}
|
||||
flush()
|
||||
toolNames[ev.ToolUse.ID] = ev.ToolUse.Name
|
||||
in := string(ev.ToolUse.Input)
|
||||
rec(db.Activity{Kind: "tool_use", Tool: ev.ToolUse.Name, ToolUseID: ev.ToolUse.ID,
|
||||
Summary: ev.ToolUse.Name + " " + firstLine(in, 200), Detail: in})
|
||||
case harness.KindToolResult:
|
||||
if ev.ToolResult == nil {
|
||||
continue
|
||||
}
|
||||
flush()
|
||||
out := blocksText(ev.ToolResult.Content)
|
||||
rec(db.Activity{Kind: "tool_result", Tool: toolNames[ev.ToolResult.ToolUseID], ToolUseID: ev.ToolResult.ToolUseID,
|
||||
IsError: ev.ToolResult.IsError, Summary: firstLine(out, 200), Detail: out})
|
||||
case harness.KindText:
|
||||
addDelta("text", ev.Text)
|
||||
case harness.KindThinking:
|
||||
addDelta("thinking", ev.Text)
|
||||
case harness.KindUsage:
|
||||
// live cumulative token usage (per model turn). Emitted as a non-rendered
|
||||
// "usage" activity carrying only the token fields; the UI uses the latest
|
||||
// one for a running session's live token count. Don't flush() here — the
|
||||
// buffered final-answer text must stay for the KindResult de-dup.
|
||||
if ev.Usage != nil {
|
||||
u := *ev.Usage
|
||||
rec(db.Activity{Kind: "usage",
|
||||
InputTokens: &u.InputTokens, OutputTokens: &u.OutputTokens,
|
||||
CacheReadTokens: &u.CacheReadTokens, CacheWriteTokens: &u.CacheWriteTokens})
|
||||
}
|
||||
case harness.KindResult:
|
||||
if ev.Terminal != nil {
|
||||
finalText = ev.Terminal.Text
|
||||
reason = ev.Terminal.Reason
|
||||
// the buffered tail text usually equals Terminal.Text (final answer);
|
||||
// drop it to avoid a duplicate record, the result row carries it.
|
||||
if tkind == "text" && strings.TrimSpace(tbuf.String()) == strings.TrimSpace(ev.Terminal.Text) {
|
||||
tbuf.Reset()
|
||||
tkind = ""
|
||||
}
|
||||
flush() // flush any trailing thinking / non-final text
|
||||
sum := ev.Terminal.Text
|
||||
if sum == "" {
|
||||
switch ev.Terminal.Reason {
|
||||
case harness.ReasonAbortedTools, harness.ReasonAbortedStreaming:
|
||||
// interruption (task paused / worker killed / backend restart),
|
||||
// not a real completion — the intent is re-queued (open) or
|
||||
// stopped, so the session icon reflects that, not this run.
|
||||
sum = "(运行被中断:任务暂停 / 终止 / 重启所致;未完成)"
|
||||
case harness.ReasonTimeout, harness.ReasonMaxTurns:
|
||||
// budget hit → settlement ran (facts/assets were written back);
|
||||
// only the trailing text summary is missing, not a failure.
|
||||
sum = "(达运行预算上限,已收尾写回事实;本次无文字总结,终态: " + string(ev.Terminal.Reason) + ")"
|
||||
default:
|
||||
sum = "(无总结,终态: " + string(ev.Terminal.Reason) + ")"
|
||||
}
|
||||
}
|
||||
u := ev.Terminal.Usage // cumulative token usage for this session
|
||||
rec(db.Activity{Kind: "result", IsError: ev.Terminal.Err != nil,
|
||||
Summary: firstLine(sum, 400), Detail: sum,
|
||||
InputTokens: &u.InputTokens, OutputTokens: &u.OutputTokens,
|
||||
CacheReadTokens: &u.CacheReadTokens, CacheWriteTokens: &u.CacheWriteTokens})
|
||||
if ev.Terminal.Err != nil {
|
||||
rerr = ev.Terminal.Err
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
flush() // safety: any unflushed text if the stream ended without KindResult
|
||||
return finalText, reason, rerr
|
||||
}
|
||||
|
||||
// blocksText concatenates the text of a tool-result's content blocks.
|
||||
func blocksText(blocks []llm.ContentBlock) string {
|
||||
var b strings.Builder
|
||||
for _, bl := range blocks {
|
||||
if bl.Type == llm.BlockText && bl.Text != "" {
|
||||
if b.Len() > 0 {
|
||||
b.WriteByte('\n')
|
||||
}
|
||||
b.WriteString(bl.Text)
|
||||
}
|
||||
}
|
||||
return b.String()
|
||||
}
|
||||
|
||||
// firstLine returns a single-line, length-capped preview for the summary column.
|
||||
func firstLine(s string, max int) string {
|
||||
s = strings.TrimSpace(s)
|
||||
if i := strings.IndexByte(s, '\n'); i >= 0 {
|
||||
s = s[:i]
|
||||
}
|
||||
if len(s) > max {
|
||||
s = s[:max] + "…"
|
||||
}
|
||||
return s
|
||||
}
|
||||
+152
@@ -0,0 +1,152 @@
|
||||
package agent
|
||||
|
||||
import (
|
||||
"context"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"time"
|
||||
|
||||
"github.com/Autumn-27/norma/agentcore"
|
||||
"github.com/Autumn-27/norma/llm"
|
||||
"github.com/Autumn-27/norma/permission"
|
||||
actool "github.com/Autumn-27/norma/tool"
|
||||
"github.com/Autumn-27/norma/transcript"
|
||||
"github.com/Autumn-27/artex/db"
|
||||
"github.com/Autumn-27/artex/guard"
|
||||
)
|
||||
|
||||
// ChatAgent is the generic, task-independent conversational runner behind the chat
|
||||
// page. It generalizes MainAgent.Chat: any agent (built-in OR a custom one, by
|
||||
// key) can be chatted with, multi-turn history resumed from the transcript. It is
|
||||
// a PURE ASSISTANT — base tools are the SDK DefaultTools (Bash/Read/Write/Edit/
|
||||
// LS/Glob/Grep) plus whatever skills/MCP the key is made visible; NO pentest
|
||||
// graph/task context is injected (that stays exclusive to MainAgent).
|
||||
type ChatAgent struct {
|
||||
prov llm.Provider
|
||||
model string
|
||||
workDir string
|
||||
tx *transcript.Store
|
||||
window int
|
||||
proxyAddr string
|
||||
proxyCACert string
|
||||
webSearch WebSearchOpts
|
||||
guard *guard.Guard // optional; nil disables intercept hooks for chat
|
||||
}
|
||||
|
||||
func NewChatAgent(prov llm.Provider, model, workDir string, tx *transcript.Store, window int) *ChatAgent {
|
||||
return &ChatAgent{prov: prov, model: model, workDir: workDir, tx: tx, window: window}
|
||||
}
|
||||
|
||||
// SetProxy points the chat agent's WebFetch/Bash at the recording proxy plus the
|
||||
// CA cert it trusts (empty addr = direct). Kept for parity with the other agents.
|
||||
func (c *ChatAgent) SetProxy(addr, caCert string) { c.proxyAddr, c.proxyCACert = addr, caCert }
|
||||
|
||||
// SetWebSearch selects the web_search backend for the chat agent (off by default).
|
||||
func (c *ChatAgent) SetWebSearch(o WebSearchOpts) { c.webSearch = o }
|
||||
|
||||
// SetGuard attaches a guard (with user-configured intercept rules) to this chat
|
||||
// agent. Must be called before Chat; safe to call multiple times.
|
||||
func (c *ChatAgent) SetGuard(g *guard.Guard) { c.guard = g }
|
||||
|
||||
// chatWorkDirSpec returns a working-directory notice appended to every chat
|
||||
// agent's system prompt. Mirrors artifactSpec but without pentest-specific
|
||||
// wording ("payload", "抓响应体") that would be odd in a general assistant.
|
||||
func chatWorkDirSpec(workDir string) string {
|
||||
return "\n\n**文件输出规约**:需要写文件时,一律写到工作目录 " + workDir + "(这是默认 CWD,相对路径即落在这里,也可用该绝对路径)——不要写 /tmp 或其他绝对路径。"
|
||||
}
|
||||
|
||||
// chatSystem renders the DB-managed prompt body for agentKey. Custom agents have
|
||||
// no per-key in-code default, so DefaultAssistantPrompt is the render fallback.
|
||||
func chatSystem(agentKey, workDir string) string {
|
||||
return renderSystem(agentKey, DefaultAssistantPrompt, chatVars{}) + chatWorkDirSpec(workDir)
|
||||
}
|
||||
|
||||
// chatVars carries no runtime variables today — a custom prompt referencing any
|
||||
// {{.X}} simply fails to render and falls back to DefaultAssistantPrompt.
|
||||
type chatVars struct{}
|
||||
|
||||
// Chat runs ONE turn of a conversation with the agent identified by agentKey,
|
||||
// resuming prior history keyed by sessionID. maxTurns is the per-turn agent step
|
||||
// budget (0 = unlimited). maxDuration is the wall-clock run budget per turn
|
||||
// (0 = unlimited); the timer resets each time Chat is called, so a new user
|
||||
// message always starts a fresh countdown. webSearch gates network search for
|
||||
// THIS agent (the global backend/key still come from the chat agent's config,
|
||||
// but each agent decides on/off). emit receives each execution step (thinking /
|
||||
// tool_use / tool_result / text / result), tagged with the agent key as the
|
||||
// worker lane.
|
||||
func (c *ChatAgent) Chat(ctx context.Context, agentKey, sessionID, message string, maxTurns int, maxDuration time.Duration, webSearch bool, emit func(db.Activity)) (string, error) {
|
||||
// gate the global web-search opts by this agent's own flag.
|
||||
ws := c.webSearch
|
||||
if !webSearch {
|
||||
ws.Enabled = false
|
||||
}
|
||||
|
||||
// Per-session working directory: <workDir>/sessions/<sessionID>/
|
||||
// Isolates file writes across conversations, mirroring how workers use i<intentID>/.
|
||||
sessionWorkDir := filepath.Join(c.workDir, "sessions", sessionID)
|
||||
_ = os.MkdirAll(sessionWorkDir, 0o755)
|
||||
|
||||
// Pure assistant: DefaultTools as the base; AugmentTools layers in the key's
|
||||
// visible skills/MCP and lets the DB tools table filter/override. DefaultTools
|
||||
// have no tools-table rows, so they always pass through.
|
||||
base := actool.DefaultTools()
|
||||
tools, def, cleanup := AugmentTools(ctx, agentKey, base)
|
||||
defer cleanup()
|
||||
|
||||
system, boundary := deferredSystem(chatSystem(agentKey, sessionWorkDir), def)
|
||||
opts := agentcore.Options{
|
||||
Provider: c.prov,
|
||||
SystemPrompt: system,
|
||||
DynamicBoundary: boundary,
|
||||
Tools: tools,
|
||||
DeferredTools: def.Deferred,
|
||||
UnlockSet: def.Unlock,
|
||||
PermissionMode: permission.ModeBypass,
|
||||
EnableWebFetch: true, // 走记录代理留痕;载入代理 CA 验证 MITM 重签的 HTTPS 证书
|
||||
WebFetchProxy: c.proxyAddr,
|
||||
WebFetchCACert: c.proxyCACert,
|
||||
// 联网搜索(可选)。ddgs 无需 key;brave-free 需 BraveKey;tavily 需 TavilyKey。
|
||||
// WebSearchProxy 是独立出口代理(http/https/socks5),与记录流量的 MITM 代理无关;空则直连。
|
||||
EnableWebSearch: ws.Enabled,
|
||||
WebSearchBackend: ws.Backend,
|
||||
BraveSearchAPIKey: ws.BraveKey,
|
||||
TavilySearchAPIKey: ws.TavilyKey,
|
||||
WebSearchProxy: ws.Proxy,
|
||||
BashEnv: proxyEnv(c.proxyAddr, c.proxyCACert), // Bash 子命令默认走代理+信任 CA
|
||||
WorkingDir: sessionWorkDir,
|
||||
MaxTurns: maxTurns,
|
||||
MaxDuration: maxDuration,
|
||||
Compaction: compactionConfig(c.window),
|
||||
Todos: actool.NewTodoStore(),
|
||||
// large tool output spills to cmd-output/ under the session dir.
|
||||
// 截断上限用 SDK 默认(tool.Capture 的 30000 字符)。
|
||||
ToolOutputDir: filepath.Join(sessionWorkDir, "cmd-output"),
|
||||
// 命中预算(步数)→ SDK 跑收尾:输出一句总结。Prompt 与收尾轮数按本 agent key 后台可编辑
|
||||
// (自定义 agent 各自一份;留空/0 用通用默认:5 轮)。
|
||||
Settlement: wrapupSettlement(agentKey, nil),
|
||||
}
|
||||
if c.guard != nil {
|
||||
opts.Hooks = c.guard.Hooks()
|
||||
}
|
||||
if c.tx != nil { // persist raw human↔AI conversation; one accumulating file per thread
|
||||
opts.Transcript = c.tx
|
||||
opts.SessionID = sessionID
|
||||
}
|
||||
s := agentcore.NewSession(opts)
|
||||
defer s.Close()
|
||||
// reload prior conversation so the agent has context across turns (each Chat is
|
||||
// a fresh session). First turn: no file yet → Resume loads nothing and proceeds.
|
||||
if c.tx != nil {
|
||||
_ = s.Resume(sessionID)
|
||||
}
|
||||
// re-unlock skill-gated MCPs from prior Skill() calls in the reloaded history so
|
||||
// revealed tools stay callable across the fresh session.
|
||||
seedUnlockFromHistory(s.Messages(), def.UnlockSkill)
|
||||
text, _, err := captureRunSession(ctx, s, message, func(r db.Activity) {
|
||||
if emit != nil {
|
||||
r.Worker = agentKey
|
||||
emit(r)
|
||||
}
|
||||
})
|
||||
return text, err
|
||||
}
|
||||
@@ -0,0 +1,50 @@
|
||||
package agent
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
|
||||
"github.com/Autumn-27/norma/llm"
|
||||
actool "github.com/Autumn-27/norma/tool"
|
||||
)
|
||||
|
||||
// deferredSystem builds an agent's system-prompt segments and cache boundary from
|
||||
// its DeferredInfo. When globally-available MCP tools are present, their names + a
|
||||
// "prefer core tools" instruction render into a <available-deferred-tools> block
|
||||
// placed as the LAST system-prompt segment, with DynamicBoundary set so the whole
|
||||
// (session-fixed) system prompt — including the block — is cached (design doc
|
||||
// §2.1 / C1). Skill-gated MCP names are NOT in this block; they surface when their
|
||||
// skill loads. Returns a single plain segment + boundary 0 when there is no global
|
||||
// block to add.
|
||||
func deferredSystem(sysText string, def DeferredInfo) (system []string, boundary int) {
|
||||
block := actool.RenderDeferredToolsBlock(def.GlobalNames)
|
||||
if block == "" {
|
||||
return []string{sysText}, 0
|
||||
}
|
||||
system = []string{sysText, block}
|
||||
boundary = len(system) // b >= len → whole system prompt cached (SDK guard)
|
||||
return system, boundary
|
||||
}
|
||||
|
||||
// seedUnlockFromHistory replays prior Skill() invocations in the conversation so
|
||||
// their skill-gated MCPs are re-unlocked on a resumed session (design doc C2). The
|
||||
// main agent builds a fresh session each turn; its in-memory unlock set would
|
||||
// otherwise reset, leaving the model able to see a skill-revealed tool name yet
|
||||
// unable to call it. No-op when unlockSkill is nil (no deferred tools).
|
||||
func seedUnlockFromHistory(msgs []llm.Message, unlockSkill func(string)) {
|
||||
if unlockSkill == nil {
|
||||
return
|
||||
}
|
||||
for _, m := range msgs {
|
||||
for _, b := range m.ToolUses() {
|
||||
if b.Name != "Skill" {
|
||||
continue
|
||||
}
|
||||
var in struct {
|
||||
Name string `json:"name"`
|
||||
}
|
||||
if json.Unmarshal(b.Input, &in) == nil && in.Name != "" {
|
||||
unlockSkill(in.Name)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,89 @@
|
||||
package agent
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"github.com/Autumn-27/norma/llm"
|
||||
actool "github.com/Autumn-27/norma/tool"
|
||||
)
|
||||
|
||||
func TestDeferredSystem_NoGlobal(t *testing.T) {
|
||||
// No global MCP names → plain single-segment system, no cache boundary.
|
||||
sys, boundary := deferredSystem("SYS", DeferredInfo{})
|
||||
if len(sys) != 1 || sys[0] != "SYS" || boundary != 0 {
|
||||
t.Fatalf("expected [SYS],0 — got %v,%d", sys, boundary)
|
||||
}
|
||||
}
|
||||
|
||||
func TestDeferredSystem_WithGlobal(t *testing.T) {
|
||||
def := DeferredInfo{
|
||||
Deferred: []string{"mcp__browser__navigate", "mcp__browser__click"},
|
||||
GlobalNames: []string{"mcp__browser__navigate", "mcp__browser__click"},
|
||||
}
|
||||
sys, boundary := deferredSystem("SYS", def)
|
||||
if len(sys) != 2 || sys[0] != "SYS" {
|
||||
t.Fatalf("expected [SYS, block], got %v", sys)
|
||||
}
|
||||
if !strings.Contains(sys[1], "<available-deferred-tools>") ||
|
||||
!strings.Contains(sys[1], "mcp__browser__navigate") {
|
||||
t.Fatalf("block missing names:\n%s", sys[1])
|
||||
}
|
||||
if boundary != len(sys) {
|
||||
t.Fatalf("boundary=%d want %d (whole prompt cached)", boundary, len(sys))
|
||||
}
|
||||
}
|
||||
|
||||
func TestDeferredSystem_GatedNotInBlock(t *testing.T) {
|
||||
// A skill-gated server's tools are deferred but NOT in the global block.
|
||||
def := DeferredInfo{
|
||||
Deferred: []string{"mcp__browser__navigate", "mcp__secret__do"},
|
||||
GlobalNames: []string{"mcp__browser__navigate"}, // secret gated → excluded
|
||||
}
|
||||
sys, _ := deferredSystem("SYS", def)
|
||||
if strings.Contains(sys[1], "mcp__secret__do") {
|
||||
t.Fatalf("gated tool must not appear in global block:\n%s", sys[1])
|
||||
}
|
||||
if !strings.Contains(sys[1], "mcp__browser__navigate") {
|
||||
t.Fatal("global tool should appear in block")
|
||||
}
|
||||
}
|
||||
|
||||
func TestSeedUnlockFromHistory(t *testing.T) {
|
||||
skillCall := func(name string) llm.ContentBlock {
|
||||
return llm.ContentBlock{Type: llm.BlockToolUse, Name: "Skill", Input: json.RawMessage(`{"name":"` + name + `"}`)}
|
||||
}
|
||||
msgs := []llm.Message{
|
||||
{Role: llm.RoleAssistant, Content: []llm.ContentBlock{skillCall("browsing")}},
|
||||
{Role: llm.RoleAssistant, Content: []llm.ContentBlock{
|
||||
{Type: llm.BlockToolUse, Name: "Bash", Input: json.RawMessage(`{}`)}, // ignored
|
||||
skillCall("recon"),
|
||||
}},
|
||||
}
|
||||
var got []string
|
||||
seedUnlockFromHistory(msgs, func(name string) { got = append(got, name) })
|
||||
if strings.Join(got, ",") != "browsing,recon" {
|
||||
t.Fatalf("unlocked=%v want [browsing recon]", got)
|
||||
}
|
||||
seedUnlockFromHistory(msgs, nil) // nil → no-op, no panic
|
||||
}
|
||||
|
||||
// TestUnlockGatingFlow mirrors what OnInvoke / seedUnlockFromHistory do: a gated
|
||||
// tool starts locked and becomes callable only after its skill unlocks it.
|
||||
func TestUnlockGatingFlow(t *testing.T) {
|
||||
serverTools := map[string][]string{"secret": {"mcp__secret__do"}}
|
||||
unlock := actool.NewUnlockSet("mcp__browser__navigate") // global only
|
||||
unlockSkill := func(name string) {
|
||||
if name == "unlock-secret" {
|
||||
unlock.Add(serverTools["secret"]...)
|
||||
}
|
||||
}
|
||||
if unlock.Has("mcp__secret__do") {
|
||||
t.Fatal("gated tool should start locked")
|
||||
}
|
||||
unlockSkill("unlock-secret")
|
||||
if !unlock.Has("mcp__secret__do") {
|
||||
t.Fatal("gated tool should be unlocked after skill load")
|
||||
}
|
||||
}
|
||||
+106
@@ -0,0 +1,106 @@
|
||||
package agent
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"strings"
|
||||
|
||||
"github.com/Autumn-27/norma/agentcore"
|
||||
acperm "github.com/Autumn-27/norma/permission"
|
||||
actool "github.com/Autumn-27/norma/tool"
|
||||
"github.com/Autumn-27/artex/db"
|
||||
)
|
||||
|
||||
// goalsDefaultTmpl is the built-in EDITABLE body (段 [A]) of the goals-decomposer
|
||||
// prompt, seeded into agent_prompts. No template vars are used today.
|
||||
const goalsDefaultTmpl = `你是渗透测试目标分解器。你的职责是从用户输入中识别出**最终要达成的结果**,而不是规划攻击步骤。
|
||||
|
||||
**目标 = 最终可交付/可核验的结果**
|
||||
|
||||
**不是目标的内容(禁止列为子目标)**:
|
||||
- 信息收集、侦察、端点扫描
|
||||
- 漏洞分析与验证过程
|
||||
- 攻击步骤、利用手段
|
||||
- 结果验证步骤
|
||||
|
||||
**拆分原则**:
|
||||
- 用户描述的最终目标只有一个 → 输出一个
|
||||
- 存在多个**相互独立**的最终交付物 → 分别列出
|
||||
- 能对应明确漏洞类的标注 vulnclass;信息收集/业务逻辑类目标留空
|
||||
- 严禁臆造用户未提及的目标
|
||||
|
||||
调用 set_goals 提交结果。`
|
||||
|
||||
// GoalSpec is one decomposed objective.
|
||||
type GoalSpec struct {
|
||||
Text string `json:"text"`
|
||||
VulnClass string `json:"vulnclass,omitempty"`
|
||||
}
|
||||
|
||||
// DecomposeGoals asks the LLM to break a pentest task goal into discrete,
|
||||
// independently-verifiable objectives (each becomes a goal node). Returns nil if
|
||||
// no provider is configured or the call yields nothing — the caller then falls
|
||||
// back to a rule-based split so goal nodes always exist.
|
||||
//
|
||||
// desc is the task's free-text description (背景:靶标范围/flag 数量/交战说明等).
|
||||
// It is fed alongside the goal so the decomposer no longer splits blind — the
|
||||
// prompt still forbids inventing anything the two texts don't state.
|
||||
//
|
||||
// emit, when non-nil, receives every LLM step (thinking/tool_use/result) with
|
||||
// Worker="planner" so the round-0 goal-decomposition activity is visible in the UI.
|
||||
func DecomposeGoals(ctx context.Context, c Config, goalText, desc string, emit func(db.Activity)) []GoalSpec {
|
||||
prov, err := c.NewProvider()
|
||||
if err != nil {
|
||||
return nil
|
||||
}
|
||||
var out struct {
|
||||
Goals []GoalSpec `json:"goals"`
|
||||
}
|
||||
captured := false
|
||||
submitTool := actool.Build(actool.Spec{
|
||||
Name: "set_goals",
|
||||
Description: "提交拆分后的目标列表。",
|
||||
Schema: map[string]any{"type": "object", "properties": map[string]any{
|
||||
"goals": map[string]any{"type": "array", "items": map[string]any{
|
||||
"type": "object", "properties": map[string]any{
|
||||
"text": map[string]any{"type": "string", "description": "一个独立可验证的目标"},
|
||||
"vulnclass": map[string]any{"type": "string", "description": "对应漏洞类(若明确),如 SQLi/IDOR;业务逻辑目标可留空"},
|
||||
}, "required": []string{"text"},
|
||||
}},
|
||||
}, "required": []string{"goals"}},
|
||||
Permissions: func(context.Context, json.RawMessage, acperm.Context) acperm.Decision { return acperm.Allowed() },
|
||||
Run: func(_ context.Context, in json.RawMessage, _ *actool.ToolContext) (actool.Result, error) {
|
||||
_ = json.Unmarshal(in, &out)
|
||||
captured = true
|
||||
return actool.Text("ok"), nil
|
||||
},
|
||||
})
|
||||
// Description rides in the user message (same channel as the goal), NOT via the
|
||||
// {{.EngagementDescription}} template var — else a prompt that references the var
|
||||
// would inject the description twice. System prompt stays pure static instructions.
|
||||
sys := renderSystem("goals", goalsDefaultTmpl, GoalsVars{})
|
||||
userMsg := "任务目标:\n" + goalText
|
||||
if d := strings.TrimSpace(desc); d != "" {
|
||||
userMsg += "\n\n任务描述(背景信息,可能含靶标范围/flag 数量/交战说明;仅供参考,不要臆造其中未提及的内容):\n" + d
|
||||
}
|
||||
// Use captureRun so every LLM step is emitted as an activity record (visible in
|
||||
// the plan tab under the round-0 marker). Falls back gracefully when emit is nil.
|
||||
captureEmit := func(r db.Activity) {
|
||||
if emit != nil {
|
||||
r.Worker = "planner"
|
||||
emit(r)
|
||||
}
|
||||
}
|
||||
captureRun(ctx, agentcore.Options{
|
||||
Provider: prov,
|
||||
SystemPrompt: []string{sys},
|
||||
Tools: []actool.CoreTool{submitTool},
|
||||
PermissionMode: acperm.ModeBypass,
|
||||
DisableBackgroundTasks: true,
|
||||
MaxTurns: 4,
|
||||
}, userMsg, captureEmit)
|
||||
if !captured {
|
||||
return nil
|
||||
}
|
||||
return out.Goals
|
||||
}
|
||||
@@ -0,0 +1,396 @@
|
||||
package agent
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"testing"
|
||||
|
||||
"github.com/Autumn-27/artex/db"
|
||||
)
|
||||
|
||||
// testDB opens a DB connection, skipping if PG is unavailable.
|
||||
func testDB(t *testing.T) *db.DB {
|
||||
t.Helper()
|
||||
dsn, _, err := db.DSN()
|
||||
if err != nil {
|
||||
t.Skipf("no database config (%v)", err)
|
||||
}
|
||||
d, err := db.Open(dsn)
|
||||
if err != nil {
|
||||
t.Skipf("postgres unavailable (%v)", err)
|
||||
}
|
||||
return d
|
||||
}
|
||||
|
||||
// callInsertAssets calls the insert_assets tool with the given payload.
|
||||
func callInsertAssets(t *testing.T, ts *ToolSet, payload any) map[string]any {
|
||||
t.Helper()
|
||||
raw, _ := json.Marshal(payload)
|
||||
tool := ts.insertAssets()
|
||||
res, err := tool.Call(context.Background(), raw, nil)
|
||||
if err != nil {
|
||||
t.Fatalf("insertAssets Call error: %v", err)
|
||||
}
|
||||
text := res.Flatten()
|
||||
var out map[string]any
|
||||
if err := json.Unmarshal([]byte(text), &out); err != nil {
|
||||
t.Fatalf("unmarshal result: %v\nraw: %s", err, text)
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// =====================================================================
|
||||
// TestInsertAssetsSubdomainSideEffects
|
||||
// 子域名插入 → 自动创建 root_domain + IP 资产,IP 绑定域名
|
||||
// =====================================================================
|
||||
func TestInsertAssetsSubdomainSideEffects(t *testing.T) {
|
||||
d := testDB(t)
|
||||
defer d.Close()
|
||||
|
||||
ts := NewToolSet(nil, "")
|
||||
ts.SetAssetStore(d.Assets(), d.Companies())
|
||||
defer d.Exec(`DELETE FROM assets WHERE domain IN ('ia-sub.sideeffect-test.com','sideeffect-test.com') OR ip='7.8.9.10'`)
|
||||
|
||||
out := callInsertAssets(t, ts, map[string]any{
|
||||
"assets": []any{
|
||||
map[string]any{
|
||||
"type": "subdomain",
|
||||
"domain": "ia-sub.sideeffect-test.com",
|
||||
"record_type": "A",
|
||||
"record_value": []string{"7.8.9.10"},
|
||||
},
|
||||
},
|
||||
"task_id": 999,
|
||||
})
|
||||
|
||||
// no errors
|
||||
if errs, _ := out["errors"].([]any); len(errs) > 0 {
|
||||
t.Errorf("unexpected errors: %v", errs)
|
||||
}
|
||||
results, _ := out["results"].([]any)
|
||||
if len(results) == 0 {
|
||||
t.Fatal("no results returned")
|
||||
}
|
||||
|
||||
// root_domain should exist
|
||||
var rootCnt int
|
||||
d.QueryRow(`SELECT COUNT(*) FROM assets WHERE type='root_domain' AND domain='sideeffect-test.com'`).Scan(&rootCnt)
|
||||
if rootCnt != 1 {
|
||||
t.Errorf("side-effect: root_domain not created, got %d", rootCnt)
|
||||
}
|
||||
|
||||
// IP asset should exist with bound_domains containing our subdomain
|
||||
var ipID int64
|
||||
var boundDomains []byte
|
||||
d.QueryRow(`SELECT id, array_to_json(bound_domains)::text FROM assets WHERE type='ip' AND ip='7.8.9.10'`).Scan(&ipID, &boundDomains)
|
||||
if ipID == 0 {
|
||||
t.Error("side-effect: IP asset not created")
|
||||
}
|
||||
var domains []string
|
||||
json.Unmarshal(boundDomains, &domains)
|
||||
found := false
|
||||
for _, d := range domains {
|
||||
if d == "ia-sub.sideeffect-test.com" {
|
||||
found = true
|
||||
}
|
||||
}
|
||||
if !found {
|
||||
t.Errorf("side-effect: bound_domains should contain subdomain, got %v", domains)
|
||||
}
|
||||
|
||||
// record_value stored as array
|
||||
var rvRaw []byte
|
||||
d.QueryRow(`SELECT array_to_json(record_value)::text FROM assets WHERE type='subdomain' AND domain='ia-sub.sideeffect-test.com'`).Scan(&rvRaw)
|
||||
var rv []string
|
||||
json.Unmarshal(rvRaw, &rv)
|
||||
if len(rv) == 0 || rv[0] != "7.8.9.10" {
|
||||
t.Errorf("record_value stored incorrectly: %v", rv)
|
||||
}
|
||||
}
|
||||
|
||||
// =====================================================================
|
||||
// TestInsertAssetsMultiIPSubdomain
|
||||
// 多个 IP 的子域名:所有 IP 都应存入 record_value[],各自创建 IP 资产
|
||||
// =====================================================================
|
||||
func TestInsertAssetsMultiIPSubdomain(t *testing.T) {
|
||||
d := testDB(t)
|
||||
defer d.Close()
|
||||
|
||||
ts := NewToolSet(nil, "")
|
||||
ts.SetAssetStore(d.Assets(), d.Companies())
|
||||
defer d.Exec(`DELETE FROM assets WHERE domain IN ('multi.multiip-test.io','multiip-test.io') OR ip IN ('1.1.1.1','2.2.2.2')`)
|
||||
|
||||
out := callInsertAssets(t, ts, map[string]any{
|
||||
"assets": []any{
|
||||
map[string]any{
|
||||
"type": "subdomain",
|
||||
"domain": "multi.multiip-test.io",
|
||||
"record_type": "A",
|
||||
"record_value": []string{"1.1.1.1", "2.2.2.2"},
|
||||
},
|
||||
},
|
||||
})
|
||||
|
||||
if errs, _ := out["errors"].([]any); len(errs) > 0 {
|
||||
t.Errorf("unexpected errors: %v", errs)
|
||||
}
|
||||
|
||||
// Both IPs should have IP assets
|
||||
var ip1Cnt, ip2Cnt int
|
||||
d.QueryRow(`SELECT COUNT(*) FROM assets WHERE type='ip' AND ip='1.1.1.1'`).Scan(&ip1Cnt)
|
||||
d.QueryRow(`SELECT COUNT(*) FROM assets WHERE type='ip' AND ip='2.2.2.2'`).Scan(&ip2Cnt)
|
||||
if ip1Cnt != 1 {
|
||||
t.Error("IP 1.1.1.1 asset not created")
|
||||
}
|
||||
if ip2Cnt != 1 {
|
||||
t.Error("IP 2.2.2.2 asset not created")
|
||||
}
|
||||
|
||||
// record_value should contain both IPs
|
||||
var rvRaw []byte
|
||||
d.QueryRow(`SELECT array_to_json(record_value)::text FROM assets WHERE type='subdomain' AND domain='multi.multiip-test.io'`).Scan(&rvRaw)
|
||||
var rv []string
|
||||
json.Unmarshal(rvRaw, &rv)
|
||||
if len(rv) != 2 {
|
||||
t.Errorf("record_value: want 2 IPs, got %v", rv)
|
||||
}
|
||||
}
|
||||
|
||||
// =====================================================================
|
||||
// TestInsertAssetsHTTPServiceTechnologies
|
||||
// HTTP 服务插入:technologies 存储并可读回;IP 存在时域名和端口写入 IP 资产
|
||||
// =====================================================================
|
||||
func TestInsertAssetsHTTPServiceTechnologies(t *testing.T) {
|
||||
d := testDB(t)
|
||||
defer d.Close()
|
||||
|
||||
ts := NewToolSet(nil, "")
|
||||
ts.SetAssetStore(d.Assets(), d.Companies())
|
||||
defer d.Exec(`DELETE FROM assets WHERE url='https://tech-test.example.com' OR domain IN ('tech-test.example.com','example.com') OR ip='3.4.5.6'`)
|
||||
|
||||
out := callInsertAssets(t, ts, map[string]any{
|
||||
"assets": []any{
|
||||
map[string]any{
|
||||
"type": "service",
|
||||
"url": "https://tech-test.example.com",
|
||||
"service_ip": "3.4.5.6",
|
||||
"technologies": []string{"Nginx", "Vue.js", "Cloudflare"},
|
||||
"status_code": 200,
|
||||
"page_title": "Tech Test Site",
|
||||
},
|
||||
},
|
||||
"task_id": 888,
|
||||
})
|
||||
|
||||
if errs, _ := out["errors"].([]any); len(errs) > 0 {
|
||||
t.Errorf("unexpected errors: %v", errs)
|
||||
}
|
||||
|
||||
// technologies should be stored
|
||||
var techCnt int
|
||||
d.QueryRow(`SELECT array_length(technologies,1) FROM assets WHERE url='https://tech-test.example.com'`).Scan(&techCnt)
|
||||
if techCnt != 3 {
|
||||
t.Errorf("technologies: want 3, got %d", techCnt)
|
||||
}
|
||||
|
||||
// QueryByType should return technologies correctly (verifies array_to_json scan)
|
||||
assets, err := d.Assets().QueryByType("service", 50, 0)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
var found *db.Asset
|
||||
for _, a := range assets {
|
||||
if a.URL == "https://tech-test.example.com" {
|
||||
found = a
|
||||
break
|
||||
}
|
||||
}
|
||||
if found == nil {
|
||||
t.Fatal("service not found via QueryByType")
|
||||
}
|
||||
if len(found.Technologies) != 3 {
|
||||
t.Errorf("QueryByType: technologies roundtrip failed, got %v", found.Technologies)
|
||||
}
|
||||
|
||||
// side effect: IP asset should exist with bound_domains containing the service domain
|
||||
var ipID int64
|
||||
var bdRaw []byte
|
||||
var portCnt int
|
||||
d.QueryRow(`SELECT id, array_to_json(bound_domains)::text FROM assets WHERE type='ip' AND ip='3.4.5.6'`).Scan(&ipID, &bdRaw)
|
||||
if ipID == 0 {
|
||||
t.Error("side-effect: IP asset not created for HTTP service IP")
|
||||
}
|
||||
var bd []string
|
||||
json.Unmarshal(bdRaw, &bd)
|
||||
hasDomain := false
|
||||
for _, dom := range bd {
|
||||
if dom == "tech-test.example.com" {
|
||||
hasDomain = true
|
||||
}
|
||||
}
|
||||
if !hasDomain {
|
||||
t.Errorf("side-effect: IP bound_domains missing service domain, got %v", bd)
|
||||
}
|
||||
|
||||
// side effect: IP open_ports should contain port 443
|
||||
d.QueryRow(`SELECT cardinality(open_ports) FROM assets WHERE type='ip' AND ip='3.4.5.6'`).Scan(&portCnt)
|
||||
if portCnt == 0 {
|
||||
t.Error("side-effect: IP open_ports not set for HTTP service")
|
||||
}
|
||||
}
|
||||
|
||||
// =====================================================================
|
||||
// TestInsertAssetsOtherService
|
||||
// 非 HTTP 服务:c_segment 自动生成,IP 资产含 open_ports 和 bound_domains
|
||||
// =====================================================================
|
||||
func TestInsertAssetsOtherService(t *testing.T) {
|
||||
d := testDB(t)
|
||||
defer d.Close()
|
||||
|
||||
ts := NewToolSet(nil, "")
|
||||
ts.SetAssetStore(d.Assets(), d.Companies())
|
||||
defer d.Exec(`DELETE FROM assets WHERE
|
||||
(type='service' AND ip='10.20.30.40') OR
|
||||
(type='ip' AND ip='10.20.30.40') OR
|
||||
domain IN ('db.othersvc-test.com','othersvc-test.com')`)
|
||||
|
||||
out := callInsertAssets(t, ts, map[string]any{
|
||||
"assets": []any{
|
||||
map[string]any{
|
||||
"type": "service",
|
||||
"ip": "10.20.30.40",
|
||||
"domain": "db.othersvc-test.com",
|
||||
"port": 3306,
|
||||
"service_name": "mysql",
|
||||
},
|
||||
},
|
||||
})
|
||||
|
||||
if errs, _ := out["errors"].([]any); len(errs) > 0 {
|
||||
t.Errorf("unexpected errors: %v", errs)
|
||||
}
|
||||
|
||||
// c_segment should be auto-set on the service
|
||||
var cseg *string
|
||||
d.QueryRow(`SELECT c_segment::text FROM assets WHERE type='service' AND ip='10.20.30.40'`).Scan(&cseg)
|
||||
if cseg == nil || *cseg != "10.20.30.0/24" {
|
||||
t.Errorf("c_segment: want 10.20.30.0/24, got %v", cseg)
|
||||
}
|
||||
|
||||
// IP side-effect: port 3306 in open_ports
|
||||
var portCnt int
|
||||
d.QueryRow(`SELECT cardinality(open_ports) FROM assets WHERE type='ip' AND ip='10.20.30.40'`).Scan(&portCnt)
|
||||
if portCnt == 0 {
|
||||
t.Error("side-effect: IP open_ports should contain port 3306")
|
||||
}
|
||||
|
||||
// IP side-effect: bound_domains contains the service domain
|
||||
var bdRaw []byte
|
||||
d.QueryRow(`SELECT array_to_json(bound_domains)::text FROM assets WHERE type='ip' AND ip='10.20.30.40'`).Scan(&bdRaw)
|
||||
var bd []string
|
||||
json.Unmarshal(bdRaw, &bd)
|
||||
hasDomain := false
|
||||
for _, dom := range bd {
|
||||
if dom == "db.othersvc-test.com" {
|
||||
hasDomain = true
|
||||
}
|
||||
}
|
||||
if !hasDomain {
|
||||
t.Errorf("side-effect: IP bound_domains missing service domain, got %v", bd)
|
||||
}
|
||||
}
|
||||
|
||||
// =====================================================================
|
||||
// TestInsertAssetsMixedBatch
|
||||
// 混合批量插入:一次调用插入多种类型
|
||||
// =====================================================================
|
||||
func TestInsertAssetsMixedBatch(t *testing.T) {
|
||||
d := testDB(t)
|
||||
defer d.Close()
|
||||
|
||||
ts := NewToolSet(nil, "")
|
||||
ts.SetAssetStore(d.Assets(), d.Companies())
|
||||
defer d.Exec(`DELETE FROM assets WHERE
|
||||
domain IN ('batch-sub.batch-test.org','batch-test.org') OR
|
||||
ip='55.66.77.88' OR
|
||||
url='https://batch-test.org/api' OR
|
||||
(type='endpoint' AND url='https://batch-test.org/api/users')`)
|
||||
|
||||
out := callInsertAssets(t, ts, map[string]any{
|
||||
"assets": []any{
|
||||
// root_domain
|
||||
map[string]any{"type": "root_domain", "domain": "batch-test.org"},
|
||||
// subdomain with A record
|
||||
map[string]any{"type": "subdomain", "domain": "batch-sub.batch-test.org", "record_type": "A", "record_value": []string{"55.66.77.88"}},
|
||||
// HTTP service
|
||||
map[string]any{"type": "service", "url": "https://batch-test.org/api", "technologies": []string{"Go", "PostgreSQL"}, "status_code": 200},
|
||||
// endpoint
|
||||
map[string]any{"type": "endpoint", "url": "https://batch-test.org/api/users", "method": "GET"},
|
||||
},
|
||||
"task_id": 777,
|
||||
})
|
||||
|
||||
if errs, _ := out["errors"].([]any); len(errs) > 0 {
|
||||
t.Errorf("unexpected errors: %v", errs)
|
||||
}
|
||||
results, _ := out["results"].([]any)
|
||||
if len(results) != 4 {
|
||||
t.Errorf("mixed batch: want 4 results, got %d", len(results))
|
||||
}
|
||||
|
||||
// verify all types exist in DB
|
||||
types := []string{"root_domain", "subdomain", "service", "endpoint"}
|
||||
for _, typ := range types {
|
||||
var cnt int
|
||||
switch typ {
|
||||
case "root_domain":
|
||||
d.QueryRow(`SELECT COUNT(*) FROM assets WHERE type=$1 AND domain='batch-test.org'`, typ).Scan(&cnt)
|
||||
case "subdomain":
|
||||
d.QueryRow(`SELECT COUNT(*) FROM assets WHERE type=$1 AND domain='batch-sub.batch-test.org'`, typ).Scan(&cnt)
|
||||
case "service":
|
||||
d.QueryRow(`SELECT COUNT(*) FROM assets WHERE type=$1 AND url='https://batch-test.org/api'`, typ).Scan(&cnt)
|
||||
case "endpoint":
|
||||
d.QueryRow(`SELECT COUNT(*) FROM assets WHERE type=$1 AND url='https://batch-test.org/api/users'`, typ).Scan(&cnt)
|
||||
}
|
||||
if cnt != 1 {
|
||||
t.Errorf("mixed batch: %s not found in DB", typ)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// =====================================================================
|
||||
// TestInsertAssetsDedup
|
||||
// 幂等写入:同一资产插入两次,返回相同 ID
|
||||
// =====================================================================
|
||||
func TestInsertAssetsDedup(t *testing.T) {
|
||||
d := testDB(t)
|
||||
defer d.Close()
|
||||
|
||||
ts := NewToolSet(nil, "")
|
||||
ts.SetAssetStore(d.Assets(), d.Companies())
|
||||
defer d.Exec(`DELETE FROM assets WHERE domain='dedup-ia.deduptest.net' OR domain='deduptest.net'`)
|
||||
|
||||
payload := map[string]any{
|
||||
"assets": []any{
|
||||
map[string]any{"type": "root_domain", "domain": "deduptest.net"},
|
||||
},
|
||||
}
|
||||
|
||||
out1 := callInsertAssets(t, ts, payload)
|
||||
out2 := callInsertAssets(t, ts, payload)
|
||||
|
||||
getID := func(out map[string]any) float64 {
|
||||
results, _ := out["results"].([]any)
|
||||
if len(results) == 0 {
|
||||
return 0
|
||||
}
|
||||
m, _ := results[0].(map[string]any)
|
||||
id, _ := m["id"].(float64)
|
||||
return id
|
||||
}
|
||||
|
||||
id1, id2 := getID(out1), getID(out2)
|
||||
if id1 == 0 || id1 != id2 {
|
||||
t.Errorf("dedup: want same ID on double insert, got %v vs %v", id1, id2)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,125 @@
|
||||
package agent
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
|
||||
"github.com/Autumn-27/norma/agentcore"
|
||||
"github.com/Autumn-27/norma/llm"
|
||||
"github.com/Autumn-27/norma/permission"
|
||||
actool "github.com/Autumn-27/norma/tool"
|
||||
"github.com/Autumn-27/norma/transcript"
|
||||
"github.com/Autumn-27/artex/db"
|
||||
)
|
||||
|
||||
// MainAgent is the thin human-interface orchestrator (docs §4.2 / §7). The human
|
||||
// chats with it; it observes (read tools), and steers by injecting hints
|
||||
// (→planner) or direct high-priority intents (→frontier). It does NOT run the
|
||||
// autonomous intent-generation loop (that is the planner's job).
|
||||
type MainAgent struct {
|
||||
prov llm.Provider
|
||||
model string
|
||||
tx *transcript.Store // raw LLM conversation persistence (nil = off)
|
||||
window int // context window in tokens (for compaction)
|
||||
maxTurns int // max agent turns per run (0 = unlimited)
|
||||
proxyAddr string // recording proxy for WebFetch (empty = direct)
|
||||
proxyCACert string // recording proxy's CA cert path (HTTPS verify)
|
||||
webSearch WebSearchOpts // web_search tool backend selection (off by default)
|
||||
workDir string // shared work dir (surfaced in prompt as artifact-output target)
|
||||
}
|
||||
|
||||
func NewMainAgent(prov llm.Provider, model, workDir string, tx *transcript.Store, window, maxTurns int) *MainAgent {
|
||||
return &MainAgent{prov: prov, model: model, workDir: workDir, tx: tx, window: window, maxTurns: maxTurns}
|
||||
}
|
||||
|
||||
// SetProxy points the main agent's WebFetch at the recording proxy plus the CA
|
||||
// cert it trusts to verify HTTPS through it (empty addr = direct).
|
||||
func (m *MainAgent) SetProxy(addr, caCert string) { m.proxyAddr, m.proxyCACert = addr, caCert }
|
||||
|
||||
// SetWebSearch selects the web_search backend for the main agent (off by default).
|
||||
func (m *MainAgent) SetWebSearch(o WebSearchOpts) { m.webSearch = o }
|
||||
|
||||
// mainAgentDefaultTmpl is the built-in EDITABLE body (段 [A]) of the main agent
|
||||
// prompt, seeded into agent_prompts. Goal is a {{.Goal}} template var; the 中间
|
||||
// 产物输出规约 tail is code-owned (artifactSpec), appended after rendering.
|
||||
const mainAgentDefaultTmpl = `你是一个授权渗透测试系统的"主 agent",是人类操作员的接口。你不亲自探索、也不自主连续生成意图(那是规划者的工作)。你的职责:
|
||||
|
||||
1. 观察:用 graph_overview / list_findings / list_facts / list_assets / get_worker_output 回答人关于当前进展的问题。
|
||||
2. 操舵(把人的意图落到系统):
|
||||
- 人想"改方向/强调某类漏洞/重点某区域" → 用 add_hint 写提示(规划者下次会读到)。
|
||||
- 人想"立刻测某个具体目标" → 用 add_intent 直接注入一条高优先级意图(priority 8-10)。
|
||||
3. 用人话简洁回复,说明你做了什么。
|
||||
|
||||
当前任务目标:{{.Goal}}
|
||||
|
||||
不要编造发现;只根据工具返回的真实数据回答。`
|
||||
|
||||
func mainAgentSystem(goal, workDir string) string {
|
||||
body := renderSystem("mainagent", mainAgentDefaultTmpl, MainVars{Goal: goal})
|
||||
return body + artifactSpec(workDir)
|
||||
}
|
||||
|
||||
// Chat handles one human message and returns the assistant reply. emit, if
|
||||
// non-nil, receives each execution step (thinking / tool_use / tool_result /
|
||||
// text / result) so the main-agent session shows its work — exactly like the
|
||||
// worker/planner sessions — not just the final answer.
|
||||
func (m *MainAgent) Chat(ctx context.Context, taskID int64, as *db.AssetStore, ts *db.ExplorationStore, goal, message string, emit func(db.Activity), notify func()) (string, error) {
|
||||
tsx := NewToolSet(ts, "human")
|
||||
if as != nil {
|
||||
tsx.SetAssetStore(as, as.Companies())
|
||||
}
|
||||
tsx.SetTaskID(taskID)
|
||||
tsx.SetNotify(notify) // add_hint wakes this task's planner (debounced)
|
||||
// 领域工具 + 基础默认工具集(Read/Write/Edit/MultiEdit/LS/Glob/Grep/Bash)
|
||||
base := append(tsx.MainAgentTools(), actool.DefaultTools()...)
|
||||
tools, def, cleanup := AugmentTools(ctx, "mainagent", base)
|
||||
defer cleanup()
|
||||
system, boundary := deferredSystem(mainAgentSystem(goal, m.workDir), def)
|
||||
opts := agentcore.Options{
|
||||
Provider: m.prov,
|
||||
SystemPrompt: system,
|
||||
DynamicBoundary: boundary,
|
||||
Tools: tools,
|
||||
DeferredTools: def.Deferred,
|
||||
UnlockSet: def.Unlock,
|
||||
PermissionMode: permission.ModeBypass,
|
||||
EnableWebFetch: true, // 走记录代理留痕;载入代理 CA 验证 MITM 重签的 HTTPS 证书
|
||||
WebFetchProxy: m.proxyAddr,
|
||||
WebFetchCACert: m.proxyCACert,
|
||||
// 联网搜索(可选)。ddgs 无需 key;brave-free 需 BraveKey;tavily 需 TavilyKey。
|
||||
// WebSearchProxy 是独立出口代理(http/https/socks5),与记录流量的 MITM 代理无关;空则直连。
|
||||
EnableWebSearch: m.webSearch.Enabled,
|
||||
WebSearchBackend: m.webSearch.Backend,
|
||||
BraveSearchAPIKey: m.webSearch.BraveKey,
|
||||
TavilySearchAPIKey: m.webSearch.TavilyKey,
|
||||
WebSearchProxy: m.webSearch.Proxy,
|
||||
BashEnv: proxyEnv(m.proxyAddr, m.proxyCACert), // Bash 子命令默认走代理+信任 CA
|
||||
MaxTurns: m.maxTurns, // 0 = unlimited (configurable in agent management)
|
||||
Compaction: compactionConfig(m.window), // long chats stay within the window
|
||||
Todos: actool.NewTodoStore(), // 会话级临时待办(TodoWrite),纯规划用,退出即丢
|
||||
// 命中预算(步数)→ SDK 跑收尾:向用户输出一句进展总结。Prompt 与收尾轮数可后台编辑(默认 5 轮)。
|
||||
Settlement: wrapupSettlement("mainagent", nil),
|
||||
}
|
||||
if m.tx != nil { // persist raw human↔AI conversation; one accumulating file per task
|
||||
opts.Transcript = m.tx
|
||||
opts.SessionID = fmt.Sprintf("exp%d-main", ts.ID())
|
||||
}
|
||||
s := agentcore.NewSession(opts)
|
||||
defer s.Close()
|
||||
// reload the prior conversation from the transcript so the agent has context
|
||||
// across turns (each Chat is a fresh session; without this it can't see earlier
|
||||
// messages). First turn: no file yet → Resume loads nothing and proceeds.
|
||||
if m.tx != nil {
|
||||
_ = s.Resume(opts.SessionID)
|
||||
}
|
||||
// C2: this session is fresh each turn; re-unlock skill-gated MCPs from prior
|
||||
// Skill() calls in the reloaded history so revealed tools stay callable.
|
||||
seedUnlockFromHistory(s.Messages(), def.UnlockSkill)
|
||||
text, _, err := captureRunSession(ctx, s, message, func(r db.Activity) {
|
||||
if emit != nil {
|
||||
r.Worker = "mainagent"
|
||||
emit(r)
|
||||
}
|
||||
})
|
||||
return text, err
|
||||
}
|
||||
@@ -0,0 +1,260 @@
|
||||
package agent
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"strings"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"github.com/Autumn-27/norma/agentcore"
|
||||
"github.com/Autumn-27/norma/llm"
|
||||
"github.com/Autumn-27/norma/permission"
|
||||
actool "github.com/Autumn-27/norma/tool"
|
||||
"github.com/Autumn-27/norma/transcript"
|
||||
"github.com/Autumn-27/artex/db"
|
||||
)
|
||||
|
||||
// Planner is the event-driven LLM planner (docs §4.3): each time the asset or
|
||||
// exploration graph changes (debounced), it reads the exploration route, queries
|
||||
// assets, judges whether the task goal is met, and emits 0..N exploration intents
|
||||
// into the frontier. It is the sole intent generator.
|
||||
type Planner struct {
|
||||
prov llm.Provider
|
||||
model string
|
||||
tx *transcript.Store // raw LLM conversation persistence (nil = off)
|
||||
window int // context window in tokens (for compaction)
|
||||
maxTurns int // max agent turns per run (0 = unlimited)
|
||||
killWork func(intentID int64) error // engine callback to terminate a running work (nil = off)
|
||||
steerWork func(intentID int64, msg string) error // engine callback to steer a running work mid-run (nil = off)
|
||||
proxyAddr string // recording proxy for WebFetch (empty = direct)
|
||||
proxyCACert string // recording proxy's CA cert path (HTTPS verify)
|
||||
webSearch WebSearchOpts // web_search tool backend selection (off by default)
|
||||
workDir string // shared work dir (surfaced in prompt as artifact-output target)
|
||||
|
||||
// todos keeps ONE plan-scratchpad per task (keyed by exploration id) so the
|
||||
// planner's multi-step plan survives across wake-ups — each Plan() is a fresh
|
||||
// session, but the shared store lets it record a serial exploit chain once and
|
||||
// dispatch it step-by-step over rounds instead of front-loading it in parallel.
|
||||
todoMu sync.Mutex
|
||||
todos map[int64]*actool.TodoStore
|
||||
}
|
||||
|
||||
func NewPlanner(prov llm.Provider, model, workDir string, tx *transcript.Store, window, maxTurns int) *Planner {
|
||||
return &Planner{prov: prov, model: model, workDir: workDir, tx: tx, window: window, maxTurns: maxTurns, todos: map[int64]*actool.TodoStore{}}
|
||||
}
|
||||
|
||||
// SetProxy points the planner's WebFetch at the recording proxy plus the CA cert
|
||||
// it trusts to verify HTTPS through it (empty addr = direct).
|
||||
func (p *Planner) SetProxy(addr, caCert string) { p.proxyAddr, p.proxyCACert = addr, caCert }
|
||||
|
||||
// SetWebSearch selects the web_search backend for the planner (off by default).
|
||||
func (p *Planner) SetWebSearch(o WebSearchOpts) { p.webSearch = o }
|
||||
|
||||
// todoFor returns the task's persistent planning todo store, creating it on first
|
||||
// use. Shared across all of this task's planner wake-ups.
|
||||
func (p *Planner) todoFor(expID int64) *actool.TodoStore {
|
||||
p.todoMu.Lock()
|
||||
defer p.todoMu.Unlock()
|
||||
s := p.todos[expID]
|
||||
if s == nil {
|
||||
s = actool.NewTodoStore()
|
||||
p.todos[expID] = s
|
||||
}
|
||||
return s
|
||||
}
|
||||
|
||||
// SetKillWork wires the engine's per-work terminate callback so the planner's
|
||||
// kill_work tool can stop a single running worker.
|
||||
func (p *Planner) SetKillWork(fn func(intentID int64) error) { p.killWork = fn }
|
||||
|
||||
// SetSteerWork wires the engine's per-work steering callback so the planner's
|
||||
// steer_work tool can inject a mid-run course-correction into a running worker.
|
||||
func (p *Planner) SetSteerWork(fn func(intentID int64, msg string) error) { p.steerWork = fn }
|
||||
|
||||
// renderPlannerTodos formats the persistent planning todo for injection into the
|
||||
// wake-up prompt (empty when there are no todos yet — first wake-up).
|
||||
func renderPlannerTodos(items []actool.Todo) string {
|
||||
if len(items) == 0 {
|
||||
return ""
|
||||
}
|
||||
var b strings.Builder
|
||||
b.WriteString("\n\n【你的规划待办(跨唤醒保留,上一轮你写的)】:\n")
|
||||
for _, it := range items {
|
||||
mark := map[actool.TodoStatus]string{actool.TodoPending: "☐", actool.TodoInProgress: "▶", actool.TodoCompleted: "✔"}[it.Status]
|
||||
if mark == "" {
|
||||
mark = "☐"
|
||||
}
|
||||
b.WriteString(fmt.Sprintf(" %s %s\n", mark, it.Content))
|
||||
}
|
||||
b.WriteString("据此推进:只对【前置步骤已完成 / 其依赖的 fact 已存在】的下一步派意图;用 TodoWrite 更新清单(把已被 fact 满足的步骤标 completed)。不要重复派已在清单里 pending/in_progress 的步骤。")
|
||||
return b.String()
|
||||
}
|
||||
|
||||
// renderDoneIntents surfaces the intent(s) whose completion triggered this round
|
||||
// so the planner looks first at their fresh yields (in the态势 below) before
|
||||
// deciding whether a new direction opened up. Empty for non-completion wakes.
|
||||
func renderDoneIntents(ids []int64) string {
|
||||
if len(ids) == 0 {
|
||||
return ""
|
||||
}
|
||||
parts := make([]string, len(ids))
|
||||
for i, id := range ids {
|
||||
parts[i] = fmt.Sprintf("#%d", id)
|
||||
}
|
||||
return "\n\n【本次触发本轮的、刚完成的意图】:" + strings.Join(parts, "、")
|
||||
}
|
||||
|
||||
// renderGraphOverview folds the pre-computed graph_overview snapshot into the
|
||||
// wake-up prompt so the planner starts each round with the full situation in
|
||||
// hand — saving the round-trip it would otherwise spend calling the tool. It is
|
||||
// the exact same JSON graph_overview would return; deeper detail is still one
|
||||
// tool call away (node_detail / list_facts / …).
|
||||
func renderGraphOverview(data map[string]any) string {
|
||||
b, err := json.Marshal(data)
|
||||
if err != nil {
|
||||
return "" // fall back to the model calling graph_overview itself
|
||||
}
|
||||
return "\n\n【本轮态势(graph_overview 预取,等同你调用该工具的返回;需要细节再按需调 node_detail/list_facts 等)】:\n" + string(b)
|
||||
}
|
||||
|
||||
// plannerDefaultTmpl is the built-in EDITABLE body (段 [A]) of the planner prompt,
|
||||
// seeded into agent_prompts. Goal is a {{.Goal}} template var; the 中间产物输出规约
|
||||
// tail is code-owned (artifactSpec) and appended by plannerSystem after rendering.
|
||||
const plannerDefaultTmpl = `你是一个授权渗透测试系统的"规划者"。你被频繁唤醒(图一变就唤醒)。你的职责:读取态势、判定目标、并且**只在确有未被覆盖的新方向时**才补充探索意图。
|
||||
|
||||
任务目标:{{.Goal}}
|
||||
|
||||
⚠️ 最重要的原则:**这一轮生成 0 个意图是完全正常、而且是最常见的结果。** 你不是"每次都要产出意图"的机器。绝大多数唤醒,frontier 里已有的意图已经覆盖了所有已知方向,你应当什么都不加、直接结束。重复/换措辞地生成已经存在的意图是严重错误。
|
||||
|
||||
决策流程(每次唤醒):
|
||||
1. **完整态势已直接附在本提示下方(就是 graph_overview 的返回,无需再调它)**:task(**原始任务标题+目标**,即根节点)、资产计数、goals 及其状态、open/running/recent_done 意图、sites_without_endpoints、findings(**确认漏洞**数)、facts(**探索事实/结论**数,与漏洞是两类)、recent_facts(最近事实的 {id, summary, confidence?},含"端口关闭/不可注入"等**否定结论**——据此别再为已探明的死路生成意图;但 confidence=inferred 的否定结论只是【推断】、证据弱,别当铁案,若该方向对目标很关键值得派一条复核意图)。**这里的探索节点(goals/意图/facts/findings)都只含本任务的**(绝不会有别的任务的目标);**资产图则全局共享**(多任务同一份,资产计数是全局在范围内的数据,非本任务独有)。需要更深的细节时才按需调:list_facts 列全部事实、list_findings 列全部漏洞、**node_detail(id)** 取某条的完整证据/详情(列表/recent_facts 只给摘要)。
|
||||
- open_intents / running_intents / recent_done_intents——"哪些方向已经有意图在覆盖/已尝试"。每个意图还带 **parents(上游:它派生自哪些事实/意图)和 yields(下游:它产生了哪些事实/发现)**——这就是探索图的**血缘关系**,据此理解"哪些事实来自哪个方向、能否综合成新方向"。recent_facts 里每个事实带 **from_intent**(由哪个意图产生)。
|
||||
- sites_without_endpoints / findings——"哪些方向【可能】需要探索"。
|
||||
- 只有需要某一片的细节时,才**按需**调 list_assets(pull 模式:可用 q 关键字搜索,可叠加 type/company_id/task_id 过滤,分页 limit/offset;或用 id/ids 直接取)、asset_neighbors、list_findings。资产图全局共享,别默认拉全量。资产中可能包含非本次任务涉及到的资产,所以需要主要出现非本次任务相关的资产时忽略这些资产。
|
||||
2. 判目标(核心职责):graph_overview 的 goals 字段已含目标与状态;对已被某发现/事实证明的未达成目标,调 prove_goal(goal_id,evidence_id,reason) 标记 met。**当你用 prove_goal 标记的这一个恰好是最后一个未完成目标时,系统会自动判定整个任务完成、无需你再调 goal_met**。goal_met 仅在你想【绕过逐个 prove_goal、直接从全局判定任务已达成】时才用。
|
||||
2.5. **(可选,仅限开局、极轻量)探测理解**:你具备 Bash 等执行能力,它的**唯一正当用途**是——当**图里几乎还没有事实**(recent_facts 基本为空、任务刚开始)、仅凭态势无法把初始意图描述具体时,对目标做**极少量、只读**的探测(如 1–2 次 curl 看首页/指纹),据此产出更精准的**初始意图**。
|
||||
⚠️ **牢记你的身份边界:你是"规划者",不是"执行者"。你在这里做的一切都只为【生成/说清意图】,绝不是【在 plan 里把活干了】。** 探测的**唯一合法产物是一句更精准的意图描述**——绝不能是漏洞的发现、验证、利用,也不能是端点/目录/参数的枚举结果。任何"我顺手把这个也测了/确认了"的念头都是越界:那是 worker 在 work 阶段该做的事,你只需把它**写成一个意图派下去**。
|
||||
**三条硬性边界,务必守住**:
|
||||
- **只要图里已经有 worker 产出的 fact(facts>0 / recent_facts 非空),就【禁止】再自己探测。** 此时一切判断都基于已有 fact,你这一轮的产出只能是"派新意图"或"结束本轮"。看到某个线索想深挖时,**正确反应是派一个意图让 worker 去查,绝不是自己 curl**。
|
||||
- 即使在开局,也**最多探几次(≤3)就收手**,只为把初始意图说清楚;**绝不要**滑进:逐个枚举端点/目录、逐个试 id、base64 解码链、反复探同一接口、任何注入/越权/漏洞的测试或验证——那些**全是 worker 的重活**,是你要生成的意图内容,不是你自己该做的事。一旦发现自己在"深入查证"而不是"快速定方向",立刻停手,把剩下的写成意图。
|
||||
- 探测只帮你"想清楚",**本身不产出意图**;能从现有事实/态势判断的,就**根本不必**探测。
|
||||
3. **决定是否补探索方向(务必克制)**:意图是【开放的探索方向】,不是固定类型/菜单——你要结合**已知事实、资产与任务目标**,自行判断"为了逼近目标,下一步还有哪些有价值、尚未被覆盖的探索方向"。把这些方向逐一与 open_intents + running_intents + recent_done_intents 比对:
|
||||
- 该方向已有 open 或 running 意图覆盖 → **不要再生成**(正在被处理)。
|
||||
- 该方向在 recent_done_intents 里已尝试过(即使没产出)→ **不要原样重试**。把它当作【已封锁路线(blocked)】:仅当出现【材料性的新机理】——新事实、新资产、新参数、或一种明显不同的打法/构造——才重新派,且新意图的 summary 里要写清"这次和上次不同在哪"。**只是换个措辞、或"再试一次说不定行"都不算新机理,禁止重试。** 反之,若某条否定结论是 confidence=inferred 的弱证据、且该方向对目标关键,用【一条复核意图】去证实或推翻它是正当的(这属于新理由)。
|
||||
- 仅对【当前完全没有任何意图覆盖的全新方向】生成。
|
||||
- 所有已知方向都已被现有意图覆盖 → **不生成任何意图,直接结束本轮。**
|
||||
- **保持路线多样、别过早收敛到一条**:当目标尚未达成时,若现有意图全都挤在【同一条攻击路线/同一类入口】上,而还存在【本质不同】的未覆盖方向(如另一种入口面、另一类资产、另一条利用链),优先补一条那样的分歧方向,而不是在同一条线上再加同义意图。理想状态是让 2–3 条彼此独立、机理不同的路线同时存活(如"从上传链打"与"从认证绕过打");只有当某条路线已交出【目标逼近】的证据时,才值得把资源集中过去。判断多样性看方向的实质差异,不看措辞。(这不与"0 意图正常"冲突:只有确存在本质不同且未覆盖的方向才补;空白路线若已被现有意图覆盖,仍旧 0 意图。)
|
||||
3.5. **串行利用链:分步派,别一次拆成并行。** 很多利用是一条**强依赖串行链**(如:①→ ② → ③)——后一步依赖前一步的**实际产出**。这种情况:
|
||||
- **不要**把 ①②③ 一次性作为多个并行意图下发(下游 worker 拿不到还不存在的前置产出,只会重复/空转);
|
||||
- **用 TodoWrite 把整条链记成待办**(每步一条),然后**本轮只派"前置已满足"的那一步**(通常是第一步);
|
||||
- 待该步产出 fact 后,下一次唤醒(提示里会带上你的待办清单)再派下一步,并把已满足的步骤标 completed;
|
||||
- 判断"同一件事"别拆两条(如"确认触发点"和"触发触发点"是同一步)。**平行的探索维度**(如同时枚举多个不相关端点)才用多意图并行。
|
||||
4. 用【一次】 add_intent 批量提交第 3 步筛出的新方向(intents 数组,最多 4 个最高价值的;不要逐条多次调):
|
||||
- summary:一句话自由描述该方向(做什么+为什么),用自然语言,**不要套用固定分类**。方向要在 summary 里写清楚,去重主要靠它与已有意图比对。
|
||||
- asset_ids:本方向要**测试/攻击的目标资产 id**(**尽量传**,0/1/多个;是 list_assets 里的资产 id)。只要方向围绕某些具体资产(某站点/接口/参数/主机)就**务必传上**——它标记「这条探索打哪些目标」,用于覆盖去重、把意图连入资产链路;跨多个资产就都传。仅当纯全局侦察、确实没有具体目标资产时才留空。
|
||||
- parent_ids:本方向由哪些【上游节点】综合得出(可选,0/1/多个)。**多个事实结合产生一个新意图,就把这些事实 id 都传上**;派生自某上游意图/发现也传其 id;顶层全新方向留空。
|
||||
|
||||
|
||||
宁可不生成,也不要重复或硬凑。简洁、克制、高效。`
|
||||
|
||||
func plannerSystem(goal, workDir string) string {
|
||||
body := renderSystem("planner", plannerDefaultTmpl, PlannerVars{Goal: goal, Now: time.Now().Format("2006-01-02 15:04:05 MST")})
|
||||
return body + artifactSpec(workDir)
|
||||
}
|
||||
|
||||
// Plan runs one planning round. emit, if non-nil, receives the planner's execution
|
||||
// steps (so users can see how it reads the situation and judges goals — the
|
||||
// planner is the intent generator and was previously a black box). Returns whether
|
||||
// the planner judged the goal met.
|
||||
// doneIntents carries the ids of the intent(s) whose completion triggered this
|
||||
// round (may be several — the engine debounces a burst; empty for non-completion
|
||||
// wakes). They are surfaced in the prompt so the planner focuses on the fresh
|
||||
// yields of what just finished.
|
||||
func (p *Planner) Plan(ctx context.Context, taskID int64, as *db.AssetStore, ts *db.ExplorationStore, goal string, doneIntents []int64, emit func(db.Activity)) (met bool, reason string, err error) {
|
||||
tsx := NewToolSet(ts, "planner")
|
||||
if as != nil {
|
||||
tsx.SetAssetStore(as, as.Companies())
|
||||
}
|
||||
tsx.SetTaskID(taskID)
|
||||
tsx.killWork = p.killWork // enable kill_work tool (nil = unavailable)
|
||||
tsx.steerWork = p.steerWork // enable steer_work tool (nil = unavailable)
|
||||
if origin, _ := ts.OriginFactID(); origin > 0 {
|
||||
tsx.SetOwnerNode(origin) // planner-side anchors default to the task root (origin fact)
|
||||
}
|
||||
// 领域工具 + 基础默认工具集(Read/Write/Edit/MultiEdit/LS/Glob/Grep/Bash)
|
||||
base := append(tsx.PlannerTools(), actool.DefaultTools()...)
|
||||
tools, def, cleanup := AugmentTools(ctx, "planner", base)
|
||||
defer cleanup()
|
||||
// 关键态势(刚完成的意图 + 预取的完整图)放进 system prompt,而不是 user 输入:
|
||||
// system 段永不被 compaction 压缩,长回合里也保证这份态势始终在场;轮内固定,还能吃到
|
||||
// system 缓存。它落在指令之后、deferred 工具块之前(deferredSystem 保证块在最后)。
|
||||
situational := renderDoneIntents(doneIntents) + renderGraphOverview(tsx.graphOverviewData())
|
||||
// 任务级 deadline / 终局模式(经 ctx 注入,见 taskclock.go)。终局那一轮把任务超时
|
||||
// planner 收尾词作为【本轮操作指令】直接注入 system,让它只做最后目标判定、不产新意图。
|
||||
tc := taskClockFrom(ctx)
|
||||
if tc.Final {
|
||||
situational += "\n\n【任务终局收尾(本轮特殊指令,覆盖上面的常规规划流程)】:" + resolveTaskTimeoutWrapup("planner")
|
||||
}
|
||||
system, boundary := deferredSystem(plannerSystem(goal, p.workDir)+situational, def)
|
||||
// planner 无自身墙钟预算;有 deadline 时把 MaxDuration 夹逼到剩余,让在跑的规划轮在
|
||||
// 任务到点时进收尾(因超时→任务超时词,因步数→per-run 词)。
|
||||
maxDur, clamped := clampMaxDuration(tc.DeadlineUnix, 0)
|
||||
settle := wrapupSettlement("planner", nil)
|
||||
if tc.DeadlineUnix > 0 {
|
||||
settle = wrapupSettlementForTask("planner", nil, clamped)
|
||||
}
|
||||
opts := agentcore.Options{
|
||||
Provider: p.prov,
|
||||
SystemPrompt: system,
|
||||
DynamicBoundary: boundary,
|
||||
Tools: tools,
|
||||
DeferredTools: def.Deferred,
|
||||
UnlockSet: def.Unlock,
|
||||
PermissionMode: permission.ModeBypass,
|
||||
EnableWebFetch: true, // 走记录代理留痕;载入代理 CA 验证 MITM 重签的 HTTPS 证书
|
||||
WebFetchProxy: p.proxyAddr,
|
||||
WebFetchCACert: p.proxyCACert,
|
||||
// 联网搜索(可选)。ddgs 无需 key;brave-free 需 BraveKey;tavily 需 TavilyKey。
|
||||
// WebSearchProxy 是独立出口代理(http/https/socks5),与记录流量的 MITM 代理无关;空则直连。
|
||||
EnableWebSearch: p.webSearch.Enabled,
|
||||
WebSearchBackend: p.webSearch.Backend,
|
||||
BraveSearchAPIKey: p.webSearch.BraveKey,
|
||||
TavilySearchAPIKey: p.webSearch.TavilyKey,
|
||||
WebSearchProxy: p.webSearch.Proxy,
|
||||
BashEnv: proxyEnv(p.proxyAddr, p.proxyCACert), // Bash 子命令默认走代理+信任 CA
|
||||
MaxTurns: p.maxTurns, // 0 = unlimited (configurable in agent management)
|
||||
MaxDuration: maxDur, // 0=不限;有 deadline 时=距 deadline 剩余
|
||||
Compaction: compactionConfig(p.window),
|
||||
// 跨唤醒共享的规划待办:让串行链在多轮之间保留(session 是新的,store 不是)。
|
||||
Todos: p.todoFor(ts.ID()),
|
||||
// 命中【本轮】步数预算→ SDK 跑收尾:把本轮已想清楚的结论落地(该派的 add_intent、
|
||||
// 能证的 prove_goal、串行链记 TodoWrite),而非停止规划——planner 之后仍会被反复唤醒。
|
||||
// clamped(被任务 deadline 夹逼)时改用 PromptByReason(见 wrapupSettlementForTask)。
|
||||
Settlement: settle,
|
||||
}
|
||||
if p.tx != nil { // persist raw LLM conversation; one accumulating file per task's planner
|
||||
opts.Transcript = p.tx
|
||||
opts.SessionID = fmt.Sprintf("exp%d-planner", ts.ID())
|
||||
}
|
||||
// 态势(刚完成的意图 + 完整图)已在上面拼进 system prompt(抗压缩、轮内常驻)。这里
|
||||
// 的 user 输入只留指令 + 跨唤醒待办(todo 是模型自己的规划便签,可再生,放 user 即可)。
|
||||
input := "图发生了变化,请规划下一步:读取上方 system 里的态势,判定目标。**本轮若无未被覆盖的新方向,直接结束即可(生成 0 个意图是正常且常见的,尤其刚派完意图在等 worker 产出时)。绝不要为了“结束本轮”去调 goal_met——goal_met 会【立即结束整个任务】,只在你确认目标已【真正达成】(已拿到目标成果/已确认目标漏洞)时才调;未达成就用 prove_goal 逐个标记、或什么都不调直接结束。**" +
|
||||
renderPlannerTodos(opts.Todos.List())
|
||||
// 有 deadline 夹逼时加硬 ctx 兜底(软预算 + grace),防单轮卡死绕过轮边界软超时。
|
||||
runCtx := ctx
|
||||
if maxDur > 0 {
|
||||
var cancel context.CancelFunc
|
||||
runCtx, cancel = context.WithTimeout(ctx, maxDur+settleHardGrace)
|
||||
defer cancel()
|
||||
}
|
||||
_, _, err = captureRun(runCtx, opts, input,
|
||||
func(r db.Activity) {
|
||||
if emit != nil {
|
||||
r.Worker = "planner" // planner activity has no intent_id (it generates them)
|
||||
emit(r)
|
||||
}
|
||||
})
|
||||
return tsx.GoalMet, tsx.Reason, err
|
||||
}
|
||||
@@ -0,0 +1,58 @@
|
||||
package agent
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"text/template"
|
||||
)
|
||||
|
||||
// PromptOverride, if set, returns the stored system-prompt template for an agent
|
||||
// key and whether one exists. The server wires it to the PG agent_prompts table.
|
||||
// When nil or no override exists, agents use their built-in default prompt — so
|
||||
// behavior is identical until a user edits a prompt in the UI.
|
||||
var PromptOverride func(agentKey string) (string, bool)
|
||||
|
||||
// Prompt-variable structs — fields mirror each agent's catalog (docs §5a) so a
|
||||
// user template referencing a catalog variable renders; referencing anything else
|
||||
// fails template execution and falls back to the built-in default.
|
||||
type PlannerVars struct{ Goal, Scope, AssetSummary, Now string }
|
||||
type WorkerVars struct{ ProxyAddr, WorkerName string }
|
||||
type MainVars struct{ Goal, AssetSummary, FindingsSummary string }
|
||||
type GoalsVars struct{ EngagementDescription, Now string }
|
||||
|
||||
// renderSystem returns the rendered system-prompt BODY (段 [A]) for agentKey.
|
||||
// Precedence: the DB-stored template (if any) over the built-in default template
|
||||
// (def). BOTH are Go templates now — the built-in default is seeded into the DB
|
||||
// verbatim, so the two paths render identically until a user edits the prompt.
|
||||
// Rendering always runs (def used to be pre-substituted plain text; it is now a
|
||||
// {{.Var}} template like the DB one). On any render error we fall back to the
|
||||
// default template, then to the raw default string — an agent never starts with a
|
||||
// half-rendered prompt. Callers append the code-owned tail (trafficTool / 中间产物
|
||||
// 输出规约) AFTER this, so those can't be edited away via the DB body.
|
||||
func renderSystem(agentKey, def string, vars any) string {
|
||||
tmpl := def
|
||||
if PromptOverride != nil {
|
||||
if t, ok := PromptOverride(agentKey); ok && t != "" {
|
||||
tmpl = t
|
||||
}
|
||||
}
|
||||
if out, err := renderTmpl(tmpl, vars); err == nil {
|
||||
return out
|
||||
}
|
||||
// DB template broke (e.g. references an out-of-catalog var) → code default.
|
||||
if out, err := renderTmpl(def, vars); err == nil {
|
||||
return out
|
||||
}
|
||||
return def
|
||||
}
|
||||
|
||||
func renderTmpl(tmpl string, vars any) (string, error) {
|
||||
t, err := template.New("p").Option("missingkey=error").Parse(tmpl)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
var b bytes.Buffer
|
||||
if err := t.Execute(&b, vars); err != nil {
|
||||
return "", err
|
||||
}
|
||||
return b.String(), nil
|
||||
}
|
||||
@@ -0,0 +1,67 @@
|
||||
package agent
|
||||
|
||||
import (
|
||||
"strings"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestRenderSystemOverrideAndFallback(t *testing.T) {
|
||||
t.Cleanup(func() { PromptOverride = nil })
|
||||
|
||||
// no override → built-in default
|
||||
PromptOverride = nil
|
||||
if got := renderSystem("planner", "DEFAULT", PlannerVars{Goal: "g"}); got != "DEFAULT" {
|
||||
t.Fatalf("no override should give default, got %q", got)
|
||||
}
|
||||
|
||||
// override → rendered with vars
|
||||
PromptOverride = func(k string) (string, bool) {
|
||||
if k == "planner" {
|
||||
return "目标:{{.Goal}} 范围:{{.Scope}}", true
|
||||
}
|
||||
return "", false
|
||||
}
|
||||
if got := renderSystem("planner", "DEFAULT", PlannerVars{Goal: "拿下X", Scope: "*.x.com"}); got != "目标:拿下X 范围:*.x.com" {
|
||||
t.Fatalf("override render: %q", got)
|
||||
}
|
||||
|
||||
// override referencing a non-catalog var → execution error → fallback to default
|
||||
PromptOverride = func(k string) (string, bool) { return "{{.NotInCatalog}}", true }
|
||||
if got := renderSystem("planner", "DEFAULT", PlannerVars{Goal: "x"}); got != "DEFAULT" {
|
||||
t.Fatalf("bad var should fall back to default, got %q", got)
|
||||
}
|
||||
|
||||
// full plannerSystem path: DB body [A] is honored, then the code-owned tail
|
||||
// [C] (中间产物输出规约) is ALWAYS appended — editing the body can't drop it.
|
||||
PromptOverride = func(k string) (string, bool) { return "PLANNER {{.Goal}}", true }
|
||||
got := plannerSystem("拿下X", "/data")
|
||||
if !strings.HasPrefix(got, "PLANNER 拿下X") {
|
||||
t.Fatalf("plannerSystem body not honored: %q", got)
|
||||
}
|
||||
if !strings.Contains(got, "中间产物输出规约") || !strings.Contains(got, "/data") {
|
||||
t.Fatalf("plannerSystem missing code-owned artifact tail: %q", got)
|
||||
}
|
||||
|
||||
// worker dual-text via {{if .ProxyAddr}} in a user template, plus the code tail:
|
||||
// [B] trafficTool present only with a proxy, [C] artifact spec always present.
|
||||
PromptOverride = func(k string) (string, bool) {
|
||||
return "{{if .ProxyAddr}}走代理 {{.ProxyAddr}}{{else}}手动{{end}}", true
|
||||
}
|
||||
withProxy := workerSystem("127.0.0.1:8080", "/data")
|
||||
if !strings.HasPrefix(withProxy, "走代理 127.0.0.1:8080") {
|
||||
t.Fatalf("worker proxy branch body: %q", withProxy)
|
||||
}
|
||||
if !strings.Contains(withProxy, "traffic_search") {
|
||||
t.Fatalf("worker with proxy should inject trafficTool: %q", withProxy)
|
||||
}
|
||||
if !strings.Contains(withProxy, "中间产物输出规约") {
|
||||
t.Fatalf("worker missing artifact tail: %q", withProxy)
|
||||
}
|
||||
noProxy := workerSystem("", "/data")
|
||||
if !strings.HasPrefix(noProxy, "手动") {
|
||||
t.Fatalf("worker no-proxy branch body: %q", noProxy)
|
||||
}
|
||||
if strings.Contains(noProxy, "traffic_search") {
|
||||
t.Fatalf("worker without proxy must NOT inject trafficTool: %q", noProxy)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,83 @@
|
||||
package agent
|
||||
|
||||
// 本文件把内置 agent 的「默认提示词正文」(段 [A]) 变成可枚举、可被服务端幂等
|
||||
// 播种进 agent_prompts 表的目录 —— 镜像 toolcatalog.go 的 BuiltinToolSeeds()。
|
||||
//
|
||||
// 只包含【可编辑正文】:段 [B] trafficTool 与段 [C] 中间产物输出规约 是代码固定
|
||||
// 注入(见 worker.go 的 workerTrafficBlock/artifactSpec),不入库、不可编辑,因此
|
||||
// 不在种子里。种子文本用 Go 模板占位({{.Goal}} 等),渲染时按运行期变量填充。
|
||||
|
||||
// autoDefaultTmpl is the built-in "Auto" platform-operator agent's prompt. Auto
|
||||
// runs via the chat page and drives the platform through tools: task ops
|
||||
// (spawn/list/pause/hint + read graph/findings/traces) and platform management
|
||||
// (create/modify skill, custom tool, MCP). It seeds into agent_prompts like the
|
||||
// other built-ins.
|
||||
const autoDefaultTmpl = `你是 **Auto**,这个渗透测试平台的「操作助手」。你不亲自渗透,而是**用工具操作平台**、按用户指令把事情办好。
|
||||
|
||||
你能做的(取决于给你开放了哪些工具):
|
||||
1. **任务操作**:list_tasks 看全局、spawn_task 起子任务、get_task_graph / list_task_findings 读某任务的进展与漏洞(含 flag)、get_task_worker_trace 看某个 work 的执行过程、pause_task 暂停、add_task_hint 给任务注入提示。
|
||||
2. **平台管理**:create_skill / update_skill 建改技能;create_custom_tool / update_custom_tool 建改自定义工具(command/script/http);create_mcp / update_mcp 建改 MCP 服务器。
|
||||
|
||||
原则:
|
||||
- 先看清现状(list_tasks / get_task_graph 等)再动手;一步到位、少空转。
|
||||
- 建/改 skill、工具、MCP 时,把用户意图翻译成正确的结构化参数(kind/exec/schema 等),字段拿不准就按最小可用填。
|
||||
- 用人话简洁汇报你做了什么、结果如何;只根据工具真实返回作答,不臆造。
|
||||
- 只在授权范围内操作。`
|
||||
|
||||
// pentestDefaultTmpl is the built-in "渗透测试" (solo pentest) agent's prompt. Unlike
|
||||
// the orchestration roles (goals/planner/worker), it runs standalone via the chat page
|
||||
// and is its own planner + executor + auditor. Default tools: list_assets / insert_assets
|
||||
// / report_finding / list_findings (bound in toolcatalog + seedPentestDefaultBindings).
|
||||
const pentestDefaultTmpl = `你是一个授权渗透测试系统的"独立渗透 agent"。你**一个人从头打到尾**:侦察 → 找攻击面 → 深入利用 → 验证 → 收尾。你同时是自己的规划者和执行者——没有别人给你派活,也没有别人替你把关,所有判断和动手都由你完成。正因如此,你要**主动切换视角**:该拓宽时像规划者一样铺开多条路线,该动手时像执行者一样把一条路走透,该验证时像审计者一样怀疑自己的结论。
|
||||
|
||||
|
||||
**只在授权范围内操作。范围外的目标一律不碰。**
|
||||
|
||||
━━ 核心心法(贯穿全程)━━
|
||||
1. **先广后聚,别隧道视野**。开局别一头扎进第一个看起来好打的点。先快速摸清目标有哪些**本质不同**的攻击面,铺开一个**多样化的路线组合**,让 2–3 条机理不同的路线并行推进(如"从上传链打"与"从认证绕过打")。只有当某条路线交出了【逼近目标】的实证,才值得把精力集中过去。单脑最容易犯的错就是过早爱上一条优雅路线而错过真正的洞。
|
||||
2. **一条路要走透再下结论**。初次受阻(一个 payload 被过滤、一个端点 404、一个注入点没回显)**不等于**此路不通——换编码、换方法、换参数、换路径,把这条方向的合理手段走完,再判"死路"。"我试了一次没成功"绝不等于"已穷尽"。
|
||||
3. **封锁路线不无理由重试**。确认走不通的方向,标记为封锁;**只有出现材料性的新机理**(新发现、新入口、新参数、明显不同的构造)才重开,且要能说清"这次和上次不同在哪"。换个措辞、"再试一次说不定行"都不算,禁止空转。
|
||||
4. **对自己的结论做对抗式自检**。这是单 agent 最关键的纪律:每当你觉得"发现漏洞了/成功了",**先切换成怀疑者**,用与首次【不同的路径或独立命令】再触发一次来证实,而不是复述原来的证据。尤其警惕这些自欺模式——把"版本号/CVE 命中"当漏洞、把"参数看起来可注入"当已利用、用与结论等价的假设循环当证据。**证伪和证实同等有价值**:自检没过就老实记为未确认,别硬认。
|
||||
5. **要具体结论,不要状态报告**。你的产出是可核验的事实、可复现的 PoC、或明确的否定结论——不是"看起来有戏""疑似存在""大概可以"这类含糊乐观。拿不准就标 inferred,别当铁案。
|
||||
6. **不轻言放弃**。一波尝试失败很正常,别就此收手。回到路线组合,换个攻击面、找新的形式化切入,继续推进;只有在目标达成、或所有合理路线都真正探尽后才停。
|
||||
|
||||
━━ 工作循环(是启发,不是死板流程)━━
|
||||
- **侦察定面**:识别指纹、入口、参数、信任边界,把目标的攻击面铺开。常被忽略的高价值面(据实际情况挑,非清单义务):输入解析/编码与字符集边界、文件上传、(反)序列化、内置路由与认证前可达面、错误处理泄露、缓存(投毒/竞态)、竞态条件、类型混淆(scalar vs array)、批量赋值,以及任何你识别出的攻击者可及面。
|
||||
- **组合与优先级**:把发现的方向排成 2–3 条独立路线,用 TodoWrite 记下来(每条一项),据"离目标多近 + 代价多大"定先后。
|
||||
- **深入利用**:挑前置已满足的路线动手,走透。**串行利用链**(①→②→③,后一步依赖前一步的**实际产出**)就一步步来:先做第一步、拿到真实产出,再据此做下一步;别在前置还不存在时就假想后续。跨代码库/跨接口把多个 gadget 在**本次会话内**串成一条可触发的链,正是单 agent 的强项——主动把已知线索的完整细节调出来综合,别停留在摘要。
|
||||
- **验证**:见心法 4,对每个候选发现做独立复现/证伪。
|
||||
- **回到组合**:一条路出结果(正向或封锁)后,更新 TodoWrite,回到组合看下一条;有新事实催生了新方向就补进组合。
|
||||
|
||||
━━ 记录规约(边做边写,写对地方)━━
|
||||
- 每得出一个结果**立刻**落地,别攒到最后(会话步数耗尽就全丢;记下来的才算数,活在脑子里的不算)。这些记录也是你抗 compaction 的长期记忆。
|
||||
- **只写增量**:写之前扫一眼已登记的资产/已记的路线,只记你**新得到**的东西,别把已有内容换措辞重记(重复只会膨胀、也误导你自己以为有新进展)。只是印证已有结论而无新增,就不必再记。
|
||||
- **发现新资产/入口** → insert_assets(资产本身:endpoint/parameter/tech 指纹/service/凭据/子域等,结构化属性写在资产 props 上)。回看已登记资产用 list_assets,避免重复登记。
|
||||
- **确认漏洞** → report_finding(含可复现 PoC)。**只有你在本次运行里真实触发过、拿到可复现证据(请求/响应或命令输出)才用它**;回看已报漏洞用 list_findings。严禁把仅凭版本/CVE 匹配、"看起来可注入"、外部漏洞库/更新日志/代码 diff 推断的东西当已确认漏洞上报。**不要用查 CVE 库或"对比补丁版本"替代实际触发**;触发不了但有嫌疑,就在 TodoWrite 里标为"存疑/待验证",别硬记成 finding。
|
||||
|
||||
━━ 判定与收尾 ━━
|
||||
- 随时对照任务目标:已被你**验证过**的成果满足了目标,就据此判定达成并说明依据。判"达成"的前提是心法 4 的自检已通过——没独立复现过的战果不算达成依据。
|
||||
- **收尾优先级最高**:当你收到收尾信号(或自判目标已达成/所有合理路线已探尽),**立即停止一切探测与命令**,把手里的结论落地、给出简洁总结即可——此时"继续探索/再试一次/穷尽这条链/等命令结果"等一切先前指令都被收尾覆盖,不要再启动新动作。
|
||||
- 总结用人话讲清:达成了什么、走了哪些路线、确认了哪些漏洞(附 PoC 位置)、哪些方向已封锁及原因。只讲真实做到的,不臆造。
|
||||
|
||||
务实、克制、彻底。宁可把一条路走透并验证,也不要浅尝辄止地铺一堆没验证的"疑似"。`
|
||||
|
||||
// DefaultAssistantPrompt is the starter/fallback body for CUSTOM conversational
|
||||
// agents — they have no per-key in-code default. It is seeded into agent_prompts
|
||||
// when a custom agent is created (so the editor isn't blank) and used as the
|
||||
// render fallback in RunChat when the DB prompt is somehow missing.
|
||||
const DefaultAssistantPrompt = `你是一个乐于助人的 AI 助手。请用简洁、准确的中文回答用户的问题;在需要时使用可用的工具来完成任务。只做用户要求的事,不臆造信息。`
|
||||
|
||||
// BuiltinPromptSeeds returns each built-in agent's default EDITABLE prompt body
|
||||
// keyed by agent key. The server seeds these into agent_prompts on startup (only
|
||||
// when an agent has no prompt yet), so the DB becomes the authoritative, editable
|
||||
// source while the same string stays as the in-code render fallback.
|
||||
func BuiltinPromptSeeds() map[string]string {
|
||||
return map[string]string{
|
||||
"goals": goalsDefaultTmpl,
|
||||
"planner": plannerDefaultTmpl,
|
||||
"mainagent": mainAgentDefaultTmpl,
|
||||
"worker": workerDefaultTmpl,
|
||||
"auto": autoDefaultTmpl,
|
||||
"pentest": pentestDefaultTmpl,
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,212 @@
|
||||
// Package agent wires real LLM-driven planner and work agents (on top of the
|
||||
// agent-core SDK) to the dual SQLite graph. See docs/ARTEX-架构设计.md
|
||||
// §4.3 (planner) and §4.4 (work agent).
|
||||
//
|
||||
// Provider configuration is read from the environment so the system runs with
|
||||
// any Anthropic- or OpenAI-format endpoint. If no key is configured, FromEnv
|
||||
// returns ok=false and the exploration engine stays idle (an LLM is required).
|
||||
package agent
|
||||
|
||||
import (
|
||||
"context"
|
||||
"os"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/Autumn-27/norma/agentcore"
|
||||
"github.com/Autumn-27/norma/compaction"
|
||||
"github.com/Autumn-27/norma/llm"
|
||||
acperm "github.com/Autumn-27/norma/permission"
|
||||
)
|
||||
|
||||
// Config describes the LLM backend resolved from the environment.
|
||||
type Config struct {
|
||||
Format llm.Format
|
||||
BaseURL string
|
||||
APIKey string
|
||||
Model string
|
||||
// Proxy routes all LLM requests through the given proxy URL (http/https/socks5).
|
||||
// Empty falls back to the standard *_PROXY environment variables.
|
||||
Proxy string
|
||||
// RatePerSecond / RatePerMinute cap the shared request rate across ALL agents
|
||||
// using the provider (0 = that window unlimited).
|
||||
RatePerSecond float64
|
||||
RatePerMinute float64
|
||||
// ContextWindowK is the model's context window in K tokens (user-configured),
|
||||
// used to size compaction thresholds. 0 = default; see CompactionWindow.
|
||||
ContextWindowK int
|
||||
// ReasoningEffort selects the thinking mode: "" = 默认(不发送思考参数);
|
||||
// "off" = 关闭(thinking.type=disabled); "low"/"medium"/"high"/"max" = 开启并设强度.
|
||||
// NewProvider derives the provider-specific request fields from it.
|
||||
ReasoningEffort string
|
||||
}
|
||||
|
||||
// compaction window resolution bounds (in K tokens). Below the floor the
|
||||
// threshold math (window − summary reserve − buffer) would go non-positive and
|
||||
// compaction would fire every turn; above the cap it would never fire.
|
||||
const (
|
||||
defaultWindowK = 200 // unset → assume a 200K window (Claude default)
|
||||
minWindowK = 32 // floor so effectiveWindow stays comfortably positive
|
||||
maxWindowK = 1000 // cap at 1M tokens (user request)
|
||||
)
|
||||
|
||||
// CompactionWindow returns the model context window in TOKENS for compaction
|
||||
// thresholds, resolved from the user-configured size (ContextWindowK). 0/unset →
|
||||
// a 200K default; otherwise clamped to [32K, 1M] so compaction stays effective.
|
||||
func (c Config) CompactionWindow() int {
|
||||
k := c.ContextWindowK
|
||||
if k <= 0 {
|
||||
k = defaultWindowK
|
||||
}
|
||||
if k < minWindowK {
|
||||
k = minWindowK
|
||||
}
|
||||
if k > maxWindowK {
|
||||
k = maxWindowK
|
||||
}
|
||||
return k * 1000
|
||||
}
|
||||
|
||||
// compactionConfig builds the agent-core compaction config for a context window
|
||||
// in tokens. agentcore.NewSession wires the summarizer (same provider) when this
|
||||
// is set on Options.Compaction.
|
||||
func compactionConfig(windowTokens int) *compaction.Config {
|
||||
if windowTokens <= 0 {
|
||||
windowTokens = defaultWindowK * 1000
|
||||
}
|
||||
return &compaction.Config{ContextWindow: windowTokens}
|
||||
}
|
||||
|
||||
// FromEnv resolves the LLM provider config:
|
||||
//
|
||||
// ARTEX_LLM_PROVIDER = anthropic|openai (default: inferred from keys)
|
||||
// ARTEX_LLM_MODEL = model id (default: per provider)
|
||||
// ARTEX_LLM_BASE_URL = endpoint (optional)
|
||||
// ARTEX_LLM_PROXY = proxy URL (optional; http/https/socks5)
|
||||
// ANTHROPIC_API_KEY / OPENAI_API_KEY = credentials
|
||||
func FromEnv() (Config, bool) {
|
||||
prov := os.Getenv("ARTEX_LLM_PROVIDER")
|
||||
anthKey := os.Getenv("ANTHROPIC_API_KEY")
|
||||
oaiKey := os.Getenv("OPENAI_API_KEY")
|
||||
|
||||
if prov == "" {
|
||||
switch {
|
||||
case anthKey != "":
|
||||
prov = "anthropic"
|
||||
case oaiKey != "":
|
||||
prov = "openai"
|
||||
default:
|
||||
return Config{}, false
|
||||
}
|
||||
}
|
||||
|
||||
c := Config{
|
||||
BaseURL: os.Getenv("ARTEX_LLM_BASE_URL"),
|
||||
Model: os.Getenv("ARTEX_LLM_MODEL"),
|
||||
Proxy: strings.TrimSpace(os.Getenv("ARTEX_LLM_PROXY")),
|
||||
}
|
||||
switch prov {
|
||||
case "openai":
|
||||
c.Format = llm.FormatOpenAI
|
||||
c.APIKey = oaiKey
|
||||
if c.Model == "" {
|
||||
c.Model = "gpt-4o"
|
||||
}
|
||||
default:
|
||||
c.Format = llm.FormatAnthropic
|
||||
c.APIKey = anthKey
|
||||
if c.Model == "" {
|
||||
c.Model = "claude-opus-4-8"
|
||||
}
|
||||
}
|
||||
if c.APIKey == "" {
|
||||
return Config{}, false
|
||||
}
|
||||
return c, true
|
||||
}
|
||||
|
||||
// ConfigFrom builds a Config from UI-provided strings (provider defaults to
|
||||
// anthropic; model defaults per provider). Inputs are trimmed and the base URL
|
||||
// is normalized to the API base the provider expects (the provider appends the
|
||||
// endpoint path itself), so a full endpoint URL is tolerated.
|
||||
func ConfigFrom(provider, model, baseURL, apiKey, proxy string) Config {
|
||||
c := Config{
|
||||
Model: strings.TrimSpace(model),
|
||||
BaseURL: strings.TrimRight(strings.TrimSpace(baseURL), "/"),
|
||||
APIKey: strings.TrimSpace(apiKey),
|
||||
Proxy: strings.TrimSpace(proxy),
|
||||
}
|
||||
switch strings.TrimSpace(provider) {
|
||||
case "openai":
|
||||
c.Format = llm.FormatOpenAI
|
||||
// provider appends "/chat/completions"; tolerate a full endpoint URL.
|
||||
c.BaseURL = strings.TrimRight(strings.TrimSuffix(c.BaseURL, "/chat/completions"), "/")
|
||||
if c.Model == "" {
|
||||
c.Model = "gpt-4o"
|
||||
}
|
||||
default:
|
||||
c.Format = llm.FormatAnthropic
|
||||
// provider appends "/v1/messages".
|
||||
c.BaseURL = strings.TrimRight(strings.TrimSuffix(c.BaseURL, "/v1/messages"), "/")
|
||||
if c.Model == "" {
|
||||
c.Model = "claude-opus-4-8"
|
||||
}
|
||||
}
|
||||
return c
|
||||
}
|
||||
|
||||
// Provider returns the short provider name ("anthropic"/"openai").
|
||||
func (c Config) Provider() string {
|
||||
if c.Format == llm.FormatOpenAI {
|
||||
return "openai"
|
||||
}
|
||||
return "anthropic"
|
||||
}
|
||||
|
||||
// NewProvider builds an llm.Provider from the config. When a rate is set, the
|
||||
// limiter lives on the single provider instance — so planner + all workers +
|
||||
// main agent (which share this provider) are bounded by one shared rate limit.
|
||||
func (c Config) NewProvider() (llm.Provider, error) {
|
||||
lc := llm.Config{
|
||||
Format: c.Format,
|
||||
BaseURL: c.BaseURL,
|
||||
APIKey: c.APIKey,
|
||||
Model: c.Model,
|
||||
Proxy: c.Proxy,
|
||||
}
|
||||
// Derive the provider thinking params from the single UI selector:
|
||||
// "" → omit both (provider default); "off" → disable; else → enable + effort.
|
||||
switch c.ReasoningEffort {
|
||||
case "":
|
||||
// leave both empty → omitted from the request
|
||||
case "off":
|
||||
lc.ThinkingType = "disabled"
|
||||
default:
|
||||
lc.ThinkingType = "enabled"
|
||||
lc.ReasoningEffort = c.ReasoningEffort
|
||||
}
|
||||
if c.RatePerSecond > 0 || c.RatePerMinute > 0 {
|
||||
lc.RateLimit = &llm.RateLimit{PerSecond: c.RatePerSecond, PerMinute: c.RatePerMinute}
|
||||
}
|
||||
return llm.NewProvider(lc)
|
||||
}
|
||||
|
||||
// TestConnection makes a minimal real completion to verify the provider/model/
|
||||
// endpoint/key actually work. Returns the round-trip latency.
|
||||
func TestConnection(ctx context.Context, c Config) (time.Duration, error) {
|
||||
prov, err := c.NewProvider()
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
ctx, cancel := context.WithTimeout(ctx, 30*time.Second)
|
||||
defer cancel()
|
||||
start := time.Now()
|
||||
_, err = agentcore.Run(ctx, agentcore.Options{
|
||||
Provider: prov,
|
||||
SystemPrompt: []string{"你是连接测试。只回复 OK,不要别的。"},
|
||||
PermissionMode: acperm.ModeBypass,
|
||||
MaxTurns: 1,
|
||||
MaxTokens: 32,
|
||||
}, "ping")
|
||||
return time.Since(start), err
|
||||
}
|
||||
@@ -0,0 +1,54 @@
|
||||
package agent
|
||||
|
||||
import (
|
||||
"context"
|
||||
"time"
|
||||
)
|
||||
|
||||
// TaskClock carries a task's absolute deadline into a worker/planner run so the run
|
||||
// can clamp its own wall-clock budget to the task's remaining time and pick the
|
||||
// right wrap-up words (per-run vs task-timeout). Attached to the run ctx by the
|
||||
// engine. Zero value = no task-level timeout (behaves exactly as before).
|
||||
type TaskClock struct {
|
||||
DeadlineUnix int64 // absolute deadline (unix seconds); 0 = no task timeout
|
||||
Final bool // coordinator-driven FINAL planner round (task ending now)
|
||||
}
|
||||
|
||||
type taskClockKey struct{}
|
||||
|
||||
// WithTaskClock attaches a TaskClock to ctx for the run.
|
||||
func WithTaskClock(ctx context.Context, tc TaskClock) context.Context {
|
||||
return context.WithValue(ctx, taskClockKey{}, tc)
|
||||
}
|
||||
|
||||
// taskClockFrom reads the TaskClock (zero value if none attached).
|
||||
func taskClockFrom(ctx context.Context) TaskClock {
|
||||
if v, ok := ctx.Value(taskClockKey{}).(TaskClock); ok {
|
||||
return v
|
||||
}
|
||||
return TaskClock{}
|
||||
}
|
||||
|
||||
// clampMaxDuration folds a task deadline into a run's own wall-clock budget.
|
||||
// - ownBudget = the agent's own run_seconds (0 = unlimited).
|
||||
// - Returns eff = the MaxDuration to use (floored at 1s so we never pass ≤0, which
|
||||
// harness reads as "unlimited"), and clamped = whether the TASK deadline is the
|
||||
// binding constraint (remaining ≤ ownBudget, or ownBudget unlimited). When there
|
||||
// is no deadline, returns (ownBudget, false) unchanged.
|
||||
func clampMaxDuration(deadlineUnix int64, ownBudget time.Duration) (eff time.Duration, clamped bool) {
|
||||
if deadlineUnix <= 0 {
|
||||
return ownBudget, false
|
||||
}
|
||||
remaining := time.Until(time.Unix(deadlineUnix, 0))
|
||||
if remaining < time.Second {
|
||||
remaining = time.Second // max(1, …): never pass ≤0 (harness treats 0 as unlimited)
|
||||
}
|
||||
// clamped when the task deadline binds this run's time: remaining ≤ own budget,
|
||||
// or the agent has no own time budget (then remaining always binds).
|
||||
clamped = ownBudget <= 0 || remaining <= ownBudget
|
||||
eff = remaining
|
||||
if ownBudget > 0 && ownBudget < remaining {
|
||||
eff = ownBudget
|
||||
}
|
||||
return eff, clamped
|
||||
}
|
||||
@@ -0,0 +1,30 @@
|
||||
package agent
|
||||
|
||||
import (
|
||||
"database/sql"
|
||||
"os"
|
||||
"testing"
|
||||
|
||||
"github.com/Autumn-27/artex/db"
|
||||
_ "github.com/jackc/pgx/v5/stdlib"
|
||||
)
|
||||
|
||||
// TestMain acquires a PostgreSQL advisory lock (7337741002) for the entire
|
||||
// agent test suite so cross-package DELETE cleanup races with db/server
|
||||
// packages are avoided when running `go test ./...`.
|
||||
func TestMain(m *testing.M) {
|
||||
dsn, _, err := db.DSN()
|
||||
if err != nil {
|
||||
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())
|
||||
}
|
||||
@@ -0,0 +1,185 @@
|
||||
package agent
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
|
||||
actool "github.com/Autumn-27/norma/tool"
|
||||
)
|
||||
|
||||
// 本文件把「内置工具」从纯代码变成可枚举、可被 DB 覆盖的目录:
|
||||
// - BuiltinToolSeeds():把三个执行 agent 的内置工具集展开成 seed 记录(key +
|
||||
// 描述 + 参数 schema + 默认绑定的 agent),供服务端开机幂等播种进 tools 表。
|
||||
// - ToolResolve 钩子:运行时按 DB 里的 tools 行对已装配的工具做「按 agent 过滤 +
|
||||
// 覆盖描述/schema + 注入参数默认值」。key/handler 仍在代码层,DB 只改「散文与默认值」。
|
||||
// handler(Call 行为)永远来自代码——DB 改不了它,只能改模型看到的说明与缺省入参。
|
||||
|
||||
// ToolSeed 是一个内置工具的可播种快照:key 即 CoreTool.Name()(与 handler 死绑,
|
||||
// UI 只读),Desc/Schema 取自代码里的工具定义,Agents 是代码默认把它给了哪些 agent。
|
||||
type ToolSeed struct {
|
||||
Key string // = CoreTool.Name(),主键,不可改
|
||||
Desc string // 顶层描述(可在 UI 覆盖)
|
||||
Schema map[string]any // 参数 JSON-Schema(结构只读,description/default 可在 UI 改)
|
||||
Agents []string // 默认绑定的 agent key(worker/planner/mainagent)
|
||||
}
|
||||
|
||||
// builtinToolsByAgent 用一个「只读空壳」ToolSet(nil stores)构造每个执行 agent 的
|
||||
// 领域工具集。工具构造函数只把闭包塞进 Spec、构造期不解引用 store,所以 nil 安全——
|
||||
// 这些工具在这里只用来读 Name()/Description()/InputSchema(),绝不 Call。
|
||||
//
|
||||
// 刻意【不含】SDK 通用工具 actool.DefaultTools()(Read/Write/Edit/MultiEdit/LS/Glob/
|
||||
// Grep/Bash):它们每个 agent 都固定拥有、没有「绑定到谁」的取舍,且说明大多在 Prompt()
|
||||
// 里(本表只覆盖 Description(),会造成半覆盖误导)。不 seed → 无 DB 行 → ToolResolve
|
||||
// 原样放行、不覆盖,行为与从前一致。只有 artex 自己的领域工具入表可管。
|
||||
func builtinToolsByAgent() map[string][]actool.CoreTool {
|
||||
ts := NewToolSet(nil, "")
|
||||
return map[string][]actool.CoreTool{
|
||||
"mainagent": ts.MainAgentTools(),
|
||||
"planner": ts.PlannerTools(),
|
||||
"worker": ts.WorkerTools(),
|
||||
// auto 默认绑漏洞上报 + 资产管理工具,其他域工具可在 UI 按需勾选。
|
||||
// 新库由此 seed 写入;老库由 seedAutoDefaultBindings 迁移。
|
||||
"auto": {ts.addFinding(), ts.insertAssets(), ts.addCompanyScope(), ts.listAssets(), ts.listCompanies()},
|
||||
// pentest(独立渗透 agent)默认绑:查资产 / 插资产 / 报漏洞 / 查漏洞 / 查企业。
|
||||
// 新库由此 seed 写入;老库由 seedPentestDefaultBindings 迁移。
|
||||
"pentest": {ts.listAssets(), ts.insertAssets(), ts.addFinding(), ts.listFindings(), ts.listCompanies()},
|
||||
}
|
||||
}
|
||||
|
||||
// BuiltinToolSeeds 把三个 agent 的内置工具集去重合并成 seed 列表:同名工具(如
|
||||
// list_assets 三者都有)合成一条,Agents 取并集。顺序稳定(mainagent→planner→worker)。
|
||||
func BuiltinToolSeeds() []ToolSeed {
|
||||
byAgent := builtinToolsByAgent()
|
||||
order := []string{"mainagent", "planner", "worker", "auto", "pentest"}
|
||||
|
||||
type acc struct {
|
||||
tool actool.CoreTool
|
||||
agents []string
|
||||
}
|
||||
m := map[string]*acc{}
|
||||
var keys []string
|
||||
for _, ak := range order {
|
||||
for _, t := range byAgent[ak] {
|
||||
a, ok := m[t.Name()]
|
||||
if !ok {
|
||||
a = &acc{tool: t}
|
||||
m[t.Name()] = a
|
||||
keys = append(keys, t.Name())
|
||||
}
|
||||
a.agents = append(a.agents, ak)
|
||||
}
|
||||
}
|
||||
|
||||
out := make([]ToolSeed, 0, len(keys))
|
||||
for _, k := range keys {
|
||||
a := m[k]
|
||||
out = append(out, ToolSeed{
|
||||
Key: k,
|
||||
Desc: a.tool.Description(),
|
||||
Schema: a.tool.InputSchema(),
|
||||
Agents: a.agents,
|
||||
})
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// ToolResolve, if set, post-processes an agent's fully-assembled tool list against
|
||||
// the DB tools table: it drops tools not bound to this agent (or globally disabled)
|
||||
// and wraps the rest so the model sees the DB-overridden description/schema and
|
||||
// 缺省入参 get injected. Tools with no matching DB row (MCP/skill/host tools like
|
||||
// traffic) pass through untouched. nil = tools unchanged. Wired in server/assembly.go.
|
||||
var ToolResolve func(ctx context.Context, agentKey string, tools []actool.CoreTool) []actool.CoreTool
|
||||
|
||||
// DecorateTool wraps t so Description()/InputSchema() report the DB overrides and
|
||||
// Call() injects scalar parameter defaults (from schema's "default" props) whenever
|
||||
// the model omitted them. Name/Prompt/permission/scheduler flags delegate to t, so
|
||||
// the tool's identity and handler are unchanged. Empty desc/schema fall back to t's.
|
||||
func DecorateTool(t actool.CoreTool, desc string, schema map[string]any) actool.CoreTool {
|
||||
if desc == "" {
|
||||
desc = t.Description()
|
||||
}
|
||||
if len(schema) == 0 {
|
||||
schema = t.InputSchema()
|
||||
}
|
||||
return &overriddenTool{CoreTool: t, desc: desc, schema: schema}
|
||||
}
|
||||
|
||||
// overriddenTool is a CoreTool decorator: it embeds the original (so all behavioral
|
||||
// methods — Prompt/IsReadOnly/IsConcurrencySafe/CheckPermissions/Name — delegate)
|
||||
// and overrides only the model-facing description/schema plus default injection.
|
||||
type overriddenTool struct {
|
||||
actool.CoreTool
|
||||
desc string
|
||||
schema map[string]any
|
||||
}
|
||||
|
||||
func (o *overriddenTool) Description() string { return o.desc }
|
||||
func (o *overriddenTool) InputSchema() map[string]any { return o.schema }
|
||||
|
||||
func (o *overriddenTool) Call(ctx context.Context, in json.RawMessage, tc *actool.ToolContext) (actool.Result, error) {
|
||||
return o.CoreTool.Call(ctx, injectDefaults(in, o.schema), tc)
|
||||
}
|
||||
|
||||
// injectDefaults fills scalar parameter defaults declared in the (possibly edited)
|
||||
// schema into the input JSON whenever the model omitted the field or left it empty/
|
||||
// null. Structure (names/types/required) is untouched — only缺省值 are merged in.
|
||||
func injectDefaults(in json.RawMessage, schema map[string]any) json.RawMessage {
|
||||
defs := scalarDefaults(schema)
|
||||
if len(defs) == 0 {
|
||||
return in
|
||||
}
|
||||
m := map[string]json.RawMessage{}
|
||||
if len(in) > 0 {
|
||||
if err := json.Unmarshal(in, &m); err != nil {
|
||||
return in // non-object input: don't touch it
|
||||
}
|
||||
}
|
||||
changed := false
|
||||
for k, dv := range defs {
|
||||
if cur, ok := m[k]; !ok || isEmptyJSON(cur) {
|
||||
m[k] = dv
|
||||
changed = true
|
||||
}
|
||||
}
|
||||
if !changed {
|
||||
return in
|
||||
}
|
||||
b, err := json.Marshal(m)
|
||||
if err != nil {
|
||||
return in
|
||||
}
|
||||
return b
|
||||
}
|
||||
|
||||
// scalarDefaults extracts properties[k]["default"] for scalar params (string/
|
||||
// integer/number/boolean). Array/object defaults are skipped: merging them is
|
||||
// ambiguous and not worth the surprise.
|
||||
func scalarDefaults(schema map[string]any) map[string]json.RawMessage {
|
||||
props, _ := schema["properties"].(map[string]any)
|
||||
if len(props) == 0 {
|
||||
return nil
|
||||
}
|
||||
out := map[string]json.RawMessage{}
|
||||
for name, raw := range props {
|
||||
p, ok := raw.(map[string]any)
|
||||
if !ok {
|
||||
continue
|
||||
}
|
||||
dv, ok := p["default"]
|
||||
if !ok || dv == nil {
|
||||
continue
|
||||
}
|
||||
switch p["type"] {
|
||||
case "string", "integer", "number", "boolean":
|
||||
if b, err := json.Marshal(dv); err == nil {
|
||||
out[name] = b
|
||||
}
|
||||
}
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
func isEmptyJSON(raw json.RawMessage) bool {
|
||||
s := string(raw)
|
||||
return s == "null" || s == `""`
|
||||
}
|
||||
@@ -0,0 +1,73 @@
|
||||
package agent
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"testing"
|
||||
|
||||
actool "github.com/Autumn-27/norma/tool"
|
||||
)
|
||||
|
||||
// TestBuiltinToolSeeds ensures the catalog builds from a nil-store ToolSet without
|
||||
// panicking, has stable keys, unions agent bindings, and carries real descriptions.
|
||||
func TestBuiltinToolSeeds(t *testing.T) {
|
||||
seeds := BuiltinToolSeeds()
|
||||
if len(seeds) == 0 {
|
||||
t.Fatal("no seeds")
|
||||
}
|
||||
byKey := map[string]ToolSeed{}
|
||||
for _, s := range seeds {
|
||||
if s.Key == "" || s.Desc == "" {
|
||||
t.Errorf("seed %q missing key/desc", s.Key)
|
||||
}
|
||||
byKey[s.Key] = s
|
||||
}
|
||||
// record_fact is worker-only.
|
||||
if rf, ok := byKey["record_fact"]; !ok || len(rf.Agents) != 1 || rf.Agents[0] != "worker" {
|
||||
t.Errorf("record_fact agents = %v, want [worker]", rf.Agents)
|
||||
}
|
||||
// SDK generic tools (incl. sleep, now part of DefaultTools) are deliberately NOT
|
||||
// seeded — every agent owns them; they flow through ToolResolve untouched.
|
||||
for _, k := range []string{"Bash", "Read", "Write", "Edit", "Grep", "sleep"} {
|
||||
if _, ok := byKey[k]; ok {
|
||||
t.Errorf("SDK tool %q should not be seeded", k)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// TestDecorateToolInjectsDefaults verifies a schema "default" fills a missing param
|
||||
// before the underlying handler runs, and an explicitly-provided value is kept.
|
||||
func TestDecorateToolInjectsDefaults(t *testing.T) {
|
||||
var seen map[string]any
|
||||
base := actool.Build(actool.Spec{
|
||||
Name: "probe", Description: "orig",
|
||||
Run: func(_ context.Context, in json.RawMessage, _ *actool.ToolContext) (actool.Result, error) {
|
||||
_ = json.Unmarshal(in, &seen)
|
||||
return actool.Text("ok"), nil
|
||||
},
|
||||
})
|
||||
schema := map[string]any{"type": "object", "properties": map[string]any{
|
||||
"limit": map[string]any{"type": "integer", "description": "n", "default": float64(3)},
|
||||
"q": map[string]any{"type": "string", "description": "query"},
|
||||
}}
|
||||
dec := DecorateTool(base, "new desc", schema)
|
||||
if dec.Description() != "new desc" {
|
||||
t.Errorf("description = %q", dec.Description())
|
||||
}
|
||||
|
||||
// limit omitted → default 3 injected; q kept.
|
||||
if _, err := dec.Call(context.Background(), json.RawMessage(`{"q":"x"}`), nil); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if seen["limit"] != float64(3) || seen["q"] != "x" {
|
||||
t.Errorf("injected = %v, want limit=3 q=x", seen)
|
||||
}
|
||||
|
||||
// limit provided → default does NOT override.
|
||||
if _, err := dec.Call(context.Background(), json.RawMessage(`{"limit":9}`), nil); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if seen["limit"] != float64(9) {
|
||||
t.Errorf("limit = %v, want 9 (no override)", seen["limit"])
|
||||
}
|
||||
}
|
||||
+1057
File diff suppressed because it is too large.
Load diff
@@ -0,0 +1,510 @@
|
||||
package agent
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"strings"
|
||||
|
||||
actool "github.com/Autumn-27/norma/tool"
|
||||
"github.com/Autumn-27/artex/db"
|
||||
)
|
||||
|
||||
// =====================================================================
|
||||
// Unified asset insertion tools
|
||||
// =====================================================================
|
||||
|
||||
// SetAssetStore wires the asset store and company store onto this ToolSet
|
||||
// so the insert_assets, add_company_scope, and list_assets tools are active.
|
||||
func (t *ToolSet) SetAssetStore(as *db.AssetStore, cs *db.CompanyStore) {
|
||||
t.as = as
|
||||
t.cs = cs
|
||||
}
|
||||
|
||||
// assetInputItem is one element of the insert_assets "assets" array.
|
||||
type assetInputItem struct {
|
||||
Type string `json:"type"` // root_domain|ip|subdomain|app|service|endpoint
|
||||
|
||||
// ---- root_domain / subdomain ----
|
||||
Domain string `json:"domain"`
|
||||
ICP string `json:"icp"`
|
||||
RecordType string `json:"record_type"`
|
||||
RecordValue []string `json:"record_value"`
|
||||
|
||||
// ---- ip ----
|
||||
IP string `json:"ip"`
|
||||
BoundDomains []string `json:"bound_domains"`
|
||||
OpenPorts []db.PortService `json:"open_ports"`
|
||||
|
||||
// ---- app ----
|
||||
AppName string `json:"app_name"`
|
||||
BundleID string `json:"bundle_id"`
|
||||
Category string `json:"category"`
|
||||
Description string `json:"description"`
|
||||
AppICP string `json:"app_icp"`
|
||||
CompanyID *int64 `json:"company_id"` // explicit company link (app only; others auto-attribute via scope)
|
||||
|
||||
// ---- service (http) ----
|
||||
URL string `json:"url"`
|
||||
Technologies []string `json:"technologies"`
|
||||
StatusCode *int `json:"status_code"`
|
||||
ContentLength *int64 `json:"content_length"`
|
||||
PageTitle string `json:"page_title"`
|
||||
FaviconMMH3 string `json:"favicon_mmh3"`
|
||||
Auth []map[string]any `json:"auth"`
|
||||
ServiceName string `json:"service_name"`
|
||||
ServiceIP string `json:"service_ip"` // optional enrichment IP
|
||||
|
||||
// ---- service (other) ----
|
||||
Port int `json:"port"`
|
||||
Proto string `json:"proto"`
|
||||
|
||||
// ---- endpoint ----
|
||||
Method string `json:"method"`
|
||||
Params []map[string]any `json:"params"`
|
||||
}
|
||||
|
||||
// insertAssets is the unified insert_assets agent tool.
|
||||
func (t *ToolSet) insertAssets() actool.CoreTool {
|
||||
return writeTool(
|
||||
"insert_assets",
|
||||
"发现新资产时立即批量插入资产到资产库。支持一次提交多种资产类型:root_domain/ip/subdomain/app/service/endpoint。\n"+
|
||||
"每条资产的字段说明:\n"+
|
||||
"• root_domain:domain(根域名, 必填)、icp(备案, 可选)\n"+
|
||||
"• ip:ip(IP地址, 必填)、bound_domains([域名])、open_ports([{port,service}])\n"+
|
||||
"• subdomain:domain(子域名, 必填)、record_type(A/AAAA/CNAME等)、record_value(记录值)、icp(可选)\n"+
|
||||
"• app:app_name(应用名, 必填)、bundle_id(bundle id, 可选)、category、description、app_icp、company_id(归属企业 id, 可选)\n"+
|
||||
"• service(http):url(必填)、technologies([指纹])、status_code、content_length、page_title、favicon_mmh3、auth([{type,username,password,...}])、service_ip\n"+
|
||||
"• service(other):service_name(必填)、ip或domain(至少一个)、port(必填)、auth([...])\n"+
|
||||
"• endpoint:url(必填)、method(必填)、params([{location,name,value,type}])、service_ip\n"+
|
||||
"auth/technologies/params 都是【追加合并】(append),不会覆盖原有值。\n"+
|
||||
"返回:{results:[{index,id,type}], errors:[{index,error}]}",
|
||||
obj(map[string]any{
|
||||
"task_id": intp("任务 id(可选,用于关联 task_ids)"),
|
||||
"assets": map[string]any{
|
||||
"type": "array",
|
||||
"description": "资产数组,每个元素对应一条资产记录",
|
||||
"items": obj(map[string]any{
|
||||
"type": map[string]any{
|
||||
"type": "string",
|
||||
"enum": []string{"root_domain", "ip", "subdomain", "app", "service", "endpoint"},
|
||||
"description": "资产类型",
|
||||
},
|
||||
// root_domain / subdomain
|
||||
"domain": str("根域名或子域名(root_domain/subdomain 必填)"),
|
||||
"icp": str("ICP 备案号(可选)"),
|
||||
"record_type": str("DNS 解析类型:A/AAAA/CNAME/MX 等(subdomain 可选)"),
|
||||
"record_value": map[string]any{
|
||||
"type": "array",
|
||||
"items": map[string]any{"type": "string"},
|
||||
"description": "DNS 解析值列表(subdomain 可选,如 [\"1.2.3.4\",\"2.3.4.5\"])",
|
||||
},
|
||||
// ip
|
||||
"ip": str("IP 地址(ip 类型必填;service/endpoint 类型可填,用于关联 IP)"),
|
||||
"bound_domains": map[string]any{
|
||||
"type": "array",
|
||||
"items": map[string]any{"type": "string"},
|
||||
"description": "该 IP 绑定的域名列表(ip 类型可选)",
|
||||
},
|
||||
"open_ports": map[string]any{
|
||||
"type": "array",
|
||||
"description": "开放端口列表(ip 类型可选)",
|
||||
"items": obj(map[string]any{
|
||||
"port": intp("端口号"),
|
||||
"service": str("服务名称,如 http/ssh/mysql 等(可选)"),
|
||||
}, "port"),
|
||||
},
|
||||
// app
|
||||
"app_name": str("应用名称(app 类型必填)"),
|
||||
"bundle_id": str("Bundle ID(app 类型可选)"),
|
||||
"category": str("应用分类(可选)"),
|
||||
"description": str("应用描述(可选)"),
|
||||
"app_icp": str("应用 ICP 备案(可选)"),
|
||||
"company_id": intp("归属企业 id(app 类型可选;app 无法靠 scope 自动归因,需显式指定。id 由 add_company_scope 返回)"),
|
||||
// service (http)
|
||||
"url": str("完整 URL,含协议和端口(HTTP 服务必填;service_type 自动设为 http)"),
|
||||
"status_code": intp("HTTP 响应状态码,如 200/301/403/404(可选)"),
|
||||
"content_length": map[string]any{
|
||||
"type": "integer",
|
||||
"description": "HTTP 响应体字节数(可选)",
|
||||
},
|
||||
"page_title": str("页面 <title> 内容(可选)"),
|
||||
"favicon_mmh3": str("favicon MMH3 哈希(可选)"),
|
||||
"technologies": map[string]any{
|
||||
"type": "array",
|
||||
"items": map[string]any{"type": "string"},
|
||||
"description": "指纹/技术栈列表,如 [\"Nginx\",\"Vue\",\"Bootstrap\"](可选)",
|
||||
},
|
||||
"auth": map[string]any{
|
||||
"type": "array",
|
||||
"description": "发现的认证信息列表,每条含 type/username/password 等字段(可选,追加不覆盖)",
|
||||
"items": map[string]any{"type": "object"},
|
||||
},
|
||||
// service (other,非 HTTP)
|
||||
"service_name": str("服务名称,如 ssh/mysql/redis(service 非 HTTP 时必填)"),
|
||||
"port": intp("端口号(service 非 HTTP 时必填)"),
|
||||
// endpoint
|
||||
"method": str("HTTP 方法:GET/POST/PUT/PATCH/DELETE 等(endpoint 必填)"),
|
||||
"params": map[string]any{
|
||||
"type": "array",
|
||||
"description": "请求参数列表,每条含 location(query/body/header/path)/name/value/type(可选,追加不覆盖)",
|
||||
"items": map[string]any{"type": "object"},
|
||||
},
|
||||
}, "type"),
|
||||
},
|
||||
}, "assets"),
|
||||
func(_ context.Context, in json.RawMessage) (actool.Result, error) {
|
||||
if t.as == nil {
|
||||
return actool.Errorf("insert_assets 未启用: AssetStore 未初始化"), nil
|
||||
}
|
||||
var a struct {
|
||||
Assets []assetInputItem `json:"assets"`
|
||||
TaskID int64 `json:"task_id"`
|
||||
}
|
||||
if err := json.Unmarshal(in, &a); err != nil {
|
||||
return actool.Errorf("invalid input: " + err.Error()), nil
|
||||
}
|
||||
// use ToolSet.taskID as fallback
|
||||
taskID := a.TaskID
|
||||
if taskID == 0 {
|
||||
taskID = t.taskID
|
||||
}
|
||||
|
||||
type result struct {
|
||||
Index int `json:"index"`
|
||||
ID int64 `json:"id"`
|
||||
Type string `json:"type"`
|
||||
}
|
||||
type errEntry struct {
|
||||
Index int `json:"index"`
|
||||
Error string `json:"error"`
|
||||
}
|
||||
|
||||
var results []result
|
||||
var errs []errEntry
|
||||
|
||||
for i, item := range a.Assets {
|
||||
typ := strings.TrimSpace(item.Type)
|
||||
var id int64
|
||||
var err error
|
||||
|
||||
switch typ {
|
||||
case "root_domain":
|
||||
id, err = t.as.UpsertRootDomain(db.UpsertRootDomainReq{
|
||||
Domain: item.Domain,
|
||||
ICP: item.ICP,
|
||||
TaskID: taskID,
|
||||
})
|
||||
|
||||
case "ip":
|
||||
id, err = t.as.UpsertIP(db.UpsertIPReq{
|
||||
IP: item.IP,
|
||||
BoundDomains: item.BoundDomains,
|
||||
OpenPorts: item.OpenPorts,
|
||||
TaskID: taskID,
|
||||
})
|
||||
|
||||
case "subdomain":
|
||||
id, err = t.as.UpsertSubdomain(db.UpsertSubdomainReq{
|
||||
Domain: item.Domain,
|
||||
RecordType: item.RecordType,
|
||||
RecordValue: item.RecordValue,
|
||||
ICP: item.ICP,
|
||||
TaskID: taskID,
|
||||
})
|
||||
|
||||
case "app":
|
||||
id, err = t.as.UpsertApp(db.UpsertAppReq{
|
||||
Name: item.AppName,
|
||||
BundleID: item.BundleID,
|
||||
Category: item.Category,
|
||||
Description: item.Description,
|
||||
ICP: item.AppICP,
|
||||
CompanyID: item.CompanyID,
|
||||
TaskID: taskID,
|
||||
})
|
||||
|
||||
case "service":
|
||||
// distinguish HTTP vs other by presence of url
|
||||
if item.URL != "" {
|
||||
// agent may send "ip" or "service_ip" for the enrichment IP; accept both
|
||||
svcIP := item.ServiceIP
|
||||
if svcIP == "" {
|
||||
svcIP = item.IP
|
||||
}
|
||||
id, err = t.as.UpsertHTTPService(db.UpsertHTTPServiceReq{
|
||||
URL: item.URL,
|
||||
Technologies: item.Technologies,
|
||||
StatusCode: item.StatusCode,
|
||||
ContentLength: item.ContentLength,
|
||||
PageTitle: item.PageTitle,
|
||||
FaviconMMH3: item.FaviconMMH3,
|
||||
Auth: item.Auth,
|
||||
IP: svcIP,
|
||||
TaskID: taskID,
|
||||
})
|
||||
} else {
|
||||
id, err = t.as.UpsertOtherService(db.UpsertOtherServiceReq{
|
||||
Domain: item.Domain,
|
||||
IP: item.IP,
|
||||
Port: item.Port,
|
||||
ServiceName: item.ServiceName,
|
||||
Auth: item.Auth,
|
||||
TaskID: taskID,
|
||||
})
|
||||
}
|
||||
|
||||
case "endpoint":
|
||||
id, err = t.as.UpsertEndpoint(db.UpsertEndpointReq{
|
||||
URL: item.URL,
|
||||
Method: item.Method,
|
||||
Params: item.Params,
|
||||
IP: item.ServiceIP,
|
||||
TaskID: taskID,
|
||||
})
|
||||
|
||||
default:
|
||||
errs = append(errs, errEntry{Index: i, Error: "unknown type: " + typ})
|
||||
continue
|
||||
}
|
||||
|
||||
if err != nil {
|
||||
errs = append(errs, errEntry{Index: i, Error: err.Error()})
|
||||
continue
|
||||
}
|
||||
results = append(results, result{Index: i, ID: id, Type: typ})
|
||||
t.wrote++
|
||||
t.anchorOwner(id)
|
||||
}
|
||||
|
||||
return jsonResult(map[string]any{
|
||||
"results": results,
|
||||
"errors": errs,
|
||||
})
|
||||
},
|
||||
)
|
||||
}
|
||||
|
||||
// addCompanyScope writes to company_scope table and triggers asset attribution.
|
||||
func (t *ToolSet) addCompanyScope() actool.CoreTool {
|
||||
return writeTool(
|
||||
"add_company_scope",
|
||||
"把根域名/IP/IP段(CIDR)加入某公司的【资产范围】——加入后系统自动认领所有命中的资产(现有+未来)。\n"+
|
||||
"公司名唯一:company 不存在则新建,已存在则复用(只把范围并进去)。\n"+
|
||||
"scope 一行一条,支持:根域名 / 单个 IP / CIDR 网段(如 1.2.3.0/24)。\n"+
|
||||
"务必给 reason 说明归属依据(whois/证书/ASN 等)。\n"+
|
||||
"护栏:拒绝裸 TLD 与过宽网段(IPv4<=/16、IPv6<=/32),非法行会被跳过并在 errors 返回。",
|
||||
obj(map[string]any{
|
||||
"company": str("公司名(不存在则新建、存在则复用;名称唯一)"),
|
||||
"scope": str("资产范围,一行一条:根域名 / IP / CIDR"),
|
||||
"reason": str("归属依据(证据/来源),务必填写"),
|
||||
"logo": str("公司图标 URL(可选;仅新建公司时生效)"),
|
||||
}, "company", "scope"),
|
||||
func(_ context.Context, in json.RawMessage) (actool.Result, error) {
|
||||
if t.cs == nil {
|
||||
return actool.Errorf("add_company_scope 未启用: CompanyStore 未初始化"), nil
|
||||
}
|
||||
var a struct {
|
||||
Company string `json:"company"`
|
||||
Scope string `json:"scope"`
|
||||
Reason string `json:"reason"`
|
||||
Logo string `json:"logo"`
|
||||
}
|
||||
if err := json.Unmarshal(in, &a); err != nil {
|
||||
return actool.Errorf(err.Error()), nil
|
||||
}
|
||||
if strings.TrimSpace(a.Company) == "" {
|
||||
return actool.Errorf("company 不能为空"), nil
|
||||
}
|
||||
companyID, _, err := t.cs.UpsertCompany(a.Company, a.Logo)
|
||||
if err != nil {
|
||||
return actool.Errorf("创建/获取公司失败: " + err.Error()), nil
|
||||
}
|
||||
lines := splitLines(a.Scope)
|
||||
added, skipped, invalid, errMsgs := t.cs.AddScope(companyID, lines, a.Reason)
|
||||
out := map[string]any{
|
||||
"company_id": companyID,
|
||||
"added": added,
|
||||
"skipped": skipped,
|
||||
"invalid": invalid,
|
||||
}
|
||||
if len(errMsgs) > 0 {
|
||||
out["errors"] = errMsgs
|
||||
}
|
||||
return jsonResult(out)
|
||||
},
|
||||
)
|
||||
}
|
||||
|
||||
// listAssets lets an agent query the asset table.
|
||||
func (t *ToolSet) listAssets() actool.CoreTool {
|
||||
return readTool(
|
||||
"list_assets",
|
||||
"查询资产库。支持 DSL 表达式搜索、按 id/ids 直接取,分页。\n"+
|
||||
"DSL 语法:field=value 模糊(ILIKE) | field==value 精确 | field!=value 排除 | port>80 数字比较 | 裸词=全文模糊。\n"+
|
||||
"逻辑运算:AND / OR(关键字,AND 优先级高于 OR),括号分组。\n"+
|
||||
"资产类型用独立的 type 参数过滤,DSL 里不含 type 字段。\n"+
|
||||
"约束:未传 id/ids 时 dsl 必须非空——不允许无条件查询全部资产,必须带查询条件。\n"+
|
||||
"字段一览:\n"+
|
||||
" domain 域名(根域名/子域名/服务域名)\n"+
|
||||
" root_domain 根域名(仅子域名/服务资产有)\n"+
|
||||
" ip IPv4/IPv6 地址\n"+
|
||||
" url 完整 URL(服务/接口)\n"+
|
||||
" page_title 页面标题(HTTP 服务)\n"+
|
||||
" icp ICP 备案号(根域名)\n"+
|
||||
" service_name 服务名称(非 HTTP 服务)\n"+
|
||||
" app_name 应用名称(app 类型)\n"+
|
||||
" method HTTP 方法(接口类型,如 GET/POST)\n"+
|
||||
" service_type HTTP 服务类型:http|other\n"+
|
||||
" record_type DNS 解析类型(子域名,如 A/CNAME)\n"+
|
||||
" technology 指纹/技术栈(数组字段,= 模糊 == 精确)\n"+
|
||||
" port 端口号(整数,支持 > >= < <=)\n"+
|
||||
" status_code HTTP 状态码(整数,支持 > >= < <=)\n"+
|
||||
" company_id 归属企业 id(整数)\n"+
|
||||
" task_id 来源任务 id(整数)\n"+
|
||||
"示例:status_code>=400 AND technology=shiro\n"+
|
||||
" method==POST AND url=/api/admin\n"+
|
||||
" icp=京 OR icp=沪\n"+
|
||||
" (port==80 OR port==443) AND technology=nginx",
|
||||
obj(map[string]any{
|
||||
"dsl": str(`DSL 查询表达式,参见工具描述。未传 id/ids 时必须非空。例:status_code>=400 technology=shiro / method==POST url=/api/admin`),
|
||||
"type": str("资产类型过滤:root_domain|ip|subdomain|app|service|endpoint(独立字段,可与 dsl 叠加;单独 type 不足以查询,仍需 dsl)"),
|
||||
"id": intp("直接按单个资产 id 取(可选,与 dsl/type 互斥)"),
|
||||
"ids": map[string]any{"type": "array", "items": map[string]any{"type": "integer"}, "description": "直接按多个资产 id 取(可选,与 dsl/type 互斥)"},
|
||||
"limit": intp("返回上限,默认 10(可选)"),
|
||||
"offset": intp("分页偏移,默认 0(可选)"),
|
||||
}),
|
||||
func(_ context.Context, in json.RawMessage) (actool.Result, error) {
|
||||
if t.as == nil {
|
||||
return actool.Errorf("list_assets 未启用: AssetStore 未初始化"), nil
|
||||
}
|
||||
var a struct {
|
||||
DSL string `json:"dsl"`
|
||||
Type string `json:"type"`
|
||||
ID int64 `json:"id"`
|
||||
IDs []int64 `json:"ids"`
|
||||
Limit int `json:"limit"`
|
||||
Offset int `json:"offset"`
|
||||
}
|
||||
_ = json.Unmarshal(in, &a)
|
||||
if a.Limit <= 0 {
|
||||
a.Limit = 10
|
||||
}
|
||||
|
||||
var assets []*db.Asset
|
||||
var err error
|
||||
switch {
|
||||
case a.ID > 0:
|
||||
assets, err = t.as.GetByIDs([]int64{a.ID})
|
||||
case len(a.IDs) > 0:
|
||||
assets, err = t.as.GetByIDs(a.IDs)
|
||||
case a.DSL != "":
|
||||
assets, err = t.as.QueryDSL(a.DSL, a.Type, a.Limit, a.Offset)
|
||||
default:
|
||||
return actool.Errorf("未传 id/ids 时 dsl 不能为空:不允许无条件查询全部资产,请提供查询条件"), nil
|
||||
}
|
||||
if err != nil {
|
||||
return actool.Errorf("DSL 错误: " + err.Error()), nil
|
||||
}
|
||||
return jsonResult(map[string]any{
|
||||
"count": len(assets),
|
||||
"assets": assets,
|
||||
})
|
||||
},
|
||||
)
|
||||
}
|
||||
|
||||
// listCompanies lets an agent enumerate companies (企业) with their scope + asset count.
|
||||
func (t *ToolSet) listCompanies() actool.CoreTool {
|
||||
return readTool(
|
||||
"list_companies",
|
||||
"列出资产库中的【企业/公司】及其资产范围(scope)与已归属资产数。用于查看有哪些公司、"+
|
||||
"拿到 company_id(insert_assets 关联 app、list_assets 按 company_id 过滤时用)。"+
|
||||
"可选 search 按公司名模糊过滤(不区分大小写),留空返回全部。",
|
||||
obj(map[string]any{
|
||||
"search": str("按公司名模糊过滤(可选,不区分大小写);留空返回全部"),
|
||||
}),
|
||||
func(_ context.Context, in json.RawMessage) (actool.Result, error) {
|
||||
if t.cs == nil {
|
||||
return actool.Errorf("list_companies 未启用: CompanyStore 未初始化"), nil
|
||||
}
|
||||
var a struct {
|
||||
Search string `json:"search"`
|
||||
}
|
||||
_ = json.Unmarshal(in, &a)
|
||||
cos, err := t.cs.ListCompanies()
|
||||
if err != nil {
|
||||
return actool.Errorf("查询公司失败: " + err.Error()), nil
|
||||
}
|
||||
q := strings.ToLower(strings.TrimSpace(a.Search))
|
||||
type companyOut struct {
|
||||
ID int64 `json:"id"`
|
||||
Name string `json:"name"`
|
||||
AssetCount int `json:"asset_count"`
|
||||
Scope []string `json:"scope"`
|
||||
}
|
||||
out := make([]companyOut, 0, len(cos))
|
||||
for _, c := range cos {
|
||||
if q != "" && !strings.Contains(strings.ToLower(c.Name), q) {
|
||||
continue
|
||||
}
|
||||
scope := make([]string, 0, len(c.Scope))
|
||||
for _, r := range c.Scope {
|
||||
scope = append(scope, r.Raw)
|
||||
}
|
||||
out = append(out, companyOut{ID: c.ID, Name: c.Name, AssetCount: c.AssetCount, Scope: scope})
|
||||
}
|
||||
return jsonResult(map[string]any{"count": len(out), "companies": out})
|
||||
},
|
||||
)
|
||||
}
|
||||
|
||||
// splitLines splits a multi-line string into non-empty trimmed lines.
|
||||
func splitLines(s string) []string {
|
||||
var out []string
|
||||
for _, line := range strings.Split(s, "\n") {
|
||||
line = strings.TrimSpace(line)
|
||||
if line != "" {
|
||||
out = append(out, line)
|
||||
}
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// WorkerTools returns the tool set for a work agent.
|
||||
func (t *ToolSet) WorkerTools() []actool.CoreTool {
|
||||
return []actool.CoreTool{
|
||||
t.listFindings(), t.listFacts(), t.nodeDetail(),
|
||||
t.addFinding(), t.recordFact(),
|
||||
// cross-work retrieval
|
||||
t.searchAllWorkerTraces(), t.listWorkerTraces(), t.getWorkerTrace(),
|
||||
// asset management (handlers guard nil store internally)
|
||||
t.insertAssets(), t.addCompanyScope(), t.listAssets(), t.listCompanies(),
|
||||
}
|
||||
}
|
||||
|
||||
// MainAgentTools returns the human-interface tool set.
|
||||
func (t *ToolSet) MainAgentTools() []actool.CoreTool {
|
||||
return []actool.CoreTool{
|
||||
t.graphOverview(), t.listFindings(), t.listFacts(), t.nodeDetail(),
|
||||
t.getWorkerOutput(), t.getWorkerTrace(), t.searchAllWorkerTraces(), t.addHint(), t.addIntent(),
|
||||
// asset management (handlers guard nil store internally)
|
||||
t.insertAssets(), t.addCompanyScope(), t.listAssets(),
|
||||
}
|
||||
}
|
||||
|
||||
// AllDomainTools returns the union of all domain tools across all agent types,
|
||||
// deduped by name (mainagent order wins). Used by the server to build a registry
|
||||
// for injecting domain tools into agents (Auto, custom) that don't own a per-task
|
||||
// ToolSet. The caller provides real stores; tools are callable at taskID=0 scope.
|
||||
func (t *ToolSet) AllDomainTools() []actool.CoreTool {
|
||||
seen := map[string]bool{}
|
||||
var out []actool.CoreTool
|
||||
all := append(append(t.MainAgentTools(), t.PlannerTools()...), t.WorkerTools()...)
|
||||
for _, tool := range all {
|
||||
if !seen[tool.Name()] {
|
||||
seen[tool.Name()] = true
|
||||
out = append(out, tool)
|
||||
}
|
||||
}
|
||||
return out
|
||||
}
|
||||
+302
@@ -0,0 +1,302 @@
|
||||
package agent
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"path/filepath"
|
||||
"time"
|
||||
|
||||
"github.com/Autumn-27/artex/db"
|
||||
"github.com/Autumn-27/norma/agentcore"
|
||||
"github.com/Autumn-27/norma/harness"
|
||||
"github.com/Autumn-27/norma/llm"
|
||||
"github.com/Autumn-27/norma/memory"
|
||||
"github.com/Autumn-27/norma/permission"
|
||||
actool "github.com/Autumn-27/norma/tool"
|
||||
"github.com/Autumn-27/norma/transcript"
|
||||
)
|
||||
|
||||
// Worker is an LLM work agent (docs §4.4): it claims ONE intent, completes it
|
||||
// with real tools (Bash: kali tooling through the recording proxy), writes the
|
||||
// FACTS it found back into the graph, and stops. It does NOT generate new
|
||||
// directions (that is the planner's job) and does NOT keep exploring toward the
|
||||
// goal on its own. Multiple workers run concurrently as goroutines.
|
||||
// WebSearchOpts is the web-search backend selection the server pushes into each
|
||||
// agent (planner/worker/main). Enabled=false leaves the web_search tool off.
|
||||
// Backend is "ddgs" (no key), "brave-free" (BraveKey required), or "tavily"
|
||||
// (TavilyKey required). It maps directly onto agentcore.Options.
|
||||
// Proxy is a dedicated egress proxy for the search request (http/https/socks5),
|
||||
// independent of the traffic-recording MITM proxy — set it when the search endpoint
|
||||
// is only reachable via a VPN/SOCKS proxy. Empty = direct.
|
||||
type WebSearchOpts struct {
|
||||
Enabled bool
|
||||
Backend string
|
||||
BraveKey string
|
||||
TavilyKey string
|
||||
Proxy string
|
||||
}
|
||||
|
||||
type Worker struct {
|
||||
prov llm.Provider
|
||||
model string
|
||||
workDir string
|
||||
proxyAddr string
|
||||
proxyCACert string // recording proxy's CA cert path (for WebFetch HTTPS verify)
|
||||
webSearch WebSearchOpts // web_search tool backend selection (off by default)
|
||||
mem *memory.Store // cross-engagement tradecraft memory (G4)
|
||||
tx *transcript.Store // raw LLM conversation persistence (nil = off)
|
||||
window int // context window in tokens (for compaction)
|
||||
maxTurns int // max agent turns per run (0 = unlimited)
|
||||
// runTimeout is the wall-clock budget for the main exploration of one intent
|
||||
// (0 = unlimited). When it fires, the run is cut and a settlement round is
|
||||
// forced so already-identified facts get written back instead of being lost.
|
||||
runTimeout time.Duration
|
||||
// extraTools are host-provided tools (e.g. traffic query, oast) appended to
|
||||
// the worker's graph write-back tools.
|
||||
extraTools []actool.CoreTool
|
||||
}
|
||||
|
||||
// SetRunTimeout configures the per-intent wall-clock budget for the main
|
||||
// exploration (0 = unlimited). When it fires, the SDK settlement phase still runs
|
||||
// so facts are never lost to a timeout. Safe to call before Execute.
|
||||
func (w *Worker) SetRunTimeout(run time.Duration) {
|
||||
w.runTimeout = run
|
||||
}
|
||||
|
||||
// settleHardGrace is how far past the soft wall-clock budget (runTimeout) the hard
|
||||
// ctx deadline sits — a backstop for a hung turn; the soft budget handles the
|
||||
// normal case at the turn boundary.
|
||||
const settleHardGrace = 90 * time.Second
|
||||
|
||||
// settleWrapUpPrompt is injected by the SDK settlement phase when a worker hits its
|
||||
// turn/time budget: stop probing, write back what was found, then end with a
|
||||
// plain-text one-liner (which becomes this run's displayed result).
|
||||
const settleWrapUpPrompt = "你即将因预算耗尽被终止。不要再运行任何命令/探测。请依次:(1) 把你上面已识别但还没写回的内容逐条写回——新资产用 insert_assets、探索结论/事实用 record_fact、确认漏洞用 report_finding;(2) **最后单独用一句话纯文本**总结你做了什么、得到哪些关键结论(这句会作为本次运行的结果展示,务必输出)。"
|
||||
|
||||
// SetMemory enables cross-engagement tradecraft memory (RecallMemory/RecordMemory
|
||||
// tools + auto-injection of relevant memories, docs §8 G4).
|
||||
func (w *Worker) SetMemory(m *memory.Store) { w.mem = m }
|
||||
|
||||
func NewWorker(prov llm.Provider, model, workDir string, tx *transcript.Store, window, maxTurns int, extra ...actool.CoreTool) *Worker {
|
||||
return &Worker{prov: prov, model: model, workDir: workDir, tx: tx, window: window, maxTurns: maxTurns, extraTools: extra}
|
||||
}
|
||||
|
||||
// SetProxy configures the recording proxy address that workers route target
|
||||
// traffic through, plus the CA cert path WebFetch trusts to verify HTTPS through
|
||||
// that MITM proxy. Empty addr disables the hint.
|
||||
func (w *Worker) SetProxy(addr, caCert string) { w.proxyAddr, w.proxyCACert = addr, caCert }
|
||||
|
||||
// SetWebSearch selects the web_search backend for this worker (off by default).
|
||||
func (w *Worker) SetWebSearch(o WebSearchOpts) { w.webSearch = o }
|
||||
|
||||
// proxyEnv builds the Bash-subprocess env that routes child-command HTTP through
|
||||
// the recording proxy and makes the common toolchain trust its MITM CA — so tools
|
||||
// need no manual -x/--proxy/-k. Each ecosystem reads a different CA var (verified
|
||||
// empirically): SSL_CERT_FILE→curl/urllib/Go/openssl, REQUESTS_CA_BUNDLE→python
|
||||
// requests (it ignores SSL_CERT_FILE), CURL_CA_BUNDLE→curl, GIT_SSL_CAINFO→git,
|
||||
// NODE_EXTRA_CA_CERTS→node; NODE_USE_ENV_PROXY makes Node 24+ honor the proxy vars.
|
||||
// Empty proxyAddr → nil (direct, unchanged env).
|
||||
func proxyEnv(proxyAddr, caCert string) []string {
|
||||
if proxyAddr == "" {
|
||||
return nil
|
||||
}
|
||||
env := []string{
|
||||
"HTTP_PROXY=" + proxyAddr, "HTTPS_PROXY=" + proxyAddr,
|
||||
"http_proxy=" + proxyAddr, "https_proxy=" + proxyAddr,
|
||||
"NODE_USE_ENV_PROXY=1", // Node 24+: honor HTTP(S)_PROXY in built-in fetch/http
|
||||
}
|
||||
if caCert != "" {
|
||||
env = append(env,
|
||||
"SSL_CERT_FILE="+caCert,
|
||||
"CURL_CA_BUNDLE="+caCert,
|
||||
"REQUESTS_CA_BUNDLE="+caCert,
|
||||
"GIT_SSL_CAINFO="+caCert,
|
||||
"NODE_EXTRA_CA_CERTS="+caCert,
|
||||
)
|
||||
}
|
||||
return env
|
||||
}
|
||||
|
||||
// workerDefaultTmpl is the built-in EDITABLE body (段 [A]) of the worker system
|
||||
// prompt, seeded into agent_prompts. The trafficTool block and the 中间产物输出规约
|
||||
// are NOT here — they are code-owned and appended by workerSystem after rendering
|
||||
// (段 [B]/[C]), so editing the DB body can never drop them.
|
||||
const workerDefaultTmpl = `你是一个授权渗透测试系统的"执行者"(work agent)。你领到【一条意图】(一句话探索方向),唯一职责是:**完成这一条意图、把发现的事实写回知识图谱,然后停止返回。**
|
||||
|
||||
铁律(务必遵守):
|
||||
1. **只做这一条意图**。意图边界就是你的红线。指纹意图就只做指纹识别,不要顺手去枚举端点、爆破目录、扒 JS 找 API、测漏洞——那些是【别的意图】的事,由规划者去派别的 worker。
|
||||
2. **你不负责"探索方向"**。发现了值得继续追的新线索,不要自己接着打;只要把它写回图,规划者会读到、自己生成新意图。生成探索方向是规划者的职责,不是你的。
|
||||
3. **穷尽后再返回,别在第一个障碍前放弃**。判"意图达成"的标准是【你已把这条方向真正探透】:初次尝试被拦(一个 payload 被过滤、一个端点 404、一个注入点没回显)不等于此路不通——先换编码/换方法/换参数/换路径把这条意图的合理手段走完,再下结论。**但边界不变**:穷尽的只是【这一条意图内部】的手段,绝不是顺手去做别的意图(枚举别的端点、测别的漏洞);那些仍是规划者派别的 worker 的事。真正探透了、或确认此路不通了,就立即写回并返回,别因为"任务总目标还没达成"就继续,也别为凑步数在已探尽的方向空转。
|
||||
4. **边发现边写回,并且写对地方**。每得出一个结果就立刻写回(别攒到最后,否则步数耗尽全丢)。结果只算写进图里的,活在你脑子/文字里的不算。**两张图分清楚**:
|
||||
- **发现新资产/资源** → insert_assets 写【资产图】(登记资产本身:endpoint / parameter / tech 指纹 / service / 凭据 / 子域 等)。资产的结构化属性写在它自己身上:站点/接口的状态码/标题/body 长度/content_type 放 props.http,技术栈登记为 type=tech 节点。多个资产用 insert_assets 的 assets 数组一次批量登记。
|
||||
- **得出探索结论/事实**(含指纹/枚举等正向结论,和"端口关闭"、"该参数不可注入"、"未发现登录入口"等否定结论)→ record_fact 写【探索图】(传 intent_id)。**一次探索的多个观察汇总成【一条】事实**:summary 写总结性一句话,detail 写相关细节(技术栈、状态码、响应特征等都塞进这一条的 detail)——**不要一个属性一条事实**,一条意图通常只产出一条事实,拆太碎会让图谱无限膨胀。真有多条【彼此不同】的结论才用 facts 数组一次写。【新增】**只写增量**:写回前先扫一眼上方【全局探索态势】里的 recent_facts——只写你这次【新得到】的结论,别把图里已有的事实换个措辞再记一遍(重复事实会让图谱膨胀、误导规划者以为有新进展)。若你的观察只是印证了已有 fact 而无新增,就不必再记一条。
|
||||
- **确认漏洞** → report_finding 写【探索图】(含 PoC,传 intent_id)。
|
||||
|
||||
可用工具:
|
||||
- insert_assets:登记新资产(资产图)。**你只传原始信息,key 与父子关联由代码算**:新接口→传完整 url+method(代码自动建 domain→site→endpoint、自动抽 URL 里的 query 参数,body/header 参数放 params);指纹→type=tech,name=技术名,on_url=站点地址,props填{version,category}。多个资产放 assets 数组一次批量登记。属性写在资产自己的 props 上,探索结论不要写这里。
|
||||
- record_fact:把探索【事实/结论】写入探索图并连到意图(传 intent_id)。正向/否定结论、观察、判断用它;**一次探索的多个观察汇总成一条事实**(summary 总结一句话 + detail 放细节),不要一个属性一条。真有多条不同结论才用 facts 数组。**只写你真实看到的**:给 evidence(一行关键证据:命令+关键输出,简洁,别粘大段——细节在 detail)、标 confidence(observed 直接看到 / inferred 推断);**否定结论**(不可注入/端口关闭等)尤其要给证据、证据弱就标 inferred,别让错的否定误导规划者放弃方向。
|
||||
- report_finding:确认漏洞 → 记录(含 PoC,传 intent_id=你领到的意图id)。**只有你在本次运行里真实触发过该漏洞、拿到了可复现的证据(请求/响应或命令输出)才用它。** 严禁把下列当作已确认漏洞上报:仅凭版本号/指纹匹配到某 CVE、仅凭"参数看起来可注入"、仅凭外部漏洞库/更新日志/代码 diff 推断。**不要用查 CVE 库或"对比补丁版本"替代实际触发。** 触发不了但确有嫌疑,就用 record_fact 记一条 confidence=inferred 的事实(描述嫌疑点+为何未能触发),交给规划者派后续意图,别硬记成 finding。
|
||||
- list_assets(查询资产,非探索节点) / asset_neighbors / list_facts(探索事实) / list_findings(漏洞) / node_detail(探索节点 id,非资产 id):按需查上下文。
|
||||
- 【必要时才会使用,大多数上下文都在本次会话中】search_all_worker_traces / list_worker_traces / get_worker_trace:**跨 work 复用信息**(别的 work 执行过程里出现过、却没写进 fact 的东西)。你看不到探索图,但可以:search_all_worker_traces(q) 按关键字搜全部 work 的过程(返回带 intent_id);list_worker_traces 看有哪些 work 跑过;get_worker_trace(intent_id) 看某 work 的步骤摘要、get_worker_trace(intent_id, step_ids=[…]) 取那几步完整内容(一次≤5个)。**仅用于复用他人观察、避免重复劳动——不改变你的任务边界(仍只做你这条意图)。**
|
||||
|
||||
只在授权范围内操作。完成本意图后用一句话总结你做了什么、写回了哪些事实。务实、克制、聚焦这一条意图。`
|
||||
|
||||
// workerTrafficBlock is 段 [B]: the traffic-tool note, code-injected only when
|
||||
// traffic capture is on (proxyAddr set). Not stored, not editable.
|
||||
func workerTrafficBlock(proxyAddr string) string {
|
||||
if proxyAddr == "" {
|
||||
return ""
|
||||
}
|
||||
return "\n\n**流量工具**:\n- traffic_search / traffic_get:目标流量已被记录代理全量落库(本地文件树 data/traffic/<host>/<访问路径>/)。回看响应、找已访问过的资源,**先查流量、不要重复 curl 同一 URL**。traffic_search **必须指定 host**、默认只回 3 条极轻量索引(id/method/url/status/resp_len,无响应内容),需要更多显式调大 limit;要看某条原文用 traffic_get(id)。"
|
||||
}
|
||||
|
||||
// artifactSpec is 段 [C]: the code-owned, non-editable tail appended to every
|
||||
// pentest agent's prompt — intermediate artifacts must land in the shared work
|
||||
// dir, never /tmp. Guaranteed present regardless of how the DB body is edited.
|
||||
func artifactSpec(workDir string) string {
|
||||
return "\n\n**中间产物输出规约**:脚本、payload、抓到的响应体、临时数据等一切中间产物,**一律写到共享工作目录 " + workDir + "**(这是所有 agent 共用的 CWD,相对路径即写在这里,也可用该绝对路径)——**不要写 /tmp、不要用其它绝对路径**。"
|
||||
}
|
||||
|
||||
// workerArtifactSubdir is the worker-only addendum to artifactSpec: put a run's
|
||||
// artifacts under an i<intentID>/ subdir to avoid concurrent name collisions.
|
||||
const workerArtifactSubdir = "为避免与其他 work 撞名,把本次产物放到子目录 i<意图id>/ 下(如 i123/exploit.py)。"
|
||||
|
||||
func workerSystem(proxyAddr, workDir string) string {
|
||||
body := renderSystem("worker", workerDefaultTmpl, WorkerVars{ProxyAddr: proxyAddr})
|
||||
return body + workerTrafficBlock(proxyAddr) + artifactSpec(workDir) + workerArtifactSubdir
|
||||
}
|
||||
|
||||
// renderIntentTask formats the claimed intent for the worker's SYSTEM prompt: the
|
||||
// intent is the worker's whole job, so — like the planner's situational block — it
|
||||
// belongs in system, where compaction can never drop it during a long run.
|
||||
func renderIntentTask(intent *db.Node) string {
|
||||
return fmt.Sprintf("\n\n【你领到的意图(本次唯一任务:只做这一条、只产生事实、做完即停)】:\n%s\n意图 id: %d(写回 record_fact / report_finding 时传它)", string(intent.Payload), intent.ID)
|
||||
}
|
||||
|
||||
// renderWorkerGraphOverview folds the global situational snapshot into the worker's
|
||||
// SYSTEM prompt for AWARENESS ONLY. The framing is deliberately strong: the overview
|
||||
// must NOT widen the worker's job — it still does only its assigned intent. Its sole
|
||||
// purpose is letting the worker read context (existing facts/assets/hints)
|
||||
// so it avoids redundant work and doesn't re-derive what others already found.
|
||||
func renderWorkerGraphOverview(data map[string]any) string {
|
||||
b, err := json.Marshal(data)
|
||||
if err != nil {
|
||||
return "" // fall back silently: the worker just won't have the global context
|
||||
}
|
||||
return "\n\n【全局探索态势(仅供你了解大局,不是你的任务清单)】:\n" +
|
||||
"下面是整个任务当前的探索概况。给你的**唯一目的**是让你了解全局动态、复用已有结论、避免重复别人已做过的事。\n" +
|
||||
"**它绝不扩大你的职责边界**:你仍然只做上面领到的那一条意图。看到这里有别的 open 意图 / 未覆盖的点 / 其它可打方向,也**绝不要自己去动手**——那些是别的 worker 的事,由规划者调度。你若发现相关新线索,最多写进 fact 让规划者知道,不要自己追。\n" +
|
||||
string(b)
|
||||
}
|
||||
|
||||
// Execute runs one intent. hooks (the per-task Guard) gates every tool call; may
|
||||
// be nil. emit, if non-nil, receives one ActivityRecord per execution step.
|
||||
// Returns the terminal reason (so the engine can distinguish completed vs
|
||||
// max_turns) and how many facts were written back (so an intent that explored but
|
||||
// persisted nothing isn't mistaken for done).
|
||||
func (w *Worker) Execute(ctx context.Context, name string, taskID int64, as *db.AssetStore, ts *db.ExplorationStore, intent *db.Node, hooks harness.HookRunner, emit func(db.Activity), enr EnrichTrigger) (harness.TerminalReason, int, error) {
|
||||
tsx := NewToolSet(ts, name)
|
||||
tsx.SetTaskID(taskID)
|
||||
if as != nil {
|
||||
tsx.SetAssetStore(as, as.Companies())
|
||||
}
|
||||
tsx.SetOwnerNode(intent.ID) // assets this worker discovers anchor to its intent → visible to the task
|
||||
tsx.SetEnrich(enr) // async DNS/HTTP auto-completion for assets this worker writes
|
||||
// base = built-in worker tools ∪ host tools (traffic) ∪ default tools (incl. Bash);
|
||||
// then augment with the agent's visible skills/MCP. During the SDK settlement
|
||||
// phase, Bash is hidden via Settlement.DisabledTools (no local gating needed).
|
||||
base := append(tsx.WorkerTools(), w.extraTools...)
|
||||
base = append(base, actool.DefaultTools()...)
|
||||
tools, def, cleanup := AugmentTools(ctx, "worker", base)
|
||||
defer cleanup()
|
||||
|
||||
// 意图是 worker 的全部任务,放进 system prompt(抗 compaction、整轮常驻),而不是
|
||||
// user 输入——与 planner 把关键态势放 system 一致。落在指令之后、deferred 块之前。
|
||||
// 全局态势(graphOverviewData)也带上,但仅供了解大局;意图放在最后、最醒目的位置。
|
||||
overview := renderWorkerGraphOverview(tsx.graphOverviewData())
|
||||
system, boundary := deferredSystem(workerSystem(w.proxyAddr, w.workDir)+overview+renderIntentTask(intent), def)
|
||||
// 任务级 deadline(经 ctx 注入)夹逼本 run 的墙钟预算 + 决定收尾词(见 taskclock.go)。
|
||||
tc := taskClockFrom(ctx)
|
||||
maxDur, clamped := clampMaxDuration(tc.DeadlineUnix, w.runTimeout)
|
||||
settle := wrapupSettlement("worker", []string{"Bash"})
|
||||
if tc.DeadlineUnix > 0 {
|
||||
settle = wrapupSettlementForTask("worker", []string{"Bash"}, clamped)
|
||||
}
|
||||
opts := agentcore.Options{
|
||||
Provider: w.prov,
|
||||
SystemPrompt: system,
|
||||
DynamicBoundary: boundary,
|
||||
Tools: tools,
|
||||
DeferredTools: def.Deferred,
|
||||
UnlockSet: def.Unlock,
|
||||
PermissionMode: permission.ModeBypass,
|
||||
// WebFetch 走记录代理,其 HTTP 与 curl 一样被留痕;载入代理 CA 让经 MITM
|
||||
// 重签的 HTTPS 证书能【正常校验通过】(而非关掉校验)。proxy 空则直连。
|
||||
EnableWebFetch: true,
|
||||
WebFetchProxy: w.proxyAddr,
|
||||
WebFetchCACert: w.proxyCACert,
|
||||
// 联网搜索(可选)。ddgs 无需 key;brave-free 需 BraveKey;tavily 需 TavilyKey。
|
||||
// WebSearchProxy 是独立的出口代理(http/https/socks5),与记录流量的 MITM 代理无关;空则直连。
|
||||
EnableWebSearch: w.webSearch.Enabled,
|
||||
WebSearchBackend: w.webSearch.Backend,
|
||||
BraveSearchAPIKey: w.webSearch.BraveKey,
|
||||
TavilySearchAPIKey: w.webSearch.TavilyKey,
|
||||
WebSearchProxy: w.webSearch.Proxy,
|
||||
// Bash 子命令的 HTTP 默认走记录代理 + 信任其 CA(工具无需 -x/-k)。
|
||||
BashEnv: proxyEnv(w.proxyAddr, w.proxyCACert),
|
||||
WorkingDir: w.workDir,
|
||||
MaxTurns: w.maxTurns, // 0 = unlimited (configurable in agent management)
|
||||
// 墙钟预算,轮边界判,不打断半路;0 = 不限。有任务级 deadline 时夹逼到 min(自身预算,
|
||||
// 距 deadline 剩余),让本 run 在任务到点时自然进收尾(见 taskclock.go)。
|
||||
MaxDuration: maxDur,
|
||||
// 命中预算(轮次 OR 时长)→ SDK 跑一轮收尾(隐藏 Bash),把已识别的写回,避免烂尾。
|
||||
// clamped(被任务 deadline 夹逼)时用 PromptByReason:因超时=任务到点→任务超时词,
|
||||
// 因步数=夹逼窗口内步数先耗尽→回落 per-run 词。非 clamped 维持纯 per-run。
|
||||
Settlement: settle,
|
||||
// large tool output spills to cmd-output/ with a head + pointer (SDK tool.Capture);
|
||||
// full output preserved on disk. 截断上限用 SDK 默认(30000 字符)。
|
||||
ToolOutputDir: filepath.Join(w.workDir, "cmd-output"),
|
||||
Compaction: compactionConfig(w.window), // long tool-heavy runs stay within the window
|
||||
Todos: actool.NewTodoStore(), // 会话级临时待办(TodoWrite),纯规划用,退出即丢
|
||||
}
|
||||
if hooks != nil { // typed-nil guard: only set when concrete (avoids harness panic)
|
||||
opts.Hooks = hooks
|
||||
}
|
||||
if w.mem != nil {
|
||||
opts.Memory = &agentcore.MemoryOptions{Store: w.mem, AutoInject: true, MaxInject: 3}
|
||||
}
|
||||
if w.tx != nil { // persist raw LLM conversation; one file per worked intent
|
||||
opts.Transcript = w.tx
|
||||
opts.SessionID = fmt.Sprintf("exp%d-worker-i%d", ts.ID(), intent.ID)
|
||||
}
|
||||
intentID := intent.ID
|
||||
emitWrap := func(r db.Activity) {
|
||||
if emit != nil {
|
||||
r.NodeID, r.Worker = &intentID, name
|
||||
emit(r)
|
||||
}
|
||||
}
|
||||
// 意图已在 system 里。user 输入只放启动指令。
|
||||
input := "开始执行 system 提示里的这条意图:只做它、只产生事实、做完即停。"
|
||||
|
||||
s := agentcore.NewSession(opts)
|
||||
defer s.Close() // release the session's background-task manager (temp dir + processes)
|
||||
|
||||
// Budgets + settlement are owned by the SDK (MaxTurns/MaxDuration + Settlement):
|
||||
// on hit it runs a wrap-up turn and finishes with ReasonMaxTurns/ReasonTimeout.
|
||||
// We only add a HARD ctx backstop past the soft wall-clock budget, for the rare
|
||||
// case a single turn hangs so the turn-boundary check never runs — the soft
|
||||
// budget handles the normal case gracefully (no mid-stream abort). pause /
|
||||
// planner kill cancel ctx; the engine distinguishes those and re-queues/stops.
|
||||
// backstop uses the EFFECTIVE (deadline-clamped) budget, not the raw runTimeout,
|
||||
// so a task-timeout run's hard deadline sits just past its clamped soft budget.
|
||||
runCtx := ctx
|
||||
if maxDur > 0 {
|
||||
var runCancel context.CancelFunc
|
||||
runCtx, runCancel = context.WithTimeout(ctx, maxDur+settleHardGrace)
|
||||
defer runCancel()
|
||||
}
|
||||
_, reason, err := captureRunSession(runCtx, s, input, emitWrap)
|
||||
return reason, tsx.Wrote(), err
|
||||
}
|
||||
+173
@@ -0,0 +1,173 @@
|
||||
package agent
|
||||
|
||||
import (
|
||||
"strings"
|
||||
|
||||
"github.com/Autumn-27/norma/harness"
|
||||
)
|
||||
|
||||
// 收尾提示词(wrap-up / settlement prompt):当 agent 因【步数耗尽(MaxTurns)】或
|
||||
// 【超时(run_seconds/MaxDuration)】被终止时,SDK 的 settlement 阶段会注入这段提示,
|
||||
// 让 agent 先把已识别但未写回的内容落库、再输出一句总结,避免烂尾。
|
||||
//
|
||||
// 每个 agent 的收尾提示词可在后台按需覆盖(存 agents.wrapup_prompt),留空则用这里的
|
||||
// 内置默认。仅【提示词正文】可编辑;禁用哪些工具、收尾自身给几轮预算属代码固定策略。
|
||||
|
||||
// WrapupOverride, if set, returns the stored wrap-up prompt for an agent key and
|
||||
// whether a non-empty one exists. Wired by the server to the agents table (like
|
||||
// PromptOverride for system prompts). nil / empty → the built-in default is used.
|
||||
var WrapupOverride func(agentKey string) (string, bool)
|
||||
|
||||
// WrapupMaxTurnsOverride, if set, returns the admin-configured turn budget for the
|
||||
// wrap-up phase of an agent and whether a positive one exists. Wired to the agents
|
||||
// table. nil / ≤0 → the built-in per-agent default (wrapupTurnDefaults) is used.
|
||||
var WrapupMaxTurnsOverride func(agentKey string) (int, bool)
|
||||
|
||||
// 内置默认收尾提示词,按 agent key 索引。worker 复用历史上硬编码的 settleWrapUpPrompt
|
||||
// (定义在 worker.go),planner/mainagent 各有一版;未命中的(自定义 agent)走通用兜底。
|
||||
var wrapupDefaults = map[string]string{
|
||||
"worker": settleWrapUpPrompt,
|
||||
"planner": plannerWrapUpDefault,
|
||||
"mainagent": mainAgentWrapUpDefault,
|
||||
}
|
||||
|
||||
// wrapupTurnDefaults: 各 agent 收尾阶段【自身】的轮数预算内置默认(可被后台 >0 覆盖)。
|
||||
// worker 写回可能占几轮(record_fact/insert_assets)故给 4;其余给 5。未命中走 genericWrapupTurns。
|
||||
var wrapupTurnDefaults = map[string]int{
|
||||
"worker": 4,
|
||||
"planner": 5,
|
||||
"mainagent": 5,
|
||||
}
|
||||
|
||||
const genericWrapupTurns = 5
|
||||
|
||||
const plannerWrapUpDefault = "你本轮规划的步数即将用尽——注意只是【这一轮】结束,系统之后仍会随态势变化再次唤醒你继续规划,并非任务终止,你无需在此收束整个规划。请把本轮已经想清楚的结论落地、别让这一轮白跑,但也【不要为了收尾硬凑意图】(本轮 0 个意图仍是完全正常的结果):(1) 若已判断出【当前就该派发】的探索方向,用一次 add_intent 批量提交(想好的别憋着不发);(2) 对已被某发现/事实证明达成的目标,调 prove_goal 标记 met(别漏判);(3) 若识别出需要分步的串行利用链,用 TodoWrite 记下,便于下次唤醒接着派。做完直接结束本轮,无需输出总结文本。"
|
||||
|
||||
const mainAgentWrapUpDefault = "你的步数即将用尽,本次交互就要结束。不要再发起新的探索/操作。请**单独用一句话纯文本**向用户总结当前进展、关键结论,以及建议的下一步。"
|
||||
|
||||
const genericWrapUpDefault = "你即将因预算耗尽被终止。请先把已完成但未落库的结果写回,再**单独用一句话纯文本**总结你做了什么、得到哪些关键结论(这句会作为本次运行的结果展示)。"
|
||||
|
||||
// WrapupDefault returns the built-in default wrap-up prompt for an agent key —
|
||||
// used by the admin UI as the "restore default" value and empty-field placeholder.
|
||||
func WrapupDefault(agentKey string) string {
|
||||
if d, ok := wrapupDefaults[agentKey]; ok {
|
||||
return d
|
||||
}
|
||||
return genericWrapUpDefault
|
||||
}
|
||||
|
||||
// WrapupTurnsDefault returns the built-in wrap-up turn budget for an agent key —
|
||||
// used by the admin UI as the "0 = default N" hint.
|
||||
func WrapupTurnsDefault(agentKey string) int {
|
||||
if n, ok := wrapupTurnDefaults[agentKey]; ok {
|
||||
return n
|
||||
}
|
||||
return genericWrapupTurns
|
||||
}
|
||||
|
||||
// resolveWrapup returns the effective wrap-up prompt: the DB override (if set and
|
||||
// non-empty) over the built-in default.
|
||||
func resolveWrapup(agentKey string) string {
|
||||
if WrapupOverride != nil {
|
||||
if t, ok := WrapupOverride(agentKey); ok && strings.TrimSpace(t) != "" {
|
||||
return t
|
||||
}
|
||||
}
|
||||
return WrapupDefault(agentKey)
|
||||
}
|
||||
|
||||
// resolveWrapupTurns returns the effective wrap-up turn budget: a positive DB
|
||||
// override over the built-in per-agent default.
|
||||
func resolveWrapupTurns(agentKey string) int {
|
||||
if WrapupMaxTurnsOverride != nil {
|
||||
if v, ok := WrapupMaxTurnsOverride(agentKey); ok && v > 0 {
|
||||
return v
|
||||
}
|
||||
}
|
||||
return WrapupTurnsDefault(agentKey)
|
||||
}
|
||||
|
||||
// wrapupSettlement builds the settlement config for an agent's run. Prompt and the
|
||||
// turn budget are admin-editable per agent; disabled tools are code-owned policy so
|
||||
// a user can't edit away the "stop probing" guardrail. Resolved fresh each run
|
||||
// (reads DB live), so edits apply on the next run without a restart.
|
||||
func wrapupSettlement(agentKey string, disabledTools []string) *harness.Settlement {
|
||||
return &harness.Settlement{
|
||||
Prompt: resolveWrapup(agentKey),
|
||||
DisabledTools: disabledTools,
|
||||
MaxTurns: resolveWrapupTurns(agentKey),
|
||||
}
|
||||
}
|
||||
|
||||
// ---------- 任务级超时收尾词(见 docs/任务级超时与收尾设计.md)----------
|
||||
//
|
||||
// 与 per-run 收尾词是【两套】:per-run 是"你这一次 run 的预算用完了";任务超时是
|
||||
// "整个任务到点、即将结束"。语义常相反(尤其 planner:per-run 说"别停继续规划",
|
||||
// 任务超时说"到点停止规划、做最后判定")。只给 worker/planner 配置。
|
||||
|
||||
// WrapupTaskTimeoutOverride / …TurnsOverride:任务超时收尾词与轮数的 DB 覆盖
|
||||
// (wire 到 agents.task_timeout_wrapup_prompt / _max_turns,仅 worker/planner)。
|
||||
var (
|
||||
WrapupTaskTimeoutOverride func(agentKey string) (string, bool)
|
||||
WrapupTaskTimeoutTurnsOverride func(agentKey string) (int, bool)
|
||||
)
|
||||
|
||||
var taskTimeoutWrapupDefaults = map[string]string{
|
||||
"worker": workerTaskTimeoutDefault,
|
||||
"planner": plannerTaskTimeoutDefault,
|
||||
}
|
||||
|
||||
const workerTaskTimeoutDefault = "**整个任务已到达超时上限,即将结束**(不是你这次 run 的预算,是整场探索到点了)。这是最后机会:(1) 把你已识别但还没写回的内容【全部】落库——新资产 insert_assets、探索结论/事实 record_fact、确认漏洞 report_finding;(2) 不要再启动任何新命令/探测;(3) **最后单独用一句话纯文本**总结你在本意图上的关键结论。"
|
||||
|
||||
const plannerTaskTimeoutDefault = "**整个任务已到达超时上限,即将结束**(不是本轮,是整个任务终止)。请基于当前【全部】事实与发现,做最后一次目标判定:对已被证据证明达成的目标调 prove_goal 标记 met(别漏判)。**不要再生成任何新意图**(此时派意图也不会再被执行)。判定完即收束,无需输出总结文本。"
|
||||
|
||||
// TaskTimeoutWrapupDefault 返回某 agent 的任务超时内置默认收尾词(供后台占位/恢复默认)。
|
||||
func TaskTimeoutWrapupDefault(agentKey string) string {
|
||||
return taskTimeoutWrapupDefaults[agentKey] // 未配置(mainagent/chat)返回空串
|
||||
}
|
||||
|
||||
// resolveTaskTimeoutWrapup:DB 覆盖(非空) > 内置默认。空串表示该 agent 无任务超时词
|
||||
// (非 worker/planner),此时调用方应回退 per-run 词。
|
||||
func resolveTaskTimeoutWrapup(agentKey string) string {
|
||||
if WrapupTaskTimeoutOverride != nil {
|
||||
if t, ok := WrapupTaskTimeoutOverride(agentKey); ok && strings.TrimSpace(t) != "" {
|
||||
return t
|
||||
}
|
||||
}
|
||||
return TaskTimeoutWrapupDefault(agentKey)
|
||||
}
|
||||
|
||||
func resolveTaskTimeoutTurns(agentKey string) int {
|
||||
if WrapupTaskTimeoutTurnsOverride != nil {
|
||||
if v, ok := WrapupTaskTimeoutTurnsOverride(agentKey); ok && v > 0 {
|
||||
return v
|
||||
}
|
||||
}
|
||||
return resolveWrapupTurns(agentKey) // 默认沿用 per-run 轮数
|
||||
}
|
||||
|
||||
// wrapupSettlementForTask builds settlement for a worker/planner run that is aware
|
||||
// of the task deadline. See §5 of the design doc:
|
||||
// - clamped=true → 本次 run 被任务 deadline 夹逼:因 Timeout 收尾=任务到点→任务超时词;
|
||||
// 因 MaxTurns 收尾=夹逼窗口内步数先耗尽、任务还剩几分钟→回落 per-run 词。
|
||||
// - clamped=false → 任务还早:两种 reason 都用 per-run 词(即退化为 wrapupSettlement)。
|
||||
//
|
||||
// 交给 harness 的 PromptByReason 在收尾时按【实际】reason 现场挑,无 build 时错配。
|
||||
func wrapupSettlementForTask(agentKey string, disabledTools []string, clamped bool) *harness.Settlement {
|
||||
perRun := resolveWrapup(agentKey)
|
||||
st := &harness.Settlement{
|
||||
Prompt: perRun, // 兜底(也是非 clamped 时两种 reason 的取值)
|
||||
DisabledTools: disabledTools,
|
||||
MaxTurns: resolveWrapupTurns(agentKey),
|
||||
}
|
||||
if clamped {
|
||||
if tt := resolveTaskTimeoutWrapup(agentKey); tt != "" {
|
||||
st.PromptByReason = map[harness.TerminalReason]string{
|
||||
harness.ReasonTimeout: tt, // 任务到点
|
||||
harness.ReasonMaxTurns: perRun, // 步数先耗尽、任务还剩时间
|
||||
}
|
||||
st.MaxTurns = resolveTaskTimeoutTurns(agentKey)
|
||||
}
|
||||
}
|
||||
return st
|
||||
}
|
||||
Reference in new issue
Block a user