diff --git a/intercept/intercept.go b/intercept/intercept.go index 20d969b..45d8f51 100644 --- a/intercept/intercept.go +++ b/intercept/intercept.go @@ -572,8 +572,9 @@ func (i *Interceptor) Match(toolName string, input []byte) (Decision, bool) { if err != nil || len(rules) == 0 { return Decision{}, false } + values := inputStrings(input) for _, r := range rules { - if ruleMatches(r, toolName, input) { + if ruleMatches(r, toolName, input, values) { msg := r.Message if msg == "" { msg = defaultMessage(r.Action, r.Name) @@ -593,20 +594,60 @@ func (i *Interceptor) Match(toolName string, input []byte) (Decision, bool) { return Decision{}, false } -func ruleMatches(r compiledRule, toolName string, input []byte) bool { - var subject string +// ruleMatches reports whether r applies to the call. tool_input rules are tried +// against the raw JSON first, then against each decoded string value in values. +// The raw form alone is not enough: JSON escapes newlines, tabs and quotes, and +// Go's encoder also escapes & < >, so a rule like `\brm\b` never fires on +// "cd /tmp\nrm -rf /var/www" (the raw bytes read `\nrm`, with no word boundary). +func ruleMatches(r compiledRule, toolName string, input []byte, values []string) bool { switch r.MatchTarget { case "tool_name": - subject = toolName + return r.matchString(toolName) case "tool_input": - subject = string(input) - default: - return false + if r.matchString(string(input)) { + return true + } + for _, v := range values { + if r.matchString(v) { + return true + } + } } + return false +} + +func (r compiledRule) matchString(s string) bool { if r.MatchType == "regex" { - return r.re != nil && r.re.MatchString(subject) + return r.re != nil && r.re.MatchString(s) } - return strings.Contains(subject, r.Pattern) + return strings.Contains(s, r.Pattern) +} + +// inputStrings returns every string value in a JSON tool input, at any depth. +// Object keys are skipped. It returns nil if the input is not valid JSON. +func inputStrings(input []byte) []string { + var doc any + if json.Unmarshal(input, &doc) != nil { + return nil + } + var out []string + var walk func(any) + walk = func(v any) { + switch v := v.(type) { + case string: + out = append(out, v) + case []any: + for _, e := range v { + walk(e) + } + case map[string]any: + for _, e := range v { + walk(e) + } + } + } + walk(doc) + return out } func defaultMessage(action, name string) string { diff --git a/intercept/rule_match_test.go b/intercept/rule_match_test.go new file mode 100644 index 0000000..39a0dfa --- /dev/null +++ b/intercept/rule_match_test.go @@ -0,0 +1,138 @@ +package intercept + +import ( + "encoding/json" + "regexp" + "testing" + + "github.com/Autumn-27/artex/db" +) + +// Patterns copied from the built-in rules seeded in db/db.go. +const ( + builtinRmRf = `(?i)\brm\b.{0,80}(?:-[a-z]*r[a-z]*f[a-z]*|-[a-z]*f[a-z]*r[a-z]*|--recursive|--no-preserve-root)` + builtinMkfs = `\bmkfs\b` + builtinDrop = `(?i)\bDROP\s+(?:DATABASE|TABLE|SCHEMA|INDEX|VIEW|TABLESPACE|USER|ROLE)\b` + builtinFlush = `(?i)\b(?:FLUSHALL|FLUSHDB)\b` + builtinV2 = `(?i)\b(rm\s+-rf\s+/|mkfs|dd\s+if=|:\(\)\s*\{|shutdown|reboot|>\s*/dev/sd)` +) + +func testRule(target, matchType, pattern string) compiledRule { + r := compiledRule{InterceptRule: db.InterceptRule{ + ID: 1, Name: "test", Enabled: true, Action: "deny", + MatchTarget: target, MatchType: matchType, Pattern: pattern, + }} + if matchType == "regex" { + r.re = regexp.MustCompile(pattern) + } + return r +} + +// interceptorWith returns an Interceptor whose rule cache is pre-filled, so +// Match runs without a database. +func interceptorWith(rules ...compiledRule) *Interceptor { + return &Interceptor{cached: rules} +} + +func toolInput(t *testing.T, v any) []byte { + t.Helper() + b, err := json.Marshal(v) + if err != nil { + t.Fatal(err) + } + return b +} + +func TestMatchToolInputSeesDecodedValues(t *testing.T) { + cases := []struct { + name string + pattern string + tool string + input any + }{ + {"rm -rf on its own line", builtinRmRf, "Bash", map[string]string{"command": "cd /tmp\nrm -rf /var/www"}}, + {"rm -rf after a tab", builtinRmRf, "Bash", map[string]string{"command": "echo done\trm -rf /var/www"}}, + {"mkfs on its own line", builtinMkfs, "Bash", map[string]string{"command": "lsblk\nmkfs.ext4 /dev/sdb"}}, + {"DROP inside a heredoc", builtinDrop, "Bash", map[string]string{"command": "psql < as >, so this only matches the decoded value. + {"redirect to a disk", `>\s*/dev/sd`, "Bash", map[string]string{"command": "cat img > /dev/sda"}}, + {"interactive shell text", builtinRmRf, "shell_send", map[string]string{"text": "ls\nrm -rf /srv/data\n"}}, + {"nested edit value", builtinDrop, "MultiEdit", map[string]any{ + "file_path": "/tmp/m.sql", + "edits": []map[string]string{{"old_string": "-- x", "new_string": "-- x\nDROP TABLE users;"}}, + }}, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + ic := interceptorWith(testRule("tool_input", "regex", tc.pattern)) + in := toolInput(t, tc.input) + if _, ok := ic.Match(tc.tool, in); !ok { + t.Fatalf("rule %q did not match input %s", tc.pattern, in) + } + }) + } +} + +func TestMatchToolInputStringRuleSeesDecodedValues(t *testing.T) { + // The raw JSON holds rm -rf \"/var/www\", so the quotes only line up once decoded. + ic := interceptorWith(testRule("tool_input", "string", `rm -rf "/var/www"`)) + in := toolInput(t, map[string]string{"command": `rm -rf "/var/www"`}) + if _, ok := ic.Match("Bash", in); !ok { + t.Fatalf("string rule did not match input %s", in) + } +} + +// Rules written against the raw JSON form must keep working. +func TestMatchToolInputStillMatchesRawJSON(t *testing.T) { + ic := interceptorWith(testRule("tool_input", "regex", `"command":"rm `)) + in := toolInput(t, map[string]string{"command": "rm notes.txt"}) + if _, ok := ic.Match("Bash", in); !ok { + t.Fatalf("raw JSON rule did not match input %s", in) + } +} + +func TestMatchToolInputInvalidJSONFallsBackToRaw(t *testing.T) { + ic := interceptorWith(testRule("tool_input", "regex", builtinMkfs)) + if _, ok := ic.Match("Bash", []byte(`not json: mkfs /dev/sdb`)); !ok { + t.Fatal("rule did not match non-JSON input") + } +} + +func TestMatchToolInputNoFalsePositives(t *testing.T) { + ic := interceptorWith( + testRule("tool_input", "regex", builtinRmRf), + testRule("tool_input", "regex", builtinMkfs), + testRule("tool_input", "regex", builtinDrop), + ) + for _, cmd := range []string{ + "cd /tmp\nls -la", + "nmap -sV 10.0.0.5\ncurl -s http://10.0.0.5/", + "echo firm\nrmdir build", // "firm" and "rmdir" must not look like rm + } { + in := toolInput(t, map[string]string{"command": cmd}) + if dec, ok := ic.Match("Bash", in); ok { + t.Fatalf("unexpected match %q for input %s", dec.RuleName, in) + } + } +} + +// The match key ("command") must not be matched as a value. +func TestMatchToolInputIgnoresObjectKeys(t *testing.T) { + ic := interceptorWith(testRule("tool_input", "regex", `^mkfs$`)) + in := toolInput(t, map[string]string{"mkfs": "ls"}) + if _, ok := ic.Match("Bash", in); ok { + t.Fatalf("object key matched as a value for input %s", in) + } +} + +func TestMatchToolNameUnchanged(t *testing.T) { + ic := interceptorWith(testRule("tool_name", "string", "WebFetch")) + if _, ok := ic.Match("WebFetch", []byte(`{"url":"http://x"}`)); !ok { + t.Fatal("tool_name rule did not match") + } + if _, ok := ic.Match("Bash", []byte(`{"command":"WebFetch"}`)); ok { + t.Fatal("tool_name rule matched tool input") + } +}