ForcePilot/backend/package/yuxi/channel/gateway/routes.py

import asyncio
import logging
from collections.abc import Callable
from typing import Any

from yuxi.channel.gateway.protocol import DeliveryMode
from yuxi.channel.gateway.webhook_security import (
    WebhookGuard,
    WebhookGuardConfig,
    WebhookGuardResult,
)
from yuxi.channel.gateway.webhook_security import (
    webhook_guard as _default_webhook_guard,
)

logger = logging.getLogger(__name__)

WebhookHandler = Callable[[dict[str, Any]], asyncio.Future[dict[str, Any]]]


class WebhookRegistry:
    """渠道 Webhook 处理器注册表。

    每个渠道插件注册自己的 webhook 处理函数，
    HTTP 层统一通过 FastAPI router 将请求分发到此处。

    集成了多层 Webhook 安全 Guard：
    - Method 检查（仅允许注册的 HTTP 方法）
    - Content-Type 验证
    - Body 大小限制
    - HMAC-SHA256 签名校验
    - 并发限制
    - 异常追踪告警
    """

    def __init__(self, webhook_guard: WebhookGuard | None = None):
        self._handlers: dict[str, WebhookHandler] = {}
        self._delivery_modes: dict[str, DeliveryMode] = {}
        self._guard = webhook_guard or _default_webhook_guard

    def register(
        self,
        channel_type: str,
        handler: WebhookHandler,
        delivery_mode: DeliveryMode = DeliveryMode.DIRECT,
        guard_config: WebhookGuardConfig | None = None,
    ) -> None:
        if channel_type in self._handlers:
            logger.warning("Overwriting webhook handler for channel: %s", channel_type)
        self._handlers[channel_type] = handler
        self._delivery_modes[channel_type] = delivery_mode
        self._guard.register_channel(channel_type, guard_config)

    def get_handler(self, channel_type: str) -> WebhookHandler | None:
        return self._handlers.get(channel_type)

    def get_delivery_mode(self, channel_type: str) -> DeliveryMode:
        return self._delivery_modes.get(channel_type, DeliveryMode.DIRECT)

    def list_channels(self) -> list[dict[str, str]]:
        return [
            {
                "channel_type": ct,
                "delivery_mode": self._delivery_modes.get(ct, DeliveryMode.DIRECT).value,
            }
            for ct in self._handlers
        ]

    def remove(self, channel_type: str) -> None:
        self._handlers.pop(channel_type, None)
        self._delivery_modes.pop(channel_type, None)
        self._guard.unregister_channel(channel_type)

    def clear(self) -> None:
        for ct in list(self._handlers):
            self._guard.unregister_channel(ct)
        self._handlers.clear()
        self._delivery_modes.clear()

    async def dispatch_guarded(
        self,
        channel_type: str,
        method: str,
        content_type: str | None,
        body: bytes,
        signature: str | None = None,
        extra_headers: dict[str, str] | None = None,
    ) -> tuple[WebhookGuardResult | None, dict[str, Any] | None]:
        guard_result = self._guard.run_pipeline(
            channel_type,
            method,
            content_type,
            body,
            signature,
        )
        if not guard_result.allowed:
            return guard_result, None

        handler = self._handlers.get(channel_type)
        if handler is None:
            self._guard.release_concurrency(channel_type)
            return WebhookGuardResult(
                allowed=False,
                error_code=None,
                error_message=f"No handler registered for channel: {channel_type}",
                http_status=404,
            ), None

        try:
            import json

            payload = json.loads(body)
        except Exception:
            self._guard.release_concurrency(channel_type)
            return WebhookGuardResult(
                allowed=False,
                error_code=None,
                error_message="Invalid JSON payload",
                http_status=400,
            ), None

        if extra_headers:
            payload["_headers"] = extra_headers
            payload["_raw_body"] = body

        try:
            result = await handler(payload)
        except Exception:
            logger.exception("Webhook handler failed for channel: %s", channel_type)
            result = None
        finally:
            self._guard.release_concurrency(channel_type)

        return None, result

    @property
    def guard(self) -> WebhookGuard:
        return self._guard


webhook_registry = WebhookRegistry()


import re

_WEBHOOK_PATH_RE = re.compile(r"^[a-zA-Z0-9_-]+$")


def build_webhook_path(channel_type: str) -> str:
    if not channel_type or not _WEBHOOK_PATH_RE.match(channel_type):
        raise ValueError(f"Invalid channel_type: {channel_type!r}")
    return f"/webhook/{channel_type}"


def build_health_response(status: str, channels: dict, summary: dict, ts: str | None = None) -> dict:
    return {
        "status": status,
        "timestamp": ts,
        "channels": channels,
        "summary": summary,
    }