lijiaoqiao/agent/hook_profile.py

"""
Hook Profile System

Manages different security profiles for command hooks.
Users can switch between profiles to adjust the level of
security enforcement.

Inspired by gstack's hook patterns (careful/freeze/guard).

Usage:
    from agent.hook_profile import HookProfileManager, get_profile_manager
    
    # Get the manager
    manager = get_profile_manager()
    
    # Get current profile
    profile = manager.get_active_profile()
    
    # Switch to strict mode
    manager.switch_profile("strict")
    
    # Check if a command would be blocked
    result = manager.check_command("rm -rf /")
"""

import yaml
import logging
import subprocess
from pathlib import Path
from typing import Dict, List, Optional, Any, Callable
from dataclasses import dataclass, field
from enum import Enum

from hermes_constants import get_hermes_home

logger = logging.getLogger(__name__)


class ProfileLevel(Enum):
    """Hook security profile levels."""
    MINIMAL = "minimal"
    STANDARD = "standard"
    STRICT = "strict"
    
    def __str__(self):
        return self.value


@dataclass
class HookDefinition:
    """Definition of a hook."""
    name: str
    hook_type: str  # "pre-command" or "post-write"
    script_path: str
    enabled: bool = True
    config_path: Optional[str] = None  # Optional YAML config
    description: str = ""
    blocks_on_match: bool = False  # Does this hook block commands?
    
    @property
    def full_path(self) -> Path:
        """Get the full path to the hook script."""
        return Path.home() / ".hermes" / "hooks" / self.hook_type / self.script_path


@dataclass
class HookProfile:
    """A named hook profile with associated hooks."""
    name: str
    level: ProfileLevel
    description: str
    hooks: List[str]  # Hook names to enable
    env_overrides: Dict[str, str] = field(default_factory=dict)
    auto_activate: bool = False  # Auto-activate based on context
    
    @property
    def hook_definitions(self) -> List[HookDefinition]:
        """Get HookDefinition objects for this profile's hooks."""
        return [BUILTIN_HOOKS.get(h) for h in self.hooks if h in BUILTIN_HOOKS]


# =============================================================================
# Built-in Hook Profiles
# =============================================================================

BUILTIN_HOOKS: Dict[str, HookDefinition] = {
    "careful": HookDefinition(
        name="careful",
        hook_type="pre-command",
        script_path="careful-hook.sh",
        blocks_on_match=False,
        description="Warn about dangerous commands but allow execution",
    ),
    "freeze": HookDefinition(
        name="freeze",
        hook_type="pre-command",
        script_path="freeze-hook.sh",
        blocks_on_match=True,
        description="Block dangerous commands entirely",
    ),
    "guard": HookDefinition(
        name="guard",
        hook_type="pre-command",
        script_path="guard-hook.sh",
        blocks_on_match=False,
        description="Context-aware guard with learning",
    ),
    "rust-fmt": HookDefinition(
        name="rust-fmt",
        hook_type="post-write",
        script_path="rust-fmt-on-rs.yaml",
        blocks_on_match=False,
        description="Format Rust code on save",
    ),
    "cargo-check": HookDefinition(
        name="cargo-check",
        hook_type="post-write",
        script_path="cargo-check-on-rs.yaml",
        blocks_on_match=True,
        description="Check Rust code on save",
    ),
    "go-fmt": HookDefinition(
        name="go-fmt",
        hook_type="post-write",
        script_path="go-fmt-on-go.yaml",
        blocks_on_match=False,
        description="Format Go code on save",
    ),
    "ruff-lint": HookDefinition(
        name="ruff-lint",
        hook_type="post-write",
        script_path="ruff-lint-py.yaml",
        blocks_on_match=True,
        description="Lint Python code on save with ruff",
    ),
    "pytest": HookDefinition(
        name="pytest",
        hook_type="post-write",
        script_path="pytest-on-py.yaml",
        blocks_on_match=True,
        description="Run pytest on Python file change",
    ),
}


BUILTIN_PROFILES: Dict[str, HookProfile] = {
    "minimal": HookProfile(
        name="minimal",
        level=ProfileLevel.MINIMAL,
        description="No security hooks - use with caution",
        hooks=[],  # No hooks
        env_overrides={"HERMES_HOOKS_ENABLED": "false"},
    ),
    "standard": HookProfile(
        name="standard",
        level=ProfileLevel.STANDARD,
        description="Warning-only hooks for dangerous commands",
        hooks=["careful"],
        env_overrides={"HERMES_HOOKS_ENABLED": "true"},
        auto_activate=False,
    ),
    "strict": HookProfile(
        name="strict",
        level=ProfileLevel.STRICT,
        description="Block dangerous commands - recommended for production",
        hooks=["careful", "freeze"],
        env_overrides={"HERMES_HOOKS_ENABLED": "true"},
        auto_activate=False,
    ),
    "developer": HookProfile(
        name="developer",
        level=ProfileLevel.STANDARD,
        description="Standard hooks + code quality checks (Rust/Go/Python)",
        hooks=["careful", "ruff-lint", "pytest", "go-fmt"],
        env_overrides={"HERMES_HOOKS_ENABLED": "true"},
        auto_activate=False,
    ),
    "rust-dev": HookProfile(
        name="rust-dev",
        level=ProfileLevel.STANDARD,
        description="Rust development profile with formatting and checks",
        hooks=["careful", "rust-fmt", "cargo-check"],
        env_overrides={"HERMES_HOOKS_ENABLED": "true"},
        auto_activate=False,
    ),
}


