ForcePilot/backend/package/yuxi/channel/extensions/flock/gateway.py

from __future__ import annotations

import asyncio
import hashlib
import hmac
import json
import logging

from fastapi import APIRouter, Request
from fastapi.responses import JSONResponse

from .constants import (
    EVENT_APP_INSTALL,
    EVENT_APP_UNINSTALL,
    EVENT_CHAT_ACTION,
    EVENT_MESSAGE_RECEIVE,
    EVENT_SLASH_COMMAND,
    WEBHOOK_PATH_EVENT_LISTENER,
    WEBHOOK_PATH_OUTGOING,
    X_FLOCK_EVENT_TOKEN,
)
from .message import (
    flock_event_to_inbound,
    inbound_to_unified,
    parse_event_listener,
    parse_outgoing_webhook,
)

logger = logging.getLogger(__name__)


class FlockGatewayAdapter:
    def __init__(self):
        self._queue: asyncio.Queue | None = None
        self._accounts_config: dict = {}
        self._router: APIRouter | None = None

    @property
    def router(self) -> APIRouter | None:
        return self._router

    async def start(self, ctx) -> object:
        self._queue = ctx.queue
        self._accounts_config = ctx.config
        self._router = self._build_router()
        logger.info("Flock gateway started, webhook routes registered")
        return {"router": self._router, "queue": self._queue}

    async def stop(self, ctx) -> None:
        self._queue = None
        self._accounts_config = {}
        self._router = None
        logger.info("Flock gateway stopped")

    def resolve_gateway_auth_bypass_paths(self, config: dict) -> list[str]:
        return [WEBHOOK_PATH_EVENT_LISTENER, WEBHOOK_PATH_OUTGOING]

    def _build_router(self) -> APIRouter:
        router = APIRouter(prefix="/api/channel/flock", tags=["flock"])

        @router.post("/events")
        async def handle_event_listener(request: Request):
            raw_body = await request.body()

            if not self._verify_event_token(request, raw_body):
                logger.warning("Flock event token verification failed")
                return JSONResponse(status_code=401, content={"error": "Invalid event token"})

            try:
                body = json.loads(raw_body.decode("utf-8"))
            except json.JSONDecodeError:
                return JSONResponse(status_code=400, content={"error": "Invalid JSON"})

            name = body.get("name", "")
            logger.info("Flock event listener received: %s", name)

            event = parse_event_listener(body)

            if name in (EVENT_APP_INSTALL, EVENT_APP_UNINSTALL):
                logger.info("Flock lifecycle event: %s", name)
                return JSONResponse(status_code=200, content={"status": "ok"})

            if name not in (EVENT_MESSAGE_RECEIVE, EVENT_SLASH_COMMAND, EVENT_CHAT_ACTION):
                logger.debug("Flock unhandled event: %s", name)
                return JSONResponse(status_code=200, content={"status": "ignored"})

            if name == EVENT_SLASH_COMMAND:
                command_name = body.get("event", {}).get("command", "")
                logger.info("Flock slash command received: /%s", command_name)

            inbound = flock_event_to_inbound(event)

            account_id = self._resolve_account_id_for_event(inbound.to)
            unified = inbound_to_unified(inbound, account_id)

            if self._queue is not None:
                await self._queue.put(unified)
                logger.info("Flock message queued: %s", inbound.message_id)
            else:
                logger.warning("Flock message queue not available, dropping message: %s", inbound.message_id)

            if name == EVENT_SLASH_COMMAND:
                return JSONResponse(status_code=200, content={
                    "text": f"收到指令 /{command_name}，正在处理...",
                    "responseType": "ephemeral",
                })

            return JSONResponse(status_code=200, content={"status": "ok"})

        @router.post("/outgoing")
        async def handle_outgoing_webhook(request: Request):
            try:
                body = await request.json()
            except json.JSONDecodeError:
                return JSONResponse(status_code=400, content={"error": "Invalid JSON"})

            inbound = parse_outgoing_webhook(body)

            account_id = self._resolve_account_id_for_event(inbound.to)
            unified = inbound_to_unified(inbound, account_id)

            if self._queue is not None:
                await self._queue.put(unified)
                logger.info("Flock outgoing webhook queued: %s", inbound.message_id)
            else:
                logger.warning("Flock message queue not available, dropping message: %s", inbound.message_id)

            return JSONResponse(status_code=200, content={"status": "ok"})

        return router

    def _verify_event_token(self, request: Request, raw_body: bytes) -> bool:
        event_token = request.headers.get(X_FLOCK_EVENT_TOKEN, "")
        if not event_token:
            return True

        accounts = self._accounts_config.get("accounts", {})
        app_secret = ""
        for data in accounts.values():
            candidate = data.get("app_secret", "")
            if candidate:
                app_secret = candidate
                break

        if not app_secret:
            logger.debug("Flock app_secret not configured, skipping event token verification")
            return True

        expected = hmac.new(
            app_secret.encode("utf-8"),
            raw_body,
            hashlib.sha256,
        ).hexdigest()

        return hmac.compare_digest(expected, event_token)

    def _resolve_account_id_for_event(self, to: str) -> str:
        accounts = self._accounts_config.get("accounts", {})
        for aid, data in accounts.items():
            event_token = data.get("event_listener_token", "")
            if event_token and event_token in to:
                return aid
        if accounts:
            return next(iter(accounts.keys()), "default")
        return "default"