mirror of
https://gitverse.ru/kpa39l/chronicle.nixg.ru.git
synced 2026-09-29 09:55:08 +00:00
Add telegram-archiver microservice (MVP)
This commit is contained in:
@@ -0,0 +1,347 @@
|
||||
"""
|
||||
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()
|
||||
Reference in New Issue
Block a user