ForcePilot/backend/package/yuxi/channel/extensions/viber/webhook.py

from __future__ import annotations

import asyncio
import hashlib
import hmac
import json
import logging
import os

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

from yuxi.channel.extensions.viber.config import ENV_VIBER_AUTH_TOKEN

logger = logging.getLogger(__name__)

VIBER_WEBHOOK_PREAUTH_MAX_BODY_BYTES = 256 * 1024
VIBER_WEBHOOK_PREAUTH_BODY_TIMEOUT_S = 5.0

router = APIRouter(prefix="/webhook/viber", tags=["viber"])

_webhook_running = False


def set_webhook_running(running: bool) -> None:
    global _webhook_running
    _webhook_running = running


def set_webhook_config(config: dict) -> None:
    from yuxi.channel.extensions.viber import viber_plugin

    viber_plugin._config.list_account_ids(config)


def verify_viber_signature(body: bytes, signature: str, auth_token: str) -> bool:
    if not signature or not auth_token:
        return False
    expected = hmac.new(
        auth_token.encode("utf-8"),
        body,
        hashlib.sha256,
    ).hexdigest()
    return hmac.compare_digest(expected, signature)


@router.post("")
async def viber_webhook(request: Request):
    if not _webhook_running:
        return JSONResponse(
            status_code=503,
            content={"status": 1, "status_message": "channel not running"},
        )

    body = await _read_body_with_limit(request)
    if body is None:
        return JSONResponse(
            status_code=413,
            content={"status": 1, "status_message": "body too large"},
        )

    signature = request.headers.get("X-Viber-Content-Signature", "")

    from yuxi.channel.extensions.viber import viber_plugin

    account, account_id = await _resolve_account_for_webhook(viber_plugin, body, signature)
    if not account:
        return JSONResponse(
            status_code=401,
            content={"status": 1, "status_message": "invalid signature"},
        )

    try:
        callback = json.loads(body)
    except json.JSONDecodeError:
        return JSONResponse(
            status_code=400,
            content={"status": 1, "status_message": "invalid JSON"},
        )

    event_type = callback.get("event", "")

    if event_type == "webhook":
        logger.info("Viber webhook verification callback received")
        return {"status": 0, "status_message": "ok"}

    um = viber_plugin._monitor.parse_event_to_unified(callback, account_id)
    if um:
        asyncio.create_task(
            _dispatch_to_agent(um),
            name=f"viber-dispatch-{um.sender.id}",
        )

    return {"status": 0, "status_message": "ok"}


async def _resolve_account_for_webhook(plugin, body: bytes, signature: str) -> tuple[dict | None, str]:
    account_ids = plugin._config.list_account_ids({})
    accounts = []
    for aid in account_ids:
        try:
            acct = await plugin._config.resolve_account(aid)
            if acct:
                accounts.append(acct)
        except Exception:
            logger.exception("Viber resolve_account error for %s", aid)

    for acct in accounts:
        auth_token = acct.get("auth_token", "")
        if auth_token and verify_viber_signature(body, signature, auth_token):
            return acct, acct.get("account_id", "default")

    env_token = os.environ.get(ENV_VIBER_AUTH_TOKEN, "")
    if env_token and verify_viber_signature(body, signature, env_token):
        return {
            "account_id": "default",
            "auth_token": env_token,
            "sender_name": "ForcePilot Bot",
            "name": "default",
        }, "default"

    return None, ""


async def _dispatch_to_agent(msg) -> None:
    try:
        from yuxi.channel.runtime.manager import gateway

        processor = getattr(gateway, "_processor", None)
        if processor is None:
            logger.warning("Viber dispatch: message processor not available")
            return

        await asyncio.wait_for(processor.process(msg), timeout=120.0)
    except TimeoutError:
        logger.error("Viber agent response timeout for user %s", msg.sender.id)
    except Exception:
        logger.exception("Viber dispatch error for user %s", msg.sender.id)


async def _read_body_with_limit(request: Request) -> bytes | None:
    try:
        body = await asyncio.wait_for(
            request.body(),
            timeout=VIBER_WEBHOOK_PREAUTH_BODY_TIMEOUT_S,
        )
    except TimeoutError:
        logger.warning("Viber webhook: body read timeout")
        return None

    if len(body) > VIBER_WEBHOOK_PREAUTH_MAX_BODY_BYTES:
        logger.warning("Viber webhook: body too large (%d bytes)", len(body))
        return None

    return body