trace/services/notify-gateway/gateway/app.py

import asyncio
import base64
import hashlib
import hmac
import logging
import os
from pathlib import Path

import yaml
from aiohttp import web

from notifiers.base import ImageAttachment
from notifiers.factory import build_notifier
from .delivery_store import DeliveryStore

LOG = logging.getLogger("notify-gateway")
MAX_IMAGES = 4
MAX_IMAGE_BYTES = 50_000_000
MAX_TOTAL_IMAGE_BYTES = 50_000_000


def parse_images(raw: object) -> list[ImageAttachment]:
    if raw in (None, []):
        return []
    if not isinstance(raw, list) or len(raw) > MAX_IMAGES:
        raise web.HTTPBadRequest(reason="images must be a list with at most 4 items")
    images: list[ImageAttachment] = []
    total_bytes = 0
    for item in raw:
        if not isinstance(item, dict):
            raise web.HTTPBadRequest(reason="invalid image item")
        content_type = str(item.get("content_type") or "")
        if content_type not in {"image/jpeg", "image/png"}:
            raise web.HTTPBadRequest(reason="unsupported image content type")
        try:
            data = base64.b64decode(str(item.get("data_base64") or ""), validate=True)
        except Exception:
            raise web.HTTPBadRequest(reason="invalid image base64")
        magic_ok = (
            content_type == "image/jpeg" and data.startswith(bytes.fromhex("ffd8ff"))
        ) or (
            content_type == "image/png"
            and data.startswith(bytes.fromhex("89504e470d0a1a0a"))
        )
        if not magic_ok or not data or len(data) > MAX_IMAGE_BYTES:
            raise web.HTTPBadRequest(reason="invalid image data or size")
        total_bytes += len(data)
        if total_bytes > MAX_TOTAL_IMAGE_BYTES:
            raise web.HTTPBadRequest(reason="total image size exceeds 50 MB")
        filename = Path(str(item.get("filename") or "image")).name
        images.append(ImageAttachment(filename, content_type, data))
    return images


class Gateway:
    def __init__(self, cfg: dict):
        self.cfg = cfg
        self.expected = os.getenv("NOTIFY_GATEWAY_TOKEN", "") or str(cfg.get("api_token", ""))
        self.channels_cfg = cfg.get("channels") or {}
        self.notifiers = {}
        self.telegram_bots = {}
        for name, channel_cfg in self.channels_cfg.items():
            if not channel_cfg.get("enabled"):
                continue
            if name == "telegram":
                bots_cfg = channel_cfg.get("bots") or {}
                if not isinstance(bots_cfg, dict) or not bots_cfg:
                    raise RuntimeError("Telegram bots are not configured")
                self.telegram_bots = {
                    bot_name: build_notifier("telegram", bot_cfg)
                    for bot_name, bot_cfg in bots_cfg.items()
                }
            else:
                self.notifiers[name] = build_notifier(name, channel_cfg)
        self.store = DeliveryStore(str(cfg.get("state_db", "/app/data/deliveries.db")))

    def _check(self, request: web.Request) -> None:
        if not self.expected:
            raise web.HTTPForbidden(reason="gateway token not configured")
        auth = request.headers.get("Authorization", "")
        if not hmac.compare_digest(auth, f"Bearer {self.expected}"):
            raise web.HTTPUnauthorized()

    async def health(self, request: web.Request):
        channels = {
            name: (bool(self.telegram_bots) if name == "telegram" else name in self.notifiers)
            for name in self.channels_cfg
        }
        return web.json_response({
            "ok": True,
            "channels": channels,
            "telegram_bots": sorted(self.telegram_bots),
        })

    async def send(self, request: web.Request):
        self._check(request)
        try:
            data = await request.json()
        except Exception:
            raise web.HTTPBadRequest(reason="invalid json")
        title = str(data.get("title", ""))
        text = str(data.get("text", ""))
        images = parse_images(data.get("images"))
        wanted = data.get("channels")
        if not text and not title and not images:
            raise web.HTTPBadRequest(reason="empty message")
        if not isinstance(wanted, list) or not wanted or not all(isinstance(x, str) and x for x in wanted):
            raise web.HTTPBadRequest(reason="channels must explicitly select at least one channel")
        if len(set(wanted)) != len(wanted):
            raise web.HTTPBadRequest(reason="channels must not contain duplicates")
        if len(title) > 500 or len(text) > 20000:
            raise web.HTTPRequestEntityTooLarge(max_size=20500, actual_size=len(title) + len(text))
        telegram_bot = data.get("telegram_bot")
        if "telegram" in wanted:
            if not isinstance(telegram_bot, str) or not telegram_bot.strip():
                raise web.HTTPBadRequest(reason="telegram_bot is required for Telegram delivery")
            telegram_bot = telegram_bot.strip()
            if telegram_bot not in self.telegram_bots:
                raise web.HTTPBadRequest(reason="unknown telegram_bot")
        elif telegram_bot is not None:
            raise web.HTTPBadRequest(reason="telegram_bot is only valid with the Telegram channel")
        key = str(data.get("idempotency_key") or "").strip()
        if not key:
            digest_input = bytearray((title + "\0" + text).encode("utf-8"))
            for image in images:
                digest_input.extend(b"\0")
                digest_input.extend(hashlib.sha256(image.data).digest())
            digest = hashlib.sha256(digest_input).hexdigest()
            key = f"content:{digest}"
        if len(key) > 200:
            raise web.HTTPBadRequest(reason="idempotency_key too long")
        results: dict[str, str] = {}
        for name in wanted:
            notifier = self.telegram_bots.get(telegram_bot) if name == "telegram" else self.notifiers.get(name)
            if not notifier:
                results[name] = "disabled"
                continue
            delivery_target = f"telegram:{telegram_bot}" if name == "telegram" else name
            if await self.store.delivered(key, delivery_target):
                results[name] = "already_delivered"
                continue
            try:
                await notifier.send(title, text, images)
                await self.store.record(key, delivery_target, True)
                results[name] = "ok"
            except Exception as exc:
                LOG.exception("gateway send failed channel=%s", name)
                await self.store.record(key, delivery_target, False, repr(exc))
                results[name] = "error"
        ok = bool(results) and all(v in {"ok", "already_delivered"} for v in results.values())
        return web.json_response({"ok": ok, "results": results})

    async def run_app(self):
        await self.store.init()
        site = web.Application(client_max_size=70 * 1024 * 1024)
        site.router.add_get("/health", self.health)
        site.router.add_post("/send", self.send)
        runner = web.AppRunner(site)
        await runner.setup()
        await web.TCPSite(runner, self.cfg.get("listen_host", "0.0.0.0"),
                          int(self.cfg.get("listen_port", 8770))).start()
        LOG.info("notify-gateway ready channels=%s",
                 [k for k, v in (self.cfg.get("channels") or {}).items() if v.get("enabled")])
        await asyncio.Event().wait()


def load_config():
    with open(os.environ.get("NOTIFY_GATEWAY_CONFIG", "/app/config/config.yaml"),
              encoding="utf-8") as f:
        return yaml.safe_load(f)


def main():
    Path("/app/logs").mkdir(parents=True, exist_ok=True)
    logging.basicConfig(level=logging.INFO,
                        format="%(asctime)s %(levelname)s %(name)s %(message)s")
    # httpx logs full request URLs; Telegram and WeCom put credentials in URLs.
    logging.getLogger("httpx").setLevel(logging.WARNING)
    logging.getLogger("httpcore").setLevel(logging.WARNING)
    asyncio.run(Gateway(load_config()).run_app())