Make Auth asynchronous

delivrance · delivrance · commit 8049c9129bfd · 2018-06-18T13:07:02.000+02:00
diff --git a/pyrogram/session/auth.py b/pyrogram/session/auth.py
@@ -67,14 +67,14 @@ def unpack(b: BytesIO):
         b.seek(20)  # Skip auth_key_id (8), message_id (8) and message_length (4)
         return Object.read(b)
 
-    def send(self, data: Object):
+    async def send(self, data: Object):
         data = self.pack(data)
-        self.connection.send(data)
-        response = BytesIO(self.connection.recv())
+        await self.connection.send(data)
+        response = BytesIO(await self.connection.recv())
 
         return self.unpack(response)
 
-    def create(self):
+    async def create(self):
         """
         https://core.telegram.org/mtproto/auth_key
         https://core.telegram.org/mtproto/samples-auth_key
@@ -89,12 +89,12 @@ def create(self):
             try:
                 log.info("Start creating a new auth key on DC{}".format(self.dc_id))
 
-                self.connection.connect()
+                await self.connection.connect()
 
                 # Step 1; Step 2
                 nonce = int.from_bytes(urandom(16), "little", signed=True)
                 log.debug("Send req_pq: {}".format(nonce))
-                res_pq = self.send(functions.ReqPqMulti(nonce))
+                res_pq = await self.send(functions.ReqPqMulti(nonce))
                 log.debug("Got ResPq: {}".format(res_pq.server_nonce))
                 log.debug("Server public key fingerprints: {}".format(res_pq.server_public_key_fingerprints))
 
@@ -138,7 +138,7 @@ def create(self):
 
                 # Step 5. TODO: Handle "server_DH_params_fail". Code assumes response is ok
                 log.debug("Send req_DH_params")
-                server_dh_params = self.send(
+                server_dh_params = await self.send(
                     functions.ReqDHParams(
                         nonce,
                         server_nonce,
@@ -198,7 +198,7 @@ def create(self):
                 encrypted_data = AES.ige256_encrypt(data_with_hash, tmp_aes_key, tmp_aes_iv)
 
                 log.debug("Send set_client_DH_params")
-                set_client_dh_params_answer = self.send(
+                set_client_dh_params_answer = await self.send(
                     functions.SetClientDHParams(
                         nonce,
                         server_nonce,
