Merge pull request #19 from kbdevs/fix/inline-python-rewrite

fix: avoid rewriting quoted inline commands

Author: edouard-claude

internal/initcmd/init.go

@@ -22,6 +22,75 @@ fi
 
 set -euo pipefail
 
+leading_ws_len() {
+  local input="$1"
+  local len=${#input}
+  local i=0
+
+  while [ $i -lt $len ]; do
+    case "${input:$i:1}" in
+      [[:space:]]) i=$((i + 1)) ;;
+      *) break ;;
+    esac
+  done
+
+  printf '%s' "$i"
+}
+
+trailing_ws_len() {
+  local input="$1"
+  local i=$((${#input} - 1))
+  local count=0
+
+  while [ $i -ge 0 ]; do
+    case "${input:$i:1}" in
+      [[:space:]])
+        count=$((count + 1))
+        i=$((i - 1))
+        ;;
+      *) break ;;
+    esac
+  done
+
+  printf '%s' "$count"
+}
+
+extract_first_segment() {
+  local input="$1"
+  local len=${#input}
+  local i=0
+  local quote=""
+  local ch
+
+  while [ $i -lt $len ]; do
+    ch="${input:$i:1}"
+
+    if [ -n "$quote" ]; then
+      if [ "$ch" = "\\" ] && [ "$quote" = '"' ]; then
+        i=$((i + 2))
+        continue
+      fi
+
+      if [ "$ch" = "$quote" ]; then
+        quote=""
+      fi
+
+      i=$((i + 1))
+      continue
+    fi
+
+    case "$ch" in
+      "'") quote="'" ;;
+      '"') quote='"' ;;
+      ';'|'|'|'&') break ;;
+    esac
+
+    i=$((i + 1))
+  done
+
+  printf '%s' "${input:0:i}"
+}
+
 INPUT=$(cat)
 CMD=$(echo "$INPUT" | jq -r '.tool_input.command // empty')
 
@@ -30,17 +99,29 @@ if [ -z "$CMD" ]; then
   exit 0
 fi
 
-# Extract the first command (before && or | or ;)
-# head -1 prevents xargs from seeing heredoc body lines with unmatched quotes
-FIRST_CMD=$(echo "$CMD" | head -1 | sed 's/[;&|].*//' | sed 's/^[[:space:]]*//;s/[[:space:]]*$//')
+# Extract the first command segment, ignoring separators inside quotes.
+# head -1 keeps heredoc bodies out of the scan.
+FIRST_LINE=$(printf '%s\n' "$CMD" | head -1)
+FIRST_SEGMENT=$(extract_first_segment "$FIRST_LINE")
+LEADING_WS_LEN=$(leading_ws_len "$FIRST_SEGMENT")
+TRAILING_WS_LEN=$(trailing_ws_len "$FIRST_SEGMENT")
+FIRST_CMD_LEN=$((${#FIRST_SEGMENT} - LEADING_WS_LEN - TRAILING_WS_LEN))
+if [ $FIRST_CMD_LEN -lt 0 ]; then
+  FIRST_CMD_LEN=0
+fi
+FIRST_PREFIX="${FIRST_SEGMENT:0:LEADING_WS_LEN}"
+FIRST_CMD="${FIRST_SEGMENT:LEADING_WS_LEN:FIRST_CMD_LEN}"
+FIRST_SUFFIX_START=$((LEADING_WS_LEN + FIRST_CMD_LEN))
+FIRST_SUFFIX="${FIRST_SEGMENT:FIRST_SUFFIX_START}"
 
 # Skip if already using snip
 case "$FIRST_CMD" in
   snip\ *|*/snip\ *) exit 0 ;;
 esac
 
 # Strip leading env var assignments (e.g. CGO_ENABLED=0 go test)
-BARE_CMD=$(echo "$FIRST_CMD" | sed 's/^[A-Za-z_][A-Za-z0-9_]*=[^ ]* *//')
+ENV_PREFIX=$(printf '%s' "$FIRST_CMD" | sed -E 's/^(([A-Za-z_][A-Za-z0-9_]*=[^[:space:]]+[[:space:]]*)*).*/\1/')
+BARE_CMD="${FIRST_CMD:${#ENV_PREFIX}}"
 
 # Extract the base command name
 BASE=$(echo "$BARE_CMD" | awk '{print $1}')
@@ -52,7 +133,8 @@ case "$BASE" in
     # Rewrite: prefix with "snip --" so flags like --help or --version in the
     # original command are passed verbatim to the underlying tool, not parsed
     # by snip itself.
-    REWRITE=$(echo "$CMD" | sed "s|$BARE_CMD|snip -- $BARE_CMD|")
+    REST="${CMD:${#FIRST_SEGMENT}}"
+    REWRITE="${FIRST_PREFIX}${ENV_PREFIX}snip -- ${BARE_CMD}${FIRST_SUFFIX}${REST}"
     ;;
 esac
 

internal/initcmd/init_test.go

@@ -196,38 +196,33 @@ func TestUnpatchPreservesOtherHooks(t *testing.T) {
 	}
 }
 
-// TestHookScriptMultilineCommand verifies that the installed hook script handles
-// multiline commands (e.g. git commit with a heredoc) without error.
-// Previously, xargs was used to trim whitespace from FIRST_CMD and would fail
-// with exit 1 on unmatched quotes present in heredoc body lines.
-func TestHookScriptMultilineCommand(t *testing.T) {
+func runHookScript(t *testing.T, cmd string) string {
+	t.Helper()
+
 	if _, err := exec.LookPath("bash"); err != nil {
 		t.Skip("bash not available")
 	}
 	if _, err := exec.LookPath("jq"); err != nil {
 		t.Skip("jq not available")
 	}
-	if _, err := exec.LookPath("snip"); err != nil {
-		t.Skip("snip not available")
-	}
 
-	// Write the hook to a temp file (simulates snip init).
 	dir := t.TempDir()
 	hookPath := filepath.Join(dir, "snip-rewrite.sh")
 	if err := os.WriteFile(hookPath, []byte(hookScript), 0755); err != nil {
 		t.Fatalf("write hook: %v", err)
 	}
+	snipPath := filepath.Join(dir, "snip")
+	if err := os.WriteFile(snipPath, []byte("#!/bin/sh\nexit 0\n"), 0755); err != nil {
+		t.Fatalf("write fake snip: %v", err)
+	}
 
-	// Simulate the JSON Claude Code sends for a heredoc-style git commit.
-	// The multiline command contains an unmatched `)"` on the last line,
-	// which caused xargs to exit 1 (unmatched double quote).
-	cmd := "git add file.go && git commit -m \"$(cat <<'EOF'\n   fix: something\n\n   Co-Authored-By: Bot <bot@example.com>\n   EOF\n   )\""
 	payload, _ := json.Marshal(map[string]any{
 		"tool_name":  "Bash",
 		"tool_input": map[string]any{"command": cmd},
 	})
 
 	proc := exec.Command("bash", hookPath)
+	proc.Env = append(os.Environ(), "PATH="+dir+string(os.PathListSeparator)+os.Getenv("PATH"))
 	proc.Stdin = strings.NewReader(string(payload))
 	output, runErr := proc.Output()
 	if runErr != nil {
@@ -242,12 +237,43 @@ func TestHookScriptMultilineCommand(t *testing.T) {
 	hookOut, _ := result["hookSpecificOutput"].(map[string]any)
 	updated, _ := hookOut["updatedInput"].(map[string]any)
 	rewritten, _ := updated["command"].(string)
+	return rewritten
+}
+
+// TestHookScriptMultilineCommand verifies that the installed hook script handles
+// multiline commands (e.g. git commit with a heredoc) without error.
+// Previously, xargs was used to trim whitespace from FIRST_CMD and would fail
+// with exit 1 on unmatched quotes present in heredoc body lines.
+func TestHookScriptMultilineCommand(t *testing.T) {
+	// Simulate the JSON Claude Code sends for a heredoc-style git commit.
+	// The multiline command contains an unmatched `)"` on the last line,
+	// which caused xargs to exit 1 (unmatched double quote).
+	cmd := "git add file.go && git commit -m \"$(cat <<'EOF'\n   fix: something\n\n   Co-Authored-By: Bot <bot@example.com>\n   EOF\n   )\""
+	rewritten := runHookScript(t, cmd)
 
 	if !strings.HasPrefix(rewritten, "snip -- git add ") {
 		t.Errorf("expected rewritten command to start with 'snip -- git add', got: %s", rewritten)
 	}
 }
 
+func TestHookScriptInlinePythonDoesNotRewriteQuotedSemicolons(t *testing.T) {
+	cmd := "git commit -m \"$(python3 -c \\\"from pathlib import Path; import sys; print(Path('.').name); print(sys.version)\\\")\" && git status"
+	rewritten := runHookScript(t, cmd)
+
+	if !strings.HasPrefix(rewritten, "snip -- git commit ") {
+		t.Fatalf("expected rewritten command to start with 'snip -- git commit', got: %s", rewritten)
+	}
+	if strings.Count(rewritten, "snip --") != 1 {
+		t.Fatalf("expected exactly one snip injection, got: %s", rewritten)
+	}
+	if strings.Contains(rewritten, "; snip") {
+		t.Fatalf("expected inline python to stay unchanged, got: %s", rewritten)
+	}
+	if strings.Contains(rewritten, "python3 snip") {
+		t.Fatalf("expected inline python command to stay unchanged, got: %s", rewritten)
+	}
+}
+
 func readSettings(t *testing.T, path string) map[string]any {
 	t.Helper()
 	data, err := os.ReadFile(path)