Merge pull request #152 from RuoJi6/fix/intercept-match-decoded-input

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 {

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 <<EOF\nDROP TABLE users;\nEOF"}},
+		{"FLUSHALL inside a heredoc", builtinFlush, "Bash", map[string]string{"command": "redis-cli <<EOF\nFLUSHALL\nEOF"}},
+		{"shutdown on its own line", builtinV2, "Bash", map[string]string{"command": "sync\nshutdown -h now"}},
+		// encoding/json writes > 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")
+	}
+}