hyperboria/library/telegram/base.py

import datetime
import logging

from aiokit import AioThing
from izihawa_utils.random import random_string
from library.logging import error_log
from telethon import (
    TelegramClient,
    connection,
    sessions,
)
from tenacity import (  # noqa
    retry,
    retry_if_exception_type,
    wait_fixed,
)

from .session_backend import AlchemySessionContainer


class BaseTelegramClient(AioThing):
    def __init__(self, app_id, app_hash, database, bot_token=None, mtproxy=None, flood_sleep_threshold: int = 60):
        AioThing.__init__(self)
        self._telegram_client = TelegramClient(
            self._get_session(database),
            app_id,
            app_hash,
            flood_sleep_threshold=flood_sleep_threshold,
            **self._get_proxy(mtproxy=mtproxy),
        )
        self.bot_token = bot_token

    def _get_session(self, database):
        if database.get('drivername') == 'postgresql':
            self.container = AlchemySessionContainer(
                f"{database['drivername']}://"
                f"{database['username']}:"
                f"{database['password']}@"
                f"{database['host']}:"
                f"{database['port']}/"
                f"{database['database']}",
                session=False,
                manage_tables=False,
            )
            return self.container.new_session(database['session_id'])
        else:
            return sessions.SQLiteSession(session_id=database['session_id'])

    def _get_proxy(self, mtproxy=None):
        if mtproxy and mtproxy.get('enabled', True):
            proxy_config = mtproxy
            return {
                'connection': connection.tcpmtproxy.ConnectionTcpMTProxyRandomizedIntermediate,
                'proxy': (proxy_config['url'], proxy_config['port'], proxy_config['secret'])
            }
        return {}

    @retry(retry=retry_if_exception_type(ConnectionError), wait=wait_fixed(5))
    async def start(self):
        await self._telegram_client.start(bot_token=self.bot_token)

    async def stop(self):
        return await self.disconnect()

    def add_event_handler(self, *args, **kwargs):
        return self._telegram_client.add_event_handler(*args, **kwargs)

    def catch_up(self):
        return self._telegram_client.catch_up()

    def delete_messages(self, *args, **kwargs):
        return self._telegram_client.delete_messages(*args, **kwargs)

    def disconnect(self):
        return self._telegram_client.disconnect()

    @property
    def disconnected(self):
        return self._telegram_client.disconnected

    def download_document(self, *args, **kwargs):
        return self._telegram_client._download_document(
            *args,
            date=datetime.datetime.now(),
            thumb=None,
            progress_callback=None,
            msg_data=None,
            **kwargs,
        )

    def edit_message(self, *args, **kwargs):
        return self._telegram_client.edit_message(*args, **kwargs)

    def edit_permissions(self, *args, **kwargs):
        return self._telegram_client.edit_permissions(*args, **kwargs)

    def forward_messages(self, *args, **kwargs):
        return self._telegram_client.forward_messages(*args, **kwargs)

    def get_entity(self, *args, **kwargs):
        return self._telegram_client.get_entity(*args, **kwargs)

    def get_input_entity(self, *args, **kwargs):
        return self._telegram_client.get_input_entity(*args, **kwargs)

    def iter_admin_log(self, *args, **kwargs):
        return self._telegram_client.iter_admin_log(*args, **kwargs)

    def iter_messages(self, *args, **kwargs):
        return self._telegram_client.iter_messages(*args, **kwargs)

    def list_event_handlers(self):
        return self._telegram_client.list_event_handlers()

    def remove_event_handlers(self):
        for handler in reversed(self.list_event_handlers()):
            self._telegram_client.remove_event_handler(*handler)

    def run_until_disconnected(self):
        return self._telegram_client.run_until_disconnected()

    def send_message(self, *args, **kwargs):
        return self._telegram_client.send_message(*args, **kwargs)

    def send_file(self, *args, **kwargs):
        return self._telegram_client.send_file(*args, **kwargs)

    def __call__(self, *args, **kwargs):
        return self._telegram_client(*args, **kwargs)


class RequestContext:
    def __init__(self, bot_name, chat, request_id: str = None, request_id_length: int = 12):
        self.bot_name = bot_name
        self.chat = chat
        self.request_id = request_id or RequestContext.generate_request_id(request_id_length)
        self.default_fields = {
            'bot_name': self.bot_name,
            'chat_id': self.chat.id,
            'request_id': self.request_id,
        }

    @staticmethod
    def generate_request_id(length):
        return random_string(length)

    def add_default_fields(self, **fields):
        self.default_fields.update(fields)

    def statbox(self, **kwargs):
        logging.getLogger('statbox').info(
            msg=dict(
                **self.default_fields,
                **kwargs,
            ),
        )

    def error_log(self, e, level=logging.ERROR, **fields):
        all_fields = {**self.default_fields, **fields}
        error_log(e, level=level, **all_fields)