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,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
|
||||
}
|
||||
Reference in new issue
Block a user