mirror of
https://gitverse.ru/kpa39l/chronicle.nixg.ru.git
synced 2026-09-29 09:55:08 +00:00
348 lines
11 KiB
Python
348 lines
11 KiB
Python
"""
|
|
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="Telegram Archiver",
|
|
app_version="1.0.0",
|
|
lang_code="en",
|
|
system_lang_code="en",
|
|
)
|
|
|
|
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:
|
|
# Send code request
|
|
await self.client.send_code_request(self.settings.phone)
|
|
logger.info(f"Code sent to {self.settings.phone}")
|
|
|
|
# Get code from user (in MVP, we'll use input())
|
|
code = input("Enter the code you received: ")
|
|
|
|
try:
|
|
await self.client.sign_in(
|
|
phone=self.settings.phone,
|
|
code=code
|
|
)
|
|
except SessionPasswordNeededError:
|
|
# 2FA is enabled
|
|
logger.info("2FA enabled. Enter password:")
|
|
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")
|
|
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
|
|
async for message in self.client.iter_messages(
|
|
entity,
|
|
limit=limit,
|
|
min_id=from_message_id, # Messages with ID > from_message_id
|
|
):
|
|
yield message
|
|
message_count += 1
|
|
|
|
if message_count % 100 == 0:
|
|
logger.debug(f"Fetched {message_count} messages...")
|
|
|
|
logger.info(f"Finished fetching. Total messages: {message_count}")
|
|
|
|
async def download_media(
|
|
self,
|
|
message: Message,
|
|
output_path: Path,
|
|
max_size: Optional[int] = None,
|
|
) -> 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)
|
|
|
|
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)
|
|
|
|
# Handle different media types
|
|
if isinstance(message.media, MessageMediaPhoto):
|
|
# Photo
|
|
info = await self._download_photo(message, output_path, max_size)
|
|
if info:
|
|
media_info.append(info)
|
|
|
|
elif isinstance(message.media, MessageMediaDocument):
|
|
# Document, video, audio, etc.
|
|
info = await self._download_document(message, output_path, max_size)
|
|
if info:
|
|
media_info.append(info)
|
|
|
|
# Note: Other media types (geo, poll, etc.) are not downloadable
|
|
|
|
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()
|