__init__.py

"""Security Assistant For Hermes Plugin.

This plugin integrates with Alibaba Cloud Agent Protection Shield (APS)
to provide security checks for tool calls and LLM requests.
"""
from __future__ import annotations

import logging
from pathlib import Path
from typing import Any

# Plugin directory for state files
_PLUGIN_DIR = Path(__file__).parent

# Module-level logger
logger = logging.getLogger("hermes.plugin.security")

# Import sibling modules using relative imports (compatible with Hermes plugin loader)
from .config import load_config
from .auth_service import auth_service
from .check import (
    check_tool_call_request,
)
from .interceptor import HttpxInterceptor
from .asset_report import asset_report_service


# ------------------------------------------------------------------
# Plugin Registration Entry Point
# ------------------------------------------------------------------


def register(ctx):
    """Plugin registration entry point - called by Hermes plugin framework.

    Idempotent: safe to call multiple times (e.g. plugin hot-reload).
    - Hook registration dedup is the framework's responsibility.
    - auth_service / asset_report_service / interceptor handle repeated
      start()/install() internally.
    """
    config = load_config()
    auth_service.start(config)
    ctx.register_hook("pre_tool_call", on_pre_tool_call)
    ctx.register_hook("post_tool_call", on_post_tool_call)
    ctx.register_hook("pre_api_request", on_pre_api_request)
    ctx.register_hook("post_api_request", on_post_api_request)
    ctx.register_hook("pre_llm_call", on_pre_llm_call)
    ctx.register_hook("on_session_start", on_session_start)
    ctx.register_hook("on_session_end", on_session_end)

    # Install httpx monkey-patch for full LLM API interception
    _interceptor = HttpxInterceptor()
    _interceptor.install(auth_service)

    # Start periodic asset report service (non-blocking background thread)
    asset_report_service.start(config)

    logger.info(
        "[SECURITY] plugin registered (endpoint=%s, interceptor=%s)",
        config.endpoint_addr,
        "enabled" if _interceptor.installed else "disabled",
    )


# ------------------------------------------------------------------
# Hook Implementations
# ------------------------------------------------------------------

def on_pre_tool_call(
    tool_name: str,
    args: dict[str, Any],
    task_id: str | None,
    session_id: str | None,
    tool_call_id: str | None,
    **kwargs,
) -> dict[str, Any] | None:
    """Pre-tool-call hook - can block tool execution."""
    try:
        if not auth_service.is_ready():
            logger.info("[SECURITY] pre_tool_call: auth not ready, allow tool=%s", tool_name)
            return None
        logger.info("[SECURITY] pre_tool_call: checking tool=%s session=%s", tool_name, session_id)
        result = check_tool_call_request(
            tool_name=tool_name,
            parameters=args,
            session_id=session_id,
            task_id=task_id,
        )
        if result and result.get("action") == "block":
            logger.warning(
                "[SECURITY] pre_tool_call: BLOCKED tool=%s reason=%s",
                tool_name,
                result.get("content", ""),
            )
            return {
                "action": "block",
                "message": result.get(
                    "content",
                    "Tool call blocked by security policy",
                ),
            }
        logger.info("[SECURITY] pre_tool_call: allowed tool=%s action=%s", tool_name, result.get("action") if result else "none")
        return None
    except Exception as exc:
        logger.error("[SECURITY] pre_tool_call hook error (fail-open): %s", exc)
        return None


def on_post_tool_call(
    tool_name: str,
    args: dict[str, Any],
    result: Any,
    task_id: str | None,
    session_id: str | None,
    **kwargs,
) -> None:
    """Post-tool-call hook - audit logging only (observer, cannot modify result)."""
    try:
        logger.info(
            "[SECURITY] post_tool_call: tool=%s session=%s task=%s",
            tool_name, session_id, task_id,
        )
    except Exception as exc:
        logger.error("[SECURITY] post_tool_call hook error: %s", exc)
    return None


def on_pre_llm_call(
    session_id: str | None,
    user_message: str,
    conversation_history: list[dict[str, Any]] | None,
    model: str,
    platform: str,
    **kwargs,
) -> None:
    """Pre-LLM-call hook - audit logging only (LLM interception handled by httpx interceptor)."""
    try:
        msg_count = len(conversation_history) if conversation_history else 0
        logger.info(
            "[SECURITY] pre_llm_call: session=%s model=%s platform=%s messages=%d",
            session_id, model, platform, msg_count,
        )
    except Exception as exc:
        logger.error("[SECURITY] pre_llm_call hook error: %s", exc)
    return None


def on_pre_api_request(
    task_id: str | None,
    session_id: str | None,
    model: str,
    provider: str,
    api_call_count: int,
    approx_input_tokens: int | None,
    **kwargs,
) -> None:
    """Pre-API-request hook - lightweight audit logging."""
    try:
        logger.debug(
            "API_REQUEST session=%s model=%s provider=%s call#=%d tokens~%d",
            session_id, model, provider, api_call_count, approx_input_tokens or 0,
        )
    except Exception as exc:
        logger.error("pre_api_request hook error: %s", exc)
    return None


def on_post_api_request(
    task_id: str | None,
    session_id: str | None,
    model: str,
    provider: str,
    api_duration: float | None,
    finish_reason: str | None,
    usage: dict[str, Any] | None,
    **kwargs,
) -> None:
    """Post-API-request hook - lightweight audit logging."""
    try:
        logger.debug(
            "API_RESPONSE session=%s model=%s duration=%.1fs finish=%s usage=%s",
            session_id, model, api_duration or 0, finish_reason, usage,
        )
    except Exception as exc:
        logger.error("post_api_request hook error: %s", exc)
    return None


def on_session_start(
    session_id: str | None,
    model: str,
    platform: str,
    **kwargs,
) -> None:
    """Session start hook - logs session initialization."""
    try:
        logger.info(
            "[SECURITY] session_start: session=%s model=%s platform=%s",
            session_id, model, platform,
        )
    except Exception as exc:
        logger.error("on_session_start hook error: %s", exc)
    return None


def on_session_end(
    session_id: str | None,
    completed: bool,
    interrupted: bool,
    **kwargs,
) -> None:
    """Session end hook - logs session termination."""
    try:
        logger.info(
            "[SECURITY] session_end: session=%s completed=%s",
            session_id, completed,
        )
    except Exception as exc:
        logger.error("on_session_end hook error: %s", exc)
    return None