diff --git a/telethon/network/__init__.py b/telethon/network/__init__.py index 6d8584bc..f4bd72d0 100644 --- a/telethon/network/__init__.py +++ b/telethon/network/__init__.py @@ -2,7 +2,7 @@ This module contains several classes regarding network, low level connection with Telegram's servers and the protocol used (TCP full, abridged, etc.). """ -from .mtproto_plain_sender import MtProtoPlainSender +from .mtprotoplainsender import MTProtoPlainSender from .authenticator import do_authentication from .mtprotosender import MTProtoSender from .connection import ( diff --git a/telethon/network/authenticator.py b/telethon/network/authenticator.py index 8635c030..02b1128d 100644 --- a/telethon/network/authenticator.py +++ b/telethon/network/authenticator.py @@ -14,56 +14,24 @@ from .. import helpers as utils from ..crypto import AES, AuthKey, Factorization, rsa from ..errors import SecurityError from ..extensions import BinaryReader -from ..network import MtProtoPlainSender from ..tl.functions import ( ReqPqMultiRequest, ReqDHParamsRequest, SetClientDHParamsRequest ) -def do_authentication(connection, retries=5): - """ - Performs the authentication steps on the given connection. - Raises an error if all attempts fail. - - :param connection: the connection to be used (must be connected). - :param retries: how many times should we retry on failure. - :return: - """ - if not retries or retries < 0: - retries = 1 - - last_error = None - while retries: - try: - return _do_authentication(connection) - except (SecurityError, AssertionError, NotImplementedError) as e: - last_error = e - retries -= 1 - raise last_error - - -def _do_authentication(connection): +async def do_authentication(sender): """ Executes the authentication process with the Telegram servers. - :param connection: the connection to be used (must be connected). + :param sender: a connected `MTProtoPlainSender`. :return: returns a (authorization key, time offset) tuple. """ - sender = MtProtoPlainSender(connection) - # Step 1 sending: PQ Request, endianness doesn't matter since it's random - req_pq_request = ReqPqMultiRequest( - nonce=int.from_bytes(os.urandom(16), 'big', signed=True) - ) - sender.send(bytes(req_pq_request)) - with BinaryReader(sender.receive()) as reader: - req_pq_request.on_response(reader) + nonce = int.from_bytes(os.urandom(16), 'big', signed=True) + res_pq = await sender.send(ReqPqMultiRequest(nonce)) + assert isinstance(res_pq, ResPQ) - res_pq = req_pq_request.result - if not isinstance(res_pq, ResPQ): - raise AssertionError(res_pq) - - if res_pq.nonce != req_pq_request.nonce: + if res_pq.nonce != nonce: raise SecurityError('Invalid nonce from server') pq = get_int(res_pq.pq) @@ -96,20 +64,14 @@ def _do_authentication(connection): ) ) - req_dh_params = ReqDHParamsRequest( + server_dh_params = await sender.send(ReqDHParamsRequest( nonce=res_pq.nonce, server_nonce=res_pq.server_nonce, p=p, q=q, public_key_fingerprint=target_fingerprint, encrypted_data=cipher_text - ) - sender.send(bytes(req_dh_params)) + )) - # Step 2 response: DH Exchange - with BinaryReader(sender.receive()) as reader: - req_dh_params.on_response(reader) - - server_dh_params = req_dh_params.result if isinstance(server_dh_params, ServerDHParamsFail): raise SecurityError('Server DH params fail: TODO') @@ -168,18 +130,12 @@ def _do_authentication(connection): client_dh_encrypted = AES.encrypt_ige(client_dh_inner_hashed, key, iv) # Prepare Set client DH params - set_client_dh = SetClientDHParamsRequest( + dh_gen = await sender.send(SetClientDHParamsRequest( nonce=res_pq.nonce, server_nonce=res_pq.server_nonce, encrypted_data=client_dh_encrypted, - ) - sender.send(bytes(set_client_dh)) + )) - # Step 3 response: Complete DH Exchange - with BinaryReader(sender.receive()) as reader: - set_client_dh.on_response(reader) - - dh_gen = set_client_dh.result if isinstance(dh_gen, DhGenOk): if dh_gen.nonce != res_pq.nonce: raise SecurityError('Invalid nonce from server') diff --git a/telethon/network/mtproto_plain_sender.py b/telethon/network/mtprotoplainsender.py similarity index 61% rename from telethon/network/mtproto_plain_sender.py rename to telethon/network/mtprotoplainsender.py index cb6d63af..82a8a18d 100644 --- a/telethon/network/mtproto_plain_sender.py +++ b/telethon/network/mtprotoplainsender.py @@ -9,12 +9,11 @@ from ..errors import BrokenAuthKeyError from ..extensions import BinaryReader -class MtProtoPlainSender: +class MTProtoPlainSender: """ MTProto Mobile Protocol plain sender (https://core.telegram.org/mtproto/description#unencrypted-messages) """ - def __init__(self, connection): """ Initializes the MTProto plain sender. @@ -26,43 +25,27 @@ class MtProtoPlainSender: self._last_msg_id = 0 self._connection = connection - def connect(self): - """Connects to Telegram's servers.""" - self._connection.connect() - - def disconnect(self): - """Disconnects from Telegram's servers.""" - self._connection.close() - - def send(self, data): + async def send(self, request): """ - Sends a plain packet (auth_key_id = 0) containing the - given message body (data). - - :param data: the data to be sent. + Sends and receives the result for the given request. """ - self._connection.send( - struct.pack(' msg_id # msg_id + assert reader.read_int() # length + # No need to read "length" bytes first, just read the object + return reader.tgread_object() def _get_new_msg_id(self): """Generates a new message ID based on the current time since epoch.""" diff --git a/telethon/network/mtprotosender.py b/telethon/network/mtprotosender.py index 87f949e1..ae89cb01 100644 --- a/telethon/network/mtprotosender.py +++ b/telethon/network/mtprotosender.py @@ -1,9 +1,13 @@ import asyncio import logging +from . import MTProtoPlainSender, authenticator from .connection import ConnectionTcpFull from .. import helpers, utils -from ..errors import BadMessageError, TypeNotFoundError, rpc_message_to_error +from ..errors import ( + BadMessageError, TypeNotFoundError, BrokenAuthKeyError, SecurityError, + rpc_message_to_error +) from ..extensions import BinaryReader from ..tl import TLMessage, MessageContainer, GzipPacked from ..tl.functions.auth import LogOutRequest @@ -100,6 +104,13 @@ class MTProtoSender: async with self._send_lock: await self._connection.connect(ip, port) self._user_connected = True + + # TODO Handle SecurityError, AssertionError, NotImplementedError + if self.session.auth_key is None: + plain = MTProtoPlainSender(self._connection) + self.session.auth_key, self.session.time_offset =\ + await authenticator.do_authentication(plain) + self._send_loop_handle = asyncio.ensure_future(self._send_loop()) self._recv_loop_handle = asyncio.ensure_future(self._recv_loop()) @@ -214,11 +225,20 @@ class MTProtoSender: body = await self._connection.recv() # TODO Check salt, session_id and sequence_number - message, remote_msg_id, remote_seq = helpers.unpack_message( - self.session, body) - - with BinaryReader(message) as reader: - await self._process_message(remote_msg_id, remote_seq, reader) + try: + message, remote_msg_id, remote_seq =\ + helpers.unpack_message(self.session, body) + except (BrokenAuthKeyError, BufferError): + # TODO Are these temporary or do we need a new key? + pass + except SecurityError: + # TODO Can we safely ignore these? Has the message + # been decoded correctly? + pass + else: + with BinaryReader(message) as reader: + await self._process_message( + remote_msg_id, remote_seq, reader) # Response Handlers