From b18095a8c9e8693a715aee1fe8a657f8d2eb7e8b Mon Sep 17 00:00:00 2001
From: Badgerbees <rapabelias@gmail.com>
Date: Fri, 3 Apr 2026 03:37:18 +0700
Subject: [PATCH] tools: tighten exec script preflight

---
 pkg/tools/shell.go      | 26 ++++++++++++++++++++------
 pkg/tools/shell_test.go | 31 ++++++++++++++++++++++++++++---
 2 files changed, 48 insertions(+), 9 deletions(-)

diff --git a/pkg/tools/shell.go b/pkg/tools/shell.go
index e219777d4..7df330ecd 100644
--- a/pkg/tools/shell.go
+++ b/pkg/tools/shell.go
@@ -113,6 +113,7 @@ var (
 	}
 
 	scriptPreflightEnvVarPattern = regexp.MustCompile(`\$[A-Z_][A-Z0-9_]{1,}`)
+	envAssignmentPattern         = regexp.MustCompile(`^[A-Za-z_][A-Za-z0-9_]*=`)
 	interpreterPipePattern       = regexp.MustCompile(
 		`(?i)(?:^|[|;&]\s*)(?:env\s+)?(?:python(?:\d+(?:\.\d+)?)?|node(?:js)?)\b`,
 	)
@@ -1155,6 +1156,10 @@ func (t *ExecTool) validateScriptFileForShellBleed(command, cwd string) string {
 			continue
 		}
 
+		if !shouldScanScriptForShellBleed(absPath) {
+			continue
+		}
+
 		if first := scriptPreflightEnvVarPattern.Find(content); len(first) > 0 {
 			return fmt.Sprintf(
 				"Command blocked by safety guard (exec preflight: detected likely shell variable injection (%s))",
@@ -1180,13 +1185,13 @@ func extractScriptTargetFromCommand(command string) []string {
 	interpreter := strings.ToLower(filepath.Base(argv[0]))
 	switch {
 	case isPythonInterpreter(interpreter):
-		target := findLastPositionalScriptArg(argv[1:], []string{".py"})
+		target := findFirstPositionalScriptArg(argv[1:], []string{".py"})
 		if target == "" {
 			return nil
 		}
 		return []string{target}
 	case isNodeInterpreter(interpreter):
-		target := findLastPositionalScriptArg(argv[1:], []string{".js"})
+		target := findFirstPositionalScriptArg(argv[1:], []string{".js"})
 		if target == "" {
 			return nil
 		}
@@ -1282,7 +1287,7 @@ func stripEnvPrefix(argv []string) []string {
 			idx++
 			continue
 		}
-		if strings.Contains(token, "=") && !strings.HasPrefix(token, "-") && !strings.ContainsAny(token, "/\\") {
+		if envAssignmentPattern.MatchString(token) && !strings.HasPrefix(token, "-") {
 			idx++
 			continue
 		}
@@ -1302,7 +1307,7 @@ func isNodeInterpreter(token string) bool {
 	return token == "node" || token == "nodejs"
 }
 
-func findLastPositionalScriptArg(tokens []string, suffixes []string) string {
+func findFirstPositionalScriptArg(tokens []string, suffixes []string) string {
 	if len(tokens) == 0 {
 		return ""
 	}
@@ -1355,6 +1360,11 @@ func hasScriptSuffix(token string, suffixes []string) bool {
 	return false
 }
 
+func shouldScanScriptForShellBleed(path string) bool {
+	ext := strings.ToLower(filepath.Ext(path))
+	return ext == ".py" || ext == ".pyw"
+}
+
 func shouldFailClosedInterpreterPreflight(command string) bool {
 	trimmed := strings.TrimSpace(command)
 	if trimmed == "" {
@@ -1367,8 +1377,12 @@ func shouldFailClosedInterpreterPreflight(command string) bool {
 	if interpreterShellWrapperPattern.MatchString(trimmed) {
 		return true
 	}
-	if interpreterPipePattern.MatchString(trimmed) && strings.ContainsAny(trimmed, "|;&") {
-		return true
+	if matches := interpreterPipePattern.FindAllStringIndex(trimmed, -1); len(matches) > 0 {
+		for _, match := range matches {
+			if match[0] > 0 {
+				return true
+			}
+		}
 	}
 
 	return false
diff --git a/pkg/tools/shell_test.go b/pkg/tools/shell_test.go
index 20959f9f2..949646775 100644
--- a/pkg/tools/shell_test.go
+++ b/pkg/tools/shell_test.go
@@ -686,12 +686,19 @@ func TestShellTool_ScriptPreflight(t *testing.T) {
 		want     string
 	}{
 		{
-			name:     "quoted script path validates content",
-			command:  `node "bad.js"`,
-			fileName: "bad.js",
+			name:     "quoted python script path validates content",
+			command:  `python "bad.py"`,
+			fileName: "bad.py",
 			content:  "const value = $DM_JSON;",
 			want:     "exec preflight: detected likely shell variable injection ($DM_JSON)",
 		},
+		{
+			name:     "env-prefixed interpreter validates content",
+			command:  `env PYTHONPATH=/tmp python bad.py`,
+			fileName: "bad.py",
+			content:  "payload = $DM_JSON",
+			want:     "exec preflight: detected likely shell variable injection ($DM_JSON)",
+		},
 		{
 			name:     "piped interpreter fails closed",
 			command:  "cat bad.py | python",
@@ -737,6 +744,24 @@ func TestShellTool_ScriptPreflight(t *testing.T) {
 	}
 }
 
+func TestShellTool_ScriptPreflight_AllowsDirectChainedInterpreter(t *testing.T) {
+	tmpDir := t.TempDir()
+	if err := os.WriteFile(filepath.Join(tmpDir, "good.py"), []byte("print('ok')"), 0o644); err != nil {
+		t.Fatalf("failed to write test script: %v", err)
+	}
+
+	tool, err := NewExecTool(tmpDir, false)
+	if err != nil {
+		t.Fatalf("unable to configure exec tool: %s", err)
+	}
+
+	result := tool.Execute(context.Background(), map[string]any{
+		"action":  "run",
+		"command": "python good.py && echo ok",
+	})
+	require.NotContains(t, result.ForLLM, "exec preflight:")
+}
+
 // TestShellTool_URLBypassPrevented verifies that a command cannot bypass the workspace
 // sandbox by smuggling a real path after a URL that contains the same //path substring.
 // e.g. "echo https://etc/passwd && cat //etc/passwd" must still be blocked.