LLM-Zabbix-HABR/alert-processor/app/matrix_notifier.py

from __future__ import annotations

import mimetypes
import uuid
from dataclasses import dataclass
from pathlib import Path
from urllib.parse import quote

import httpx

from app.matrix_token_manager import MatrixTokenManager


@dataclass
class MatrixSendResult:
    ok: bool
    status_code: int | None
    event_id: str | None
    error: str | None
    content_uri: str | None = None


@dataclass
class MatrixUploadResult:
    ok: bool
    status_code: int | None
    content_uri: str | None
    error: str | None


class MatrixNotifier:
    def __init__(
        self,
        homeserver_url: str,
        room_id: str,
        token_manager: MatrixTokenManager,
        message_type: str = "m.notice",
        timeout_seconds: float = 10,
        verify_tls: bool = True,
    ) -> None:
        self.homeserver_url = homeserver_url.rstrip("/")
        self.room_id = room_id
        self.token_manager = token_manager
        self.message_type = message_type
        self.client = httpx.AsyncClient(
            timeout=timeout_seconds,
            verify=verify_tls,
        )

    async def close(self) -> None:
        await self.client.aclose()

    async def send_message(self, body: str) -> MatrixSendResult:
        payload = {
            "msgtype": self.message_type,
            "body": body,
        }
        return await self._send_room_message(payload)

    async def send_image(
        self,
        file_path: str,
        body: str | None = None,
    ) -> MatrixSendResult:
        path = Path(file_path)
        if not path.exists() or not path.is_file():
            return MatrixSendResult(
                ok=False,
                status_code=None,
                event_id=None,
                error=f"Image file not found: {file_path}",
                content_uri=None,
            )

        mime_type, _ = mimetypes.guess_type(path.name)
        mime_type = mime_type or "image/png"
        data = path.read_bytes()

        upload_result = await self._upload_bytes(
            filename=path.name,
            data=data,
            content_type=mime_type,
        )
        if not upload_result.ok:
            return MatrixSendResult(
                ok=False,
                status_code=upload_result.status_code,
                event_id=None,
                error=upload_result.error,
                content_uri=None,
            )

        payload = {
            "msgtype": "m.image",
            "body": body or path.name,
            "url": upload_result.content_uri,
            "info": {
                "mimetype": mime_type,
                "size": len(data),
            },
        }

        send_result = await self._send_room_message(payload)
        send_result.content_uri = upload_result.content_uri
        return send_result

    async def _send_room_message(self, payload: dict) -> MatrixSendResult:
        txn_id = str(uuid.uuid4())
        room_id_encoded = quote(self.room_id, safe="")
        url = (
            f"{self.homeserver_url}"
            f"/_matrix/client/v3/rooms/{room_id_encoded}"
            f"/send/m.room.message/{txn_id}"
        )

        try:
            response = await self._request_with_refresh(
                method="PUT",
                url=url,
                json=payload,
                headers={"Content-Type": "application/json"},
            )
            data = response.json() if response.content else {}

            if response.is_success:
                return MatrixSendResult(
                    ok=True,
                    status_code=response.status_code,
                    event_id=data.get("event_id"),
                    error=None,
                    content_uri=None,
                )

            return MatrixSendResult(
                ok=False,
                status_code=response.status_code,
                event_id=None,
                error=str(data),
                content_uri=None,
            )
        except Exception as exc:
            return MatrixSendResult(
                ok=False,
                status_code=None,
                event_id=None,
                error=str(exc),
                content_uri=None,
            )

    async def _upload_bytes(
        self,
        filename: str,
        data: bytes,
        content_type: str,
    ) -> MatrixUploadResult:
        url = f"{self.homeserver_url}/_matrix/media/v3/upload"

        try:
            response = await self._request_with_refresh(
                method="POST",
                url=url,
                params={"filename": filename},
                content=data,
                headers={"Content-Type": content_type},
            )
            payload = response.json() if response.content else {}

            if response.is_success:
                return MatrixUploadResult(
                    ok=True,
                    status_code=response.status_code,
                    content_uri=payload.get("content_uri"),
                    error=None,
                )

            return MatrixUploadResult(
                ok=False,
                status_code=response.status_code,
                content_uri=None,
                error=str(payload),
            )
        except Exception as exc:
            return MatrixUploadResult(
                ok=False,
                status_code=None,
                content_uri=None,
                error=str(exc),
            )

    async def _request_with_refresh(
        self,
        method: str,
        url: str,
        **kwargs,
    ) -> httpx.Response:
        response = await self._request_once(method, url, **kwargs)

        if response.status_code == 401:
            await self.token_manager.refresh_if_needed(force=True)
            response = await self._request_once(method, url, **kwargs)

        return response

    async def _request_once(
        self,
        method: str,
        url: str,
        **kwargs,
    ) -> httpx.Response:
        token = await self.token_manager.get_access_token()

        headers = dict(kwargs.pop("headers", {}) or {})
        headers["Authorization"] = f"Bearer {token}"

        return await self.client.request(
            method=method,
            url=url,
            headers=headers,
            **kwargs,
        )