Create auth_key if not present

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 (

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')

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('<QQi', 0, self._get_new_msg_id(), len(data)) + data
+        body = bytes(request)
+        msg_id = self._get_new_msg_id()
+        await self._connection.send(
+            struct.pack('<QQi', 0, msg_id, len(body)) + body
         )
 
-    def receive(self):
-        """
-        Receives a plain packet from the network.
-
-        :return: the response body.
-        """
-        body = self._connection.recv()
+        body = await self._connection.recv()
         if body == b'l\xfe\xff\xff':  # -404 little endian signed
             # Broken authorization, must reset the auth key
             raise BrokenAuthKeyError()
 
         with BinaryReader(body) as reader:
-            reader.read_long()  # auth_key_id
-            reader.read_long()  # msg_id
-            message_length = reader.read_int()
-
-            response = reader.read(message_length)
-            return response
+            assert reader.read_long() == 0  # auth_key_id
+            assert reader.read_long() > 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."""

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