hooks.py

"""Hook system and permission / safety checks.

Hooks are callbacks registered for lifecycle events (PreToolUse, PostToolUse,
Stop).  The permission layer uses deny-lists and rule-based checks to block
dangerous tool calls before execution.
"""

from collections import defaultdict
from enum import Enum
from typing import Callable, List

from config import WORKDIR, get_logger

logger = get_logger(__name__)


# ── Deny list ──

DENY_LIST = ["rm -rf /", "sudo", "shutdown", "reboot", "mkfs", "dd if=", "> /dev/sda"]


def check_deny_list(command: str) -> str | None:
    """Return an error message if *command* contains a forbidden substring."""
    for pattern in DENY_LIST:
        if pattern in command:
            return f"Error: Dangerous command blocked: {command}"
    return None


# ── Permission rules ──

PERMISSION_RULES = [
    {
        "tools": ["run_write_file", "run_edit_file"],
        "check": lambda args: not args.get("path", "")
        or not (WORKDIR / args.get("path", ""))
        .resolve()
        .is_relative_to(WORKDIR),
        "message": "Writing outside workspace or empty path",
    },
    {
        "tools": ["run_bash"],
        "check": lambda args: any(
            kw in args.get("command", "") for kw in ["rm ", "> /etc/", "chmod 777"]
        ),
        "message": "Potentially destructive command",
    },
]


def check_rules(tool_name: str, args: dict) -> str | None:
    """Return the first matching permission-rule message, or None."""
    for rule in PERMISSION_RULES:
        if tool_name in rule["tools"] and rule["check"](args):
            return rule["message"]
    return None


def ask_user(tool_name: str, args: dict, reason: str) -> str:
    """Prompt the user in the terminal to allow or deny a tool call."""
    logger.warning(f"⚠  {reason}")
    logger.warning(f"   Tool: {tool_name}({args})")
    choice = input("   Allow? [y/N] ").strip().lower()
    return "allow" if choice in ("y", "yes") else "deny"


# ── Hook infrastructure ──


class HookType(Enum):
    UserPromptSubmit = "UserPromptSubmit"
    PreToolUse = "PreToolUse"
    PostToolUse = "PostToolUse"
    Stop = "Stop"


class HookSystem:
    """Central hook registry and dispatcher.

    Hooks are callbacks registered for lifecycle events.  Built-in hooks
    (permission checks, logging, summary) are registered automatically.

    Usage::

        hooks = HookSystem()
        hooks.trigger(HookType.PreToolUse, tool_call)
    """

    def __init__(self) -> None:
        self._hooks: dict[HookType, List[Callable]] = defaultdict(list)
        self._register_builtins()

    def register(self, event: HookType, callback: Callable) -> None:
        """Register *callback* to fire on *event*."""
        self._hooks[event].append(callback)

    def trigger(self, event: HookType, *args) -> object | None:
        """Fire all callbacks for *event* in registration order.

        Returns the first non-None result (if any), otherwise None.
        """
        for callback in self._hooks[event]:
            result = callback(*args)
            if result is not None:
                return result
        return None

    # ── Built-in registration ───────────────────────────────────────────

    def _register_builtins(self) -> None:
        """Register the default set of hook callbacks."""
        self.register(HookType.PreToolUse, _check_permission_hook)
        self.register(HookType.PostToolUse, _log_output_hook)
        self.register(HookType.Stop, _summary_hook)


# ── Built-in hook callbacks ──


def _check_permission_hook(block) -> str | None:
    """PreToolUse hook: deny dangerous commands and prompt user for risky ones."""
    import json

    args = json.loads(block.function.arguments)
    if block.function.name == "run_bash":
        reason = check_deny_list(args.get("command", ""))
        if reason:
            logger.error(f"⛔ {reason}")
            return "Permission denied by deny list"
    reason = check_rules(block.function.name, args)
    if reason:
        decision = ask_user(block.function.name, args, reason)
        if decision == "deny":
            return "Permission denied by user"
    return None


def _log_output_hook(block, output) -> None:
    """PostToolUse hook: log tool calls and their results."""
    name = block.function.name
    args = block.function.arguments
    logger.info(f"[HOOK] {name}({args}) -> {output}")
    return None


def _summary_hook(message: list) -> None:
    """Stop hook: print a summary of tool usage for the session."""
    tool_count = sum(1 for msg in message if msg.get("role") == "tool")
    logger.info(f"[HOOK] Stop: session used {tool_count} tool calls")
    return None