LonamiWebs/Telethon
1
1
Fork 0
mirror of https://github.com/LonamiWebs/Telethon.git synced 2025-02-03 13:14:31 +03:00
Code Issues Packages Projects Releases Wiki Activity

Add ._get_cdn_client as alternative ._get_exported_client version

Browse Source
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
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

32
telethon/crypto/cdn_decrypter.py
View File

@ -10,7 +10,7 @@ from ..errors import CdnFileTamperedError
class CdnDecrypter: class CdnDecrypter:
"""Used when downloading a file results in a 'FileCdnRedirect' to """Used when downloading a file results in a 'FileCdnRedirect' to
both prepare the redirect, decrypt the file as it downloads, and 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): def __init__(self, cdn_client, file_token, cdn_aes, cdn_file_hashes):
self.client = cdn_client self.client = cdn_client
@ -19,46 +19,26 @@ class CdnDecrypter:
self.cdn_file_hashes = cdn_file_hashes self.cdn_file_hashes = cdn_file_hashes
@staticmethod @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). """Prepares a CDN decrypter, returning (decrypter, file data).
'client' should be the original TelegramBareClient that 'client' should be an existing client not connected to a CDN.
tried to download the file. 'cdn_client' should be an already-connected TelegramBareClient
with the auth key already created.
'client_cls' should be the class of the TelegramBareClient.
""" """
# TODO Avoid the need for 'client_cls=TelegramBareClient'
# https://core.telegram.org/cdn
cdn_aes = AESModeCTR( cdn_aes = AESModeCTR(
key=cdn_redirect.encryption_key, key=cdn_redirect.encryption_key,
# 12 first bytes of the IV..4 bytes of the offset (0, big endian) # 12 first bytes of the IV..4 bytes of the offset (0, big endian)
iv=cdn_redirect.encryption_iv[:12] + bytes(4) 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, # 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. # 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( decrypter = CdnDecrypter(
cdn_client, cdn_redirect.file_token, cdn_client, cdn_redirect.file_token,
cdn_aes, cdn_redirect.cdn_file_hashes cdn_aes, cdn_redirect.cdn_file_hashes
) )
cdn_file = client(GetCdnFileRequest( cdn_file = cdn_client(GetCdnFileRequest(
file_token=cdn_redirect.file_token, file_token=cdn_redirect.file_token,
offset=cdn_redirect.cdn_file_hashes[0].offset, offset=cdn_redirect.cdn_file_hashes[0].offset,
limit=cdn_redirect.cdn_file_hashes[0].limit limit=cdn_redirect.cdn_file_hashes[0].limit

37
telethon/telegram_bare_client.py
View File

@ -154,7 +154,7 @@ class TelegramBareClient:
# region Connecting # 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 """Connects to the Telegram servers, executing authentication if
required. Note that authenticating to the Telegram servers is required. Note that authenticating to the Telegram servers is
not the same as authenticating the desired user itself, which 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 will FAIL if the client is not connected to the user's
native data center, raising a "UserMigrateError", and native data center, raising a "UserMigrateError", and
calling .disconnect() in the process. 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() self._main_thread_ident = threading.get_ident()
@ -195,7 +198,7 @@ class TelegramBareClient:
self._init_connection(ImportAuthorizationRequest( self._init_connection(ImportAuthorizationRequest(
_exported_auth.id, _exported_auth.bytes _exported_auth.id, _exported_auth.bytes
)) ))
else: elif not cdn:
TelegramBareClient._dc_options = \ TelegramBareClient._dc_options = \
self._init_connection(GetConfigRequest()).dc_options self._init_connection(GetConfigRequest()).dc_options
@ -204,7 +207,7 @@ class TelegramBareClient:
_exported_auth.id, _exported_auth.bytes _exported_auth.id, _exported_auth.bytes
)) ))
if TelegramBareClient._dc_options is None: if TelegramBareClient._dc_options is None and not cdn:
TelegramBareClient._dc_options = \ TelegramBareClient._dc_options = \
self(GetConfigRequest()).dc_options self(GetConfigRequest()).dc_options
@ -213,7 +216,7 @@ class TelegramBareClient:
# another data center and this would raise UserMigrateError) # another data center and this would raise UserMigrateError)
# to also assert whether the user is logged in or not. # to also assert whether the user is logged in or not.
self._user_connected = True self._user_connected = True
if _sync_updates: if _sync_updates and not cdn:
try: try:
self.sync_updates() self.sync_updates()
self._set_connected_and_authorized() self._set_connected_and_authorized()
@ -347,6 +350,7 @@ class TelegramBareClient:
export_auth = None # Already bound with the auth key export_auth = None # Already bound with the auth key
else: else:
# TODO Add a lock, don't allow two threads to create an auth key # 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. # for the same data center.
dc = self._get_dc(dc_id) dc = self._get_dc(dc_id)
@ -371,6 +375,29 @@ class TelegramBareClient:
client.connect(_exported_auth=export_auth, _sync_updates=False) client.connect(_exported_auth=export_auth, _sync_updates=False)
return client 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 # endregion
# region Invoking Telegram requests # region Invoking Telegram requests
@ -651,7 +678,7 @@ class TelegramBareClient:
if isinstance(result, FileCdnRedirect): if isinstance(result, FileCdnRedirect):
cdn_decrypter, result = \ cdn_decrypter, result = \
CdnDecrypter.prepare_decrypter( CdnDecrypter.prepare_decrypter(
client, TelegramBareClient, result client, self._get_cdn_client(result), result
) )
except FileMigrateError as e: except FileMigrateError as e: