Files
drive/telegram_client.py
T
2026-10-18 23:48:22 +05:30

346 lines
14 KiB
Python

"""Telegram access through Telethon (MTProto).
Both operating modes use MTProto, not the HTTP Bot API:
* TELEGRAM_MODE=user — your own account (TELEGRAM_SESSION_STRING or a local session
file). Can read full channel history the account can see; downloads up to the
largest file size Telegram allows.
* TELEGRAM_MODE=bot — a bot (TELEGRAM_BOT_TOKEN) that is an admin of the channel.
Receives new posts live. Bots cannot list history; past posts are fetched by
probing message IDs (best effort). Because this is MTProto rather than the HTTP
Bot API, the HTTP API's 20 MB getFile limit does not apply.
"""
from __future__ import annotations
import asyncio
import logging
import os
import re
from datetime import datetime
from typing import AsyncIterator
from errors import BotError, PermanentError, TransientError
from link_parser import marked_channel_id
from models import MediaInfo, SourceMessage
from naming import VIDEO_EXTENSIONS, split_video_extension
log = logging.getLogger(__name__)
REQUEST_SIZE = 512 * 1024 # Telegram's maximum upload.getFile chunk
ACCESS_HASH_KEY = "channel_access_hash:{}"
def classify_media(message, include_documents: bool) -> MediaInfo | None:
"""Return MediaInfo for a supported video message, else None.
Excludes photos, stickers, GIF animations, round video notes, voice/audio and
non-video documents (unless they are video files and documents are enabled).
"""
from telethon.tl import types
media = getattr(message, "media", None)
if not isinstance(media, types.MessageMediaDocument):
return None
doc = media.document
if not isinstance(doc, types.Document):
return None
video_attr = None
file_name = None
for attr in doc.attributes:
if isinstance(attr, (types.DocumentAttributeAnimated, types.DocumentAttributeSticker)):
return None
if isinstance(attr, types.DocumentAttributeVideo):
video_attr = attr
elif isinstance(attr, types.DocumentAttributeFilename):
file_name = attr.file_name
mime = (doc.mime_type or "").lower()
if video_attr is not None:
if getattr(video_attr, "round_message", False):
return None
kind = "video"
elif include_documents and (
mime.startswith("video/") or (file_name and split_video_extension(file_name)[1].lower() in VIDEO_EXTENSIONS)
):
kind = "document"
else:
return None
return MediaInfo(
kind=kind,
size=doc.size,
mime_type=doc.mime_type,
file_name=file_name,
media_id=doc.id,
duration=getattr(video_attr, "duration", None),
)
def to_source_message(message, channel_id: int, include_documents: bool) -> SourceMessage:
return SourceMessage(
channel_id=channel_id,
message_id=message.id,
caption=getattr(message, "message", None) or "",
media=classify_media(message, include_documents),
date=getattr(message, "date", None),
)
def map_telegram_error(exc: BaseException) -> BotError:
from telethon import errors
if isinstance(exc, errors.FloodWaitError):
return TransientError(
f"Telegram flood wait {exc.seconds}s", retry_after=exc.seconds, rate_limited=True,
public=f"Telegram rate limit; waiting {exc.seconds}s",
)
if isinstance(exc, (errors.ChannelPrivateError, errors.ChatAdminRequiredError, errors.ChannelInvalidError)):
return PermanentError(str(exc), public="No access to the source channel (check membership/admin rights)")
if isinstance(exc, (errors.AuthKeyError, errors.UnauthorizedError)):
return PermanentError(str(exc), public="Telegram session is no longer authorized; log in again")
if isinstance(exc, errors.RPCError):
return TransientError(f"Telegram error: {exc}", public=f"Telegram error ({type(exc).__name__})")
return TransientError(f"Telegram connection error: {exc}", public="Telegram connection problem")
def build_client(config, role: str):
"""Create a Telethon client. role: 'user' or 'bot'."""
from telethon import TelegramClient
from telethon.sessions import StringSession
if role == "user" and config.session_string:
session = StringSession(config.session_string)
else:
config.session_dir.mkdir(parents=True, exist_ok=True)
os.chmod(config.session_dir, 0o700)
session = str(config.session_dir / role)
return TelegramClient(
session,
config.api_id,
config.api_hash,
flood_sleep_threshold=config.flood_sleep_threshold,
connection_retries=10,
retry_delay=5,
auto_reconnect=True,
)
def parse_channel_reference(ref: str) -> str | int:
ref = ref.strip()
if re.fullmatch(r"-?\d+", ref):
return int(ref)
private = re.match(r"(?:https?://)?t\.me/c/(\d+)", ref)
if private:
return marked_channel_id(int(private.group(1)))
public = re.match(r"(?:https?://)?t\.me/([A-Za-z0-9_]+)", ref)
if public:
return public.group(1)
return ref.lstrip("@")
async def resolve_channel(client, ref: str, *, is_bot: bool, db=None):
"""Resolve the configured channel to a Telethon entity.
The channel's access hash is cached in the database, so a bot (which cannot
list dialogs) still resolves a private channel after a restart on an ephemeral
filesystem.
"""
from telethon import utils
from telethon.tl import types
target = parse_channel_reference(ref)
try:
entity = await client.get_entity(target)
except (ValueError, TypeError) as first_error:
entity = None
if isinstance(target, int):
real_id, _ = utils.resolve_id(target)
cached = db.get_kv(ACCESS_HASH_KEY.format(real_id)) if db else None
if cached:
entity = await client.get_entity(types.InputPeerChannel(real_id, int(cached)))
elif not is_bot:
await client.get_dialogs() # fills the entity cache for chats the account is in
entity = await client.get_entity(target)
if entity is None:
raise PermanentError(
f"cannot resolve channel {ref!r}: {first_error}",
public=(
"Cannot resolve SOURCE_CHANNEL_ID. Use the @username for public channels. For a private channel "
"in bot mode, make the bot an admin and post one message in the channel so the bot learns it."
),
) from first_error
if db is not None and getattr(entity, "access_hash", None):
db.set_kv(ACCESS_HASH_KEY.format(entity.id), str(entity.access_hash))
return entity
class TelethonSource:
def __init__(self, client, entity, *, channel_id: int, mode: str, config, db=None):
self.client = client
self.entity = entity
self.channel_id = channel_id
self.mode = mode
self.config = config
self.db = db
self._entities: dict[int, object] = {}
@property
def channel_title(self) -> str | None:
return getattr(self.entity, "title", None)
def convert(self, message, channel_id: int | None = None) -> SourceMessage:
return to_source_message(message, channel_id or self.channel_id, self.config.include_document_videos)
async def _entity_for(self, channel_id: int | None):
if channel_id is None or channel_id == self.channel_id:
if self.entity is not None:
return self.entity
if channel_id in self._entities:
return self._entities[channel_id]
entity = await resolve_channel(self.client, str(channel_id), is_bot=self.mode == "bot", db=self.db)
self._entities[channel_id] = entity
return entity
# -- /upload task support ----------------------------------------------------------
async def resolve_link(self, link) -> tuple[int, str | None]:
"""Resolve a MessageLink's channel through the authenticated client.
Raises PermanentError if this account/bot cannot access the channel.
"""
from telethon import utils
try:
entity = await resolve_channel(self.client, link.channel_ref, is_bot=self.mode == "bot", db=self.db)
except BotError:
raise
except Exception as exc:
raise PermanentError(
f"cannot resolve {link.url}: {exc}",
public="This channel is not accessible to the bot's Telegram account (join it, or make the bot an admin).",
) from exc
channel_id = utils.get_peer_id(entity)
self._entities[channel_id] = entity
return channel_id, getattr(entity, "title", None)
async def message_accessible(self, channel_id: int, message_id: int) -> bool:
entity = await self._entity_for(channel_id)
try:
message = await self.client.get_messages(entity, ids=message_id)
except Exception as exc:
raise map_telegram_error(exc) from exc
return message is not None and message.__class__.__name__ != "MessageEmpty"
async def iter_range(self, channel_id: int, start_id: int, end_id: int) -> AsyncIterator[tuple[int, SourceMessage | None]]:
"""Yield (message_id, SourceMessage or None) for every ID in [start_id, end_id], ascending.
Uses messages.getMessages / channels.getMessages by ID, which both user
accounts and bots (channel admins) may call; no history listing is needed.
"""
entity = await self._entity_for(channel_id)
current = start_id
try:
while current <= end_id:
ids = list(range(current, min(current + 100, end_id + 1)))
messages = await self.client.get_messages(entity, ids=ids)
for message_id, message in zip(ids, messages):
if message is None or message.__class__.__name__ in ("MessageEmpty", "MessageService"):
yield message_id, None
else:
yield message_id, self.convert(message, channel_id)
current = ids[-1] + 1
except (asyncio.CancelledError, GeneratorExit):
raise
except BotError:
raise
except Exception as exc:
raise map_telegram_error(exc) from exc
# -- discovery for the configured channel ------------------------------------------
async def iter_history(self, *, min_id: int, start_date: datetime | None) -> AsyncIterator[SourceMessage]:
try:
if self.mode == "user":
kwargs = {"reverse": True, "min_id": max(min_id - 1, 0)}
if start_date:
kwargs["offset_date"] = start_date
async for message in self.client.iter_messages(self.entity, **kwargs):
yield self.convert(message)
else:
async for message in self._probe_history(min_id, start_date):
yield self.convert(message)
except (asyncio.CancelledError, GeneratorExit):
raise
except BotError:
raise
except Exception as exc:
raise map_telegram_error(exc) from exc
async def _probe_history(self, start_id: int, start_date: datetime | None):
"""Bots cannot call messages.getHistory, but may fetch messages by ID."""
current, empty = max(start_id, 1), 0
end = self.config.bot_scan_end_id
while True:
ids = list(range(current, current + 100))
if end:
ids = [i for i in ids if i <= end]
if not ids:
return
messages = await self.client.get_messages(self.entity, ids=ids)
for message in messages:
if message is None or message.__class__.__name__ == "MessageEmpty":
empty += 1
continue
empty = 0
if start_date and message.date and message.date < start_date:
continue
yield message
if not end and empty >= self.config.bot_scan_max_empty:
return
current += 100
async def _fetch(self, message_id: int, channel_id: int | None = None):
entity = await self._entity_for(channel_id)
try:
message = await self.client.get_messages(entity, ids=message_id)
except Exception as exc:
raise map_telegram_error(exc) from exc
if message is None or getattr(message, "media", None) is None or not getattr(message, "document", None):
raise PermanentError(
f"message {message_id} no longer exists or has no media", public="Source message was deleted or has no video"
)
return message
async def iter_media_chunks(
self, message_id: int, *, offset: int, expected_size: int | None, channel_id: int | None = None
) -> AsyncIterator[bytes]:
from telethon import errors
message = await self._fetch(message_id, channel_id)
position, refreshes = offset, 0
while True:
try:
async for chunk in self.client.iter_download(
message.document, offset=position, request_size=REQUEST_SIZE, file_size=message.document.size
):
position += len(chunk)
yield chunk
return
except errors.FileReferenceExpiredError:
refreshes += 1
if refreshes > 3:
raise TransientError("file reference kept expiring", public="Telegram file reference expired")
message = await self._fetch(message_id, channel_id)
except (asyncio.CancelledError, GeneratorExit):
raise
except BotError:
raise
except OSError:
raise
except Exception as exc:
raise map_telegram_error(exc) from exc
async def latest_message_id(self) -> int | None:
if self.mode != "user":
return None
messages = await self.client.get_messages(self.entity, limit=1)
return messages[0].id if messages else None