386 lines
13 KiB
Python
Executable file
386 lines
13 KiB
Python
Executable file
#!/usr/bin/env python3
|
|
"""PreToolUse hook that blocks recursive ``grep``.
|
|
|
|
Recursive grep (``grep -r``/``-R``/``--recursive``/``-d recurse``/``rgrep`` ...)
|
|
walks an entire directory tree into memory and can OOM-kill the agent process on
|
|
large repos. The system prompt only *asks* the model to avoid it; this hook
|
|
turns that into a hard, deterministic block.
|
|
|
|
Protocol: read the PreToolUse envelope as JSON on stdin and signal the decision
|
|
to the runner:
|
|
|
|
recursive grep -> deny: exit 2 + a deny JSON on stdout (+ reason on stderr)
|
|
anything else -> allow: exit 0, nothing on stdout
|
|
|
|
Any unexpected condition falls through to "allow" (fail-open), matching the
|
|
runner's contract -- only an explicit deny blocks the tool call.
|
|
|
|
Detection is a pure function (``command_is_recursive``) with no I/O, so it is
|
|
trivially unit-testable -- run ``no-recursive-grep-guard.py --self-test``.
|
|
Parsing is a single quote-aware lexer (``lex``) shared by every stage:
|
|
|
|
* a quoted span is one operand and is never read as a flag
|
|
(``grep "rm -rf" log`` and ``grep "-r" file`` are not recursive);
|
|
* shell ``#`` line-comments are dropped;
|
|
* live command substitutions ``$(...)`` / backticks (unquoted or inside
|
|
double quotes -- single quotes suppress them) are recursed into;
|
|
* pipeline/compound operators (``| & ; ( ) { }`` + newlines) split segments,
|
|
so a recursive flag must belong to grep, not e.g. ``ls -R | grep``;
|
|
* ``sh``/``bash -c "<script>"`` inner scripts are recursed into;
|
|
* transparent wrappers (``sudo``/``env``/``xargs``/...) are peeled, looking
|
|
past their own flags/args for grep.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
import re
|
|
import sys
|
|
from typing import NamedTuple
|
|
|
|
SHELLS = {"sh", "bash", "dash", "zsh", "ash", "ksh", "mksh"}
|
|
# Wrappers whose real command may sit behind their own flags/args.
|
|
WRAPPERS = {
|
|
"sudo",
|
|
"doas",
|
|
"command",
|
|
"env",
|
|
"time",
|
|
"nice",
|
|
"nohup",
|
|
"stdbuf",
|
|
"exec",
|
|
"xargs",
|
|
"setsid",
|
|
}
|
|
GREPS = {"grep", "egrep", "fgrep"}
|
|
# grep options whose value is the following token (skip it so it is not a flag).
|
|
ARG_OPTS = {"-e", "-f", "-m", "-A", "-B", "-C", "-D"}
|
|
# Operators that separate pipeline segments / compounds (brace groups included).
|
|
OPERATOR_CHARS = set("|&;(){}\n\r")
|
|
ASSIGNMENT = re.compile(r"[A-Za-z_]\w*=")
|
|
# A short-flag cluster containing r/R -- grep's only r/R short flags both recurse.
|
|
SHORT_RECURSIVE = re.compile(r"-[A-Za-z]*[rR]")
|
|
MAX_DEPTH = 5
|
|
|
|
DENY_REASON = (
|
|
"Blocked: recursive grep (grep -r/-R/--recursive/rgrep) can read an entire "
|
|
"directory tree into memory and OOM-kill the agent process on large repos. Use the "
|
|
"dedicated search tool instead, which streams ripgrep results safely."
|
|
)
|
|
|
|
|
|
class Tok(NamedTuple):
|
|
"""One lexed token."""
|
|
|
|
text: str # token text, with surrounding quotes removed
|
|
quoted: bool # any part came from inside quotes -> operand, never a flag
|
|
op: bool # True if this is a shell operator that separates segments
|
|
|
|
|
|
def _read_paren_subst(s: str, i: int) -> tuple[int, str]:
|
|
"""Read a ``$(...)`` body starting at the ``$`` (index ``i``). Parens are
|
|
matched quote-aware so a ``)`` inside a string does not close it early.
|
|
Returns ``(index_after_closing_paren, body)``."""
|
|
n = len(s)
|
|
i += 2 # skip "$("
|
|
start = i
|
|
depth = 1
|
|
quote = None
|
|
while i < n and depth > 0:
|
|
c = s[i]
|
|
if quote:
|
|
if c == quote:
|
|
quote = None
|
|
elif c in ("'", '"'):
|
|
quote = c
|
|
elif c == "(":
|
|
depth += 1
|
|
elif c == ")":
|
|
depth -= 1
|
|
if depth == 0:
|
|
break
|
|
i += 1
|
|
return i + 1, s[start:i]
|
|
|
|
|
|
def _read_backtick_subst(s: str, i: int) -> tuple[int, str]:
|
|
"""Read a backtick substitution body starting at the opening backtick."""
|
|
n = len(s)
|
|
i += 1
|
|
start = i
|
|
while i < n and s[i] != "`":
|
|
i += 1
|
|
return i + 1, s[start:i]
|
|
|
|
|
|
def lex(s: str) -> tuple[list[Tok], list[str]]:
|
|
"""Single quote-aware pass over a command string. Returns the list of tokens
|
|
plus the bodies of any *live* command substitutions (to be recursed into)."""
|
|
tokens: list[Tok] = []
|
|
subst: list[str] = []
|
|
buf: list[str] = []
|
|
quoted = False
|
|
quote = None
|
|
i, n = 0, len(s)
|
|
|
|
def flush() -> None:
|
|
nonlocal buf, quoted
|
|
if buf or quoted:
|
|
tokens.append(Tok("".join(buf), quoted, False))
|
|
buf, quoted = [], False
|
|
|
|
while i < n:
|
|
c = s[i]
|
|
if quote == "'": # single quotes: everything literal, no substitution
|
|
if c == "'":
|
|
quote = None
|
|
else:
|
|
buf.append(c)
|
|
i += 1
|
|
continue
|
|
if quote == '"': # double quotes: literal text, but substitutions live
|
|
if c == '"':
|
|
quote = None
|
|
i += 1
|
|
elif c == "$" and i + 1 < n and s[i + 1] == "(":
|
|
i, body = _read_paren_subst(s, i)
|
|
subst.append(body)
|
|
elif c == "`":
|
|
i, body = _read_backtick_subst(s, i)
|
|
subst.append(body)
|
|
else:
|
|
buf.append(c)
|
|
i += 1
|
|
continue
|
|
# unquoted
|
|
if c in ("'", '"'):
|
|
quote = c
|
|
quoted = True
|
|
i += 1
|
|
elif c == "#" and not buf and not quoted:
|
|
break # comment to end of line
|
|
elif c == "$" and i + 1 < n and s[i + 1] == "(":
|
|
i, body = _read_paren_subst(s, i)
|
|
subst.append(body)
|
|
elif c == "`":
|
|
i, body = _read_backtick_subst(s, i)
|
|
subst.append(body)
|
|
elif c in " \t":
|
|
flush()
|
|
i += 1
|
|
elif c in OPERATOR_CHARS:
|
|
flush()
|
|
tokens.append(Tok(c, False, True))
|
|
i += 1
|
|
else:
|
|
buf.append(c)
|
|
i += 1
|
|
flush()
|
|
return tokens, subst
|
|
|
|
|
|
def _split_segments(tokens: list[Tok]) -> list[list[Tok]]:
|
|
"""Cut a token list into pipeline segments on operator tokens."""
|
|
segments: list[list[Tok]] = []
|
|
current: list[Tok] = []
|
|
for tok in tokens:
|
|
if tok.op:
|
|
if current:
|
|
segments.append(current)
|
|
current = []
|
|
else:
|
|
current.append(tok)
|
|
if current:
|
|
segments.append(current)
|
|
return segments
|
|
|
|
|
|
def _basename(text: str) -> str:
|
|
return text.rsplit("/", 1)[-1]
|
|
|
|
|
|
def _grep_args_recursive(args: list[Tok]) -> bool:
|
|
"""True if grep's argument list requests recursion."""
|
|
skip_arg = False # previous option consumes this token as its value
|
|
dir_arg = False # that consumed value belongs to -d / --directories
|
|
for arg in args:
|
|
if skip_arg:
|
|
if dir_arg and arg.text == "recurse":
|
|
return True
|
|
skip_arg = dir_arg = False
|
|
continue
|
|
if arg.quoted:
|
|
continue # a quoted operand is literal text, never a flag
|
|
t = arg.text
|
|
if t == "--":
|
|
return False # end of options; nothing after it is a flag
|
|
if t in ("-d", "--directories"):
|
|
skip_arg = dir_arg = True
|
|
elif t.startswith("--directories="):
|
|
if t.endswith("=recurse"):
|
|
return True
|
|
elif t in ARG_OPTS:
|
|
skip_arg = True
|
|
elif t in ("--recursive", "--dereference-recursive") or SHORT_RECURSIVE.match(t):
|
|
return True
|
|
return False
|
|
|
|
|
|
def _shell_script_recursive(args: list[Tok], depth: int) -> bool:
|
|
"""Re-inspect a shell interpreter's ``-c`` script (and trailing words)."""
|
|
i = 0
|
|
while i < len(args) and not args[i].quoted and args[i].text.startswith("-"):
|
|
i += 1 # skip the leading run of shell flags (-c, -lc, ...)
|
|
script = " ".join(a.text for a in args[i:])
|
|
return bool(script) and command_is_recursive(script, depth + 1)
|
|
|
|
|
|
def _segment_recursive(words: list[Tok], depth: int) -> bool:
|
|
"""True if one pipeline segment runs grep recursively."""
|
|
idx = 0
|
|
n = len(words)
|
|
while idx < n and not words[idx].quoted and ASSIGNMENT.match(words[idx].text):
|
|
idx += 1 # skip leading VAR=value assignments
|
|
if idx >= n:
|
|
return False
|
|
base = _basename(words[idx].text)
|
|
# Peel transparent wrappers: the real command may sit behind the wrapper's
|
|
# own flags/args, so seek the next grep/shell among the following words.
|
|
while base in WRAPPERS:
|
|
idx += 1
|
|
while idx < n:
|
|
b = _basename(words[idx].text)
|
|
if b == "rgrep" or b in GREPS or b in SHELLS:
|
|
break
|
|
idx += 1
|
|
if idx >= n:
|
|
return False
|
|
base = _basename(words[idx].text)
|
|
if base == "rgrep":
|
|
return True
|
|
if base in SHELLS:
|
|
return _shell_script_recursive(words[idx + 1 :], depth)
|
|
if base in GREPS:
|
|
return _grep_args_recursive(words[idx + 1 :])
|
|
return False
|
|
|
|
|
|
def command_is_recursive(cmd: str, depth: int = 0) -> bool:
|
|
"""Pure predicate: True if the shell command runs grep recursively."""
|
|
if depth >= MAX_DEPTH:
|
|
return False
|
|
tokens, substitutions = lex(cmd)
|
|
for body in substitutions:
|
|
if command_is_recursive(body, depth + 1):
|
|
return True
|
|
return any(_segment_recursive(seg, depth) for seg in _split_segments(tokens))
|
|
|
|
|
|
def extract_command(envelope: object) -> str | None:
|
|
"""Pull toolInput.command out of the PreToolUse envelope, or None."""
|
|
if not isinstance(envelope, dict):
|
|
return None
|
|
# Accept both the camelCase (toolInput) and snake_case (tool_input) shapes.
|
|
tool_input = envelope.get("toolInput") or envelope.get("tool_input")
|
|
if not isinstance(tool_input, dict):
|
|
return None
|
|
command = tool_input.get("command")
|
|
return command if isinstance(command, str) and command else None
|
|
|
|
|
|
# Allow/deny cases exercised by --self-test; keep in sync with the README.
|
|
SELF_TEST_CASES: list[tuple[str, bool]] = [
|
|
# --- recursive (deny) ---
|
|
("grep -r foo .", True),
|
|
("grep -R foo .", True),
|
|
("grep --recursive foo .", True),
|
|
("grep -rn TODO src", True),
|
|
("rgrep foo .", True),
|
|
("/usr/bin/grep -r x", True),
|
|
("egrep -r x .", True),
|
|
("FOO=bar grep -r x", True),
|
|
("cat x | grep -r y", True),
|
|
("sudo grep -r x", True),
|
|
("sudo -u root grep -r x", True),
|
|
("xargs -0 grep -r .", True),
|
|
("nice -n 10 grep -r x", True),
|
|
("grep -d recurse x", True),
|
|
('grep -d "recurse" x', True),
|
|
("grep --directories=recurse x", True),
|
|
('grep --include="*.rs" -r .', True),
|
|
('grep -e "p" -r .', True),
|
|
('bash -c "grep -r x"', True),
|
|
("bash -c grep -r x", True),
|
|
('echo "$(grep -r x)"', True),
|
|
("foo=$(grep -r x)", True),
|
|
("cat <(grep -r x)", True),
|
|
("{ grep -r x; }", True),
|
|
("grep -r x . # note", True),
|
|
# --- not recursive (allow) ---
|
|
("grep foo file", False),
|
|
("grep -n foo file", False),
|
|
("ls -R | grep foo", False),
|
|
("grep -e -r file", False),
|
|
("grep -- -r file", False),
|
|
("grep -A 3 foo file", False),
|
|
("grep -d skip foo file", False),
|
|
('grep "rm -rf" log', False),
|
|
('grep "-r" file', False),
|
|
('grep --include="*.rs" foo file', False),
|
|
("grep foo file # uses -r", False),
|
|
("echo '$(grep -r x)'", False),
|
|
('echo "$(ls)"', False),
|
|
("sudo ls -R", False),
|
|
("echo grep -r as text", False),
|
|
("xargs grep foo", False),
|
|
("{ echo hi; }", False),
|
|
("", False),
|
|
]
|
|
|
|
|
|
def self_test() -> int:
|
|
failures = 0
|
|
for cmd, want in SELF_TEST_CASES:
|
|
got = command_is_recursive(cmd)
|
|
if got != want:
|
|
failures += 1
|
|
print(f"FAIL: {cmd!r} -> {got} (want {want})")
|
|
total = len(SELF_TEST_CASES)
|
|
print(f"{total - failures}/{total} passed")
|
|
return 1 if failures else 0
|
|
|
|
|
|
def main() -> None:
|
|
if "--self-test" in sys.argv[1:]:
|
|
sys.exit(self_test())
|
|
try:
|
|
envelope = json.load(sys.stdin)
|
|
except (ValueError, OSError):
|
|
sys.exit(0) # unparseable input -> fail open (allow)
|
|
command = extract_command(envelope)
|
|
if command is None or not command_is_recursive(command):
|
|
sys.exit(0) # nothing to block -> silent allow
|
|
# Deny. Emit the grok-native decision (read by this repo's runner) and the
|
|
# Claude-style hookSpecificOutput for forward-compatibility, put the reason
|
|
# on stderr for runners that surface it there, and exit 2 so any exit-code
|
|
# based runner blocks too.
|
|
print(
|
|
json.dumps(
|
|
{
|
|
"decision": "deny",
|
|
"reason": DENY_REASON,
|
|
"hookSpecificOutput": {
|
|
"hookEventName": "PreToolUse",
|
|
"permissionDecision": "deny",
|
|
"permissionDecisionReason": DENY_REASON,
|
|
},
|
|
}
|
|
)
|
|
)
|
|
print(DENY_REASON, file=sys.stderr)
|
|
sys.exit(2)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|