Expose key and iv parameter in downloads/uploads (#1397)

This commit is contained in:
painor 2020-03-04 16:12:34 +01:00 committed by GitHub
parent 1ec38aa5b2
commit 0e0052888f
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 28 additions and 2 deletions

View File

@ -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

View File

@ -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: