chronicle.nixg.ru/telegram-archiver/app/telethon_client.py

"""
Telethon client wrapper for Telegram API interaction.
Handles connection, authentication, and message fetching.
"""

import asyncio
from typing import AsyncGenerator, Optional
from pathlib import Path
import logging

from telethon import TelegramClient
from telethon.sessions import StringSession
from telethon.tl.types import (
    Message,
    Channel,
    Chat,
    User,
    MessageMediaPhoto,
    MessageMediaDocument,
    MessageMediaWebPage,
    MessageMediaGeo,
    MessageMediaPoll,
    MessageMediaDice,
    MessageMediaInvoice,
    MessageMediaContact,
    MessageMediaGame,
    MessageMediaGeoLive,
    MessageMediaVenue,
    MessageMediaGiveaway,
    MessageMediaUnsupported,
    DocumentAttributeFilename,
)
from telethon.errors import SessionPasswordNeededError, FloodWaitError

from config import Settings, get_settings

logger = logging.getLogger(__name__)


class TelethonArchiver:
    """
    Wrapper around Telethon client for archiving purposes.
    Manages connection, authentication, and message iteration.
    """

    def __init__(self, settings: Optional[Settings] = None):
        self.settings = settings or get_settings()
        self.client: Optional[TelegramClient] = None
        self._connected = False

    async def connect(self) -> None:
        """
        Initialize and connect the Telegram client.
        Handles phone number verification and 2FA if needed.
        """
        if self._connected:
            logger.debug("Already connected to Telegram")
            return

        logger.info(f"Connecting to Telegram as {self.settings.phone}...")

        # Create client with session file
        self.client = TelegramClient(
            session=self.settings.session_path,
            api_id=self.settings.api_id,
            api_hash=self.settings.api_hash,
            device_model="Desktop (Ubuntu)",  # Helps with code delivery
            system_version="Ubuntu 22.04",
            app_version="1.0.0",
            lang_code="ru",  # Russian language
            system_lang_code="ru",
        )

        await self.client.connect()

        # Check if authorized
        if not await self.client.is_user_authorized():
            logger.info("Not authorized. Starting authentication...")
            await self._authenticate()

        self._connected = True
        logger.info("Connected to Telegram successfully")

    async def _authenticate(self) -> None:
        """Handle the authentication flow."""
        try:
            # Try app code first
            logger.info("Requesting auth code via Telegram app...")
            sent_code = await self.client.send_code_request(
                self.settings.phone,
                force_sms=False
            )
            logger.info(f"Code request sent. Type: {sent_code.type}")

            # Get code from user
            print("\n" + "=" * 50)
            print("AUTHENTICATION REQUIRED")
            print("=" * 50)
            print("Check your Telegram app (NOT SMS) for the code.")
            print("Look for a message from 'Telegram' with the code.")
            print("=" * 50)
            print("\nOptions:")
            print("  1. Enter the code from Telegram app")
            print("  2. Type 'sms' to receive code via SMS")
            print("  3. Press Ctrl+C to cancel")
            print("=" * 50)
            
            user_input = input("\nEnter code (or 'sms'): ").strip()
            
            # If user requests SMS
            if user_input.lower() == 'sms':
                logger.info("Requesting SMS code...")
                print("\n📱 Requesting SMS code...")
                sent_code = await self.client.send_code_request(
                    self.settings.phone,
                    force_sms=True
                )
                print("SMS sent! Enter the code:")
                user_input = input("SMS Code: ").strip()

            code = user_input

            try:
                await self.client.sign_in(
                    phone=self.settings.phone,
                    code=code,
                    phone_code_hash=sent_code.phone_code_hash
                )
            except SessionPasswordNeededError:
                # 2FA is enabled
                logger.info("2FA enabled. Enter password:")
                print("\n🔐 2FA enabled")
                password = input("2FA Password: ")
                await self.client.sign_in(password=password)

        except FloodWaitError as e:
            logger.error(f"Flood wait: must wait {e.seconds} seconds")
            print(f"\n⚠️  You must wait {e.seconds} seconds before trying again.")
            raise
        except Exception as e:
            logger.error(f"Authentication failed: {e}")
            raise

    async def disconnect(self) -> None:
        """Disconnect from Telegram."""
        if self.client:
            await self.client.disconnect()
            self._connected = False
            logger.info("Disconnected from Telegram")

    async def get_channel_info(self, channel: str) -> Channel | Chat | User:
        """
        Get channel/chat/user information.

        Args:
            channel: Channel username (with or without @) or ID

        Returns:
            Channel, Chat, or User object
        """
        if not self._connected or not self.client:
            raise RuntimeError("Not connected to Telegram")

        entity = await self.client.get_entity(channel)
        logger.info(f"Found channel: {getattr(entity, 'title', 'Unknown')}")
        return entity

    async def fetch_messages(
        self,
        channel: str,
        limit: Optional[int] = None,
        from_message_id: Optional[int] = None,
    ) -> AsyncGenerator[Message, None]:
        """
        Fetch messages from a channel.

        Args:
            channel: Channel username or ID
            limit: Maximum number of messages to fetch (None = all)
            from_message_id: Start from this message ID (None = latest)

        Yields:
            Message objects from newest to oldest
        """
        if not self._connected or not self.client:
            raise RuntimeError("Not connected to Telegram")

        entity = await self.get_channel_info(channel)

        # Get total message count for progress
        total = entity.messages_count if hasattr(entity, 'messages_count') else None
        if limit:
            total = min(total, limit) if total else limit

        logger.info(f"Fetching messages from {getattr(entity, 'title', channel)}...")
        if from_message_id:
            logger.info(f"Starting from message ID: {from_message_id}")

        # Iterate messages (newest first)
        message_count = 0
        error_count = 0
        
        # Build iter_messages kwargs
        iter_kwargs = {"limit": limit}
        if from_message_id:
            iter_kwargs["min_id"] = from_message_id
        
        async for message in self.client.iter_messages(entity, **iter_kwargs):
            try:
                yield message
                message_count += 1

                if message_count % 50 == 0:
                    logger.info(f"Progress: {message_count} messages fetched...")
            except Exception as e:
                error_count += 1
                logger.error(f"Error processing message {message.id}: {e}")
                continue  # Skip problematic messages

        logger.info(f"Finished fetching. Total: {message_count}, Errors: {error_count}")

    async def download_media(
        self,
        message: Message,
        output_path: Path,
        max_size: Optional[int] = None,
        timeout: int = 60,  # Timeout in seconds
    ) -> list[dict]:
        """
        Download all media from a message.

        Args:
            message: Telegram message with media
            output_path: Directory to save media files
            max_size: Maximum file size to download (bytes)
            timeout: Timeout for each file download in seconds

        Returns:
            List of dicts with file info (filename, type, size, path, is_too_large)
        """
        if not message.media:
            return []

        max_size = max_size or self.settings.max_file_size
        media_info = []

        # Ensure output directory exists
        output_path.mkdir(parents=True, exist_ok=True)

        try:
            # Handle different media types with timeout
            if isinstance(message.media, MessageMediaPhoto):
                # Photo
                info = await asyncio.wait_for(
                    self._download_photo(message, output_path, max_size),
                    timeout=timeout
                )
                if info:
                    media_info.append(info)

            elif isinstance(message.media, MessageMediaDocument):
                # Document, video, audio, etc.
                info = await asyncio.wait_for(
                    self._download_document(message, output_path, max_size),
                    timeout=timeout
                )
                if info:
                    media_info.append(info)

        except asyncio.TimeoutError:
            logger.warning(f"Timeout downloading media for message {message.id}")
            # Add info about skipped file
            media_info.append({
                "filename": "timeout_skipped.bin",
                "type": "timeout",
                "caption": None,
                "size": 0,
                "path": None,
                "is_too_large": False,
                "timeout": True,
            })
        except Exception as e:
            logger.error(f"Error downloading media for message {message.id}: {e}")

        return media_info

    async def _download_photo(
        self,
        message: Message,
        output_path: Path,
        max_size: int,
    ) -> Optional[dict]:
        """Download photo from message."""
        try:
            # Get file size
            photo = message.media.photo
            # Get highest resolution
            size = max(photo.sizes, key=lambda s: getattr(s, 'w', 0) * getattr(s, 'h', 0))
            file_size = getattr(size, 'size', 0)

            # Generate filename
            timestamp = message.date.strftime("%Y%m%d_%H%M%S")
            filename = f"photo_{timestamp}.jpg"
            filepath = output_path / filename

            info = {
                "filename": filename,
                "type": "photo",
                "size": file_size,
                "path": str(filepath),
                "is_too_large": file_size > max_size,
                "caption": message.message,
            }

            if file_size > max_size:
                logger.debug(f"Photo too large: {file_size} > {max_size}")
                return info

            # Download
            await self.client.download_media(message.photo, filepath)
            logger.debug(f"Downloaded photo: {filename}")
            return info

        except Exception as e:
            logger.error(f"Failed to download photo: {e}")
            return None

    async def _download_document(
        self,
        message: Message,
        output_path: Path,
        max_size: int,
    ) -> Optional[dict]:
        """Download document/video/audio from message."""
        try:
            doc = message.media.document
            file_size = doc.size

            # Get original filename
            filename = None
            for attr in doc.attributes:
                if isinstance(attr, DocumentAttributeFilename):
                    filename = attr.file_name
                    break

            if not filename:
                # Generate filename based on type
                ext = self._get_extension(doc)
                timestamp = message.date.strftime("%Y%m%d_%H%M%S")
                filename = f"file_{timestamp}{ext}"

            filepath = output_path / filename

            # Determine media type
            media_type = self._get_media_type(doc)

            info = {
                "filename": filename,
                "type": media_type,
                "size": file_size,
                "path": str(filepath),
                "is_too_large": file_size > max_size,
                "caption": message.message,
            }

            if file_size > max_size:
                logger.debug(f"File too large: {file_size} > {max_size}")
                return info

            # Download
            await self.client.download_media(doc, filepath)
            logger.debug(f"Downloaded {media_type}: {filename}")
            return info

        except Exception as e:
            logger.error(f"Failed to download document: {e}")
            return None

    def _get_extension(self, doc) -> str:
        """Get file extension from document attributes."""
        for attr in doc.attributes:
            if isinstance(attr, DocumentAttributeFilename):
                name = attr.file_name
                if "." in name:
                    return "." + name.rsplit(".", 1)[-1]
        return ".bin"

    def _get_media_type(self, doc) -> str:
        """Determine media type from document attributes."""
        mime_type = getattr(doc, 'mime_type', '')

        if mime_type.startswith('video/'):
            return 'video'
        elif mime_type.startswith('audio/'):
            return 'audio'
        elif mime_type.startswith('image/'):
            return 'photo'
        elif mime_type == 'application/x-tgsticker':
            return 'sticker'
        else:
            return 'document'

    async def __aenter__(self):
        """Async context manager entry."""
        await self.connect()
        return self

    async def __aexit__(self, exc_type, exc_val, exc_tb):
        """Async context manager exit."""
        await self.disconnect()