diff --git a/config/config.example.json b/config/config.example.json index c43ac6497..6acec1ef6 100644 --- a/config/config.example.json +++ b/config/config.example.json @@ -343,6 +343,7 @@ "enabled": true, "risk_threshold": "medium", "risk_overrides": {}, + "arg_profiles": {}, "arg_modifiers": {}, "env_allowlist": [], "env_set": {} diff --git a/docs/design/DDR-shell-tool-hardening.md b/docs/design/DDR-shell-tool-hardening.md index b0098eac4..cb396afe6 100644 --- a/docs/design/DDR-shell-tool-hardening.md +++ b/docs/design/DDR-shell-tool-hardening.md @@ -94,7 +94,9 @@ 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`). All matching modifiers MUST be scanned and the highest level that exceeds the base level is applied (highest-match-wins, not first-match-wins). -5a. `risk_overrides` sets the **base level** for a command (replacing the built-in table entry). Argument modifiers MUST still be applied on top of the overridden base level and can elevate it further. This means `risk_overrides: {"rm": "medium"}` allows plain `rm` but `rm -rf` is still elevated to `critical` by the built-in modifier. +5a. Before modifier matching, argv SHOULD be normalized using command-specific flag profiles so equivalent CLI forms are classified identically. The implementation MAY normalize grouped short flags (e.g. `-rf` → `-r`, `-f`), long `--flag=value` forms, and explicitly whitelisted attached short-value forms (e.g. `curl -XPOST` → `-X`, `POST`). + +5b. `risk_overrides` sets the **base level** for a command (replacing the built-in table entry). Argument modifiers MUST still be applied on top of the overridden base level and can elevate it further. This means `risk_overrides: {"rm": "medium"}` allows plain `rm` but `rm -rf` is still elevated to `critical` by the built-in modifier. 6. When a command is blocked, the `ToolResult` MUST include: - Risk level of the command. @@ -123,19 +125,20 @@ 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. +9. The existing regex-based guard (`defaultDenyPatterns`, `guardCommand()`) MUST be removed entirely. The deprecated `ExecConfig` fields `EnableDenyPatterns`, `CustomDenyPatterns`, and `CustomAllowPatterns` MAY remain in the struct for backward-compatible config parsing, but they MUST be ignored at runtime and only emit migration warnings. 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 + ArgProfiles map[string]ArgProfileConfig `json:"arg_profiles"` // command → argv normalization profile (extends built-ins) 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) + 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). + `ArgModifierConfig` is `struct { Args []string; Level string }`. `ArgProfileConfig` configures argv normalization (`split_combined_short`, `split_long_equals`, `short_attached_value_flags`, `separate_value_flags`) using a fixed set of transforms (`identity`, `upper`, `lower`). User-defined modifiers are checked alongside built-ins using highest-match-wins semantics (the maximum level across all matching modifiers is applied). Built-in flag profiles are merged with config-supplied profiles; code remains the source of truth when the document lags. 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. @@ -209,7 +212,7 @@ | `pkg/tools/shell_tool.go` | Thin adapter: `ExecTool` struct implementing `Tool` + `AsyncExecutor`. Delegates to `pkg/tools/shell/` subpackage. Background results delivered via registry `AsyncCallback`. | | `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/config.go` (`ExecConfig`) | Risk config expanded with `risk_threshold`, `risk_overrides`, `arg_profiles`, `arg_modifiers`, env controls; deprecated regex-era fields remain for backward-compatible parsing but are ignored. | | `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. | @@ -222,16 +225,16 @@ | 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/risk.go` | Risk classifier: command→level mapping, flag-profile-based argv normalization, argument modifiers (built-in + configurable), `ClassifyCommand` / `ClassifyCommandWithProfiles` with highest-match-wins semantics. | +| `pkg/tools/shell/risk_test.go` | Table-driven tests for all 4 levels, argument modifiers, overrides, extra modifiers, and profile-driven argv normalization. | | `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(workDir, restrict, cfg)`, `AsyncExecutor` impl, arg modifier wiring. | -| `pkg/tools/shell_tool_test.go` | `ExecTool` sync/async tests, interface compliance checks. | +| `pkg/tools/shell_tool.go` | Adapter: `ExecTool` struct, `NewExecToolWithConfig(workDir, restrict, cfg)`, `AsyncExecutor` impl, arg modifier + arg profile wiring. | +| `pkg/tools/shell_tool_test.go` | `ExecTool` sync/async tests, interface compliance checks, deprecated-config warnings, and config→flag-profile parsing coverage. | | `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 diff --git a/docs/tools_configuration.md b/docs/tools_configuration.md index 6e66d40c4..8d153342d 100644 --- a/docs/tools_configuration.md +++ b/docs/tools_configuration.md @@ -54,6 +54,7 @@ AST-based risk classification, environment sanitization, and file-access sandbox | ---------------- | ------ | ---------- | ------------------------------------------------------------------------------- | | `risk_threshold` | string | `"medium"` | Maximum allowed risk level: `"low"`, `"medium"`, `"high"`, `"critical"` | | `risk_overrides` | object | `{}` | Per-command base risk level (command name → level); modifiers can still elevate | +| `arg_profiles` | object | `{}` | Per-command argv normalization rules used before argument modifier matching | | `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 | @@ -91,6 +92,54 @@ must all be present (order-independent) and the resulting level: The **highest matching** modifier wins (built-in and custom are merged). +#### Argument profiles + +Argument profiles normalize argv before modifier matching. This is useful when a +tool accepts multiple flag syntaxes for the same semantic operation, such as +`curl -XPOST`, `curl -X POST`, and `curl --request=POST`. + +Each command profile can enable: + +- `split_combined_short`: split grouped flags like `-rf` into `-r`, `-f` +- `split_long_equals`: split `--flag=value` into `--flag`, `value` +- `short_attached_value_flags`: split attached short-value forms like `-XPOST` +- `separate_value_flags`: normalize the token after a flag like `-X post` + +Supported value transforms are: + +- `identity`: keep the value unchanged +- `upper`: uppercase the value +- `lower`: lowercase the value + +Example: + +```json +{ + "arg_profiles": { + "curl": { + "split_combined_short": true, + "split_long_equals": true, + "short_attached_value_flags": { + "-X": "upper", + "-d": "identity", + "-T": "identity" + }, + "separate_value_flags": { + "-X": "upper", + "--request": "upper", + "-d": "identity", + "--data": "identity", + "-T": "identity", + "--upload-file": "identity" + } + } + } +} +``` + +Built-in profiles are applied first; config profiles extend or override the +command's flag maps. + #### Precedence The final risk level is computed as: @@ -155,6 +204,22 @@ execution. "ffmpeg": "low", "terraform": "critical" }, + "arg_profiles": { + "curl": { + "split_combined_short": true, + "split_long_equals": true, + "short_attached_value_flags": { + "-X": "upper", + "-d": "identity" + }, + "separate_value_flags": { + "-X": "upper", + "--request": "upper", + "-d": "identity", + "--data": "identity" + } + } + }, "arg_modifiers": { "curl": [{ "args": ["--upload-file"], "level": "high" }] }, diff --git a/pkg/config/config.go b/pkg/config/config.go index 38d343022..e17246aaf 100644 --- a/pkg/config/config.go +++ b/pkg/config/config.go @@ -608,10 +608,20 @@ type ArgModifierConfig struct { Level string `json:"level"` // target risk level: "low"|"medium"|"high"|"critical" } +// ArgProfileConfig describes how a command's argv should be normalized before +// argument-aware risk modifiers are matched. +type ArgProfileConfig struct { + SplitCombinedShort bool `json:"split_combined_short"` + SplitLongEquals bool `json:"split_long_equals"` + ShortAttachedValue map[string]string `json:"short_attached_value_flags"` + SeparateValueFlags map[string]string `json:"separate_value_flags"` +} + type ExecConfig struct { ToolConfig ` envPrefix:"PICOCLAW_TOOLS_EXEC_"` 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 + ArgProfiles map[string]ArgProfileConfig ` json:"arg_profiles"` // command → argv normalization profile (extends built-ins) 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 diff --git a/pkg/tools/shell/risk.go b/pkg/tools/shell/risk.go index 0d1b2f376..bba792f49 100644 --- a/pkg/tools/shell/risk.go +++ b/pkg/tools/shell/risk.go @@ -269,6 +269,59 @@ type ArgModifier struct { Level RiskLevel } +type FlagValueTransform string + +const ( + FlagValueIdentity FlagValueTransform = "identity" + FlagValueUpper FlagValueTransform = "upper" + FlagValueLower FlagValueTransform = "lower" +) + +func ParseFlagValueTransform(s string) (FlagValueTransform, error) { + switch FlagValueTransform(strings.ToLower(s)) { + case "", FlagValueIdentity: + return FlagValueIdentity, nil + case FlagValueUpper: + return FlagValueUpper, nil + case FlagValueLower: + return FlagValueLower, nil + default: + return "", fmt.Errorf("unknown flag value transform %q, must be one of: identity, upper, lower", s) + } +} + +type FlagProfile struct { + SplitCombinedShort bool + SplitLongEquals bool + ShortAttachedValue map[string]FlagValueTransform + SeparateValueFlags map[string]FlagValueTransform +} + +var defaultFlagProfile = FlagProfile{ + SplitCombinedShort: true, + SplitLongEquals: true, +} + +var commandFlagProfiles = map[string]FlagProfile{ + "curl": { + SplitCombinedShort: true, + SplitLongEquals: true, + ShortAttachedValue: map[string]FlagValueTransform{ + "-X": FlagValueUpper, + "-d": FlagValueIdentity, + "-T": FlagValueIdentity, + }, + SeparateValueFlags: map[string]FlagValueTransform{ + "-X": FlagValueUpper, + "--request": FlagValueUpper, + "-d": FlagValueIdentity, + "--data": FlagValueIdentity, + "-T": FlagValueIdentity, + "--upload-file": FlagValueIdentity, + }, + }, +} + // argumentModifiers maps command names to their argument-aware risk adjustments. // All matching modifiers are scanned; the highest level wins. // @@ -380,6 +433,24 @@ var argumentModifiers = map[string][]ArgModifier{ // This means risk_overrides: {"rm": "medium"} allows plain `rm` but // `rm -rf` is still elevated to critical by the built-in modifier. func ClassifyCommand(args []string, overrides map[string]string, extraModifiers ...map[string][]ArgModifier) RiskLevel { + return classifyCommand(args, overrides, nil, extraModifiers...) +} + +func ClassifyCommandWithProfiles( + args []string, + overrides map[string]string, + extraProfiles map[string]FlagProfile, + extraModifiers ...map[string][]ArgModifier, +) RiskLevel { + return classifyCommand(args, overrides, extraProfiles, extraModifiers...) +} + +func classifyCommand( + args []string, + overrides map[string]string, + extraProfiles map[string]FlagProfile, + extraModifiers ...map[string][]ArgModifier, +) RiskLevel { if len(args) == 0 { return RiskMedium } @@ -404,7 +475,8 @@ func ClassifyCommand(args []string, overrides map[string]string, extraModifiers // 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:]) + profile := flagProfileFor(cmdName, extraProfiles) + normalizedArgs := normalizeFlagsWithProfile(profile, args[1:]) // Apply built-in modifiers, then user-supplied. Keep the highest match // across all sources. Modifiers can only elevate, never lower. @@ -517,6 +589,42 @@ func NormalizeCommandKeys[V any](m map[string]V) map[string]V { return normalized } +func MergeFlagProfiles(base, override FlagProfile) FlagProfile { + merged := base + merged.SplitCombinedShort = base.SplitCombinedShort || override.SplitCombinedShort + merged.SplitLongEquals = base.SplitLongEquals || override.SplitLongEquals + merged.ShortAttachedValue = mergeFlagTransformMaps(base.ShortAttachedValue, override.ShortAttachedValue) + merged.SeparateValueFlags = mergeFlagTransformMaps(base.SeparateValueFlags, override.SeparateValueFlags) + return merged +} + +func mergeFlagTransformMaps(base, override map[string]FlagValueTransform) map[string]FlagValueTransform { + if len(base) == 0 && len(override) == 0 { + return nil + } + merged := make(map[string]FlagValueTransform, len(base)+len(override)) + for k, v := range base { + merged[k] = v + } + for k, v := range override { + merged[k] = v + } + return merged +} + +func flagProfileFor(cmdName string, extraProfiles map[string]FlagProfile) FlagProfile { + profile := defaultFlagProfile + if builtIn, ok := commandFlagProfiles[cmdName]; ok { + profile = MergeFlagProfiles(profile, builtIn) + } + if extraProfiles != nil { + if extra, ok := extraProfiles[cmdName]; ok { + profile = MergeFlagProfiles(profile, extra) + } + } + return profile +} + // baseCommand extracts the basename from a command path. // On Windows, it additionally lowercases the name and strips known // executable extensions (.exe, .cmd, .bat, .com) so that @@ -540,24 +648,107 @@ func baseCommand(cmd string) string { return lower } -// normalizeFlags expands combined short flags (e.g., "-rf" → "-r", "-f") -// so that modifier matching works regardless of how flags are grouped. -// Long flags (--flag), non-flag arguments, and slash flags (/s, /MIR) are -// passed through unchanged. Single-dash flags longer than 3 characters -// (e.g., "-urlcache") are treated as long flags and NOT expanded, since -// no standard tool uses 4+ combined single-letter flags. -func normalizeFlags(args []string) []string { +// normalizeFlags expands/normalizes flag forms according to the command's +// flag profile so modifier matching works regardless of grouping or +// flag-value style. +func normalizeFlags(cmdName string, args []string) []string { + profile := flagProfileFor(cmdName, nil) + return normalizeFlagsWithProfile(profile, args) +} + +func normalizeFlagsWithProfile(profile FlagProfile, args []string) []string { result := make([]string, 0, len(args)*2) for _, a := range args { - if len(a) > 2 && len(a) <= 4 && a[0] == '-' && a[1] != '-' { + if flag, value, ok := splitShortAttachedValue(profile, a); ok { + result = append(result, flag) + if value != "" { + result = append(result, normalizeFlagValue(profile, flag, value)) + } + continue + } + + if profile.SplitLongEquals && len(a) > 2 && strings.HasPrefix(a, "--") { + if idx := strings.IndexByte(a, '='); idx > 2 { + flag := a[:idx] + result = append(result, flag) + if idx+1 < len(a) { + result = append(result, normalizeFlagValue(profile, flag, a[idx+1:])) + } + continue + } + } + + if profile.SplitCombinedShort && len(a) > 2 && len(a) <= 4 && a[0] == '-' && a[1] != '-' { for _, ch := range a[1:] { result = append(result, "-"+string(ch)) } - } else { - result = append(result, a) + continue } + + result = append(result, a) } - return result + + return normalizeSeparateFlagValues(profile, result) +} + +func splitShortAttachedValue(profile FlagProfile, arg string) (string, string, bool) { + if len(arg) <= 3 || arg[0] != '-' || arg[1] == '-' { + return "", "", false + } + + flag := arg[:2] + if _, ok := profile.ShortAttachedValue[flag]; !ok { + return "", "", false + } + + if arg[2] == '=' { + if len(arg) == 3 { + return flag, "", true + } + return flag, arg[3:], true + } + + return flag, arg[2:], true +} + +func normalizeFlagValue(profile FlagProfile, flag, value string) string { + if transform, ok := profile.SeparateValueFlags[flag]; ok { + return applyFlagValueTransform(transform, value) + } + if transform, ok := profile.ShortAttachedValue[flag]; ok { + return applyFlagValueTransform(transform, value) + } + return value +} + +func applyFlagValueTransform(transform FlagValueTransform, value string) string { + switch transform { + case FlagValueUpper: + return strings.ToUpper(value) + case FlagValueLower: + return strings.ToLower(value) + case FlagValueIdentity, "": + return value + default: + return value + } +} + +func normalizeSeparateFlagValues(profile FlagProfile, args []string) []string { + if len(profile.SeparateValueFlags) == 0 { + return args + } + + normalized := append([]string(nil), args...) + for i := 0; i+1 < len(normalized); i++ { + transform, ok := profile.SeparateValueFlags[normalized[i]] + if !ok { + continue + } + normalized[i+1] = applyFlagValueTransform(transform, normalized[i+1]) + i++ + } + return normalized } // matchArgs checks if ALL pattern tokens are present in args (order-independent). diff --git a/pkg/tools/shell/risk_test.go b/pkg/tools/shell/risk_test.go index a21def24f..118527543 100644 --- a/pkg/tools/shell/risk_test.go +++ b/pkg/tools/shell/risk_test.go @@ -80,10 +80,18 @@ func TestClassifyCommand_ArgumentModifiers(t *testing.T) { {"curl GET (default)", []string{"curl", "https://example.com"}, RiskMedium}, {"curl POST", []string{"curl", "-X", "POST", "https://example.com"}, RiskHigh}, + {"curl post lowercase", []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 --data=payload", []string{"curl", "--data=payload", "url"}, RiskHigh}, {"curl -X DELETE", []string{"curl", "-X", "DELETE", "url"}, RiskHigh}, + {"curl -XDELETE", []string{"curl", "-XDELETE", "url"}, RiskHigh}, + {"curl -X=DELETE", []string{"curl", "-X=DELETE", "url"}, RiskHigh}, {"curl --request POST", []string{"curl", "--request", "POST", "url"}, RiskHigh}, + {"curl --request=post lowercase", []string{"curl", "--request=post", "url"}, RiskHigh}, + {"curl --request=POST", []string{"curl", "--request=POST", "url"}, RiskHigh}, + {"curl -dDATA", []string{"curl", "-dDATA", "https://example.com"}, RiskHigh}, + {"curl -d=DATA", []string{"curl", "-d=DATA", "https://example.com"}, RiskHigh}, {"rm file (no flags)", []string{"rm", "file.txt"}, RiskHigh}, {"rm -rf", []string{"rm", "-rf", "/"}, RiskCritical}, @@ -353,15 +361,22 @@ func TestNormalizeFlags(t *testing.T) { }{ {[]string{"-rf"}, []string{"-r", "-f"}}, {[]string{"-fr"}, []string{"-f", "-r"}}, + {[]string{"-XPOST"}, []string{"-XPOST"}}, + {[]string{"-X=POST"}, []string{"-X=POST"}}, + {[]string{"-dDATA"}, []string{"-dDATA"}}, + {[]string{"-d=DATA"}, []string{"-d=DATA"}}, + {[]string{"--request=POST"}, []string{"--request", "POST"}}, + {[]string{"--data=body"}, []string{"--data", "body"}}, {[]string{"-r", "-f"}, []string{"-r", "-f"}}, {[]string{"--force"}, []string{"--force"}}, {[]string{"-f"}, []string{"-f"}}, + {[]string{"-9"}, []string{"-9"}}, {[]string{"push"}, []string{"push"}}, {[]string{"-rf", "dir", "--verbose"}, []string{"-r", "-f", "dir", "--verbose"}}, } for _, tt := range tests { - got := normalizeFlags(tt.input) + got := normalizeFlags("git", tt.input) if len(got) != len(tt.want) { t.Errorf("normalizeFlags(%v) = %v, want %v", tt.input, got, tt.want) continue @@ -375,6 +390,59 @@ func TestNormalizeFlags(t *testing.T) { } } +func TestNormalizeFlags_CurlAttachedShortValues(t *testing.T) { + tests := []struct { + input []string + want []string + }{ + {[]string{"-XPOST"}, []string{"-X", "POST"}}, + {[]string{"-Xpost"}, []string{"-X", "POST"}}, + {[]string{"-X=POST"}, []string{"-X", "POST"}}, + {[]string{"-dDATA"}, []string{"-d", "DATA"}}, + {[]string{"-d=DATA"}, []string{"-d", "DATA"}}, + {[]string{"-Tfile.txt"}, []string{"-T", "file.txt"}}, + {[]string{"--request=post"}, []string{"--request", "POST"}}, + {[]string{"-X", "post"}, []string{"-X", "POST"}}, + } + + for _, tt := range tests { + got := normalizeFlags("curl", tt.input) + if len(got) != len(tt.want) { + t.Errorf("normalizeFlags(curl, %v) = %v, want %v", tt.input, got, tt.want) + continue + } + for i := range got { + if got[i] != tt.want[i] { + t.Errorf("normalizeFlags(curl, %v) = %v, want %v", tt.input, got, tt.want) + break + } + } + } +} + +func TestClassifyCommandWithProfiles_CustomAttachedValue(t *testing.T) { + profiles := map[string]FlagProfile{ + "http": { + ShortAttachedValue: map[string]FlagValueTransform{ + "-m": FlagValueUpper, + }, + SeparateValueFlags: map[string]FlagValueTransform{ + "-m": FlagValueUpper, + }, + }, + } + modifiers := map[string][]ArgModifier{ + "http": { + {Args: []string{"-m", "POST"}, Level: RiskHigh}, + }, + } + + got := ClassifyCommandWithProfiles([]string{"http", "-mpost", "https://example.com"}, nil, profiles, modifiers) + if got != RiskHigh { + t.Fatalf("ClassifyCommandWithProfiles(custom attached value) = %s, want %s", got, RiskHigh) + } +} + func TestMatchArgs_OrderIndependent(t *testing.T) { tests := []struct { name string @@ -493,21 +561,21 @@ func TestClassifyCommand_ShellWrapperFullPath(t *testing.T) { func TestApplyModifiers_HighestMatchWins(t *testing.T) { // When multiple modifiers match, the highest level should win. // Scenario: git push matches both ["push"] → High and ["push", "-f"] → Critical - args := normalizeFlags([]string{"push", "-f", "origin"}) + args := normalizeFlags("git", []string{"push", "-f", "origin"}) result := applyModifiers(args, "git", RiskMedium, argumentModifiers) if result != RiskCritical { t.Errorf("git push -f should resolve to critical (highest match), got %s", result) } // Only ["push"] matches → High - args2 := normalizeFlags([]string{"push", "origin"}) + args2 := normalizeFlags("git", []string{"push", "origin"}) result2 := applyModifiers(args2, "git", RiskMedium, argumentModifiers) if result2 != RiskHigh { t.Errorf("git push (no -f) should resolve to high, got %s", result2) } // No modifier matches → base level unchanged - args3 := normalizeFlags([]string{"status"}) + args3 := normalizeFlags("git", []string{"status"}) result3 := applyModifiers(args3, "git", RiskMedium, argumentModifiers) if result3 != RiskMedium { t.Errorf("git status should stay medium, got %s", result3) diff --git a/pkg/tools/shell/runner.go b/pkg/tools/shell/runner.go index 6c2378fc2..c3fd8234d 100644 --- a/pkg/tools/shell/runner.go +++ b/pkg/tools/shell/runner.go @@ -26,6 +26,7 @@ type RunConfig struct { RiskThreshold RiskLevel RiskOverrides map[string]string + ExtraFlagProfiles map[string]FlagProfile ExtraArgModifiers map[string][]ArgModifier // user-defined, appended after built-ins EnvAllowlist []string EnvSet map[string]string @@ -74,7 +75,7 @@ func Run(ctx context.Context, cfg RunConfig) RunResult { opts, interp.ExecHandlers( pathAwareExecHandler(env), - riskExecHandler(cfg.RiskThreshold, cfg.RiskOverrides, cfg.ExtraArgModifiers), + riskExecHandler(cfg.RiskThreshold, cfg.RiskOverrides, cfg.ExtraFlagProfiles, cfg.ExtraArgModifiers), ), ) @@ -134,6 +135,7 @@ func Run(ctx context.Context, cfg RunConfig) RunResult { func riskExecHandler( threshold RiskLevel, overrides map[string]string, + extraProfiles map[string]FlagProfile, extraMods map[string][]ArgModifier, ) func(next interp.ExecHandlerFunc) interp.ExecHandlerFunc { return func(next interp.ExecHandlerFunc) interp.ExecHandlerFunc { @@ -142,7 +144,7 @@ func riskExecHandler( return next(ctx, args) } - level := ClassifyCommand(args, overrides, extraMods) + level := ClassifyCommandWithProfiles(args, overrides, extraProfiles, extraMods) if !IsAllowed(level, threshold) { return NewBlockedError(args, level, threshold, "command risk exceeds configured threshold") } diff --git a/pkg/tools/shell_tool.go b/pkg/tools/shell_tool.go index e433461f1..ea9393ab4 100644 --- a/pkg/tools/shell_tool.go +++ b/pkg/tools/shell_tool.go @@ -31,6 +31,7 @@ type ExecTool struct { riskThreshold shell.RiskLevel riskOverrides map[string]string + argProfiles map[string]shell.FlagProfile argModifiers map[string][]shell.ArgModifier envAllowlist []string envSet map[string]string @@ -66,6 +67,7 @@ func NewExecToolWithConfig( } } t.riskOverrides = shell.NormalizeCommandKeys(execCfg.RiskOverrides) + t.argProfiles = shell.NormalizeCommandKeys(parseArgProfiles(execCfg.ArgProfiles)) t.argModifiers = shell.NormalizeCommandKeys(parseArgModifiers(execCfg.ArgModifiers)) t.envAllowlist = execCfg.EnvAllowlist t.envSet = execCfg.EnvSet @@ -214,6 +216,7 @@ func (t *ExecTool) buildConfig(args map[string]any) (shell.RunConfig, *ToolResul WorkspaceDir: t.workingDir, RiskThreshold: t.riskThreshold, RiskOverrides: t.riskOverrides, + ExtraFlagProfiles: t.argProfiles, ExtraArgModifiers: t.argModifiers, EnvAllowlist: t.envAllowlist, EnvSet: t.envSet, @@ -258,3 +261,61 @@ func parseArgModifiers(raw map[string][]config.ArgModifierConfig) map[string][]s } return out } + +func parseArgProfiles(raw map[string]config.ArgProfileConfig) map[string]shell.FlagProfile { + if len(raw) == 0 { + return nil + } + out := make(map[string]shell.FlagProfile, len(raw)) + for cmd, profile := range raw { + parsed := shell.FlagProfile{ + SplitCombinedShort: profile.SplitCombinedShort, + SplitLongEquals: profile.SplitLongEquals, + } + + if transforms := parseFlagTransforms( + cmd, + "short_attached_value_flags", + profile.ShortAttachedValue, + ); len( + transforms, + ) > 0 { + parsed.ShortAttachedValue = transforms + } + if transforms := parseFlagTransforms( + cmd, + "separate_value_flags", + profile.SeparateValueFlags, + ); len( + transforms, + ) > 0 { + parsed.SeparateValueFlags = transforms + } + + out[cmd] = parsed + } + return out +} + +func parseFlagTransforms(cmd, field string, raw map[string]string) map[string]shell.FlagValueTransform { + if len(raw) == 0 { + return nil + } + out := make(map[string]shell.FlagValueTransform, len(raw)) + for flag, name := range raw { + transform, err := shell.ParseFlagValueTransform(name) + if err != nil { + fmt.Printf( + "Warning: invalid %s transform %q for command %q flag %q: %v. Skipping this flag.\n", + field, + name, + cmd, + flag, + err, + ) + continue + } + out[flag] = transform + } + return out +} diff --git a/pkg/tools/shell_tool_test.go b/pkg/tools/shell_tool_test.go index e4536442c..fa79681b8 100644 --- a/pkg/tools/shell_tool_test.go +++ b/pkg/tools/shell_tool_test.go @@ -9,6 +9,7 @@ import ( "testing" "github.com/sipeed/picoclaw/pkg/config" + "github.com/sipeed/picoclaw/pkg/tools/shell" ) func TestExecTool_SyncExecution(t *testing.T) { @@ -175,12 +176,12 @@ func captureStdout(t *testing.T, fn func()) string { return buf.String() } -func boolPtr(b bool) *bool { return &b } +func ptr[T any](v T) *T { return &v } func TestWarnDeprecatedExecConfig_EnableDenyPatternsFalse(t *testing.T) { out := captureStdout(t, func() { warnDeprecatedExecConfig(config.ExecConfig{ - EnableDenyPatterns: boolPtr(false), + EnableDenyPatterns: ptr(false), }) }) @@ -195,7 +196,7 @@ func TestWarnDeprecatedExecConfig_EnableDenyPatternsFalse(t *testing.T) { func TestWarnDeprecatedExecConfig_EnableDenyPatternsTrue(t *testing.T) { out := captureStdout(t, func() { warnDeprecatedExecConfig(config.ExecConfig{ - EnableDenyPatterns: boolPtr(true), + EnableDenyPatterns: ptr(true), }) }) @@ -236,7 +237,7 @@ func TestWarnDeprecatedExecConfig_CustomPatterns(t *testing.T) { func TestWarnDeprecatedExecConfig_AllDeprecatedFields(t *testing.T) { out := captureStdout(t, func() { warnDeprecatedExecConfig(config.ExecConfig{ - EnableDenyPatterns: boolPtr(false), + EnableDenyPatterns: ptr(false), CustomDenyPatterns: []string{"rm"}, CustomAllowPatterns: []string{"ls"}, }) @@ -257,7 +258,7 @@ func TestWarnDeprecatedExecConfig_AllDeprecatedFields(t *testing.T) { func TestNewExecToolWithConfig_EnableDenyPatternsFalseWarning(t *testing.T) { out := captureStdout(t, func() { cfg := &config.Config{} - cfg.Tools.Exec.EnableDenyPatterns = boolPtr(false) + cfg.Tools.Exec.EnableDenyPatterns = ptr(false) _, err := NewExecToolWithConfig(t.TempDir(), false, cfg) if err != nil { t.Fatal(err) @@ -268,3 +269,45 @@ func TestNewExecToolWithConfig_EnableDenyPatternsFalseWarning(t *testing.T) { t.Errorf("expected warning in NewExecToolWithConfig output: %s", out) } } + +func TestParseArgProfiles(t *testing.T) { + profiles := parseArgProfiles(map[string]config.ArgProfileConfig{ + "curl": { + SplitCombinedShort: true, + SplitLongEquals: true, + ShortAttachedValue: map[string]string{"-X": "upper"}, + SeparateValueFlags: map[string]string{"--request": "upper"}, + }, + }) + + profile, ok := profiles["curl"] + if !ok { + t.Fatal("expected curl profile") + } + if !profile.SplitCombinedShort || !profile.SplitLongEquals { + t.Fatal("expected split flags to be enabled") + } + if got := profile.ShortAttachedValue["-X"]; got != shell.FlagValueUpper { + t.Fatalf("ShortAttachedValue[-X] = %q, want %q", got, shell.FlagValueUpper) + } + if got := profile.SeparateValueFlags["--request"]; got != shell.FlagValueUpper { + t.Fatalf("SeparateValueFlags[--request] = %q, want %q", got, shell.FlagValueUpper) + } +} + +func TestParseArgProfiles_InvalidTransformWarning(t *testing.T) { + out := captureStdout(t, func() { + profiles := parseArgProfiles(map[string]config.ArgProfileConfig{ + "curl": { + ShortAttachedValue: map[string]string{"-X": "bogus"}, + }, + }) + if got := len(profiles["curl"].ShortAttachedValue); got != 0 { + t.Fatalf("expected invalid transform to be skipped, got %d entries", got) + } + }) + + if !strings.Contains(out, "invalid short_attached_value_flags transform") { + t.Fatalf("expected invalid transform warning, got: %s", out) + } +}