oncall-engine/engine/apps/telegram/client.py

from typing import Optional, Tuple, Union
from urllib.parse import urljoin

from telegram import Bot, InlineKeyboardMarkup, Message, ParseMode
from telegram.error import InvalidToken, Unauthorized
from telegram.utils.request import Request

from apps.alerts.models import AlertGroup
from apps.base.utils import live_settings
from apps.telegram.models import TelegramMessage
from apps.telegram.renderers.keyboard import TelegramKeyboardRenderer
from apps.telegram.renderers.message import TelegramMessageRenderer


class TelegramClient:
    ALLOWED_UPDATES = ("message", "callback_query")
    PARSE_MODE = ParseMode.HTML

    def __init__(self, token: Optional[str] = None):
        self.token = token or live_settings.TELEGRAM_TOKEN

        if self.token is None:
            raise InvalidToken()

    @property
    def api_client(self) -> Bot:
        return Bot(self.token, request=Request(read_timeout=15))

    def is_chat_member(self, chat_id: Union[int, str]) -> bool:
        try:
            self.api_client.get_chat(chat_id=chat_id)
            return True
        except Unauthorized:
            return False

    def register_webhook(self, webhook_url: Optional[str] = None) -> None:
        webhook_url = webhook_url or urljoin(live_settings.TELEGRAM_WEBHOOK_HOST, "/telegram/")

        if webhook_url is None:
            webhook_url = live_settings.TELEGRAM_WEBHOOK_URL

        webhook_info = self.api_client.get_webhook_info()
        if webhook_info.url == webhook_url:
            return

        self.api_client.set_webhook(webhook_url, allowed_updates=self.ALLOWED_UPDATES)

    def send_message(
        self,
        chat_id: Union[int, str],
        message_type: int,
        alert_group: AlertGroup,
        reply_to_message_id: Optional[int] = None,
    ) -> TelegramMessage:
        text, keyboard = self._get_message_and_keyboard(message_type=message_type, alert_group=alert_group)

        raw_message = self.send_raw_message(
            chat_id=chat_id, text=text, keyboard=keyboard, reply_to_message_id=reply_to_message_id
        )
        message = TelegramMessage.create_from_message(
            message=raw_message, alert_group=alert_group, message_type=message_type
        )

        return message

    def send_raw_message(
        self,
        chat_id: Union[int, str],
        text: str,
        keyboard: Optional[InlineKeyboardMarkup] = None,
        reply_to_message_id: Optional[int] = None,
    ) -> Message:
        message = self.api_client.send_message(
            chat_id=chat_id,
            text=text,
            reply_markup=keyboard,
            reply_to_message_id=reply_to_message_id,
            parse_mode=self.PARSE_MODE,
            disable_web_page_preview=False,
        )
        return message

    def edit_message(self, message: TelegramMessage) -> TelegramMessage:
        text, keyboard = self._get_message_and_keyboard(
            message_type=message.message_type, alert_group=message.alert_group
        )

        self.edit_raw_message(chat_id=message.chat_id, message_id=message.message_id, text=text, keyboard=keyboard)
        return message

    def edit_raw_message(
        self,
        chat_id: Union[int, str],
        message_id: Union[int, str],
        text: str,
        keyboard: Optional[InlineKeyboardMarkup] = None,
    ) -> Message:
        message = self.api_client.edit_message_text(
            chat_id=chat_id,
            message_id=message_id,
            text=text,
            reply_markup=keyboard,
            parse_mode=self.PARSE_MODE,
            disable_web_page_preview=False,
        )
        return message

    @staticmethod
    def _get_message_and_keyboard(
        message_type: int, alert_group: AlertGroup
    ) -> Tuple[str, Optional[InlineKeyboardMarkup]]:
        message_renderer = TelegramMessageRenderer(alert_group=alert_group)
        keyboard_renderer = TelegramKeyboardRenderer(alert_group=alert_group)

        if message_type == TelegramMessage.ALERT_GROUP_MESSAGE:
            text = message_renderer.render_alert_group_message()
            keyboard = None
        elif message_type == TelegramMessage.LOG_MESSAGE:
            text = message_renderer.render_log_message()
            keyboard = None
        elif message_type == TelegramMessage.ACTIONS_MESSAGE:
            text = message_renderer.render_actions_message()
            keyboard = keyboard_renderer.render_actions_keyboard()
        elif message_type == TelegramMessage.PERSONAL_MESSAGE:
            text = message_renderer.render_personal_message()
            keyboard = keyboard_renderer.render_actions_keyboard()
        elif message_type == TelegramMessage.FORMATTING_ERROR:
            text = message_renderer.render_formatting_error_message()
            keyboard = None
        elif message_type in (
            TelegramMessage.LINK_TO_CHANNEL_MESSAGE,
            TelegramMessage.LINK_TO_CHANNEL_MESSAGE_WITHOUT_TITLE,
        ):
            alert_group_message = alert_group.telegram_messages.filter(
                chat_id__startswith="-",
                message_type__in=[TelegramMessage.ALERT_GROUP_MESSAGE, TelegramMessage.FORMATTING_ERROR],
            ).first()

            if alert_group_message is None:
                raise Exception("No alert group message found, probably it is not saved to database yet")

            include_title = message_type == TelegramMessage.LINK_TO_CHANNEL_MESSAGE
            link = alert_group_message.link

            text = message_renderer.render_link_to_channel_message(include_title=include_title)
            keyboard = keyboard_renderer.render_link_to_channel_keyboard(link=link)
        else:
            raise Exception(f"_get_message_and_keyboard with type {message_type} is not implemented")

        return text, keyboard