192 lines
6.4 KiB
Go
192 lines
6.4 KiB
Go
// Package guard implements the safety boundary layer (docs §11): an audit log,
|
|
// user-configured intercept-rule evaluation, 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.) Destructive/exfil gating is no longer hard-coded here — it lives in the
|
|
// DB intercept rules (seeded as ordinary [内置] rules, so users can disable or
|
|
// delete them), evaluated via applyIntercept.
|
|
package guard
|
|
|
|
import (
|
|
"context"
|
|
"encoding/json"
|
|
"regexp"
|
|
"sync"
|
|
"time"
|
|
|
|
"github.com/Autumn-27/artex/intercept"
|
|
"github.com/Autumn-27/norma/hook"
|
|
)
|
|
|
|
// 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 }
|
|
|
|
func (g *Guard) preToolUse(ctx context.Context, ev hook.Event) hook.Result {
|
|
// Extract the shell-command surface for the audit log: Bash + the interactive-shell
|
|
// tools (shell_open's command, shell_send's text). Destructive/exfil gating is no
|
|
// longer hard-coded here — it now lives in the DB intercept rules, evaluated by
|
|
// applyIntercept below. Other tools record an empty command.
|
|
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
|
|
}
|
|
g.record(ev.ToolName, "allow", "", cmd)
|
|
return g.applyIntercept(ctx, ev, cmd)
|
|
}
|
|
|
|
// applyIntercept evaluates user-configured intercept rules against the tool call.
|
|
// It is called after all built-in safety checks pass. cmd is the extracted shell
|
|
// surface (empty for non-shell tools), passed to the LLM fallback judge.
|
|
func (g *Guard) applyIntercept(ctx context.Context, ev hook.Event, cmd string) hook.Result {
|
|
if g.interceptor == nil {
|
|
return hook.Result{}
|
|
}
|
|
if !g.interceptor.IsToolEnabled(ev.ToolName) {
|
|
return hook.Result{}
|
|
}
|
|
ctx = intercept.WithCall(ctx, ev.ToolName, ev.Input)
|
|
dec, matched := g.interceptor.Match(ev.ToolName, ev.Input)
|
|
if !matched {
|
|
// No rule matched. Ask the LLM fallback judge (if enabled); when it is off
|
|
// or unwired, keep current behavior and allow.
|
|
d, judged := g.interceptor.Judge(ctx, ev.ToolName, judgeSubject(cmd, ev.Input))
|
|
if !judged {
|
|
return hook.Result{}
|
|
}
|
|
dec = d
|
|
}
|
|
switch dec.Action {
|
|
case "deny":
|
|
// 观测:deny 命中不阻塞审批,直接记一条 denied（历史/任务拦截页可见）。
|
|
g.interceptor.Log(ctx, intercept.ConvIDFromContext(ctx), dec, ev.ToolName, ev.Input, "denied")
|
|
return g.block(ev.ToolName, dec.Message, "")
|
|
case "allow":
|
|
// Record explicit rule and model approvals so review details remain auditable.
|
|
g.interceptor.Log(ctx, intercept.ConvIDFromContext(ctx), dec, ev.ToolName, ev.Input, "allowed")
|
|
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{}
|
|
}
|
|
|
|
// judgeSubject builds the text the LLM fallback judge evaluates: the shell
|
|
// command for shell tools, or the raw tool input JSON for others (Write/Edit/…)
|
|
// so the model sees what is being written/requested.
|
|
func judgeSubject(cmd string, input []byte) string {
|
|
if cmd != "" {
|
|
return cmd
|
|
}
|
|
return string(input)
|
|
}
|
|
|
|
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
|
|
}
|