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

from typing import Optional, Tuple, Union

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
from common.api_helpers.utils import create_engine_url


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 create_engine_url("/telegram/", override_base=live_settings.TELEGRAM_WEBHOOK_HOST)

        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