feat(tools): replace shell exec with AST-based risk classifier
Replace regex deny-patterns with in-process POSIX interpreter (mvdan.cc/sh/v3) and four-tier risk classification. Add env sanitization, file-access sandboxing, and argument-aware risk modifiers. Cron commands now share the same security path. Signed-off-by: Boris Bliznioukov <blib@mail.com>
This commit is contained in:
parent
1945436dd4
commit
78e3508b79
23 changed files with 2513 additions and 992 deletions
|
|
@ -340,10 +340,11 @@
|
|||
}
|
||||
},
|
||||
"exec": {
|
||||
"enabled": true,
|
||||
"enable_deny_patterns": true,
|
||||
"custom_deny_patterns": null,
|
||||
"custom_allow_patterns": null
|
||||
"risk_threshold": "medium",
|
||||
"risk_overrides": {},
|
||||
"arg_modifiers": {},
|
||||
"env_allowlist": [],
|
||||
"env_set": {}
|
||||
},
|
||||
"skills": {
|
||||
"enabled": true,
|
||||
|
|
|
|||
277
docs/design/DDR-shell-tool-hardening.md
Normal file
277
docs/design/DDR-shell-tool-hardening.md
Normal file
|
|
@ -0,0 +1,277 @@
|
|||
# DDR: Shell Tool Hardening — AST-Based Guard + Interpreter Execution
|
||||
|
||||
| Field | Value |
|
||||
| --------------- | --------------------- |
|
||||
| **Status** | IMPLEMENTED |
|
||||
| **Date** | 2026-03-04 |
|
||||
| **Author** | — |
|
||||
| **Review date** | 2026-06-04 (3 months) |
|
||||
|
||||
---
|
||||
|
||||
## 1. Problem Frame
|
||||
|
||||
- **System**: PicoClaw `ExecTool` (`pkg/tools/shell.go`) — the shell command execution tool exposed to the LLM agent.
|
||||
- **Audience**: Maintainers, security reviewers, contributors.
|
||||
- **Baseline rule**: Commands are guarded by ~40 regex patterns matched against the raw command string (`guardCommand()`). Execution is via `sh -c <command>` (Unix) or `powershell` (Windows). Environment is fully inherited (`cmd.Env` is never set). Directory restriction uses regex-based absolute-path detection.
|
||||
- **Baseline config** (`ExecConfig`):
|
||||
```go
|
||||
EnableDenyPatterns bool // default: true
|
||||
CustomDenyPatterns []string // default: []
|
||||
CustomAllowPatterns []string // default: []
|
||||
```
|
||||
- **The problem**:
|
||||
1. **Regex guards are trivially bypassable.** Variable indirection (`x=rm; $x -rf /`), command substitution (`` `echo rm` -rf / ``), quoting tricks, IFS manipulation, hex/octal encoding, and Unicode escapes all evade string-level pattern matching. The guard operates on surface syntax, not semantics.
|
||||
2. **Full environment inheritance.** Child processes receive every env var from the parent — `OPENAI_API_KEY`, `ANTHROPIC_API_KEY`, provider tokens, database credentials. An LLM-generated `curl` command can exfiltrate these.
|
||||
3. **No structured risk classification.** The system is binary (blocked / allowed) with no granularity, no feedback to the LLM about _why_ a command was blocked, and no user-configurable risk threshold.
|
||||
4. **Path validation is regex-based.** Absolute path detection via regex is incomplete and bypassable via symlinks, relative path tricks, and shell expansion.
|
||||
|
||||
- **Proposed change**: Replace the regex guard and `sh -c` execution with (a) a proper shell parser (`mvdan.cc/sh/v3/syntax`), (b) an in-process shell interpreter (`mvdan.cc/sh/v3/interp`) with `ExecHandlers` middleware for runtime command interception, (c) a 4-level risk classifier, (d) env-var allowlisting, and (e) `OpenHandler`-based file-access sandboxing. Hard switch — no fallback to the old system.
|
||||
|
||||
- **Constraints**:
|
||||
- MUST NOT introduce a new external runtime dependency (no Docker, chroot, or OS-level sandboxing required).
|
||||
- MUST cover every command currently blocked by regex deny patterns.
|
||||
- MUST preserve the `Tool` interface contract (`Name`, `Description`, `Parameters`, `Execute`).
|
||||
- MUST work cross-platform (interpreter replaces both `sh -c` and `powershell` paths).
|
||||
- Config changes MUST be additive where possible; removed fields MUST produce a logged migration warning.
|
||||
|
||||
- **Non-goals**:
|
||||
- OS-level sandboxing (namespaces, seccomp, chroot).
|
||||
- Full bash compatibility (job control, complete `shopt`/`set` support).
|
||||
- Web-UI or approval workflows for blocked commands.
|
||||
|
||||
- **⚠ Hard limitation — external program file access**:
|
||||
|
||||
> **Executed programs can read and write any file the OS process has access to.** The `OpenHandler` only intercepts file opens from _shell redirections_ (`>`, `<`, `>>`). When the interpreter execs an external binary (e.g., `cat /etc/shadow`, `python script.py`, `tar czf /tmp/x.tar /`), that binary runs as a normal OS process with the full filesystem permissions of the picoclaw user. **This design does NOT sandbox external program I/O.** Mitigating this requires OS-level isolation (namespaces, seccomp, chroot) which is explicitly out of scope. The risk classifier and env sanitization reduce the _likelihood_ of damage but do not eliminate filesystem access.
|
||||
|
||||
- **Stakeholders**: Maintainers (implementation), self-hosters (config migration), LLM agent loop (structured error contract), channel integrations (no change).
|
||||
|
||||
- **Evidence**:
|
||||
- Current regex bypass is demonstrable: `x=rm; $x -rf /` passes all 40 deny patterns.
|
||||
- `cmd.Env` is never set — confirmed by code audit of `shell.go` L198–L212.
|
||||
- `mvdan.cc/sh/v3`: 244 importers, BSD-3-Clause, v3.12.0, actively maintained (powers `shfmt`).
|
||||
- `go-safe-cmd-runner` risk model reviewed: 4-level classification (low/medium/high/critical), env allowlisting, path validation. Concept adopted; no code dependency (0-star repo, 3/4 contributors are bots).
|
||||
|
||||
- **Acceptance criteria**:
|
||||
- [x] AC-1: All 40+ existing regex deny-pattern test cases pass under the new system (no security regression).
|
||||
- [x] AC-2: Variable indirection (`x=rm; $x -rf /`), command substitution, and quoting tricks are blocked at runtime via `ExecHandlers`.
|
||||
- [x] AC-3: Child processes receive only allowlisted env vars; `*_KEY`, `*_TOKEN`, `*_SECRET` vars are absent.
|
||||
- [x] AC-4: Commands classified above the configured `risk_threshold` return a structured `ToolResult` containing risk level, command name, and reason.
|
||||
- [x] AC-5: Shell redirections (`> /etc/passwd`) to paths outside the workspace are blocked by `OpenHandler`.
|
||||
- [x] AC-6: Common LLM-generated commands (`grep`, `find`, `cat`, `python -c`, `wc`, `jq`) execute successfully.
|
||||
- [x] AC-7: Deprecated config fields (`enable_deny_patterns`, `custom_deny_patterns`, `custom_allow_patterns`) produce a logged warning with migration guidance.
|
||||
- [x] AC-8: Cron tool uses the same guard system.
|
||||
|
||||
---
|
||||
|
||||
## 2. Decision
|
||||
|
||||
### 2.1 Definitions
|
||||
|
||||
- **Risk level**: One of `low`, `medium`, `high`, `critical` — a classification of the potential damage a resolved shell command can cause.
|
||||
- **Risk threshold**: The maximum risk level the system will allow to execute. Commands above this level are blocked.
|
||||
- **Env allowlist**: The set of environment variable names permitted to propagate to child processes. Everything else is stripped.
|
||||
- **Resolved command**: The actual binary path and arguments after all shell expansion (variables, globs, substitution) has been performed by the interpreter.
|
||||
|
||||
### 2.2 Rule Changes
|
||||
|
||||
1. The `ExecTool` MUST parse all command strings using `mvdan.cc/sh/v3/syntax` with `Variant(syntax.LangBash)`. Unparseable input MUST be rejected with a structured error.
|
||||
|
||||
2. The `ExecTool` MUST execute commands using `mvdan.cc/sh/v3/interp.Runner` instead of `os/exec` with `sh -c` or `powershell`. The Windows-specific PowerShell code path MUST be removed.
|
||||
|
||||
3. The `interp.Runner` MUST be configured with an `ExecHandlers` middleware that intercepts every resolved external command and classifies it against a risk table before execution.
|
||||
|
||||
4. The risk classifier MUST implement four levels:
|
||||
|
||||
| Level | Default action | Characterization |
|
||||
| ---------- | -------------- | ------------------------------------------------------- |
|
||||
| `low` | Allow | Read-only, informational (ls, cat, grep, find, wc) |
|
||||
| `medium` | Allow (logged) | File modification, network read (cp, mv, curl GET) |
|
||||
| `high` | Block | Destructive, system-modifying (rm, chmod, git push) |
|
||||
| `critical` | Block | Privilege escalation, always dangerous (sudo, dd, eval) |
|
||||
|
||||
5. The risk classifier MUST apply argument-aware modifiers (e.g., `curl` is `medium`, but `curl -X POST` is `high`; `git` is `medium`, but `git push` is `high`).
|
||||
|
||||
6. When a command is blocked, the `ToolResult` MUST include:
|
||||
- Risk level of the command.
|
||||
- The blocked command name and arguments.
|
||||
- The configured threshold.
|
||||
- A human-readable reason string.
|
||||
|
||||
7. The `ExecTool` MUST construct a sanitized environment for the `interp.Runner` using a strict allowlist. The default allowlist MUST be:
|
||||
|
||||
```
|
||||
PATH, HOME, USER, LANG, SHELL, TERM, PWD, OLDPWD,
|
||||
HOSTNAME, LOGNAME, TZ, DISPLAY, TMPDIR, EDITOR, PAGER
|
||||
```
|
||||
|
||||
Plus all variables matching the `LC_*` prefix.
|
||||
|
||||
8. The `interp.Runner` MUST be configured with an `OpenHandler` that validates all file-open paths (from shell redirections) resolve within the configured workspace directory. Paths to safe pseudo-devices (`/dev/null`, `/dev/zero`, `/dev/urandom`, `/dev/stdin`, `/dev/stdout`, `/dev/stderr`) MUST be exempted.
|
||||
|
||||
9. The existing regex-based guard (`defaultDenyPatterns`, `guardCommand()`) MUST be removed entirely. The `ExecConfig` fields `EnableDenyPatterns`, `CustomDenyPatterns`, and `CustomAllowPatterns` MUST be removed from the struct.
|
||||
|
||||
10. The `ExecConfig` struct MUST be extended with:
|
||||
|
||||
```go
|
||||
RiskThreshold string `json:"risk_threshold"` // "low"|"medium"|"high"|"critical"; default "medium"
|
||||
RiskOverrides map[string]string `json:"risk_overrides"` // command → level override
|
||||
EnvAllowlist []string `json:"env_allowlist"` // extra vars to pass (extends defaults)
|
||||
EnvSet map[string]string `json:"env_set"` // explicit var=value pairs
|
||||
ArgModifiers map[string][]ArgModifierConfig `json:"arg_modifiers"` // command → argument-aware risk adjustments (extends built-ins)
|
||||
```
|
||||
|
||||
`ArgModifierConfig` is `struct { Args []string; Level string }`. User-defined modifiers are checked alongside built-ins using highest-match-wins semantics (the maximum level across all matching modifiers is applied).
|
||||
|
||||
11. If deprecated config fields (`enable_deny_patterns`, `custom_deny_patterns`, `custom_allow_patterns`) are present in user config, the system SHOULD log a warning with migration instructions. The system MUST NOT fail to start.
|
||||
|
||||
12. The cron tool (`pkg/tools/cron.go`) MUST use the same `ExecTool` with the same guard system.
|
||||
|
||||
13. The `ExecTool` MUST implement the `AsyncTool` interface (`SetCallback(AsyncCallback)`). When the LLM passes `background=true` and a callback is registered, the command MUST be launched in a goroutine; the result is delivered via the callback. When `background=true` but no callback is registered, execution falls through to synchronous mode. Compile-time interface check: `var _ AsyncTool = (*ExecTool)(nil)`.
|
||||
|
||||
14. The implementation MUST use a subpackage structure: `pkg/tools/shell/` contains the core logic (risk classifier, env sanitizer, sandbox, runner) and `pkg/tools/shell_tool.go` is the thin adapter implementing the `Tool` + `AsyncTool` interfaces. Tests live alongside their source in both locations.
|
||||
|
||||
### 2.3 Migration
|
||||
|
||||
- **Config**: Old fields are silently ignored with a logged warning. No config version bump required. Add a migration note to `docs/tools_configuration.md`.
|
||||
- **Behavioral**: Commands that previously passed regex checks but are genuinely dangerous (variable indirection bypasses) will now be blocked. This is intentional and constitutes the security fix.
|
||||
- **Backward compatibility**: The `Tool` interface (`Name`, `Description`, `Parameters`, `Execute`) is unchanged. Callers (`ToolRegistry`, `RunToolLoop`, agent instance) require no changes.
|
||||
|
||||
---
|
||||
|
||||
## 3. Rationale
|
||||
|
||||
### Core reasons
|
||||
|
||||
- Regex pattern matching against raw shell strings is a fundamentally wrong abstraction level. Shell is a programming language; security analysis requires parsing it as one.
|
||||
- `mvdan.cc/sh/v3` provides a battle-tested parser (powers `shfmt`, 244+ importers) and an interpreter with middleware hooks (`ExecHandlers`, `OpenHandler`, `Env`) designed exactly for sandboxed execution.
|
||||
- Runtime interception via `ExecHandlers` catches dynamic command construction (variable expansion, command substitution, arithmetic evaluation) that static analysis cannot.
|
||||
- Env inheritance is a silent data-exfiltration vector. The LLM generates arbitrary commands; any `curl`/`wget`/`nc` call can leak every secret in the parent process environment.
|
||||
- A 4-level risk classifier with structured LLM feedback is strictly better than binary block/allow — it lets the agent retry with safer alternatives.
|
||||
|
||||
### Alternatives considered
|
||||
|
||||
| # | Alternative | Disposition | Why |
|
||||
| --- | ------------------------------------------------------------------------------------------ | ----------------------------- | --------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- |
|
||||
| 1 | **Do nothing** | Rejected | Regex guards are demonstrably bypassable. Env leak is an active vulnerability. |
|
||||
| 2 | **Improve regexes** (add more patterns, normalize input before matching) | Rejected | Arms race against shell syntax. Every new evasion requires a new regex. Fundamentally wrong abstraction. Input normalization is itself a parsing problem. |
|
||||
| 3 | **Static AST analysis only** (parse with `mvdan.cc/sh`, walk tree, keep `sh -c` execution) | Rejected | Cannot resolve dynamic constructs (`$x`, `$(...)`) at parse time. Catches structural patterns but misses the primary bypass vector (variable indirection). |
|
||||
| 4 | **Hybrid: static AST + `sh -c` with improved guards** | Rejected | Two systems to maintain, still bypassable at exec time. Adds complexity without solving the core problem. |
|
||||
| 5 | **OS-level sandboxing** (namespaces, seccomp-bpf, chroot) | Rejected as primary mechanism | Requires elevated privileges, platform-specific, heavy. Out of scope. Could be layered on top later. |
|
||||
| 6 | **Import `go-safe-cmd-runner` as dependency** | Rejected | Overkill framework (TOML config, SUID privilege management, ULID audit trails). 0 stars, 3/4 contributors are bots. The useful concept (risk levels + env allowlist) is simple enough to own. |
|
||||
|
||||
### Trade-offs
|
||||
|
||||
| Dimension | Cost | Benefit |
|
||||
| ----------------- | ---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | -------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- |
|
||||
| **Complexity** | New dependency (`mvdan.cc/sh/v3`), new risk table to maintain | Eliminates regex arms race, single enforcement point |
|
||||
| **Compatibility** | Interpreter is not 100% bash (no job control, limited shopt) | Covers >95% of LLM-generated commands; bash-isms parse and execute |
|
||||
| **Performance** | Parse step adds <1ms per command; interpreter overhead comparable to `sh -c` | No meaningful regression for typical agent commands |
|
||||
| **Binary size** | `mvdan.cc/sh/v3` adds ~450 KB to the compiled binary (parser + interpreter + syntax tables) | Acceptable for a security-critical component; no runtime memory cost beyond parse/exec |
|
||||
| **Safety** | **`OpenHandler` only catches shell redirections.** External programs (`cat`, `python`, `tar`, etc.) can still read/write any file the process user has access to — this is an OS-level limitation, not solvable without namespace/seccomp isolation. | Significant improvement over regex path detection for shell-level I/O. Env sanitization prevents secret exfiltration via env. Risk classifier blocks known-dangerous commands. But **filesystem access by executed binaries remains unsandboxed.** |
|
||||
| **Usability** | Some previously-allowed commands may be blocked by risk classifier | Configurable threshold + per-command overrides; structured feedback enables LLM retry |
|
||||
| **Migration** | Config fields removed; users with custom patterns lose them | Clear migration path; new system is strictly more expressive |
|
||||
|
||||
### Risks and mitigations
|
||||
|
||||
| Risk | Likelihood | Impact | Mitigation |
|
||||
| --------------------------------------------------- | ---------- | --------------------------------------- | ---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- |
|
||||
| Interpreter fails on valid bash edge case | Medium | Medium (command fails) | Hard switch makes failures visible immediately. LangBash covers most LLM output. Users can file bugs against `mvdan/sh`. |
|
||||
| Risk table misclassifies a command (false positive) | Medium | Low (command blocked, LLM retries) | Per-command `risk_overrides` config. Structured error tells LLM what happened. |
|
||||
| Risk table misclassifies a command (false negative) | Low | High (dangerous command runs) | Defense in depth: env sanitization limits damage. OpenHandler blocks shell-level file access. Risk table starts conservative. |
|
||||
| **External programs read/write arbitrary files** | **High** | **High** (data exfil, data destruction) | **Not mitigated by this design.** `cat`, `python`, `curl --upload-file`, `tar` etc. run as normal OS processes with full FS access. Risk classifier blocks _known_ dangerous commands but cannot prevent a novel `python -c "open('/etc/passwd').read()"`. OS-level sandboxing (future work) is the only real fix. Env sanitization reduces secret exposure but not file exposure. |
|
||||
| `mvdan.cc/sh/v3` dependency becomes unmaintained | Low | Medium | Widely used (shfmt), BSD-licensed, vendorable. Could fork if needed. |
|
||||
| Binary size increase (~450 KB) | Certain | Low | One-time cost. Acceptable for security infrastructure. No runtime memory overhead beyond parse/exec. |
|
||||
| Performance regression on command-heavy agent loops | Low | Low | Parser is <1ms. Benchmark before/after in CI. |
|
||||
|
||||
---
|
||||
|
||||
## 4. Consequences
|
||||
|
||||
### Immediate impacts
|
||||
|
||||
| Who / what | Impact |
|
||||
| ------------------------------------------------------------- | -------------------------------------------------------------------------------------------------------------- |
|
||||
| `pkg/tools/shell_tool.go` | Thin adapter: `ExecTool` struct implementing `Tool` + `AsyncTool`. Delegates to `pkg/tools/shell/` subpackage. |
|
||||
| `pkg/tools/shell_tool_test.go` | Tests for `ExecTool` sync/async behavior and interface compliance. |
|
||||
| `pkg/tools/shell_process_unix.go`, `shell_process_windows.go` | Removed. Interpreter manages process lifecycle. |
|
||||
| `pkg/config/config.go` (`ExecConfig`) | Three fields removed, four fields added. |
|
||||
| `pkg/config/defaults.go` | Default `RiskThreshold` set to `"medium"`. |
|
||||
| `pkg/tools/cron.go` | Updated to use new `ExecTool` constructor. |
|
||||
| `docs/tools_configuration.md` | Rewritten for new config fields and risk model. |
|
||||
| `config/config.example.json` | Updated. |
|
||||
| `go.mod` | `mvdan.cc/sh/v3` added. |
|
||||
| Self-hosters with `custom_deny_patterns` | Logged warning on startup. Patterns no longer functional. Must migrate to `risk_overrides`. |
|
||||
| LLM agent loop | No code change. Receives richer error messages on blocked commands. |
|
||||
|
||||
### New files
|
||||
|
||||
| File | Purpose |
|
||||
| --------------------------------- | ------------------------------------------------------------------------------------------------------------------------------------------ |
|
||||
| `pkg/tools/shell/risk.go` | Risk classifier: command→level mapping, argument modifiers (built-in + configurable), `ClassifyCommand` with highest-match-wins semantics. |
|
||||
| `pkg/tools/shell/risk_test.go` | Table-driven tests for all 4 levels, argument modifiers, overrides, extra modifiers. |
|
||||
| `pkg/tools/shell/env.go` | Env sanitization: allowlist builder, default list, `LC_*` prefix matching, `env_set` application. |
|
||||
| `pkg/tools/shell/env_test.go` | Verify secrets stripped, allowlist passed, `env_set` applied. |
|
||||
| `pkg/tools/shell/sandbox.go` | `OpenHandler` implementation, path-within-workspace validation, pseudo-device exemptions. |
|
||||
| `pkg/tools/shell/sandbox_test.go` | Redirect inside/outside workspace, symlink escape, safe-path exemption. |
|
||||
| `pkg/tools/shell/runner.go` | `Run` function: parser + interpreter + `ExecHandlers` middleware integration. |
|
||||
| `pkg/tools/shell/runner_test.go` | End-to-end runner tests: timeout, working dir, env sanitization, pipelines. |
|
||||
| `pkg/tools/shell_tool.go` | Adapter: `ExecTool` struct, `NewExecToolWithConfig`, `AsyncTool` impl, arg modifier wiring. |
|
||||
| `pkg/tools/shell_tool_test.go` | `ExecTool` sync/async tests, interface compliance checks. |
|
||||
| `pkg/tools/cron_exec_test.go` | AC-8: cron-originated `ExecTool` blocks dangerous commands identically to agent-created one; safe commands pass. |
|
||||
|
||||
### Follow-up tasks
|
||||
|
||||
| # | Task | Owner (role) | Blocks | Status |
|
||||
| --- | ------------------------------------------------------------------------------------------- | --------------- | ------ | -------- |
|
||||
| 1 | Implement risk classifier (`risk.go`) with command table | Maintainer | 3, 4 | **Done** |
|
||||
| 2 | Implement env sanitization (`env.go`) | Maintainer | 4 | **Done** |
|
||||
| 3 | Implement sandbox OpenHandler (`sandbox.go`) | Maintainer | 4 | **Done** |
|
||||
| 4 | Rewrite shell tool: parser + interpreter + middleware (`shell/runner.go` + `shell_tool.go`) | Maintainer | 5, 6 | **Done** |
|
||||
| 5 | Port all existing test cases + add bypass tests | Maintainer | 7 | **Done** |
|
||||
| 6 | Update `ExecConfig`, defaults, migration warning | Maintainer | 7 | **Done** |
|
||||
| 7 | Implement `AsyncTool` on `ExecTool` | Maintainer | — | **Done** |
|
||||
| 8 | Implement configurable `ArgModifiers` (user-defined, highest-match-wins) | Maintainer | — | **Done** |
|
||||
| 9 | Fix runner_test PATH resolution for external binaries in sandboxed interpreter | Maintainer | — | **Done** |
|
||||
| 10 | Update cron tool, docs, config example | Maintainer | — | **Done** |
|
||||
| 11 | CI: add benchmark comparison (old vs new exec latency) | Maintainer | — | Open |
|
||||
| 12 | Run against corpus of real LLM-generated commands (compatibility) | QA / Maintainer | — | Open |
|
||||
|
||||
### Test / verification plan
|
||||
|
||||
| AC | Test | Method |
|
||||
| ---- | ------------------------------------------------------ | --------------------------------------------------------------------------------------------------------------------------------------------------------- |
|
||||
| AC-1 | Port all 40+ regex deny-pattern test cases | Unit tests in `shell_test.go` — same inputs, same block expectations |
|
||||
| AC-2 | Variable indirection, cmd substitution, quoting bypass | New unit tests: `x=rm; $x -rf /`, `` `echo rm` -rf / ``, `'r'"m"' -rf /` |
|
||||
| AC-3 | Env sanitization | Unit test: set `OPENAI_API_KEY=secret` in parent, verify absent in child env. Verify `PATH`, `HOME` present. |
|
||||
| AC-4 | Structured error on blocked command | Unit test: execute `sudo ls`, verify `ToolResult.ForLLM` contains risk level, command, threshold, reason. Verify `IsError == true`. |
|
||||
| AC-5 | OpenHandler path validation | Unit test: `echo test > /etc/shadow` — blocked. `echo test > ./output.txt` — allowed. Symlink to outside — blocked. `/dev/null` — allowed. |
|
||||
| AC-6 | Compatibility with common commands | Integration test: `grep -r pattern .`, `find . -name '*.go'`, `cat file.txt`, `python3 -c "print(1)"`, `wc -l file`, `jq .field file.json` — all succeed. |
|
||||
| AC-7 | Config migration warning | Unit test: load config with `enable_deny_patterns: false` — verify logged warning, no startup failure. |
|
||||
| AC-8 | Cron uses same system | Unit test: cron-created `ExecTool` blocks `sudo ls` identically to agent-created one. |
|
||||
|
||||
### Rollback plan
|
||||
|
||||
- The change is contained within `pkg/tools/` and `pkg/config/`. Git revert of the implementation commits restores the regex-based system.
|
||||
- Assumption: no other packages depend on `ExecTool` internals (only the `Tool` interface is public contract). Verified: only `pkg/agent/instance.go` and `pkg/tools/cron.go` call `NewExecToolWithConfig`.
|
||||
- If the interpreter causes widespread command failures post-deploy, revert and reconsider the hybrid approach (Alternative #4).
|
||||
|
||||
### Decision review date
|
||||
|
||||
**2026-06-04** — Review:
|
||||
|
||||
- False-positive rate (commands incorrectly blocked).
|
||||
- False-negative rate (dangerous commands that slipped through).
|
||||
- Interpreter compatibility issues reported.
|
||||
- Whether `OpenHandler` limitations warrant OS-level sandboxing investment.
|
||||
|
||||
---
|
||||
|
||||
## Assumptions
|
||||
|
||||
- **Assumption**: >95% of LLM-generated shell commands are compatible with `mvdan.cc/sh/v3/interp` in LangBash mode. Based on: LLMs generate simple pipelines, not job-control or advanced bash internals.
|
||||
- **Assumption**: The `ExecHandlers` middleware sees the fully-resolved command (after variable expansion, globbing, etc.) before execution. Based on: `mvdan.cc/sh/v3/interp` documentation for `ExecHandlerFunc`.
|
||||
- **Assumption**: No downstream consumers depend on the `ExecConfig` fields being removed. Based on: config is JSON-deserialized; unknown fields are ignored by default in Go's `encoding/json`.
|
||||
- **Assumption**: `mvdan.cc/sh/v3` will remain maintained for the foreseeable future. Based on: BSD-3-Clause, powers `shfmt`, 244+ importers, latest release v3.12.0 (Jul 2025).
|
||||
- **Assumption**: `mvdan.cc/sh/v3` adds ~450 KB to compiled binary size. Based on: typical Go module size for parser+interpreter; to be verified with `go build` size comparison before/after.
|
||||
- **Assumption**: Users accept that executed external programs have full filesystem access. This is inherent to exec-based execution and is not a regression — the current `sh -c` system has the same property. The new system makes this limitation _explicit_ rather than hiding it behind regex theater.
|
||||
|
|
@ -45,34 +45,70 @@ Web tools are used for web search and fetching.
|
|||
|
||||
## Exec Tool
|
||||
|
||||
The exec tool is used to execute shell commands.
|
||||
The exec tool executes shell commands using an in-process POSIX interpreter with
|
||||
AST-based risk classification, environment sanitization, and file-access sandboxing.
|
||||
|
||||
### Configuration
|
||||
|
||||
| Config | Type | Default | Description |
|
||||
| ---------------------- | ----- | ------- | ------------------------------------------ |
|
||||
| `enable_deny_patterns` | bool | true | Enable default dangerous command blocking |
|
||||
| `custom_deny_patterns` | array | [] | Custom deny patterns (regular expressions) |
|
||||
| ---------------- | ------ | ---------- | ----------------------------------------------------------------------- |
|
||||
| `risk_threshold` | string | `"medium"` | Maximum allowed risk level: `"low"`, `"medium"`, `"high"`, `"critical"` |
|
||||
| `risk_overrides` | object | `{}` | Per-command risk level overrides (command name → level) |
|
||||
| `arg_modifiers` | object | `{}` | Per-command argument patterns that adjust risk level |
|
||||
| `env_allowlist` | array | `[]` | Extra environment variables to expose (extends built-in defaults) |
|
||||
| `env_set` | object | `{}` | Explicit `VAR=value` pairs injected into every command |
|
||||
|
||||
### Functionality
|
||||
### Risk Classification
|
||||
|
||||
- **`enable_deny_patterns`**: Set to `false` to completely disable the default dangerous command blocking patterns
|
||||
- **`custom_deny_patterns`**: Add custom deny regex patterns; commands matching these will be blocked
|
||||
Every command is parsed into an AST before execution. Each resolved binary is
|
||||
looked up in a built-in risk table with four levels:
|
||||
|
||||
### Default Blocked Command Patterns
|
||||
| Level | Meaning | Examples |
|
||||
| ---------- | ------------------------------------- | ---------------------------------- |
|
||||
| `low` | Read-only / informational | `echo`, `cat`, `ls`, `date` |
|
||||
| `medium` | Writes files but limited blast radius | `cp`, `mv`, `mkdir`, `tee` |
|
||||
| `high` | System-wide side effects | `apt`, `brew`, `docker`, `mount` |
|
||||
| `critical` | Destructive / privilege escalation | `sudo`, `rm -rf`, `shutdown`, `dd` |
|
||||
|
||||
By default, PicoClaw blocks the following dangerous commands:
|
||||
Commands with a risk level **above** `risk_threshold` are blocked before execution.
|
||||
|
||||
- Delete commands: `rm -rf`, `del /f/q`, `rmdir /s`
|
||||
- Disk operations: `format`, `mkfs`, `diskpart`, `dd if=`, writing to `/dev/sd*`
|
||||
- System operations: `shutdown`, `reboot`, `poweroff`
|
||||
- Command substitution: `$()`, `${}`, backticks
|
||||
- Pipe to shell: `| sh`, `| bash`
|
||||
- Privilege escalation: `sudo`, `chmod`, `chown`
|
||||
- Process control: `pkill`, `killall`, `kill -9`
|
||||
- Remote operations: `curl | sh`, `wget | sh`, `ssh`
|
||||
- Package management: `apt`, `yum`, `dnf`, `npm install -g`, `pip install --user`
|
||||
- Containers: `docker run`, `docker exec`
|
||||
- Git: `git push`, `git force`
|
||||
- Other: `eval`, `source *.sh`
|
||||
#### Argument modifiers
|
||||
|
||||
Some commands change risk depending on their arguments. For example, `rm` is
|
||||
`medium` by default but becomes `critical` when called with `-rf`.
|
||||
|
||||
You can add custom argument modifiers via config. Each entry lists tokens that
|
||||
must all be present (order-independent) and the resulting level:
|
||||
|
||||
```json
|
||||
{
|
||||
"arg_modifiers": {
|
||||
"curl": [{ "args": ["--upload-file"], "level": "high" }],
|
||||
"git": [{ "args": ["push", "--force"], "level": "critical" }]
|
||||
}
|
||||
}
|
||||
```
|
||||
|
||||
The **highest matching** modifier wins (built-in and custom are merged).
|
||||
|
||||
### Environment Sanitization
|
||||
|
||||
The shell interpreter runs with a sanitized environment. Only a safe allowlist
|
||||
of variables is exposed (e.g., `PATH`, `HOME`, `LANG`, `TERM`).
|
||||
|
||||
- **`env_allowlist`**: Extend the defaults with additional variable names.
|
||||
- **`env_set`**: Inject fixed `VAR=value` pairs (overrides real env).
|
||||
|
||||
### File-Access Sandboxing
|
||||
|
||||
When `restrict_to_workspace` is enabled (the default), the interpreter's
|
||||
`OpenHandler` blocks reads and writes outside the configured workspace directory.
|
||||
|
||||
### Cron Integration
|
||||
|
||||
The cron tool creates its own `ExecTool` via `NewExecToolWithConfig`, so
|
||||
scheduled commands go through the same risk classifier, env sanitization, and
|
||||
sandbox as agent-originated commands.
|
||||
|
||||
### Configuration Example
|
||||
|
||||
|
|
@ -80,8 +116,18 @@ By default, PicoClaw blocks the following dangerous commands:
|
|||
{
|
||||
"tools": {
|
||||
"exec": {
|
||||
"enable_deny_patterns": true,
|
||||
"custom_deny_patterns": ["\\brm\\s+-r\\b", "\\bkillall\\s+python"]
|
||||
"risk_threshold": "medium",
|
||||
"risk_overrides": {
|
||||
"ffmpeg": "low",
|
||||
"terraform": "critical"
|
||||
},
|
||||
"arg_modifiers": {
|
||||
"curl": [{ "args": ["--upload-file"], "level": "high" }]
|
||||
},
|
||||
"env_allowlist": ["GOPATH", "JAVA_HOME"],
|
||||
"env_set": {
|
||||
"NODE_ENV": "production"
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
|
|||
1
go.mod
1
go.mod
|
|
@ -28,6 +28,7 @@ require (
|
|||
golang.org/x/time v0.14.0
|
||||
google.golang.org/protobuf v1.36.11
|
||||
modernc.org/sqlite v1.46.1
|
||||
mvdan.cc/sh/v3 v3.12.0
|
||||
)
|
||||
|
||||
require (
|
||||
|
|
|
|||
12
go.sum
12
go.sum
|
|
@ -40,6 +40,8 @@ github.com/coder/websocket v1.8.14/go.mod h1:NX3SzP+inril6yawo5CQXx8+fk145lPDC6p
|
|||
github.com/coreos/go-systemd/v22 v22.5.0/go.mod h1:Y58oyj3AT4RCenI/lSvhwexgC+NSVTIJ3seZv2GcEnc=
|
||||
github.com/cpuguy83/go-md2man/v2 v2.0.6/go.mod h1:oOW0eioCTA6cOiMLiUPZOpcVxMig6NIQQ7OS05n1F4g=
|
||||
github.com/creack/pty v1.1.9/go.mod h1:oKZEueFk5CKHvIhNR5MUki03XCEU+Q6VDXinZuGJ33E=
|
||||
github.com/creack/pty v1.1.24 h1:bJrF4RRfyJnbTJqzRLHzcGaZK1NeM5kTC9jGgovnR1s=
|
||||
github.com/creack/pty v1.1.24/go.mod h1:08sCNb52WyoAwi2QDyzUCTgcvVFhUzewun7wtTfvcwE=
|
||||
github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
|
||||
github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c=
|
||||
github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
|
||||
|
|
@ -56,6 +58,8 @@ github.com/gdamore/tcell/v2 v2.13.8 h1:Mys/Kl5wfC/GcC5Cx4C2BIQH9dbnhnkPgS9/wF3Rl
|
|||
github.com/gdamore/tcell/v2 v2.13.8/go.mod h1:+Wfe208WDdB7INEtCsNrAN6O2m+wsTPk1RAovjaILlo=
|
||||
github.com/github/copilot-sdk/go v0.1.23 h1:uExtO/inZQndCZMiSAA1hvXINiz9tqo/MZgQzFzurxw=
|
||||
github.com/github/copilot-sdk/go v0.1.23/go.mod h1:GdwwBfMbm9AABLEM3x5IZKw4ZfwCYxZ1BgyytmZenQ0=
|
||||
github.com/go-quicktest/qt v1.101.0 h1:O1K29Txy5P2OK0dGo59b7b0LR6wKfIhttaAhHUyn7eI=
|
||||
github.com/go-quicktest/qt v1.101.0/go.mod h1:14Bz/f7NwaXPtdYEgzsx46kqSxVwTbzVZsDC26tQJow=
|
||||
github.com/go-redis/redis/v8 v8.11.4/go.mod h1:2Z2wHZXdQpCDXEGzqMockDpNyYvi2l4Pxt6RJr792+w=
|
||||
github.com/go-resty/resty/v2 v2.6.0/go.mod h1:PwvJS6hvaPkjtjNg9ph+VrSD92bi5Zq73w/BIH7cC3Q=
|
||||
github.com/go-resty/resty/v2 v2.17.1 h1:x3aMpHK1YM9e4va/TMDRlusDDoZiQ+ViDu/WpA6xTM4=
|
||||
|
|
@ -113,8 +117,9 @@ github.com/klauspost/cpuid/v2 v2.3.0 h1:S4CRMLnYUhGeDFDqkGriYKdfoFlDnMtqTiI/sFzh
|
|||
github.com/klauspost/cpuid/v2 v2.3.0/go.mod h1:hqwkgyIinND0mEev00jJYCxPNVRVXFQeu1XKlok6oO0=
|
||||
github.com/kr/pretty v0.1.0/go.mod h1:dAy3ld7l9f0ibDNOQOHHMYYIIbhfbHSm3C4ZsoJORNo=
|
||||
github.com/kr/pretty v0.2.1/go.mod h1:ipq/a2n7PKx3OHsz4KJII5eveXtPO4qwEXGdVfWzfnI=
|
||||
github.com/kr/pretty v0.3.0 h1:WgNl7dwNpEZ6jJ9k1snq4pZsg7DOEN8hP9Xw0Tsjwk0=
|
||||
github.com/kr/pretty v0.3.0/go.mod h1:640gp4NfQd8pI5XOwp5fnNeVWj67G7CFk/SaSQn7NBk=
|
||||
github.com/kr/pretty v0.3.1 h1:flRD4NNwYAUpkphVc1HcthR4KEIFJ65n8Mw5qdRn3LE=
|
||||
github.com/kr/pretty v0.3.1/go.mod h1:hoEshYVHaxMs3cyo3Yncou5ZscifuDolrwPKZanG3xk=
|
||||
github.com/kr/pty v1.1.1/go.mod h1:pFQYn66WHrOpPYNljwOMqo10TkYh1fy3cYio2l3bCsQ=
|
||||
github.com/kr/text v0.1.0/go.mod h1:4Jbv+DJW3UT/LiOwJeYQe1efqtUx/iVham/4vfdArNI=
|
||||
github.com/kr/text v0.2.0 h1:5Nx0Ya0ZqY2ygV366QzturHI13Jq95ApcVaJBhpS+AY=
|
||||
|
|
@ -165,8 +170,9 @@ github.com/rivo/tview v0.42.0/go.mod h1:cSfIYfhpSGCjp3r/ECJb+GKS7cGJnqV8vfjQPwoX
|
|||
github.com/rivo/uniseg v0.4.7 h1:WUdvkW8uEhrYfLC4ZzdpI2ztxP1I582+49Oc5Mq64VQ=
|
||||
github.com/rivo/uniseg v0.4.7/go.mod h1:FN3SvrM+Zdj16jyLfmOkMNblXMcoc8DfTHruCPUcx88=
|
||||
github.com/rogpeppe/go-internal v1.6.1/go.mod h1:xXDCJY+GAPziupqXw64V24skbSoqbTEfhy4qGm1nDQc=
|
||||
github.com/rogpeppe/go-internal v1.9.0 h1:73kH8U+JUqXU8lRuOHeVHaa/SZPifC7BkcraZVejAe8=
|
||||
github.com/rogpeppe/go-internal v1.9.0/go.mod h1:WtVeX8xhTBvf0smdhujwtBcq4Qrzq/fJaraNFVN+nFs=
|
||||
github.com/rogpeppe/go-internal v1.14.1 h1:UQB4HGPB6osV0SQTLymcB4TgvyWu6ZyliaW0tI/otEQ=
|
||||
github.com/rogpeppe/go-internal v1.14.1/go.mod h1:MaRKkUm5W0goXpeCfT7UZI6fk/L7L7so1lCWt35ZSgc=
|
||||
github.com/rs/xid v1.6.0/go.mod h1:7XoLgs4eV+QndskICGsho+ADou8ySMSjJKDIan90Nz0=
|
||||
github.com/rs/zerolog v1.34.0 h1:k43nTLIwcTVQAncfCw4KZ2VY6ukYoZaBPNOE8txlOeY=
|
||||
github.com/rs/zerolog v1.34.0/go.mod h1:bJsvje4Z08ROH4Nhs5iH600c3IkWhwp44iRc54W6wYQ=
|
||||
|
|
@ -383,5 +389,7 @@ modernc.org/strutil v1.2.1 h1:UneZBkQA+DX2Rp35KcM69cSsNES9ly8mQWD71HKlOA0=
|
|||
modernc.org/strutil v1.2.1/go.mod h1:EHkiggD70koQxjVdSBM3JKM7k6L0FbGE5eymy9i3B9A=
|
||||
modernc.org/token v1.1.0 h1:Xl7Ap9dKaEs5kLoOQeQmPWevfnk/DM5qcLcYlA8ys6Y=
|
||||
modernc.org/token v1.1.0/go.mod h1:UGzOrNV1mAFSEB63lOFHIpNRUVMvYTc6yu1SMY/XTDM=
|
||||
mvdan.cc/sh/v3 v3.12.0 h1:ejKUR7ONP5bb+UGHGEG/k9V5+pRVIyD+LsZz7o8KHrI=
|
||||
mvdan.cc/sh/v3 v3.12.0/go.mod h1:Se6Cj17eYSn+sNooLZiEUnNNmNxg0imoYlTu4CyaGyg=
|
||||
rsc.io/qr v0.2.0 h1:6vBLea5/NRMVTz8V66gipeLycZMl/+UlFmk8DvqQ6WY=
|
||||
rsc.io/qr v0.2.0/go.mod h1:IF+uZjkb9fqyeF/4tlBoynqmQxUoPfWEKh921coOuXs=
|
||||
|
|
|
|||
|
|
@ -602,12 +602,23 @@ type CronToolsConfig struct {
|
|||
ExecTimeoutMinutes int ` env:"PICOCLAW_TOOLS_CRON_EXEC_TIMEOUT_MINUTES" json:"exec_timeout_minutes"` // 0 means no timeout
|
||||
}
|
||||
|
||||
// ArgModifierConfig describes an argument pattern that elevates a command's risk.
|
||||
type ArgModifierConfig struct {
|
||||
Args []string `json:"args"` // tokens that must all be present (order-independent)
|
||||
Level string `json:"level"` // target risk level: "low"|"medium"|"high"|"critical"
|
||||
}
|
||||
|
||||
type ExecConfig struct {
|
||||
ToolConfig ` envPrefix:"PICOCLAW_TOOLS_EXEC_"`
|
||||
EnableDenyPatterns bool ` env:"PICOCLAW_TOOLS_EXEC_ENABLE_DENY_PATTERNS" json:"enable_deny_patterns"`
|
||||
CustomDenyPatterns []string ` env:"PICOCLAW_TOOLS_EXEC_CUSTOM_DENY_PATTERNS" json:"custom_deny_patterns"`
|
||||
CustomAllowPatterns []string ` env:"PICOCLAW_TOOLS_EXEC_CUSTOM_ALLOW_PATTERNS" json:"custom_allow_patterns"`
|
||||
TimeoutSeconds int ` env:"PICOCLAW_TOOLS_EXEC_TIMEOUT_SECONDS" json:"timeout_seconds"` // 0 means use default (60s)
|
||||
RiskThreshold string `json:"risk_threshold" env:"PICOCLAW_TOOLS_EXEC_RISK_THRESHOLD"` // "low"|"medium"|"high"|"critical"; default "medium"
|
||||
RiskOverrides map[string]string `json:"risk_overrides" env:"PICOCLAW_TOOLS_EXEC_RISK_OVERRIDES"` // command → level override
|
||||
ArgModifiers map[string][]ArgModifierConfig `json:"arg_modifiers" env:"PICOCLAW_TOOLS_EXEC_ARG_MODIFIERS"` // command → argument-aware risk adjustments (extends built-ins)
|
||||
EnvAllowlist []string `json:"env_allowlist" env:"PICOCLAW_TOOLS_EXEC_ENV_ALLOWLIST"` // extra env vars to pass (extends defaults)
|
||||
EnvSet map[string]string `json:"env_set" env:"PICOCLAW_TOOLS_EXEC_ENV_SET"` // explicit var=value pairs
|
||||
|
||||
// Deprecated: these fields are ignored. See risk_threshold and risk_overrides.
|
||||
EnableDenyPatterns bool `json:"enable_deny_patterns,omitempty" env:"PICOCLAW_TOOLS_EXEC_ENABLE_DENY_PATTERNS"`
|
||||
CustomDenyPatterns []string `json:"custom_deny_patterns,omitempty" env:"PICOCLAW_TOOLS_EXEC_CUSTOM_DENY_PATTERNS"`
|
||||
CustomAllowPatterns []string `json:"custom_allow_patterns,omitempty" env:"PICOCLAW_TOOLS_EXEC_CUSTOM_ALLOW_PATTERNS"`
|
||||
}
|
||||
|
||||
type SkillsToolsConfig struct {
|
||||
|
|
|
|||
|
|
@ -382,11 +382,7 @@ func DefaultConfig() *Config {
|
|||
ExecTimeoutMinutes: 5,
|
||||
},
|
||||
Exec: ExecConfig{
|
||||
ToolConfig: ToolConfig{
|
||||
Enabled: true,
|
||||
},
|
||||
EnableDenyPatterns: true,
|
||||
TimeoutSeconds: 60,
|
||||
RiskThreshold: "medium",
|
||||
},
|
||||
Skills: SkillsToolsConfig{
|
||||
ToolConfig: ToolConfig{
|
||||
|
|
|
|||
153
pkg/tools/cron_exec_test.go
Normal file
153
pkg/tools/cron_exec_test.go
Normal file
|
|
@ -0,0 +1,153 @@
|
|||
package tools
|
||||
|
||||
import (
|
||||
"context"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/sipeed/picoclaw/pkg/bus"
|
||||
"github.com/sipeed/picoclaw/pkg/config"
|
||||
"github.com/sipeed/picoclaw/pkg/cron"
|
||||
)
|
||||
|
||||
// readOutbound reads the next outbound message with a timeout.
|
||||
func readOutbound(t *testing.T, msgBus *bus.MessageBus, timeout time.Duration) (bus.OutboundMessage, bool) {
|
||||
t.Helper()
|
||||
ctx, cancel := context.WithTimeout(context.Background(), timeout)
|
||||
defer cancel()
|
||||
msg, ok := msgBus.SubscribeOutbound(ctx)
|
||||
if !ok {
|
||||
t.Fatal("timed out or bus closed waiting for outbound message")
|
||||
}
|
||||
return msg, ok
|
||||
}
|
||||
|
||||
// stubJobExecutor satisfies JobExecutor for testing without a real agent.
|
||||
type stubJobExecutor struct{}
|
||||
|
||||
func (s *stubJobExecutor) ProcessDirectWithChannel(_ context.Context, _, _, _, _ string) (string, error) {
|
||||
return "ok", nil
|
||||
}
|
||||
|
||||
// TestCronTool_ExecuteJob_BlocksDangerousCommands verifies that commands
|
||||
// executed via the cron scheduler go through the same risk classifier
|
||||
// as agent-originated ExecTool calls (acceptance criterion AC-8).
|
||||
func TestCronTool_ExecuteJob_BlocksDangerousCommands(t *testing.T) {
|
||||
tmpDir := t.TempDir()
|
||||
storePath := filepath.Join(tmpDir, "cron.json")
|
||||
|
||||
cronSvc := cron.NewCronService(storePath, func(_ *cron.CronJob) (string, error) {
|
||||
return "ok", nil
|
||||
})
|
||||
|
||||
msgBus := bus.NewMessageBus()
|
||||
defer msgBus.Close()
|
||||
|
||||
workspace := t.TempDir()
|
||||
|
||||
ct, err := NewCronTool(
|
||||
cronSvc,
|
||||
&stubJobExecutor{},
|
||||
msgBus,
|
||||
workspace,
|
||||
true,
|
||||
5*time.Second,
|
||||
&config.Config{}, // default config → risk threshold = medium
|
||||
)
|
||||
if err != nil {
|
||||
t.Fatalf("NewCronTool: %v", err)
|
||||
}
|
||||
|
||||
tests := []struct {
|
||||
name string
|
||||
command string
|
||||
}{
|
||||
{"sudo", "sudo ls"},
|
||||
{"rm -rf", "rm -rf /"},
|
||||
{"shutdown", "shutdown -h now"},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
job := &cron.CronJob{
|
||||
ID: "test-" + tt.name,
|
||||
Name: tt.name,
|
||||
Payload: cron.CronPayload{
|
||||
Command: tt.command,
|
||||
Channel: "test",
|
||||
To: "test-chat",
|
||||
},
|
||||
}
|
||||
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
|
||||
defer cancel()
|
||||
|
||||
ct.ExecuteJob(ctx, job)
|
||||
|
||||
msg, _ := readOutbound(t, msgBus, 3*time.Second)
|
||||
if !strings.Contains(msg.Content, "Error executing scheduled command") {
|
||||
t.Errorf("expected blocked error, got: %s", msg.Content)
|
||||
}
|
||||
if !strings.Contains(msg.Content, "Command blocked by risk classifier") {
|
||||
t.Errorf("expected risk classifier message, got: %s", msg.Content)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// TestCronTool_ExecuteJob_AllowsSafeCommands verifies harmless commands
|
||||
// still execute normally through the cron path.
|
||||
func TestCronTool_ExecuteJob_AllowsSafeCommands(t *testing.T) {
|
||||
tmpDir := t.TempDir()
|
||||
storePath := filepath.Join(tmpDir, "cron.json")
|
||||
|
||||
cronSvc := cron.NewCronService(storePath, func(_ *cron.CronJob) (string, error) {
|
||||
return "ok", nil
|
||||
})
|
||||
|
||||
msgBus := bus.NewMessageBus()
|
||||
defer msgBus.Close()
|
||||
|
||||
workspace := t.TempDir()
|
||||
// Create a file to read so the command has visible output.
|
||||
os.WriteFile(filepath.Join(workspace, "hello.txt"), []byte("cron-test"), 0o644)
|
||||
|
||||
ct, err := NewCronTool(
|
||||
cronSvc,
|
||||
&stubJobExecutor{},
|
||||
msgBus,
|
||||
workspace,
|
||||
true,
|
||||
5*time.Second,
|
||||
&config.Config{},
|
||||
)
|
||||
if err != nil {
|
||||
t.Fatalf("NewCronTool: %v", err)
|
||||
}
|
||||
|
||||
job := &cron.CronJob{
|
||||
ID: "safe-echo",
|
||||
Name: "safe echo",
|
||||
Payload: cron.CronPayload{
|
||||
Command: "echo hello",
|
||||
Channel: "test",
|
||||
To: "test-chat",
|
||||
},
|
||||
}
|
||||
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
|
||||
defer cancel()
|
||||
|
||||
ct.ExecuteJob(ctx, job)
|
||||
|
||||
msg, _ := readOutbound(t, msgBus, 3*time.Second)
|
||||
if strings.Contains(msg.Content, "Error executing scheduled command") {
|
||||
t.Errorf("safe command should not be blocked: %s", msg.Content)
|
||||
}
|
||||
if !strings.Contains(msg.Content, "hello") {
|
||||
t.Errorf("expected 'hello' in output, got: %s", msg.Content)
|
||||
}
|
||||
}
|
||||
|
|
@ -1,383 +0,0 @@
|
|||
package tools
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"errors"
|
||||
"fmt"
|
||||
"os"
|
||||
"os/exec"
|
||||
"path/filepath"
|
||||
"regexp"
|
||||
"runtime"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/sipeed/picoclaw/pkg/config"
|
||||
)
|
||||
|
||||
type ExecTool struct {
|
||||
workingDir string
|
||||
timeout time.Duration
|
||||
denyPatterns []*regexp.Regexp
|
||||
allowPatterns []*regexp.Regexp
|
||||
customAllowPatterns []*regexp.Regexp
|
||||
restrictToWorkspace bool
|
||||
}
|
||||
|
||||
var (
|
||||
defaultDenyPatterns = []*regexp.Regexp{
|
||||
regexp.MustCompile(`\brm\s+-[rf]{1,2}\b`),
|
||||
regexp.MustCompile(`\bdel\s+/[fq]\b`),
|
||||
regexp.MustCompile(`\brmdir\s+/s\b`),
|
||||
// Match disk wiping commands (must be followed by space/args)
|
||||
regexp.MustCompile(
|
||||
`\b(format|mkfs|diskpart)\b\s`,
|
||||
),
|
||||
regexp.MustCompile(`\bdd\s+if=`),
|
||||
// Block writes to block devices (all common naming schemes).
|
||||
regexp.MustCompile(
|
||||
`>\s*/dev/(sd[a-z]|hd[a-z]|vd[a-z]|xvd[a-z]|nvme\d|mmcblk\d|loop\d|dm-\d|md\d|sr\d|nbd\d)`,
|
||||
),
|
||||
regexp.MustCompile(`\b(shutdown|reboot|poweroff)\b`),
|
||||
regexp.MustCompile(`:\(\)\s*\{.*\};\s*:`),
|
||||
regexp.MustCompile(`\$\([^)]+\)`),
|
||||
regexp.MustCompile(`\$\{[^}]+\}`),
|
||||
regexp.MustCompile("`[^`]+`"),
|
||||
regexp.MustCompile(`\|\s*sh\b`),
|
||||
regexp.MustCompile(`\|\s*bash\b`),
|
||||
regexp.MustCompile(`;\s*rm\s+-[rf]`),
|
||||
regexp.MustCompile(`&&\s*rm\s+-[rf]`),
|
||||
regexp.MustCompile(`\|\|\s*rm\s+-[rf]`),
|
||||
regexp.MustCompile(`<<\s*EOF`),
|
||||
regexp.MustCompile(`\$\(\s*cat\s+`),
|
||||
regexp.MustCompile(`\$\(\s*curl\s+`),
|
||||
regexp.MustCompile(`\$\(\s*wget\s+`),
|
||||
regexp.MustCompile(`\$\(\s*which\s+`),
|
||||
regexp.MustCompile(`\bsudo\b`),
|
||||
regexp.MustCompile(`\bchmod\s+[0-7]{3,4}\b`),
|
||||
regexp.MustCompile(`\bchown\b`),
|
||||
regexp.MustCompile(`\bpkill\b`),
|
||||
regexp.MustCompile(`\bkillall\b`),
|
||||
regexp.MustCompile(`\bkill\b`),
|
||||
regexp.MustCompile(`\bcurl\b.*\|\s*(sh|bash)`),
|
||||
regexp.MustCompile(`\bwget\b.*\|\s*(sh|bash)`),
|
||||
regexp.MustCompile(`\bnpm\s+install\s+-g\b`),
|
||||
regexp.MustCompile(`\bpip\s+install\s+--user\b`),
|
||||
regexp.MustCompile(`\bapt\s+(install|remove|purge)\b`),
|
||||
regexp.MustCompile(`\byum\s+(install|remove)\b`),
|
||||
regexp.MustCompile(`\bdnf\s+(install|remove)\b`),
|
||||
regexp.MustCompile(`\bdocker\s+run\b`),
|
||||
regexp.MustCompile(`\bdocker\s+exec\b`),
|
||||
regexp.MustCompile(`\bgit\s+push\b`),
|
||||
regexp.MustCompile(`\bgit\s+force\b`),
|
||||
regexp.MustCompile(`\bssh\b.*@`),
|
||||
regexp.MustCompile(`\beval\b`),
|
||||
regexp.MustCompile(`\bsource\s+.*\.sh\b`),
|
||||
}
|
||||
|
||||
// absolutePathPattern matches absolute file paths in commands (Unix and Windows).
|
||||
absolutePathPattern = regexp.MustCompile(`[A-Za-z]:\\[^\\\"']+|/[^\s\"']+`)
|
||||
|
||||
// safePaths are kernel pseudo-devices that are always safe to reference in
|
||||
// commands, regardless of workspace restriction. They contain no user data
|
||||
// and cannot cause destructive writes.
|
||||
safePaths = map[string]bool{
|
||||
"/dev/null": true,
|
||||
"/dev/zero": true,
|
||||
"/dev/random": true,
|
||||
"/dev/urandom": true,
|
||||
"/dev/stdin": true,
|
||||
"/dev/stdout": true,
|
||||
"/dev/stderr": true,
|
||||
}
|
||||
)
|
||||
|
||||
func NewExecTool(workingDir string, restrict bool) (*ExecTool, error) {
|
||||
return NewExecToolWithConfig(workingDir, restrict, nil)
|
||||
}
|
||||
|
||||
func NewExecToolWithConfig(workingDir string, restrict bool, config *config.Config) (*ExecTool, error) {
|
||||
denyPatterns := make([]*regexp.Regexp, 0)
|
||||
customAllowPatterns := make([]*regexp.Regexp, 0)
|
||||
|
||||
if config != nil {
|
||||
execConfig := config.Tools.Exec
|
||||
enableDenyPatterns := execConfig.EnableDenyPatterns
|
||||
if enableDenyPatterns {
|
||||
denyPatterns = append(denyPatterns, defaultDenyPatterns...)
|
||||
if len(execConfig.CustomDenyPatterns) > 0 {
|
||||
fmt.Printf("Using custom deny patterns: %v\n", execConfig.CustomDenyPatterns)
|
||||
for _, pattern := range execConfig.CustomDenyPatterns {
|
||||
re, err := regexp.Compile(pattern)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("invalid custom deny pattern %q: %w", pattern, err)
|
||||
}
|
||||
denyPatterns = append(denyPatterns, re)
|
||||
}
|
||||
}
|
||||
} else {
|
||||
// If deny patterns are disabled, we won't add any patterns, allowing all commands.
|
||||
fmt.Println("Warning: deny patterns are disabled. All commands will be allowed.")
|
||||
}
|
||||
for _, pattern := range execConfig.CustomAllowPatterns {
|
||||
re, err := regexp.Compile(pattern)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("invalid custom allow pattern %q: %w", pattern, err)
|
||||
}
|
||||
customAllowPatterns = append(customAllowPatterns, re)
|
||||
}
|
||||
} else {
|
||||
denyPatterns = append(denyPatterns, defaultDenyPatterns...)
|
||||
}
|
||||
|
||||
timeout := 60 * time.Second
|
||||
if config != nil && config.Tools.Exec.TimeoutSeconds > 0 {
|
||||
timeout = time.Duration(config.Tools.Exec.TimeoutSeconds) * time.Second
|
||||
}
|
||||
|
||||
return &ExecTool{
|
||||
workingDir: workingDir,
|
||||
timeout: timeout,
|
||||
denyPatterns: denyPatterns,
|
||||
allowPatterns: nil,
|
||||
customAllowPatterns: customAllowPatterns,
|
||||
restrictToWorkspace: restrict,
|
||||
}, nil
|
||||
}
|
||||
|
||||
func (t *ExecTool) Name() string {
|
||||
return "exec"
|
||||
}
|
||||
|
||||
func (t *ExecTool) Description() string {
|
||||
return "Execute a shell command and return its output. Use with caution."
|
||||
}
|
||||
|
||||
func (t *ExecTool) Parameters() map[string]any {
|
||||
return map[string]any{
|
||||
"type": "object",
|
||||
"properties": map[string]any{
|
||||
"command": map[string]any{
|
||||
"type": "string",
|
||||
"description": "The shell command to execute",
|
||||
},
|
||||
"working_dir": map[string]any{
|
||||
"type": "string",
|
||||
"description": "Optional working directory for the command",
|
||||
},
|
||||
},
|
||||
"required": []string{"command"},
|
||||
}
|
||||
}
|
||||
|
||||
func (t *ExecTool) Execute(ctx context.Context, args map[string]any) *ToolResult {
|
||||
command, ok := args["command"].(string)
|
||||
if !ok {
|
||||
return ErrorResult("command is required")
|
||||
}
|
||||
|
||||
cwd := t.workingDir
|
||||
if wd, ok := args["working_dir"].(string); ok && wd != "" {
|
||||
if t.restrictToWorkspace && t.workingDir != "" {
|
||||
resolvedWD, err := validatePath(wd, t.workingDir, true)
|
||||
if err != nil {
|
||||
return ErrorResult("Command blocked by safety guard (" + err.Error() + ")")
|
||||
}
|
||||
cwd = resolvedWD
|
||||
} else {
|
||||
cwd = wd
|
||||
}
|
||||
}
|
||||
|
||||
if cwd == "" {
|
||||
wd, err := os.Getwd()
|
||||
if err == nil {
|
||||
cwd = wd
|
||||
}
|
||||
}
|
||||
|
||||
if guardError := t.guardCommand(command, cwd); guardError != "" {
|
||||
return ErrorResult(guardError)
|
||||
}
|
||||
|
||||
// timeout == 0 means no timeout
|
||||
var cmdCtx context.Context
|
||||
var cancel context.CancelFunc
|
||||
if t.timeout > 0 {
|
||||
cmdCtx, cancel = context.WithTimeout(ctx, t.timeout)
|
||||
} else {
|
||||
cmdCtx, cancel = context.WithCancel(ctx)
|
||||
}
|
||||
defer cancel()
|
||||
|
||||
var cmd *exec.Cmd
|
||||
if runtime.GOOS == "windows" {
|
||||
cmd = exec.CommandContext(cmdCtx, "powershell", "-NoProfile", "-NonInteractive", "-Command", command)
|
||||
} else {
|
||||
cmd = exec.CommandContext(cmdCtx, "sh", "-c", command)
|
||||
}
|
||||
if cwd != "" {
|
||||
cmd.Dir = cwd
|
||||
}
|
||||
|
||||
prepareCommandForTermination(cmd)
|
||||
|
||||
var stdout, stderr bytes.Buffer
|
||||
cmd.Stdout = &stdout
|
||||
cmd.Stderr = &stderr
|
||||
|
||||
if err := cmd.Start(); err != nil {
|
||||
return ErrorResult(fmt.Sprintf("failed to start command: %v", err))
|
||||
}
|
||||
|
||||
done := make(chan error, 1)
|
||||
go func() {
|
||||
done <- cmd.Wait()
|
||||
}()
|
||||
|
||||
var err error
|
||||
select {
|
||||
case err = <-done:
|
||||
case <-cmdCtx.Done():
|
||||
_ = terminateProcessTree(cmd)
|
||||
select {
|
||||
case err = <-done:
|
||||
case <-time.After(2 * time.Second):
|
||||
if cmd.Process != nil {
|
||||
_ = cmd.Process.Kill()
|
||||
}
|
||||
err = <-done
|
||||
}
|
||||
}
|
||||
|
||||
output := stdout.String()
|
||||
if stderr.Len() > 0 {
|
||||
output += "\nSTDERR:\n" + stderr.String()
|
||||
}
|
||||
|
||||
if err != nil {
|
||||
if errors.Is(cmdCtx.Err(), context.DeadlineExceeded) {
|
||||
msg := fmt.Sprintf("Command timed out after %v", t.timeout)
|
||||
return &ToolResult{
|
||||
ForLLM: msg,
|
||||
ForUser: msg,
|
||||
IsError: true,
|
||||
}
|
||||
}
|
||||
output += fmt.Sprintf("\nExit code: %v", err)
|
||||
}
|
||||
|
||||
if output == "" {
|
||||
output = "(no output)"
|
||||
}
|
||||
|
||||
maxLen := 10000
|
||||
if len(output) > maxLen {
|
||||
output = output[:maxLen] + fmt.Sprintf("\n... (truncated, %d more chars)", len(output)-maxLen)
|
||||
}
|
||||
|
||||
if err != nil {
|
||||
return &ToolResult{
|
||||
ForLLM: output,
|
||||
ForUser: output,
|
||||
IsError: true,
|
||||
}
|
||||
}
|
||||
|
||||
return &ToolResult{
|
||||
ForLLM: output,
|
||||
ForUser: output,
|
||||
IsError: false,
|
||||
}
|
||||
}
|
||||
|
||||
func (t *ExecTool) guardCommand(command, cwd string) string {
|
||||
cmd := strings.TrimSpace(command)
|
||||
lower := strings.ToLower(cmd)
|
||||
|
||||
// Custom allow patterns exempt a command from deny checks.
|
||||
explicitlyAllowed := false
|
||||
for _, pattern := range t.customAllowPatterns {
|
||||
if pattern.MatchString(lower) {
|
||||
explicitlyAllowed = true
|
||||
break
|
||||
}
|
||||
}
|
||||
|
||||
if !explicitlyAllowed {
|
||||
for _, pattern := range t.denyPatterns {
|
||||
if pattern.MatchString(lower) {
|
||||
return "Command blocked by safety guard (dangerous pattern detected)"
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if len(t.allowPatterns) > 0 {
|
||||
allowed := false
|
||||
for _, pattern := range t.allowPatterns {
|
||||
if pattern.MatchString(lower) {
|
||||
allowed = true
|
||||
break
|
||||
}
|
||||
}
|
||||
if !allowed {
|
||||
return "Command blocked by safety guard (not in allowlist)"
|
||||
}
|
||||
}
|
||||
|
||||
if t.restrictToWorkspace {
|
||||
if strings.Contains(cmd, "..\\") || strings.Contains(cmd, "../") {
|
||||
return "Command blocked by safety guard (path traversal detected)"
|
||||
}
|
||||
|
||||
cwdPath, err := filepath.Abs(cwd)
|
||||
if err != nil {
|
||||
return ""
|
||||
}
|
||||
|
||||
matches := absolutePathPattern.FindAllString(cmd, -1)
|
||||
|
||||
for _, raw := range matches {
|
||||
p, err := filepath.Abs(raw)
|
||||
if err != nil {
|
||||
continue
|
||||
}
|
||||
|
||||
if safePaths[p] {
|
||||
continue
|
||||
}
|
||||
|
||||
rel, err := filepath.Rel(cwdPath, p)
|
||||
if err != nil {
|
||||
continue
|
||||
}
|
||||
|
||||
if strings.HasPrefix(rel, "..") {
|
||||
return "Command blocked by safety guard (path outside working dir)"
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
return ""
|
||||
}
|
||||
|
||||
func (t *ExecTool) SetTimeout(timeout time.Duration) {
|
||||
t.timeout = timeout
|
||||
}
|
||||
|
||||
func (t *ExecTool) SetRestrictToWorkspace(restrict bool) {
|
||||
t.restrictToWorkspace = restrict
|
||||
}
|
||||
|
||||
func (t *ExecTool) SetAllowPatterns(patterns []string) error {
|
||||
t.allowPatterns = make([]*regexp.Regexp, 0, len(patterns))
|
||||
for _, p := range patterns {
|
||||
re, err := regexp.Compile(p)
|
||||
if err != nil {
|
||||
return fmt.Errorf("invalid allow pattern %q: %w", p, err)
|
||||
}
|
||||
t.allowPatterns = append(t.allowPatterns, re)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
108
pkg/tools/shell/env.go
Normal file
108
pkg/tools/shell/env.go
Normal file
|
|
@ -0,0 +1,108 @@
|
|||
package shell
|
||||
|
||||
import (
|
||||
"os"
|
||||
"strings"
|
||||
|
||||
"mvdan.cc/sh/v3/expand"
|
||||
)
|
||||
|
||||
// DefaultEnvAllowlist is the set of environment variable names that are safe
|
||||
// to propagate to child processes. Everything else is stripped to prevent
|
||||
// accidental credential leakage.
|
||||
var DefaultEnvAllowlist = map[string]bool{
|
||||
"PATH": true,
|
||||
"HOME": true,
|
||||
"USER": true,
|
||||
"LANG": true,
|
||||
"SHELL": true,
|
||||
"TERM": true,
|
||||
"PWD": true,
|
||||
"OLDPWD": true,
|
||||
"HOSTNAME": true,
|
||||
"LOGNAME": true,
|
||||
"TZ": true,
|
||||
"DISPLAY": true,
|
||||
"TMPDIR": true,
|
||||
"EDITOR": true,
|
||||
"PAGER": true,
|
||||
}
|
||||
|
||||
// defaultEnvAllowPrefixes are env var prefixes that are always allowed.
|
||||
var defaultEnvAllowPrefixes = []string{
|
||||
"LC_",
|
||||
}
|
||||
|
||||
// BuildSanitizedEnv constructs an expand.Environ from the current process
|
||||
// environment, filtering to only allowlisted variables.
|
||||
//
|
||||
// extraAllowlist adds additional variable names to the default allowlist.
|
||||
// envSet provides explicit key=value pairs that override any inherited value.
|
||||
func BuildSanitizedEnv(extraAllowlist []string, envSet map[string]string) expand.Environ {
|
||||
allowed := make(map[string]bool, len(DefaultEnvAllowlist)+len(extraAllowlist))
|
||||
for k := range DefaultEnvAllowlist {
|
||||
allowed[k] = true
|
||||
}
|
||||
for _, k := range extraAllowlist {
|
||||
allowed[k] = true
|
||||
}
|
||||
|
||||
vars := make(map[string]string, len(allowed)+len(envSet))
|
||||
|
||||
for _, entry := range os.Environ() {
|
||||
k, v, ok := strings.Cut(entry, "=")
|
||||
if !ok {
|
||||
continue
|
||||
}
|
||||
if allowed[k] || isAllowedPrefix(k) {
|
||||
vars[k] = v
|
||||
}
|
||||
}
|
||||
|
||||
for k, v := range envSet {
|
||||
vars[k] = v
|
||||
}
|
||||
|
||||
return &sanitizedEnv{vars: vars}
|
||||
}
|
||||
|
||||
func isAllowedPrefix(name string) bool {
|
||||
for _, prefix := range defaultEnvAllowPrefixes {
|
||||
if strings.HasPrefix(name, prefix) {
|
||||
return true
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
// sanitizedEnv implements expand.Environ backed by a simple map.
|
||||
type sanitizedEnv struct {
|
||||
vars map[string]string
|
||||
}
|
||||
|
||||
func (e *sanitizedEnv) Get(name string) expand.Variable {
|
||||
val, ok := e.vars[name]
|
||||
if !ok {
|
||||
return expand.Variable{}
|
||||
}
|
||||
return expand.Variable{
|
||||
Set: true,
|
||||
Exported: true,
|
||||
Kind: expand.String,
|
||||
Str: val,
|
||||
}
|
||||
}
|
||||
|
||||
func (e *sanitizedEnv) Each(fn func(name string, vr expand.Variable) bool) {
|
||||
for k, v := range e.vars {
|
||||
vr := expand.Variable{
|
||||
Set: true,
|
||||
Exported: true,
|
||||
Kind: expand.String,
|
||||
Str: v,
|
||||
}
|
||||
if !fn(k, vr) {
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
112
pkg/tools/shell/env_test.go
Normal file
112
pkg/tools/shell/env_test.go
Normal file
|
|
@ -0,0 +1,112 @@
|
|||
package shell
|
||||
|
||||
import (
|
||||
"testing"
|
||||
|
||||
"mvdan.cc/sh/v3/expand"
|
||||
)
|
||||
|
||||
func TestBuildSanitizedEnv_FiltersSecrets(t *testing.T) {
|
||||
// Set some secret env vars.
|
||||
t.Setenv("OPENAI_API_KEY", "sk-secret")
|
||||
t.Setenv("ANTHROPIC_API_KEY", "anthro-secret")
|
||||
t.Setenv("AWS_SECRET_ACCESS_KEY", "aws-secret")
|
||||
t.Setenv("DATABASE_URL", "postgres://secret")
|
||||
// Set something that should pass.
|
||||
t.Setenv("PATH", "/usr/bin")
|
||||
t.Setenv("HOME", "/home/test")
|
||||
t.Setenv("LC_ALL", "en_US.UTF-8")
|
||||
|
||||
env := BuildSanitizedEnv(nil, nil)
|
||||
|
||||
assertEnvPresent(t, env, "PATH")
|
||||
assertEnvPresent(t, env, "HOME")
|
||||
assertEnvPresent(t, env, "LC_ALL")
|
||||
|
||||
assertEnvAbsent(t, env, "OPENAI_API_KEY")
|
||||
assertEnvAbsent(t, env, "ANTHROPIC_API_KEY")
|
||||
assertEnvAbsent(t, env, "AWS_SECRET_ACCESS_KEY")
|
||||
assertEnvAbsent(t, env, "DATABASE_URL")
|
||||
}
|
||||
|
||||
func TestBuildSanitizedEnv_ExtraAllowlist(t *testing.T) {
|
||||
t.Setenv("MY_CUSTOM_VAR", "hello")
|
||||
|
||||
env := BuildSanitizedEnv([]string{"MY_CUSTOM_VAR"}, nil)
|
||||
assertEnvPresent(t, env, "MY_CUSTOM_VAR")
|
||||
}
|
||||
|
||||
func TestBuildSanitizedEnv_EnvSet(t *testing.T) {
|
||||
env := BuildSanitizedEnv(nil, map[string]string{
|
||||
"INJECTED": "value123",
|
||||
})
|
||||
|
||||
v := env.Get("INJECTED")
|
||||
if !v.IsSet() {
|
||||
t.Fatal("expected INJECTED to be present")
|
||||
}
|
||||
if v.Str != "value123" {
|
||||
t.Errorf("INJECTED = %q, want %q", v.Str, "value123")
|
||||
}
|
||||
}
|
||||
|
||||
func TestBuildSanitizedEnv_EnvSetOverridesInherited(t *testing.T) {
|
||||
t.Setenv("PATH", "/original")
|
||||
|
||||
env := BuildSanitizedEnv(nil, map[string]string{
|
||||
"PATH": "/overridden",
|
||||
})
|
||||
|
||||
v := env.Get("PATH")
|
||||
if v.Str != "/overridden" {
|
||||
t.Errorf("PATH = %q, want %q", v.Str, "/overridden")
|
||||
}
|
||||
}
|
||||
|
||||
func TestBuildSanitizedEnv_DefaultAllowlist(t *testing.T) {
|
||||
for name := range DefaultEnvAllowlist {
|
||||
t.Setenv(name, "test-"+name)
|
||||
}
|
||||
|
||||
env := BuildSanitizedEnv(nil, nil)
|
||||
|
||||
for name := range DefaultEnvAllowlist {
|
||||
v := env.Get(name)
|
||||
if !v.IsSet() {
|
||||
t.Errorf("expected %s to be in sanitized env", name)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestBuildSanitizedEnv_Each(t *testing.T) {
|
||||
env := BuildSanitizedEnv(nil, map[string]string{
|
||||
"TEST_A": "a",
|
||||
"TEST_B": "b",
|
||||
})
|
||||
|
||||
found := make(map[string]bool)
|
||||
env.Each(func(name string, vr expand.Variable) bool {
|
||||
found[name] = true
|
||||
return true
|
||||
})
|
||||
|
||||
if !found["TEST_A"] || !found["TEST_B"] {
|
||||
t.Errorf("Each did not iterate over all vars: %v", found)
|
||||
}
|
||||
}
|
||||
|
||||
func assertEnvPresent(t *testing.T, env expand.Environ, name string) {
|
||||
t.Helper()
|
||||
v := env.Get(name)
|
||||
if !v.IsSet() {
|
||||
t.Errorf("expected %s to be present in sanitized env", name)
|
||||
}
|
||||
}
|
||||
|
||||
func assertEnvAbsent(t *testing.T, env expand.Environ, name string) {
|
||||
t.Helper()
|
||||
v := env.Get(name)
|
||||
if v.IsSet() {
|
||||
t.Errorf("expected %s to be absent from sanitized env, got %q", name, v.Str)
|
||||
}
|
||||
}
|
||||
435
pkg/tools/shell/risk.go
Normal file
435
pkg/tools/shell/risk.go
Normal file
|
|
@ -0,0 +1,435 @@
|
|||
package shell
|
||||
|
||||
import "fmt"
|
||||
|
||||
// RiskLevel represents the potential danger of a shell command.
|
||||
type RiskLevel int
|
||||
|
||||
const (
|
||||
RiskLow RiskLevel = iota // Read-only, informational
|
||||
RiskMedium // File modification, network read
|
||||
RiskHigh // Destructive, system-modifying
|
||||
RiskCritical // Privilege escalation, always dangerous
|
||||
)
|
||||
|
||||
func (r RiskLevel) String() string {
|
||||
switch r {
|
||||
case RiskLow:
|
||||
return "low"
|
||||
case RiskMedium:
|
||||
return "medium"
|
||||
case RiskHigh:
|
||||
return "high"
|
||||
case RiskCritical:
|
||||
return "critical"
|
||||
default:
|
||||
return "unknown"
|
||||
}
|
||||
}
|
||||
|
||||
// ParseRiskLevel converts a string to a RiskLevel.
|
||||
// Returns an error if the string is unrecognized.
|
||||
func ParseRiskLevel(s string) (RiskLevel, error) {
|
||||
switch s {
|
||||
case "low":
|
||||
return RiskLow, nil
|
||||
case "medium":
|
||||
return RiskMedium, nil
|
||||
case "high":
|
||||
return RiskHigh, nil
|
||||
case "critical":
|
||||
return RiskCritical, nil
|
||||
default:
|
||||
return RiskMedium, fmt.Errorf("unknown risk level %q, must be one of: low, medium, high, critical", s)
|
||||
}
|
||||
}
|
||||
|
||||
// commandRiskTable maps base command names to their default risk level.
|
||||
// Commands not in this table default to RiskMedium.
|
||||
var commandRiskTable = map[string]RiskLevel{
|
||||
// Low — read-only, informational
|
||||
"ls": RiskLow,
|
||||
"cat": RiskLow,
|
||||
"head": RiskLow,
|
||||
"tail": RiskLow,
|
||||
"grep": RiskLow,
|
||||
"egrep": RiskLow,
|
||||
"fgrep": RiskLow,
|
||||
"find": RiskLow,
|
||||
"wc": RiskLow,
|
||||
"echo": RiskLow,
|
||||
"printf": RiskLow,
|
||||
"pwd": RiskLow,
|
||||
"whoami": RiskLow,
|
||||
"id": RiskLow,
|
||||
"date": RiskLow,
|
||||
"uname": RiskLow,
|
||||
"hostname": RiskLow,
|
||||
"env": RiskLow,
|
||||
"printenv": RiskLow,
|
||||
"which": RiskLow,
|
||||
"type": RiskLow,
|
||||
"file": RiskLow,
|
||||
"stat": RiskLow,
|
||||
"readlink": RiskLow,
|
||||
"realpath": RiskLow,
|
||||
"basename": RiskLow,
|
||||
"dirname": RiskLow,
|
||||
"sort": RiskLow,
|
||||
"uniq": RiskLow,
|
||||
"cut": RiskLow,
|
||||
"tr": RiskLow,
|
||||
"awk": RiskLow,
|
||||
"sed": RiskLow,
|
||||
"diff": RiskLow,
|
||||
"md5sum": RiskLow,
|
||||
"sha256sum": RiskLow,
|
||||
"sha1sum": RiskLow,
|
||||
"xxd": RiskLow,
|
||||
"od": RiskLow,
|
||||
"hexdump": RiskLow,
|
||||
"strings": RiskLow,
|
||||
"tee": RiskLow,
|
||||
"xargs": RiskLow,
|
||||
"true": RiskLow,
|
||||
"false": RiskLow,
|
||||
"test": RiskLow,
|
||||
"[": RiskLow,
|
||||
"seq": RiskLow,
|
||||
"yes": RiskLow,
|
||||
"sleep": RiskLow,
|
||||
"du": RiskLow,
|
||||
"df": RiskLow,
|
||||
"free": RiskLow,
|
||||
"top": RiskLow,
|
||||
"ps": RiskLow,
|
||||
"uptime": RiskLow,
|
||||
"lsof": RiskLow,
|
||||
"tree": RiskLow,
|
||||
"less": RiskLow,
|
||||
"more": RiskLow,
|
||||
"jq": RiskLow,
|
||||
"yq": RiskLow,
|
||||
"column": RiskLow,
|
||||
"fold": RiskLow,
|
||||
"fmt": RiskLow,
|
||||
"rev": RiskLow,
|
||||
"tac": RiskLow,
|
||||
"nl": RiskLow,
|
||||
"comm": RiskLow,
|
||||
"join": RiskLow,
|
||||
"paste": RiskLow,
|
||||
"expand": RiskLow,
|
||||
"unexpand": RiskLow,
|
||||
|
||||
// Medium — file modification, network reads, build tools
|
||||
"cp": RiskMedium,
|
||||
"mv": RiskMedium,
|
||||
"mkdir": RiskMedium,
|
||||
"touch": RiskMedium,
|
||||
"ln": RiskMedium,
|
||||
"tar": RiskMedium,
|
||||
"zip": RiskMedium,
|
||||
"unzip": RiskMedium,
|
||||
"gzip": RiskMedium,
|
||||
"gunzip": RiskMedium,
|
||||
"bzip2": RiskMedium,
|
||||
"xz": RiskMedium,
|
||||
"curl": RiskMedium,
|
||||
"wget": RiskMedium,
|
||||
"git": RiskMedium,
|
||||
"make": RiskMedium,
|
||||
"go": RiskMedium,
|
||||
"python": RiskMedium,
|
||||
"python3": RiskMedium,
|
||||
"node": RiskMedium,
|
||||
"npm": RiskMedium,
|
||||
"npx": RiskMedium,
|
||||
"yarn": RiskMedium,
|
||||
"pnpm": RiskMedium,
|
||||
"pip": RiskMedium,
|
||||
"pip3": RiskMedium,
|
||||
"cargo": RiskMedium,
|
||||
"rustc": RiskMedium,
|
||||
"gcc": RiskMedium,
|
||||
"g++": RiskMedium,
|
||||
"clang": RiskMedium,
|
||||
"javac": RiskMedium,
|
||||
"java": RiskMedium,
|
||||
"ruby": RiskMedium,
|
||||
"perl": RiskMedium,
|
||||
"php": RiskMedium,
|
||||
"patch": RiskMedium,
|
||||
|
||||
// High — destructive, system-modifying
|
||||
"rm": RiskHigh,
|
||||
"rmdir": RiskHigh,
|
||||
"chmod": RiskHigh,
|
||||
"chown": RiskHigh,
|
||||
"chgrp": RiskHigh,
|
||||
"kill": RiskHigh,
|
||||
"pkill": RiskHigh,
|
||||
"killall": RiskHigh,
|
||||
"ssh": RiskHigh,
|
||||
"scp": RiskHigh,
|
||||
"rsync": RiskHigh,
|
||||
"docker": RiskHigh,
|
||||
"kubectl": RiskHigh,
|
||||
"systemctl": RiskHigh,
|
||||
"service": RiskHigh,
|
||||
|
||||
// Critical — privilege escalation, always dangerous
|
||||
"sudo": RiskCritical,
|
||||
"su": RiskCritical,
|
||||
"dd": RiskCritical,
|
||||
"mkfs": RiskCritical,
|
||||
"fdisk": RiskCritical,
|
||||
"parted": RiskCritical,
|
||||
"mount": RiskCritical,
|
||||
"umount": RiskCritical,
|
||||
"shutdown": RiskCritical,
|
||||
"reboot": RiskCritical,
|
||||
"poweroff": RiskCritical,
|
||||
"halt": RiskCritical,
|
||||
"init": RiskCritical,
|
||||
"insmod": RiskCritical,
|
||||
"rmmod": RiskCritical,
|
||||
"modprobe": RiskCritical,
|
||||
"iptables": RiskCritical,
|
||||
"nft": RiskCritical,
|
||||
"eval": RiskCritical,
|
||||
"exec": RiskCritical,
|
||||
"source": RiskCritical,
|
||||
".": RiskCritical,
|
||||
"format": RiskCritical,
|
||||
"diskpart": RiskCritical,
|
||||
}
|
||||
|
||||
// ArgModifier describes a condition that elevates a command's risk level.
|
||||
// All tokens in Args must be present in the command (order-independent, after
|
||||
// flag normalization).
|
||||
type ArgModifier struct {
|
||||
Args []string
|
||||
Level RiskLevel
|
||||
}
|
||||
|
||||
// argumentModifiers maps command names to their argument-aware risk adjustments.
|
||||
// Checked in order; first match wins.
|
||||
//
|
||||
// Patterns use individual flags (e.g., "-r", "-f") rather than combined forms
|
||||
// ("-rf") because normalizeFlags splits combined flags before matching. This
|
||||
// means "rm -rf", "rm -fr", "rm -r -f", and "rm -f -r" all match correctly.
|
||||
var argumentModifiers = map[string][]ArgModifier{
|
||||
"git": {
|
||||
{Args: []string{"push", "--force"}, Level: RiskCritical},
|
||||
{Args: []string{"push", "-f"}, Level: RiskCritical},
|
||||
{Args: []string{"push"}, Level: RiskHigh},
|
||||
{Args: []string{"reset", "--hard"}, Level: RiskHigh},
|
||||
{Args: []string{"clean", "-f", "-d"}, Level: RiskHigh},
|
||||
{Args: []string{"clean", "-f"}, Level: RiskHigh},
|
||||
},
|
||||
"curl": {
|
||||
{Args: []string{"-X", "POST"}, Level: RiskHigh},
|
||||
{Args: []string{"-X", "PUT"}, Level: RiskHigh},
|
||||
{Args: []string{"-X", "DELETE"}, Level: RiskHigh},
|
||||
{Args: []string{"--request", "POST"}, Level: RiskHigh},
|
||||
{Args: []string{"--request", "PUT"}, Level: RiskHigh},
|
||||
{Args: []string{"--request", "DELETE"}, Level: RiskHigh},
|
||||
{Args: []string{"--data"}, Level: RiskHigh},
|
||||
{Args: []string{"-d"}, Level: RiskHigh},
|
||||
{Args: []string{"--upload-file"}, Level: RiskHigh},
|
||||
{Args: []string{"-T"}, Level: RiskHigh},
|
||||
},
|
||||
"wget": {
|
||||
{Args: []string{"--post-data"}, Level: RiskHigh},
|
||||
{Args: []string{"--post-file"}, Level: RiskHigh},
|
||||
},
|
||||
"npm": {
|
||||
{Args: []string{"install", "-g"}, Level: RiskHigh},
|
||||
{Args: []string{"install", "--global"}, Level: RiskHigh},
|
||||
{Args: []string{"publish"}, Level: RiskHigh},
|
||||
},
|
||||
"pip": {
|
||||
{Args: []string{"install", "--user"}, Level: RiskHigh},
|
||||
},
|
||||
"pip3": {
|
||||
{Args: []string{"install", "--user"}, Level: RiskHigh},
|
||||
},
|
||||
"docker": {
|
||||
{Args: []string{"run"}, Level: RiskHigh},
|
||||
{Args: []string{"exec"}, Level: RiskHigh},
|
||||
{Args: []string{"rm"}, Level: RiskHigh},
|
||||
{Args: []string{"rmi"}, Level: RiskHigh},
|
||||
},
|
||||
"apt": {
|
||||
{Args: []string{"install"}, Level: RiskHigh},
|
||||
{Args: []string{"remove"}, Level: RiskHigh},
|
||||
{Args: []string{"purge"}, Level: RiskCritical},
|
||||
},
|
||||
"apt-get": {
|
||||
{Args: []string{"install"}, Level: RiskHigh},
|
||||
{Args: []string{"remove"}, Level: RiskHigh},
|
||||
{Args: []string{"purge"}, Level: RiskCritical},
|
||||
},
|
||||
"yum": {
|
||||
{Args: []string{"install"}, Level: RiskHigh},
|
||||
{Args: []string{"remove"}, Level: RiskHigh},
|
||||
},
|
||||
"dnf": {
|
||||
{Args: []string{"install"}, Level: RiskHigh},
|
||||
{Args: []string{"remove"}, Level: RiskHigh},
|
||||
},
|
||||
"rm": {
|
||||
{Args: []string{"-r", "-f"}, Level: RiskCritical},
|
||||
},
|
||||
"kill": {
|
||||
{Args: []string{"-9"}, Level: RiskCritical},
|
||||
{Args: []string{"-KILL"}, Level: RiskCritical},
|
||||
{Args: []string{"-SIGKILL"}, Level: RiskCritical},
|
||||
},
|
||||
}
|
||||
|
||||
// ClassifyCommand determines the risk level of a resolved command.
|
||||
// args[0] is the command name (basename), args[1:] are the arguments.
|
||||
// overrides allows per-command level overrides from config.
|
||||
// extraModifiers are checked after built-in argumentModifiers.
|
||||
// The highest matching level across all sources wins.
|
||||
func ClassifyCommand(args []string, overrides map[string]string, extraModifiers ...map[string][]ArgModifier) RiskLevel {
|
||||
if len(args) == 0 {
|
||||
return RiskMedium
|
||||
}
|
||||
|
||||
cmdName := baseCommand(args[0])
|
||||
|
||||
if overrides != nil {
|
||||
if levelStr, ok := overrides[cmdName]; ok {
|
||||
level, err := ParseRiskLevel(levelStr)
|
||||
if err == nil {
|
||||
return level
|
||||
}
|
||||
// Invalid risk level in override: fall through to default classification.
|
||||
// The parse error is descriptive, but we can't log from here without
|
||||
// injecting a logger. Config-time validation catches user errors.
|
||||
}
|
||||
}
|
||||
|
||||
level, known := commandRiskTable[cmdName]
|
||||
if !known {
|
||||
level = RiskMedium
|
||||
}
|
||||
|
||||
// Normalize args: expand combined short flags like -rf → -r, -f
|
||||
// so that modifiers match regardless of how flags were grouped or ordered.
|
||||
normalizedArgs := normalizeFlags(args[1:])
|
||||
|
||||
// Check built-in modifiers, then user-supplied. Keep the highest match.
|
||||
if elevated, ok := applyModifiers(normalizedArgs, cmdName, level, argumentModifiers); ok {
|
||||
level = elevated
|
||||
}
|
||||
for _, extra := range extraModifiers {
|
||||
if extra == nil {
|
||||
continue
|
||||
}
|
||||
if elevated, ok := applyModifiers(normalizedArgs, cmdName, level, extra); ok {
|
||||
level = elevated
|
||||
}
|
||||
}
|
||||
|
||||
return level
|
||||
}
|
||||
|
||||
// applyModifiers checks whether any modifier for cmdName matches the
|
||||
// normalised args and would elevate the risk. Returns (newLevel, true)
|
||||
// on first match, or (0, false) if nothing matched.
|
||||
func applyModifiers(
|
||||
normalizedArgs []string,
|
||||
cmdName string,
|
||||
baseLevel RiskLevel,
|
||||
mods map[string][]ArgModifier,
|
||||
) (RiskLevel, bool) {
|
||||
if entries, ok := mods[cmdName]; ok {
|
||||
for _, mod := range entries {
|
||||
if matchArgs(normalizedArgs, mod.Args) && mod.Level > baseLevel {
|
||||
return mod.Level, true
|
||||
}
|
||||
}
|
||||
}
|
||||
return 0, false
|
||||
}
|
||||
|
||||
// IsAllowed returns true if the given risk level is at or below the threshold.
|
||||
func IsAllowed(level, threshold RiskLevel) bool {
|
||||
return level <= threshold
|
||||
}
|
||||
|
||||
// BlockedCommandError formats a structured error message for the LLM.
|
||||
func BlockedCommandError(args []string, level, threshold RiskLevel, reason string) string {
|
||||
cmd := ""
|
||||
if len(args) > 0 {
|
||||
cmd = args[0]
|
||||
if len(args) > 1 {
|
||||
end := len(args)
|
||||
if end > 5 {
|
||||
end = 5
|
||||
}
|
||||
for _, a := range args[1:end] {
|
||||
cmd += " " + a
|
||||
}
|
||||
if len(args) > 5 {
|
||||
cmd += " ..."
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
return fmt.Sprintf(
|
||||
"Command blocked by risk classifier: command=%q risk_level=%s threshold=%s reason=%s",
|
||||
cmd, level, threshold, reason,
|
||||
)
|
||||
}
|
||||
|
||||
// baseCommand extracts the basename from a command path.
|
||||
func baseCommand(cmd string) string {
|
||||
for i := len(cmd) - 1; i >= 0; i-- {
|
||||
if cmd[i] == '/' {
|
||||
return cmd[i+1:]
|
||||
}
|
||||
}
|
||||
return cmd
|
||||
}
|
||||
|
||||
// normalizeFlags expands combined short flags (e.g., "-rf" → "-r", "-f")
|
||||
// so that modifier matching works regardless of how flags are grouped.
|
||||
// Long flags (--flag) and non-flag arguments are passed through unchanged.
|
||||
func normalizeFlags(args []string) []string {
|
||||
result := make([]string, 0, len(args)*2)
|
||||
for _, a := range args {
|
||||
if len(a) > 2 && a[0] == '-' && a[1] != '-' {
|
||||
for _, ch := range a[1:] {
|
||||
result = append(result, "-"+string(ch))
|
||||
}
|
||||
} else {
|
||||
result = append(result, a)
|
||||
}
|
||||
}
|
||||
return result
|
||||
}
|
||||
|
||||
// matchArgs checks if ALL pattern tokens are present in args (order-independent).
|
||||
// "git push -x -f" matches pattern ["push", "-f"] because both tokens exist.
|
||||
// "git -f push" also matches ["push", "-f"]. Order does not matter.
|
||||
func matchArgs(args, pattern []string) bool {
|
||||
if len(pattern) == 0 {
|
||||
return true
|
||||
}
|
||||
argSet := make(map[string]int, len(args))
|
||||
for _, a := range args {
|
||||
argSet[a]++
|
||||
}
|
||||
for _, p := range pattern {
|
||||
if argSet[p] <= 0 {
|
||||
return false
|
||||
}
|
||||
argSet[p]--
|
||||
}
|
||||
return true
|
||||
}
|
||||
297
pkg/tools/shell/risk_test.go
Normal file
297
pkg/tools/shell/risk_test.go
Normal file
|
|
@ -0,0 +1,297 @@
|
|||
package shell
|
||||
|
||||
import "testing"
|
||||
|
||||
func TestClassifyCommand_BaseTable(t *testing.T) {
|
||||
tests := []struct {
|
||||
args []string
|
||||
level RiskLevel
|
||||
}{
|
||||
{[]string{"ls", "-la"}, RiskLow},
|
||||
{[]string{"cat", "file.txt"}, RiskLow},
|
||||
{[]string{"grep", "-r", "pattern", "."}, RiskLow},
|
||||
{[]string{"find", ".", "-name", "*.go"}, RiskLow},
|
||||
{[]string{"wc", "-l"}, RiskLow},
|
||||
{[]string{"echo", "hello"}, RiskLow},
|
||||
{[]string{"jq", ".field", "data.json"}, RiskLow},
|
||||
|
||||
{[]string{"cp", "a", "b"}, RiskMedium},
|
||||
{[]string{"mv", "a", "b"}, RiskMedium},
|
||||
{[]string{"python3", "-c", "print(1)"}, RiskMedium},
|
||||
{[]string{"git", "status"}, RiskMedium},
|
||||
{[]string{"curl", "https://example.com"}, RiskMedium},
|
||||
|
||||
{[]string{"rm", "file.txt"}, RiskHigh},
|
||||
{[]string{"chmod", "755", "script.sh"}, RiskHigh},
|
||||
{[]string{"docker", "ps"}, RiskHigh},
|
||||
{[]string{"ssh", "user@host"}, RiskHigh},
|
||||
|
||||
{[]string{"sudo", "ls"}, RiskCritical},
|
||||
{[]string{"dd", "if=/dev/zero", "of=/dev/sda"}, RiskCritical},
|
||||
{[]string{"shutdown", "-h", "now"}, RiskCritical},
|
||||
{[]string{"eval", "echo hi"}, RiskCritical},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
got := ClassifyCommand(tt.args, nil)
|
||||
if got != tt.level {
|
||||
t.Errorf("ClassifyCommand(%v) = %s, want %s", tt.args, got, tt.level)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestClassifyCommand_ArgumentModifiers(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
args []string
|
||||
level RiskLevel
|
||||
}{
|
||||
{"git push", []string{"git", "push", "origin", "main"}, RiskHigh},
|
||||
{"git push --force", []string{"git", "push", "--force", "origin"}, RiskCritical},
|
||||
{"git push -f", []string{"git", "push", "-f"}, RiskCritical},
|
||||
{"git -f push (reordered)", []string{"git", "-f", "push"}, RiskCritical},
|
||||
{"git push -x -f (extra flags)", []string{"git", "push", "-x", "-f"}, RiskCritical},
|
||||
{"git reset --hard", []string{"git", "reset", "--hard", "HEAD~1"}, RiskHigh},
|
||||
{"git clean -f", []string{"git", "clean", "-f"}, RiskHigh},
|
||||
{"git clean -fd", []string{"git", "clean", "-fd"}, RiskHigh},
|
||||
{"git clean -d -f (reordered)", []string{"git", "clean", "-d", "-f"}, RiskHigh},
|
||||
|
||||
{"curl GET (default)", []string{"curl", "https://example.com"}, RiskMedium},
|
||||
{"curl POST", []string{"curl", "-X", "POST", "https://example.com"}, RiskHigh},
|
||||
{"curl -d data", []string{"curl", "-d", "data", "https://example.com"}, RiskHigh},
|
||||
{"curl --data data", []string{"curl", "--data", "data", "url"}, RiskHigh},
|
||||
{"curl -X DELETE", []string{"curl", "-X", "DELETE", "url"}, RiskHigh},
|
||||
{"curl --request POST", []string{"curl", "--request", "POST", "url"}, RiskHigh},
|
||||
|
||||
{"rm file (no flags)", []string{"rm", "file.txt"}, RiskHigh},
|
||||
{"rm -rf", []string{"rm", "-rf", "/"}, RiskCritical},
|
||||
{"rm -fr", []string{"rm", "-fr", "/"}, RiskCritical},
|
||||
{"rm -r -f (separate)", []string{"rm", "-r", "-f", "dir"}, RiskCritical},
|
||||
{"rm -f -r (reordered)", []string{"rm", "-f", "-r", "dir"}, RiskCritical},
|
||||
{"rm -f -x -r (extra flags between)", []string{"rm", "-f", "-x", "-r", "dir"}, RiskCritical},
|
||||
|
||||
{"kill (no signal)", []string{"kill", "1234"}, RiskHigh},
|
||||
{"kill -9", []string{"kill", "-9", "1234"}, RiskCritical},
|
||||
|
||||
{"npm install (local)", []string{"npm", "install", "lodash"}, RiskMedium},
|
||||
{"npm install -g", []string{"npm", "install", "-g", "lodash"}, RiskHigh},
|
||||
{"npm publish", []string{"npm", "publish"}, RiskHigh},
|
||||
|
||||
{"docker ps (no modifier)", []string{"docker", "ps"}, RiskHigh},
|
||||
{"docker run", []string{"docker", "run", "ubuntu"}, RiskHigh},
|
||||
{"docker exec", []string{"docker", "exec", "-it", "cnt", "bash"}, RiskHigh},
|
||||
|
||||
{"apt install", []string{"apt", "install", "vim"}, RiskHigh},
|
||||
{"apt purge", []string{"apt", "purge", "vim"}, RiskCritical},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
got := ClassifyCommand(tt.args, nil)
|
||||
if got != tt.level {
|
||||
t.Errorf("ClassifyCommand(%v) = %s, want %s", tt.args, got, tt.level)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestClassifyCommand_Overrides(t *testing.T) {
|
||||
overrides := map[string]string{
|
||||
"rm": "low",
|
||||
"curl": "critical",
|
||||
}
|
||||
|
||||
got := ClassifyCommand([]string{"rm", "-rf", "/"}, overrides)
|
||||
if got != RiskLow {
|
||||
t.Errorf("override rm to low: got %s", got)
|
||||
}
|
||||
|
||||
got = ClassifyCommand([]string{"curl", "https://example.com"}, overrides)
|
||||
if got != RiskCritical {
|
||||
t.Errorf("override curl to critical: got %s", got)
|
||||
}
|
||||
|
||||
got = ClassifyCommand([]string{"ls"}, overrides)
|
||||
if got != RiskLow {
|
||||
t.Errorf("ls (no override) should be low: got %s", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestClassifyCommand_UnknownCommand(t *testing.T) {
|
||||
got := ClassifyCommand([]string{"some_unknown_tool", "--flag"}, nil)
|
||||
if got != RiskMedium {
|
||||
t.Errorf("unknown command should default to medium, got %s", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestClassifyCommand_FullPath(t *testing.T) {
|
||||
got := ClassifyCommand([]string{"/usr/bin/rm", "-rf", "/"}, nil)
|
||||
if got != RiskCritical {
|
||||
t.Errorf("/usr/bin/rm -rf should be critical, got %s", got)
|
||||
}
|
||||
|
||||
got = ClassifyCommand([]string{"/bin/ls", "-la"}, nil)
|
||||
if got != RiskLow {
|
||||
t.Errorf("/bin/ls should be low, got %s", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestIsAllowed(t *testing.T) {
|
||||
tests := []struct {
|
||||
level RiskLevel
|
||||
threshold RiskLevel
|
||||
allowed bool
|
||||
}{
|
||||
{RiskLow, RiskMedium, true},
|
||||
{RiskMedium, RiskMedium, true},
|
||||
{RiskHigh, RiskMedium, false},
|
||||
{RiskCritical, RiskMedium, false},
|
||||
{RiskCritical, RiskCritical, true},
|
||||
{RiskLow, RiskLow, true},
|
||||
{RiskMedium, RiskLow, false},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
got := IsAllowed(tt.level, tt.threshold)
|
||||
if got != tt.allowed {
|
||||
t.Errorf("IsAllowed(%s, %s) = %v, want %v", tt.level, tt.threshold, got, tt.allowed)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestParseRiskLevel(t *testing.T) {
|
||||
tests := []struct {
|
||||
input string
|
||||
want RiskLevel
|
||||
wantErr bool
|
||||
}{
|
||||
{"low", RiskLow, false},
|
||||
{"medium", RiskMedium, false},
|
||||
{"high", RiskHigh, false},
|
||||
{"critical", RiskCritical, false},
|
||||
{"bogus", RiskMedium, true},
|
||||
{"", RiskMedium, true},
|
||||
{"invalid", RiskMedium, true},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
got, err := ParseRiskLevel(tt.input)
|
||||
if (err != nil) != tt.wantErr {
|
||||
t.Errorf("ParseRiskLevel(%q) error = %v, wantErr %v", tt.input, err, tt.wantErr)
|
||||
continue
|
||||
}
|
||||
if !tt.wantErr && got != tt.want {
|
||||
t.Errorf("ParseRiskLevel(%q) = %s, want %s", tt.input, got, tt.want)
|
||||
}
|
||||
// For error cases, verify the default is RiskMedium
|
||||
if tt.wantErr && got != RiskMedium {
|
||||
t.Errorf("ParseRiskLevel(%q) = %s on error, should default to RiskMedium", tt.input, got)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestNormalizeFlags(t *testing.T) {
|
||||
tests := []struct {
|
||||
input []string
|
||||
want []string
|
||||
}{
|
||||
{[]string{"-rf"}, []string{"-r", "-f"}},
|
||||
{[]string{"-fr"}, []string{"-f", "-r"}},
|
||||
{[]string{"-r", "-f"}, []string{"-r", "-f"}},
|
||||
{[]string{"--force"}, []string{"--force"}},
|
||||
{[]string{"-f"}, []string{"-f"}},
|
||||
{[]string{"push"}, []string{"push"}},
|
||||
{[]string{"-rf", "dir", "--verbose"}, []string{"-r", "-f", "dir", "--verbose"}},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
got := normalizeFlags(tt.input)
|
||||
if len(got) != len(tt.want) {
|
||||
t.Errorf("normalizeFlags(%v) = %v, want %v", tt.input, got, tt.want)
|
||||
continue
|
||||
}
|
||||
for i := range got {
|
||||
if got[i] != tt.want[i] {
|
||||
t.Errorf("normalizeFlags(%v) = %v, want %v", tt.input, got, tt.want)
|
||||
break
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestMatchArgs_OrderIndependent(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
args []string
|
||||
pattern []string
|
||||
match bool
|
||||
}{
|
||||
{"exact match", []string{"push", "-f"}, []string{"push", "-f"}, true},
|
||||
{"reversed order", []string{"-f", "push"}, []string{"push", "-f"}, true},
|
||||
{"extra flags between", []string{"push", "-x", "-f"}, []string{"push", "-f"}, true},
|
||||
{"missing token", []string{"push", "-x"}, []string{"push", "-f"}, false},
|
||||
{"empty pattern", []string{"push"}, []string{}, true},
|
||||
{"empty args", []string{}, []string{"push"}, false},
|
||||
{"superset ok", []string{"push", "-f", "--verbose", "origin"}, []string{"push", "-f"}, true},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
got := matchArgs(tt.args, tt.pattern)
|
||||
if got != tt.match {
|
||||
t.Errorf("matchArgs(%v, %v) = %v, want %v", tt.args, tt.pattern, got, tt.match)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestClassifyCommand_ExtraArgModifiers(t *testing.T) {
|
||||
extra := map[string][]ArgModifier{
|
||||
// Custom: "make deploy" should be critical
|
||||
"make": {
|
||||
{Args: []string{"deploy"}, Level: RiskCritical},
|
||||
},
|
||||
// Custom: "git push --mirror" should be critical (not in built-ins)
|
||||
"git": {
|
||||
{Args: []string{"push", "--mirror"}, Level: RiskCritical},
|
||||
},
|
||||
}
|
||||
|
||||
tests := []struct {
|
||||
name string
|
||||
args []string
|
||||
level RiskLevel
|
||||
}{
|
||||
{"make deploy (extra)", []string{"make", "deploy"}, RiskCritical},
|
||||
{"make build (no extra)", []string{"make", "build"}, RiskMedium},
|
||||
{"git push --mirror (extra)", []string{"git", "push", "--mirror"}, RiskCritical},
|
||||
// Built-in still works: git push -f is critical even without extra
|
||||
{"git push -f (built-in)", []string{"git", "push", "-f"}, RiskCritical},
|
||||
// Extra doesn't override a higher built-in: git push --force is already critical
|
||||
{"git push --force (built-in wins)", []string{"git", "push", "--force"}, RiskCritical},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
got := ClassifyCommand(tt.args, nil, extra)
|
||||
if got != tt.level {
|
||||
t.Errorf("ClassifyCommand(%v, nil, extra) = %s, want %s", tt.args, got, tt.level)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestClassifyCommand_ExtraArgModifiers_NoOverrideBuiltIn(t *testing.T) {
|
||||
// Extra modifier tries to set "rm -rf" to medium, but built-in already
|
||||
// elevates to critical and built-in is checked first.
|
||||
extra := map[string][]ArgModifier{
|
||||
"rm": {
|
||||
{Args: []string{"-r", "-f"}, Level: RiskMedium},
|
||||
},
|
||||
}
|
||||
|
||||
got := ClassifyCommand([]string{"rm", "-rf", "/"}, nil, extra)
|
||||
if got != RiskCritical {
|
||||
t.Errorf("built-in should win over extra for rm -rf: got %s", got)
|
||||
}
|
||||
}
|
||||
210
pkg/tools/shell/runner.go
Normal file
210
pkg/tools/shell/runner.go
Normal file
|
|
@ -0,0 +1,210 @@
|
|||
package shell
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"fmt"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"mvdan.cc/sh/v3/expand"
|
||||
"mvdan.cc/sh/v3/interp"
|
||||
"mvdan.cc/sh/v3/syntax"
|
||||
)
|
||||
|
||||
// RunConfig holds all parameters for a single command execution.
|
||||
type RunConfig struct {
|
||||
Command string
|
||||
Dir string
|
||||
Timeout time.Duration
|
||||
Restrict bool
|
||||
WorkspaceDir string
|
||||
|
||||
RiskThreshold RiskLevel
|
||||
RiskOverrides map[string]string
|
||||
ExtraArgModifiers map[string][]ArgModifier // user-defined, appended after built-ins
|
||||
EnvAllowlist []string
|
||||
EnvSet map[string]string
|
||||
}
|
||||
|
||||
// RunResult contains the output of a command execution.
|
||||
type RunResult struct {
|
||||
Output string
|
||||
IsError bool
|
||||
}
|
||||
|
||||
// Run parses and executes a shell command using the in-process interpreter.
|
||||
func Run(ctx context.Context, cfg RunConfig) RunResult {
|
||||
parser := syntax.NewParser(syntax.Variant(syntax.LangBash))
|
||||
prog, err := parser.Parse(strings.NewReader(cfg.Command), "")
|
||||
if err != nil {
|
||||
return RunResult{
|
||||
Output: fmt.Sprintf("Failed to parse command: %v", err),
|
||||
IsError: true,
|
||||
}
|
||||
}
|
||||
|
||||
var runCtx context.Context
|
||||
var cancel context.CancelFunc
|
||||
if cfg.Timeout > 0 {
|
||||
runCtx, cancel = context.WithTimeout(ctx, cfg.Timeout)
|
||||
} else {
|
||||
runCtx, cancel = context.WithCancel(ctx)
|
||||
}
|
||||
defer cancel()
|
||||
|
||||
env := BuildSanitizedEnv(cfg.EnvAllowlist, cfg.EnvSet)
|
||||
|
||||
var stdout, stderr bytes.Buffer
|
||||
|
||||
opts := []interp.RunnerOption{
|
||||
interp.Env(env),
|
||||
interp.StdIO(nil, &stdout, &stderr),
|
||||
}
|
||||
|
||||
if cfg.Dir != "" {
|
||||
opts = append(opts, interp.Dir(cfg.Dir))
|
||||
}
|
||||
|
||||
opts = append(
|
||||
opts,
|
||||
interp.ExecHandlers(
|
||||
pathAwareExecHandler(env),
|
||||
riskExecHandler(cfg.RiskThreshold, cfg.RiskOverrides, cfg.ExtraArgModifiers),
|
||||
),
|
||||
)
|
||||
|
||||
if cfg.Restrict && cfg.WorkspaceDir != "" {
|
||||
opts = append(opts, interp.OpenHandler(SandboxedOpenHandler(cfg.WorkspaceDir)))
|
||||
}
|
||||
|
||||
runner, err := interp.New(opts...)
|
||||
if err != nil {
|
||||
return RunResult{
|
||||
Output: fmt.Sprintf("Failed to create interpreter: %v", err),
|
||||
IsError: true,
|
||||
}
|
||||
}
|
||||
|
||||
err = runner.Run(runCtx, prog)
|
||||
|
||||
output := stdout.String()
|
||||
if stderr.Len() > 0 {
|
||||
output += "\nSTDERR:\n" + stderr.String()
|
||||
}
|
||||
|
||||
if err != nil {
|
||||
if runCtx.Err() == context.DeadlineExceeded {
|
||||
msg := fmt.Sprintf("Command timed out after %v", cfg.Timeout)
|
||||
return RunResult{Output: msg, IsError: true}
|
||||
}
|
||||
|
||||
errStr := err.Error()
|
||||
if strings.Contains(errStr, "Command blocked by risk classifier") {
|
||||
return RunResult{Output: errStr, IsError: true}
|
||||
}
|
||||
|
||||
output += fmt.Sprintf("\nExit code: %v", err)
|
||||
}
|
||||
|
||||
if output == "" {
|
||||
output = "(no output)"
|
||||
}
|
||||
|
||||
maxLen := 10000
|
||||
if len(output) > maxLen {
|
||||
output = output[:maxLen] + fmt.Sprintf("\n... (truncated, %d more chars)", len(output)-maxLen)
|
||||
}
|
||||
|
||||
return RunResult{
|
||||
Output: output,
|
||||
IsError: err != nil,
|
||||
}
|
||||
}
|
||||
|
||||
// riskExecHandler returns an ExecHandlers middleware that classifies each resolved
|
||||
// command against the risk table before delegating to the default handler.
|
||||
func riskExecHandler(
|
||||
threshold RiskLevel,
|
||||
overrides map[string]string,
|
||||
extraMods map[string][]ArgModifier,
|
||||
) func(next interp.ExecHandlerFunc) interp.ExecHandlerFunc {
|
||||
return func(next interp.ExecHandlerFunc) interp.ExecHandlerFunc {
|
||||
return func(ctx context.Context, args []string) error {
|
||||
if len(args) == 0 {
|
||||
return next(ctx, args)
|
||||
}
|
||||
|
||||
level := ClassifyCommand(args, overrides, extraMods)
|
||||
if !IsAllowed(level, threshold) {
|
||||
reason := "command risk exceeds configured threshold"
|
||||
return fmt.Errorf("%s", BlockedCommandError(args, level, threshold, reason))
|
||||
}
|
||||
|
||||
return next(ctx, args)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// pathAwareExecHandler returns an ExecHandlerFunc that looks up commands
|
||||
// using the interpreter's environment (not os.Getenv). This ensures that
|
||||
// PATH from the sanitized environment is used for command resolution.
|
||||
func pathAwareExecHandler(env expand.Environ) func(next interp.ExecHandlerFunc) interp.ExecHandlerFunc {
|
||||
return func(next interp.ExecHandlerFunc) interp.ExecHandlerFunc {
|
||||
return func(ctx context.Context, args []string) error {
|
||||
if len(args) == 0 {
|
||||
return next(ctx, args)
|
||||
}
|
||||
|
||||
// Look up command using interpreter's PATH
|
||||
path, err := lookPath(env, args[0])
|
||||
if err != nil {
|
||||
// Command not found in PATH, try the default handler
|
||||
return next(ctx, args)
|
||||
}
|
||||
|
||||
// Replace command with full path and continue
|
||||
fullArgs := append([]string{path}, args[1:]...)
|
||||
return next(ctx, fullArgs)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// lookPath searches for an executable named cmd in the directories
|
||||
// listed in the PATH variable from the given environment.
|
||||
func lookPath(env expand.Environ, cmd string) (string, error) {
|
||||
// If command contains a slash, it's a path - return as-is
|
||||
if strings.Contains(cmd, "/") {
|
||||
return cmd, nil
|
||||
}
|
||||
|
||||
// Get PATH from interpreter environment
|
||||
pathVar := env.Get("PATH")
|
||||
if !pathVar.Set {
|
||||
// PATH not set in environment - let default handler try
|
||||
return "", fmt.Errorf("PATH not set")
|
||||
}
|
||||
if pathVar.Str == "" {
|
||||
return "", fmt.Errorf("PATH is empty")
|
||||
}
|
||||
|
||||
// Search each directory in PATH
|
||||
for _, dir := range filepath.SplitList(pathVar.Str) {
|
||||
if dir == "" {
|
||||
dir = "."
|
||||
}
|
||||
fullPath := filepath.Join(dir, cmd)
|
||||
// Check if file exists and is executable
|
||||
if stat, err := os.Stat(fullPath); err == nil && !stat.IsDir() {
|
||||
// On Unix, check executable bit
|
||||
if stat.Mode()&0o111 != 0 {
|
||||
return fullPath, nil
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Not found - let the default handler handle it
|
||||
return "", fmt.Errorf("command %q not found in PATH", cmd)
|
||||
}
|
||||
243
pkg/tools/shell/runner_test.go
Normal file
243
pkg/tools/shell/runner_test.go
Normal file
|
|
@ -0,0 +1,243 @@
|
|||
package shell
|
||||
|
||||
import (
|
||||
"context"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
|
||||
func TestRun_Success(t *testing.T) {
|
||||
result := Run(context.Background(), RunConfig{
|
||||
Command: "echo 'hello world'",
|
||||
Dir: t.TempDir(),
|
||||
Timeout: 5 * time.Second,
|
||||
RiskThreshold: RiskMedium,
|
||||
})
|
||||
|
||||
if result.IsError {
|
||||
t.Fatalf("expected success, got error: %s", result.Output)
|
||||
}
|
||||
if !strings.Contains(result.Output, "hello world") {
|
||||
t.Errorf("expected 'hello world' in output, got: %s", result.Output)
|
||||
}
|
||||
}
|
||||
|
||||
func TestRun_BlocksDangerousCommand(t *testing.T) {
|
||||
result := Run(context.Background(), RunConfig{
|
||||
Command: "rm -rf /",
|
||||
Dir: t.TempDir(),
|
||||
Timeout: 5 * time.Second,
|
||||
RiskThreshold: RiskMedium,
|
||||
})
|
||||
|
||||
if !result.IsError {
|
||||
t.Fatal("expected rm -rf to be blocked")
|
||||
}
|
||||
if !strings.Contains(result.Output, "blocked") {
|
||||
t.Errorf("expected 'blocked' in output: %s", result.Output)
|
||||
}
|
||||
}
|
||||
|
||||
func TestRun_BlocksSudo(t *testing.T) {
|
||||
result := Run(context.Background(), RunConfig{
|
||||
Command: "sudo ls",
|
||||
Dir: t.TempDir(),
|
||||
Timeout: 5 * time.Second,
|
||||
RiskThreshold: RiskMedium,
|
||||
})
|
||||
|
||||
if !result.IsError {
|
||||
t.Fatal("expected sudo to be blocked")
|
||||
}
|
||||
if !strings.Contains(result.Output, "risk_level=critical") {
|
||||
t.Errorf("expected risk_level=critical in output: %s", result.Output)
|
||||
}
|
||||
}
|
||||
|
||||
func TestRun_BlocksVariableIndirection(t *testing.T) {
|
||||
// x=rm; $x -rf / — the old regex system couldn't catch this.
|
||||
result := Run(context.Background(), RunConfig{
|
||||
Command: `x=rm; $x -rf /`,
|
||||
Dir: t.TempDir(),
|
||||
Timeout: 5 * time.Second,
|
||||
RiskThreshold: RiskMedium,
|
||||
})
|
||||
|
||||
if !result.IsError {
|
||||
t.Fatal("expected variable indirection bypass to be blocked")
|
||||
}
|
||||
if !strings.Contains(result.Output, "blocked") {
|
||||
t.Errorf("expected 'blocked' in output: %s", result.Output)
|
||||
}
|
||||
}
|
||||
|
||||
func TestRun_BlocksCommandSubstitution(t *testing.T) {
|
||||
result := Run(context.Background(), RunConfig{
|
||||
Command: "$(echo rm) -rf /tmp/safe_test_dir",
|
||||
Dir: t.TempDir(),
|
||||
Timeout: 5 * time.Second,
|
||||
RiskThreshold: RiskMedium,
|
||||
})
|
||||
|
||||
if !result.IsError {
|
||||
t.Fatal("expected command substitution bypass to be blocked")
|
||||
}
|
||||
}
|
||||
|
||||
func TestRun_Timeout(t *testing.T) {
|
||||
result := Run(context.Background(), RunConfig{
|
||||
Command: "sleep 60",
|
||||
Dir: t.TempDir(),
|
||||
Timeout: 200 * time.Millisecond,
|
||||
RiskThreshold: RiskMedium,
|
||||
})
|
||||
|
||||
if !result.IsError {
|
||||
t.Fatal("expected timeout error")
|
||||
}
|
||||
if !strings.Contains(result.Output, "timed out") {
|
||||
t.Errorf("expected 'timed out' in output: %s", result.Output)
|
||||
}
|
||||
}
|
||||
|
||||
func TestRun_WorkingDir(t *testing.T) {
|
||||
tmpDir := t.TempDir()
|
||||
testFile := filepath.Join(tmpDir, "test.txt")
|
||||
os.WriteFile(testFile, []byte("test content"), 0o644)
|
||||
|
||||
result := Run(context.Background(), RunConfig{
|
||||
Command: "cat test.txt",
|
||||
Dir: tmpDir,
|
||||
Timeout: 5 * time.Second,
|
||||
RiskThreshold: RiskMedium,
|
||||
})
|
||||
|
||||
if result.IsError {
|
||||
t.Fatalf("expected success: %s", result.Output)
|
||||
}
|
||||
if !strings.Contains(result.Output, "test content") {
|
||||
t.Errorf("expected 'test content' in output: %s", result.Output)
|
||||
}
|
||||
}
|
||||
|
||||
func TestRun_ParseError(t *testing.T) {
|
||||
result := Run(context.Background(), RunConfig{
|
||||
Command: "echo 'unclosed",
|
||||
Dir: t.TempDir(),
|
||||
Timeout: 5 * time.Second,
|
||||
RiskThreshold: RiskMedium,
|
||||
})
|
||||
|
||||
if !result.IsError {
|
||||
t.Fatal("expected parse error")
|
||||
}
|
||||
if !strings.Contains(result.Output, "parse") {
|
||||
t.Errorf("expected 'parse' in error: %s", result.Output)
|
||||
}
|
||||
}
|
||||
|
||||
func TestRun_StderrCapture(t *testing.T) {
|
||||
result := Run(context.Background(), RunConfig{
|
||||
Command: "echo stdout_msg; echo stderr_msg >&2",
|
||||
Dir: t.TempDir(),
|
||||
Timeout: 5 * time.Second,
|
||||
RiskThreshold: RiskMedium,
|
||||
})
|
||||
|
||||
if !strings.Contains(result.Output, "stdout_msg") {
|
||||
t.Errorf("expected stdout in output: %s", result.Output)
|
||||
}
|
||||
if !strings.Contains(result.Output, "stderr_msg") {
|
||||
t.Errorf("expected stderr in output: %s", result.Output)
|
||||
}
|
||||
}
|
||||
|
||||
func TestRun_HighThresholdAllowsRm(t *testing.T) {
|
||||
tmpDir := t.TempDir()
|
||||
testFile := filepath.Join(tmpDir, "delete_me.txt")
|
||||
os.WriteFile(testFile, []byte("bye"), 0o644)
|
||||
|
||||
result := Run(context.Background(), RunConfig{
|
||||
Command: "rm delete_me.txt",
|
||||
Dir: tmpDir,
|
||||
Timeout: 5 * time.Second,
|
||||
RiskThreshold: RiskHigh,
|
||||
})
|
||||
|
||||
if result.IsError {
|
||||
t.Fatalf("with threshold=high, rm should be allowed: %s", result.Output)
|
||||
}
|
||||
if _, err := os.Stat(testFile); err == nil {
|
||||
t.Error("file should have been deleted")
|
||||
}
|
||||
}
|
||||
|
||||
func TestRun_EnvSanitization(t *testing.T) {
|
||||
t.Setenv("OPENAI_API_KEY", "sk-secret-test")
|
||||
t.Setenv("PATH", os.Getenv("PATH"))
|
||||
|
||||
result := Run(context.Background(), RunConfig{
|
||||
Command: "env",
|
||||
Dir: t.TempDir(),
|
||||
Timeout: 5 * time.Second,
|
||||
RiskThreshold: RiskMedium,
|
||||
})
|
||||
|
||||
if result.IsError {
|
||||
t.Fatalf("expected env command to succeed: %s", result.Output)
|
||||
}
|
||||
if strings.Contains(result.Output, "OPENAI_API_KEY") {
|
||||
t.Error("OPENAI_API_KEY should not be in child environment")
|
||||
}
|
||||
if !strings.Contains(result.Output, "PATH=") {
|
||||
t.Error("PATH should be in child environment")
|
||||
}
|
||||
}
|
||||
|
||||
func TestRun_PipelineCommand(t *testing.T) {
|
||||
result := Run(context.Background(), RunConfig{
|
||||
Command: "echo 'line1\nline2\nline3' | wc -l",
|
||||
Dir: t.TempDir(),
|
||||
Timeout: 5 * time.Second,
|
||||
RiskThreshold: RiskMedium,
|
||||
})
|
||||
|
||||
if result.IsError {
|
||||
t.Fatalf("expected pipeline to succeed: %s", result.Output)
|
||||
}
|
||||
}
|
||||
|
||||
func TestRun_DevNullRedirection(t *testing.T) {
|
||||
result := Run(context.Background(), RunConfig{
|
||||
Command: "echo hello 2>/dev/null",
|
||||
Dir: t.TempDir(),
|
||||
Timeout: 5 * time.Second,
|
||||
Restrict: true,
|
||||
WorkspaceDir: t.TempDir(),
|
||||
RiskThreshold: RiskMedium,
|
||||
})
|
||||
|
||||
if result.IsError && strings.Contains(result.Output, "sandbox") {
|
||||
t.Errorf("/dev/null should not be blocked: %s", result.Output)
|
||||
}
|
||||
}
|
||||
|
||||
func TestRun_RiskOverrides(t *testing.T) {
|
||||
// Override rm to low so it passes with threshold=medium.
|
||||
result := Run(context.Background(), RunConfig{
|
||||
Command: "rm nonexistent_file_xyz 2>/dev/null; echo done",
|
||||
Dir: t.TempDir(),
|
||||
Timeout: 5 * time.Second,
|
||||
RiskThreshold: RiskMedium,
|
||||
RiskOverrides: map[string]string{"rm": "low"},
|
||||
})
|
||||
|
||||
// rm should be allowed because of override. The command may fail
|
||||
// (file doesn't exist) but it should not be _blocked_.
|
||||
if result.IsError && strings.Contains(result.Output, "blocked") {
|
||||
t.Errorf("rm should be allowed with override: %s", result.Output)
|
||||
}
|
||||
}
|
||||
79
pkg/tools/shell/sandbox.go
Normal file
79
pkg/tools/shell/sandbox.go
Normal file
|
|
@ -0,0 +1,79 @@
|
|||
package shell
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"io"
|
||||
"os"
|
||||
"path/filepath"
|
||||
|
||||
"mvdan.cc/sh/v3/interp"
|
||||
)
|
||||
|
||||
// SafePaths are kernel pseudo-devices that are always safe to open,
|
||||
// regardless of workspace restriction.
|
||||
var SafePaths = map[string]bool{
|
||||
"/dev/null": true,
|
||||
"/dev/zero": true,
|
||||
"/dev/random": true,
|
||||
"/dev/urandom": true,
|
||||
"/dev/stdin": true,
|
||||
"/dev/stdout": true,
|
||||
"/dev/stderr": true,
|
||||
}
|
||||
|
||||
// SandboxedOpenHandler returns an interp.OpenHandlerFunc that restricts
|
||||
// shell redirections (>, <, >>) to files within the workspace directory.
|
||||
//
|
||||
// NOTE: This only intercepts opens from the shell interpreter for
|
||||
// redirections. External programs open files via their own syscalls
|
||||
// and are NOT restricted by this handler.
|
||||
func SandboxedOpenHandler(workspaceDir string) interp.OpenHandlerFunc {
|
||||
absWorkspace, err := filepath.Abs(workspaceDir)
|
||||
if err != nil {
|
||||
absWorkspace = workspaceDir
|
||||
}
|
||||
// Resolve workspace symlinks for accurate escape detection.
|
||||
absWorkspace, err = filepath.EvalSymlinks(absWorkspace)
|
||||
if err != nil {
|
||||
// Non-fatal: continue with absolute path. Realpath failures
|
||||
// are caught per-file when the sandbox is actually used.
|
||||
}
|
||||
|
||||
return func(ctx context.Context, path string, flag int, perm os.FileMode) (io.ReadWriteCloser, error) {
|
||||
if SafePaths[path] {
|
||||
return os.OpenFile(path, flag, perm)
|
||||
}
|
||||
|
||||
absPath, err := filepath.Abs(path)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("sandbox: cannot resolve path %q: %w", path, err)
|
||||
}
|
||||
|
||||
// Resolve symlinks to prevent escape.
|
||||
// If the file doesn't exist yet, resolve the parent.
|
||||
var resolved string
|
||||
if _, err := os.Lstat(absPath); err == nil {
|
||||
resolved, err = filepath.EvalSymlinks(absPath)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("sandbox: cannot resolve symlink %q: %w", path, err)
|
||||
}
|
||||
} else {
|
||||
parentDir := filepath.Dir(absPath)
|
||||
resolvedParent, err := filepath.EvalSymlinks(parentDir)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("sandbox: cannot resolve parent dir %q: %w", parentDir, err)
|
||||
}
|
||||
resolved = filepath.Join(resolvedParent, filepath.Base(absPath))
|
||||
}
|
||||
|
||||
if absWorkspace != "" {
|
||||
rel, err := filepath.Rel(absWorkspace, resolved)
|
||||
if err != nil || !filepath.IsLocal(rel) {
|
||||
return nil, fmt.Errorf("sandbox: path %q resolves outside workspace", path)
|
||||
}
|
||||
}
|
||||
|
||||
return os.OpenFile(resolved, flag, perm)
|
||||
}
|
||||
}
|
||||
123
pkg/tools/shell/sandbox_test.go
Normal file
123
pkg/tools/shell/sandbox_test.go
Normal file
|
|
@ -0,0 +1,123 @@
|
|||
package shell
|
||||
|
||||
import (
|
||||
"context"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestSandboxedOpenHandler_AllowsInsideWorkspace(t *testing.T) {
|
||||
workspace := t.TempDir()
|
||||
handler := SandboxedOpenHandler(workspace)
|
||||
|
||||
testFile := filepath.Join(workspace, "test.txt")
|
||||
os.WriteFile(testFile, []byte("hello"), 0o644)
|
||||
|
||||
f, err := handler(context.Background(), testFile, os.O_RDONLY, 0)
|
||||
if err != nil {
|
||||
t.Fatalf("expected open inside workspace to succeed: %v", err)
|
||||
}
|
||||
f.Close()
|
||||
}
|
||||
|
||||
func TestSandboxedOpenHandler_BlocksOutsideWorkspace(t *testing.T) {
|
||||
workspace := t.TempDir()
|
||||
handler := SandboxedOpenHandler(workspace)
|
||||
|
||||
outsideFile := filepath.Join(t.TempDir(), "secret.txt")
|
||||
os.WriteFile(outsideFile, []byte("secret"), 0o644)
|
||||
|
||||
_, err := handler(context.Background(), outsideFile, os.O_RDONLY, 0)
|
||||
if err == nil {
|
||||
t.Fatal("expected open outside workspace to be blocked")
|
||||
}
|
||||
}
|
||||
|
||||
func TestSandboxedOpenHandler_AllowsSafePaths(t *testing.T) {
|
||||
workspace := t.TempDir()
|
||||
handler := SandboxedOpenHandler(workspace)
|
||||
|
||||
f, err := handler(context.Background(), "/dev/null", os.O_WRONLY, 0)
|
||||
if err != nil {
|
||||
t.Fatalf("expected /dev/null to be allowed: %v", err)
|
||||
}
|
||||
f.Close()
|
||||
}
|
||||
|
||||
func TestSandboxedOpenHandler_BlocksSymlinkEscape(t *testing.T) {
|
||||
root := t.TempDir()
|
||||
workspace := filepath.Join(root, "workspace")
|
||||
secretDir := filepath.Join(root, "secret")
|
||||
os.MkdirAll(workspace, 0o755)
|
||||
os.MkdirAll(secretDir, 0o755)
|
||||
os.WriteFile(filepath.Join(secretDir, "data.txt"), []byte("secret"), 0o644)
|
||||
|
||||
link := filepath.Join(workspace, "escape")
|
||||
if err := os.Symlink(secretDir, link); err != nil {
|
||||
t.Skipf("symlinks not supported: %v", err)
|
||||
}
|
||||
|
||||
handler := SandboxedOpenHandler(workspace)
|
||||
|
||||
target := filepath.Join(link, "data.txt")
|
||||
_, err := handler(context.Background(), target, os.O_RDONLY, 0)
|
||||
if err == nil {
|
||||
t.Fatal("expected symlink escape to be blocked")
|
||||
}
|
||||
}
|
||||
|
||||
func TestSandboxedOpenHandler_AllowsNewFileInWorkspace(t *testing.T) {
|
||||
workspace := t.TempDir()
|
||||
handler := SandboxedOpenHandler(workspace)
|
||||
|
||||
newFile := filepath.Join(workspace, "new_output.txt")
|
||||
f, err := handler(context.Background(), newFile, os.O_WRONLY|os.O_CREATE|os.O_TRUNC, 0o644)
|
||||
if err != nil {
|
||||
t.Fatalf("expected new file creation inside workspace to succeed: %v", err)
|
||||
}
|
||||
f.Close()
|
||||
os.Remove(newFile)
|
||||
}
|
||||
|
||||
// TestSandboxedOpenHandler_AllowsDottedFiles verifies that files with
|
||||
// names starting with ".." (like ".../file", "....txt", "..something")
|
||||
// are NOT incorrectly blocked by the escape check.
|
||||
// Regression test for bug where rel[:2] == ".." would match these.
|
||||
func TestSandboxedOpenHandler_AllowsDottedFiles(t *testing.T) {
|
||||
workspace := t.TempDir()
|
||||
handler := SandboxedOpenHandler(workspace)
|
||||
|
||||
// Create files with names that start with ".." but don't escape
|
||||
testCases := []string{
|
||||
".../test.txt",
|
||||
"....txt",
|
||||
"..something.txt",
|
||||
}
|
||||
|
||||
for _, tc := range testCases {
|
||||
t.Run(tc, func(t *testing.T) {
|
||||
fullPath := filepath.Join(workspace, tc)
|
||||
|
||||
// Create parent directories if needed
|
||||
if err := os.MkdirAll(filepath.Dir(fullPath), 0o755); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
// Create the file
|
||||
if err := os.WriteFile(fullPath, []byte("test"), 0o644); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer os.Remove(fullPath)
|
||||
|
||||
// Try to open via sandbox handler
|
||||
f, err := handler(context.Background(), fullPath, os.O_RDONLY, 0)
|
||||
if err != nil {
|
||||
t.Errorf("file %q should be allowed within workspace, got error: %v", tc, err)
|
||||
}
|
||||
if f != nil {
|
||||
f.Close()
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
|
@ -1,32 +0,0 @@
|
|||
//go:build !windows
|
||||
|
||||
package tools
|
||||
|
||||
import (
|
||||
"os/exec"
|
||||
"syscall"
|
||||
)
|
||||
|
||||
func prepareCommandForTermination(cmd *exec.Cmd) {
|
||||
if cmd == nil {
|
||||
return
|
||||
}
|
||||
cmd.SysProcAttr = &syscall.SysProcAttr{Setpgid: true}
|
||||
}
|
||||
|
||||
func terminateProcessTree(cmd *exec.Cmd) error {
|
||||
if cmd == nil || cmd.Process == nil {
|
||||
return nil
|
||||
}
|
||||
|
||||
pid := cmd.Process.Pid
|
||||
if pid <= 0 {
|
||||
return nil
|
||||
}
|
||||
|
||||
// Kill the entire process group spawned by the shell command.
|
||||
_ = syscall.Kill(-pid, syscall.SIGKILL)
|
||||
// Fallback kill on the shell process itself.
|
||||
_ = cmd.Process.Kill()
|
||||
return nil
|
||||
}
|
||||
|
|
@ -1,27 +0,0 @@
|
|||
//go:build windows
|
||||
|
||||
package tools
|
||||
|
||||
import (
|
||||
"os/exec"
|
||||
"strconv"
|
||||
)
|
||||
|
||||
func prepareCommandForTermination(cmd *exec.Cmd) {
|
||||
// no-op on Windows
|
||||
}
|
||||
|
||||
func terminateProcessTree(cmd *exec.Cmd) error {
|
||||
if cmd == nil || cmd.Process == nil {
|
||||
return nil
|
||||
}
|
||||
|
||||
pid := cmd.Process.Pid
|
||||
if pid <= 0 {
|
||||
return nil
|
||||
}
|
||||
|
||||
_ = exec.Command("taskkill", "/T", "/F", "/PID", strconv.Itoa(pid)).Run()
|
||||
_ = cmd.Process.Kill()
|
||||
return nil
|
||||
}
|
||||
|
|
@ -1,445 +0,0 @@
|
|||
package tools
|
||||
|
||||
import (
|
||||
"context"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/sipeed/picoclaw/pkg/config"
|
||||
)
|
||||
|
||||
// TestShellTool_Success verifies successful command execution
|
||||
func TestShellTool_Success(t *testing.T) {
|
||||
tool, err := NewExecTool("", false)
|
||||
if err != nil {
|
||||
t.Errorf("unable to configure exec tool: %s", err)
|
||||
}
|
||||
|
||||
ctx := context.Background()
|
||||
args := map[string]any{
|
||||
"command": "echo 'hello world'",
|
||||
}
|
||||
|
||||
result := tool.Execute(ctx, args)
|
||||
|
||||
// Success should not be an error
|
||||
if result.IsError {
|
||||
t.Errorf("Expected success, got IsError=true: %s", result.ForLLM)
|
||||
}
|
||||
|
||||
// ForUser should contain command output
|
||||
if !strings.Contains(result.ForUser, "hello world") {
|
||||
t.Errorf("Expected ForUser to contain 'hello world', got: %s", result.ForUser)
|
||||
}
|
||||
|
||||
// ForLLM should contain full output
|
||||
if !strings.Contains(result.ForLLM, "hello world") {
|
||||
t.Errorf("Expected ForLLM to contain 'hello world', got: %s", result.ForLLM)
|
||||
}
|
||||
}
|
||||
|
||||
// TestShellTool_Failure verifies failed command execution
|
||||
func TestShellTool_Failure(t *testing.T) {
|
||||
tool, err := NewExecTool("", false)
|
||||
if err != nil {
|
||||
t.Errorf("unable to configure exec tool: %s", err)
|
||||
}
|
||||
|
||||
ctx := context.Background()
|
||||
args := map[string]any{
|
||||
"command": "ls /nonexistent_directory_12345",
|
||||
}
|
||||
|
||||
result := tool.Execute(ctx, args)
|
||||
|
||||
// Failure should be marked as error
|
||||
if !result.IsError {
|
||||
t.Errorf("Expected error for failed command, got IsError=false")
|
||||
}
|
||||
|
||||
// ForUser should contain error information
|
||||
if result.ForUser == "" {
|
||||
t.Errorf("Expected ForUser to contain error info, got empty string")
|
||||
}
|
||||
|
||||
// ForLLM should contain exit code or error
|
||||
if !strings.Contains(result.ForLLM, "Exit code") && result.ForUser == "" {
|
||||
t.Errorf("Expected ForLLM to contain exit code or error, got: %s", result.ForLLM)
|
||||
}
|
||||
}
|
||||
|
||||
// TestShellTool_Timeout verifies command timeout handling
|
||||
func TestShellTool_Timeout(t *testing.T) {
|
||||
tool, err := NewExecTool("", false)
|
||||
if err != nil {
|
||||
t.Errorf("unable to configure exec tool: %s", err)
|
||||
}
|
||||
|
||||
tool.SetTimeout(100 * time.Millisecond)
|
||||
|
||||
ctx := context.Background()
|
||||
args := map[string]any{
|
||||
"command": "sleep 10",
|
||||
}
|
||||
|
||||
result := tool.Execute(ctx, args)
|
||||
|
||||
// Timeout should be marked as error
|
||||
if !result.IsError {
|
||||
t.Errorf("Expected error for timeout, got IsError=false")
|
||||
}
|
||||
|
||||
// Should mention timeout
|
||||
if !strings.Contains(result.ForLLM, "timed out") && !strings.Contains(result.ForUser, "timed out") {
|
||||
t.Errorf("Expected timeout message, got ForLLM: %s, ForUser: %s", result.ForLLM, result.ForUser)
|
||||
}
|
||||
}
|
||||
|
||||
// TestShellTool_WorkingDir verifies custom working directory
|
||||
func TestShellTool_WorkingDir(t *testing.T) {
|
||||
// Create temp directory
|
||||
tmpDir := t.TempDir()
|
||||
testFile := filepath.Join(tmpDir, "test.txt")
|
||||
os.WriteFile(testFile, []byte("test content"), 0o644)
|
||||
|
||||
tool, err := NewExecTool("", false)
|
||||
if err != nil {
|
||||
t.Errorf("unable to configure exec tool: %s", err)
|
||||
}
|
||||
|
||||
ctx := context.Background()
|
||||
args := map[string]any{
|
||||
"command": "cat test.txt",
|
||||
"working_dir": tmpDir,
|
||||
}
|
||||
|
||||
result := tool.Execute(ctx, args)
|
||||
|
||||
if result.IsError {
|
||||
t.Errorf("Expected success in custom working dir, got error: %s", result.ForLLM)
|
||||
}
|
||||
|
||||
if !strings.Contains(result.ForUser, "test content") {
|
||||
t.Errorf("Expected output from custom dir, got: %s", result.ForUser)
|
||||
}
|
||||
}
|
||||
|
||||
// TestShellTool_DangerousCommand verifies safety guard blocks dangerous commands
|
||||
func TestShellTool_DangerousCommand(t *testing.T) {
|
||||
tool, err := NewExecTool("", false)
|
||||
if err != nil {
|
||||
t.Errorf("unable to configure exec tool: %s", err)
|
||||
}
|
||||
|
||||
ctx := context.Background()
|
||||
args := map[string]any{
|
||||
"command": "rm -rf /",
|
||||
}
|
||||
|
||||
result := tool.Execute(ctx, args)
|
||||
|
||||
// Dangerous command should be blocked
|
||||
if !result.IsError {
|
||||
t.Errorf("Expected dangerous command to be blocked (IsError=true)")
|
||||
}
|
||||
|
||||
if !strings.Contains(result.ForLLM, "blocked") && !strings.Contains(result.ForUser, "blocked") {
|
||||
t.Errorf("Expected 'blocked' message, got ForLLM: %s, ForUser: %s", result.ForLLM, result.ForUser)
|
||||
}
|
||||
}
|
||||
|
||||
func TestShellTool_DangerousCommand_KillBlocked(t *testing.T) {
|
||||
tool, err := NewExecTool("", false)
|
||||
if err != nil {
|
||||
t.Errorf("unable to configure exec tool: %s", err)
|
||||
}
|
||||
|
||||
ctx := context.Background()
|
||||
args := map[string]any{
|
||||
"command": "kill 12345",
|
||||
}
|
||||
|
||||
result := tool.Execute(ctx, args)
|
||||
if !result.IsError {
|
||||
t.Errorf("Expected kill command to be blocked")
|
||||
}
|
||||
if !strings.Contains(result.ForLLM, "blocked") && !strings.Contains(result.ForUser, "blocked") {
|
||||
t.Errorf("Expected blocked message, got ForLLM: %s, ForUser: %s", result.ForLLM, result.ForUser)
|
||||
}
|
||||
}
|
||||
|
||||
// TestShellTool_MissingCommand verifies error handling for missing command
|
||||
func TestShellTool_MissingCommand(t *testing.T) {
|
||||
tool, err := NewExecTool("", false)
|
||||
if err != nil {
|
||||
t.Errorf("unable to configure exec tool: %s", err)
|
||||
}
|
||||
|
||||
ctx := context.Background()
|
||||
args := map[string]any{}
|
||||
|
||||
result := tool.Execute(ctx, args)
|
||||
|
||||
// Should return error result
|
||||
if !result.IsError {
|
||||
t.Errorf("Expected error when command is missing")
|
||||
}
|
||||
}
|
||||
|
||||
// TestShellTool_StderrCapture verifies stderr is captured and included
|
||||
func TestShellTool_StderrCapture(t *testing.T) {
|
||||
tool, err := NewExecTool("", false)
|
||||
if err != nil {
|
||||
t.Errorf("unable to configure exec tool: %s", err)
|
||||
}
|
||||
|
||||
ctx := context.Background()
|
||||
args := map[string]any{
|
||||
"command": "sh -c 'echo stdout; echo stderr >&2'",
|
||||
}
|
||||
|
||||
result := tool.Execute(ctx, args)
|
||||
|
||||
// Both stdout and stderr should be in output
|
||||
if !strings.Contains(result.ForLLM, "stdout") {
|
||||
t.Errorf("Expected stdout in output, got: %s", result.ForLLM)
|
||||
}
|
||||
if !strings.Contains(result.ForLLM, "stderr") {
|
||||
t.Errorf("Expected stderr in output, got: %s", result.ForLLM)
|
||||
}
|
||||
}
|
||||
|
||||
// TestShellTool_OutputTruncation verifies long output is truncated
|
||||
func TestShellTool_OutputTruncation(t *testing.T) {
|
||||
tool, err := NewExecTool("", false)
|
||||
if err != nil {
|
||||
t.Errorf("unable to configure exec tool: %s", err)
|
||||
}
|
||||
|
||||
ctx := context.Background()
|
||||
// Generate long output (>10000 chars)
|
||||
args := map[string]any{
|
||||
"command": "python3 -c \"print('x' * 20000)\" || echo " + strings.Repeat("x", 20000),
|
||||
}
|
||||
|
||||
result := tool.Execute(ctx, args)
|
||||
|
||||
// Should have truncation message or be truncated
|
||||
if len(result.ForLLM) > 15000 {
|
||||
t.Errorf("Expected output to be truncated, got length: %d", len(result.ForLLM))
|
||||
}
|
||||
}
|
||||
|
||||
// TestShellTool_WorkingDir_OutsideWorkspace verifies that working_dir cannot escape the workspace directly
|
||||
func TestShellTool_WorkingDir_OutsideWorkspace(t *testing.T) {
|
||||
root := t.TempDir()
|
||||
workspace := filepath.Join(root, "workspace")
|
||||
outsideDir := filepath.Join(root, "outside")
|
||||
if err := os.MkdirAll(workspace, 0o755); err != nil {
|
||||
t.Fatalf("failed to create workspace: %v", err)
|
||||
}
|
||||
if err := os.MkdirAll(outsideDir, 0o755); err != nil {
|
||||
t.Fatalf("failed to create outside dir: %v", err)
|
||||
}
|
||||
|
||||
tool, err := NewExecTool(workspace, true)
|
||||
if err != nil {
|
||||
t.Errorf("unable to configure exec tool: %s", err)
|
||||
}
|
||||
|
||||
result := tool.Execute(context.Background(), map[string]any{
|
||||
"command": "pwd",
|
||||
"working_dir": outsideDir,
|
||||
})
|
||||
|
||||
if !result.IsError {
|
||||
t.Fatalf("expected working_dir outside workspace to be blocked, got output: %s", result.ForLLM)
|
||||
}
|
||||
if !strings.Contains(result.ForLLM, "blocked") {
|
||||
t.Errorf("expected 'blocked' in error, got: %s", result.ForLLM)
|
||||
}
|
||||
}
|
||||
|
||||
// TestShellTool_WorkingDir_SymlinkEscape verifies that a symlink inside the workspace
|
||||
// pointing outside cannot be used as working_dir to escape the sandbox.
|
||||
func TestShellTool_WorkingDir_SymlinkEscape(t *testing.T) {
|
||||
root := t.TempDir()
|
||||
workspace := filepath.Join(root, "workspace")
|
||||
secretDir := filepath.Join(root, "secret")
|
||||
if err := os.MkdirAll(workspace, 0o755); err != nil {
|
||||
t.Fatalf("failed to create workspace: %v", err)
|
||||
}
|
||||
if err := os.MkdirAll(secretDir, 0o755); err != nil {
|
||||
t.Fatalf("failed to create secret dir: %v", err)
|
||||
}
|
||||
os.WriteFile(filepath.Join(secretDir, "secret.txt"), []byte("top secret"), 0o644)
|
||||
|
||||
// symlink lives inside the workspace but resolves to secretDir outside it
|
||||
link := filepath.Join(workspace, "escape")
|
||||
if err := os.Symlink(secretDir, link); err != nil {
|
||||
t.Skipf("symlinks not supported in this environment: %v", err)
|
||||
}
|
||||
|
||||
tool, err := NewExecTool(workspace, true)
|
||||
if err != nil {
|
||||
t.Errorf("unable to configure exec tool: %s", err)
|
||||
}
|
||||
|
||||
result := tool.Execute(context.Background(), map[string]any{
|
||||
"command": "cat secret.txt",
|
||||
"working_dir": link,
|
||||
})
|
||||
|
||||
if !result.IsError {
|
||||
t.Fatalf("expected symlink working_dir escape to be blocked, got output: %s", result.ForLLM)
|
||||
}
|
||||
if !strings.Contains(result.ForLLM, "blocked") {
|
||||
t.Errorf("expected 'blocked' in error, got: %s", result.ForLLM)
|
||||
}
|
||||
}
|
||||
|
||||
// TestShellTool_RestrictToWorkspace verifies workspace restriction
|
||||
func TestShellTool_RestrictToWorkspace(t *testing.T) {
|
||||
tmpDir := t.TempDir()
|
||||
tool, err := NewExecTool(tmpDir, false)
|
||||
if err != nil {
|
||||
t.Errorf("unable to configure exec tool: %s", err)
|
||||
}
|
||||
|
||||
tool.SetRestrictToWorkspace(true)
|
||||
|
||||
ctx := context.Background()
|
||||
args := map[string]any{
|
||||
"command": "cat ../../etc/passwd",
|
||||
}
|
||||
|
||||
result := tool.Execute(ctx, args)
|
||||
|
||||
// Path traversal should be blocked
|
||||
if !result.IsError {
|
||||
t.Errorf("Expected path traversal to be blocked with restrictToWorkspace=true")
|
||||
}
|
||||
|
||||
if !strings.Contains(result.ForLLM, "blocked") && !strings.Contains(result.ForUser, "blocked") {
|
||||
t.Errorf(
|
||||
"Expected 'blocked' message for path traversal, got ForLLM: %s, ForUser: %s",
|
||||
result.ForLLM,
|
||||
result.ForUser,
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
// TestShellTool_DevNullAllowed verifies that /dev/null redirections are not blocked (issue #964).
|
||||
func TestShellTool_DevNullAllowed(t *testing.T) {
|
||||
tmpDir := t.TempDir()
|
||||
tool, err := NewExecTool(tmpDir, true)
|
||||
if err != nil {
|
||||
t.Fatalf("unable to configure exec tool: %s", err)
|
||||
}
|
||||
|
||||
commands := []string{
|
||||
"echo hello 2>/dev/null",
|
||||
"echo hello >/dev/null",
|
||||
"echo hello > /dev/null",
|
||||
"echo hello 2> /dev/null",
|
||||
"echo hello >/dev/null 2>&1",
|
||||
"find " + tmpDir + " -name '*.go' 2>/dev/null",
|
||||
}
|
||||
|
||||
for _, cmd := range commands {
|
||||
result := tool.Execute(context.Background(), map[string]any{"command": cmd})
|
||||
if result.IsError && strings.Contains(result.ForLLM, "blocked") {
|
||||
t.Errorf("command should not be blocked: %s\n error: %s", cmd, result.ForLLM)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// TestShellTool_BlockDevices verifies that writes to block devices are blocked (issue #965).
|
||||
func TestShellTool_BlockDevices(t *testing.T) {
|
||||
tool, err := NewExecTool("", false)
|
||||
if err != nil {
|
||||
t.Fatalf("unable to configure exec tool: %s", err)
|
||||
}
|
||||
|
||||
blocked := []string{
|
||||
"echo x > /dev/sda",
|
||||
"echo x > /dev/hda",
|
||||
"echo x > /dev/vda",
|
||||
"echo x > /dev/xvda",
|
||||
"echo x > /dev/nvme0n1",
|
||||
"echo x > /dev/mmcblk0",
|
||||
"echo x > /dev/loop0",
|
||||
"echo x > /dev/dm-0",
|
||||
"echo x > /dev/md0",
|
||||
"echo x > /dev/sr0",
|
||||
"echo x > /dev/nbd0",
|
||||
}
|
||||
|
||||
for _, cmd := range blocked {
|
||||
result := tool.Execute(context.Background(), map[string]any{"command": cmd})
|
||||
if !result.IsError {
|
||||
t.Errorf("expected block device write to be blocked: %s", cmd)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// TestShellTool_SafePathsInWorkspaceRestriction verifies that safe kernel pseudo-devices
|
||||
// are allowed even when workspace restriction is active.
|
||||
func TestShellTool_SafePathsInWorkspaceRestriction(t *testing.T) {
|
||||
tmpDir := t.TempDir()
|
||||
tool, err := NewExecTool(tmpDir, true)
|
||||
if err != nil {
|
||||
t.Fatalf("unable to configure exec tool: %s", err)
|
||||
}
|
||||
|
||||
// These reference paths outside workspace but should be allowed via safePaths.
|
||||
commands := []string{
|
||||
"cat /dev/urandom | head -c 16 | od",
|
||||
"echo test > /dev/null",
|
||||
"dd if=/dev/zero bs=1 count=1",
|
||||
}
|
||||
|
||||
for _, cmd := range commands {
|
||||
result := tool.Execute(context.Background(), map[string]any{"command": cmd})
|
||||
if result.IsError && strings.Contains(result.ForLLM, "path outside working dir") {
|
||||
t.Errorf("safe path should not be blocked by workspace check: %s\n error: %s", cmd, result.ForLLM)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// TestShellTool_CustomAllowPatterns verifies that custom allow patterns exempt
|
||||
// commands from deny pattern checks.
|
||||
func TestShellTool_CustomAllowPatterns(t *testing.T) {
|
||||
cfg := &config.Config{
|
||||
Tools: config.ToolsConfig{
|
||||
Exec: config.ExecConfig{
|
||||
EnableDenyPatterns: true,
|
||||
CustomAllowPatterns: []string{`\bgit\s+push\s+origin\b`},
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
tool, err := NewExecToolWithConfig("", false, cfg)
|
||||
if err != nil {
|
||||
t.Fatalf("unable to configure exec tool: %s", err)
|
||||
}
|
||||
|
||||
// "git push origin main" should be allowed by custom allow pattern.
|
||||
result := tool.Execute(context.Background(), map[string]any{
|
||||
"command": "git push origin main",
|
||||
})
|
||||
if result.IsError && strings.Contains(result.ForLLM, "blocked") {
|
||||
t.Errorf("custom allow pattern should exempt 'git push origin main', got: %s", result.ForLLM)
|
||||
}
|
||||
|
||||
// "git push upstream main" should still be blocked (does not match allow pattern).
|
||||
result = tool.Execute(context.Background(), map[string]any{
|
||||
"command": "git push upstream main",
|
||||
})
|
||||
if !result.IsError {
|
||||
t.Errorf("'git push upstream main' should still be blocked by deny pattern")
|
||||
}
|
||||
}
|
||||
|
|
@ -1,65 +0,0 @@
|
|||
//go:build !windows
|
||||
|
||||
package tools
|
||||
|
||||
import (
|
||||
"context"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strconv"
|
||||
"strings"
|
||||
"syscall"
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
|
||||
func processExists(pid int) bool {
|
||||
if pid <= 0 {
|
||||
return false
|
||||
}
|
||||
err := syscall.Kill(pid, 0)
|
||||
return err == nil || err == syscall.EPERM
|
||||
}
|
||||
|
||||
func TestShellTool_TimeoutKillsChildProcess(t *testing.T) {
|
||||
tool, err := NewExecTool(t.TempDir(), false)
|
||||
if err != nil {
|
||||
t.Errorf("unable to configure exec tool: %s", err)
|
||||
}
|
||||
|
||||
tool.SetTimeout(500 * time.Millisecond)
|
||||
|
||||
args := map[string]any{
|
||||
// Spawn a child process that would outlive the shell unless process-group kill is used.
|
||||
"command": "sleep 60 & echo $! > child.pid; wait",
|
||||
}
|
||||
|
||||
result := tool.Execute(context.Background(), args)
|
||||
if !result.IsError {
|
||||
t.Fatalf("expected timeout error, got success: %s", result.ForLLM)
|
||||
}
|
||||
if !strings.Contains(result.ForLLM, "timed out") {
|
||||
t.Fatalf("expected timeout message, got: %s", result.ForLLM)
|
||||
}
|
||||
|
||||
childPIDPath := filepath.Join(tool.workingDir, "child.pid")
|
||||
data, err := os.ReadFile(childPIDPath)
|
||||
if err != nil {
|
||||
t.Fatalf("failed to read child pid file: %v", err)
|
||||
}
|
||||
|
||||
childPID, err := strconv.Atoi(strings.TrimSpace(string(data)))
|
||||
if err != nil {
|
||||
t.Fatalf("failed to parse child pid: %v", err)
|
||||
}
|
||||
|
||||
deadline := time.Now().Add(2 * time.Second)
|
||||
for time.Now().Before(deadline) {
|
||||
if !processExists(childPID) {
|
||||
return
|
||||
}
|
||||
time.Sleep(50 * time.Millisecond)
|
||||
}
|
||||
|
||||
t.Fatalf("child process %d is still running after timeout", childPID)
|
||||
}
|
||||
228
pkg/tools/shell_tool.go
Normal file
228
pkg/tools/shell_tool.go
Normal file
|
|
@ -0,0 +1,228 @@
|
|||
package tools
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"os"
|
||||
"time"
|
||||
|
||||
"github.com/sipeed/picoclaw/pkg/config"
|
||||
"github.com/sipeed/picoclaw/pkg/tools/shell"
|
||||
)
|
||||
|
||||
// Compile-time interface checks.
|
||||
var (
|
||||
_ Tool = (*ExecTool)(nil)
|
||||
_ AsyncTool = (*ExecTool)(nil)
|
||||
)
|
||||
|
||||
// ExecTool executes shell commands using an in-process interpreter
|
||||
// with AST-based risk classification, env sanitization, and file-access sandboxing.
|
||||
//
|
||||
// ExecTool implements AsyncTool. When the LLM passes background=true the
|
||||
// command runs in a goroutine and the result is delivered via the callback
|
||||
// injected by the tool registry.
|
||||
type ExecTool struct {
|
||||
workingDir string
|
||||
timeout time.Duration
|
||||
restrictToWorkspace bool
|
||||
|
||||
riskThreshold shell.RiskLevel
|
||||
riskOverrides map[string]string
|
||||
argModifiers map[string][]shell.ArgModifier
|
||||
envAllowlist []string
|
||||
envSet map[string]string
|
||||
|
||||
callback AsyncCallback
|
||||
}
|
||||
|
||||
func NewExecTool(workingDir string, restrict bool) (*ExecTool, error) {
|
||||
return NewExecToolWithConfig(workingDir, restrict, nil)
|
||||
}
|
||||
|
||||
func NewExecToolWithConfig(workingDir string, restrict bool, cfg *config.Config) (*ExecTool, error) {
|
||||
t := &ExecTool{
|
||||
workingDir: workingDir,
|
||||
timeout: 60 * time.Second,
|
||||
restrictToWorkspace: restrict,
|
||||
riskThreshold: shell.RiskMedium,
|
||||
}
|
||||
|
||||
if cfg != nil {
|
||||
execCfg := cfg.Tools.Exec
|
||||
|
||||
warnDeprecatedExecConfig(execCfg)
|
||||
|
||||
if execCfg.RiskThreshold != "" {
|
||||
level, err := shell.ParseRiskLevel(execCfg.RiskThreshold)
|
||||
if err != nil {
|
||||
fmt.Printf("Warning: invalid risk_threshold %q: %v. Using medium.\n", execCfg.RiskThreshold, err)
|
||||
} else {
|
||||
t.riskThreshold = level
|
||||
}
|
||||
}
|
||||
t.riskOverrides = execCfg.RiskOverrides
|
||||
t.argModifiers = parseArgModifiers(execCfg.ArgModifiers)
|
||||
t.envAllowlist = execCfg.EnvAllowlist
|
||||
t.envSet = execCfg.EnvSet
|
||||
}
|
||||
|
||||
return t, nil
|
||||
}
|
||||
|
||||
func warnDeprecatedExecConfig(cfg config.ExecConfig) {
|
||||
if !cfg.EnableDenyPatterns {
|
||||
fmt.Println("Warning: 'enable_deny_patterns' is deprecated and ignored. " +
|
||||
"The new shell tool uses AST-based risk classification. " +
|
||||
"Use 'risk_threshold' to control command blocking.")
|
||||
}
|
||||
if len(cfg.CustomDenyPatterns) > 0 {
|
||||
fmt.Println("Warning: 'custom_deny_patterns' is deprecated and ignored. " +
|
||||
"Use 'risk_overrides' to adjust per-command risk levels.")
|
||||
}
|
||||
if len(cfg.CustomAllowPatterns) > 0 {
|
||||
fmt.Println("Warning: 'custom_allow_patterns' is deprecated and ignored. " +
|
||||
"Use 'risk_overrides' to lower the risk level of specific commands.")
|
||||
}
|
||||
}
|
||||
|
||||
func (t *ExecTool) Name() string {
|
||||
return "exec"
|
||||
}
|
||||
|
||||
func (t *ExecTool) Description() string {
|
||||
return "Execute a shell command and return its output. Use with caution."
|
||||
}
|
||||
|
||||
// SetCallback implements AsyncTool.
|
||||
func (t *ExecTool) SetCallback(cb AsyncCallback) {
|
||||
t.callback = cb
|
||||
}
|
||||
|
||||
func (t *ExecTool) Parameters() map[string]any {
|
||||
return map[string]any{
|
||||
"type": "object",
|
||||
"properties": map[string]any{
|
||||
"command": map[string]any{
|
||||
"type": "string",
|
||||
"description": "The shell command to execute",
|
||||
},
|
||||
"working_dir": map[string]any{
|
||||
"type": "string",
|
||||
"description": "Optional working directory for the command",
|
||||
},
|
||||
"background": map[string]any{
|
||||
"type": "boolean",
|
||||
"description": "Run the command in the background. Returns immediately; result is delivered asynchronously.",
|
||||
},
|
||||
},
|
||||
"required": []string{"command"},
|
||||
}
|
||||
}
|
||||
|
||||
func (t *ExecTool) Execute(ctx context.Context, args map[string]any) *ToolResult {
|
||||
command, ok := args["command"].(string)
|
||||
if !ok {
|
||||
return ErrorResult("command is required")
|
||||
}
|
||||
|
||||
cwd := t.workingDir
|
||||
if wd, ok := args["working_dir"].(string); ok && wd != "" {
|
||||
if t.restrictToWorkspace && t.workingDir != "" {
|
||||
resolvedWD, err := validatePath(wd, t.workingDir, true)
|
||||
if err != nil {
|
||||
return ErrorResult("Command blocked by safety guard (" + err.Error() + ")")
|
||||
}
|
||||
cwd = resolvedWD
|
||||
} else {
|
||||
cwd = wd
|
||||
}
|
||||
}
|
||||
|
||||
if cwd == "" {
|
||||
wd, err := os.Getwd()
|
||||
if err == nil {
|
||||
cwd = wd
|
||||
}
|
||||
}
|
||||
|
||||
cfg := shell.RunConfig{
|
||||
Command: command,
|
||||
Dir: cwd,
|
||||
Timeout: t.timeout,
|
||||
Restrict: t.restrictToWorkspace,
|
||||
WorkspaceDir: t.workingDir,
|
||||
RiskThreshold: t.riskThreshold,
|
||||
RiskOverrides: t.riskOverrides,
|
||||
ExtraArgModifiers: t.argModifiers,
|
||||
EnvAllowlist: t.envAllowlist,
|
||||
EnvSet: t.envSet,
|
||||
}
|
||||
|
||||
background, _ := args["background"].(bool)
|
||||
if background && t.callback != nil {
|
||||
return t.executeAsync(ctx, cfg)
|
||||
}
|
||||
|
||||
result := shell.Run(ctx, cfg)
|
||||
return &ToolResult{
|
||||
ForLLM: result.Output,
|
||||
ForUser: result.Output,
|
||||
IsError: result.IsError,
|
||||
}
|
||||
}
|
||||
|
||||
// executeAsync launches the command in a goroutine and delivers the result
|
||||
// through the AsyncCallback. The parent ctx is used for cancellation so the
|
||||
// goroutine respects agent shutdown.
|
||||
func (t *ExecTool) executeAsync(ctx context.Context, cfg shell.RunConfig) *ToolResult {
|
||||
cb := t.callback // capture before goroutine
|
||||
go func() {
|
||||
result := shell.Run(ctx, cfg)
|
||||
cb(ctx, &ToolResult{
|
||||
ForLLM: result.Output,
|
||||
ForUser: result.Output,
|
||||
IsError: result.IsError,
|
||||
})
|
||||
}()
|
||||
return AsyncResult(fmt.Sprintf("Running `%s` in background", cfg.Command))
|
||||
}
|
||||
|
||||
func (t *ExecTool) SetTimeout(timeout time.Duration) {
|
||||
t.timeout = timeout
|
||||
}
|
||||
|
||||
func (t *ExecTool) SetRestrictToWorkspace(restrict bool) {
|
||||
t.restrictToWorkspace = restrict
|
||||
}
|
||||
|
||||
func (t *ExecTool) SetRiskThreshold(level shell.RiskLevel) {
|
||||
t.riskThreshold = level
|
||||
}
|
||||
|
||||
// parseArgModifiers converts config.ArgModifierConfig entries into shell.ArgModifier.
|
||||
func parseArgModifiers(raw map[string][]config.ArgModifierConfig) map[string][]shell.ArgModifier {
|
||||
if len(raw) == 0 {
|
||||
return nil
|
||||
}
|
||||
out := make(map[string][]shell.ArgModifier, len(raw))
|
||||
for cmd, entries := range raw {
|
||||
for _, e := range entries {
|
||||
level, err := shell.ParseRiskLevel(e.Level)
|
||||
if err != nil {
|
||||
fmt.Printf(
|
||||
"Warning: invalid risk level %q for command %q: %v. Skipping this modifier.\n",
|
||||
e.Level,
|
||||
cmd,
|
||||
err,
|
||||
)
|
||||
continue
|
||||
}
|
||||
out[cmd] = append(out[cmd], shell.ArgModifier{
|
||||
Args: e.Args,
|
||||
Level: level,
|
||||
})
|
||||
}
|
||||
}
|
||||
return out
|
||||
}
|
||||
145
pkg/tools/shell_tool_test.go
Normal file
145
pkg/tools/shell_tool_test.go
Normal file
|
|
@ -0,0 +1,145 @@
|
|||
package tools
|
||||
|
||||
import (
|
||||
"context"
|
||||
"sync"
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
|
||||
func TestExecTool_SyncExecution(t *testing.T) {
|
||||
tool, err := NewExecTool(t.TempDir(), false)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
result := tool.Execute(context.Background(), map[string]any{
|
||||
"command": "echo sync_output",
|
||||
})
|
||||
|
||||
if result.IsError {
|
||||
t.Fatalf("expected success: %s", result.ForLLM)
|
||||
}
|
||||
if result.Async {
|
||||
t.Error("sync execution should not be async")
|
||||
}
|
||||
}
|
||||
|
||||
func TestExecTool_BackgroundWithoutCallback(t *testing.T) {
|
||||
tool, err := NewExecTool(t.TempDir(), false)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
// background=true but no callback set → falls through to sync
|
||||
result := tool.Execute(context.Background(), map[string]any{
|
||||
"command": "echo fallback",
|
||||
"background": true,
|
||||
})
|
||||
|
||||
if result.Async {
|
||||
t.Error("should fall back to sync when no callback is set")
|
||||
}
|
||||
if result.IsError {
|
||||
t.Fatalf("expected success: %s", result.ForLLM)
|
||||
}
|
||||
}
|
||||
|
||||
func TestExecTool_BackgroundWithCallback(t *testing.T) {
|
||||
tool, err := NewExecTool(t.TempDir(), false)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
var (
|
||||
mu sync.Mutex
|
||||
received *ToolResult
|
||||
)
|
||||
done := make(chan struct{})
|
||||
|
||||
tool.SetCallback(func(_ context.Context, r *ToolResult) {
|
||||
mu.Lock()
|
||||
received = r
|
||||
mu.Unlock()
|
||||
close(done)
|
||||
})
|
||||
|
||||
result := tool.Execute(context.Background(), map[string]any{
|
||||
"command": "echo async_output",
|
||||
"background": true,
|
||||
})
|
||||
|
||||
if !result.Async {
|
||||
t.Fatal("expected async result")
|
||||
}
|
||||
|
||||
select {
|
||||
case <-done:
|
||||
case <-time.After(10 * time.Second):
|
||||
t.Fatal("timed out waiting for async callback")
|
||||
}
|
||||
|
||||
mu.Lock()
|
||||
defer mu.Unlock()
|
||||
|
||||
if received == nil {
|
||||
t.Fatal("callback was never invoked")
|
||||
}
|
||||
if received.IsError {
|
||||
t.Fatalf("async command failed: %s", received.ForLLM)
|
||||
}
|
||||
}
|
||||
|
||||
func TestExecTool_BackgroundBlockedCommand(t *testing.T) {
|
||||
tool, err := NewExecTool(t.TempDir(), false)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
var (
|
||||
mu sync.Mutex
|
||||
received *ToolResult
|
||||
)
|
||||
done := make(chan struct{})
|
||||
|
||||
tool.SetCallback(func(_ context.Context, r *ToolResult) {
|
||||
mu.Lock()
|
||||
received = r
|
||||
mu.Unlock()
|
||||
close(done)
|
||||
})
|
||||
|
||||
result := tool.Execute(context.Background(), map[string]any{
|
||||
"command": "sudo rm -rf /",
|
||||
"background": true,
|
||||
})
|
||||
|
||||
if !result.Async {
|
||||
t.Fatal("expected async result even for blocked commands")
|
||||
}
|
||||
|
||||
select {
|
||||
case <-done:
|
||||
case <-time.After(10 * time.Second):
|
||||
t.Fatal("timed out waiting for async callback")
|
||||
}
|
||||
|
||||
mu.Lock()
|
||||
defer mu.Unlock()
|
||||
|
||||
if received == nil {
|
||||
t.Fatal("callback was never invoked")
|
||||
}
|
||||
if !received.IsError {
|
||||
t.Error("blocked command should report error via callback")
|
||||
}
|
||||
}
|
||||
|
||||
func TestExecTool_ImplementsAsyncTool(t *testing.T) {
|
||||
tool, err := NewExecTool(t.TempDir(), false)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
var _ AsyncTool = tool // compile-time check
|
||||
}
|
||||
Loading…
Add table
Reference in a new issue