LonamiWebs/Telethon
1
1
Fork 0
mirror of https://github.com/LonamiWebs/Telethon.git synced 2024-11-26 03:13:45 +03:00
Code Issues Packages Projects Releases Wiki Activity

Tidy up TelegramBareClient.download_file CDN's mess

Browse Source
This commit is contained in:
Lonami Exo 2017-08-24 18:00:27 +02:00
parent d5323c5506
commit 4384ccb132
2 changed files with 51 additions and 38 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

85
telethon/telegram_bare_client.py
View File

@ -458,7 +458,6 @@ class TelegramBareClient:
takes two parameters, (bytes_downloaded, total_bytes). Note that
'total_bytes' simply equals 'file_size', and may be None.
"""
# TODO Clean up the CDN mess below
if not part_size_kb:
if not file_size:
part_size_kb = 64 # Reasonable default
@ -487,49 +486,33 @@ class TelegramBareClient:
try:
offset_index = 0
cdn_redirect = None
cdn_aes = None
cdn_file_token = None
def encrypt_method(x):
return x # Defaults to no-op
while True:
offset = offset_index * part_size
try:
if cdn_redirect:
if cdn_file_token:
result = client(GetCdnFileRequest(
cdn_redirect.file_token, offset, part_size
cdn_file_token, offset, part_size
))
else:
result = client(GetFileRequest(
input_location, offset, part_size
))
if isinstance(result, FileCdnRedirect):
# https://core.telegram.org/cdn
cdn_redirect = result
cdn_aes = pyaes.AESModeOfOperationCTR(
result.encryption_key
)
# The returned IV is the counter used on CTR
cdn_aes._counter._counter = list(
result.encryption_iv[:12] +
(offset >> 4).to_bytes(4, 'big')
)
if isinstance(result, FileCdnRedirect):
client, cdn_file_token, encrypt_method, result = \
self._prepare_cdn_redirect(
result, offset, part_size
)
client, cdn_file = self._get_cdn_client(
result.dc_id,
GetCdnFileRequest(
cdn_redirect.file_token, offset, part_size
)
)
if isinstance(cdn_file, CdnFileReuploadNeeded):
# We need to use the original client here
self(ReuploadCdnFileRequest(
file_token=cdn_redirect.file_token,
request_token=cdn_file.request_token
))
# TODO else: we have the first file bytes already,
# avoid a redundant call
continue
if result is None:
# File was not ready on the CDN yet
continue
except FileMigrateError as e:
client = self._get_exported_client(e.new_dc)
@ -543,16 +526,42 @@ class TelegramBareClient:
# Return some extra information, unless it's a cdn file
return getattr(result, 'type', '')
if cdn_redirect:
# We first need to decrypt the result
# TODO Use libssl if available
result.bytes = cdn_aes.encrypt(result.bytes)
f.write(result.bytes)
f.write(encrypt_method(result.bytes))
if progress_callback:
progress_callback(f.tell(), file_size)
finally:
if isinstance(file, str):
f.close()
def _prepare_cdn_redirect(self, cdn_redirect, offset, part_size):
"""Returns (client, cdn_file_token, encrypt_method, result)"""
# https://core.telegram.org/cdn
# TODO Use libssl if available
cdn_aes = pyaes.AESModeOfOperationCTR(
cdn_redirect.encryption_key
)
# The returned IV is the counter used on CTR
cdn_aes._counter._counter = list(
cdn_redirect.encryption_iv[:12] +
(offset >> 4).to_bytes(4, 'big')
)
client, cdn_file = self._get_cdn_client(
cdn_redirect.dc_id,
GetCdnFileRequest(
cdn_redirect.file_token, offset, part_size
)
)
if isinstance(cdn_file, CdnFileReuploadNeeded):
# We need to use the original client here
self(ReuploadCdnFileRequest(
file_token=cdn_redirect.file_token,
request_token=cdn_file.request_token
))
return client, cdn_redirect.file_token, cdn_aes.encrypt, None
else:
# We have the first bytes for the file
return client, cdn_redirect.file_token, cdn_aes.encrypt, cdn_file
# endregion

4
telethon_tests/higher_level_test.py
View File

@ -39,6 +39,10 @@ class HigherLevelTests(unittest.TestCase):
client.download_media(msg, out)
assert sha256(data).digest() == sha256(out.getvalue()).digest()
out = BytesIO()
client.download_media(msg, out) # Won't redirect
assert sha256(data).digest() == sha256(out.getvalue()).digest()
client.log_out()
finally:
client.disconnect()