import asyncio import hashlib import os import time from typing import Iterable, Optional import click from telethon import TelegramClient, utils, helpers, custom from telethon.crypto import AES from telethon.errors import RPCError, FloodWaitError, InvalidBufferError, FloodError from telethon.tl import types, functions, TLRequest from telethon.utils import pack_bot_file_id from telegram_upload.client.progress_bar import get_progress_bar from telegram_upload.exceptions import TelegramUploadDataLoss, MissingFileError from telegram_upload.upload_files import File from telegram_upload.utils import grouper, async_to_sync, get_environment_integer PARALLEL_UPLOAD_BLOCKS = get_environment_integer('TELEGRAM_UPLOAD_PARALLEL_UPLOAD_BLOCKS', 8) ALBUM_FILES = 10 RETRIES = 10 MAX_RECONNECT_RETRIES = get_environment_integer('TELEGRAM_UPLOAD_MAX_RECONNECT_RETRIES', 10) RECONNECT_TIMEOUT = get_environment_integer('TELEGRAM_UPLOAD_RECONNECT_TIMEOUT', 10) MIN_RECONNECT_WAIT = get_environment_integer('TELEGRAM_UPLOAD_MIN_RECONNECT_WAIT', 10) class TelegramUploadClient(TelegramClient): parallel_upload_blocks = PARALLEL_UPLOAD_BLOCKS def __init__(self, *args, **kwargs): self.reconnecting_lock = asyncio.Lock() self.upload_semaphore = asyncio.Semaphore(self.parallel_upload_blocks) super().__init__(*args, **kwargs) def forward_to(self, message, destinations): for destination in destinations: self.forward_messages(destination, [message]) async def _send_album_media(self, entity, media): entity = await self.get_input_entity(entity) request = functions.messages.SendMultiMediaRequest( entity, multi_media=media, silent=None, schedule_date=None, clear_draft=None ) result = await self(request) random_ids = [m.random_id for m in media] return self._get_response_message(random_ids, result, entity) def send_files_as_album(self, entity, files, delete_on_success=False, print_file_id=False, forward=()): for files_group in grouper(ALBUM_FILES, files): media = self.send_files(entity, files_group, delete_on_success, print_file_id, forward, send_as_media=True) async_to_sync(self._send_album_media(entity, media)) def _send_file_message(self, entity, file, thumb, progress): message = self.send_file(entity, file, thumb=thumb, file_size=file.file_size if isinstance(file, File) else None, caption=file.file_caption, force_document=file.force_file, progress_callback=progress, attributes=file.file_attributes) if hasattr(message.media, 'document') and file.file_size != message.media.document.size: raise TelegramUploadDataLoss( 'Remote document size: {} bytes (local file size: {} bytes)'.format( message.media.document.size, file.file_size)) return message async def _send_media(self, entity, file: File, progress): entity = await self.get_input_entity(entity) supports_streaming = False # TODO fh, fm, _ = await self._file_to_media( file, supports_streaming=file, progress_callback=progress) if isinstance(fm, types.InputMediaUploadedPhoto): r = await self(functions.messages.UploadMediaRequest( entity, media=fm )) fm = utils.get_input_media(r.photo) elif isinstance(fm, types.InputMediaUploadedDocument): r = await self(functions.messages.UploadMediaRequest( entity, media=fm )) fm = utils.get_input_media( r.document, supports_streaming=supports_streaming) return types.InputSingleMedia( fm, message=file.short_name, entities=None, # random_id is autogenerated ) def send_one_file(self, entity, file: File, send_as_media: bool = False, thumb: Optional[str] = None, retries=RETRIES): message = None progress, bar = get_progress_bar('Uploading', file.file_name, file.file_size) try: try: # TODO: remove distinction? if send_as_media: message = async_to_sync(self._send_media(entity, file, progress)) else: message = self._send_file_message(entity, file, thumb, progress) finally: bar.render_finish() except (FloodWaitError,FloodError) as e: click.echo(f'{e}. Waiting for {e.seconds} seconds.', err=True) time.sleep(e.seconds) message = self.send_one_file(entity, file, send_as_media, thumb, retries) except RPCError as e: if retries > 0: click.echo(f'The file "{file.file_name}" could not be uploaded: {e}. Retrying...', err=True) message = self.send_one_file(entity, file, send_as_media, thumb, retries - 1) else: click.echo(f'The file "{file.file_name}" could not be uploaded: {e}. It will not be retried.', err=True) return message def send_files(self, entity, files: Iterable[File], delete_on_success=False, print_file_id=False, forward=(), send_as_media: bool = False): has_files = False messages = [] for file in files: has_files = True thumb = file.get_thumbnail() try: message = self.send_one_file(entity, file, send_as_media, thumb=thumb) finally: if thumb and not file.is_custom_thumbnail and os.path.lexists(thumb): os.remove(thumb) if message is None: click.echo('Failed to upload file "{}"'.format(file.file_name), err=True) if message and print_file_id: click.echo('Uploaded successfully "{}" (file_id {})'.format(file.file_name, pack_bot_file_id(message.media))) if message and delete_on_success: click.echo('Deleting "{}"'.format(file)) os.remove(file.path) if message: self.forward_to(message, forward) messages.append(message) if not has_files: raise MissingFileError('Files do not exist.') return messages async def upload_file( self: 'TelegramClient', file: 'hints.FileLike', *, part_size_kb: float = None, file_size: int = None, file_name: str = None, use_cache: type = None, key: bytes = None, iv: bytes = None, progress_callback: 'hints.ProgressCallback' = None) -> 'types.TypeInputFile': """ Uploads a file to Telegram's servers, without sending it. .. note:: Generally, you want to use `send_file` instead. This method returns a handle (an instance of :tl:`InputFile` or :tl:`InputFileBig`, as required) which can be later used before it expires (they are usable during less than a day). Uploading a file will simply return a "handle" to the file stored remotely in the Telegram servers, which can be later used on. This will **not** upload the file to your own chat or any chat at all. Arguments file (`str` | `bytes` | `file`): The path of the file, byte array, or stream that will be sent. Note that if a byte array or a stream is given, a filename or its type won't be inferred, and it will be sent as an "unnamed application/octet-stream". part_size_kb (`int`, optional): Chunk size when uploading files. The larger, the less requests will be made (up to 512KB maximum). file_size (`int`, optional): The size of the file to be uploaded, which will be determined automatically if not specified. If the file size can't be determined beforehand, the entire file will be read in-memory to find out how large it is. file_name (`str`, optional): The file name which will be used on the resulting InputFile. If not specified, the name will be taken from the ``file`` and if this is not a `str`, it will be ``"unnamed"``. use_cache (`type`, optional): This parameter currently does nothing, but is kept for backward-compatibility (and it may get its use back in the future). key ('bytes', optional): In case of an encrypted upload (secret chats) a key is supplied iv ('bytes', optional): In case of an encrypted upload (secret chats) an iv is supplied progress_callback (`callable`, optional): A callback function accepting two parameters: ``(sent bytes, total)``. When sending an album, the callback will receive a number between 0 and the amount of files as the "sent" parameter, and the amount of files as the "total". Note that the first parameter will be a floating point number to indicate progress within a file (e.g. ``2.5`` means it has sent 50% of the third file, because it's between 2 and 3). Returns :tl:`InputFileBig` if the file size is larger than 10MB, `InputSizedFile ` (subclass of :tl:`InputFile`) otherwise. Example .. code-block:: python # Photos as photo and document file = await client.upload_file('photo.jpg') await client.send_file(chat, file) # sends as photo await client.send_file(chat, file, force_document=True) # sends as document file.name = 'not a photo.jpg' await client.send_file(chat, file, force_document=True) # document, new name # As song or as voice note file = await client.upload_file('song.ogg') await client.send_file(chat, file) # sends as song await client.send_file(chat, file, voice_note=True) # sends as voice note """ if isinstance(file, (types.InputFile, types.InputFileBig)): return file # Already uploaded async with helpers._FileStream(file, file_size=file_size) as stream: # Opening the stream will determine the correct file size file_size = stream.file_size if not part_size_kb: part_size_kb = utils.get_appropriated_part_size(file_size) if part_size_kb > 512: raise ValueError('The part size must be less or equal to 512KB') part_size = int(part_size_kb * 1024) if part_size % 1024 != 0: raise ValueError( 'The part size must be evenly divisible by 1024') # Set a default file name if None was specified file_id = helpers.generate_random_long() if not file_name: file_name = stream.name or str(file_id) # If the file name lacks extension, add it if possible. # Else Telegram complains with `PHOTO_EXT_INVALID_ERROR` # even if the uploaded image is indeed a photo. if not os.path.splitext(file_name)[-1]: file_name += utils._get_extension(stream) # Determine whether the file is too big (over 10MB) or not # Telegram does make a distinction between smaller or larger files is_big = file_size > 10 * 1024 * 1024 hash_md5 = hashlib.md5() part_count = (file_size + part_size - 1) // part_size self._log[__name__].info('Uploading file of %d bytes in %d chunks of %d', file_size, part_count, part_size) pos = 0 for part_index in range(part_count): # Read the file by in chunks of size part_size part = await helpers._maybe_await(stream.read(part_size)) if not isinstance(part, bytes): raise TypeError( 'file descriptor returned {}, not bytes (you must ' 'open the file in bytes mode)'.format(type(part))) # `file_size` could be wrong in which case `part` may not be # `part_size` before reaching the end. if len(part) != part_size and part_index < part_count - 1: raise ValueError( 'read less than {} before reaching the end; either ' '`file_size` or `read` are wrong'.format(part_size)) pos += len(part) # Encryption part if needed if key and iv: part = AES.encrypt_ige(part, key, iv) if not is_big: # Bit odd that MD5 is only needed for small files and not # big ones with more chance for corruption, but that's # what Telegram wants. hash_md5.update(part) # The SavePartRequest is different depending on whether # the file is too large or not (over or less than 10MB) if is_big: request = functions.upload.SaveBigFilePartRequest( file_id, part_index, part_count, part) else: request = functions.upload.SaveFilePartRequest( file_id, part_index, part) await self.upload_semaphore.acquire() self.loop.create_task( self._send_file_part(request, part_index, part_count, pos, file_size, progress_callback), name=f"telegram-upload-file-{part_index}" ) # Wait for all tasks to finish await asyncio.wait([ task for task in asyncio.all_tasks() if task.get_name().startswith(f"telegram-upload-file-") ]) if is_big: return types.InputFileBig(file_id, part_count, file_name) else: return custom.InputSizedFile( file_id, part_count, file_name, md5=hash_md5, size=file_size ) # endregion async def _send_file_part(self, request: TLRequest, part_index: int, part_count: int, pos: int, file_size: int, progress_callback: Optional['hints.ProgressCallback'] = None, retry: int = 0) -> None: """ Submit the file request part to Telegram. This method waits for the request to be executed, logs the upload, and releases the semaphore to allow further uploading. :param request: SaveBigFilePartRequest or SaveFilePartRequest. This request will be awaited. :param part_index: Part index as integer. Used in logging. :param part_count: Total parts count as integer. Used in logging. :param pos: Number of part as integer. Used for progress bar. :param file_size: Total file size. Used for progress bar. :param progress_callback: Callback to use after submit the request. Optional. :return: None """ result = None try: result = await self(request) except InvalidBufferError as e: if e.code == 429: # Too many connections click.echo(f'Too many connections to Telegram servers.', err=True) else: raise except ConnectionError: # Retry to send the file part click.echo(f'Detected connection error. Retrying...', err=True) except FloodError as e: print(e) else: self.upload_semaphore.release() if result is None and retry < MAX_RECONNECT_RETRIES: # An error occurred, retry await asyncio.sleep(max(MIN_RECONNECT_WAIT, retry * MIN_RECONNECT_WAIT)) await self.reconnect() await self._send_file_part( request, part_index, part_count, pos, file_size, progress_callback, retry + 1 ) elif result: self._log[__name__].debug('Uploaded %d/%d', part_index + 1, part_count) if progress_callback: await helpers._maybe_await(progress_callback(pos, file_size)) else: raise RuntimeError( 'Failed to upload file part {}.'.format(part_index)) def decrease_upload_semaphore(self): """ Decreases the upload semaphore by one. This method is used to reduce the number of parallel uploads. :return: """ if self.parallel_upload_blocks > 1: self.parallel_upload_blocks -= 1 self.loop.create_task(self.upload_semaphore.acquire()) async def reconnect(self): """ Reconnects to Telegram servers. :return: None """ await self.reconnecting_lock.acquire() if self.is_connected(): # Reconnected in another task self.reconnecting_lock.release() return self.decrease_upload_semaphore() try: click.echo(f'Reconnecting to Telegram servers...') await asyncio.wait_for(self.connect(), RECONNECT_TIMEOUT) click.echo(f'Reconnected to Telegram servers.') except InvalidBufferError: click.echo(f'InvalidBufferError connecting to Telegram servers.', err=True) except asyncio.TimeoutError: click.echo(f'Timeout connecting to Telegram servers.', err=True) finally: self.reconnecting_lock.release()