feat: add configurable security policy engine with IM-based approval
Introduce a three-mode security policy system (off/block/approve) for exec guard, SSRF protection, path validation, and skill validation. All modes default to "off" for full backward compatibility. In "approve" mode, violations pause tool execution and send an approval request to the user via their active IM channel, supporting multi-language approval/denial keywords (EN/ZH/JA). Key changes: - Add PolicyEngine and approval workflow (pkg/security/) - Add MessageBus inbound interceptor for approval message routing - Integrate policy checks into ExecTool, WebFetchTool, filesystem tools, edit tools, and SkillInstaller - Wire PolicyEngine through agent loop and cron tool setup - Add SecurityConfig with env var support to pkg/config - Update config.example.json with new security section - Update README.md, README.zh.md, README.ja.md with security policy modes and IM approval documentation - Add comprehensive tests for all new functionality
This commit is contained in:
parent
e087e314f1
commit
f9a579260b
25 changed files with 1574 additions and 105 deletions
63
README.ja.md
63
README.ja.md
|
|
@ -548,7 +548,7 @@ PicoClaw はデフォルトでサンドボックス環境で実行されます
|
|||
|
||||
#### 組み込み Exec 保護ルール
|
||||
|
||||
`restrict_to_workspace: false` でも、`exec` ツールには**常に有効**な組み込み拒否ルールがあり、設定で上書きすることはできません:
|
||||
`security.exec_guard` が `"block"` または `"approve"` に設定されている場合、`exec` ツールには設定で上書きできない組み込み拒否ルールがあります:
|
||||
|
||||
| カテゴリ | ブロックパターン | 説明 |
|
||||
|---------|----------------|------|
|
||||
|
|
@ -574,7 +574,7 @@ PicoClaw はデフォルトでサンドボックス環境で実行されます
|
|||
|
||||
#### SSRF 保護
|
||||
|
||||
すべてのアウトバウンド HTTP リクエスト(`web_fetch` ツールおよびチャットチャネルからのファイルダウンロード)は SSRF 攻撃に対して検証されます:
|
||||
`security.ssrf_protection` が `"block"` または `"approve"` に設定されている場合、すべてのアウトバウンド HTTP リクエスト(`web_fetch` ツールおよびファイルダウンロード)は SSRF 攻撃に対して検証されます:
|
||||
|
||||
- プライベート IP 範囲(`10.0.0.0/8`, `172.16.0.0/12`, `192.168.0.0/16`)がブロックされます
|
||||
- ループバックアドレス(`127.0.0.0/8`, `::1`)がブロックされます
|
||||
|
|
@ -630,6 +630,65 @@ export PICOCLAW_AGENTS_DEFAULTS_RESTRICT_TO_WORKSPACE=false
|
|||
|
||||
すべてのパスで同じワークスペース制限が適用されます — サブエージェント、Cron ジョブ、またはスケジュールタスクを通じてセキュリティ境界をバイパスする方法はありません。
|
||||
|
||||
#### セキュリティポリシーモード
|
||||
|
||||
すべてのセキュリティチェック(コマンドガード、SSRF 保護、パス検証、スキル検証)は、3つの設定可能なモードをサポートしています。デフォルトではすべてのモードが `"off"` です — セキュリティ機能は**オプトイン**であり、明示的に設定しない限り動作は変更されません。
|
||||
|
||||
| モード | 動作 |
|
||||
|--------|------|
|
||||
| `off` | セキュリティチェック無効(デフォルト)。制限なし。 |
|
||||
| `block` | 違反が検出されると即座にエラーで拒否。 |
|
||||
| `approve` | 違反が検出されると実行を一時停止し、IM 経由でユーザーに承認リクエストを送信、返信を待機。 |
|
||||
|
||||
<details>
|
||||
<summary><b>セキュリティ設定</b></summary>
|
||||
|
||||
```json
|
||||
{
|
||||
"security": {
|
||||
"exec_guard": "off",
|
||||
"ssrf_protection": "off",
|
||||
"path_validation": "off",
|
||||
"skill_validation": "off",
|
||||
"approval_timeout": 300
|
||||
}
|
||||
}
|
||||
```
|
||||
|
||||
| オプション | デフォルト | 説明 |
|
||||
|-----------|-----------|------|
|
||||
| `exec_guard` | `"off"` | コマンドの拒否/許可パターンチェックのモード |
|
||||
| `ssrf_protection` | `"off"` | アウトバウンド URL 検証(プライベート IP、メタデータエンドポイント) |
|
||||
| `path_validation` | `"off"` | シンボリックリンク対応の強化パス制限 |
|
||||
| `skill_validation` | `"off"` | スキルインストール時のリポジトリ形式チェック |
|
||||
| `approval_timeout` | `300` | ユーザー承認のタイムアウト(秒)。タイムアウト時は自動拒否 |
|
||||
|
||||
環境変数もサポートされています(例: `PICOCLAW_SECURITY_EXEC_GUARD=approve`)。
|
||||
|
||||
</details>
|
||||
|
||||
#### IM ベースの承認メカニズム
|
||||
|
||||
セキュリティチェックが `"approve"` モードに設定されている場合、PicoClaw はコマンドを即座に拒否する代わりに:
|
||||
|
||||
1. ツールの実行を**一時停止**
|
||||
2. 現在の IM チャネル(Telegram、Feishu、DingTalk、Slack など)経由でユーザーに承認リクエストを**送信**
|
||||
3. ユーザーの承認または拒否キーワードの返信を**待機**
|
||||
4. 承認された場合は実行を**再開**、拒否またはタイムアウトの場合はエラーを返却
|
||||
|
||||
**サポートされている承認キーワード:**
|
||||
|
||||
| アクション | 英語 | 中国語 | 日本語 |
|
||||
|-----------|------|--------|--------|
|
||||
| 承認 | approve, yes, allow, ok, y | 批准, 允许, 通过, 是 | 承認, 許可, はい |
|
||||
| 拒否 | deny, no, reject, block, n | 拒绝, 否决, 不 | 拒否, いいえ |
|
||||
|
||||
**注意事項:**
|
||||
- CLI モードでは、非同期 IM チャネルがないため、`"approve"` は `"block"` にフォールバックします。
|
||||
- Cron ジョブの承認リクエストは、最後にアクティブだった IM チャネルに送信されます。利用可能なチャネルがない場合は `"block"` にフォールバックします。
|
||||
- 承認待機中にユーザーが送信した承認キーワード以外のメッセージは、通常通りエージェントに渡されます。
|
||||
- `approval_timeout` 秒以内に返信がない場合、リクエストは自動的に拒否されます。
|
||||
|
||||
### ハートビート(定期タスク)
|
||||
|
||||
PicoClaw は自動的に定期タスクを実行できます。ワークスペースに `HEARTBEAT.md` ファイルを作成します:
|
||||
|
|
|
|||
63
README.md
63
README.md
|
|
@ -574,7 +574,7 @@ You can customize the `exec` tool's security behavior through configuration:
|
|||
|
||||
#### Built-in Exec Protection
|
||||
|
||||
Even with `restrict_to_workspace: false`, the `exec` tool has built-in deny rules that are **always active** and cannot be overridden by configuration:
|
||||
When `security.exec_guard` is set to `"block"` or `"approve"`, the `exec` tool has built-in deny rules that cannot be overridden by configuration:
|
||||
|
||||
| Category | Blocked Pattern | Description |
|
||||
|----------|----------------|-------------|
|
||||
|
|
@ -600,7 +600,7 @@ When `restrict_to_workspace: true`, additional restrictions apply:
|
|||
|
||||
#### SSRF Protection
|
||||
|
||||
All outbound HTTP requests (via `web_fetch` tool and file downloads from chat channels) are validated against SSRF attacks:
|
||||
When `security.ssrf_protection` is set to `"block"` or `"approve"`, all outbound HTTP requests (via `web_fetch` tool and file downloads) are validated against SSRF attacks:
|
||||
|
||||
* Private IP ranges (`10.0.0.0/8`, `172.16.0.0/12`, `192.168.0.0/16`) are blocked
|
||||
* Loopback addresses (`127.0.0.0/8`, `::1`) are blocked
|
||||
|
|
@ -658,6 +658,65 @@ The `restrict_to_workspace` setting applies consistently across all execution pa
|
|||
|
||||
All paths share the same workspace restriction — there's no way to bypass the security boundary through subagents, cron jobs, or scheduled tasks.
|
||||
|
||||
#### Security Policy Modes
|
||||
|
||||
All security checks (exec guard, SSRF protection, path validation, skill validation) support three configurable modes. By default, all modes are set to `"off"` — security features are **opt-in** and do not change behavior unless explicitly configured.
|
||||
|
||||
| Mode | Behavior |
|
||||
|------|----------|
|
||||
| `off` | Security check disabled (default). No enforcement. |
|
||||
| `block` | Violations are immediately rejected with an error. |
|
||||
| `approve` | Violations pause execution, send an approval request to the user via IM, and wait for a reply. |
|
||||
|
||||
<details>
|
||||
<summary><b>Security Configuration</b></summary>
|
||||
|
||||
```json
|
||||
{
|
||||
"security": {
|
||||
"exec_guard": "off",
|
||||
"ssrf_protection": "off",
|
||||
"path_validation": "off",
|
||||
"skill_validation": "off",
|
||||
"approval_timeout": 300
|
||||
}
|
||||
}
|
||||
```
|
||||
|
||||
| Option | Default | Description |
|
||||
|--------|---------|-------------|
|
||||
| `exec_guard` | `"off"` | Mode for command deny/allow pattern checks |
|
||||
| `ssrf_protection` | `"off"` | Mode for outbound URL validation (private IP, metadata endpoints) |
|
||||
| `path_validation` | `"off"` | Mode for enhanced symlink-aware path restriction |
|
||||
| `skill_validation` | `"off"` | Mode for skill installation repository format checks |
|
||||
| `approval_timeout` | `300` | Seconds to wait for user approval before auto-deny |
|
||||
|
||||
Environment variables are also supported (e.g. `PICOCLAW_SECURITY_EXEC_GUARD=approve`).
|
||||
|
||||
</details>
|
||||
|
||||
#### IM-based Approval
|
||||
|
||||
When a security check is set to `"approve"` mode, instead of immediately rejecting the command, PicoClaw will:
|
||||
|
||||
1. **Pause** the tool execution
|
||||
2. **Send** an approval request to the user via the current IM channel (Telegram, Feishu, DingTalk, Slack, etc.)
|
||||
3. **Wait** for the user to reply with an approval or denial keyword
|
||||
4. **Resume** execution if approved, or return an error if denied or timed out
|
||||
|
||||
**Supported approval keywords:**
|
||||
|
||||
| Action | English | Chinese | Japanese |
|
||||
|--------|---------|---------|----------|
|
||||
| Approve | approve, yes, allow, ok, y | 批准, 允许, 通过, 是 | 承認, 許可, はい |
|
||||
| Deny | deny, no, reject, block, n | 拒绝, 否决, 不 | 拒否, いいえ |
|
||||
|
||||
**Notes:**
|
||||
- In CLI mode, `"approve"` falls back to `"block"` since there is no async IM channel.
|
||||
- For cron jobs, the approval request is sent to the last active IM channel; if none is available, it falls back to `"block"`.
|
||||
- Non-approval messages sent during an active approval request are passed through to the agent normally.
|
||||
- If no reply is received within `approval_timeout` seconds, the request is auto-denied.
|
||||
|
||||
### Heartbeat (Periodic Tasks)
|
||||
|
||||
PicoClaw can perform periodic tasks automatically. Create a `HEARTBEAT.md` file in your workspace:
|
||||
|
|
|
|||
63
README.zh.md
63
README.zh.md
|
|
@ -536,7 +536,7 @@ PicoClaw 默认在沙盒环境中运行。Agent 只能访问配置的工作区
|
|||
|
||||
#### 内置 Exec 防护规则
|
||||
|
||||
即使 `restrict_to_workspace: false`,`exec` 工具拥有**始终生效**的内置拦截规则,不可通过配置覆盖:
|
||||
当 `security.exec_guard` 设为 `"block"` 或 `"approve"` 时,`exec` 工具拥有不可通过配置覆盖的内置拦截规则:
|
||||
|
||||
| 类别 | 拦截模式 | 说明 |
|
||||
|------|---------|------|
|
||||
|
|
@ -562,7 +562,7 @@ PicoClaw 默认在沙盒环境中运行。Agent 只能访问配置的工作区
|
|||
|
||||
#### SSRF 防护
|
||||
|
||||
所有出站 HTTP 请求(通过 `web_fetch` 工具和聊天渠道的文件下载)均会进行 SSRF 攻击验证:
|
||||
当 `security.ssrf_protection` 设为 `"block"` 或 `"approve"` 时,所有出站 HTTP 请求(通过 `web_fetch` 工具和文件下载)均会进行 SSRF 攻击验证:
|
||||
|
||||
* 私有 IP 段(`10.0.0.0/8`, `172.16.0.0/12`, `192.168.0.0/16`)被拦截
|
||||
* 回环地址(`127.0.0.0/8`, `::1`)被拦截
|
||||
|
|
@ -608,6 +608,65 @@ export PICOCLAW_AGENTS_DEFAULTS_RESTRICT_TO_WORKSPACE=false
|
|||
|
||||
所有路径共享相同的工作区限制 — 无法通过子 Agent、Cron 定时任务或其他调度方式绕过安全边界。
|
||||
|
||||
#### 安全策略模式
|
||||
|
||||
所有安全检查(命令防护、SSRF 防护、路径验证、技能验证)均支持三种可配置模式。默认所有模式均为 `"off"` — 安全功能为**可选启用**,未明确配置时不会改变任何行为。
|
||||
|
||||
| 模式 | 行为 |
|
||||
|------|------|
|
||||
| `off` | 安全检查禁用(默认)。不进行任何拦截。 |
|
||||
| `block` | 检测到违规时立即拒绝并返回错误。 |
|
||||
| `approve` | 检测到违规时暂停执行,通过 IM 向用户发送审批请求,等待用户回复。 |
|
||||
|
||||
<details>
|
||||
<summary><b>安全配置</b></summary>
|
||||
|
||||
```json
|
||||
{
|
||||
"security": {
|
||||
"exec_guard": "off",
|
||||
"ssrf_protection": "off",
|
||||
"path_validation": "off",
|
||||
"skill_validation": "off",
|
||||
"approval_timeout": 300
|
||||
}
|
||||
}
|
||||
```
|
||||
|
||||
| 选项 | 默认值 | 描述 |
|
||||
|------|--------|------|
|
||||
| `exec_guard` | `"off"` | 命令拦截/放行模式检查 |
|
||||
| `ssrf_protection` | `"off"` | 出站 URL 验证(私有 IP、元数据端点) |
|
||||
| `path_validation` | `"off"` | 增强的符号链接感知路径限制 |
|
||||
| `skill_validation` | `"off"` | 技能安装仓库格式检查 |
|
||||
| `approval_timeout` | `300` | 等待用户审批的超时时间(秒),超时自动拒绝 |
|
||||
|
||||
也支持环境变量(如 `PICOCLAW_SECURITY_EXEC_GUARD=approve`)。
|
||||
|
||||
</details>
|
||||
|
||||
#### 基于 IM 的审批机制
|
||||
|
||||
当安全检查设为 `"approve"` 模式时,PicoClaw 不会直接拒绝命令,而是:
|
||||
|
||||
1. **暂停**工具执行
|
||||
2. **发送**审批请求到用户当前的 IM 渠道(Telegram、飞书、钉钉、Slack 等)
|
||||
3. **等待**用户回复批准或拒绝关键词
|
||||
4. 批准则**继续执行**,拒绝或超时则返回错误
|
||||
|
||||
**支持的审批关键词:**
|
||||
|
||||
| 操作 | 英文 | 中文 | 日文 |
|
||||
|------|------|------|------|
|
||||
| 批准 | approve, yes, allow, ok, y | 批准, 允许, 通过, 是 | 承認, 許可, はい |
|
||||
| 拒绝 | deny, no, reject, block, n | 拒绝, 否决, 不 | 拒否, いいえ |
|
||||
|
||||
**说明:**
|
||||
- CLI 模式下,`"approve"` 模式会退化为 `"block"`,因为 CLI 没有异步 IM 渠道。
|
||||
- Cron 定时任务的审批请求会发送到最近活跃的 IM 渠道;若无可用渠道则退化为 `"block"`。
|
||||
- 审批等待期间,用户发送的非审批关键词消息会正常传递给 Agent。
|
||||
- 若在 `approval_timeout` 秒内未收到回复,请求将自动拒绝。
|
||||
|
||||
### 心跳 / 周期性任务 (Heartbeat)
|
||||
|
||||
PicoClaw 可以自动执行周期性任务。在工作区创建 `HEARTBEAT.md` 文件:
|
||||
|
|
|
|||
|
|
@ -32,6 +32,7 @@ import (
|
|||
"github.com/sipeed/picoclaw/pkg/logger"
|
||||
"github.com/sipeed/picoclaw/pkg/migrate"
|
||||
"github.com/sipeed/picoclaw/pkg/providers"
|
||||
"github.com/sipeed/picoclaw/pkg/security"
|
||||
"github.com/sipeed/picoclaw/pkg/skills"
|
||||
"github.com/sipeed/picoclaw/pkg/state"
|
||||
"github.com/sipeed/picoclaw/pkg/tools"
|
||||
|
|
@ -560,7 +561,7 @@ func gatewayCmd() {
|
|||
})
|
||||
|
||||
// Setup cron tool and service
|
||||
cronService := setupCronTool(agentLoop, msgBus, cfg.WorkspacePath(), cfg.Agents.Defaults.RestrictToWorkspace)
|
||||
cronService := setupCronTool(agentLoop, msgBus, cfg.WorkspacePath(), cfg.Agents.Defaults.RestrictToWorkspace, cfg)
|
||||
|
||||
heartbeatService := heartbeat.NewHeartbeatService(
|
||||
cfg.WorkspacePath(),
|
||||
|
|
@ -973,14 +974,23 @@ func getConfigPath() string {
|
|||
return filepath.Join(home, ".picoclaw", "config.json")
|
||||
}
|
||||
|
||||
func setupCronTool(agentLoop *agent.AgentLoop, msgBus *bus.MessageBus, workspace string, restrict bool) *cron.CronService {
|
||||
func setupCronTool(agentLoop *agent.AgentLoop, msgBus *bus.MessageBus, workspace string, restrict bool, cfg *config.Config) *cron.CronService {
|
||||
cronStorePath := filepath.Join(workspace, "cron", "jobs.json")
|
||||
|
||||
// Create cron service
|
||||
cronService := cron.NewCronService(cronStorePath, nil)
|
||||
|
||||
// Create and register CronTool (uses same workspace restriction as main agent)
|
||||
cronTool := tools.NewCronTool(cronService, agentLoop, msgBus, workspace, restrict)
|
||||
// Create PolicyEngine for cron exec tool
|
||||
pe := security.NewPolicyEngine(&cfg.Security, msgBus)
|
||||
|
||||
// Create and register CronTool (uses same workspace restriction and security policy as main agent)
|
||||
cronTool := tools.NewCronToolWithConfig(cronService, agentLoop, msgBus, workspace, restrict, tools.ExecToolConfig{
|
||||
DenyPatterns: cfg.Tools.Exec.DenyPatterns,
|
||||
AllowPatterns: cfg.Tools.Exec.AllowPatterns,
|
||||
MaxTimeout: cfg.Tools.Exec.MaxTimeout,
|
||||
PolicyEngine: pe,
|
||||
ExecGuardMode: pe.GetMode("exec_guard"),
|
||||
})
|
||||
agentLoop.RegisterTool(cronTool)
|
||||
|
||||
// Set the onJob handler
|
||||
|
|
|
|||
|
|
@ -122,6 +122,13 @@
|
|||
"max_timeout": 60
|
||||
}
|
||||
},
|
||||
"security": {
|
||||
"exec_guard": "off",
|
||||
"ssrf_protection": "off",
|
||||
"path_validation": "off",
|
||||
"skill_validation": "off",
|
||||
"approval_timeout": 300
|
||||
},
|
||||
"heartbeat": {
|
||||
"enabled": true,
|
||||
"interval": 30
|
||||
|
|
|
|||
|
|
@ -23,6 +23,7 @@ import (
|
|||
"github.com/sipeed/picoclaw/pkg/constants"
|
||||
"github.com/sipeed/picoclaw/pkg/logger"
|
||||
"github.com/sipeed/picoclaw/pkg/providers"
|
||||
"github.com/sipeed/picoclaw/pkg/security"
|
||||
"github.com/sipeed/picoclaw/pkg/session"
|
||||
"github.com/sipeed/picoclaw/pkg/state"
|
||||
"github.com/sipeed/picoclaw/pkg/tools"
|
||||
|
|
@ -61,18 +62,28 @@ type processOptions struct {
|
|||
func createToolRegistry(workspace string, restrict bool, cfg *config.Config, msgBus *bus.MessageBus) *tools.ToolRegistry {
|
||||
registry := tools.NewToolRegistry()
|
||||
|
||||
// Create shared PolicyEngine from security config
|
||||
pe := security.NewPolicyEngine(&cfg.Security, msgBus)
|
||||
|
||||
pathOpts := tools.PathPolicyOpts{
|
||||
PathMode: pe.GetMode("path_validation"),
|
||||
PolicyEngine: pe,
|
||||
}
|
||||
|
||||
// File system tools
|
||||
registry.Register(tools.NewReadFileTool(workspace, restrict))
|
||||
registry.Register(tools.NewWriteFileTool(workspace, restrict))
|
||||
registry.Register(tools.NewListDirTool(workspace, restrict))
|
||||
registry.Register(tools.NewEditFileTool(workspace, restrict))
|
||||
registry.Register(tools.NewAppendFileTool(workspace, restrict))
|
||||
registry.Register(tools.NewReadFileToolWithPolicy(workspace, restrict, pathOpts))
|
||||
registry.Register(tools.NewWriteFileToolWithPolicy(workspace, restrict, pathOpts))
|
||||
registry.Register(tools.NewListDirToolWithPolicy(workspace, restrict, pathOpts))
|
||||
registry.Register(tools.NewEditFileToolWithPolicy(workspace, restrict, pathOpts))
|
||||
registry.Register(tools.NewAppendFileToolWithPolicy(workspace, restrict, pathOpts))
|
||||
|
||||
// Shell execution
|
||||
registry.Register(tools.NewExecToolWithConfig(workspace, restrict, tools.ExecToolConfig{
|
||||
DenyPatterns: cfg.Tools.Exec.DenyPatterns,
|
||||
AllowPatterns: cfg.Tools.Exec.AllowPatterns,
|
||||
MaxTimeout: cfg.Tools.Exec.MaxTimeout,
|
||||
PolicyEngine: pe,
|
||||
ExecGuardMode: pe.GetMode("exec_guard"),
|
||||
}))
|
||||
|
||||
if searchTool := tools.NewWebSearchTool(tools.WebSearchToolOptions{
|
||||
|
|
@ -84,7 +95,11 @@ func createToolRegistry(workspace string, restrict bool, cfg *config.Config, msg
|
|||
}); searchTool != nil {
|
||||
registry.Register(searchTool)
|
||||
}
|
||||
registry.Register(tools.NewWebFetchTool(50000))
|
||||
registry.Register(tools.NewWebFetchToolWithPolicy(tools.WebFetchToolOptions{
|
||||
MaxChars: 50000,
|
||||
PolicyEngine: pe,
|
||||
SSRFMode: pe.GetMode("ssrf"),
|
||||
}))
|
||||
|
||||
// Hardware tools (I2C, SPI) - Linux only, returns error on other platforms
|
||||
registry.Register(tools.NewI2CTool())
|
||||
|
|
|
|||
|
|
@ -3,12 +3,20 @@ package bus
|
|||
import (
|
||||
"context"
|
||||
"sync"
|
||||
"sync/atomic"
|
||||
)
|
||||
|
||||
type interceptorEntry struct {
|
||||
id uint64
|
||||
fn InboundInterceptor
|
||||
}
|
||||
|
||||
type MessageBus struct {
|
||||
inbound chan InboundMessage
|
||||
outbound chan OutboundMessage
|
||||
handlers map[string]MessageHandler
|
||||
interceptors []*interceptorEntry
|
||||
nextID uint64
|
||||
mu sync.RWMutex
|
||||
}
|
||||
|
||||
|
|
@ -20,7 +28,39 @@ func NewMessageBus() *MessageBus {
|
|||
}
|
||||
}
|
||||
|
||||
// AddInterceptor registers an interceptor that inspects inbound messages before
|
||||
// they reach the main consumer queue. Returns a removal function.
|
||||
func (mb *MessageBus) AddInterceptor(fn InboundInterceptor) func() {
|
||||
id := atomic.AddUint64(&mb.nextID, 1)
|
||||
entry := &interceptorEntry{id: id, fn: fn}
|
||||
|
||||
mb.mu.Lock()
|
||||
mb.interceptors = append(mb.interceptors, entry)
|
||||
mb.mu.Unlock()
|
||||
|
||||
return func() {
|
||||
mb.mu.Lock()
|
||||
defer mb.mu.Unlock()
|
||||
for i, e := range mb.interceptors {
|
||||
if e.id == id {
|
||||
mb.interceptors = append(mb.interceptors[:i], mb.interceptors[i+1:]...)
|
||||
break
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (mb *MessageBus) PublishInbound(msg InboundMessage) {
|
||||
mb.mu.RLock()
|
||||
interceptors := make([]*interceptorEntry, len(mb.interceptors))
|
||||
copy(interceptors, mb.interceptors)
|
||||
mb.mu.RUnlock()
|
||||
|
||||
for _, entry := range interceptors {
|
||||
if entry.fn(msg) {
|
||||
return
|
||||
}
|
||||
}
|
||||
mb.inbound <- msg
|
||||
}
|
||||
|
||||
|
|
|
|||
178
pkg/bus/bus_test.go
Normal file
178
pkg/bus/bus_test.go
Normal file
|
|
@ -0,0 +1,178 @@
|
|||
package bus
|
||||
|
||||
import (
|
||||
"context"
|
||||
"sync"
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
|
||||
func TestMessageBus_InterceptorConsumesMessage(t *testing.T) {
|
||||
mb := NewMessageBus()
|
||||
|
||||
consumed := make(chan bool, 1)
|
||||
mb.AddInterceptor(func(msg InboundMessage) bool {
|
||||
if msg.Content == "intercept-me" {
|
||||
consumed <- true
|
||||
return true
|
||||
}
|
||||
return false
|
||||
})
|
||||
|
||||
mb.PublishInbound(InboundMessage{Content: "intercept-me"})
|
||||
|
||||
select {
|
||||
case <-consumed:
|
||||
// ok
|
||||
case <-time.After(time.Second):
|
||||
t.Fatal("interceptor did not consume the message")
|
||||
}
|
||||
|
||||
// Message should not have reached the main queue
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 100*time.Millisecond)
|
||||
defer cancel()
|
||||
_, ok := mb.ConsumeInbound(ctx)
|
||||
if ok {
|
||||
t.Error("intercepted message should not reach main consumer")
|
||||
}
|
||||
}
|
||||
|
||||
func TestMessageBus_InterceptorPassesThrough(t *testing.T) {
|
||||
mb := NewMessageBus()
|
||||
|
||||
mb.AddInterceptor(func(msg InboundMessage) bool {
|
||||
return false // never consume
|
||||
})
|
||||
|
||||
mb.PublishInbound(InboundMessage{Content: "pass-through"})
|
||||
|
||||
ctx, cancel := context.WithTimeout(context.Background(), time.Second)
|
||||
defer cancel()
|
||||
msg, ok := mb.ConsumeInbound(ctx)
|
||||
if !ok {
|
||||
t.Fatal("message should have passed through to main consumer")
|
||||
}
|
||||
if msg.Content != "pass-through" {
|
||||
t.Errorf("expected 'pass-through', got %q", msg.Content)
|
||||
}
|
||||
}
|
||||
|
||||
func TestMessageBus_InterceptorRemoval(t *testing.T) {
|
||||
mb := NewMessageBus()
|
||||
count := 0
|
||||
|
||||
remove := mb.AddInterceptor(func(msg InboundMessage) bool {
|
||||
count++
|
||||
return true
|
||||
})
|
||||
|
||||
mb.PublishInbound(InboundMessage{Content: "first"})
|
||||
if count != 1 {
|
||||
t.Fatalf("expected count=1, got %d", count)
|
||||
}
|
||||
|
||||
remove() // unregister
|
||||
|
||||
mb.PublishInbound(InboundMessage{Content: "second"})
|
||||
|
||||
// Message should now reach main queue
|
||||
ctx, cancel := context.WithTimeout(context.Background(), time.Second)
|
||||
defer cancel()
|
||||
msg, ok := mb.ConsumeInbound(ctx)
|
||||
if !ok || msg.Content != "second" {
|
||||
t.Error("after removal, message should reach main consumer")
|
||||
}
|
||||
if count != 1 {
|
||||
t.Errorf("interceptor should not have been called after removal, count=%d", count)
|
||||
}
|
||||
}
|
||||
|
||||
func TestMessageBus_MultipleInterceptors(t *testing.T) {
|
||||
mb := NewMessageBus()
|
||||
order := []string{}
|
||||
var mu sync.Mutex
|
||||
|
||||
mb.AddInterceptor(func(msg InboundMessage) bool {
|
||||
mu.Lock()
|
||||
order = append(order, "first")
|
||||
mu.Unlock()
|
||||
return false
|
||||
})
|
||||
|
||||
mb.AddInterceptor(func(msg InboundMessage) bool {
|
||||
mu.Lock()
|
||||
order = append(order, "second")
|
||||
mu.Unlock()
|
||||
return msg.Content == "stop-at-second"
|
||||
})
|
||||
|
||||
// This message passes through both interceptors
|
||||
mb.PublishInbound(InboundMessage{Content: "pass"})
|
||||
ctx, cancel := context.WithTimeout(context.Background(), time.Second)
|
||||
defer cancel()
|
||||
_, ok := mb.ConsumeInbound(ctx)
|
||||
if !ok {
|
||||
t.Error("message should reach main consumer")
|
||||
}
|
||||
|
||||
mu.Lock()
|
||||
if len(order) != 2 || order[0] != "first" || order[1] != "second" {
|
||||
t.Errorf("expected [first, second], got %v", order)
|
||||
}
|
||||
order = nil
|
||||
mu.Unlock()
|
||||
|
||||
// This message is consumed by second interceptor
|
||||
mb.PublishInbound(InboundMessage{Content: "stop-at-second"})
|
||||
|
||||
ctx2, cancel2 := context.WithTimeout(context.Background(), 100*time.Millisecond)
|
||||
defer cancel2()
|
||||
_, ok = mb.ConsumeInbound(ctx2)
|
||||
if ok {
|
||||
t.Error("message consumed by second interceptor should not reach main consumer")
|
||||
}
|
||||
}
|
||||
|
||||
func TestMessageBus_InterceptorConcurrency(t *testing.T) {
|
||||
mb := NewMessageBus()
|
||||
var wg sync.WaitGroup
|
||||
intercepted := int32(0)
|
||||
|
||||
remove := mb.AddInterceptor(func(msg InboundMessage) bool {
|
||||
if msg.Content == "catch" {
|
||||
return true
|
||||
}
|
||||
return false
|
||||
})
|
||||
|
||||
// Publish many messages concurrently
|
||||
for i := 0; i < 50; i++ {
|
||||
wg.Add(1)
|
||||
go func(i int) {
|
||||
defer wg.Done()
|
||||
if i%2 == 0 {
|
||||
mb.PublishInbound(InboundMessage{Content: "catch"})
|
||||
} else {
|
||||
mb.PublishInbound(InboundMessage{Content: "pass"})
|
||||
}
|
||||
}(i)
|
||||
}
|
||||
|
||||
wg.Wait()
|
||||
remove()
|
||||
|
||||
// Drain remaining messages
|
||||
for {
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 100*time.Millisecond)
|
||||
msg, ok := mb.ConsumeInbound(ctx)
|
||||
cancel()
|
||||
if !ok {
|
||||
break
|
||||
}
|
||||
if msg.Content == "catch" {
|
||||
t.Error("intercepted message should not reach consumer")
|
||||
}
|
||||
intercepted++
|
||||
}
|
||||
_ = intercepted // just ensure no panics
|
||||
}
|
||||
|
|
@ -17,3 +17,7 @@ type OutboundMessage struct {
|
|||
}
|
||||
|
||||
type MessageHandler func(InboundMessage) error
|
||||
|
||||
// InboundInterceptor inspects an inbound message before it reaches the main consumer.
|
||||
// Returns true if the message was consumed and should not be enqueued.
|
||||
type InboundInterceptor func(msg InboundMessage) bool
|
||||
|
|
|
|||
|
|
@ -49,6 +49,7 @@ type Config struct {
|
|||
Providers ProvidersConfig `json:"providers"`
|
||||
Gateway GatewayConfig `json:"gateway"`
|
||||
Tools ToolsConfig `json:"tools"`
|
||||
Security SecurityConfig `json:"security"`
|
||||
Heartbeat HeartbeatConfig `json:"heartbeat"`
|
||||
Devices DevicesConfig `json:"devices"`
|
||||
mu sync.RWMutex
|
||||
|
|
@ -221,6 +222,17 @@ type ToolsConfig struct {
|
|||
Exec ExecConfig `json:"exec"`
|
||||
}
|
||||
|
||||
// SecurityConfig controls optional security features.
|
||||
// All modes default to "off" to preserve pre-security-modification behavior.
|
||||
// Supported modes: "off" (disabled), "block" (reject), "approve" (IM-based approval).
|
||||
type SecurityConfig struct {
|
||||
ExecGuard string `json:"exec_guard" env:"PICOCLAW_SECURITY_EXEC_GUARD"` // "off" | "block" | "approve"
|
||||
SSRFProtection string `json:"ssrf_protection" env:"PICOCLAW_SECURITY_SSRF_PROTECTION"` // "off" | "block" | "approve"
|
||||
PathValidation string `json:"path_validation" env:"PICOCLAW_SECURITY_PATH_VALIDATION"` // "off" | "block" | "approve"
|
||||
SkillValidation string `json:"skill_validation" env:"PICOCLAW_SECURITY_SKILL_VALIDATION"` // "off" | "block" | "approve"
|
||||
ApprovalTimeout int `json:"approval_timeout" env:"PICOCLAW_SECURITY_APPROVAL_TIMEOUT"` // seconds, default 300
|
||||
}
|
||||
|
||||
func DefaultConfig() *Config {
|
||||
return &Config{
|
||||
Agents: AgentsConfig{
|
||||
|
|
@ -334,6 +346,13 @@ func DefaultConfig() *Config {
|
|||
MaxTimeout: 60,
|
||||
},
|
||||
},
|
||||
Security: SecurityConfig{
|
||||
ExecGuard: "off",
|
||||
SSRFProtection: "off",
|
||||
PathValidation: "off",
|
||||
SkillValidation: "off",
|
||||
ApprovalTimeout: 300,
|
||||
},
|
||||
Heartbeat: HeartbeatConfig{
|
||||
Enabled: true,
|
||||
Interval: 30, // default 30 minutes
|
||||
|
|
|
|||
124
pkg/security/approval.go
Normal file
124
pkg/security/approval.go
Normal file
|
|
@ -0,0 +1,124 @@
|
|||
package security
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/sipeed/picoclaw/pkg/bus"
|
||||
)
|
||||
|
||||
// ApprovalResult carries the user's decision on a security approval request.
|
||||
type ApprovalResult struct {
|
||||
Approved bool
|
||||
Reason string
|
||||
}
|
||||
|
||||
// requestApproval sends an approval notification via IM and blocks until the
|
||||
// user responds with an approval/denial keyword or the timeout expires.
|
||||
func (pe *PolicyEngine) requestApproval(ctx context.Context, v Violation, channel, chatID string) error {
|
||||
resultCh := make(chan ApprovalResult, 1)
|
||||
|
||||
// Register an interceptor to capture the approval reply from the same chat
|
||||
removeInterceptor := pe.bus.AddInterceptor(func(msg bus.InboundMessage) bool {
|
||||
if msg.Channel != channel || msg.ChatID != chatID {
|
||||
return false
|
||||
}
|
||||
content := strings.TrimSpace(msg.Content)
|
||||
lower := strings.ToLower(content)
|
||||
if isApproveKeyword(lower) || isApproveKeywordCJK(content) {
|
||||
resultCh <- ApprovalResult{Approved: true}
|
||||
return true
|
||||
}
|
||||
if isDenyKeyword(lower) || isDenyKeywordCJK(content) {
|
||||
resultCh <- ApprovalResult{Approved: false, Reason: "denied by user"}
|
||||
return true
|
||||
}
|
||||
return false // not an approval keyword, pass through
|
||||
})
|
||||
defer removeInterceptor()
|
||||
|
||||
// Send approval request notification to the user via IM
|
||||
pe.bus.PublishOutbound(bus.OutboundMessage{
|
||||
Channel: channel,
|
||||
ChatID: chatID,
|
||||
Content: formatApprovalMessage(v, pe.config.ApprovalTimeout),
|
||||
})
|
||||
|
||||
timeout := time.Duration(pe.config.ApprovalTimeout) * time.Second
|
||||
if timeout <= 0 {
|
||||
timeout = 300 * time.Second
|
||||
}
|
||||
|
||||
select {
|
||||
case result := <-resultCh:
|
||||
if result.Approved {
|
||||
return nil
|
||||
}
|
||||
return fmt.Errorf("denied by user: %s", result.Reason)
|
||||
case <-time.After(timeout):
|
||||
return fmt.Errorf("approval timed out after %v", timeout)
|
||||
case <-ctx.Done():
|
||||
return ctx.Err()
|
||||
}
|
||||
}
|
||||
|
||||
// formatApprovalMessage builds a human-readable approval notification.
|
||||
func formatApprovalMessage(v Violation, timeoutSec int) string {
|
||||
var b strings.Builder
|
||||
b.WriteString("⚠️ Security Approval Required / 安全审批请求\n\n")
|
||||
b.WriteString(fmt.Sprintf("Category: %s\n", v.Category))
|
||||
if v.Tool != "" {
|
||||
b.WriteString(fmt.Sprintf("Tool: %s\n", v.Tool))
|
||||
}
|
||||
if v.Action != "" {
|
||||
b.WriteString(fmt.Sprintf("Action: %s\n", v.Action))
|
||||
}
|
||||
b.WriteString(fmt.Sprintf("Reason: %s\n", v.Reason))
|
||||
if v.RuleName != "" {
|
||||
b.WriteString(fmt.Sprintf("Rule: %s\n", v.RuleName))
|
||||
}
|
||||
b.WriteString(fmt.Sprintf("\nReply \"approve\" to allow or \"deny\" to block.\n"))
|
||||
b.WriteString(fmt.Sprintf("回复 \"批准\" 允许执行,回复 \"拒绝\" 阻止执行。\n"))
|
||||
if timeoutSec > 0 {
|
||||
b.WriteString(fmt.Sprintf("Auto-deny in %d seconds.\n", timeoutSec))
|
||||
}
|
||||
return b.String()
|
||||
}
|
||||
|
||||
// isApproveKeyword checks lowercase ASCII approval keywords.
|
||||
func isApproveKeyword(lower string) bool {
|
||||
switch lower {
|
||||
case "approve", "yes", "allow", "ok", "y":
|
||||
return true
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
// isApproveKeywordCJK checks CJK approval keywords (case-sensitive).
|
||||
func isApproveKeywordCJK(s string) bool {
|
||||
switch s {
|
||||
case "批准", "允许", "通过", "是", "承認", "許可", "はい":
|
||||
return true
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
// isDenyKeyword checks lowercase ASCII denial keywords.
|
||||
func isDenyKeyword(lower string) bool {
|
||||
switch lower {
|
||||
case "deny", "no", "reject", "block", "n":
|
||||
return true
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
// isDenyKeywordCJK checks CJK denial keywords (case-sensitive).
|
||||
func isDenyKeywordCJK(s string) bool {
|
||||
switch s {
|
||||
case "拒绝", "否决", "不", "拒否", "いいえ":
|
||||
return true
|
||||
}
|
||||
return false
|
||||
}
|
||||
107
pkg/security/approval_test.go
Normal file
107
pkg/security/approval_test.go
Normal file
|
|
@ -0,0 +1,107 @@
|
|||
package security
|
||||
|
||||
import (
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestIsApproveKeyword(t *testing.T) {
|
||||
approve := []string{"approve", "yes", "allow", "ok", "y"}
|
||||
for _, w := range approve {
|
||||
if !isApproveKeyword(w) {
|
||||
t.Errorf("expected %q to be an approve keyword", w)
|
||||
}
|
||||
}
|
||||
notApprove := []string{"deny", "no", "hello", ""}
|
||||
for _, w := range notApprove {
|
||||
if isApproveKeyword(w) {
|
||||
t.Errorf("expected %q to NOT be an approve keyword", w)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestIsApproveKeywordCJK(t *testing.T) {
|
||||
approve := []string{"批准", "允许", "通过", "是", "承認", "許可", "はい"}
|
||||
for _, w := range approve {
|
||||
if !isApproveKeywordCJK(w) {
|
||||
t.Errorf("expected %q to be a CJK approve keyword", w)
|
||||
}
|
||||
}
|
||||
notApprove := []string{"拒绝", "否决", "hello", ""}
|
||||
for _, w := range notApprove {
|
||||
if isApproveKeywordCJK(w) {
|
||||
t.Errorf("expected %q to NOT be a CJK approve keyword", w)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestIsDenyKeyword(t *testing.T) {
|
||||
deny := []string{"deny", "no", "reject", "block", "n"}
|
||||
for _, w := range deny {
|
||||
if !isDenyKeyword(w) {
|
||||
t.Errorf("expected %q to be a deny keyword", w)
|
||||
}
|
||||
}
|
||||
notDeny := []string{"approve", "yes", "hello", ""}
|
||||
for _, w := range notDeny {
|
||||
if isDenyKeyword(w) {
|
||||
t.Errorf("expected %q to NOT be a deny keyword", w)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestIsDenyKeywordCJK(t *testing.T) {
|
||||
deny := []string{"拒绝", "否决", "不", "拒否", "いいえ"}
|
||||
for _, w := range deny {
|
||||
if !isDenyKeywordCJK(w) {
|
||||
t.Errorf("expected %q to be a CJK deny keyword", w)
|
||||
}
|
||||
}
|
||||
notDeny := []string{"批准", "允许", "hello", ""}
|
||||
for _, w := range notDeny {
|
||||
if isDenyKeywordCJK(w) {
|
||||
t.Errorf("expected %q to NOT be a CJK deny keyword", w)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestFormatApprovalMessage(t *testing.T) {
|
||||
msg := formatApprovalMessage(Violation{
|
||||
Category: "exec_guard",
|
||||
Tool: "exec",
|
||||
Action: "rm -rf /tmp",
|
||||
Reason: "dangerous pattern detected",
|
||||
RuleName: `\brm\s+-[rf]`,
|
||||
}, 300)
|
||||
|
||||
// Check essential fields are present
|
||||
checks := []string{
|
||||
"Approval Required",
|
||||
"exec_guard",
|
||||
"exec",
|
||||
"rm -rf /tmp",
|
||||
"dangerous pattern",
|
||||
"300 seconds",
|
||||
"approve",
|
||||
"deny",
|
||||
"批准",
|
||||
"拒绝",
|
||||
}
|
||||
for _, c := range checks {
|
||||
if !containsSubstring(msg, c) {
|
||||
t.Errorf("approval message missing %q:\n%s", c, msg)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func containsSubstring(s, sub string) bool {
|
||||
return len(s) >= len(sub) && (s == sub || len(s) > 0 && containsHelper(s, sub))
|
||||
}
|
||||
|
||||
func containsHelper(s, sub string) bool {
|
||||
for i := 0; i <= len(s)-len(sub); i++ {
|
||||
if s[i:i+len(sub)] == sub {
|
||||
return true
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
92
pkg/security/policy.go
Normal file
92
pkg/security/policy.go
Normal file
|
|
@ -0,0 +1,92 @@
|
|||
package security
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
|
||||
"github.com/sipeed/picoclaw/pkg/bus"
|
||||
"github.com/sipeed/picoclaw/pkg/config"
|
||||
)
|
||||
|
||||
// PolicyMode represents the security enforcement mode.
|
||||
type PolicyMode string
|
||||
|
||||
const (
|
||||
ModeOff PolicyMode = "off" // Security check disabled
|
||||
ModeBlock PolicyMode = "block" // Reject on violation
|
||||
ModeApprove PolicyMode = "approve" // Pause and request IM approval
|
||||
)
|
||||
|
||||
// IsOff returns true when the mode means "no enforcement".
|
||||
// Both the zero value ("") and the explicit "off" string are treated as off.
|
||||
func (m PolicyMode) IsOff() bool {
|
||||
return m == "" || m == ModeOff
|
||||
}
|
||||
|
||||
// Violation describes a security event detected by a guard.
|
||||
type Violation struct {
|
||||
Category string // e.g. "exec_guard", "ssrf", "path_validation", "skill_validation"
|
||||
Tool string // tool name that triggered the violation
|
||||
Action string // the action that was attempted (command, URL, path, etc.)
|
||||
Reason string // human-readable explanation
|
||||
RuleName string // name/pattern of the matched rule
|
||||
}
|
||||
|
||||
// PolicyEngine centralises security policy decisions.
|
||||
type PolicyEngine struct {
|
||||
config *config.SecurityConfig
|
||||
bus *bus.MessageBus
|
||||
}
|
||||
|
||||
// NewPolicyEngine creates a PolicyEngine from configuration and message bus.
|
||||
func NewPolicyEngine(cfg *config.SecurityConfig, msgBus *bus.MessageBus) *PolicyEngine {
|
||||
return &PolicyEngine{
|
||||
config: cfg,
|
||||
bus: msgBus,
|
||||
}
|
||||
}
|
||||
|
||||
// GetMode returns the configured PolicyMode for a given security category.
|
||||
func (pe *PolicyEngine) GetMode(category string) PolicyMode {
|
||||
var raw string
|
||||
switch category {
|
||||
case "exec_guard":
|
||||
raw = pe.config.ExecGuard
|
||||
case "ssrf":
|
||||
raw = pe.config.SSRFProtection
|
||||
case "path_validation":
|
||||
raw = pe.config.PathValidation
|
||||
case "skill_validation":
|
||||
raw = pe.config.SkillValidation
|
||||
default:
|
||||
return ModeOff
|
||||
}
|
||||
switch PolicyMode(raw) {
|
||||
case ModeBlock:
|
||||
return ModeBlock
|
||||
case ModeApprove:
|
||||
return ModeApprove
|
||||
default:
|
||||
return ModeOff
|
||||
}
|
||||
}
|
||||
|
||||
// Evaluate checks a violation against the given mode and returns nil to allow
|
||||
// or an error to deny. In "approve" mode it sends an IM approval request and
|
||||
// blocks until the user responds or the timeout expires.
|
||||
func (pe *PolicyEngine) Evaluate(ctx context.Context, mode PolicyMode, v Violation, channel, chatID string) error {
|
||||
switch {
|
||||
case mode.IsOff():
|
||||
return nil
|
||||
case mode == ModeBlock:
|
||||
return fmt.Errorf("blocked by security policy [%s]: %s", v.Category, v.Reason)
|
||||
case mode == ModeApprove:
|
||||
// CLI channel has no async IM listener; fall back to block
|
||||
if channel == "" || channel == "cli" {
|
||||
return fmt.Errorf("blocked by security policy [%s]: %s (approve mode unavailable in CLI)", v.Category, v.Reason)
|
||||
}
|
||||
return pe.requestApproval(ctx, v, channel, chatID)
|
||||
default:
|
||||
return nil
|
||||
}
|
||||
}
|
||||
241
pkg/security/policy_test.go
Normal file
241
pkg/security/policy_test.go
Normal file
|
|
@ -0,0 +1,241 @@
|
|||
package security
|
||||
|
||||
import (
|
||||
"context"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/sipeed/picoclaw/pkg/bus"
|
||||
"github.com/sipeed/picoclaw/pkg/config"
|
||||
)
|
||||
|
||||
func TestPolicyMode_IsOff(t *testing.T) {
|
||||
tests := []struct {
|
||||
mode PolicyMode
|
||||
want bool
|
||||
}{
|
||||
{ModeOff, true},
|
||||
{"off", true},
|
||||
{"", true},
|
||||
{ModeBlock, false},
|
||||
{ModeApprove, false},
|
||||
{"unknown", false},
|
||||
}
|
||||
for _, tt := range tests {
|
||||
if got := tt.mode.IsOff(); got != tt.want {
|
||||
t.Errorf("PolicyMode(%q).IsOff() = %v, want %v", tt.mode, got, tt.want)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestPolicyEngine_GetMode(t *testing.T) {
|
||||
cfg := &config.SecurityConfig{
|
||||
ExecGuard: "block",
|
||||
SSRFProtection: "approve",
|
||||
PathValidation: "off",
|
||||
SkillValidation: "",
|
||||
}
|
||||
pe := NewPolicyEngine(cfg, nil)
|
||||
|
||||
tests := []struct {
|
||||
category string
|
||||
want PolicyMode
|
||||
}{
|
||||
{"exec_guard", ModeBlock},
|
||||
{"ssrf", ModeApprove},
|
||||
{"path_validation", ModeOff},
|
||||
{"skill_validation", ModeOff},
|
||||
{"unknown_category", ModeOff},
|
||||
}
|
||||
for _, tt := range tests {
|
||||
if got := pe.GetMode(tt.category); got != tt.want {
|
||||
t.Errorf("GetMode(%q) = %v, want %v", tt.category, got, tt.want)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestPolicyEngine_Evaluate_Off(t *testing.T) {
|
||||
pe := NewPolicyEngine(&config.SecurityConfig{}, nil)
|
||||
err := pe.Evaluate(context.Background(), ModeOff, Violation{Reason: "test"}, "telegram", "chat1")
|
||||
if err != nil {
|
||||
t.Errorf("ModeOff should allow, got: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestPolicyEngine_Evaluate_Block(t *testing.T) {
|
||||
pe := NewPolicyEngine(&config.SecurityConfig{}, nil)
|
||||
err := pe.Evaluate(context.Background(), ModeBlock, Violation{
|
||||
Category: "exec_guard",
|
||||
Reason: "dangerous pattern",
|
||||
}, "telegram", "chat1")
|
||||
if err == nil {
|
||||
t.Fatal("ModeBlock should reject")
|
||||
}
|
||||
if !strings.Contains(err.Error(), "blocked") {
|
||||
t.Errorf("error should contain 'blocked', got: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestPolicyEngine_Evaluate_Approve_CLIFallback(t *testing.T) {
|
||||
pe := NewPolicyEngine(&config.SecurityConfig{ApprovalTimeout: 5}, bus.NewMessageBus())
|
||||
err := pe.Evaluate(context.Background(), ModeApprove, Violation{
|
||||
Category: "exec_guard",
|
||||
Reason: "test",
|
||||
}, "cli", "direct")
|
||||
if err == nil {
|
||||
t.Fatal("CLI should fall back to block")
|
||||
}
|
||||
if !strings.Contains(err.Error(), "unavailable in CLI") {
|
||||
t.Errorf("error should mention CLI, got: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestPolicyEngine_Evaluate_Approve_Approved(t *testing.T) {
|
||||
msgBus := bus.NewMessageBus()
|
||||
pe := NewPolicyEngine(&config.SecurityConfig{ApprovalTimeout: 5}, msgBus)
|
||||
|
||||
// Start approval in goroutine
|
||||
errCh := make(chan error, 1)
|
||||
go func() {
|
||||
errCh <- pe.Evaluate(context.Background(), ModeApprove, Violation{
|
||||
Category: "exec_guard",
|
||||
Tool: "exec",
|
||||
Action: "rm -rf /tmp/test",
|
||||
Reason: "dangerous pattern",
|
||||
}, "telegram", "chat123")
|
||||
}()
|
||||
|
||||
// Consume the outbound approval notification
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
|
||||
defer cancel()
|
||||
outMsg, ok := msgBus.SubscribeOutbound(ctx)
|
||||
if !ok {
|
||||
t.Fatal("expected outbound approval message")
|
||||
}
|
||||
if !strings.Contains(outMsg.Content, "Approval Required") {
|
||||
t.Errorf("approval message should contain 'Approval Required', got: %s", outMsg.Content)
|
||||
}
|
||||
|
||||
// Send approval reply
|
||||
time.Sleep(50 * time.Millisecond) // small delay to ensure interceptor is registered
|
||||
msgBus.PublishInbound(bus.InboundMessage{
|
||||
Channel: "telegram",
|
||||
ChatID: "chat123",
|
||||
Content: "approve",
|
||||
})
|
||||
|
||||
select {
|
||||
case err := <-errCh:
|
||||
if err != nil {
|
||||
t.Errorf("expected approval to succeed, got: %v", err)
|
||||
}
|
||||
case <-time.After(3 * time.Second):
|
||||
t.Fatal("approval timed out")
|
||||
}
|
||||
}
|
||||
|
||||
func TestPolicyEngine_Evaluate_Approve_Denied(t *testing.T) {
|
||||
msgBus := bus.NewMessageBus()
|
||||
pe := NewPolicyEngine(&config.SecurityConfig{ApprovalTimeout: 5}, msgBus)
|
||||
|
||||
errCh := make(chan error, 1)
|
||||
go func() {
|
||||
errCh <- pe.Evaluate(context.Background(), ModeApprove, Violation{
|
||||
Category: "ssrf",
|
||||
Reason: "private IP",
|
||||
}, "feishu", "chat456")
|
||||
}()
|
||||
|
||||
// Drain outbound message
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
|
||||
defer cancel()
|
||||
msgBus.SubscribeOutbound(ctx)
|
||||
|
||||
time.Sleep(50 * time.Millisecond)
|
||||
msgBus.PublishInbound(bus.InboundMessage{
|
||||
Channel: "feishu",
|
||||
ChatID: "chat456",
|
||||
Content: "deny",
|
||||
})
|
||||
|
||||
select {
|
||||
case err := <-errCh:
|
||||
if err == nil {
|
||||
t.Fatal("expected denial error")
|
||||
}
|
||||
if !strings.Contains(err.Error(), "denied") {
|
||||
t.Errorf("error should contain 'denied', got: %v", err)
|
||||
}
|
||||
case <-time.After(3 * time.Second):
|
||||
t.Fatal("timed out waiting for denial")
|
||||
}
|
||||
}
|
||||
|
||||
func TestPolicyEngine_Evaluate_Approve_Timeout(t *testing.T) {
|
||||
msgBus := bus.NewMessageBus()
|
||||
pe := NewPolicyEngine(&config.SecurityConfig{ApprovalTimeout: 1}, msgBus)
|
||||
|
||||
errCh := make(chan error, 1)
|
||||
go func() {
|
||||
errCh <- pe.Evaluate(context.Background(), ModeApprove, Violation{
|
||||
Category: "exec_guard",
|
||||
Reason: "test",
|
||||
}, "telegram", "chat789")
|
||||
}()
|
||||
|
||||
// Drain outbound
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
|
||||
defer cancel()
|
||||
msgBus.SubscribeOutbound(ctx)
|
||||
|
||||
// Don't reply - let it timeout
|
||||
select {
|
||||
case err := <-errCh:
|
||||
if err == nil {
|
||||
t.Fatal("expected timeout error")
|
||||
}
|
||||
if !strings.Contains(err.Error(), "timed out") {
|
||||
t.Errorf("error should contain 'timed out', got: %v", err)
|
||||
}
|
||||
case <-time.After(5 * time.Second):
|
||||
t.Fatal("test itself timed out")
|
||||
}
|
||||
}
|
||||
|
||||
func TestPolicyEngine_Evaluate_Approve_NonMatchingMessages(t *testing.T) {
|
||||
msgBus := bus.NewMessageBus()
|
||||
pe := NewPolicyEngine(&config.SecurityConfig{ApprovalTimeout: 2}, msgBus)
|
||||
|
||||
errCh := make(chan error, 1)
|
||||
go func() {
|
||||
errCh <- pe.Evaluate(context.Background(), ModeApprove, Violation{
|
||||
Category: "exec_guard",
|
||||
Reason: "test",
|
||||
}, "telegram", "chat100")
|
||||
}()
|
||||
|
||||
// Drain outbound
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
|
||||
defer cancel()
|
||||
msgBus.SubscribeOutbound(ctx)
|
||||
|
||||
time.Sleep(50 * time.Millisecond)
|
||||
|
||||
// Send non-matching messages (different chat, random text)
|
||||
msgBus.PublishInbound(bus.InboundMessage{Channel: "telegram", ChatID: "other-chat", Content: "approve"})
|
||||
msgBus.PublishInbound(bus.InboundMessage{Channel: "telegram", ChatID: "chat100", Content: "hello"})
|
||||
|
||||
// Now send actual approval
|
||||
time.Sleep(50 * time.Millisecond)
|
||||
msgBus.PublishInbound(bus.InboundMessage{Channel: "telegram", ChatID: "chat100", Content: "approve"})
|
||||
|
||||
select {
|
||||
case err := <-errCh:
|
||||
if err != nil {
|
||||
t.Errorf("expected approval after non-matching messages, got: %v", err)
|
||||
}
|
||||
case <-time.After(3 * time.Second):
|
||||
t.Fatal("timed out")
|
||||
}
|
||||
}
|
||||
|
|
@ -11,6 +11,8 @@ import (
|
|||
"regexp"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/sipeed/picoclaw/pkg/security"
|
||||
)
|
||||
|
||||
// repoNamePattern validates GitHub repository format: "owner/repo"
|
||||
|
|
@ -18,6 +20,8 @@ var repoNamePattern = regexp.MustCompile(`^[a-zA-Z0-9_.-]+/[a-zA-Z0-9_.-]+$`)
|
|||
|
||||
type SkillInstaller struct {
|
||||
workspace string
|
||||
policyEngine *security.PolicyEngine
|
||||
skillMode security.PolicyMode
|
||||
}
|
||||
|
||||
type AvailableSkill struct {
|
||||
|
|
@ -40,10 +44,35 @@ func NewSkillInstaller(workspace string) *SkillInstaller {
|
|||
}
|
||||
}
|
||||
|
||||
// NewSkillInstallerWithPolicy creates a SkillInstaller with security policy support.
|
||||
func NewSkillInstallerWithPolicy(workspace string, pe *security.PolicyEngine, mode security.PolicyMode) *SkillInstaller {
|
||||
return &SkillInstaller{
|
||||
workspace: workspace,
|
||||
policyEngine: pe,
|
||||
skillMode: mode,
|
||||
}
|
||||
}
|
||||
|
||||
func (si *SkillInstaller) InstallFromGitHub(ctx context.Context, repo string) error {
|
||||
// Validate repo format to prevent URL injection
|
||||
// Validate repo format to prevent URL injection (mode-aware)
|
||||
if !si.skillMode.IsOff() {
|
||||
if !repoNamePattern.MatchString(repo) {
|
||||
return fmt.Errorf("invalid repository format: must be 'owner/repo' (got %q)", repo)
|
||||
reason := fmt.Sprintf("invalid repository format: must be 'owner/repo' (got %q)", repo)
|
||||
if si.policyEngine != nil {
|
||||
err := si.policyEngine.Evaluate(ctx, si.skillMode, security.Violation{
|
||||
Category: "skill_validation",
|
||||
Tool: "skill_install",
|
||||
Action: repo,
|
||||
Reason: reason,
|
||||
}, "", "")
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
// approved by user, continue
|
||||
} else {
|
||||
return fmt.Errorf("%s", reason)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
skillDir := filepath.Join(si.workspace, "skills", filepath.Base(repo))
|
||||
|
|
|
|||
|
|
@ -30,11 +30,16 @@ type CronTool struct {
|
|||
// NewCronTool creates a new CronTool.
|
||||
// The restrict parameter controls whether cron commands are restricted to the workspace.
|
||||
func NewCronTool(cronService *cron.CronService, executor JobExecutor, msgBus *bus.MessageBus, workspace string, restrict bool) *CronTool {
|
||||
return NewCronToolWithConfig(cronService, executor, msgBus, workspace, restrict, ExecToolConfig{})
|
||||
}
|
||||
|
||||
// NewCronToolWithConfig creates a CronTool with explicit ExecToolConfig (including PolicyEngine).
|
||||
func NewCronToolWithConfig(cronService *cron.CronService, executor JobExecutor, msgBus *bus.MessageBus, workspace string, restrict bool, execCfg ExecToolConfig) *CronTool {
|
||||
return &CronTool{
|
||||
cronService: cronService,
|
||||
executor: executor,
|
||||
msgBus: msgBus,
|
||||
execTool: NewExecTool(workspace, restrict),
|
||||
execTool: NewExecToolWithConfig(workspace, restrict, execCfg),
|
||||
}
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -5,6 +5,8 @@ import (
|
|||
"fmt"
|
||||
"os"
|
||||
"strings"
|
||||
|
||||
"github.com/sipeed/picoclaw/pkg/security"
|
||||
)
|
||||
|
||||
// EditFileTool edits a file by replacing old_text with new_text.
|
||||
|
|
@ -12,6 +14,10 @@ import (
|
|||
type EditFileTool struct {
|
||||
allowedDir string
|
||||
restrict bool
|
||||
pathMode security.PolicyMode
|
||||
policyEngine *security.PolicyEngine
|
||||
channel string
|
||||
chatID string
|
||||
}
|
||||
|
||||
// NewEditFileTool creates a new EditFileTool with optional directory restriction.
|
||||
|
|
@ -22,6 +28,15 @@ func NewEditFileTool(allowedDir string, restrict bool) *EditFileTool {
|
|||
}
|
||||
}
|
||||
|
||||
func NewEditFileToolWithPolicy(allowedDir string, restrict bool, opts PathPolicyOpts) *EditFileTool {
|
||||
return &EditFileTool{allowedDir: allowedDir, restrict: restrict, pathMode: opts.PathMode, policyEngine: opts.PolicyEngine}
|
||||
}
|
||||
|
||||
func (t *EditFileTool) SetContext(channel, chatID string) {
|
||||
t.channel = channel
|
||||
t.chatID = chatID
|
||||
}
|
||||
|
||||
func (t *EditFileTool) Name() string {
|
||||
return "edit_file"
|
||||
}
|
||||
|
|
@ -67,7 +82,7 @@ func (t *EditFileTool) Execute(ctx context.Context, args map[string]interface{})
|
|||
return ErrorResult("new_text is required")
|
||||
}
|
||||
|
||||
resolvedPath, err := validatePath(path, t.allowedDir, t.restrict)
|
||||
resolvedPath, err := validatePathWithMode(path, t.allowedDir, t.restrict, t.pathMode, t.policyEngine, t.channel, t.chatID)
|
||||
if err != nil {
|
||||
return ErrorResult(err.Error())
|
||||
}
|
||||
|
|
@ -111,12 +126,25 @@ func (t *EditFileTool) Execute(ctx context.Context, args map[string]interface{})
|
|||
type AppendFileTool struct {
|
||||
workspace string
|
||||
restrict bool
|
||||
pathMode security.PolicyMode
|
||||
policyEngine *security.PolicyEngine
|
||||
channel string
|
||||
chatID string
|
||||
}
|
||||
|
||||
func NewAppendFileTool(workspace string, restrict bool) *AppendFileTool {
|
||||
return &AppendFileTool{workspace: workspace, restrict: restrict}
|
||||
}
|
||||
|
||||
func NewAppendFileToolWithPolicy(workspace string, restrict bool, opts PathPolicyOpts) *AppendFileTool {
|
||||
return &AppendFileTool{workspace: workspace, restrict: restrict, pathMode: opts.PathMode, policyEngine: opts.PolicyEngine}
|
||||
}
|
||||
|
||||
func (t *AppendFileTool) SetContext(channel, chatID string) {
|
||||
t.channel = channel
|
||||
t.chatID = chatID
|
||||
}
|
||||
|
||||
func (t *AppendFileTool) Name() string {
|
||||
return "append_file"
|
||||
}
|
||||
|
|
@ -153,7 +181,7 @@ func (t *AppendFileTool) Execute(ctx context.Context, args map[string]interface{
|
|||
return ErrorResult("content is required")
|
||||
}
|
||||
|
||||
resolvedPath, err := validatePath(path, t.workspace, t.restrict)
|
||||
resolvedPath, err := validatePathWithMode(path, t.workspace, t.restrict, t.pathMode, t.policyEngine, t.channel, t.chatID)
|
||||
if err != nil {
|
||||
return ErrorResult(err.Error())
|
||||
}
|
||||
|
|
|
|||
|
|
@ -6,6 +6,8 @@ import (
|
|||
"path/filepath"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"github.com/sipeed/picoclaw/pkg/security"
|
||||
)
|
||||
|
||||
// TestEditTool_EditFile_Success verifies successful file editing
|
||||
|
|
@ -287,3 +289,35 @@ func TestEditTool_AppendFile_MissingContent(t *testing.T) {
|
|||
t.Errorf("Expected error when content is missing")
|
||||
}
|
||||
}
|
||||
|
||||
func TestEditFileTool_SetContext(t *testing.T) {
|
||||
tool := NewEditFileToolWithPolicy("", false, PathPolicyOpts{PathMode: security.ModeBlock})
|
||||
tool.SetContext("slack", "C999")
|
||||
if tool.channel != "slack" || tool.chatID != "C999" {
|
||||
t.Errorf("SetContext failed: channel=%q, chatID=%q", tool.channel, tool.chatID)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAppendFileTool_SetContext(t *testing.T) {
|
||||
tool := NewAppendFileToolWithPolicy("", false, PathPolicyOpts{PathMode: security.ModeApprove})
|
||||
tool.SetContext("telegram", "chat-42")
|
||||
if tool.channel != "telegram" || tool.chatID != "chat-42" {
|
||||
t.Errorf("SetContext failed: channel=%q, chatID=%q", tool.channel, tool.chatID)
|
||||
}
|
||||
}
|
||||
|
||||
func TestNewEditFileToolWithPolicy(t *testing.T) {
|
||||
opts := PathPolicyOpts{PathMode: security.ModeBlock}
|
||||
tool := NewEditFileToolWithPolicy("/dir", true, opts)
|
||||
if tool.allowedDir != "/dir" || !tool.restrict || tool.pathMode != security.ModeBlock {
|
||||
t.Error("WithPolicy constructor did not set fields correctly")
|
||||
}
|
||||
}
|
||||
|
||||
func TestNewAppendFileToolWithPolicy(t *testing.T) {
|
||||
opts := PathPolicyOpts{PathMode: security.ModeApprove}
|
||||
tool := NewAppendFileToolWithPolicy("/ws", true, opts)
|
||||
if tool.workspace != "/ws" || !tool.restrict || tool.pathMode != security.ModeApprove {
|
||||
t.Error("WithPolicy constructor did not set fields correctly")
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -6,12 +6,19 @@ import (
|
|||
"os"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
|
||||
"github.com/sipeed/picoclaw/pkg/security"
|
||||
)
|
||||
|
||||
// validatePath ensures the given path is within the workspace if restrict is true.
|
||||
// It resolves symlinks to prevent symlink-based path traversal attacks and uses
|
||||
// a trailing separator comparison to avoid prefix collisions (e.g. /workspace vs /workspace2).
|
||||
// When pathMode is "off", only basic prefix check is performed (no symlink resolution).
|
||||
// When pathMode is "block" or "approve", enhanced symlink resolution is used.
|
||||
func validatePath(path, workspace string, restrict bool) (string, error) {
|
||||
return validatePathWithMode(path, workspace, restrict, security.ModeOff, nil, "", "")
|
||||
}
|
||||
|
||||
// validatePathWithMode is the full-featured path validator with policy support.
|
||||
func validatePathWithMode(path, workspace string, restrict bool, pathMode security.PolicyMode, pe *security.PolicyEngine, channel, chatID string) (string, error) {
|
||||
if workspace == "" {
|
||||
return path, nil
|
||||
}
|
||||
|
|
@ -32,28 +39,44 @@ func validatePath(path, workspace string, restrict bool) (string, error) {
|
|||
}
|
||||
|
||||
if restrict {
|
||||
// Resolve symlinks for the workspace path
|
||||
useSymlinkResolution := !pathMode.IsOff()
|
||||
|
||||
realWorkspace := absWorkspace
|
||||
if useSymlinkResolution {
|
||||
if resolved, err := filepath.EvalSymlinks(absWorkspace); err == nil {
|
||||
realWorkspace = resolved
|
||||
}
|
||||
// Ensure workspace path ends with separator for safe prefix comparison
|
||||
}
|
||||
workspacePrefix := filepath.Clean(realWorkspace) + string(filepath.Separator)
|
||||
|
||||
// Resolve symlinks for the target path.
|
||||
// For non-existent files, resolve the parent directory's symlinks and append the filename.
|
||||
realPath := absPath
|
||||
if useSymlinkResolution {
|
||||
if resolved, err := filepath.EvalSymlinks(absPath); err == nil {
|
||||
realPath = resolved
|
||||
} else if resolved, err := filepath.EvalSymlinks(filepath.Dir(absPath)); err == nil {
|
||||
realPath = filepath.Join(resolved, filepath.Base(absPath))
|
||||
}
|
||||
}
|
||||
|
||||
// Check: realPath must be the workspace itself or inside it
|
||||
if realPath != filepath.Clean(realWorkspace) &&
|
||||
!strings.HasPrefix(realPath+string(filepath.Separator), workspacePrefix) &&
|
||||
!strings.HasPrefix(realPath, workspacePrefix) {
|
||||
return "", fmt.Errorf("access denied: path is outside the workspace")
|
||||
violation := fmt.Errorf("access denied: path is outside the workspace")
|
||||
if pe != nil && pathMode == security.ModeApprove {
|
||||
ctx := context.Background()
|
||||
pErr := pe.Evaluate(ctx, pathMode, security.Violation{
|
||||
Category: "path_validation",
|
||||
Tool: "filesystem",
|
||||
Action: path,
|
||||
Reason: violation.Error(),
|
||||
}, channel, chatID)
|
||||
if pErr != nil {
|
||||
return "", pErr
|
||||
}
|
||||
// approved
|
||||
} else {
|
||||
return "", violation
|
||||
}
|
||||
}
|
||||
|
||||
absPath = realPath
|
||||
|
|
@ -62,15 +85,34 @@ func validatePath(path, workspace string, restrict bool) (string, error) {
|
|||
return absPath, nil
|
||||
}
|
||||
|
||||
// PathPolicyOpts holds optional security policy settings for filesystem tools.
|
||||
type PathPolicyOpts struct {
|
||||
PathMode security.PolicyMode
|
||||
PolicyEngine *security.PolicyEngine
|
||||
}
|
||||
|
||||
type ReadFileTool struct {
|
||||
workspace string
|
||||
restrict bool
|
||||
pathMode security.PolicyMode
|
||||
policyEngine *security.PolicyEngine
|
||||
channel string
|
||||
chatID string
|
||||
}
|
||||
|
||||
func NewReadFileTool(workspace string, restrict bool) *ReadFileTool {
|
||||
return &ReadFileTool{workspace: workspace, restrict: restrict}
|
||||
}
|
||||
|
||||
func NewReadFileToolWithPolicy(workspace string, restrict bool, opts PathPolicyOpts) *ReadFileTool {
|
||||
return &ReadFileTool{workspace: workspace, restrict: restrict, pathMode: opts.PathMode, policyEngine: opts.PolicyEngine}
|
||||
}
|
||||
|
||||
func (t *ReadFileTool) SetContext(channel, chatID string) {
|
||||
t.channel = channel
|
||||
t.chatID = chatID
|
||||
}
|
||||
|
||||
func (t *ReadFileTool) Name() string {
|
||||
return "read_file"
|
||||
}
|
||||
|
|
@ -98,7 +140,7 @@ func (t *ReadFileTool) Execute(ctx context.Context, args map[string]interface{})
|
|||
return ErrorResult("path is required")
|
||||
}
|
||||
|
||||
resolvedPath, err := validatePath(path, t.workspace, t.restrict)
|
||||
resolvedPath, err := validatePathWithMode(path, t.workspace, t.restrict, t.pathMode, t.policyEngine, t.channel, t.chatID)
|
||||
if err != nil {
|
||||
return ErrorResult(err.Error())
|
||||
}
|
||||
|
|
@ -114,12 +156,25 @@ func (t *ReadFileTool) Execute(ctx context.Context, args map[string]interface{})
|
|||
type WriteFileTool struct {
|
||||
workspace string
|
||||
restrict bool
|
||||
pathMode security.PolicyMode
|
||||
policyEngine *security.PolicyEngine
|
||||
channel string
|
||||
chatID string
|
||||
}
|
||||
|
||||
func NewWriteFileTool(workspace string, restrict bool) *WriteFileTool {
|
||||
return &WriteFileTool{workspace: workspace, restrict: restrict}
|
||||
}
|
||||
|
||||
func NewWriteFileToolWithPolicy(workspace string, restrict bool, opts PathPolicyOpts) *WriteFileTool {
|
||||
return &WriteFileTool{workspace: workspace, restrict: restrict, pathMode: opts.PathMode, policyEngine: opts.PolicyEngine}
|
||||
}
|
||||
|
||||
func (t *WriteFileTool) SetContext(channel, chatID string) {
|
||||
t.channel = channel
|
||||
t.chatID = chatID
|
||||
}
|
||||
|
||||
func (t *WriteFileTool) Name() string {
|
||||
return "write_file"
|
||||
}
|
||||
|
|
@ -156,7 +211,7 @@ func (t *WriteFileTool) Execute(ctx context.Context, args map[string]interface{}
|
|||
return ErrorResult("content is required")
|
||||
}
|
||||
|
||||
resolvedPath, err := validatePath(path, t.workspace, t.restrict)
|
||||
resolvedPath, err := validatePathWithMode(path, t.workspace, t.restrict, t.pathMode, t.policyEngine, t.channel, t.chatID)
|
||||
if err != nil {
|
||||
return ErrorResult(err.Error())
|
||||
}
|
||||
|
|
@ -176,12 +231,25 @@ func (t *WriteFileTool) Execute(ctx context.Context, args map[string]interface{}
|
|||
type ListDirTool struct {
|
||||
workspace string
|
||||
restrict bool
|
||||
pathMode security.PolicyMode
|
||||
policyEngine *security.PolicyEngine
|
||||
channel string
|
||||
chatID string
|
||||
}
|
||||
|
||||
func NewListDirTool(workspace string, restrict bool) *ListDirTool {
|
||||
return &ListDirTool{workspace: workspace, restrict: restrict}
|
||||
}
|
||||
|
||||
func NewListDirToolWithPolicy(workspace string, restrict bool, opts PathPolicyOpts) *ListDirTool {
|
||||
return &ListDirTool{workspace: workspace, restrict: restrict, pathMode: opts.PathMode, policyEngine: opts.PolicyEngine}
|
||||
}
|
||||
|
||||
func (t *ListDirTool) SetContext(channel, chatID string) {
|
||||
t.channel = channel
|
||||
t.chatID = chatID
|
||||
}
|
||||
|
||||
func (t *ListDirTool) Name() string {
|
||||
return "list_dir"
|
||||
}
|
||||
|
|
@ -209,7 +277,7 @@ func (t *ListDirTool) Execute(ctx context.Context, args map[string]interface{})
|
|||
path = "."
|
||||
}
|
||||
|
||||
resolvedPath, err := validatePath(path, t.workspace, t.restrict)
|
||||
resolvedPath, err := validatePathWithMode(path, t.workspace, t.restrict, t.pathMode, t.policyEngine, t.channel, t.chatID)
|
||||
if err != nil {
|
||||
return ErrorResult(err.Error())
|
||||
}
|
||||
|
|
|
|||
|
|
@ -6,6 +6,8 @@ import (
|
|||
"path/filepath"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"github.com/sipeed/picoclaw/pkg/security"
|
||||
)
|
||||
|
||||
// TestFilesystemTool_ReadFile_Success verifies successful file reading
|
||||
|
|
@ -249,6 +251,7 @@ func TestFilesystemTool_ListDir_DefaultPath(t *testing.T) {
|
|||
}
|
||||
|
||||
// TestValidatePath_SymlinkEscape verifies that symlinks pointing outside workspace are blocked
|
||||
// when path validation mode is "block" (enhanced symlink resolution).
|
||||
func TestValidatePath_SymlinkEscape(t *testing.T) {
|
||||
// Create workspace and an outside directory
|
||||
workspace := t.TempDir()
|
||||
|
|
@ -262,8 +265,8 @@ func TestValidatePath_SymlinkEscape(t *testing.T) {
|
|||
t.Skipf("Cannot create symlink: %v", err)
|
||||
}
|
||||
|
||||
// Try to access the symlinked file with restrict=true
|
||||
_, err := validatePath("escape/secret.txt", workspace, true)
|
||||
// Try to access the symlinked file with restrict=true and path validation enabled
|
||||
_, err := validatePathWithMode("escape/secret.txt", workspace, true, security.ModeBlock, nil, "", "")
|
||||
if err == nil {
|
||||
t.Error("Expected symlink escape to be blocked, but it was allowed")
|
||||
}
|
||||
|
|
@ -288,14 +291,86 @@ func TestValidatePath_PrefixCollision(t *testing.T) {
|
|||
func TestValidatePath_AllowsWorkspaceItself(t *testing.T) {
|
||||
workspace := t.TempDir()
|
||||
|
||||
path, err := validatePath(".", workspace, true)
|
||||
// With block mode (enhanced symlink resolution), path should be resolved
|
||||
path, err := validatePathWithMode(".", workspace, true, security.ModeBlock, nil, "", "")
|
||||
if err != nil {
|
||||
t.Errorf("Expected workspace root access to be allowed, got error: %v", err)
|
||||
}
|
||||
|
||||
// Resolve both to compare
|
||||
expectedPath, _ := filepath.EvalSymlinks(workspace)
|
||||
if path != expectedPath {
|
||||
t.Errorf("Expected path to be %s, got %s", expectedPath, path)
|
||||
}
|
||||
}
|
||||
|
||||
func TestValidatePath_EmptyWorkspace(t *testing.T) {
|
||||
path, err := validatePath("/some/path", "", true)
|
||||
if err != nil {
|
||||
t.Errorf("Expected no error with empty workspace, got: %v", err)
|
||||
}
|
||||
if path != "/some/path" {
|
||||
t.Errorf("Expected path returned as-is, got: %s", path)
|
||||
}
|
||||
}
|
||||
|
||||
func TestValidatePath_ModeOff_NoSymlinkResolution(t *testing.T) {
|
||||
workspace := t.TempDir()
|
||||
testFile := filepath.Join(workspace, "file.txt")
|
||||
os.WriteFile(testFile, []byte("data"), 0644)
|
||||
|
||||
path, err := validatePathWithMode("file.txt", workspace, true, security.ModeOff, nil, "", "")
|
||||
if err != nil {
|
||||
t.Errorf("Expected success, got: %v", err)
|
||||
}
|
||||
if path == "" {
|
||||
t.Error("Expected non-empty path")
|
||||
}
|
||||
}
|
||||
|
||||
func TestReadFileTool_SetContext(t *testing.T) {
|
||||
tool := NewReadFileToolWithPolicy("", false, PathPolicyOpts{})
|
||||
tool.SetContext("telegram", "chat-1")
|
||||
if tool.channel != "telegram" || tool.chatID != "chat-1" {
|
||||
t.Errorf("SetContext failed: channel=%q, chatID=%q", tool.channel, tool.chatID)
|
||||
}
|
||||
}
|
||||
|
||||
func TestWriteFileTool_SetContext(t *testing.T) {
|
||||
tool := NewWriteFileToolWithPolicy("", false, PathPolicyOpts{})
|
||||
tool.SetContext("feishu", "chat-2")
|
||||
if tool.channel != "feishu" || tool.chatID != "chat-2" {
|
||||
t.Errorf("SetContext failed: channel=%q, chatID=%q", tool.channel, tool.chatID)
|
||||
}
|
||||
}
|
||||
|
||||
func TestListDirTool_SetContext(t *testing.T) {
|
||||
tool := NewListDirToolWithPolicy("", false, PathPolicyOpts{})
|
||||
tool.SetContext("dingtalk", "chat-3")
|
||||
if tool.channel != "dingtalk" || tool.chatID != "chat-3" {
|
||||
t.Errorf("SetContext failed: channel=%q, chatID=%q", tool.channel, tool.chatID)
|
||||
}
|
||||
}
|
||||
|
||||
func TestNewReadFileToolWithPolicy(t *testing.T) {
|
||||
opts := PathPolicyOpts{PathMode: security.ModeBlock}
|
||||
tool := NewReadFileToolWithPolicy("/workspace", true, opts)
|
||||
if tool.workspace != "/workspace" || !tool.restrict || tool.pathMode != security.ModeBlock {
|
||||
t.Error("WithPolicy constructor did not set fields correctly")
|
||||
}
|
||||
}
|
||||
|
||||
func TestNewWriteFileToolWithPolicy(t *testing.T) {
|
||||
opts := PathPolicyOpts{PathMode: security.ModeApprove}
|
||||
tool := NewWriteFileToolWithPolicy("/ws", true, opts)
|
||||
if tool.workspace != "/ws" || !tool.restrict || tool.pathMode != security.ModeApprove {
|
||||
t.Error("WithPolicy constructor did not set fields correctly")
|
||||
}
|
||||
}
|
||||
|
||||
func TestNewListDirToolWithPolicy(t *testing.T) {
|
||||
opts := PathPolicyOpts{PathMode: security.ModeBlock}
|
||||
tool := NewListDirToolWithPolicy("/ws", false, opts)
|
||||
if tool.workspace != "/ws" || tool.restrict || tool.pathMode != security.ModeBlock {
|
||||
t.Error("WithPolicy constructor did not set fields correctly")
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -11,6 +11,8 @@ import (
|
|||
"runtime"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/sipeed/picoclaw/pkg/security"
|
||||
)
|
||||
|
||||
// ExecToolConfig holds configurable options for ExecTool.
|
||||
|
|
@ -18,6 +20,8 @@ type ExecToolConfig struct {
|
|||
DenyPatterns []string // Additional regex deny patterns from config
|
||||
AllowPatterns []string // If set, only matching commands are allowed
|
||||
MaxTimeout int // Seconds, default 60
|
||||
PolicyEngine *security.PolicyEngine
|
||||
ExecGuardMode security.PolicyMode
|
||||
}
|
||||
|
||||
type ExecTool struct {
|
||||
|
|
@ -26,6 +30,10 @@ type ExecTool struct {
|
|||
denyPatterns []*regexp.Regexp
|
||||
allowPatterns []*regexp.Regexp
|
||||
restrictToWorkspace bool
|
||||
policyEngine *security.PolicyEngine
|
||||
execGuardMode security.PolicyMode
|
||||
channel string
|
||||
chatID string
|
||||
}
|
||||
|
||||
func NewExecTool(workingDir string, restrict bool) *ExecTool {
|
||||
|
|
@ -86,9 +94,18 @@ func NewExecToolWithConfig(workingDir string, restrict bool, cfg ExecToolConfig)
|
|||
denyPatterns: denyPatterns,
|
||||
allowPatterns: allowPatterns,
|
||||
restrictToWorkspace: restrict,
|
||||
policyEngine: cfg.PolicyEngine,
|
||||
execGuardMode: cfg.ExecGuardMode,
|
||||
}
|
||||
}
|
||||
|
||||
// SetContext implements ContextualTool so the ExecTool receives the current
|
||||
// IM channel and chatID for approval requests.
|
||||
func (t *ExecTool) SetContext(channel, chatID string) {
|
||||
t.channel = channel
|
||||
t.chatID = chatID
|
||||
}
|
||||
|
||||
func (t *ExecTool) Name() string {
|
||||
return "exec"
|
||||
}
|
||||
|
|
@ -132,7 +149,7 @@ func (t *ExecTool) Execute(ctx context.Context, args map[string]interface{}) *To
|
|||
}
|
||||
}
|
||||
|
||||
if guardError := t.guardCommand(command, cwd); guardError != "" {
|
||||
if guardError := t.guardCommand(ctx, command, cwd); guardError != "" {
|
||||
return ErrorResult(guardError)
|
||||
}
|
||||
|
||||
|
|
@ -195,16 +212,24 @@ func (t *ExecTool) Execute(ctx context.Context, args map[string]interface{}) *To
|
|||
}
|
||||
}
|
||||
|
||||
func (t *ExecTool) guardCommand(command, cwd string) string {
|
||||
func (t *ExecTool) guardCommand(ctx context.Context, command, cwd string) string {
|
||||
mode := t.execGuardMode
|
||||
cmd := strings.TrimSpace(command)
|
||||
lower := strings.ToLower(cmd)
|
||||
|
||||
// Deny-pattern check (mode-aware)
|
||||
if !mode.IsOff() {
|
||||
for _, pattern := range t.denyPatterns {
|
||||
if pattern.MatchString(lower) {
|
||||
return "Command blocked by safety guard (dangerous pattern detected)"
|
||||
reason := "dangerous pattern detected: " + pattern.String()
|
||||
if err := t.evaluatePolicy(ctx, mode, command, reason, pattern.String()); err != nil {
|
||||
return err.Error()
|
||||
}
|
||||
break // approved by user, continue
|
||||
}
|
||||
}
|
||||
|
||||
// Allow-pattern check
|
||||
if len(t.allowPatterns) > 0 {
|
||||
allowed := false
|
||||
for _, pattern := range t.allowPatterns {
|
||||
|
|
@ -214,16 +239,20 @@ func (t *ExecTool) guardCommand(command, cwd string) string {
|
|||
}
|
||||
}
|
||||
if !allowed {
|
||||
return "Command blocked by safety guard (not in allowlist)"
|
||||
reason := "command not in allowlist"
|
||||
if err := t.evaluatePolicy(ctx, mode, command, reason, "allowlist"); err != nil {
|
||||
return err.Error()
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Workspace restriction checks (always active when restrictToWorkspace is true)
|
||||
if t.restrictToWorkspace {
|
||||
if strings.Contains(cmd, "..\\") || strings.Contains(cmd, "../") {
|
||||
return "Command blocked by safety guard (path traversal detected)"
|
||||
}
|
||||
|
||||
// Block access to sensitive system paths
|
||||
sensitivePathPatterns := []*regexp.Regexp{
|
||||
regexp.MustCompile(`\b/etc/`),
|
||||
regexp.MustCompile(`\b/var/`),
|
||||
|
|
@ -267,6 +296,23 @@ func (t *ExecTool) guardCommand(command, cwd string) string {
|
|||
return ""
|
||||
}
|
||||
|
||||
// evaluatePolicy delegates to the PolicyEngine when available.
|
||||
func (t *ExecTool) evaluatePolicy(ctx context.Context, mode security.PolicyMode, action, reason, ruleName string) error {
|
||||
if t.policyEngine == nil {
|
||||
if mode == security.ModeOff {
|
||||
return nil
|
||||
}
|
||||
return fmt.Errorf("blocked by safety guard: %s", reason)
|
||||
}
|
||||
return t.policyEngine.Evaluate(ctx, mode, security.Violation{
|
||||
Category: "exec_guard",
|
||||
Tool: "exec",
|
||||
Action: action,
|
||||
Reason: reason,
|
||||
RuleName: ruleName,
|
||||
}, t.channel, t.chatID)
|
||||
}
|
||||
|
||||
func (t *ExecTool) SetTimeout(timeout time.Duration) {
|
||||
t.timeout = timeout
|
||||
}
|
||||
|
|
|
|||
|
|
@ -7,6 +7,8 @@ import (
|
|||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/sipeed/picoclaw/pkg/security"
|
||||
)
|
||||
|
||||
// TestShellTool_Success verifies successful command execution
|
||||
|
|
@ -113,8 +115,9 @@ func TestShellTool_WorkingDir(t *testing.T) {
|
|||
}
|
||||
|
||||
// TestShellTool_DangerousCommand verifies safety guard blocks dangerous commands
|
||||
// when exec_guard mode is "block".
|
||||
func TestShellTool_DangerousCommand(t *testing.T) {
|
||||
tool := NewExecTool("", false)
|
||||
tool := NewExecToolWithConfig("", false, ExecToolConfig{ExecGuardMode: "block"})
|
||||
|
||||
ctx := context.Background()
|
||||
args := map[string]interface{}{
|
||||
|
|
@ -210,8 +213,9 @@ func TestShellTool_RestrictToWorkspace(t *testing.T) {
|
|||
}
|
||||
|
||||
// TestShellTool_DataExfiltrationBlocked verifies data exfiltration patterns are blocked
|
||||
// when exec_guard mode is "block".
|
||||
func TestShellTool_DataExfiltrationBlocked(t *testing.T) {
|
||||
tool := NewExecTool("", false)
|
||||
tool := NewExecToolWithConfig("", false, ExecToolConfig{ExecGuardMode: "block"})
|
||||
ctx := context.Background()
|
||||
|
||||
dangerousCmds := []string{
|
||||
|
|
@ -257,6 +261,7 @@ func TestShellTool_WithConfig(t *testing.T) {
|
|||
cfg := ExecToolConfig{
|
||||
DenyPatterns: []string{`\bmy_custom_blocked\b`},
|
||||
MaxTimeout: 30,
|
||||
ExecGuardMode: "block",
|
||||
}
|
||||
tool := NewExecToolWithConfig("", false, cfg)
|
||||
|
||||
|
|
@ -271,3 +276,94 @@ func TestShellTool_WithConfig(t *testing.T) {
|
|||
t.Errorf("Expected timeout to be 30s, got %v", tool.timeout)
|
||||
}
|
||||
}
|
||||
|
||||
func TestExecTool_SetContext(t *testing.T) {
|
||||
tool := NewExecTool("", false)
|
||||
tool.SetContext("telegram", "chat-123")
|
||||
|
||||
if tool.channel != "telegram" {
|
||||
t.Errorf("Expected channel 'telegram', got %q", tool.channel)
|
||||
}
|
||||
if tool.chatID != "chat-123" {
|
||||
t.Errorf("Expected chatID 'chat-123', got %q", tool.chatID)
|
||||
}
|
||||
}
|
||||
|
||||
func TestExecTool_AllowPatternBlocked(t *testing.T) {
|
||||
cfg := ExecToolConfig{
|
||||
AllowPatterns: []string{`^echo\b`, `^ls\b`},
|
||||
ExecGuardMode: "block",
|
||||
}
|
||||
tool := NewExecToolWithConfig("", false, cfg)
|
||||
|
||||
ctx := context.Background()
|
||||
result := tool.Execute(ctx, map[string]interface{}{"command": "rm -rf /"})
|
||||
if !result.IsError {
|
||||
t.Error("Expected command not in allowlist to be blocked")
|
||||
}
|
||||
if !strings.Contains(result.ForLLM, "not in allowlist") && !strings.Contains(result.ForLLM, "blocked") {
|
||||
t.Errorf("Expected allowlist error, got: %s", result.ForLLM)
|
||||
}
|
||||
}
|
||||
|
||||
func TestExecTool_AllowPatternAllowed(t *testing.T) {
|
||||
cfg := ExecToolConfig{
|
||||
AllowPatterns: []string{`^echo\b`},
|
||||
ExecGuardMode: "block",
|
||||
}
|
||||
tool := NewExecToolWithConfig("", false, cfg)
|
||||
|
||||
ctx := context.Background()
|
||||
result := tool.Execute(ctx, map[string]interface{}{"command": "echo hello"})
|
||||
if result.IsError {
|
||||
t.Errorf("Expected allowed command to succeed, got error: %s", result.ForLLM)
|
||||
}
|
||||
}
|
||||
|
||||
func TestExecTool_EvaluatePolicy_NilEngine_ModeOff(t *testing.T) {
|
||||
tool := NewExecTool("", false)
|
||||
err := tool.evaluatePolicy(context.Background(), security.ModeOff, "test", "reason", "rule")
|
||||
if err != nil {
|
||||
t.Errorf("Expected nil error for ModeOff with nil engine, got: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestExecTool_EvaluatePolicy_NilEngine_ModeBlock(t *testing.T) {
|
||||
tool := NewExecTool("", false)
|
||||
err := tool.evaluatePolicy(context.Background(), security.ModeBlock, "test", "reason", "rule")
|
||||
if err == nil {
|
||||
t.Error("Expected error for ModeBlock with nil engine")
|
||||
}
|
||||
if !strings.Contains(err.Error(), "blocked by safety guard") {
|
||||
t.Errorf("Expected 'blocked by safety guard' error, got: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestExecTool_SetAllowPatterns(t *testing.T) {
|
||||
tool := NewExecTool("", false)
|
||||
err := tool.SetAllowPatterns([]string{`^git\b`, `^ls\b`})
|
||||
if err != nil {
|
||||
t.Fatalf("Expected no error, got: %v", err)
|
||||
}
|
||||
if len(tool.allowPatterns) != 2 {
|
||||
t.Errorf("Expected 2 allow patterns, got %d", len(tool.allowPatterns))
|
||||
}
|
||||
}
|
||||
|
||||
func TestExecTool_SetAllowPatterns_Invalid(t *testing.T) {
|
||||
tool := NewExecTool("", false)
|
||||
err := tool.SetAllowPatterns([]string{`[invalid`})
|
||||
if err == nil {
|
||||
t.Error("Expected error for invalid regex")
|
||||
}
|
||||
}
|
||||
|
||||
func TestExecTool_GuardOff_DangerousAllowed(t *testing.T) {
|
||||
tool := NewExecToolWithConfig("", false, ExecToolConfig{})
|
||||
|
||||
ctx := context.Background()
|
||||
msg := tool.guardCommand(ctx, "rm -rf /", "")
|
||||
if msg != "" {
|
||||
t.Errorf("Expected dangerous command to pass through when exec_guard is off, got: %s", msg)
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -11,6 +11,7 @@ import (
|
|||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/sipeed/picoclaw/pkg/security"
|
||||
"github.com/sipeed/picoclaw/pkg/utils"
|
||||
)
|
||||
|
||||
|
|
@ -269,18 +270,40 @@ func (t *WebSearchTool) Execute(ctx context.Context, args map[string]interface{}
|
|||
|
||||
type WebFetchTool struct {
|
||||
maxChars int
|
||||
skipSSRFCheck bool // for testing only
|
||||
ssrfMode security.PolicyMode
|
||||
policyEngine *security.PolicyEngine
|
||||
channel string
|
||||
chatID string
|
||||
}
|
||||
|
||||
// WebFetchToolOptions configures optional security policy for WebFetchTool.
|
||||
type WebFetchToolOptions struct {
|
||||
MaxChars int
|
||||
PolicyEngine *security.PolicyEngine
|
||||
SSRFMode security.PolicyMode
|
||||
}
|
||||
|
||||
func NewWebFetchTool(maxChars int) *WebFetchTool {
|
||||
if maxChars <= 0 {
|
||||
maxChars = 50000
|
||||
return NewWebFetchToolWithPolicy(WebFetchToolOptions{MaxChars: maxChars})
|
||||
}
|
||||
|
||||
func NewWebFetchToolWithPolicy(opts WebFetchToolOptions) *WebFetchTool {
|
||||
if opts.MaxChars <= 0 {
|
||||
opts.MaxChars = 50000
|
||||
}
|
||||
return &WebFetchTool{
|
||||
maxChars: maxChars,
|
||||
maxChars: opts.MaxChars,
|
||||
policyEngine: opts.PolicyEngine,
|
||||
ssrfMode: opts.SSRFMode,
|
||||
}
|
||||
}
|
||||
|
||||
// SetContext implements ContextualTool for IM-based approval.
|
||||
func (t *WebFetchTool) SetContext(channel, chatID string) {
|
||||
t.channel = channel
|
||||
t.chatID = chatID
|
||||
}
|
||||
|
||||
func (t *WebFetchTool) Name() string {
|
||||
return "web_fetch"
|
||||
}
|
||||
|
|
@ -313,12 +336,25 @@ func (t *WebFetchTool) Execute(ctx context.Context, args map[string]interface{})
|
|||
return ErrorResult("url is required")
|
||||
}
|
||||
|
||||
// Validate URL against SSRF attacks (blocks private IPs, localhost, metadata endpoints)
|
||||
if !t.skipSSRFCheck {
|
||||
// SSRF protection (mode-aware)
|
||||
if !t.ssrfMode.IsOff() {
|
||||
if err := utils.ValidateURL(urlStr); err != nil {
|
||||
if t.policyEngine != nil {
|
||||
pErr := t.policyEngine.Evaluate(ctx, t.ssrfMode, security.Violation{
|
||||
Category: "ssrf",
|
||||
Tool: "web_fetch",
|
||||
Action: urlStr,
|
||||
Reason: err.Error(),
|
||||
}, t.channel, t.chatID)
|
||||
if pErr != nil {
|
||||
return ErrorResult(fmt.Sprintf("URL blocked: %v", pErr))
|
||||
}
|
||||
// approved by user
|
||||
} else {
|
||||
return ErrorResult(fmt.Sprintf("URL blocked: %v", err))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
parsedURL, err := url.Parse(urlStr)
|
||||
if err != nil {
|
||||
|
|
@ -355,10 +391,11 @@ func (t *WebFetchTool) Execute(ctx context.Context, args map[string]interface{})
|
|||
if len(via) >= 5 {
|
||||
return fmt.Errorf("stopped after 5 redirects")
|
||||
}
|
||||
// Validate redirect target against SSRF
|
||||
if !t.ssrfMode.IsOff() {
|
||||
if err := utils.ValidateURL(req.URL.String()); err != nil {
|
||||
return fmt.Errorf("redirect blocked: %w", err)
|
||||
}
|
||||
}
|
||||
return nil
|
||||
},
|
||||
}
|
||||
|
|
|
|||
|
|
@ -11,9 +11,8 @@ import (
|
|||
|
||||
// newTestWebFetchTool creates a WebFetchTool with SSRF check disabled for local test servers.
|
||||
func newTestWebFetchTool(maxChars int) *WebFetchTool {
|
||||
tool := NewWebFetchTool(maxChars)
|
||||
tool.skipSSRFCheck = true
|
||||
return tool
|
||||
return NewWebFetchToolWithPolicy(WebFetchToolOptions{MaxChars: maxChars})
|
||||
// ssrfMode defaults to ModeOff (""), which disables SSRF checks
|
||||
}
|
||||
|
||||
// TestWebTool_WebFetch_Success verifies successful URL fetching
|
||||
|
|
@ -116,8 +115,8 @@ func TestWebTool_WebFetch_UnsupportedScheme(t *testing.T) {
|
|||
t.Errorf("Expected error for unsupported URL scheme")
|
||||
}
|
||||
|
||||
// Should mention only http/https allowed
|
||||
if !strings.Contains(result.ForLLM, "http/https") && !strings.Contains(result.ForUser, "http/https") {
|
||||
// Should mention scheme issue (http/https or unsupported protocol scheme)
|
||||
if !strings.Contains(result.ForLLM, "http/https") && !strings.Contains(result.ForLLM, "unsupported protocol scheme") {
|
||||
t.Errorf("Expected scheme error message, got ForLLM: %s", result.ForLLM)
|
||||
}
|
||||
}
|
||||
|
|
@ -239,7 +238,7 @@ func TestWebTool_WebFetch_HTMLExtraction(t *testing.T) {
|
|||
|
||||
// TestWebTool_WebFetch_SSRFBlocked verifies that SSRF attempts are blocked
|
||||
func TestWebTool_WebFetch_SSRFBlocked(t *testing.T) {
|
||||
tool := NewWebFetchTool(50000) // Note: NOT using newTestWebFetchTool, SSRF check is active
|
||||
tool := NewWebFetchToolWithPolicy(WebFetchToolOptions{MaxChars: 50000, SSRFMode: "block"})
|
||||
ctx := context.Background()
|
||||
|
||||
ssrfURLs := []string{
|
||||
|
|
@ -282,3 +281,38 @@ func TestWebTool_WebFetch_MissingDomain(t *testing.T) {
|
|||
t.Errorf("Expected domain/host error message, got ForLLM: %s", result.ForLLM)
|
||||
}
|
||||
}
|
||||
|
||||
func TestWebFetchTool_SetContext(t *testing.T) {
|
||||
tool := NewWebFetchTool(50000)
|
||||
tool.SetContext("slack", "C12345")
|
||||
|
||||
if tool.channel != "slack" {
|
||||
t.Errorf("Expected channel 'slack', got %q", tool.channel)
|
||||
}
|
||||
if tool.chatID != "C12345" {
|
||||
t.Errorf("Expected chatID 'C12345', got %q", tool.chatID)
|
||||
}
|
||||
}
|
||||
|
||||
func TestWebFetchTool_SSRFBlocked_NoPolicyEngine(t *testing.T) {
|
||||
tool := NewWebFetchToolWithPolicy(WebFetchToolOptions{
|
||||
MaxChars: 50000,
|
||||
SSRFMode: "block",
|
||||
})
|
||||
ctx := context.Background()
|
||||
|
||||
result := tool.Execute(ctx, map[string]interface{}{"url": "http://127.0.0.1/secret"})
|
||||
if !result.IsError {
|
||||
t.Error("Expected SSRF block without policy engine")
|
||||
}
|
||||
if !strings.Contains(result.ForLLM, "blocked") {
|
||||
t.Errorf("Expected 'blocked' in error, got: %s", result.ForLLM)
|
||||
}
|
||||
}
|
||||
|
||||
func TestWebFetchTool_DefaultMaxChars(t *testing.T) {
|
||||
tool := NewWebFetchToolWithPolicy(WebFetchToolOptions{MaxChars: 0})
|
||||
if tool.maxChars != 50000 {
|
||||
t.Errorf("Expected default maxChars 50000, got %d", tool.maxChars)
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -51,6 +51,7 @@ type DownloadOptions struct {
|
|||
Timeout time.Duration
|
||||
ExtraHeaders map[string]string
|
||||
LoggerPrefix string
|
||||
SkipURLValidation bool // When true, skip SSRF validation (caller handles it)
|
||||
}
|
||||
|
||||
// DownloadFile downloads a file from URL to a local temp directory.
|
||||
|
|
@ -64,7 +65,8 @@ func DownloadFile(url, filename string, opts DownloadOptions) string {
|
|||
opts.LoggerPrefix = "utils"
|
||||
}
|
||||
|
||||
// Validate URL against SSRF attacks
|
||||
// Validate URL against SSRF attacks (skipped when caller handles policy)
|
||||
if !opts.SkipURLValidation {
|
||||
if err := ValidateURL(url); err != nil {
|
||||
logger.ErrorCF(opts.LoggerPrefix, "URL validation failed", map[string]interface{}{
|
||||
"error": err.Error(),
|
||||
|
|
@ -72,6 +74,7 @@ func DownloadFile(url, filename string, opts DownloadOptions) string {
|
|||
})
|
||||
return ""
|
||||
}
|
||||
}
|
||||
|
||||
mediaDir := filepath.Join(os.TempDir(), "picoclaw_media")
|
||||
if err := os.MkdirAll(mediaDir, 0700); err != nil {
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue