05_hooks.py

"""05: Hooks. PreToolUse blocks writes to .env files; PostToolUse logs every tool result."""

import asyncio

from claude_agent_sdk import ClaudeAgentOptions, ClaudeSDKClient, HookMatcher, ResultMessage, TextBlock, AssistantMessage, UserMessage, ToolResultBlock
from demo_util import PROJECT

audit_log = []


async def protect_env_files(input_data, tool_use_id, context):
    file_path = input_data["tool_input"].get("file_path", "")
    if file_path.split("/")[-1] in (".env", ".env.local"):
        return {
            "hookSpecificOutput": {
                "hookEventName": "PreToolUse",
                "permissionDecision": "deny",
                "permissionDecisionReason": "Policy: .env files may not be created or modified by the agent",
            }
        }
    return {}  # empty dict = no opinion, let the normal permission flow decide


async def audit(input_data, tool_use_id, context):
    audit_log.append((input_data["tool_name"], tool_use_id))
    return {}


async def main():
    options = ClaudeAgentOptions(
        cwd=PROJECT,
        permission_mode="acceptEdits",
        allowed_tools=["Read", "Write", "Edit", "Glob"],
        max_turns=6,
        hooks={
            "PreToolUse": [HookMatcher(matcher="Write|Edit", hooks=[protect_env_files])],
            "PostToolUse": [HookMatcher(hooks=[audit])],  # no matcher: every tool
        },
    )
    async with ClaudeSDKClient(options=options) as client:
        await client.query("Create a .env file containing DATABASE_URL=sqlite:///users.db, "
                           "then create config.example.txt with the same line.")
        async for message in client.receive_response():
            if isinstance(message, AssistantMessage):
                for block in message.content:
                    if isinstance(block, TextBlock):
                        print(f"Claude: {block.text}")
            elif isinstance(message, UserMessage) and isinstance(message.content, list):
                for block in message.content:
                    if isinstance(block, ToolResultBlock):
                        print(f"[tool_result]{' ERROR' if block.is_error else ''} {str(block.content)[:200]}")
            elif isinstance(message, ResultMessage):
                print(f"[result] {message.subtype}, denials={message.permission_denials}")

    print("\nAudit log (PostToolUse):", audit_log)


asyncio.run(main())