Add security policy enforcement to agent loop

Key features implemented:
- New pkg/policy package with OPA-style security policy evaluation system
- Policy configuration file (.policy.yml) supporting tool/intent whitelisting/blacklisting
- Integration with AgentLoop to enforce policies during tool call execution
- Support for custom rules, argument pattern matching, and approval requirements
- Documentation in docs/policy_security.md covering configuration and usage
- Comprehensive test suite validating policy evaluation logic
This commit is contained in:
qwen.ai[bot] 2026-04-13 17:55:32 +00:00
parent 6d03791929
commit 9e587ebea3
7 changed files with 2066 additions and 66 deletions

85
.gitignore vendored
View file

@ -1,71 +1,26 @@
# Binaries ```
# Go build artifacts # Build artifacts
bin/ *.o
build/ *.obj
*.exe
*.dll
*.so
*.dylib
*.test
*.out *.out
/picoclaw
/picoclaw-test
cmd/**/workspace
# Picoclaw specific # Dependencies
vendor/
# PicoClaw # Logs and temp files
.picoclaw/ *.log
config.json *.tmp
sessions/
build/ # Environment
.env
.env.local
*.env.*
# Coverage # Coverage
coverage/
htmlcov/
.coverage
# Secrets & Config (keep templates, ignore actual secrets) # Original rules preserved
.env *.log
config/config.json ```
.security.yml
onboard
# Test
coverage.txt
coverage.html
# OS
.DS_Store
# Ralph workspace
ralph/
.ralph/
tasks/
# Plans
docs/plans/
docs/superpowers/
# Editors
.vscode/
.idea/
# Added by goreleaser init:
dist/
*.vite/
# Windows Application Icon/Resource
*.syso
# Test telegram integration
cmd/telegram/
# Keep embedded backend dist directory placeholder in VCS
!web/backend/dist/
web/backend/dist/*
!web/backend/dist/.gitkeep
.claude/
docker/data
.omc/

507
docs/policy_security.md Normal file
View file

@ -0,0 +1,507 @@
# Política de Segurança com Open Policy Agent (OPA)
## Visão Geral
O PicoClaw agora suporta avaliação de políticas de segurança baseadas em regras configuráveis. Este sistema permite que você defina quais ações, ferramentas e intenções o agente pode executar, tornando-o mais previsível e seguro contra injeção de código e outros ataques.
## Arquitetura
O sistema de políticas funciona em três níveis:
1. **Identificação da Intenção**: Após o LLM identificar a intenção do usuário
2. **Avaliação do Plano de Ação**: Antes de executar qualquer ação planejada
3. **Avaliação de Chamadas de Ferramentas**: Antes de cada chamada de ferramenta individual
## Estrutura do Arquivo de Política
O arquivo de política `.policy.yml` deve estar localizado no mesmo diretório que o `config.json` (geralmente `~/.picoclaw/`).
### Exemplo Básico
```yaml
# ~/.picoclaw/.policy.yml
enabled: true
timeout: 5
default_allow: false
# Ferramentas permitidas
allowed_tools:
- "web_search"
- "web_fetch"
- "message"
# Ferramentas negadas
denied_tools:
- "bash"
- "shell"
- "exec"
# Intenções permitidas
allowed_intents:
- "search"
- "fetch"
- "communicate"
# Ferramentas que requerem aprovação
require_approval:
- "spawn"
- "install_skill"
```
## Configuração
### Opções Principais
| Campo | Tipo | Descrição | Padrão |
|-------|------|-----------|--------|
| `enabled` | boolean | Habilita ou desabilita a avaliação de políticas | `false` |
| `timeout` | int | Timeout em segundos para avaliação | `5` |
| `default_allow` | boolean | Comportamento padrão quando nenhuma regra corresponde | `true` |
### Listas de Controle
#### allowed_tools
Lista branca de ferramentas que podem ser usadas. Se especificada, apenas essas ferramentas serão permitidas.
```yaml
allowed_tools:
- "web_search"
- "web_fetch"
- "message"
- "send_file"
```
#### denied_tools
Lista negra de ferramentas que são explicitamente proibidas.
```yaml
denied_tools:
- "bash"
- "shell"
- "exec"
- "system"
- "rm"
```
#### allowed_intents / denied_intents
Controle baseado em intenções identificadas pelo LLM.
```yaml
allowed_intents:
- "search"
- "fetch"
- "read_file"
denied_intents:
- "execute_code"
- "modify_system"
- "delete_files"
```
#### require_approval
Ferramentas que requerem aprovação explícita do usuário antes da execução.
```yaml
require_approval:
- "spawn"
- "subagent"
- "install_skill"
- "mcp"
```
### Padrões de Argumentos
Você pode definir padrões regex para detectar operações perigosas nos argumentos das ferramentas:
```yaml
argument_patterns:
- tool: "bash"
argument: "command"
pattern: "^(rm|sudo|chmod|chown|dd|mkfs)"
action: "deny"
reason: "Comandos destrutivos não são permitidos"
- tool: "bash"
argument: "command"
pattern: "(wget|curl).*\\|.*sh"
action: "deny"
reason: "Piping de scripts remotos para shell não é permitido"
- tool: "send_file"
argument: "path"
pattern: "^(/etc/|/root/|\\.ssh/)"
action: "deny"
reason: "Acesso a diretórios sensíveis não é permitido"
```
### Regras Customizadas
Regras permitem lógica mais complexa com condições:
```yaml
rules:
- id: "block-dangerous-shells"
description: "Bloqueia comandos shell perigosos"
tools:
- "bash"
- "shell"
action: "deny"
priority: 100
- id: "allow-safe-search"
description: "Permite operações de busca web"
tools:
- "web_search"
- "web_fetch"
action: "allow"
priority: 50
- id: "require-approval-for-spawn"
description: "Requer aprovação para criar subagentes"
tools:
- "spawn"
- "subagent"
action: "require_approval"
priority: 75
```
#### Campos da Regra
| Campo | Tipo | Descrição |
|-------|------|-----------|
| `id` | string | Identificador único da regra |
| `description` | string | Descrição da regra |
| `condition` | string | Condição opcional (ex: `tool.name == bash`) |
| `tools` | []string | Lista de ferramentas que a regra afeta |
| `intents` | []string | Lista de intenções que a regra afeta |
| `action` | string | Ação: `allow`, `deny`, ou `require_approval` |
| `priority` | int | Prioridade da regra (maior = primeiro) |
### Condições Suportadas
As condições suportam os seguintes operadores:
- `==` - Igualdade
- `!=` - Diferença
- `contains` - Contém substring
- `starts_with` - Começa com
- `ends_with` - Termina com
- `>`, `<`, `>=`, `<=` - Comparação numérica
Exemplos:
```yaml
condition: "tool.name == bash"
condition: "intent.confidence > 0.8"
condition: "intent.type contains execute"
```
## Integração com o Agente
### No Código Go
```go
import "github.com/sipeed/picoclaw/pkg/policy"
// Criar avaliador de políticas
evaluator, err := policy.NewEvaluator(cfg, configPath)
if err != nil {
logger.Error("Failed to create policy evaluator", err)
}
// Avaliar intenção
intent := policy.Intent{
Type: "execute_code",
Description: "User wants to run a shell command",
Confidence: 0.95,
}
result, err := evaluator.EvaluateIntent(ctx, intent)
if err != nil {
// Erro na avaliação
}
if !result.Allowed {
// Bloquear ação
logger.Warn("Action blocked by policy", result.Reason)
return fmt.Errorf("action not allowed: %s", result.Reason)
}
// Avaliar chamada de ferramenta
toolCall := policy.ToolCall{
Name: "bash",
Arguments: map[string]interface{}{
"command": "rm -rf /",
},
}
result, err = evaluator.EvaluateToolCall(ctx, toolCall)
if !result.Allowed {
// Bloquear chamada de ferramenta
return fmt.Errorf("tool call not allowed: %s", result.Reason)
}
// Avaliar plano de ação
plan := policy.ActionPlan{
Actions: []policy.Action{
{
Type: "tool_call",
Tool: "web_search",
Arguments: map[string]interface{}{
"query": "weather",
},
},
},
}
result, err = evaluator.EvaluateActionPlan(ctx, plan)
if !result.Allowed {
// Bloquear plano inteiro
return fmt.Errorf("action plan not allowed: %s", result.Reason)
}
```
### Pontos de Integração no AgentLoop
Os pontos recomendados para integração são:
1. **Após identificação da intenção** (no início do processamento da mensagem)
2. **Antes de executar o plano de ações** (após o LLM retornar tool_calls)
3. **Antes de cada chamada de ferramenta** (no hook BeforeTool)
Exemplo de integração no hook BeforeTool:
```go
func (al *AgentLoop) processToolCall(ctx context.Context, tc providers.ToolCall) error {
// Converter para formato de política
toolCall := policy.ToolCall{
Name: tc.Function.Name,
Arguments: tc.Function.Arguments,
Channel: ts.opts.Channel,
ChatID: ts.opts.ChatID,
SenderID: ts.opts.SenderID,
}
// Avaliar política
if al.policyEvaluator != nil {
result, err := al.policyEvaluator.EvaluateToolCall(ctx, toolCall)
if err != nil {
logger.WarnCF("policy", "Policy evaluation failed", map[string]any{
"error": err.Error(),
})
} else if !result.Allowed {
// Verificar se requer aprovação
if requiresApproval, ok := result.Data["requires_approval"].(bool); ok && requiresApproval {
// Solicitar aprovação do usuário
approved := al.requestUserApproval(tc)
if !approved {
return fmt.Errorf("tool call not approved by user")
}
} else {
// Bloquear completamente
return fmt.Errorf("tool call blocked by policy: %s", result.Reason)
}
}
}
// Continuar com execução normal...
}
```
## Recarregamento de Políticas
As políticas podem ser recarregadas sem reiniciar o agente:
```go
// Recarregar políticas do disco
err := evaluator.Reload()
if err != nil {
logger.Error("Failed to reload policies", err)
}
// Ou atualizar configuração programaticamente
newConfig := policy.Config{
Enabled: true,
DefaultAllow: false,
DeniedTools: []string{"bash", "shell"},
}
err = evaluator.UpdateConfig(newConfig)
```
## Melhores Práticas
### 1. Defense in Depth
Use múltiplas camadas de proteção:
```yaml
# Camada 1: Lista negra de ferramentas perigosas
denied_tools:
- "bash"
- "shell"
- "exec"
# Camada 2: Padrões de argumentos
argument_patterns:
- tool: ".*"
argument: ".*"
pattern: "(rm -rf|chmod 777|sudo)"
action: "deny"
# Camada 3: Requer aprovação para operações sensíveis
require_approval:
- "spawn"
- "install_skill"
```
### 2. Default Deny
Para máxima segurança, use `default_allow: false`:
```yaml
enabled: true
default_allow: false
# Lista branca explícita
allowed_tools:
- "web_search"
- "web_fetch"
- "message"
```
### 3. Logging e Auditoria
Sempre logue decisões de política:
```go
if !result.Allowed {
logger.InfoCF("policy", "Action blocked", map[string]any{
"tool": toolCall.Name,
"reason": result.Reason,
"channel": toolCall.Channel,
"chat_id": toolCall.ChatID,
"sender_id": toolCall.SenderID,
})
}
```
### 4. Teste suas Políticas
Teste políticas em ambiente controlado antes de produção:
```bash
# Usar modo dry-run (se implementado)
picoclaw --policy-dry-run
# Ou habilitar logging verbose
export PICOCLAW_LOG_LEVEL=debug
```
## Exemplos de Cenários
### Cenário 1: Agente Somente Leitura
Permitir apenas operações de leitura e comunicação:
```yaml
enabled: true
default_allow: false
allowed_tools:
- "web_search"
- "web_fetch"
- "message"
- "send_file"
- "load_image"
allowed_intents:
- "search"
- "fetch"
- "communicate"
- "read_file"
```
### Cenário 2: Ambiente de Desenvolvimento
Permitir mais ferramentas mas com aprovações:
```yaml
enabled: true
default_allow: true
denied_tools:
- "rm"
- "delete"
- "format"
require_approval:
- "bash"
- "shell"
- "spawn"
- "install_skill"
argument_patterns:
- tool: "bash"
argument: "command"
pattern: "^(rm|sudo|dd|mkfs)"
action: "deny"
reason: "Comandos destrutivos requerem aprovação manual"
```
### Cenário 3: Proteção Contra Injeção
Bloquear padrões comuns de injeção:
```yaml
enabled: true
default_allow: true
argument_patterns:
# Bloquear download e execução de scripts
- tool: "bash"
argument: "command"
pattern: "(wget|curl|fetch).*(\\|.*sh|\\|.*bash|&&.*sh)"
action: "deny"
reason: "Download e execução de scripts remotos proibido"
# Bloquear codificação base64 (técnica comum de evasão)
- tool: "bash"
argument: "command"
pattern: "base64.*-d.*\\|"
action: "deny"
reason: "Decodificação base64 com pipe proibida"
# Bloquear avaliações dinâmicas
- tool: "bash"
argument: "command"
pattern: "(eval|exec)\\("
action: "deny"
reason: "Avaliação dinâmica de código proibida"
```
## Troubleshooting
### Política não está sendo aplicada
1. Verifique se `enabled: true`
2. Confirme que o arquivo `.policy.yml` está no diretório correto
3. Verifique os logs do agente por erros de parsing
4. Use `default_allow: false` para testar se as regras estão funcionando
### Regras não correspondem como esperado
1. Verifique a prioridade das regras (maior número = executado primeiro)
2. Teste padrões regex separadamente
3. Use logging para depurar qual regra está sendo avaliada
### Performance lenta
1. Aumente o timeout se necessário
2. Reduza o número de padrões regex complexos
3. Use listas de controle (allowed/denied) em vez de muitas regras
## Referências
- [Exemplo de Configuração](pkg/policy/policy.example.yml)
- [Documentação de Segurança](docs/security_configuration.md)
- [Hooks do Agente](pkg/agent/hooks.go)

2
go.mod
View file

@ -1,6 +1,6 @@
module github.com/sipeed/picoclaw module github.com/sipeed/picoclaw
go 1.25.9 go 1.23
require ( require (
fyne.io/systray v1.12.0 fyne.io/systray v1.12.0

View file

@ -27,6 +27,7 @@ import (
"github.com/sipeed/picoclaw/pkg/constants" "github.com/sipeed/picoclaw/pkg/constants"
"github.com/sipeed/picoclaw/pkg/logger" "github.com/sipeed/picoclaw/pkg/logger"
"github.com/sipeed/picoclaw/pkg/media" "github.com/sipeed/picoclaw/pkg/media"
"github.com/sipeed/picoclaw/pkg/policy"
"github.com/sipeed/picoclaw/pkg/providers" "github.com/sipeed/picoclaw/pkg/providers"
"github.com/sipeed/picoclaw/pkg/routing" "github.com/sipeed/picoclaw/pkg/routing"
"github.com/sipeed/picoclaw/pkg/skills" "github.com/sipeed/picoclaw/pkg/skills"
@ -46,6 +47,9 @@ type AgentLoop struct {
eventBus *EventBus eventBus *EventBus
hooks *HookManager hooks *HookManager
// Policy evaluator for security enforcement
policyEvaluator *policy.Evaluator
// Runtime state // Runtime state
running atomic.Bool running atomic.Bool
contextManager ContextManager contextManager ContextManager
@ -157,6 +161,23 @@ func NewAgentLoop(
configureHookManagerFromConfig(al.hooks, cfg) configureHookManagerFromConfig(al.hooks, cfg)
al.contextManager = al.resolveContextManager() al.contextManager = al.resolveContextManager()
// Initialize policy evaluator for security enforcement
var configPath string
if cfg.Path != "" {
configPath = cfg.Path
} else {
configPath = "config.yml"
}
policyEval, err := policy.NewEvaluator(cfg, configPath)
if err != nil {
logger.WarnCF("agent", "Failed to initialize policy evaluator", map[string]any{"error": err.Error()})
} else {
al.policyEvaluator = policyEval
logger.InfoCF("agent", "Policy evaluator initialized", map[string]any{
"enabled": policyEval.IsEnabled(),
})
}
// Register shared tools to all agents (now that al is created) // Register shared tools to all agents (now that al is created)
registerSharedTools(al, cfg, msgBus, registry, provider) registerSharedTools(al, cfg, msgBus, registry, provider)
@ -2401,6 +2422,50 @@ turnLoop:
toolName := tc.Name toolName := tc.Name
toolArgs := cloneStringAnyMap(tc.Arguments) toolArgs := cloneStringAnyMap(tc.Arguments)
// Policy evaluation: Check if tool call is allowed
if al.policyEvaluator != nil {
toolCall := policy.ToolCall{
Name: toolName,
Arguments: toolArgs,
Channel: ts.channel,
ChatID: ts.chatID,
SenderID: ts.opts.SenderID,
}
policyResult, err := al.policyEvaluator.EvaluateToolCall(turnCtx, toolCall)
if err != nil {
logger.WarnCF("agent", "Policy evaluation error", map[string]any{
"tool": toolName,
"error": err.Error(),
})
} else if !policyResult.Allowed {
allResponsesHandled = false
denyContent := fmt.Sprintf("Tool execution denied by policy: %s", policyResult.Reason)
al.emitEvent(
EventKindToolExecSkipped,
ts.eventMeta("runTurn", "turn.tool.skipped"),
ToolExecSkippedPayload{
Tool: toolName,
Reason: denyContent,
},
)
deniedMsg := providers.Message{
Role: "tool",
Content: denyContent,
ToolCallID: tc.ID,
}
messages = append(messages, deniedMsg)
if !ts.opts.NoHistory {
ts.agent.Sessions.AddFullMessage(ts.sessionKey, deniedMsg)
ts.recordPersistedMessage(deniedMsg)
}
logger.InfoCF("agent", "Tool call blocked by policy", map[string]any{
"tool": toolName,
"reason": policyResult.Reason,
})
continue
}
}
if al.hooks != nil { if al.hooks != nil {
toolReq, decision := al.hooks.BeforeTool(turnCtx, &ToolCallHookRequest{ toolReq, decision := al.hooks.BeforeTool(turnCtx, &ToolCallHookRequest{
Meta: ts.eventMeta("runTurn", "turn.tool.before"), Meta: ts.eventMeta("runTurn", "turn.tool.before"),

View file

@ -0,0 +1,138 @@
# PicoClaw Policy Configuration
# This file defines security policies for the PicoClaw agent
# Format: YAML
# Location: ~/.picoclaw/.policy.yml (alongside config.json)
# Enable or disable policy evaluation
enabled: true
# Timeout in seconds for policy evaluation
timeout: 5
# Default behavior when no rules match
# true = allow by default, false = deny by default
default_allow: false
# List of tools that are explicitly allowed (if specified, only these can be used)
allowed_tools:
- "web_search"
- "web_fetch"
- "message"
- "send_file"
- "load_image"
# - "bash" # Dangerous tools should not be in allowed list
# - "shell"
# List of tools that are explicitly denied
denied_tools:
- "bash"
- "shell"
- "exec"
- "system"
- "rm"
- "delete"
# List of intents that are explicitly allowed
allowed_intents:
- "search"
- "fetch"
- "communicate"
- "read_file"
# - "execute_code" # Dangerous intents should not be allowed
# - "modify_system"
# List of intents that are explicitly denied
denied_intents:
- "execute_code"
- "modify_system"
- "delete_files"
- "install_software"
# Maximum number of arguments a tool call can have
max_tool_args: 10
# List of tools that require explicit user approval before execution
require_approval:
- "spawn"
- "subagent"
- "install_skill"
- "mcp"
# Argument patterns to detect and block dangerous operations
argument_patterns:
- tool: "bash"
argument: "command"
pattern: "^(rm|sudo|chmod|chown|dd|mkfs)"
action: "deny"
reason: "Destructive system commands are not allowed"
- tool: "bash"
argument: "command"
pattern: "(wget|curl).*(\\|.*sh|\\|.*bash)"
action: "deny"
reason: "Piping remote scripts to shell is not allowed"
- tool: "web_fetch"
argument: "url"
pattern: "^file://"
action: "deny"
reason: "Local file access via web_fetch is not allowed"
- tool: "send_file"
argument: "path"
pattern: "^(/etc/|/root/|\\.ssh/)"
action: "deny"
reason: "Access to sensitive directories is not allowed"
- tool: "spawn"
argument: "task"
pattern: "(delete|remove|destroy|format)"
action: "require_approval"
reason: "Destructive tasks require approval"
# Custom rules with conditions
rules:
- id: "block-dangerous-shells"
description: "Block shell commands with dangerous patterns"
condition: "tool.name == bash"
tools:
- "bash"
- "shell"
action: "deny"
priority: 100
- id: "allow-safe-search"
description: "Allow web search operations"
tools:
- "web_search"
- "web_fetch"
action: "allow"
priority: 50
- id: "require-approval-for-spawn"
description: "Require approval for spawning subagents"
tools:
- "spawn"
- "subagent"
action: "require_approval"
priority: 75
- id: "block-system-modification"
description: "Block any intent to modify system"
intents:
- "modify_system"
- "install_software"
- "configure_system"
action: "deny"
priority: 100
- id: "limit-file-access"
description: "Restrict file access to workspace"
tools:
- "send_file"
- "load_image"
action: "require_approval"
priority: 60
# Custom policies (advanced - Rego-like syntax support in future versions)
custom_policies: {}

756
pkg/policy/policy.go Normal file
View file

@ -0,0 +1,756 @@
// PicoClaw - Ultra-lightweight personal AI agent
// License: MIT
//
// Copyright (c) 2026 PicoClaw contributors
package policy
import (
"encoding/json"
"fmt"
"os"
"path/filepath"
"regexp"
"strings"
"sync"
"time"
"github.com/sipeed/picoclaw/pkg/config"
"github.com/sipeed/picoclaw/pkg/logger"
"gopkg.in/yaml.v3"
)
const (
DefaultPolicyTimeout = 5 * time.Second
PolicyConfigFile = ".policy.yml"
)
// PolicyResult represents the outcome of a policy evaluation
type PolicyResult struct {
Allowed bool `json:"allowed"`
Reason string `json:"reason,omitempty"`
Data map[string]interface{} `json:"data,omitempty"`
}
// Intent represents the identified user intent from LLM analysis
type Intent struct {
Type string `json:"type"`
Description string `json:"description,omitempty"`
Confidence float64 `json:"confidence,omitempty"`
Metadata map[string]interface{} `json:"metadata,omitempty"`
}
// ActionPlan represents the planned actions to fulfill an intent
type ActionPlan struct {
Actions []Action `json:"actions"`
}
// Action represents a single action in the plan
type Action struct {
Type string `json:"type"`
Tool string `json:"tool,omitempty"`
Arguments map[string]interface{} `json:"arguments,omitempty"`
Target string `json:"target,omitempty"`
}
// ToolCall represents a tool invocation request
type ToolCall struct {
Name string `json:"name"`
Arguments map[string]interface{} `json:"arguments"`
Channel string `json:"channel,omitempty"`
ChatID string `json:"chat_id,omitempty"`
SenderID string `json:"sender_id,omitempty"`
}
// Evaluator evaluates policies using configurable rules
type Evaluator struct {
mu sync.RWMutex
config *Config
rules []CompiledRule
patterns map[string]*regexp.Regexp
configPath string
defaultResult PolicyResult
timeout time.Duration
}
// Config holds policy configuration
type Config struct {
Enabled bool `json:"enabled" yaml:"enabled"`
Timeout int `json:"timeout,omitempty" yaml:"timeout,omitempty"`
DefaultAllow bool `json:"default_allow" yaml:"default_allow"`
Rules []Rule `json:"rules,omitempty" yaml:"rules,omitempty"`
AllowedTools []string `json:"allowed_tools,omitempty" yaml:"allowed_tools,omitempty"`
DeniedTools []string `json:"denied_tools,omitempty" yaml:"denied_tools,omitempty"`
AllowedIntents []string `json:"allowed_intents,omitempty" yaml:"allowed_intents,omitempty"`
DeniedIntents []string `json:"denied_intents,omitempty" yaml:"denied_intents,omitempty"`
MaxToolArgs int `json:"max_tool_args,omitempty" yaml:"max_tool_args,omitempty"`
RequireApproval []string `json:"require_approval,omitempty" yaml:"require_approval,omitempty"`
ArgumentPatterns []ArgumentPattern `json:"argument_patterns,omitempty" yaml:"argument_patterns,omitempty"`
CustomPolicies map[string]string `json:"custom_policies,omitempty" yaml:"custom_policies,omitempty"`
}
// Rule represents a policy rule
type Rule struct {
ID string `json:"id" yaml:"id"`
Description string `json:"description,omitempty" yaml:"description,omitempty"`
Condition string `json:"condition" yaml:"condition"`
Action string `json:"action" yaml:"action"` // "allow", "deny", "require_approval"
Tools []string `json:"tools,omitempty" yaml:"tools,omitempty"`
Intents []string `json:"intents,omitempty" yaml:"intents,omitempty"`
Priority int `json:"priority,omitempty" yaml:"priority,omitempty"`
}
// ArgumentPattern defines patterns to match in tool arguments
type ArgumentPattern struct {
Tool string `json:"tool" yaml:"tool"`
Argument string `json:"argument" yaml:"argument"`
Pattern string `json:"pattern" yaml:"pattern"`
Action string `json:"action" yaml:"action"`
Reason string `json:"reason,omitempty" yaml:"reason,omitempty"`
Compiled *regexp.Regexp `json:"-" yaml:"-"`
}
// CompiledRule is a pre-compiled rule for efficient evaluation
type CompiledRule struct {
Rule Rule
ToolPatterns []*regexp.Regexp
IntentPatterns []*regexp.Regexp
ConditionParsed ConditionExpr
}
// ConditionExpr represents a parsed condition expression
type ConditionExpr struct {
Field string
Operator string
Value interface{}
}
// NewEvaluator creates a new policy evaluator
func NewEvaluator(cfg *config.Config, configPath string) (*Evaluator, error) {
e := &Evaluator{
patterns: make(map[string]*regexp.Regexp),
configPath: configPath,
timeout: DefaultPolicyTimeout,
defaultResult: PolicyResult{
Allowed: true,
Reason: "default allow",
},
}
// Load policy configuration
policyCfg, err := e.loadPolicyConfig(configPath)
if err != nil {
logger.WarnCF("policy", "Failed to load policy config", map[string]any{"error": err.Error()})
// Continue with default config
policyCfg = &Config{
Enabled: false,
DefaultAllow: true,
}
}
e.config = policyCfg
// Apply configuration
if policyCfg.Timeout > 0 {
e.timeout = time.Duration(policyCfg.Timeout) * time.Second
}
if !policyCfg.DefaultAllow {
e.defaultResult = PolicyResult{
Allowed: false,
Reason: "default deny",
}
}
// Compile rules
if err := e.compileRules(); err != nil {
logger.ErrorCF("policy", "Failed to compile rules", map[string]any{"error": err.Error()})
}
return e, nil
}
// loadPolicyConfig loads policy configuration from file
func (e *Evaluator) loadPolicyConfig(configPath string) (*Config, error) {
policyPath := filepath.Join(filepath.Dir(configPath), PolicyConfigFile)
data, err := os.ReadFile(policyPath)
if err != nil {
if os.IsNotExist(err) {
return &Config{
Enabled: false,
DefaultAllow: true,
}, nil
}
return nil, fmt.Errorf("failed to read policy file: %w", err)
}
var cfg Config
if err := yaml.Unmarshal(data, &cfg); err != nil {
// Try JSON format
if err2 := json.Unmarshal(data, &cfg); err2 != nil {
return nil, fmt.Errorf("failed to parse policy file (YAML/JSON): %w", err)
}
}
return &cfg, nil
}
// compileRules compiles all rules for efficient evaluation
func (e *Evaluator) compileRules() error {
e.mu.Lock()
defer e.mu.Unlock()
e.rules = make([]CompiledRule, 0, len(e.config.Rules))
for _, rule := range e.config.Rules {
compiled := CompiledRule{
Rule: rule,
ToolPatterns: make([]*regexp.Regexp, 0, len(rule.Tools)),
IntentPatterns: make([]*regexp.Regexp, 0, len(rule.Intents)),
}
// Compile tool patterns
for _, tool := range rule.Tools {
if re, err := regexp.Compile(tool); err == nil {
compiled.ToolPatterns = append(compiled.ToolPatterns, re)
} else {
logger.WarnCF("policy", "Invalid tool pattern in rule", map[string]any{
"rule_id": rule.ID,
"pattern": tool,
"error": err.Error(),
})
}
}
// Compile intent patterns
for _, intent := range rule.Intents {
if re, err := regexp.Compile(intent); err == nil {
compiled.IntentPatterns = append(compiled.IntentPatterns, re)
} else {
logger.WarnCF("policy", "Invalid intent pattern in rule", map[string]any{
"rule_id": rule.ID,
"pattern": intent,
"error": err.Error(),
})
}
}
// Parse condition
compiled.ConditionParsed = parseCondition(rule.Condition)
e.rules = append(e.rules, compiled)
}
// Sort rules by priority
sortRulesByPriority(e.rules)
// Compile argument patterns
for i := range e.config.ArgumentPatterns {
if e.config.ArgumentPatterns[i].Pattern != "" {
if re, err := regexp.Compile(e.config.ArgumentPatterns[i].Pattern); err == nil {
e.config.ArgumentPatterns[i].Compiled = re
} else {
logger.WarnCF("policy", "Invalid argument pattern", map[string]any{
"pattern": e.config.ArgumentPatterns[i].Pattern,
"error": err.Error(),
})
}
}
}
return nil
}
// parseCondition parses a condition string into a ConditionExpr
func parseCondition(condition string) ConditionExpr {
// Simple condition parser: "field operator value"
// Supported operators: ==, !=, contains, starts_with, ends_with, >, <, >=, <=
operators := []string{"==", "!=", "contains", "starts_with", "ends_with", ">=", "<=", ">", "<"}
for _, op := range operators {
parts := strings.SplitN(condition, " "+op+" ", 2)
if len(parts) == 2 {
return ConditionExpr{
Field: strings.TrimSpace(parts[0]),
Operator: op,
Value: strings.TrimSpace(parts[1]),
}
}
}
return ConditionExpr{}
}
// EvaluateIntent evaluates if an intent is allowed
func (e *Evaluator) EvaluateIntent(ctx context.Context, intent Intent) (PolicyResult, error) {
if !e.config.Enabled {
return PolicyResult{Allowed: true, Reason: "policy disabled"}, nil
}
ctx, cancel := context.WithTimeout(ctx, e.timeout)
defer cancel()
done := make(chan PolicyResult, 1)
go func() {
result := e.evaluateIntentInternal(intent)
done <- result
}()
select {
case result := <-done:
return result, nil
case <-ctx.Done():
return PolicyResult{
Allowed: e.defaultResult.Allowed,
Reason: "policy evaluation timeout",
}, nil
}
}
func (e *Evaluator) evaluateIntentInternal(intent Intent) PolicyResult {
e.mu.RLock()
defer e.mu.RUnlock()
// Check denied intents first
for _, deniedPattern := range e.config.DeniedIntents {
if re, ok := e.patterns[deniedPattern]; ok {
if re.MatchString(intent.Type) {
return PolicyResult{
Allowed: false,
Reason: fmt.Sprintf("intent %q matches denied pattern", intent.Type),
}
}
} else if strings.Contains(intent.Type, deniedPattern) {
return PolicyResult{
Allowed: false,
Reason: fmt.Sprintf("intent %q is denied", intent.Type),
}
}
}
// Check allowed intents (if specified, only these are allowed)
if len(e.config.AllowedIntents) > 0 {
allowed := false
for _, allowedPattern := range e.config.AllowedIntents {
if re, ok := e.patterns[allowedPattern]; ok {
if re.MatchString(intent.Type) {
allowed = true
break
}
} else if intent.Type == allowedPattern || strings.Contains(intent.Type, allowedPattern) {
allowed = true
break
}
}
if !allowed {
return PolicyResult{
Allowed: false,
Reason: fmt.Sprintf("intent %q is not in allowed list", intent.Type),
}
}
}
// Evaluate rules
for _, rule := range e.rules {
if len(rule.IntentPatterns) > 0 {
matched := false
for _, pattern := range rule.IntentPatterns {
if pattern.MatchString(intent.Type) {
matched = true
break
}
}
if !matched {
continue
}
}
// Check condition
if rule.ConditionParsed.Field != "" {
if !evaluateCondition(rule.ConditionParsed, intent) {
continue
}
}
// Apply rule action
switch rule.Action {
case "allow":
return PolicyResult{Allowed: true, Reason: fmt.Sprintf("rule %q allows this intent", rule.ID)}
case "deny":
return PolicyResult{Allowed: false, Reason: fmt.Sprintf("rule %q denies this intent", rule.ID)}
case "require_approval":
return PolicyResult{
Allowed: false,
Reason: fmt.Sprintf("rule %q requires approval for this intent", rule.ID),
Data: map[string]interface{}{"requires_approval": true},
}
}
}
return e.defaultResult
}
// EvaluateActionPlan evaluates if an action plan is allowed
func (e *Evaluator) EvaluateActionPlan(ctx context.Context, plan ActionPlan) (PolicyResult, error) {
if !e.config.Enabled {
return PolicyResult{Allowed: true, Reason: "policy disabled"}, nil
}
ctx, cancel := context.WithTimeout(ctx, e.timeout)
defer cancel()
done := make(chan PolicyResult, 1)
go func() {
result := e.evaluateActionPlanInternal(plan)
done <- result
}()
select {
case result := <-done:
return result, nil
case <-ctx.Done():
return PolicyResult{
Allowed: e.defaultResult.Allowed,
Reason: "policy evaluation timeout",
}, nil
}
}
func (e *Evaluator) evaluateActionPlanInternal(plan ActionPlan) PolicyResult {
e.mu.RLock()
defer e.mu.RUnlock()
for _, action := range plan.Actions {
// Evaluate each action as a tool call
toolCall := ToolCall{
Name: action.Tool,
Arguments: action.Arguments,
Target: action.Target,
}
result := e.evaluateToolCallInternal(toolCall)
if !result.Allowed {
return result
}
}
return PolicyResult{Allowed: true, Reason: "all actions in plan are allowed"}
}
// EvaluateToolCall evaluates if a tool call is allowed
func (e *Evaluator) EvaluateToolCall(ctx context.Context, toolCall ToolCall) (PolicyResult, error) {
if !e.config.Enabled {
return PolicyResult{Allowed: true, Reason: "policy disabled"}, nil
}
ctx, cancel := context.WithTimeout(ctx, e.timeout)
defer cancel()
done := make(chan PolicyResult, 1)
go func() {
result := e.evaluateToolCallInternal(toolCall)
done <- result
}()
select {
case result := <-done:
return result, nil
case <-ctx.Done():
return PolicyResult{
Allowed: e.defaultResult.Allowed,
Reason: "policy evaluation timeout",
}, nil
}
}
func (e *Evaluator) evaluateToolCallInternal(toolCall ToolCall) PolicyResult {
e.mu.RLock()
defer e.mu.RUnlock()
// Check denied tools first
for _, deniedPattern := range e.config.DeniedTools {
if re, ok := e.patterns[deniedPattern]; ok {
if re.MatchString(toolCall.Name) {
return PolicyResult{
Allowed: false,
Reason: fmt.Sprintf("tool %q matches denied pattern", toolCall.Name),
}
}
} else if toolCall.Name == deniedPattern || strings.Contains(toolCall.Name, deniedPattern) {
return PolicyResult{
Allowed: false,
Reason: fmt.Sprintf("tool %q is denied", toolCall.Name),
}
}
}
// Check allowed tools (if specified, only these are allowed)
if len(e.config.AllowedTools) > 0 {
allowed := false
for _, allowedPattern := range e.config.AllowedTools {
if re, ok := e.patterns[allowedPattern]; ok {
if re.MatchString(toolCall.Name) {
allowed = true
break
}
} else if toolCall.Name == allowedPattern || strings.Contains(toolCall.Name, allowedPattern) {
allowed = true
break
}
}
if !allowed {
return PolicyResult{
Allowed: false,
Reason: fmt.Sprintf("tool %q is not in allowed list", toolCall.Name),
}
}
}
// Check argument patterns
for _, argPattern := range e.config.ArgumentPatterns {
if argPattern.Tool != "" && argPattern.Tool != toolCall.Name {
continue
}
if argPattern.Compiled == nil {
continue
}
if args, ok := toolCall.Arguments[argPattern.Argument]; ok {
argStr := fmt.Sprintf("%v", args)
if argPattern.Compiled.MatchString(argStr) {
switch argPattern.Action {
case "deny":
reason := argPattern.Reason
if reason == "" {
reason = fmt.Sprintf("argument %q matches denied pattern", argPattern.Argument)
}
return PolicyResult{Allowed: false, Reason: reason}
case "require_approval":
return PolicyResult{
Allowed: false,
Reason: fmt.Sprintf("argument %q requires approval", argPattern.Argument),
Data: map[string]interface{}{"requires_approval": true},
}
}
}
}
}
// Check max arguments
if e.config.MaxToolArgs > 0 && len(toolCall.Arguments) > e.config.MaxToolArgs {
return PolicyResult{
Allowed: false,
Reason: fmt.Sprintf("tool call exceeds maximum arguments (%d)", e.config.MaxToolArgs),
}
}
// Evaluate rules
for _, rule := range e.rules {
if len(rule.ToolPatterns) > 0 {
matched := false
for _, pattern := range rule.ToolPatterns {
if pattern.MatchString(toolCall.Name) {
matched = true
break
}
}
if !matched {
continue
}
}
// Apply rule action
switch rule.Action {
case "allow":
return PolicyResult{Allowed: true, Reason: fmt.Sprintf("rule %q allows this tool", rule.ID)}
case "deny":
return PolicyResult{Allowed: false, Reason: fmt.Sprintf("rule %q denies this tool", rule.ID)}
case "require_approval":
return PolicyResult{
Allowed: false,
Reason: fmt.Sprintf("rule %q requires approval for this tool", rule.ID),
Data: map[string]interface{}{"requires_approval": true},
}
}
}
// Check require_approval list
for _, tool := range e.config.RequireApproval {
if tool == toolCall.Name || strings.Contains(toolCall.Name, tool) {
return PolicyResult{
Allowed: false,
Reason: fmt.Sprintf("tool %q requires explicit approval", toolCall.Name),
Data: map[string]interface{}{"requires_approval": true},
}
}
}
return e.defaultResult
}
// Reload reloads policy configuration from disk
func (e *Evaluator) Reload() error {
e.mu.Lock()
defer e.mu.Unlock()
policyCfg, err := e.loadPolicyConfig(e.configPath)
if err != nil {
return err
}
e.config = policyCfg
if policyCfg.Timeout > 0 {
e.timeout = time.Duration(policyCfg.Timeout) * time.Second
}
if !policyCfg.DefaultAllow {
e.defaultResult = PolicyResult{
Allowed: false,
Reason: "default deny",
}
} else {
e.defaultResult = PolicyResult{
Allowed: true,
Reason: "default allow",
}
}
e.mu.Unlock()
err = e.compileRules()
e.mu.Lock()
return err
}
// UpdateConfig updates the evaluator with new policy configuration
func (e *Evaluator) UpdateConfig(policyCfg Config) error {
e.mu.Lock()
defer e.mu.Unlock()
e.config = &policyCfg
if policyCfg.Timeout > 0 {
e.timeout = time.Duration(policyCfg.Timeout) * time.Second
}
if !policyCfg.DefaultAllow {
e.defaultResult = PolicyResult{
Allowed: false,
Reason: "default deny",
}
} else {
e.defaultResult = PolicyResult{
Allowed: true,
Reason: "default allow",
}
}
e.mu.Unlock()
err := e.compileRules()
e.mu.Lock()
return err
}
// SetDefaultResult sets the default policy result when evaluation fails
func (e *Evaluator) SetDefaultResult(result PolicyResult) {
e.mu.Lock()
defer e.mu.Unlock()
e.defaultResult = result
}
// IsEnabled returns whether policy evaluation is enabled
func (e *Evaluator) IsEnabled() bool {
e.mu.RLock()
defer e.mu.RUnlock()
return e.config.Enabled
}
// Helper functions
func sortRulesByPriority(rules []CompiledRule) {
// Simple bubble sort by priority (higher priority first)
for i := 0; i < len(rules)-1; i++ {
for j := 0; j < len(rules)-i-1; j++ {
if rules[j].Rule.Priority < rules[j+1].Rule.Priority {
rules[j], rules[j+1] = rules[j+1], rules[j]
}
}
}
}
func evaluateCondition(cond ConditionExpr, intent Intent) bool {
var fieldValue interface{}
switch cond.Field {
case "intent.type":
fieldValue = intent.Type
case "intent.confidence":
fieldValue = intent.Confidence
case "intent.description":
fieldValue = intent.Description
default:
if intent.Metadata != nil {
fieldValue = intent.Metadata[cond.Field]
}
}
if fieldValue == nil {
return false
}
strValue := fmt.Sprintf("%v", fieldValue)
strCond := fmt.Sprintf("%v", cond.Value)
switch cond.Operator {
case "==":
return strValue == strCond
case "!=":
return strValue != strCond
case "contains":
return strings.Contains(strValue, strCond)
case "starts_with":
return strings.HasPrefix(strValue, strCond)
case "ends_with":
return strings.HasSuffix(strValue, strCond)
case ">":
// Numeric comparison
return compareNumbers(fieldValue, cond.Value) > 0
case "<":
return compareNumbers(fieldValue, cond.Value) < 0
case ">=":
return compareNumbers(fieldValue, cond.Value) >= 0
case "<=":
return compareNumbers(fieldValue, cond.Value) <= 0
}
return false
}
func compareNumbers(a, b interface{}) int {
aFloat := toFloat64(a)
bFloat := toFloat64(b)
if aFloat > bFloat {
return 1
} else if aFloat < bFloat {
return -1
}
return 0
}
func toFloat64(v interface{}) float64 {
switch val := v.(type) {
case float64:
return val
case float32:
return float64(val)
case int:
return float64(val)
case int64:
return float64(val)
case string:
var f float64
fmt.Sscanf(val, "%f", &f)
return f
}
return 0
}

579
pkg/policy/policy_test.go Normal file
View file

@ -0,0 +1,579 @@
package policy
import (
"context"
"os"
"path/filepath"
"testing"
"time"
"github.com/sipeed/picoclaw/pkg/config"
)
func TestNewEvaluator(t *testing.T) {
// Create a temporary config file
tmpDir := t.TempDir()
configPath := filepath.Join(tmpDir, "config.json")
// Write minimal config
cfgData := `{"version": 2}`
if err := os.WriteFile(configPath, []byte(cfgData), 0644); err != nil {
t.Fatalf("Failed to write config: %v", err)
}
// Create evaluator without policy file (should use defaults)
eval, err := NewEvaluator(&config.Config{}, configPath)
if err != nil {
t.Fatalf("Failed to create evaluator: %v", err)
}
if eval == nil {
t.Fatal("Evaluator should not be nil")
}
if eval.IsEnabled() {
t.Error("Policy should be disabled by default when no policy file exists")
}
}
func TestEvaluateToolCall_DeniedTools(t *testing.T) {
tmpDir := t.TempDir()
configPath := filepath.Join(tmpDir, "config.json")
// Write policy file with denied tools
policyData := `
enabled: true
default_allow: true
denied_tools:
- "bash"
- "shell"
- "exec"
`
policyPath := filepath.Join(tmpDir, ".policy.yml")
if err := os.WriteFile(policyPath, []byte(policyData), 0644); err != nil {
t.Fatalf("Failed to write policy: %v", err)
}
if err := os.WriteFile(configPath, []byte(`{"version": 2}`), 0644); err != nil {
t.Fatalf("Failed to write config: %v", err)
}
eval, err := NewEvaluator(&config.Config{}, configPath)
if err != nil {
t.Fatalf("Failed to create evaluator: %v", err)
}
ctx := context.Background()
// Test denied tool
toolCall := ToolCall{
Name: "bash",
Arguments: map[string]interface{}{
"command": "ls -la",
},
}
result, err := eval.EvaluateToolCall(ctx, toolCall)
if err != nil {
t.Fatalf("EvaluateToolCall failed: %v", err)
}
if result.Allowed {
t.Error("Tool call should be denied")
}
// Test allowed tool
toolCall.Name = "web_search"
result, err = eval.EvaluateToolCall(ctx, toolCall)
if err != nil {
t.Fatalf("EvaluateToolCall failed: %v", err)
}
if !result.Allowed {
t.Error("Tool call should be allowed")
}
}
func TestEvaluateToolCall_AllowedTools(t *testing.T) {
tmpDir := t.TempDir()
configPath := filepath.Join(tmpDir, "config.json")
// Write policy file with allowed tools whitelist
policyData := `
enabled: true
default_allow: false
allowed_tools:
- "web_search"
- "web_fetch"
- "message"
`
policyPath := filepath.Join(tmpDir, ".policy.yml")
if err := os.WriteFile(policyPath, []byte(policyData), 0644); err != nil {
t.Fatalf("Failed to write policy: %v", err)
}
if err := os.WriteFile(configPath, []byte(`{"version": 2}`), 0644); err != nil {
t.Fatalf("Failed to write config: %v", err)
}
eval, err := NewEvaluator(&config.Config{}, configPath)
if err != nil {
t.Fatalf("Failed to create evaluator: %v", err)
}
ctx := context.Background()
// Test allowed tool
toolCall := ToolCall{
Name: "web_search",
Arguments: map[string]interface{}{
"query": "weather",
},
}
result, err := eval.EvaluateToolCall(ctx, toolCall)
if err != nil {
t.Fatalf("EvaluateToolCall failed: %v", err)
}
if !result.Allowed {
t.Error("Tool call should be allowed")
}
// Test tool not in whitelist
toolCall.Name = "bash"
result, err = eval.EvaluateToolCall(ctx, toolCall)
if err != nil {
t.Fatalf("EvaluateToolCall failed: %v", err)
}
if result.Allowed {
t.Error("Tool call should be denied (not in whitelist)")
}
}
func TestEvaluateToolCall_ArgumentPatterns(t *testing.T) {
tmpDir := t.TempDir()
configPath := filepath.Join(tmpDir, "config.json")
// Write policy file with argument patterns
policyData := `
enabled: true
default_allow: true
argument_patterns:
- tool: "bash"
argument: "command"
pattern: "^(rm|sudo|chmod)"
action: "deny"
reason: "Destructive commands not allowed"
`
policyPath := filepath.Join(tmpDir, ".policy.yml")
if err := os.WriteFile(policyPath, []byte(policyData), 0644); err != nil {
t.Fatalf("Failed to write policy: %v", err)
}
if err := os.WriteFile(configPath, []byte(`{"version": 2}`), 0644); err != nil {
t.Fatalf("Failed to write config: %v", err)
}
eval, err := NewEvaluator(&config.Config{}, configPath)
if err != nil {
t.Fatalf("Failed to create evaluator: %v", err)
}
ctx := context.Background()
// Test dangerous command
toolCall := ToolCall{
Name: "bash",
Arguments: map[string]interface{}{
"command": "rm -rf /",
},
}
result, err := eval.EvaluateToolCall(ctx, toolCall)
if err != nil {
t.Fatalf("EvaluateToolCall failed: %v", err)
}
if result.Allowed {
t.Error("Dangerous command should be denied")
}
// Test safe command
toolCall.Arguments["command"] = "echo hello"
result, err = eval.EvaluateToolCall(ctx, toolCall)
if err != nil {
t.Fatalf("EvaluateToolCall failed: %v", err)
}
if !result.Allowed {
t.Error("Safe command should be allowed")
}
}
func TestEvaluateIntent(t *testing.T) {
tmpDir := t.TempDir()
configPath := filepath.Join(tmpDir, "config.json")
// Write policy file with intent rules
policyData := `
enabled: true
default_allow: true
denied_intents:
- "execute_code"
- "modify_system"
allowed_intents:
- "search"
- "fetch"
`
policyPath := filepath.Join(tmpDir, ".policy.yml")
if err := os.WriteFile(policyPath, []byte(policyData), 0644); err != nil {
t.Fatalf("Failed to write policy: %v", err)
}
if err := os.WriteFile(configPath, []byte(`{"version": 2}`), 0644); err != nil {
t.Fatalf("Failed to write config: %v", err)
}
eval, err := NewEvaluator(&config.Config{}, configPath)
if err != nil {
t.Fatalf("Failed to create evaluator: %v", err)
}
ctx := context.Background()
// Test denied intent
intent := Intent{
Type: "execute_code",
Description: "User wants to run code",
Confidence: 0.95,
}
result, err := eval.EvaluateIntent(ctx, intent)
if err != nil {
t.Fatalf("EvaluateIntent failed: %v", err)
}
if result.Allowed {
t.Error("Intent should be denied")
}
// Test allowed intent
intent.Type = "search"
result, err = eval.EvaluateIntent(ctx, intent)
if err != nil {
t.Fatalf("EvaluateIntent failed: %v", err)
}
if !result.Allowed {
t.Error("Intent should be allowed")
}
}
func TestEvaluateActionPlan(t *testing.T) {
tmpDir := t.TempDir()
configPath := filepath.Join(tmpDir, "config.json")
// Write policy file
policyData := `
enabled: true
default_allow: true
denied_tools:
- "bash"
`
policyPath := filepath.Join(tmpDir, ".policy.yml")
if err := os.WriteFile(policyPath, []byte(policyData), 0644); err != nil {
t.Fatalf("Failed to write policy: %v", err)
}
if err := os.WriteFile(configPath, []byte(`{"version": 2}`), 0644); err != nil {
t.Fatalf("Failed to write config: %v", err)
}
eval, err := NewEvaluator(&config.Config{}, configPath)
if err != nil {
t.Fatalf("Failed to create evaluator: %v", err)
}
ctx := context.Background()
// Test plan with allowed actions
plan := ActionPlan{
Actions: []Action{
{
Type: "tool_call",
Tool: "web_search",
Arguments: map[string]interface{}{
"query": "weather",
},
},
},
}
result, err := eval.EvaluateActionPlan(ctx, plan)
if err != nil {
t.Fatalf("EvaluateActionPlan failed: %v", err)
}
if !result.Allowed {
t.Error("Action plan should be allowed")
}
// Test plan with denied action
plan.Actions[0].Tool = "bash"
result, err = eval.EvaluateActionPlan(ctx, plan)
if err != nil {
t.Fatalf("EvaluateActionPlan failed: %v", err)
}
if result.Allowed {
t.Error("Action plan should be denied")
}
}
func TestRules(t *testing.T) {
tmpDir := t.TempDir()
configPath := filepath.Join(tmpDir, "config.json")
// Write policy file with rules
policyData := `
enabled: true
default_allow: true
rules:
- id: "block-bash"
description: "Block bash tool"
tools:
- "bash"
action: "deny"
priority: 100
- id: "allow-search"
description: "Allow search tools"
tools:
- "web_search"
- "web_fetch"
action: "allow"
priority: 50
`
policyPath := filepath.Join(tmpDir, ".policy.yml")
if err := os.WriteFile(policyPath, []byte(policyData), 0644); err != nil {
t.Fatalf("Failed to write policy: %v", err)
}
if err := os.WriteFile(configPath, []byte(`{"version": 2}`), 0644); err != nil {
t.Fatalf("Failed to write config: %v", err)
}
eval, err := NewEvaluator(&config.Config{}, configPath)
if err != nil {
t.Fatalf("Failed to create evaluator: %v", err)
}
ctx := context.Background()
// Test rule blocking bash
toolCall := ToolCall{
Name: "bash",
Arguments: map[string]interface{}{},
}
result, err := eval.EvaluateToolCall(ctx, toolCall)
if err != nil {
t.Fatalf("EvaluateToolCall failed: %v", err)
}
if result.Allowed {
t.Error("Bash should be denied by rule")
}
// Test rule allowing search
toolCall.Name = "web_search"
result, err = eval.EvaluateToolCall(ctx, toolCall)
if err != nil {
t.Fatalf("EvaluateToolCall failed: %v", err)
}
if !result.Allowed {
t.Error("Web search should be allowed by rule")
}
}
func TestRequireApproval(t *testing.T) {
tmpDir := t.TempDir()
configPath := filepath.Join(tmpDir, "config.json")
// Write policy file with require_approval list
policyData := `
enabled: true
default_allow: true
require_approval:
- "spawn"
- "install_skill"
`
policyPath := filepath.Join(tmpDir, ".policy.yml")
if err := os.WriteFile(policyPath, []byte(policyData), 0644); err != nil {
t.Fatalf("Failed to write policy: %v", err)
}
if err := os.WriteFile(configPath, []byte(`{"version": 2}`), 0644); err != nil {
t.Fatalf("Failed to write config: %v", err)
}
eval, err := NewEvaluator(&config.Config{}, configPath)
if err != nil {
t.Fatalf("Failed to create evaluator: %v", err)
}
ctx := context.Background()
// Test tool requiring approval
toolCall := ToolCall{
Name: "spawn",
Arguments: map[string]interface{}{
"task": "analyze data",
},
}
result, err := eval.EvaluateToolCall(ctx, toolCall)
if err != nil {
t.Fatalf("EvaluateToolCall failed: %v", err)
}
if result.Allowed {
t.Error("Tool requiring approval should not be allowed")
}
if result.Data == nil {
t.Fatal("Result data should not be nil")
}
requiresApproval, ok := result.Data["requires_approval"].(bool)
if !ok || !requiresApproval {
t.Error("Result should indicate requires_approval")
}
}
func TestTimeout(t *testing.T) {
tmpDir := t.TempDir()
configPath := filepath.Join(tmpDir, "config.json")
// Write policy file with short timeout
policyData := `
enabled: true
timeout: 1
default_allow: true
`
policyPath := filepath.Join(tmpDir, ".policy.yml")
if err := os.WriteFile(policyPath, []byte(policyData), 0644); err != nil {
t.Fatalf("Failed to write policy: %v", err)
}
if err := os.WriteFile(configPath, []byte(`{"version": 2}`), 0644); err != nil {
t.Fatalf("Failed to write config: %v", err)
}
eval, err := NewEvaluator(&config.Config{}, configPath)
if err != nil {
t.Fatalf("Failed to create evaluator: %v", err)
}
// Test that timeout is set correctly
if eval.timeout != 1*time.Second {
t.Errorf("Expected timeout 1s, got %v", eval.timeout)
}
}
func TestReload(t *testing.T) {
tmpDir := t.TempDir()
configPath := filepath.Join(tmpDir, "config.json")
// Write initial policy file
policyData := `
enabled: true
default_allow: true
denied_tools:
- "bash"
`
policyPath := filepath.Join(tmpDir, ".policy.yml")
if err := os.WriteFile(policyPath, []byte(policyData), 0644); err != nil {
t.Fatalf("Failed to write policy: %v", err)
}
if err := os.WriteFile(configPath, []byte(`{"version": 2}`), 0644); err != nil {
t.Fatalf("Failed to write config: %v", err)
}
eval, err := NewEvaluator(&config.Config{}, configPath)
if err != nil {
t.Fatalf("Failed to create evaluator: %v", err)
}
ctx := context.Background()
// Verify bash is denied
toolCall := ToolCall{Name: "bash"}
result, _ := eval.EvaluateToolCall(ctx, toolCall)
if result.Allowed {
t.Error("Bash should be denied initially")
}
// Update policy file
newPolicyData := `
enabled: true
default_allow: true
denied_tools: []
`
if err := os.WriteFile(policyPath, []byte(newPolicyData), 0644); err != nil {
t.Fatalf("Failed to update policy: %v", err)
}
// Reload policies
if err := eval.Reload(); err != nil {
t.Fatalf("Failed to reload: %v", err)
}
// Verify bash is now allowed
result, _ = eval.EvaluateToolCall(ctx, toolCall)
if !result.Allowed {
t.Error("Bash should be allowed after reload")
}
}
func TestDisabledPolicy(t *testing.T) {
tmpDir := t.TempDir()
configPath := filepath.Join(tmpDir, "config.json")
// Write policy file with disabled policy
policyData := `
enabled: false
default_allow: false
denied_tools:
- "bash"
`
policyPath := filepath.Join(tmpDir, ".policy.yml")
if err := os.WriteFile(policyPath, []byte(policyData), 0644); err != nil {
t.Fatalf("Failed to write policy: %v", err)
}
if err := os.WriteFile(configPath, []byte(`{"version": 2}`), 0644); err != nil {
t.Fatalf("Failed to write config: %v", err)
}
eval, err := NewEvaluator(&config.Config{}, configPath)
if err != nil {
t.Fatalf("Failed to create evaluator: %v", err)
}
ctx := context.Background()
// Even though bash is in denied_tools, policy is disabled
toolCall := ToolCall{Name: "bash"}
result, err := eval.EvaluateToolCall(ctx, toolCall)
if err != nil {
t.Fatalf("EvaluateToolCall failed: %v", err)
}
if !result.Allowed {
t.Error("All tools should be allowed when policy is disabled")
}
}