pyrogram-bridge/telegram_client.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-

# flake8: noqa
# pylint: disable=broad-exception-raised, raise-missing-from, too-many-arguments, redefined-outer-name
# pylint: disable=multiple-statements, logging-fstring-interpolation, trailing-whitespace, line-too-long
# pylint: disable=broad-exception-caught, missing-function-docstring, missing-class-docstring
# pylint: disable=f-string-without-interpolation
# pylance: disable=reportMissingImports, reportMissingModuleSource

import logging
import os
import asyncio
import sys
import signal
import uvloop
import time
asyncio.set_event_loop_policy(uvloop.EventLoopPolicy())

from pyrogram import Client
from pyrogram.handlers import DisconnectHandler
from config import get_settings

logger = logging.getLogger(__name__)
settings = get_settings()

class TelegramClient:
    def __init__(self):
        self._ensure_session_directory()
        self.client = Client(
            name="pyro_bridge",
            api_id=settings["tg_api_id"],
            api_hash=settings["tg_api_hash"],
            workdir=settings["session_path"],
            proxy=settings["proxy"],  # MTProto proxy config, None if not set
        )
        self.disconnect_count = 0
        self.max_disconnects = 3
        self._shutting_down = False  # Guard to prevent re-triggering restart during shutdown
        self._setup_connection_handlers()

    def _ensure_session_directory(self):
        try:
            os.makedirs(settings["session_path"], exist_ok=True)
            logger.debug(f'Session directory created/verified: {settings["session_path"]}')
        except Exception as e:
            logger.error(f'Failed to create session directory {settings["session_path"]}: {str(e)}')
            raise

    def _setup_connection_handlers(self):
        """Sets up connection/disconnection handlers"""
        self.client.add_handler(DisconnectHandler(self._on_disconnect))
        logger.info("connection_handlers: connection handlers set up")

    async def _on_disconnect(self, _client, _session=None):
        """Handles disconnection from Telegram servers"""
        # Ignore disconnects that are caused by the intentional shutdown sequence
        if self._shutting_down:
            logger.debug("connection_handler: ignoring disconnect event during shutdown")
            return
        self.disconnect_count += 1
        logger.warning(f"connection_handler: connection lost (#{self.disconnect_count})")

        # Wait briefly to allow Pyrogram's internal reconnect to complete
        await asyncio.sleep(10)

        # Re-check shutdown flag — it may have been set while we were sleeping
        # (e.g. another concurrent _on_disconnect already triggered _restart_app)
        if self._shutting_down:
            logger.debug("connection_handler: shutdown started during sleep, suppressing further action")
            return

        # If the client reconnected successfully, reset the counter and return
        if self.client.is_connected:
            logger.info(f"connection_handler: reconnected successfully, resetting disconnect counter")
            self.disconnect_count = 0
            return

        if self.disconnect_count >= self.max_disconnects:
            logger.critical(f"connection_handler: reached disconnect limit ({self.max_disconnects})")
            self._restart_app()

    async def start(self):
        try:
            if not self.client.is_connected:
                await self.client.start()
                logger.info("Telegram client connected successfully")
                # Reset disconnect counter on successful connection
                self.disconnect_count = 0
                logger.info("connection_handler: connection established")
        except Exception as e:
            logger.error(f"Failed to start Telegram client: {str(e)}")
            raise

    def _restart_app(self):
        """Restarts the application by sending SIGTERM to the process"""
        self._shutting_down = True  # Set flag before sending signal to suppress subsequent disconnect events
        logger.warning("connection_handler: restarting application by sending SIGTERM")
        try:
            # Use SIGTERM for proper Docker container restart
            logger.critical("connection_handler: sending SIGTERM signal")
            os.kill(os.getpid(), signal.SIGTERM)
        except Exception as e:
            logger.error(f"connection_handler: error during restart: {str(e)}")
            # Emergency termination
            os._exit(1)

    async def safe_get_messages(self, channel_id, post_id, max_retries=2):
        """Wrapper with retry logic for auth errors"""
        for attempt in range(max_retries):
            try:
                return await asyncio.wait_for(
                    self.client.get_messages(channel_id, post_id),
                    timeout=30.0
                )
            except Exception as e:
                if isinstance(e, KeyError) and attempt < max_retries - 1:
                    logger.warning(f"Auth error on attempt {attempt + 1}, retrying in 5s...")
                    await asyncio.sleep(5)
                    continue
                raise

    async def safe_download_media(self, file_id, file_name, max_retries=2):
        """Wrapper with retry logic for download errors"""
        for attempt in range(max_retries):
            try:
                return await asyncio.wait_for(
                    self.client.download_media(file_id, file_name=file_name),
                    timeout=120.0
                )
            except Exception as e:
                if isinstance(e, KeyError) and attempt < max_retries - 1:
                    logger.warning(f"Download auth error on attempt {attempt + 1}, retrying...")
                    await asyncio.sleep(5)
                    continue
                raise

    async def stop(self):
        if self.client.is_connected:
            await self.client.stop()
            logger.info("Telegram client disconnected")