From 4ff0756ffca049580cdb1088d10d6492723bbc4d Mon Sep 17 00:00:00 2001 From: Lonami Exo Date: Sun, 10 Jun 2018 12:04:23 +0200 Subject: [PATCH] Separate download requests from the TelegramClient --- telethon/client/downloads.py | 430 +++++++++++++++++++++++++++++ telethon/client/telegramclient.py | 432 +----------------------------- telethon/crypto/cdn_decrypter.py | 6 +- 3 files changed, 436 insertions(+), 432 deletions(-) create mode 100644 telethon/client/downloads.py diff --git a/telethon/client/downloads.py b/telethon/client/downloads.py new file mode 100644 index 00000000..80b5b1e6 --- /dev/null +++ b/telethon/client/downloads.py @@ -0,0 +1,430 @@ +import datetime +import io +import logging +import os + +from .users import UserMethods +from .. import utils, helpers, errors +from ..crypto import CdnDecrypter +from ..tl import TLObject, types, functions + +__log__ = logging.getLogger(__name__) + + +class DownloadMethods(UserMethods): + + # region Public methods + + async def download_profile_photo( + self, entity, file=None, download_big=True): + """ + Downloads the profile photo of the given entity (user/chat/channel). + + Args: + entity (`entity`): + From who the photo will be downloaded. + + file (`str` | `file`, optional): + The output file path, directory, or stream-like object. + If the path exists and is a file, it will be overwritten. + + download_big (`bool`, optional): + Whether to use the big version of the available photos. + + Returns: + ``None`` if no photo was provided, or if it was Empty. On success + the file path is returned since it may differ from the one given. + """ + # hex(crc32(x.encode('ascii'))) for x in + # ('User', 'Chat', 'UserFull', 'ChatFull') + ENTITIES = (0x2da17977, 0xc5af5d94, 0x1f4661b9, 0xd49a2697) + # ('InputPeer', 'InputUser', 'InputChannel') + INPUTS = (0xc91c90b6, 0xe669bf46, 0x40f202fd) + if not isinstance(entity, TLObject) or entity.SUBCLASS_OF_ID in INPUTS: + entity = await self.get_entity(entity) + + possible_names = [] + if entity.SUBCLASS_OF_ID not in ENTITIES: + photo = entity + else: + if not hasattr(entity, 'photo'): + # Special case: may be a ChatFull with photo:Photo + # This is different from a normal UserProfilePhoto and Chat + if not hasattr(entity, 'chat_photo'): + return None + + return await self._download_photo( + entity.chat_photo, file, date=None, progress_callback=None) + + for attr in ('username', 'first_name', 'title'): + possible_names.append(getattr(entity, attr, None)) + + photo = entity.photo + + if isinstance(photo, (types.UserProfilePhoto, types.ChatPhoto)): + loc = photo.photo_big if download_big else photo.photo_small + else: + try: + loc = utils.get_input_location(photo) + except TypeError: + return None + + file = self._get_proper_filename( + file, 'profile_photo', '.jpg', + possible_names=possible_names + ) + + try: + await self.download_file(loc, file) + return file + except errors.LocationInvalidError: + # See issue #500, Android app fails as of v4.6.0 (1155). + # The fix seems to be using the full channel chat photo. + ie = await self.get_input_entity(entity) + if isinstance(ie, types.InputPeerChannel): + full = await self(functions.channels.GetFullChannelRequest(ie)) + return await self._download_photo( + full.full_chat.chat_photo, file, + date=None, progress_callback=None + ) + else: + # Until there's a report for chats, no need to. + return None + + async def download_media(self, message, file=None, progress_callback=None): + """ + Downloads the given media, or the media from a specified Message. + + Note that if the download is too slow, you should consider installing + ``cryptg`` (through ``pip install cryptg``) so that decrypting the + received data is done in C instead of Python (much faster). + + message (:tl:`Message` | :tl:`Media`): + The media or message containing the media that will be downloaded. + + file (`str` | `file`, optional): + The output file path, directory, or stream-like object. + If the path exists and is a file, it will be overwritten. + + progress_callback (`callable`, optional): + A callback function accepting two parameters: + ``(received bytes, total)``. + + Returns: + ``None`` if no media was provided, or if it was Empty. On success + the file path is returned since it may differ from the one given. + """ + # TODO This won't work for messageService + if isinstance(message, types.Message): + date = message.date + media = message.media + else: + date = datetime.datetime.now() + media = message + + if isinstance(media, types.MessageMediaWebPage): + if isinstance(media.webpage, types.WebPage): + media = media.webpage.document or media.webpage.photo + + if isinstance(media, (types.MessageMediaPhoto, types.Photo, + types.PhotoSize, types.PhotoCachedSize)): + return await self._download_photo( + media, file, date, progress_callback + ) + elif isinstance(media, (types.MessageMediaDocument, types.Document)): + return await self._download_document( + media, file, date, progress_callback + ) + elif isinstance(media, types.MessageMediaContact): + return self._download_contact( + media, file + ) + + async def download_file( + self, input_location, file=None, part_size_kb=None, + file_size=None, progress_callback=None): + """ + Downloads the given input location to a file. + + Args: + input_location (:tl:`FileLocation` | :tl:`InputFileLocation`): + The file location from which the file will be downloaded. + See `telethon.utils.get_input_location` source for a complete + list of supported types. + + file (`str` | `file`, optional): + The output file path, directory, or stream-like object. + If the path exists and is a file, it will be overwritten. + + If the file path is ``None``, then the result will be + saved in memory and returned as `bytes`. + + part_size_kb (`int`, optional): + Chunk size when downloading files. The larger, the less + requests will be made (up to 512KB maximum). + + file_size (`int`, optional): + The file size that is about to be downloaded, if known. + Only used if ``progress_callback`` is specified. + + progress_callback (`callable`, optional): + A callback function accepting two parameters: + ``(downloaded bytes, total)``. Note that the + ``total`` is the provided ``file_size``. + """ + if not part_size_kb: + if not file_size: + part_size_kb = 64 # Reasonable default + else: + part_size_kb = utils.get_appropriated_part_size(file_size) + + part_size = int(part_size_kb * 1024) + # https://core.telegram.org/api/files says: + # > part_size % 1024 = 0 (divisible by 1KB) + # + # But https://core.telegram.org/cdn (more recent) says: + # > limit must be divisible by 4096 bytes + # So we just stick to the 4096 limit. + if part_size % 4096 != 0: + raise ValueError( + 'The part size must be evenly divisible by 4096.') + + in_memory = file is None + if in_memory: + f = io.BytesIO() + elif isinstance(file, str): + # Ensure that we'll be able to download the media + helpers.ensure_parent_dir_exists(file) + f = open(file, 'wb') + else: + f = file + + # The used client will change if FileMigrateError occurs + client = self + cdn_decrypter = None + input_location = utils.get_input_location(input_location) + + __log__.info('Downloading file in chunks of %d bytes', part_size) + try: + offset = 0 + while True: + try: + if cdn_decrypter: + result = cdn_decrypter.get_file() + else: + result = client(functions.upload.GetFileRequest( + input_location, offset, part_size + )) + + if isinstance(result, types.upload.FileCdnRedirect): + __log__.info('File lives in a CDN') + cdn_decrypter, result = \ + await CdnDecrypter.prepare_decrypter( + client, await self._get_cdn_client(result), + result + ) + + except errors.FileMigrateError as e: + __log__.info('File lives in another DC') + client = await self._get_exported_client(e.new_dc) + continue + + offset += part_size + + # If we have received no data (0 bytes), the file is over + # So there is nothing left to download and write + if not result.bytes: + # Return some extra information, unless it's a CDN file + if in_memory: + f.flush() + return f.getvalue() + else: + return getattr(result, 'type', '') + + f.write(result.bytes) + __log__.debug('Saved %d more bytes', len(result.bytes)) + if progress_callback: + progress_callback(f.tell(), file_size) + finally: + if client != self: + await client.disconnect() + if cdn_decrypter: + await cdn_decrypter.client.disconnect() + if isinstance(file, str) or in_memory: + f.close() + + # endregion + + # region Private methods + + async def _download_photo(self, photo, file, date, progress_callback): + """Specialized version of .download_media() for photos""" + # Determine the photo and its largest size + if isinstance(photo, types.MessageMediaPhoto): + photo = photo.photo + if isinstance(photo, types.Photo): + for size in reversed(photo.sizes): + if not isinstance(size, types.PhotoSizeEmpty): + photo = size + break + else: + return + if not isinstance(photo, (types.PhotoSize, types.PhotoCachedSize)): + return + + file = self._get_proper_filename(file, 'photo', '.jpg', date=date) + if isinstance(photo, types.PhotoCachedSize): + # No need to download anything, simply write the bytes + if isinstance(file, str): + helpers.ensure_parent_dir_exists(file) + f = open(file, 'wb') + else: + f = file + try: + f.write(photo.bytes) + finally: + if isinstance(file, str): + f.close() + return file + + await self.download_file( + photo.location, file, file_size=photo.size, + progress_callback=progress_callback) + return file + + async def _download_document( + self, document, file, date, progress_callback): + """Specialized version of .download_media() for documents.""" + if isinstance(document, types.MessageMediaDocument): + document = document.document + if not isinstance(document, types.Document): + return + + file_size = document.size + + kind = 'document' + possible_names = [] + for attr in document.attributes: + if isinstance(attr, types.DocumentAttributeFilename): + possible_names.insert(0, attr.file_name) + + elif isinstance(attr, types.DocumentAttributeAudio): + kind = 'audio' + if attr.performer and attr.title: + possible_names.append('{} - {}'.format( + attr.performer, attr.title + )) + elif attr.performer: + possible_names.append(attr.performer) + elif attr.title: + possible_names.append(attr.title) + elif attr.voice: + kind = 'voice' + + file = self._get_proper_filename( + file, kind, utils.get_extension(document), + date=date, possible_names=possible_names + ) + + await self.download_file( + document, file, file_size=file_size, + progress_callback=progress_callback) + return file + + @classmethod + def _download_contact(cls, mm_contact, file): + """ + Specialized version of .download_media() for contacts. + Will make use of the vCard 4.0 format. + """ + first_name = mm_contact.first_name + last_name = mm_contact.last_name + phone_number = mm_contact.phone_number + + if isinstance(file, str): + file = cls._get_proper_filename( + file, 'contact', '.vcard', + possible_names=[first_name, phone_number, last_name] + ) + f = open(file, 'w', encoding='utf-8') + else: + f = file + + try: + # Remove these pesky characters + first_name = first_name.replace(';', '') + last_name = (last_name or '').replace(';', '') + f.write('BEGIN:VCARD\n') + f.write('VERSION:4.0\n') + f.write('N:{};{};;;\n'.format(first_name, last_name)) + f.write('FN:{} {}\n'.format(first_name, last_name)) + f.write('TEL;TYPE=cell;VALUE=uri:tel:+{}\n'.format(phone_number)) + f.write('END:VCARD\n') + finally: + # Only close the stream if we opened it + if isinstance(file, str): + f.close() + + return file + + @staticmethod + def _get_proper_filename(file, kind, extension, + date=None, possible_names=None): + """Gets a proper filename for 'file', if this is a path. + + 'kind' should be the kind of the output file (photo, document...) + 'extension' should be the extension to be added to the file if + the filename doesn't have any yet + 'date' should be when this file was originally sent, if known + 'possible_names' should be an ordered list of possible names + + If no modification is made to the path, any existing file + will be overwritten. + If any modification is made to the path, this method will + ensure that no existing file will be overwritten. + """ + if file is not None and not isinstance(file, str): + # Probably a stream-like object, we cannot set a filename here + return file + + if file is None: + file = '' + elif os.path.isfile(file): + # Make no modifications to valid existing paths + return file + + if os.path.isdir(file) or not file: + try: + name = None if possible_names is None else next( + x for x in possible_names if x + ) + except StopIteration: + name = None + + if not name: + if not date: + date = datetime.datetime.now() + name = '{}_{}-{:02}-{:02}_{:02}-{:02}-{:02}'.format( + kind, + date.year, date.month, date.day, + date.hour, date.minute, date.second, + ) + file = os.path.join(file, name) + + directory, name = os.path.split(file) + name, ext = os.path.splitext(name) + if not ext: + ext = extension + + result = os.path.join(directory, name + ext) + if not os.path.isfile(result): + return result + + i = 1 + while True: + result = os.path.join(directory, '{} ({}){}'.format(name, i, ext)) + if not os.path.isfile(result): + return result + i += 1 + + # endregion diff --git a/telethon/client/telegramclient.py b/telethon/client/telegramclient.py index 084fce13..30352217 100644 --- a/telethon/client/telegramclient.py +++ b/telethon/client/telegramclient.py @@ -1,20 +1,14 @@ import getpass import hashlib -import io import logging import sys import warnings -from ..crypto import CdnDecrypter from ..tl.functions.help import AcceptTermsOfServiceRequest from ..tl.functions.updates import GetDifferenceRequest -from ..tl.functions.upload import ( - GetFileRequest -) from ..tl.types.updates import ( DifferenceSlice, DifferenceEmpty, Difference, DifferenceTooLong ) -from ..tl.types.upload import FileCdnRedirect try: import socks @@ -26,8 +20,8 @@ from .telegrambaseclient import TelegramBaseClient from .. import helpers, events from ..errors import ( PhoneCodeEmptyError, PhoneCodeExpiredError, - PhoneCodeHashEmptyError, PhoneCodeInvalidError, LocationInvalidError, - SessionPasswordNeededError, FileMigrateError, PhoneNumberUnoccupiedError, + PhoneCodeHashEmptyError, PhoneCodeInvalidError, SessionPasswordNeededError, + PhoneNumberUnoccupiedError, PhoneNumberOccupiedError ) from ..tl.functions.account import ( @@ -38,25 +32,15 @@ from ..tl.functions.auth import ( SignUpRequest, ResendCodeRequest, ImportBotAuthorizationRequest ) -from ..tl.functions.channels import ( - GetFullChannelRequest -) from ..tl.types import ( - DocumentAttributeAudio, DocumentAttributeFilename, - Message, MessageMediaContact, MessageMediaDocument, MessageMediaPhoto, - UserProfilePhoto, ChatPhoto, UpdateNewMessage, InputPeerChannel, Photo, - Document, Updates, - MessageMediaWebPage, PhotoSize, PhotoCachedSize, - PhotoSizeEmpty, WebPage + UpdateNewMessage, Updates ) from ..tl.types.account import PasswordInputSettings, NoPassword __log__ = logging.getLogger(__name__) import os -from datetime import datetime from .. import utils from ..errors import RPCError -from ..tl import TLObject class TelegramClient(TelegramBaseClient): @@ -395,416 +379,6 @@ class TelegramClient(TelegramBaseClient): # region Downloading media requests - def download_profile_photo(self, entity, file=None, download_big=True): - """ - Downloads the profile photo of the given entity (user/chat/channel). - - Args: - entity (`entity`): - From who the photo will be downloaded. - - file (`str` | `file`, optional): - The output file path, directory, or stream-like object. - If the path exists and is a file, it will be overwritten. - - download_big (`bool`, optional): - Whether to use the big version of the available photos. - - Returns: - ``None`` if no photo was provided, or if it was Empty. On success - the file path is returned since it may differ from the one given. - """ - # hex(crc32(x.encode('ascii'))) for x in - # ('User', 'Chat', 'UserFull', 'ChatFull') - ENTITIES = (0x2da17977, 0xc5af5d94, 0x1f4661b9, 0xd49a2697) - # ('InputPeer', 'InputUser', 'InputChannel') - INPUTS = (0xc91c90b6, 0xe669bf46, 0x40f202fd) - if not isinstance(entity, TLObject) or entity.SUBCLASS_OF_ID in INPUTS: - entity = self.get_entity(entity) - - possible_names = [] - if entity.SUBCLASS_OF_ID not in ENTITIES: - photo = entity - else: - if not hasattr(entity, 'photo'): - # Special case: may be a ChatFull with photo:Photo - # This is different from a normal UserProfilePhoto and Chat - if not hasattr(entity, 'chat_photo'): - return None - - return self._download_photo(entity.chat_photo, file, - date=None, progress_callback=None) - - for attr in ('username', 'first_name', 'title'): - possible_names.append(getattr(entity, attr, None)) - - photo = entity.photo - - if isinstance(photo, (UserProfilePhoto, ChatPhoto)): - loc = photo.photo_big if download_big else photo.photo_small - else: - try: - loc = utils.get_input_location(photo) - except TypeError: - return None - - file = self._get_proper_filename( - file, 'profile_photo', '.jpg', - possible_names=possible_names - ) - - try: - self.download_file(loc, file) - return file - except LocationInvalidError: - # See issue #500, Android app fails as of v4.6.0 (1155). - # The fix seems to be using the full channel chat photo. - ie = self.get_input_entity(entity) - if isinstance(ie, InputPeerChannel): - full = self(GetFullChannelRequest(ie)) - return self._download_photo( - full.full_chat.chat_photo, file, - date=None, progress_callback=None - ) - else: - # Until there's a report for chats, no need to. - return None - - def download_media(self, message, file=None, progress_callback=None): - """ - Downloads the given media, or the media from a specified Message. - - Note that if the download is too slow, you should consider installing - ``cryptg`` (through ``pip install cryptg``) so that decrypting the - received data is done in C instead of Python (much faster). - - message (:tl:`Message` | :tl:`Media`): - The media or message containing the media that will be downloaded. - - file (`str` | `file`, optional): - The output file path, directory, or stream-like object. - If the path exists and is a file, it will be overwritten. - - progress_callback (`callable`, optional): - A callback function accepting two parameters: - ``(received bytes, total)``. - - Returns: - ``None`` if no media was provided, or if it was Empty. On success - the file path is returned since it may differ from the one given. - """ - # TODO This won't work for messageService - if isinstance(message, Message): - date = message.date - media = message.media - else: - date = datetime.now() - media = message - - if isinstance(media, MessageMediaWebPage): - if isinstance(media.webpage, WebPage): - media = media.webpage.document or media.webpage.photo - - if isinstance(media, (MessageMediaPhoto, Photo, - PhotoSize, PhotoCachedSize)): - return self._download_photo( - media, file, date, progress_callback - ) - elif isinstance(media, (MessageMediaDocument, Document)): - return self._download_document( - media, file, date, progress_callback - ) - elif isinstance(media, MessageMediaContact): - return self._download_contact( - media, file - ) - - def _download_photo(self, photo, file, date, progress_callback): - """Specialized version of .download_media() for photos""" - # Determine the photo and its largest size - if isinstance(photo, MessageMediaPhoto): - photo = photo.photo - if isinstance(photo, Photo): - for size in reversed(photo.sizes): - if not isinstance(size, PhotoSizeEmpty): - photo = size - break - else: - return - if not isinstance(photo, (PhotoSize, PhotoCachedSize)): - return - - file = self._get_proper_filename(file, 'photo', '.jpg', date=date) - if isinstance(photo, PhotoCachedSize): - # No need to download anything, simply write the bytes - if isinstance(file, str): - helpers.ensure_parent_dir_exists(file) - f = open(file, 'wb') - else: - f = file - try: - f.write(photo.bytes) - finally: - if isinstance(file, str): - f.close() - return file - - self.download_file(photo.location, file, file_size=photo.size, - progress_callback=progress_callback) - return file - - def _download_document(self, document, file, date, progress_callback): - """Specialized version of .download_media() for documents.""" - if isinstance(document, MessageMediaDocument): - document = document.document - if not isinstance(document, Document): - return - - file_size = document.size - - kind = 'document' - possible_names = [] - for attr in document.attributes: - if isinstance(attr, DocumentAttributeFilename): - possible_names.insert(0, attr.file_name) - - elif isinstance(attr, DocumentAttributeAudio): - kind = 'audio' - if attr.performer and attr.title: - possible_names.append('{} - {}'.format( - attr.performer, attr.title - )) - elif attr.performer: - possible_names.append(attr.performer) - elif attr.title: - possible_names.append(attr.title) - elif attr.voice: - kind = 'voice' - - file = self._get_proper_filename( - file, kind, utils.get_extension(document), - date=date, possible_names=possible_names - ) - - self.download_file(document, file, file_size=file_size, - progress_callback=progress_callback) - return file - - @staticmethod - def _download_contact(mm_contact, file): - """Specialized version of .download_media() for contacts. - Will make use of the vCard 4.0 format. - """ - first_name = mm_contact.first_name - last_name = mm_contact.last_name - phone_number = mm_contact.phone_number - - if isinstance(file, str): - file = TelegramClient._get_proper_filename( - file, 'contact', '.vcard', - possible_names=[first_name, phone_number, last_name] - ) - f = open(file, 'w', encoding='utf-8') - else: - f = file - - try: - # Remove these pesky characters - first_name = first_name.replace(';', '') - last_name = (last_name or '').replace(';', '') - f.write('BEGIN:VCARD\n') - f.write('VERSION:4.0\n') - f.write('N:{};{};;;\n'.format(first_name, last_name)) - f.write('FN:{} {}\n'.format(first_name, last_name)) - f.write('TEL;TYPE=cell;VALUE=uri:tel:+{}\n'.format(phone_number)) - f.write('END:VCARD\n') - finally: - # Only close the stream if we opened it - if isinstance(file, str): - f.close() - - return file - - @staticmethod - def _get_proper_filename(file, kind, extension, - date=None, possible_names=None): - """Gets a proper filename for 'file', if this is a path. - - 'kind' should be the kind of the output file (photo, document...) - 'extension' should be the extension to be added to the file if - the filename doesn't have any yet - 'date' should be when this file was originally sent, if known - 'possible_names' should be an ordered list of possible names - - If no modification is made to the path, any existing file - will be overwritten. - If any modification is made to the path, this method will - ensure that no existing file will be overwritten. - """ - if file is not None and not isinstance(file, str): - # Probably a stream-like object, we cannot set a filename here - return file - - if file is None: - file = '' - elif os.path.isfile(file): - # Make no modifications to valid existing paths - return file - - if os.path.isdir(file) or not file: - try: - name = None if possible_names is None else next( - x for x in possible_names if x - ) - except StopIteration: - name = None - - if not name: - if not date: - date = datetime.now() - name = '{}_{}-{:02}-{:02}_{:02}-{:02}-{:02}'.format( - kind, - date.year, date.month, date.day, - date.hour, date.minute, date.second, - ) - file = os.path.join(file, name) - - directory, name = os.path.split(file) - name, ext = os.path.splitext(name) - if not ext: - ext = extension - - result = os.path.join(directory, name + ext) - if not os.path.isfile(result): - return result - - i = 1 - while True: - result = os.path.join(directory, '{} ({}){}'.format(name, i, ext)) - if not os.path.isfile(result): - return result - i += 1 - - def download_file(self, - input_location, - file=None, - part_size_kb=None, - file_size=None, - progress_callback=None): - """ - Downloads the given input location to a file. - - Args: - input_location (:tl:`FileLocation` | :tl:`InputFileLocation`): - The file location from which the file will be downloaded. - See `telethon.utils.get_input_location` source for a complete - list of supported types. - - file (`str` | `file`, optional): - The output file path, directory, or stream-like object. - If the path exists and is a file, it will be overwritten. - - If the file path is ``None``, then the result will be - saved in memory and returned as `bytes`. - - part_size_kb (`int`, optional): - Chunk size when downloading files. The larger, the less - requests will be made (up to 512KB maximum). - - file_size (`int`, optional): - The file size that is about to be downloaded, if known. - Only used if ``progress_callback`` is specified. - - progress_callback (`callable`, optional): - A callback function accepting two parameters: - ``(downloaded bytes, total)``. Note that the - ``total`` is the provided ``file_size``. - """ - if not part_size_kb: - if not file_size: - part_size_kb = 64 # Reasonable default - else: - part_size_kb = utils.get_appropriated_part_size(file_size) - - part_size = int(part_size_kb * 1024) - # https://core.telegram.org/api/files says: - # > part_size % 1024 = 0 (divisible by 1KB) - # - # But https://core.telegram.org/cdn (more recent) says: - # > limit must be divisible by 4096 bytes - # So we just stick to the 4096 limit. - if part_size % 4096 != 0: - raise ValueError( - 'The part size must be evenly divisible by 4096.') - - in_memory = file is None - if in_memory: - f = io.BytesIO() - elif isinstance(file, str): - # Ensure that we'll be able to download the media - helpers.ensure_parent_dir_exists(file) - f = open(file, 'wb') - else: - f = file - - # The used client will change if FileMigrateError occurs - client = self - cdn_decrypter = None - input_location = utils.get_input_location(input_location) - - __log__.info('Downloading file in chunks of %d bytes', part_size) - try: - offset = 0 - while True: - try: - if cdn_decrypter: - result = cdn_decrypter.get_file() - else: - result = client(GetFileRequest( - input_location, offset, part_size - )) - - if isinstance(result, FileCdnRedirect): - __log__.info('File lives in a CDN') - cdn_decrypter, result = \ - CdnDecrypter.prepare_decrypter( - client, self._get_cdn_client(result), - result - ) - - except FileMigrateError as e: - __log__.info('File lives in another DC') - client = self._get_exported_client(e.new_dc) - continue - - offset += part_size - - # If we have received no data (0 bytes), the file is over - # So there is nothing left to download and write - if not result.bytes: - # Return some extra information, unless it's a CDN file - if in_memory: - f.flush() - return f.getvalue() - else: - return getattr(result, 'type', '') - - f.write(result.bytes) - __log__.debug('Saved %d more bytes', len(result.bytes)) - if progress_callback: - progress_callback(f.tell(), file_size) - finally: - if client != self: - client.disconnect() - - if cdn_decrypter: - try: - cdn_decrypter.client.disconnect() - except: - pass - if isinstance(file, str) or in_memory: - f.close() - # endregion # endregion diff --git a/telethon/crypto/cdn_decrypter.py b/telethon/crypto/cdn_decrypter.py index 24a4bb49..dd615a5a 100644 --- a/telethon/crypto/cdn_decrypter.py +++ b/telethon/crypto/cdn_decrypter.py @@ -30,7 +30,7 @@ class CdnDecrypter: self.cdn_file_hashes = cdn_file_hashes @staticmethod - def prepare_decrypter(client, cdn_client, cdn_redirect): + async def prepare_decrypter(client, cdn_client, cdn_redirect): """ Prepares a new CDN decrypter. @@ -52,14 +52,14 @@ class CdnDecrypter: cdn_aes, cdn_redirect.cdn_file_hashes ) - cdn_file = cdn_client(GetCdnFileRequest( + cdn_file = await cdn_client(GetCdnFileRequest( file_token=cdn_redirect.file_token, offset=cdn_redirect.cdn_file_hashes[0].offset, limit=cdn_redirect.cdn_file_hashes[0].limit )) if isinstance(cdn_file, CdnFileReuploadNeeded): # We need to use the original client here - client(ReuploadCdnFileRequest( + await client(ReuploadCdnFileRequest( file_token=cdn_redirect.file_token, request_token=cdn_file.request_token ))