From 8afcd0b91f3f96b90407cae380ef846995249b9e Mon Sep 17 00:00:00 2001 From: Lonami Exo Date: Mon, 28 Aug 2017 16:25:10 +0200 Subject: [PATCH] Turn HashChecker into CdnDecrypter to abstract CDN-specific aspects --- telethon/crypto/__init__.py | 2 +- telethon/crypto/cdn_decrypter.py | 116 +++++++++++++++++++++++++++++++ telethon/crypto/hash_checker.py | 39 ----------- telethon/telegram_bare_client.py | 84 ++++------------------ 4 files changed, 129 insertions(+), 112 deletions(-) create mode 100644 telethon/crypto/cdn_decrypter.py delete mode 100644 telethon/crypto/hash_checker.py diff --git a/telethon/crypto/__init__.py b/telethon/crypto/__init__.py index 6dbcdc36..9b4fdb73 100644 --- a/telethon/crypto/__init__.py +++ b/telethon/crypto/__init__.py @@ -1,4 +1,4 @@ from .aes import AES from .auth_key import AuthKey from .factorization import Factorization -from .hash_checker import HashChecker +from .cdn_decrypter import CdnDecrypter diff --git a/telethon/crypto/cdn_decrypter.py b/telethon/crypto/cdn_decrypter.py new file mode 100644 index 00000000..e36fc135 --- /dev/null +++ b/telethon/crypto/cdn_decrypter.py @@ -0,0 +1,116 @@ +from hashlib import sha256 +import pyaes + +from ..tl import JsonSession +from ..tl.functions.upload import GetCdnFileRequest, ReuploadCdnFileRequest +from ..tl.types.upload import CdnFileReuploadNeeded +from ..errors import CdnFileTamperedError + + +class CdnDecrypter: + """Used when downloading a file results in a 'FileCdnRedirect' to + both prepare the redirect, decrypt the file as it downloads, and + ensure the file hasn't been tampered. + """ + def __init__(self, cdn_client, file_token, cdn_aes, cdn_file_hashes): + self.client = cdn_client + self.file_token = file_token + self.cdn_aes = cdn_aes + self.cdn_file_hashes = cdn_file_hashes + self.shaes = [sha256() for _ in range(len(cdn_file_hashes))] + + @staticmethod + def prepare_decrypter(client, client_cls, cdn_redirect, offset, part_size): + """Prepares a CDN decrypter, returning (decrypter, file data). + 'client' should be the original TelegramBareClient that + tried to download the file. + + 'client_cls' should be the class of the TelegramBareClient. + """ + # TODO Avoid the need for 'client_cls=TelegramBareClient' + # 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') + ) + + # Create a new client on said CDN + dc = client._get_dc(cdn_redirect.dc_id, cdn=True) + session = JsonSession(client.session) + session.server_address = dc.ip_address + session.port = dc.port + cdn_client = client_cls( # Avoid importing TelegramBareClient + session, client.api_id, client.api_hash, + timeout=client._timeout + ) + # This will make use of the new RSA keys for this specific CDN + cdn_file = cdn_client.connect(initial_query=GetCdnFileRequest( + cdn_redirect.file_token, offset, part_size + )) + + # CDN client is ready, create the resulting CdnDecrypter + decrypter = CdnDecrypter( + cdn_client, cdn_redirect.file_token, + cdn_aes, cdn_redirect.cdn_file_hashes + ) + + if isinstance(cdn_file, CdnFileReuploadNeeded): + # We need to use the original client here + client(ReuploadCdnFileRequest( + file_token=cdn_redirect.file_token, + request_token=cdn_file.request_token + )) + + # We want to always return a valid upload.CdnFile + cdn_file = decrypter.get_file(offset, part_size) + else: + cdn_file.bytes = decrypter.cdn_aes.encrypt(cdn_file.bytes) + decrypter.check(offset, cdn_file.bytes) + + return decrypter, cdn_file + + def get_file(self, offset, limit): + """Calls GetCdnFileRequest and decrypts its bytes. + Also ensures that the file hasn't been tampered. + """ + result = self.client(GetCdnFileRequest(self.file_token, offset, limit)) + result.bytes = self.cdn_aes.encrypt(result.bytes) + self.check(offset, result.bytes) + return result + + def check(self, offset, data): + """Checks the integrity of the given data""" + for cdn_hash, sha in zip(self.cdn_file_hashes, self.shaes): + inter = self.intersect( + cdn_hash.offset, cdn_hash.offset + cdn_hash.limit, + offset, offset + len(data) + ) + if inter: + x1, x2 = inter[0] - offset, inter[1] - offset + sha.update(data[x1:x2]) + elif offset > cdn_hash.offset: + if cdn_hash.hash == sha.digest(): + self.cdn_file_hashes.remove(cdn_hash) + self.shaes.remove(sha) + else: + raise CdnFileTamperedError() + + def finish_check(self): + """Similar to the check method, but for all unchecked hashes""" + for cdn_hash, sha in zip(self.cdn_file_hashes, self.shaes): + if cdn_hash.hash != sha.digest(): + raise CdnFileTamperedError() + + self.cdn_file_hashes.clear() + self.shaes.clear() + + @staticmethod + def intersect(x1, x2, z1, z2): + if x1 > z1: + return None if x1 > z2 else (x1, min(x2, z2)) + else: + return (z1, min(x2, z2)) if x2 > z1 else None diff --git a/telethon/crypto/hash_checker.py b/telethon/crypto/hash_checker.py deleted file mode 100644 index 5cdcf73f..00000000 --- a/telethon/crypto/hash_checker.py +++ /dev/null @@ -1,39 +0,0 @@ -from hashlib import sha256 -from ..errors import CdnFileTamperedError - - -class HashChecker: - def __init__(self, cdn_file_hashes): - self.cdn_file_hashes = cdn_file_hashes - self.shaes = [sha256() for _ in range(len(cdn_file_hashes))] - - def check(self, offset, data): - for cdn_hash, sha in zip(self.cdn_file_hashes, self.shaes): - inter = self.intersect( - cdn_hash.offset, cdn_hash.offset + cdn_hash.limit, - offset, offset + len(data) - ) - if inter: - x1, x2 = inter[0] - offset, inter[1] - offset - sha.update(data[x1:x2]) - elif offset > cdn_hash.offset: - if cdn_hash.hash == sha.digest(): - self.cdn_file_hashes.remove(cdn_hash) - self.shaes.remove(sha) - else: - raise CdnFileTamperedError() - - def finish_check(self): - for cdn_hash, sha in zip(self.cdn_file_hashes, self.shaes): - if cdn_hash.hash != sha.digest(): - raise CdnFileTamperedError() - - self.cdn_file_hashes.clear() - self.shaes.clear() - - @staticmethod - def intersect(x1, x2, z1, z2): - if x1 > z1: - return None if x1 > z2 else (x1, min(x2, z2)) - else: - return (z1, min(x2, z2)) if x2 > z1 else None diff --git a/telethon/telegram_bare_client.py b/telethon/telegram_bare_client.py index 24b60f6f..dbd12e41 100644 --- a/telethon/telegram_bare_client.py +++ b/telethon/telegram_bare_client.py @@ -12,7 +12,7 @@ from .errors import ( ) from .network import authenticator, MtProtoSender, TcpTransport from .utils import get_appropriated_part_size -from .crypto import rsa, HashChecker +from .crypto import rsa, CdnDecrypter # For sending and receiving requests from .tl import TLObject, JsonSession @@ -298,21 +298,6 @@ class TelegramBareClient: self._cached_clients[dc_id] = client return client - def _get_cdn_client(self, dc_id, query): - """_get_exported_client counterpart for CDNs. - Returns a tuple of (client, query result) - """ - dc = self._get_dc(dc_id, cdn=True) - session = JsonSession(self.session) - session.server_address = dc.ip_address - session.port = dc.port - client = TelegramBareClient( - session, self.api_id, self.api_hash, - timeout=self._timeout - ) - # This will make use of the new RSA keys for this specific CDN - return client, client.connect(initial_query=query) - # endregion # region Invoking Telegram requests @@ -485,38 +470,25 @@ class TelegramBareClient: try: offset_index = 0 - cdn_file_token = None - hash_checker = None - - def encrypt_method(x): - return x # Defaults to no-op + cdn_decrypter = None while True: offset = offset_index * part_size try: - if cdn_file_token: - result = client(GetCdnFileRequest( - cdn_file_token, offset, part_size - )) + if cdn_decrypter: + result = cdn_decrypter.get_file(offset, part_size) else: result = client(GetFileRequest( input_location, offset, part_size )) if isinstance(result, FileCdnRedirect): - cdn_file_token = result.file_token - hash_checker = HashChecker( - result.cdn_file_hashes + cdn_decrypter, result = \ + CdnDecrypter.prepare_decrypter( + client, TelegramBareClient, result, + offset, part_size ) - client, encrypt_method, result = \ - self._prepare_cdn_redirect( - result, offset, part_size - ) - - 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) @@ -527,12 +499,11 @@ class TelegramBareClient: # 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 - hash_checker.finish_check() - return getattr(result, 'type', '') + if cdn_decrypter: + cdn_decrypter.finish_check() - result.bytes = encrypt_method(result.bytes) - hash_checker.check(offset, result.bytes) + # Return some extra information, unless it's a CDN file + return getattr(result, 'type', '') f.write(result.bytes) if progress_callback: @@ -541,35 +512,4 @@ class TelegramBareClient: if isinstance(file, str): f.close() - def _prepare_cdn_redirect(self, cdn_redirect, offset, part_size): - """Returns (client, 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_aes.encrypt, None - else: - # We have the first bytes for the file - return client, cdn_aes.encrypt, cdn_file - # endregion