ITWorx_MCP_Hub-Public/packages/hub_policy/governance.py

"""Shared decision, rate-limit, and audit lifecycle used at every invocation boundary."""

from __future__ import annotations

import inspect
import time
from collections.abc import Callable, Mapping
from functools import wraps
from typing import Any, TypeVar, cast

from .audit import AuditEvent, AuditLog
from .engine import Decision, DecisionTrace, InvocationContext, PolicyEngine
from .rate_limit import RateLimitExceeded, SlidingWindowRateLimiter


class PolicyDenied(PermissionError):
    def __init__(self, trace: DecisionTrace) -> None:
        super().__init__(trace.reason_code)
        self.trace = trace
        self.code = trace.reason_code


T = TypeVar("T", bound=Callable[..., Any])


def governed_tool(tool: str) -> Callable[[T], T]:
    """Decorate direct connector tool methods with the shared governance lifecycle."""

    def decorator(method: T) -> T:
        @wraps(method)
        def wrapped(self: Any, context: Any, *args: Any, **kwargs: Any) -> Any:
            governance = getattr(self, "governance", None)
            if governance is None or getattr(self, "_governance_dispatch", False):
                return method(self, context, *args, **kwargs)
            bound = inspect.signature(method).bind(self, context, *args, **kwargs)
            arguments = {
                key: value
                for key, value in bound.arguments.items()
                if key not in {"self", "context"}
            }
            policy_context = governance.context(
                workspace=str(context.workspace_id),
                project=context.project_slug,
                environment=context.environment,
                principal=str(context.client_id or "anonymous"),
                client_id=context.client_id,
                binding_present=context.client_id is not None,
                client_enabled=context.client_id is not None,
                tool=tool,
                connector=getattr(self, "connector_type", "connector"),
                arguments=arguments,
            )
            self._governance_dispatch = True
            try:
                return governance.invoke(
                    policy_context, lambda: method(self, context, *args, **kwargs)
                )
            finally:
                self._governance_dispatch = False

        return cast(T, wrapped)

    return decorator


class Governance:
    def __init__(
        self,
        engine: PolicyEngine,
        audit: AuditLog,
        *,
        rate_limiter: SlidingWindowRateLimiter | None = None,
    ) -> None:
        self.engine = engine
        self.audit = audit
        self.rate_limiter = rate_limiter

    def begin(self, context: InvocationContext) -> tuple[DecisionTrace, float]:
        started = time.monotonic()
        trace = self.engine.evaluate(context)
        if trace.decision in {Decision.DENY, Decision.DENY_AND_ALERT}:
            self.audit.append(
                AuditEvent.from_context(
                    context, trace, outcome="denied", started=started, error_code=trace.reason_code
                )
            )
            raise PolicyDenied(trace)
        if self.rate_limiter is not None:
            try:
                self.rate_limiter.check(
                    client_id=context.client_id or context.principal,
                    project=context.project,
                    tool=context.tool,
                )
            except RateLimitExceeded:
                self.audit.append(
                    AuditEvent.from_context(
                        context,
                        trace,
                        outcome="rate_limited",
                        started=started,
                        error_code="RATE_LIMITED",
                    )
                )
                raise
        return trace, started

    def finish(
        self,
        context: InvocationContext,
        trace: DecisionTrace,
        started: float,
        *,
        outcome: str = "success",
        error_code: str | None = None,
        result_summary: dict[str, Any] | None = None,
    ) -> None:
        self.audit.append(
            AuditEvent.from_context(
                context, trace, outcome=outcome, started=started, error_code=error_code,
                result_summary=result_summary,
            )
        )

    def invoke(self, context: InvocationContext, operation: Callable[[], Any]) -> Any:
        trace, started = self.begin(context)
        try:
            result = operation()
        except Exception as exc:
            self.finish(
                context,
                trace,
                started,
                outcome="error",
                error_code=getattr(exc, "code", type(exc).__name__),
            )
            raise
        summary = None
        if context.tool == "gitea.pullrequest.merge" and isinstance(result, dict):
            data = result.get("data") or {}
            if isinstance(data, dict):
                summary = data.get("auditEvidence")
        self.finish(context, trace, started, result_summary=summary)
        return result

    @staticmethod
    def context(
        *,
        workspace: str,
        project: str,
        environment: str,
        principal: str,
        tool: str,
        connector: str,
        arguments: Mapping[str, Any],
        client_id: Any = None,
        binding_present: bool = True,
        client_enabled: bool = True,
        **kwargs: Any,
    ) -> InvocationContext:
        return InvocationContext(
            workspace=workspace,
            project=project,
            environment=environment,
            principal=principal,
            tool=tool,
            connector=connector,
            arguments=dict(arguments),
            client_id=client_id,
            binding_present=binding_present,
            client_enabled=client_enabled,
            **kwargs,
        )