first commit
This commit is contained in:
gosangam
@@ -0,0 +1,345 @@
|
||||
"""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
|
||||
Reference in New Issue
Block a user