# =============================================================================
# Hook Profile Manager
# =============================================================================

class HookProfileManager:
    """
    Manages hook security profiles.
    
    Allows switching between different hook configurations
    to adjust the level of security enforcement.
    """
    
    def __init__(self, config_path: Optional[Path] = None):
        """
        Initialize the hook profile manager.
        
        Args:
            config_path: Path to store profile config.
                        Defaults to ~/.hermes/hook-profile.yaml
        """
        if config_path is None:
            config_path = get_hermes_home() / "hook-profile.yaml"
        
        self.config_path = config_path
        self._active_profile_name: Optional[str] = "standard"
        self._custom_profiles: Dict[str, HookProfile] = {}
        self._load()
    
    def _load(self) -> None:
        """Load profile configuration from disk."""
        if self.config_path.exists():
            try:
                data = yaml.safe_load(self.config_path.read_text())
                self._active_profile_name = data.get("active_profile", "standard")
                
                # Load custom profiles
                custom = data.get("custom_profiles", {})
                for name, profile_data in custom.items():
                    self._custom_profiles[name] = HookProfile(
                        name=profile_data.get("name", name),
                        level=ProfileLevel(profile_data.get("level", "standard")),
                        description=profile_data.get("description", ""),
                        hooks=profile_data.get("hooks", []),
                        env_overrides=profile_data.get("env_overrides", {}),
                        auto_activate=profile_data.get("auto_activate", False),
                    )
                
                logger.info(f"Loaded hook profile config: active={self._active_profile_name}")
            except Exception as e:
                logger.warning(f"Failed to load hook profile config: {e}")
                self._active_profile_name = "standard"
    
    def _save(self) -> None:
        """Save profile configuration to disk."""
        try:
            data = {
                "active_profile": self._active_profile_name,
                "custom_profiles": {},
            }
            
            for name, profile in self._custom_profiles.items():
                data["custom_profiles"][name] = {
                    "name": profile.name,
                    "level": profile.level.value,
                    "description": profile.description,
                    "hooks": profile.hooks,
                    "env_overrides": profile.env_overrides,
                    "auto_activate": profile.auto_activate,
                }
            
            self.config_path.parent.mkdir(parents=True, exist_ok=True)
            self.config_path.write_text(yaml.dump(data))
            logger.debug(f"Saved hook profile config")
        except Exception as e:
            logger.error(f"Failed to save hook profile config: {e}")
    
    def get_active_profile(self) -> HookProfile:
        """Get the currently active profile."""
        if self._active_profile_name in BUILTIN_PROFILES:
            return BUILTIN_PROFILES[self._active_profile_name]
        if self._active_profile_name in self._custom_profiles:
            return self._custom_profiles[self._active_profile_name]
        
        # Fallback to standard
        return BUILTIN_PROFILES["standard"]
    
    def get_profile(self, name: str) -> Optional[HookProfile]:
        """Get a profile by name."""
        if name in BUILTIN_PROFILES:
            return BUILTIN_PROFILES[name]
        return self._custom_profiles.get(name)
    
    def list_profiles(self) -> List[str]:
        """List all available profile names."""
        profiles = list(BUILTIN_PROFILES.keys())
        profiles.extend(self._custom_profiles.keys())
        return profiles
    
    def switch_profile(self, name: str) -> bool:
        """
        Switch to a different hook profile.
        
        Args:
            name: Profile name to switch to
            
        Returns:
            True if switch successful, False if profile doesn't exist
        """
        if name not in BUILTIN_PROFILES and name not in self._custom_profiles:
            logger.warning(f"Hook profile '{name}' not found")
            return False
        
        old_profile = self._active_profile_name
        self._active_profile_name = name
        self._save()
        
        logger.info(f"Switched hook profile: {old_profile} → {name}")
        return True
    
    def check_command(self, command: str, full_command: Optional[str] = None) -> Dict[str, Any]:
        """
        Check a command against the active profile's hooks.
        
        Args:
            command: The command to check
            full_command: Optional full command with args
            
        Returns:
            Dict with keys:
                - allowed: bool
                - blocked: bool
                - warned: bool
                - messages: List[str]
                - hook_name: str or None
        """
        if full_command is None:
            full_command = command
        
        profile = self.get_active_profile()
        
        result = {
            "allowed": True,
            "blocked": False,
            "warned": False,
            "messages": [],
            "hook_name": None,
        }
        
        for hook_name in profile.hooks:
            hook_def = BUILTIN_HOOKS.get(hook_name)
            if not hook_def or not hook_def.enabled:
                continue
            
            hook_result = self._run_hook(hook_def, command, full_command)
            
            if hook_result is None:
                continue
            
            if hook_result["blocked"]:
                result["blocked"] = True
                result["allowed"] = False
                result["hook_name"] = hook_name
                result["messages"].extend(hook_result["messages"])
                break
            
            if hook_result["warned"]:
                result["warned"] = True
                result["messages"].extend(hook_result["messages"])
        
        return result
    
    def _run_hook(self, hook: HookDefinition, command: str, full_command: str) -> Optional[Dict[str, Any]]:
        """
        Run a single hook and return the result.
        
        Args:
            hook: Hook definition
            command: Short command (first word)
            full_command: Full command string
            
        Returns:
            Dict with blocked/warned/messages, or None if hook not found
        """
        if not hook.full_path.exists():
            logger.warning(f"Hook script not found: {hook.full_path}")
            return None
        
        try:
            proc = subprocess.run(
                [str(hook.full_path), command, full_command],
                capture_output=True,
                text=True,
                timeout=5,
            )
            
            output = proc.stdout.strip()
            stderr = proc.stderr.strip()
            
            if output:
                messages = output.split("\n")
            else:
                messages = []
            
            if stderr:
                messages.append(f"STDERR: {stderr}")
            
            return {
                "blocked": proc.returncode != 0 and hook.blocks_on_match,
                "warned": proc.returncode == 0 and bool(output),
                "messages": messages,
            }
        except subprocess.TimeoutExpired:
            logger.error(f"Hook timed out: {hook.name}")
            return {"blocked": False, "warned": True, "messages": ["Hook timed out"]}
        except Exception as e:
            logger.error(f"Hook failed: {hook.name}: {e}")
            return {"blocked": False, "warned": True, "messages": [f"Hook error: {e}"]}
    
    def add_profile(self, profile: HookProfile) -> None:
        """
        Add a custom profile.
        
        Args:
            profile: Profile to add
        """
        self._custom_profiles[profile.name] = profile
        self._save()
    
    def remove_profile(self, name: str) -> bool:
        """
        Remove a custom profile.
        
        Args:
            name: Profile name to remove
            
        Returns:
            True if removed, False if not found or is builtin
        """
        if name in BUILTIN_PROFILES:
            logger.warning(f"Cannot remove builtin profile: {name}")
            return False
        
        if name in self._custom_profiles:
            del self._custom_profiles[name]
            if self._active_profile_name == name:
                self._active_profile_name = "standard"
            self._save()
            return True
        
        return False
    
    def auto_select_profile(self, context: str) -> HookProfile:
        """
        Automatically select profile based on context.

        Args:
            context: Conversation context

        Returns:
            Selected profile
        """
        context_lower = context.lower()

        # Check for keywords that suggest different profiles
        if any(kw in context_lower for kw in ["production", "prod", "deploy", "release"]):
            self.switch_profile("strict")
            return self.get_active_profile()

        if any(kw in context_lower for kw in ["develop", "dev", "local", "test"]):
            self.switch_profile("developer")
            return self.get_active_profile()

        if any(kw in context_lower for kw in ["rust", "cargo", "rs "]):
            self.switch_profile("rust-dev")
            return self.get_active_profile()

        # Default to current profile
        return self.get_active_profile()
    
    def get_enabled_hooks(self) -> List[HookDefinition]:
        """Get list of enabled hooks for the active profile."""
        profile = self.get_active_profile()
        enabled = []
        
        for hook_name in profile.hooks:
            hook_def = BUILTIN_HOOKS.get(hook_name)
            if hook_def:
                enabled.append(hook_def)
        
        return enabled


# =============================================================================
# Singleton
# =============================================================================

_hook_profile_manager: Optional[HookProfileManager] = None


def get_profile_manager() -> HookProfileManager:
    """Get the singleton HookProfileManager instance."""
    global _hook_profile_manager
    if _hook_profile_manager is None:
        _hook_profile_manager = HookProfileManager()
    return _hook_profile_manager


def check_command(command: str, full_command: Optional[str] = None) -> Dict[str, Any]:
    """
    Check a command against the active profile's hooks.
    
    Args:
        command: The command to check
        full_command: Optional full command with args
        
    Returns:
        Dict with allowed/blocked/warned/messages
    """
    return get_profile_manager().check_command(command, full_command)


def switch_hook_profile(name: str) -> bool:
    """
    Switch to a different hook profile.
    
    Args:
        name: Profile name
        
    Returns:
        True if successful
    """
    return get_profile_manager().switch_profile(name)


def get_active_hook_profile() -> HookProfile:
    """Get the currently active hook profile."""
    return get_profile_manager().get_active_profile()