diff --git a/telethon/client/downloads.py b/telethon/client/downloads.py index 6f1f0fdb..23441a18 100644 --- a/telethon/client/downloads.py +++ b/telethon/client/downloads.py @@ -5,6 +5,8 @@ import pathlib import typing import inspect +from ..crypto import AES + from .. import utils, helpers, errors, hints from ..requestiter import RequestIter from ..tl import TLObject, types, functions @@ -17,7 +19,6 @@ except ImportError: if typing.TYPE_CHECKING: from .telegramclient import TelegramClient - # Chunk sizes for upload.getFile must be multiples of the smallest size MIN_CHUNK_SIZE = 4096 MAX_CHUNK_SIZE = 512 * 1024 @@ -376,7 +377,9 @@ class DownloadMethods: part_size_kb: float = None, file_size: int = None, progress_callback: 'hints.ProgressCallback' = None, - dc_id: int = None) -> typing.Optional[bytes]: + dc_id: int = None, + key: bytes = None, + iv: bytes = None) -> typing.Optional[bytes]: """ Low-level method to download files from their input location. @@ -415,6 +418,13 @@ class DownloadMethods: The data center the library should connect to in order to download the file. You shouldn't worry about this. + 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 + + Example .. code-block:: python @@ -446,6 +456,8 @@ class DownloadMethods: try: async for chunk in self.iter_download( input_location, request_size=part_size, dc_id=dc_id): + if iv and key: + chunk = AES.decrypt_ige(chunk, key, iv) r = f.write(chunk) if inspect.isawaitable(r): await r diff --git a/telethon/client/uploads.py b/telethon/client/uploads.py index 57262068..c5d2e3f4 100644 --- a/telethon/client/uploads.py +++ b/telethon/client/uploads.py @@ -8,6 +8,8 @@ import typing import inspect from io import BytesIO +from ..crypto import AES + from .. import utils, helpers, hints from ..tl import types, functions, custom @@ -439,6 +441,8 @@ class UploadMethods: part_size_kb: float = 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. @@ -476,6 +480,12 @@ class UploadMethods: 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)``. @@ -586,6 +596,10 @@ class UploadMethods: # Read the file by in chunks of size part_size part = stream.read(part_size) + # encryption part if needed + if key and iv: + part = AES.encrypt_ige(part, key, iv) + # The SavePartRequest is different depending on whether # the file is too large or not (over or less than 10MB) if is_large: