- set_goals 从目标拆解器内联提交工具,改为入 tools 目录的领域工具:直接写 goal 节点、 连任务根(RelSpawns)、notify 唤醒 planner;支持单条/批量;origin 按来源区分 (拆解器=goals、主 agent 运行时=human)。新 key 走 SeedTool 首插入,旧库无需迁移。 - 默认绑定 goals + mainagent:目录 seed 新增 goals 项并入 order,MainAgentTools 纳入; DecomposeGoals 统一改用同一受管工具直接落库,createGoals 只保留无产出兜底。 - 主 agent 通过 set_goals 新增目标时,把终态(done/failed/timeout)/暂停的任务拉回 running 继续跑:新增 ToolSet.resumeTask 回调(仅 mainagent 接线),经 resumeTaskForGoal 置 running + engine.Run(幂等/重启后重起循环)+ 解除暂停。拆解器/worker 不触发。 - mainagent 提示词补充 set_goals 用法。
140 lines
6.9 KiB
Go
140 lines
6.9 KiB
Go
package agent
|
||
|
||
import (
|
||
"context"
|
||
"encoding/json"
|
||
"strings"
|
||
|
||
"github.com/Autumn-27/artex/db"
|
||
"github.com/Autumn-27/norma/agentcore"
|
||
acperm "github.com/Autumn-27/norma/permission"
|
||
actool "github.com/Autumn-27/norma/tool"
|
||
)
|
||
|
||
// 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 提交结果。`
|
||
|
||
// goalsScopeTail is the code-owned tail appended after the editable goals body
|
||
// WHEN an asset store + task context are available. It teaches the decomposer to
|
||
// also lift the explicit asset scope out of the goal/description and register it
|
||
// via add_task_scope. Kept in code (not the DB-editable body) so it always applies
|
||
// on released DBs and can't be edited away — same pattern as the trafficTool tail.
|
||
const goalsScopeTail = `
|
||
|
||
**额外职责:登记测试资产范围**
|
||
除拆分目标外,你还要从「任务目标 / 任务描述」里识别出**明确给出的测试资产范围**,调用 add_task_scope 登记(资产测试覆盖度的分母,也是授权边界)。**最小范围原则:只登记用户明确点到的那一个目标,绝不擅自放大。**
|
||
- 目标是 URL 或带主机名的地址(如 https://xxx.example.com/path、app.example.com)→ 取其**完整主机名**,kind=subdomain,value=完整主机名。
|
||
例:目标 https://a1b2c3.lab.example.net/path → kind=subdomain,value=a1b2c3.lab.example.net(**不是** example.net)。
|
||
**严禁**把带子域的主机名缩成根域名——看到 xxx.example.com 就登记整个 example.com 会把范围扩到用户目标之外,违背最小范围原则。
|
||
- 仅当用户给的就是**裸根域名、且不含任何子域**(如直接写 example.com),或明确说“整个站点 / 所有子域 / 全域名” → 才用 kind=root_domain,value=example.com。
|
||
- 纯 IP 或网段 → kind=ip / cidr,value=IP 或 CIDR。
|
||
- **不要**登记公司范围(company)——任务刚建立、资产系统里通常还没有这家公司,登记不上,公司级范围交由后续 plan 阶段处理。
|
||
其它规则:
|
||
- 只登记**目标/描述里明确写出**的范围;严禁臆造或推断未提及的域名/IP。
|
||
- reason 简述依据来自哪句话,便于审计。
|
||
- 若目标/描述中没有任何明确资产范围,则**不要**调用 add_task_scope。
|
||
先用 add_task_scope 登记范围(如有),再调用 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.
|
||
//
|
||
// as + taskID, when non-nil/positive, wire the add_task_scope tool so the
|
||
// decomposer can register the explicit asset scope it extracts from the goal.
|
||
//
|
||
// ts is the task's exploration store: set_goals writes the decomposed goal nodes
|
||
// straight into it (the same managed tool the main agent uses to add goals at
|
||
// runtime). The returned specs are read back from the store so callers can emit
|
||
// per-goal activity and detect the "LLM produced nothing" case for their fallback.
|
||
func DecomposeGoals(ctx context.Context, c Config, dataDir, goalText, desc string, as *db.AssetStore, ts *db.ExplorationStore, taskID int64, emit func(db.Activity)) []GoalSpec {
|
||
prov, err := c.NewProvider()
|
||
if err != nil {
|
||
return nil
|
||
}
|
||
// worker="goals" tags the goal nodes' provenance; ts/taskID let set_goals link
|
||
// each goal under the task root. This is the catalog's real set_goals tool, so a
|
||
// web-edited description/schema on it applies here too.
|
||
tsx := &ToolSet{as: as, ts: ts, taskID: taskID, worker: "goals"}
|
||
// 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{DataDir: dataDir, Now: nowStr()})
|
||
tools := []actool.CoreTool{tsx.setGoals()}
|
||
// Wire add_task_scope only when we have a real asset store + task to write to.
|
||
// The scope-extraction tail is appended in lockstep so the prompt never asks for
|
||
// a tool that isn't present.
|
||
if as != nil && taskID > 0 {
|
||
tools = append(tools, tsx.addTaskScope())
|
||
sys += goalsScopeTail
|
||
}
|
||
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: tools,
|
||
PermissionMode: acperm.ModeBypass,
|
||
DisableBackgroundTasks: true,
|
||
MaxTurns: 6,
|
||
}, userMsg, captureEmit)
|
||
// set_goals persisted the goals directly; read them back so the caller sees what
|
||
// was written (empty slice ⇒ the LLM produced nothing ⇒ caller falls back).
|
||
if ts == nil {
|
||
return nil
|
||
}
|
||
nodes, _ := ts.ListByKind(db.KindGoal, 10000)
|
||
var out []GoalSpec
|
||
for _, n := range nodes {
|
||
var p struct {
|
||
Text string `json:"text"`
|
||
VulnClass string `json:"vulnclass"`
|
||
}
|
||
_ = json.Unmarshal(n.Payload, &p)
|
||
if strings.TrimSpace(p.Text) != "" {
|
||
out = append(out, GoalSpec{Text: p.Text, VulnClass: p.VulnClass})
|
||
}
|
||
}
|
||
return out
|
||
}
|