Add ._get_cdn_client as alternative ._get_exported_client version

This commit is contained in:
Lonami Exo 2017-09-30 17:51:07 +02:00
parent a35c4b15db
commit d28f370ab6
2 changed files with 38 additions and 31 deletions

View File

@ -10,7 +10,7 @@ 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.
ensure the file hasn't been tampered. https://core.telegram.org/cdn
"""
def __init__(self, cdn_client, file_token, cdn_aes, cdn_file_hashes):
self.client = cdn_client
@ -19,46 +19,26 @@ class CdnDecrypter:
self.cdn_file_hashes = cdn_file_hashes
@staticmethod
def prepare_decrypter(client, client_cls, cdn_redirect):
def prepare_decrypter(client, cdn_client, cdn_redirect):
"""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.
'client' should be an existing client not connected to a CDN.
'cdn_client' should be an already-connected TelegramBareClient
with the auth key already created.
"""
# TODO Avoid the need for 'client_cls=TelegramBareClient'
# https://core.telegram.org/cdn
cdn_aes = AESModeCTR(
key=cdn_redirect.encryption_key,
# 12 first bytes of the IV..4 bytes of the offset (0, big endian)
iv=cdn_redirect.encryption_iv[:12] + bytes(4)
)
# Create a new client on said CDN
dc = client._get_dc(cdn_redirect.dc_id, cdn=True)
session = Session(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._sender.connection.get_timeout()
)
# This will make use of the new RSA keys for this specific CDN.
#
# We assume that cdn_redirect.cdn_file_hashes are ordered by offset,
# and that there will be enough of these to retrieve the whole file.
#
# This relies on the fact that TelegramBareClient._dc_options is
# static and it won't be called from this DC (it would fail).
cdn_client.connect()
# CDN client is ready, create the resulting CdnDecrypter
decrypter = CdnDecrypter(
cdn_client, cdn_redirect.file_token,
cdn_aes, cdn_redirect.cdn_file_hashes
)
cdn_file = client(GetCdnFileRequest(
cdn_file = cdn_client(GetCdnFileRequest(
file_token=cdn_redirect.file_token,
offset=cdn_redirect.cdn_file_hashes[0].offset,
limit=cdn_redirect.cdn_file_hashes[0].limit

View File

@ -154,7 +154,7 @@ class TelegramBareClient:
# region Connecting
def connect(self, _exported_auth=None, _sync_updates=True):
def connect(self, _exported_auth=None, _sync_updates=True, cdn=False):
"""Connects to the Telegram servers, executing authentication if
required. Note that authenticating to the Telegram servers is
not the same as authenticating the desired user itself, which
@ -170,6 +170,9 @@ class TelegramBareClient:
will FAIL if the client is not connected to the user's
native data center, raising a "UserMigrateError", and
calling .disconnect() in the process.
If 'cdn' is False, methods that are not allowed on such data
centers won't be invoked.
"""
self._main_thread_ident = threading.get_ident()
@ -195,7 +198,7 @@ class TelegramBareClient:
self._init_connection(ImportAuthorizationRequest(
_exported_auth.id, _exported_auth.bytes
))
else:
elif not cdn:
TelegramBareClient._dc_options = \
self._init_connection(GetConfigRequest()).dc_options
@ -204,7 +207,7 @@ class TelegramBareClient:
_exported_auth.id, _exported_auth.bytes
))
if TelegramBareClient._dc_options is None:
if TelegramBareClient._dc_options is None and not cdn:
TelegramBareClient._dc_options = \
self(GetConfigRequest()).dc_options
@ -213,7 +216,7 @@ class TelegramBareClient:
# another data center and this would raise UserMigrateError)
# to also assert whether the user is logged in or not.
self._user_connected = True
if _sync_updates:
if _sync_updates and not cdn:
try:
self.sync_updates()
self._set_connected_and_authorized()
@ -347,6 +350,7 @@ class TelegramBareClient:
export_auth = None # Already bound with the auth key
else:
# TODO Add a lock, don't allow two threads to create an auth key
# (when calling .connect() if there wasn't a previous session).
# for the same data center.
dc = self._get_dc(dc_id)
@ -371,6 +375,29 @@ class TelegramBareClient:
client.connect(_exported_auth=export_auth, _sync_updates=False)
return client
def _get_cdn_client(self, cdn_redirect):
"""Similar to ._get_exported_client, but for CDNs"""
session = self._exported_sessions.get(cdn_redirect.dc_id)
if not session:
dc = self._get_dc(cdn_redirect.dc_id, cdn=True)
session = Session(self.session)
session.server_address = dc.ip_address
session.port = dc.port
self._exported_sessions[cdn_redirect.dc_id] = session
client = TelegramBareClient(
session, self.api_id, self.api_hash,
proxy=self._sender.connection.conn.proxy,
timeout=self._sender.connection.get_timeout()
)
# This will make use of the new RSA keys for this specific CDN.
#
# This relies on the fact that TelegramBareClient._dc_options is
# static and it won't be called from this DC (it would fail).
client.connect(cdn=True) # Avoid invoking non-CDN specific methods
return client
# endregion
# region Invoking Telegram requests
@ -651,7 +678,7 @@ class TelegramBareClient:
if isinstance(result, FileCdnRedirect):
cdn_decrypter, result = \
CdnDecrypter.prepare_decrypter(
client, TelegramBareClient, result
client, self._get_cdn_client(result), result
)
except FileMigrateError as e: