From 19b1bbb94297a54ece9bb310ea274796e9121e8a Mon Sep 17 00:00:00 2001 From: Eric Blundell Date: Tue, 20 Mar 2018 07:04:35 -0500 Subject: [PATCH 01/24] Allow download_media to download media to anywhere Remove the use of a temporary file in the programs working directory. --- pyrogram/client/client.py | 180 ++++++++++++++++++++------------------ 1 file changed, 97 insertions(+), 83 deletions(-) diff --git a/pyrogram/client/client.py b/pyrogram/client/client.py index 20ac58bc..a20ba521 100644 --- a/pyrogram/client/client.py +++ b/pyrogram/client/client.py @@ -55,6 +55,8 @@ from pyrogram.session.internals import MsgId from .input_media import InputMedia from .style import Markdown, HTML +from typing import Any + log = logging.getLogger(__name__) ApiKey = namedtuple("ApiKey", ["api_id", "api_hash"]) @@ -509,7 +511,6 @@ class Client: try: media, file_name, done, progress, path = media - tmp_file_name = None if isinstance(media, types.MessageMediaDocument): document = media.document @@ -535,13 +536,14 @@ class Client: elif isinstance(i, types.DocumentAttributeAnimated): file_name = file_name.replace("doc", "gif") - tmp_file_name = self.get_file( + self.get_file( dc_id=document.dc_id, id=document.id, access_hash=document.access_hash, version=document.version, size=document.size, - progress=progress + progress=progress, + file_out=file_name ) elif isinstance(media, (types.MessageMediaPhoto, types.Photo)): if isinstance(media, types.MessageMediaPhoto): @@ -558,37 +560,23 @@ class Client: photo_loc = photo.sizes[-1].location - tmp_file_name = self.get_file( + self.get_file( dc_id=photo_loc.dc_id, volume_id=photo_loc.volume_id, local_id=photo_loc.local_id, secret=photo_loc.secret, size=photo.sizes[-1].size, - progress=progress + progress=progress, + file_out=file_name ) if file_name is not None: - path[0] = "downloads/{}".format(file_name) - - try: - os.remove("downloads/{}".format(file_name)) - except OSError: - pass - finally: - try: - os.renames("{}".format(tmp_file_name), "downloads/{}".format(file_name)) - except OSError: - pass + path[0] = file_name except Exception as e: log.error(e, exc_info=True) finally: done.set() - try: - os.remove("{}".format(tmp_file_name)) - except OSError: - pass - log.debug("{} stopped".format(name)) def updates_worker(self): @@ -2177,7 +2165,9 @@ class Client: secret: int = None, version: int = 0, size: int = None, - progress: callable = None) -> str: + progress: callable = None, + file_out: Any = None) -> str: + if dc_id != self.dc_id: exported_auth = self.send( functions.auth.ExportAuthorization( @@ -2225,10 +2215,13 @@ class Client: version=version ) - file_name = "download_{}.temp".format(MsgId()) limit = 1024 * 1024 offset = 0 + # file object being written + f = None + close_file, call_flush, call_fsync = False, False, False + try: r = session.send( functions.upload.GetFile( @@ -2238,30 +2231,49 @@ class Client: ) ) + if file_out is None: + f = open("download_{}.temp".format(MsgId(), 'wb')) + close_file = True + + elif isinstance(file_out, str): + f = open(file_out, 'wb') + elif hasattr(file_out, 'write'): + f = file_out + + if hasattr(file_out, 'flush'): + call_flush = True + if hasattr(file_out, 'fileno'): + call_fsync = True + else: + raise ValueError('file_out argument of client.get_file must at least implement a write method if not a ' + 'string.') + if isinstance(r, types.upload.File): - with open(file_name, "wb") as f: - while True: - chunk = r.bytes + while True: + chunk = r.bytes - if not chunk: - break + if not chunk: + break - f.write(chunk) + f.write(chunk) + + if call_flush: f.flush() + if call_fsync: os.fsync(f.fileno()) - offset += limit + offset += limit - if progress: - progress(min(offset, size), size) + if progress: + progress(min(offset, size), size) - r = session.send( - functions.upload.GetFile( - location=location, - offset=offset, - limit=limit - ) + r = session.send( + functions.upload.GetFile( + location=location, + offset=offset, + limit=limit ) + ) if isinstance(r, types.upload.FileCdnRedirect): cdn_session = Session( @@ -2276,63 +2288,65 @@ class Client: cdn_session.start() try: - with open(file_name, "wb") as f: - while True: - r2 = cdn_session.send( - functions.upload.GetCdnFile( - location=location, - file_token=r.file_token, - offset=offset, - limit=limit - ) + while True: + r2 = cdn_session.send( + functions.upload.GetCdnFile( + location=location, + file_token=r.file_token, + offset=offset, + limit=limit ) + ) - if isinstance(r2, types.upload.CdnFileReuploadNeeded): - try: - session.send( - functions.upload.ReuploadCdnFile( - file_token=r.file_token, - request_token=r2.request_token - ) + if isinstance(r2, types.upload.CdnFileReuploadNeeded): + try: + session.send( + functions.upload.ReuploadCdnFile( + file_token=r.file_token, + request_token=r2.request_token ) - except VolumeLocNotFound: - break - else: - continue + ) + except VolumeLocNotFound: + break + else: + continue - chunk = r2.bytes + chunk = r2.bytes - # https://core.telegram.org/cdn#decrypting-files - decrypted_chunk = AES.ctr_decrypt( - chunk, - r.encryption_key, - r.encryption_iv, + # https://core.telegram.org/cdn#decrypting-files + decrypted_chunk = AES.ctr_decrypt( + chunk, + r.encryption_key, + r.encryption_iv, + offset + ) + + hashes = session.send( + functions.upload.GetCdnFileHashes( + r.file_token, offset ) + ) - hashes = session.send( - functions.upload.GetCdnFileHashes( - r.file_token, - offset - ) - ) + # https://core.telegram.org/cdn#verifying-files + for i, h in enumerate(hashes): + cdn_chunk = decrypted_chunk[h.limit * i: h.limit * (i + 1)] + assert h.hash == sha256(cdn_chunk).digest(), "Invalid CDN hash part {}".format(i) - # https://core.telegram.org/cdn#verifying-files - for i, h in enumerate(hashes): - cdn_chunk = decrypted_chunk[h.limit * i: h.limit * (i + 1)] - assert h.hash == sha256(cdn_chunk).digest(), "Invalid CDN hash part {}".format(i) + f.write(decrypted_chunk) - f.write(decrypted_chunk) + if call_flush: f.flush() + if call_fsync: os.fsync(f.fileno()) - offset += limit + offset += limit - if progress: - progress(min(offset, size), size) + if progress: + progress(min(offset, size), size) - if len(chunk) < limit: - break + if len(chunk) < limit: + break except Exception as e: log.error(e) finally: @@ -2340,8 +2354,10 @@ class Client: except Exception as e: log.error(e) else: - return file_name + return file_out finally: + if close_file and f and hasattr(f, 'close'): + f.close() session.stop() def join_chat(self, chat_id: str): @@ -2602,8 +2618,6 @@ class Client: progress: callable = None): """Use this method to download the media from a Message. - Files are saved in the *downloads* folder. - Args: message (:obj:`Message `): The Message containing the media. From 6bb004fc83d92dfc328bd1cdeb1e26bc1856f336 Mon Sep 17 00:00:00 2001 From: Eric Blundell Date: Tue, 20 Mar 2018 07:34:38 -0500 Subject: [PATCH 02/24] Add file_dir parameter to client.download_media --- pyrogram/client/client.py | 17 +++++++++++++++-- 1 file changed, 15 insertions(+), 2 deletions(-) diff --git a/pyrogram/client/client.py b/pyrogram/client/client.py index a20ba521..89b256c5 100644 --- a/pyrogram/client/client.py +++ b/pyrogram/client/client.py @@ -510,7 +510,7 @@ class Client: break try: - media, file_name, done, progress, path = media + media, file_dir, file_name, done, progress, path = media if isinstance(media, types.MessageMediaDocument): document = media.document @@ -536,6 +536,8 @@ class Client: elif isinstance(i, types.DocumentAttributeAnimated): file_name = file_name.replace("doc", "gif") + file_name = os.path.join(file_dir if file_dir is not None else '', file_name) + self.get_file( dc_id=document.dc_id, id=document.id, @@ -558,6 +560,8 @@ class Client: self.rnd_id() ) + file_name = os.path.join(file_dir if file_dir is not None else '', file_name) + photo_loc = photo.sizes[-1].location self.get_file( @@ -2614,6 +2618,7 @@ class Client: def download_media(self, message: types.Message, file_name: str = None, + file_dir: str = 'downloads', block: bool = True, progress: callable = None): """Use this method to download the media from a Message. @@ -2624,6 +2629,14 @@ class Client: file_name (:obj:`str`, optional): Specify a custom *file_name* to be used instead of the one provided by Telegram. + This parameter is expected to be a full file path to the location you want the + file to be placed. If not specified, the file will be put into the directory + specified by *file_dir* with a generated name. + + file_dir (:obj:`str`, optional): + Specify a directory to place the file in if no *file_name* is specified. + If *file_dir* is *None*, the current working directory is used. The default + value is the "downloads" folder in the current working directory. block (:obj:`bool`, optional): Blocks the code execution until the file has been downloaded. @@ -2656,7 +2669,7 @@ class Client: media = message if media is not None: - self.download_queue.put((media, file_name, done, progress, path)) + self.download_queue.put((media, file_dir, file_name, done, progress, path)) else: return From b9f623921dead51872a424d6d251cdb89b13e8ad Mon Sep 17 00:00:00 2001 From: Eric Blundell Date: Tue, 20 Mar 2018 07:47:38 -0500 Subject: [PATCH 03/24] Make file_name and file_dir mutually exclusive --- pyrogram/client/client.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/pyrogram/client/client.py b/pyrogram/client/client.py index 89b256c5..4862b9e8 100644 --- a/pyrogram/client/client.py +++ b/pyrogram/client/client.py @@ -2659,6 +2659,10 @@ class Client: Raises: :class:`pyrogram.Error` """ + + if file_name is not None and file_dir is not None: + ValueError('file_name and file_dir may not be specified together.') + if isinstance(message, (types.Message, types.Photo)): done = Event() path = [None] From 4ae9a5ad38ea240eee60d052cf3b9ddd366f63b3 Mon Sep 17 00:00:00 2001 From: Eric Blundell Date: Tue, 20 Mar 2018 08:05:41 -0500 Subject: [PATCH 04/24] Make sure file_dir is created --- pyrogram/client/client.py | 25 +++++++++++++++++++++---- 1 file changed, 21 insertions(+), 4 deletions(-) diff --git a/pyrogram/client/client.py b/pyrogram/client/client.py index 4862b9e8..2f2d1a23 100644 --- a/pyrogram/client/client.py +++ b/pyrogram/client/client.py @@ -512,6 +512,10 @@ class Client: try: media, file_dir, file_name, done, progress, path = media + if file_dir is not None: + # Make file_dir if it was specified + os.makedirs(file_dir, exist_ok=True) + if isinstance(media, types.MessageMediaDocument): document = media.document @@ -2620,7 +2624,8 @@ class Client: file_name: str = None, file_dir: str = 'downloads', block: bool = True, - progress: callable = None): + progress: callable = None + ): """Use this method to download the media from a Message. Args: @@ -2636,7 +2641,8 @@ class Client: file_dir (:obj:`str`, optional): Specify a directory to place the file in if no *file_name* is specified. If *file_dir* is *None*, the current working directory is used. The default - value is the "downloads" folder in the current working directory. + value is the "downloads" folder in the current working directory. The + directory tree will be created if it does not exist. block (:obj:`bool`, optional): Blocks the code execution until the file has been downloaded. @@ -2658,6 +2664,7 @@ class Client: Raises: :class:`pyrogram.Error` + :class:`ValueError` if both file_name and file_dir are specified. """ if file_name is not None and file_dir is not None: @@ -2685,6 +2692,7 @@ class Client: def download_photo(self, photo: types.Photo or types.UserProfilePhoto or types.ChatPhoto, file_name: str = None, + file_dir: str = None, block: bool = True): """Use this method to download a photo not contained inside a Message. For example, a photo of a User or a Chat/Channel. @@ -2696,7 +2704,16 @@ class Client: The photo object. file_name (:obj:`str`, optional): - Specify a custom *file_name* to be used. + Specify a custom *file_name* to be used instead of the one provided by Telegram. + This parameter is expected to be a full file path to the location you want the + photo to be placed. If not specified, the photo will be put into the directory + specified by *file_dir* with a generated name. + + file_dir (:obj:`str`, optional): + Specify a directory to place the photo in if no *file_name* is specified. + If *file_dir* is *None*, the current working directory is used. The default + value is the "downloads" folder in the current working directory. The + directory tree will be created if it does not exist. block (:obj:`bool`, optional): Blocks the code execution until the photo has been downloaded. @@ -2722,7 +2739,7 @@ class Client: )] ) - return self.download_media(photo, file_name, block) + return self.download_media(photo, file_name, file_dir, block) def add_contacts(self, contacts: list): """Use this method to add contacts to your Telegram address book. From 19854a5d4f3cae71bf4b32eebf7603882c3d7065 Mon Sep 17 00:00:00 2001 From: Eric Blundell Date: Tue, 20 Mar 2018 08:10:24 -0500 Subject: [PATCH 05/24] Actually raise mutually exclusive arg error --- pyrogram/client/client.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pyrogram/client/client.py b/pyrogram/client/client.py index 2f2d1a23..6c8ee1b5 100644 --- a/pyrogram/client/client.py +++ b/pyrogram/client/client.py @@ -2668,7 +2668,7 @@ class Client: """ if file_name is not None and file_dir is not None: - ValueError('file_name and file_dir may not be specified together.') + raise ValueError('file_name and file_dir may not be specified together.') if isinstance(message, (types.Message, types.Photo)): done = Event() From c0212a7b104c8cf03a67cd7e685f13401a50f7dd Mon Sep 17 00:00:00 2001 From: Eric Blundell Date: Tue, 20 Mar 2018 08:20:03 -0500 Subject: [PATCH 06/24] Correct default file_dir value behaviour --- pyrogram/client/client.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/pyrogram/client/client.py b/pyrogram/client/client.py index 6c8ee1b5..9955dd77 100644 --- a/pyrogram/client/client.py +++ b/pyrogram/client/client.py @@ -2622,7 +2622,7 @@ class Client: def download_media(self, message: types.Message, file_name: str = None, - file_dir: str = 'downloads', + file_dir: str = None, block: bool = True, progress: callable = None ): @@ -2670,6 +2670,9 @@ class Client: if file_name is not None and file_dir is not None: raise ValueError('file_name and file_dir may not be specified together.') + if file_name is None and file_dir is None: + file_dir = 'downloads' + if isinstance(message, (types.Message, types.Photo)): done = Event() path = [None] From db80c72b08506ba03c6c33c5a2e56db8a88f349e Mon Sep 17 00:00:00 2001 From: Eric Blundell Date: Tue, 20 Mar 2018 08:27:44 -0500 Subject: [PATCH 07/24] Create file_name directory trees in download_worker --- pyrogram/client/client.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/pyrogram/client/client.py b/pyrogram/client/client.py index 9955dd77..c2854796 100644 --- a/pyrogram/client/client.py +++ b/pyrogram/client/client.py @@ -516,6 +516,9 @@ class Client: # Make file_dir if it was specified os.makedirs(file_dir, exist_ok=True) + if file_name is not None: + os.makedirs(os.path.dirname(file_name), exist_ok=True) + if isinstance(media, types.MessageMediaDocument): document = media.document From 0694480a461337bb5764f900d7294487b0c73492 Mon Sep 17 00:00:00 2001 From: Eric Blundell Date: Tue, 20 Mar 2018 08:33:14 -0500 Subject: [PATCH 08/24] allow file objects be passed to file_name arg of client.download_media --- pyrogram/client/client.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pyrogram/client/client.py b/pyrogram/client/client.py index c2854796..a53225a0 100644 --- a/pyrogram/client/client.py +++ b/pyrogram/client/client.py @@ -516,7 +516,7 @@ class Client: # Make file_dir if it was specified os.makedirs(file_dir, exist_ok=True) - if file_name is not None: + if isinstance(file_name, str) and file_name is not None: os.makedirs(os.path.dirname(file_name), exist_ok=True) if isinstance(media, types.MessageMediaDocument): From bd1234f227a27e5ab506c15782a32d8b5055ec35 Mon Sep 17 00:00:00 2001 From: Eric Blundell Date: Tue, 20 Mar 2018 09:02:17 -0500 Subject: [PATCH 09/24] fix open file leak in client.download_media --- pyrogram/client/client.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/pyrogram/client/client.py b/pyrogram/client/client.py index a53225a0..57351887 100644 --- a/pyrogram/client/client.py +++ b/pyrogram/client/client.py @@ -2248,6 +2248,8 @@ class Client: elif isinstance(file_out, str): f = open(file_out, 'wb') + close_file = True + elif hasattr(file_out, 'write'): f = file_out @@ -2367,7 +2369,7 @@ class Client: else: return file_out finally: - if close_file and f and hasattr(f, 'close'): + if close_file and f is not None: f.close() session.stop() From 62831001b799039e149a53e826ce4aa21932d064 Mon Sep 17 00:00:00 2001 From: Eric Blundell Date: Tue, 20 Mar 2018 09:39:58 -0500 Subject: [PATCH 10/24] Slight amendment to client.download_(media/photo) doc --- pyrogram/client/client.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/pyrogram/client/client.py b/pyrogram/client/client.py index 57351887..765de03b 100644 --- a/pyrogram/client/client.py +++ b/pyrogram/client/client.py @@ -2640,8 +2640,8 @@ class Client: file_name (:obj:`str`, optional): Specify a custom *file_name* to be used instead of the one provided by Telegram. This parameter is expected to be a full file path to the location you want the - file to be placed. If not specified, the file will be put into the directory - specified by *file_dir* with a generated name. + file to be placed, or a file like object. If not specified, the file will + be put into the directory specified by *file_dir* with a generated name. file_dir (:obj:`str`, optional): Specify a directory to place the file in if no *file_name* is specified. @@ -2714,8 +2714,8 @@ class Client: file_name (:obj:`str`, optional): Specify a custom *file_name* to be used instead of the one provided by Telegram. This parameter is expected to be a full file path to the location you want the - photo to be placed. If not specified, the photo will be put into the directory - specified by *file_dir* with a generated name. + photo to be placed, or a file like object. If not specified, the photo will + be put into the directory specified by *file_dir* with a generated name. file_dir (:obj:`str`, optional): Specify a directory to place the photo in if no *file_name* is specified. From 5bc10b45a33fb61ff169ba126930d38de1de642f Mon Sep 17 00:00:00 2001 From: Eric Blundell Date: Tue, 20 Mar 2018 15:20:04 -0500 Subject: [PATCH 11/24] Use OS temp file, specific path download via path seperator inspection --- pyrogram/client/client.py | 262 +++++++++++++++++--------------------- 1 file changed, 117 insertions(+), 145 deletions(-) diff --git a/pyrogram/client/client.py b/pyrogram/client/client.py index 765de03b..9ced004f 100644 --- a/pyrogram/client/client.py +++ b/pyrogram/client/client.py @@ -34,6 +34,11 @@ from hashlib import sha256, md5 from queue import Queue from signal import signal, SIGINT, SIGTERM, SIGABRT from threading import Event, Thread +import tempfile + +import shutil + +import errno from pyrogram.api import functions, types from pyrogram.api.core import Object @@ -55,8 +60,6 @@ from pyrogram.session.internals import MsgId from .input_media import InputMedia from .style import Markdown, HTML -from typing import Any - log = logging.getLogger(__name__) ApiKey = namedtuple("ApiKey", ["api_id", "api_hash"]) @@ -510,14 +513,18 @@ class Client: break try: - media, file_dir, file_name, done, progress, path = media + media, file_name, done, progress, path = media + tmp_file_name = None - if file_dir is not None: - # Make file_dir if it was specified - os.makedirs(file_dir, exist_ok=True) + download_directory = "downloads" - if isinstance(file_name, str) and file_name is not None: - os.makedirs(os.path.dirname(file_name), exist_ok=True) + if file_name.endswith('/') or file_name.endswith('\\'): + # treat the file name as a directory + download_directory = file_name + file_name = None + elif '/' in file_name or '\\' in file_name: + # use file_name as a full path instead + download_directory = '' if isinstance(media, types.MessageMediaDocument): document = media.document @@ -543,16 +550,13 @@ class Client: elif isinstance(i, types.DocumentAttributeAnimated): file_name = file_name.replace("doc", "gif") - file_name = os.path.join(file_dir if file_dir is not None else '', file_name) - - self.get_file( + tmp_file_name = self.get_file( dc_id=document.dc_id, id=document.id, access_hash=document.access_hash, version=document.version, size=document.size, - progress=progress, - file_out=file_name + progress=progress ) elif isinstance(media, (types.MessageMediaPhoto, types.Photo)): if isinstance(media, types.MessageMediaPhoto): @@ -567,27 +571,46 @@ class Client: self.rnd_id() ) - file_name = os.path.join(file_dir if file_dir is not None else '', file_name) - photo_loc = photo.sizes[-1].location - self.get_file( + tmp_file_name = self.get_file( dc_id=photo_loc.dc_id, volume_id=photo_loc.volume_id, local_id=photo_loc.local_id, secret=photo_loc.secret, size=photo.sizes[-1].size, - progress=progress, - file_out=file_name + progress=progress ) if file_name is not None: - path[0] = file_name + path[0] = os.path.join(download_directory, file_name) + + try: + os.remove(os.path.join(download_directory, file_name)) + except OSError: + pass + finally: + try: + if download_directory: + os.makedirs(download_directory, exist_ok=True) + else: + os.makedirs(os.path.dirname(file_name), exist_ok=True) + + # avoid errors moving between drives on windows + shutil.move(tmp_file_name, os.path.join(download_directory, file_name)) + except OSError as e: + log.error(e, exc_info=True) except Exception as e: log.error(e, exc_info=True) finally: done.set() + try: + os.remove(tmp_file_name) + except OSError as e: + if not e.errno == errno.ENOENT: + log.error(e, exc_info=True) + log.debug("{} stopped".format(name)) def updates_worker(self): @@ -2176,9 +2199,7 @@ class Client: secret: int = None, version: int = 0, size: int = None, - progress: callable = None, - file_out: Any = None) -> str: - + progress: callable = None) -> str: if dc_id != self.dc_id: exported_auth = self.send( functions.auth.ExportAuthorization( @@ -2226,13 +2247,11 @@ class Client: version=version ) + fd, file_name = tempfile.mkstemp() + limit = 1024 * 1024 offset = 0 - # file object being written - f = None - close_file, call_flush, call_fsync = False, False, False - try: r = session.send( functions.upload.GetFile( @@ -2242,51 +2261,30 @@ class Client: ) ) - if file_out is None: - f = open("download_{}.temp".format(MsgId(), 'wb')) - close_file = True - - elif isinstance(file_out, str): - f = open(file_out, 'wb') - close_file = True - - elif hasattr(file_out, 'write'): - f = file_out - - if hasattr(file_out, 'flush'): - call_flush = True - if hasattr(file_out, 'fileno'): - call_fsync = True - else: - raise ValueError('file_out argument of client.get_file must at least implement a write method if not a ' - 'string.') - if isinstance(r, types.upload.File): - while True: - chunk = r.bytes + with os.fdopen(fd, "wb") as f: + while True: + chunk = r.bytes - if not chunk: - break + if not chunk: + break - f.write(chunk) - - if call_flush: + f.write(chunk) f.flush() - if call_fsync: os.fsync(f.fileno()) - offset += limit + offset += limit - if progress: - progress(min(offset, size), size) + if progress: + progress(min(offset, size), size) - r = session.send( - functions.upload.GetFile( - location=location, - offset=offset, - limit=limit + r = session.send( + functions.upload.GetFile( + location=location, + offset=offset, + limit=limit + ) ) - ) if isinstance(r, types.upload.FileCdnRedirect): cdn_session = Session( @@ -2301,76 +2299,77 @@ class Client: cdn_session.start() try: - while True: - r2 = cdn_session.send( - functions.upload.GetCdnFile( - location=location, - file_token=r.file_token, - offset=offset, - limit=limit - ) - ) - - if isinstance(r2, types.upload.CdnFileReuploadNeeded): - try: - session.send( - functions.upload.ReuploadCdnFile( - file_token=r.file_token, - request_token=r2.request_token - ) + with os.fdopen(fd, "wb") as f: + while True: + r2 = cdn_session.send( + functions.upload.GetCdnFile( + location=location, + file_token=r.file_token, + offset=offset, + limit=limit ) - except VolumeLocNotFound: - break - else: - continue + ) - chunk = r2.bytes + if isinstance(r2, types.upload.CdnFileReuploadNeeded): + try: + session.send( + functions.upload.ReuploadCdnFile( + file_token=r.file_token, + request_token=r2.request_token + ) + ) + except VolumeLocNotFound: + break + else: + continue - # https://core.telegram.org/cdn#decrypting-files - decrypted_chunk = AES.ctr_decrypt( - chunk, - r.encryption_key, - r.encryption_iv, - offset - ) + chunk = r2.bytes - hashes = session.send( - functions.upload.GetCdnFileHashes( - r.file_token, + # https://core.telegram.org/cdn#decrypting-files + decrypted_chunk = AES.ctr_decrypt( + chunk, + r.encryption_key, + r.encryption_iv, offset ) - ) - # https://core.telegram.org/cdn#verifying-files - for i, h in enumerate(hashes): - cdn_chunk = decrypted_chunk[h.limit * i: h.limit * (i + 1)] - assert h.hash == sha256(cdn_chunk).digest(), "Invalid CDN hash part {}".format(i) + hashes = session.send( + functions.upload.GetCdnFileHashes( + r.file_token, + offset + ) + ) - f.write(decrypted_chunk) + # https://core.telegram.org/cdn#verifying-files + for i, h in enumerate(hashes): + cdn_chunk = decrypted_chunk[h.limit * i: h.limit * (i + 1)] + assert h.hash == sha256(cdn_chunk).digest(), "Invalid CDN hash part {}".format(i) - if call_flush: + f.write(decrypted_chunk) f.flush() - if call_fsync: os.fsync(f.fileno()) - offset += limit + offset += limit - if progress: - progress(min(offset, size), size) + if progress: + progress(min(offset, size), size) - if len(chunk) < limit: - break + if len(chunk) < limit: + break except Exception as e: - log.error(e) + raise e finally: cdn_session.stop() except Exception as e: - log.error(e) + log.error(e, exc_info=True) + + try: + os.remove(file_name) + except OSError: + pass else: - return file_out + return file_name finally: - if close_file and f is not None: - f.close() session.stop() def join_chat(self, chat_id: str): @@ -2627,27 +2626,18 @@ class Client: def download_media(self, message: types.Message, file_name: str = None, - file_dir: str = None, block: bool = True, - progress: callable = None - ): + progress: callable = None): """Use this method to download the media from a Message. + Files are saved in the *downloads* folder. + Args: message (:obj:`Message `): The Message containing the media. file_name (:obj:`str`, optional): Specify a custom *file_name* to be used instead of the one provided by Telegram. - This parameter is expected to be a full file path to the location you want the - file to be placed, or a file like object. If not specified, the file will - be put into the directory specified by *file_dir* with a generated name. - - file_dir (:obj:`str`, optional): - Specify a directory to place the file in if no *file_name* is specified. - If *file_dir* is *None*, the current working directory is used. The default - value is the "downloads" folder in the current working directory. The - directory tree will be created if it does not exist. block (:obj:`bool`, optional): Blocks the code execution until the file has been downloaded. @@ -2669,15 +2659,7 @@ class Client: Raises: :class:`pyrogram.Error` - :class:`ValueError` if both file_name and file_dir are specified. """ - - if file_name is not None and file_dir is not None: - raise ValueError('file_name and file_dir may not be specified together.') - - if file_name is None and file_dir is None: - file_dir = 'downloads' - if isinstance(message, (types.Message, types.Photo)): done = Event() path = [None] @@ -2688,7 +2670,7 @@ class Client: media = message if media is not None: - self.download_queue.put((media, file_dir, file_name, done, progress, path)) + self.download_queue.put((media, file_name, done, progress, path)) else: return @@ -2700,7 +2682,6 @@ class Client: def download_photo(self, photo: types.Photo or types.UserProfilePhoto or types.ChatPhoto, file_name: str = None, - file_dir: str = None, block: bool = True): """Use this method to download a photo not contained inside a Message. For example, a photo of a User or a Chat/Channel. @@ -2712,16 +2693,7 @@ class Client: The photo object. file_name (:obj:`str`, optional): - Specify a custom *file_name* to be used instead of the one provided by Telegram. - This parameter is expected to be a full file path to the location you want the - photo to be placed, or a file like object. If not specified, the photo will - be put into the directory specified by *file_dir* with a generated name. - - file_dir (:obj:`str`, optional): - Specify a directory to place the photo in if no *file_name* is specified. - If *file_dir* is *None*, the current working directory is used. The default - value is the "downloads" folder in the current working directory. The - directory tree will be created if it does not exist. + Specify a custom *file_name* to be used. block (:obj:`bool`, optional): Blocks the code execution until the photo has been downloaded. @@ -2747,7 +2719,7 @@ class Client: )] ) - return self.download_media(photo, file_name, file_dir, block) + return self.download_media(photo, file_name, block) def add_contacts(self, contacts: list): """Use this method to add contacts to your Telegram address book. From cd0e585d0d3c5555a2496e402b044b7e4bedbbe1 Mon Sep 17 00:00:00 2001 From: Eric Blundell Date: Tue, 20 Mar 2018 15:42:31 -0500 Subject: [PATCH 12/24] Avoid calling fdopen on closed descriptor --- pyrogram/client/client.py | 11 +++++++++++ 1 file changed, 11 insertions(+) diff --git a/pyrogram/client/client.py b/pyrogram/client/client.py index 9ced004f..b29792d6 100644 --- a/pyrogram/client/client.py +++ b/pyrogram/client/client.py @@ -2299,6 +2299,17 @@ class Client: cdn_session.start() try: + # cant fdopen the closed file descriptor from above + # which is closed due to the with statement in the branch just above + # make a new temp file to write to + + try: + os.remove(file_name) + except OSError: + pass + + fd, file_name = tempfile.mkstemp() + with os.fdopen(fd, "wb") as f: while True: r2 = cdn_session.send( From 4c9e4df53291876ece62610401c7a94dbdd343b9 Mon Sep 17 00:00:00 2001 From: Eric Blundell Date: Tue, 20 Mar 2018 16:18:32 -0500 Subject: [PATCH 13/24] Amendment to comment on fdopen usage in get_file --- pyrogram/client/client.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/pyrogram/client/client.py b/pyrogram/client/client.py index b29792d6..b63d0b4f 100644 --- a/pyrogram/client/client.py +++ b/pyrogram/client/client.py @@ -2299,8 +2299,8 @@ class Client: cdn_session.start() try: - # cant fdopen the closed file descriptor from above - # which is closed due to the with statement in the branch just above + # cant fdopen the closed file descriptor which could be closed due + # to the with statement in the branch just above. # make a new temp file to write to try: From f0c00c88013ba88a5c9e98b43f1446945cbbbe3d Mon Sep 17 00:00:00 2001 From: Eric Blundell Date: Tue, 20 Mar 2018 16:30:48 -0500 Subject: [PATCH 14/24] move first mkstemp to exception safe location in get_file --- pyrogram/client/client.py | 14 ++++++++------ 1 file changed, 8 insertions(+), 6 deletions(-) diff --git a/pyrogram/client/client.py b/pyrogram/client/client.py index b63d0b4f..529873ff 100644 --- a/pyrogram/client/client.py +++ b/pyrogram/client/client.py @@ -2247,10 +2247,9 @@ class Client: version=version ) - fd, file_name = tempfile.mkstemp() - limit = 1024 * 1024 offset = 0 + file_name = None try: r = session.send( @@ -2261,6 +2260,8 @@ class Client: ) ) + fd, file_name = tempfile.mkstemp() + if isinstance(r, types.upload.File): with os.fdopen(fd, "wb") as f: while True: @@ -2374,10 +2375,11 @@ class Client: except Exception as e: log.error(e, exc_info=True) - try: - os.remove(file_name) - except OSError: - pass + if file_name: + try: + os.remove(file_name) + except OSError: + pass else: return file_name finally: From 8796e857af44918a05cf37fb3b5f42cbca6fb7c4 Mon Sep 17 00:00:00 2001 From: Eric Blundell Date: Tue, 20 Mar 2018 23:20:08 -0500 Subject: [PATCH 15/24] Amend comment on shutil.move in download_worker os.renames cannot move across drives/partitions on any platform. that is why shutil.move is used, because the OS allotted temp file could possibly be on another drive or partition. Also fix code formatting on new import statements. --- pyrogram/client/client.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/pyrogram/client/client.py b/pyrogram/client/client.py index 211a23d5..26d5e9df 100644 --- a/pyrogram/client/client.py +++ b/pyrogram/client/client.py @@ -35,9 +35,7 @@ from queue import Queue from signal import signal, SIGINT, SIGTERM, SIGABRT from threading import Event, Thread import tempfile - import shutil - import errno from pyrogram.api import functions, types @@ -599,7 +597,7 @@ class Client: else: os.makedirs(os.path.dirname(file_name), exist_ok=True) - # avoid errors moving between drives on windows + # avoid errors moving between drives/partitions etc. shutil.move(tmp_file_name, os.path.join(download_directory, file_name)) except OSError as e: log.error(e, exc_info=True) From b6a42aa8cd662d7839ada632e2e70ac5d63e5fee Mon Sep 17 00:00:00 2001 From: Dan <14043624+delivrance@users.noreply.github.com> Date: Wed, 21 Mar 2018 09:01:18 +0100 Subject: [PATCH 16/24] Do not mkstemp twice Also use elif to make it less confusing --- pyrogram/client/client.py | 26 ++++++-------------------- 1 file changed, 6 insertions(+), 20 deletions(-) diff --git a/pyrogram/client/client.py b/pyrogram/client/client.py index 26d5e9df..0776c8d3 100644 --- a/pyrogram/client/client.py +++ b/pyrogram/client/client.py @@ -2250,7 +2250,7 @@ class Client: limit = 1024 * 1024 offset = 0 - file_name = None + fd, file_name = tempfile.mkstemp() try: r = session.send( @@ -2261,8 +2261,6 @@ class Client: ) ) - fd, file_name = tempfile.mkstemp() - if isinstance(r, types.upload.File): with os.fdopen(fd, "wb") as f: while True: @@ -2288,7 +2286,7 @@ class Client: ) ) - if isinstance(r, types.upload.FileCdnRedirect): + elif isinstance(r, types.upload.FileCdnRedirect): cdn_session = Session( r.dc_id, self.test_mode, @@ -2301,17 +2299,6 @@ class Client: cdn_session.start() try: - # cant fdopen the closed file descriptor which could be closed due - # to the with statement in the branch just above. - # make a new temp file to write to - - try: - os.remove(file_name) - except OSError: - pass - - fd, file_name = tempfile.mkstemp() - with os.fdopen(fd, "wb") as f: while True: r2 = cdn_session.send( @@ -2376,11 +2363,10 @@ class Client: except Exception as e: log.error(e, exc_info=True) - if file_name: - try: - os.remove(file_name) - except OSError: - pass + try: + os.remove(file_name) + except OSError: + pass else: return file_name finally: From b45960212b86a4277d155e74e9c8fbb7c2e54503 Mon Sep 17 00:00:00 2001 From: Eric Blundell Date: Wed, 21 Mar 2018 03:19:09 -0500 Subject: [PATCH 17/24] Simplify branch in download_worker exception handler --- pyrogram/client/client.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pyrogram/client/client.py b/pyrogram/client/client.py index 0776c8d3..cf4deb85 100644 --- a/pyrogram/client/client.py +++ b/pyrogram/client/client.py @@ -609,7 +609,7 @@ class Client: try: os.remove(tmp_file_name) except OSError as e: - if not e.errno == errno.ENOENT: + if e.errno != errno.ENOENT: log.error(e, exc_info=True) log.debug("{} stopped".format(name)) From 0f4e29584ac46af38e8605f9720e3b0e8189693d Mon Sep 17 00:00:00 2001 From: Eric Blundell Date: Wed, 21 Mar 2018 04:07:55 -0500 Subject: [PATCH 18/24] Make use of tempfile.NamedTemporaryFile in getfile, use context managers --- pyrogram/client/client.py | 9 ++++++--- 1 file changed, 6 insertions(+), 3 deletions(-) diff --git a/pyrogram/client/client.py b/pyrogram/client/client.py index cf4deb85..50009ba4 100644 --- a/pyrogram/client/client.py +++ b/pyrogram/client/client.py @@ -2250,7 +2250,7 @@ class Client: limit = 1024 * 1024 offset = 0 - fd, file_name = tempfile.mkstemp() + file_name = None try: r = session.send( @@ -2262,7 +2262,9 @@ class Client: ) if isinstance(r, types.upload.File): - with os.fdopen(fd, "wb") as f: + with tempfile.NamedTemporaryFile('wb', delete=False) as f: + file_name = f.name + while True: chunk = r.bytes @@ -2299,7 +2301,8 @@ class Client: cdn_session.start() try: - with os.fdopen(fd, "wb") as f: + with tempfile.NamedTemporaryFile('wb', delete=False) as f: + file_name = f.name while True: r2 = cdn_session.send( functions.upload.GetCdnFile( From f6ea3e9b424538c1d7995852877a056400a7d790 Mon Sep 17 00:00:00 2001 From: Dan <14043624+delivrance@users.noreply.github.com> Date: Wed, 21 Mar 2018 13:39:23 +0100 Subject: [PATCH 19/24] Cleaner code and some little changes TODO: "" or None for faulty download, which is better? --- pyrogram/client/client.py | 71 ++++++++++++++++----------------------- 1 file changed, 29 insertions(+), 42 deletions(-) diff --git a/pyrogram/client/client.py b/pyrogram/client/client.py index 50009ba4..ea77ff75 100644 --- a/pyrogram/client/client.py +++ b/pyrogram/client/client.py @@ -25,6 +25,7 @@ import mimetypes import os import re import struct +import tempfile import threading import time from collections import namedtuple @@ -34,9 +35,6 @@ from hashlib import sha256, md5 from queue import Queue from signal import signal, SIGINT, SIGTERM, SIGABRT from threading import Event, Thread -import tempfile -import shutil -import errno from pyrogram.api import functions, types from pyrogram.api.core import Object @@ -506,23 +504,17 @@ class Client: while True: media = self.download_queue.get() + temp_file_path = "" + final_file_path = "" if media is None: break try: media, file_name, done, progress, path = media - tmp_file_name = None - download_directory = "downloads" - - if file_name.endswith('/') or file_name.endswith('\\'): - # treat the file name as a directory - download_directory = file_name - file_name = None - elif '/' in file_name or '\\' in file_name: - # use file_name as a full path instead - download_directory = '' + directory, file_name = os.path.split(file_name) + directory = directory or "downloads" if isinstance(media, types.MessageMediaDocument): document = media.document @@ -548,7 +540,7 @@ class Client: elif isinstance(i, types.DocumentAttributeAnimated): file_name = file_name.replace("doc", "gif") - tmp_file_name = self.get_file( + temp_file_path = self.get_file( dc_id=document.dc_id, id=document.id, access_hash=document.access_hash, @@ -571,7 +563,7 @@ class Client: photo_loc = photo.sizes[-1].location - tmp_file_name = self.get_file( + temp_file_path = self.get_file( dc_id=photo_loc.dc_id, volume_id=photo_loc.volume_id, local_id=photo_loc.local_id, @@ -580,37 +572,29 @@ class Client: progress=progress ) - if tmp_file_name is None: - return None + if temp_file_path: + final_file_path = os.path.join(directory, file_name) - if file_name is not None: - path[0] = os.path.join(download_directory, file_name) - - try: - os.remove(os.path.join(download_directory, file_name)) - except OSError: - pass - finally: try: - if download_directory: - os.makedirs(download_directory, exist_ok=True) - else: - os.makedirs(os.path.dirname(file_name), exist_ok=True) + os.remove(final_file_path) + except OSError: + pass - # avoid errors moving between drives/partitions etc. - shutil.move(tmp_file_name, os.path.join(download_directory, file_name)) - except OSError as e: - log.error(e, exc_info=True) + os.renames(temp_file_path, final_file_path) except Exception as e: log.error(e, exc_info=True) - finally: - done.set() try: - os.remove(tmp_file_name) - except OSError as e: - if e.errno != errno.ENOENT: - log.error(e, exc_info=True) + os.remove(temp_file_path) + except OSError: + pass + else: + # TODO: "" or None for faulty download, which is better? + # os.path methods return "" in case something does not exist, I prefer this. + # For now let's keep None + path[0] = final_file_path or None + finally: + done.set() log.debug("{} stopped".format(name)) @@ -2250,7 +2234,7 @@ class Client: limit = 1024 * 1024 offset = 0 - file_name = None + file_name = "" try: r = session.send( @@ -2303,6 +2287,7 @@ class Client: try: with tempfile.NamedTemporaryFile('wb', delete=False) as f: file_name = f.name + while True: r2 = cdn_session.send( functions.upload.GetCdnFile( @@ -2370,6 +2355,8 @@ class Client: os.remove(file_name) except OSError: pass + + return "" else: return file_name finally: @@ -2628,7 +2615,7 @@ class Client: def download_media(self, message: types.Message, - file_name: str = None, + file_name: str = "", block: bool = True, progress: callable = None): """Use this method to download the media from a Message. @@ -2684,7 +2671,7 @@ class Client: def download_photo(self, photo: types.Photo or types.UserProfilePhoto or types.ChatPhoto, - file_name: str = None, + file_name: str = "", block: bool = True): """Use this method to download a photo not contained inside a Message. For example, a photo of a User or a Chat/Channel. From 76ad29ae11099bce3137633cfba9d1a141838347 Mon Sep 17 00:00:00 2001 From: Dan <14043624+delivrance@users.noreply.github.com> Date: Wed, 21 Mar 2018 15:42:32 +0100 Subject: [PATCH 20/24] Fix saving files on another drive (windows) @EriHoss --- pyrogram/client/client.py | 9 ++------- 1 file changed, 2 insertions(+), 7 deletions(-) diff --git a/pyrogram/client/client.py b/pyrogram/client/client.py index ea77ff75..dd0d64fe 100644 --- a/pyrogram/client/client.py +++ b/pyrogram/client/client.py @@ -574,13 +574,8 @@ class Client: if temp_file_path: final_file_path = os.path.join(directory, file_name) - - try: - os.remove(final_file_path) - except OSError: - pass - - os.renames(temp_file_path, final_file_path) + os.makedirs(directory, exist_ok=True) + shutil.move(temp_file_path, final_file_path) except Exception as e: log.error(e, exc_info=True) From 40e7d72e873239b7f11d88ce0b8345010627a01e Mon Sep 17 00:00:00 2001 From: Dan <14043624+delivrance@users.noreply.github.com> Date: Wed, 21 Mar 2018 15:43:58 +0100 Subject: [PATCH 21/24] Make paths good looking --- pyrogram/client/client.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pyrogram/client/client.py b/pyrogram/client/client.py index dd0d64fe..746d1929 100644 --- a/pyrogram/client/client.py +++ b/pyrogram/client/client.py @@ -573,7 +573,7 @@ class Client: ) if temp_file_path: - final_file_path = os.path.join(directory, file_name) + final_file_path = re.sub("\\\\", "/", os.path.join(directory, file_name)) os.makedirs(directory, exist_ok=True) shutil.move(temp_file_path, final_file_path) except Exception as e: From fa6af8695e9c43ca41a982ab0eff47680c294624 Mon Sep 17 00:00:00 2001 From: Eric Blundell Date: Wed, 21 Mar 2018 10:13:45 -0500 Subject: [PATCH 22/24] Fix missing shutil import --- pyrogram/client/client.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/pyrogram/client/client.py b/pyrogram/client/client.py index 746d1929..f024db96 100644 --- a/pyrogram/client/client.py +++ b/pyrogram/client/client.py @@ -28,6 +28,8 @@ import struct import tempfile import threading import time +import shutil + from collections import namedtuple from configparser import ConfigParser from datetime import datetime From 569ab1696ac4245c96521ae6d48c3225c45c1f79 Mon Sep 17 00:00:00 2001 From: Dan <14043624+delivrance@users.noreply.github.com> Date: Wed, 21 Mar 2018 16:17:13 +0100 Subject: [PATCH 23/24] Return the good looking absolute path instead of an ugly relative one #37 --- pyrogram/client/client.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pyrogram/client/client.py b/pyrogram/client/client.py index f024db96..869ca39d 100644 --- a/pyrogram/client/client.py +++ b/pyrogram/client/client.py @@ -575,7 +575,7 @@ class Client: ) if temp_file_path: - final_file_path = re.sub("\\\\", "/", os.path.join(directory, file_name)) + final_file_path = os.path.abspath(re.sub("\\\\", "/", os.path.join(directory, file_name))) os.makedirs(directory, exist_ok=True) shutil.move(temp_file_path, final_file_path) except Exception as e: From e4642266084c40ffbb1f58cc7c57d717d2f832cc Mon Sep 17 00:00:00 2001 From: Dan <14043624+delivrance@users.noreply.github.com> Date: Wed, 21 Mar 2018 17:39:53 +0100 Subject: [PATCH 24/24] Update docs --- pyrogram/client/client.py | 21 +++++++++++---------- 1 file changed, 11 insertions(+), 10 deletions(-) diff --git a/pyrogram/client/client.py b/pyrogram/client/client.py index 869ca39d..8277debb 100644 --- a/pyrogram/client/client.py +++ b/pyrogram/client/client.py @@ -24,12 +24,11 @@ import math import mimetypes import os import re +import shutil import struct import tempfile import threading import time -import shutil - from collections import namedtuple from configparser import ConfigParser from datetime import datetime @@ -2617,14 +2616,15 @@ class Client: progress: callable = None): """Use this method to download the media from a Message. - Files are saved in the *downloads* folder. - Args: message (:obj:`Message `): The Message containing the media. file_name (:obj:`str`, optional): - Specify a custom *file_name* to be used instead of the one provided by Telegram. + A custom *file_name* to be used instead of the one provided by Telegram. + By default, all files are downloaded in the *downloads* folder in your working directory. + You can also specify a path for downloading files in a custom location: paths that end with "/" + are considered directories. All non-existent folders will be created automatically. block (:obj:`bool`, optional): Blocks the code execution until the file has been downloaded. @@ -2642,7 +2642,7 @@ class Client: The size of the file. Returns: - The relative path of the downloaded file. + On success, the absolute path of the downloaded file as string is returned, None otherwise. Raises: :class:`pyrogram.Error` @@ -2673,21 +2673,22 @@ class Client: """Use this method to download a photo not contained inside a Message. For example, a photo of a User or a Chat/Channel. - Photos are saved in the *downloads* folder. - Args: photo (:obj:`Photo ` | :obj:`UserProfilePhoto ` | :obj:`ChatPhoto `): The photo object. file_name (:obj:`str`, optional): - Specify a custom *file_name* to be used. + A custom *file_name* to be used instead of the one provided by Telegram. + By default, all photos are downloaded in the *downloads* folder in your working directory. + You can also specify a path for downloading photos in a custom location: paths that end with "/" + are considered directories. All non-existent folders will be created automatically. block (:obj:`bool`, optional): Blocks the code execution until the photo has been downloaded. Defaults to True. Returns: - The relative path of the downloaded photo. + On success, the absolute path of the downloaded photo as string is returned, None otherwise. Raises: :class:`pyrogram.Error`