Problem-2
diff --git a/‎pyrogram/session/auth.py‎
Lines changed: 172 additions & 167 deletions b/‎pyrogram/session/auth.py‎
Lines changed: 172 additions & 167 deletions
@@ -31,9 +31,6 @@
 log = logging.getLogger(__name__)
 
 
-# TODO: When using TCP connection mode, the server may close it at any time, causing the Auth key creation to fail
-# The above is true when dealing with temporary keys, although for perm keys it didn't happened, yet.
-
 class Auth:
     CURRENT_DH_PRIME = int(
         "C71CAEB9C6B1C9048E6C522F70F13F73980D40238E3E21C14934D037563D930F"
@@ -79,168 +76,176 @@ def create(self):
         https://core.telegram.org/mtproto/auth_key
         https://core.telegram.org/mtproto/samples-auth_key
         """
-        log.info("Start creating a new auth key on DC{}".format(self.dc_id))
-
-        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.ReqPq(nonce))
-        log.debug("Got ResPq: {}".format(res_pq.server_nonce))
-
-        # Step 3
-        pq = int.from_bytes(res_pq.pq, "big")
-        log.debug("Start PQ factorization: {}".format(pq))
-        start = time.time()
-        g = Prime.decompose(pq)
-        p, q = sorted((g, pq // g))  # p < q
-        log.debug("Done PQ factorization ({}s): {} {}".format(round(time.time() - start, 3), p, q))
-
-        # Step 4
-        server_nonce = res_pq.server_nonce
-        new_nonce = int.from_bytes(urandom(32), "little", signed=True)
-
-        data = types.PQInnerData(
-            res_pq.pq,
-            int.to_bytes(p, 4, "big"),
-            int.to_bytes(q, 4, "big"),
-            nonce,
-            server_nonce,
-            new_nonce,
-        ).write()
-
-        sha = sha1(data).digest()
-        padding = urandom(- (len(data) + len(sha)) % 255)
-        data_with_hash = sha + data + padding
-        encrypted_data = RSA.encrypt(data_with_hash, res_pq.server_public_key_fingerprints[0])
-
-        log.debug("Done encrypt data with RSA")
-
-        # Step 5. TODO: Handle "server_DH_params_fail". Code assumes response is ok
-        log.debug("Send req_DH_params")
-        server_dh_params = self.send(
-            functions.ReqDhParams(
-                nonce,
-                server_nonce,
-                int.to_bytes(p, 4, "big"),
-                int.to_bytes(q, 4, "big"),
-                res_pq.server_public_key_fingerprints[0],
-                encrypted_data
-            )
-        )
-
-        encrypted_answer = server_dh_params.encrypted_answer
-
-        server_nonce = int.to_bytes(server_nonce, 16, "little", signed=True)
-        new_nonce = int.to_bytes(new_nonce, 32, "little", signed=True)
-
-        tmp_aes_key = (
-            sha1(new_nonce + server_nonce).digest()
-            + sha1(server_nonce + new_nonce).digest()[:12]
-        )
-
-        tmp_aes_iv = (
-            sha1(server_nonce + new_nonce).digest()[12:]
-            + sha1(new_nonce + new_nonce).digest() + new_nonce[:4]
-        )
-
-        server_nonce = int.from_bytes(server_nonce, "little", signed=True)
-
-        answer_with_hash = IGE.decrypt(encrypted_answer, tmp_aes_key, tmp_aes_iv)
-        answer = answer_with_hash[20:]
-
-        server_dh_inner_data = Object.read(BytesIO(answer))
-
-        log.debug("Done decrypting answer")
-
-        dh_prime = int.from_bytes(server_dh_inner_data.dh_prime, "big")
-        delta_time = server_dh_inner_data.server_time - time.time()
-
-        log.debug("Delta time: {}".format(round(delta_time, 3)))
-
-        # Step 6
-        g = server_dh_inner_data.g
-        b = int.from_bytes(urandom(256), "big")
-        g_b = int.to_bytes(pow(g, b, dh_prime), 256, "big")
-
-        retry_id = 0
-
-        data = types.ClientDhInnerData(
-            nonce,
-            server_nonce,
-            retry_id,
-            g_b
-        ).write()
-
-        sha = sha1(data).digest()
-        padding = urandom(- (len(data) + len(sha)) % 16)
-        data_with_hash = sha + data + padding
-        encrypted_data = IGE.encrypt(data_with_hash, tmp_aes_key, tmp_aes_iv)
-
-        log.debug("Send set_client_DH_params")
-        set_client_dh_params_answer = self.send(
-            functions.SetClientDhParams(
-                nonce,
-                server_nonce,
-                encrypted_data
-            )
-        )
-
-        # TODO: Handle "auth_key_aux_hash" if the previous step fails
-
-        # Step 7; Step 8
-        g_a = int.from_bytes(server_dh_inner_data.g_a, "big")
-        auth_key = int.to_bytes(pow(g_a, b, dh_prime), 256, "big")
-        server_nonce = int.to_bytes(server_nonce, 16, "little", signed=True)
-
-        # TODO: Handle errors
-
-        #######################
-        # Security checks
-        #######################
-
-        assert dh_prime == self.CURRENT_DH_PRIME
-        log.debug("DH parameters check: OK")
-
-        # https://core.telegram.org/mtproto/security_guidelines#g-a-and-g-b-validation
-        g_b = int.from_bytes(g_b, "big")
-        assert 1 < g < dh_prime - 1
-        assert 1 < g_a < dh_prime - 1
-        assert 1 < g_b < dh_prime - 1
-        assert 2 ** (2048 - 64) < g_a < dh_prime - 2 ** (2048 - 64)
-        assert 2 ** (2048 - 64) < g_b < dh_prime - 2 ** (2048 - 64)
-        log.debug("g_a and g_b validation: OK")
-
-        # https://core.telegram.org/mtproto/security_guidelines#checking-sha1-hash-values
-        answer = server_dh_inner_data.write()  # Call .write() to remove padding
-        assert answer_with_hash[:20] == sha1(answer).digest()
-        log.debug("SHA1 hash values check: OK")
-
-        # https://core.telegram.org/mtproto/security_guidelines#checking-nonce-server-nonce-and-new-nonce-fields
-        # 1st message
-        assert nonce == res_pq.nonce
-        # 2nd message
-        server_nonce = int.from_bytes(server_nonce, "little", signed=True)
-        assert nonce == server_dh_params.nonce
-        assert server_nonce == server_dh_params.server_nonce
-        # 3rd message
-        assert nonce == set_client_dh_params_answer.nonce
-        assert server_nonce == set_client_dh_params_answer.server_nonce
-        server_nonce = int.to_bytes(server_nonce, 16, "little", signed=True)
-        log.debug("Nonce fields check: OK")
-
-        # Step 9
-        server_salt = IGE.xor(new_nonce[:8], server_nonce[:8])
-
-        log.debug("Server salt: {}".format(int.from_bytes(server_salt, "little")))
-
-        log.info(
-            "Done auth key exchange: {}".format(
-                set_client_dh_params_answer.__class__.__name__
-            )
-        )
-
-        self.connection.close()
 
-        return auth_key
+        # The server may close the connection at any time, causing the auth key creation to fail.
+        # If that happens, just try again until it succeed.
+        while True:
+            try:
+                log.info("Start creating a new auth key on DC{}".format(self.dc_id))
+
+                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.ReqPq(nonce))
+                log.debug("Got ResPq: {}".format(res_pq.server_nonce))
+
+                # Step 3
+                pq = int.from_bytes(res_pq.pq, "big")
+                log.debug("Start PQ factorization: {}".format(pq))
+                start = time.time()
+                g = Prime.decompose(pq)
+                p, q = sorted((g, pq // g))  # p < q
+                log.debug("Done PQ factorization ({}s): {} {}".format(round(time.time() - start, 3), p, q))
+
+                # Step 4
+                server_nonce = res_pq.server_nonce
+                new_nonce = int.from_bytes(urandom(32), "little", signed=True)
+
+                data = types.PQInnerData(
+                    res_pq.pq,
+                    int.to_bytes(p, 4, "big"),
+                    int.to_bytes(q, 4, "big"),
+                    nonce,
+                    server_nonce,
+                    new_nonce,
+                ).write()
+
+                sha = sha1(data).digest()
+                padding = urandom(- (len(data) + len(sha)) % 255)
+                data_with_hash = sha + data + padding
+                encrypted_data = RSA.encrypt(data_with_hash, res_pq.server_public_key_fingerprints[0])
+
+                log.debug("Done encrypt data with RSA")
+
+                # Step 5. TODO: Handle "server_DH_params_fail". Code assumes response is ok
+                log.debug("Send req_DH_params")
+                server_dh_params = self.send(
+                    functions.ReqDhParams(
+                        nonce,
+                        server_nonce,
+                        int.to_bytes(p, 4, "big"),
+                        int.to_bytes(q, 4, "big"),
+                        res_pq.server_public_key_fingerprints[0],
+                        encrypted_data
+                    )
+                )
+
+                encrypted_answer = server_dh_params.encrypted_answer
+
+                server_nonce = int.to_bytes(server_nonce, 16, "little", signed=True)
+                new_nonce = int.to_bytes(new_nonce, 32, "little", signed=True)
+
+                tmp_aes_key = (
+                    sha1(new_nonce + server_nonce).digest()
+                    + sha1(server_nonce + new_nonce).digest()[:12]
+                )
+
+                tmp_aes_iv = (
+                    sha1(server_nonce + new_nonce).digest()[12:]
+                    + sha1(new_nonce + new_nonce).digest() + new_nonce[:4]
+                )
+
+                server_nonce = int.from_bytes(server_nonce, "little", signed=True)
+
+                answer_with_hash = IGE.decrypt(encrypted_answer, tmp_aes_key, tmp_aes_iv)
+                answer = answer_with_hash[20:]
+
+                server_dh_inner_data = Object.read(BytesIO(answer))
+
+                log.debug("Done decrypting answer")
+
+                dh_prime = int.from_bytes(server_dh_inner_data.dh_prime, "big")
+                delta_time = server_dh_inner_data.server_time - time.time()
+
+                log.debug("Delta time: {}".format(round(delta_time, 3)))
+
+                # Step 6
+                g = server_dh_inner_data.g
+                b = int.from_bytes(urandom(256), "big")
+                g_b = int.to_bytes(pow(g, b, dh_prime), 256, "big")
+
+                retry_id = 0
+
+                data = types.ClientDhInnerData(
+                    nonce,
+                    server_nonce,
+                    retry_id,
+                    g_b
+                ).write()
+
+                sha = sha1(data).digest()
+                padding = urandom(- (len(data) + len(sha)) % 16)
+                data_with_hash = sha + data + padding
+                encrypted_data = IGE.encrypt(data_with_hash, tmp_aes_key, tmp_aes_iv)
+
+                log.debug("Send set_client_DH_params")
+                set_client_dh_params_answer = self.send(
+                    functions.SetClientDhParams(
+                        nonce,
+                        server_nonce,
+                        encrypted_data
+                    )
+                )
+
+                # TODO: Handle "auth_key_aux_hash" if the previous step fails
+
+                # Step 7; Step 8
+                g_a = int.from_bytes(server_dh_inner_data.g_a, "big")
+                auth_key = int.to_bytes(pow(g_a, b, dh_prime), 256, "big")
+                server_nonce = int.to_bytes(server_nonce, 16, "little", signed=True)
+
+                # TODO: Handle errors
+
+                #######################
+                # Security checks
+                #######################
+
+                assert dh_prime == self.CURRENT_DH_PRIME
+                log.debug("DH parameters check: OK")
+
+                # https://core.telegram.org/mtproto/security_guidelines#g-a-and-g-b-validation
+                g_b = int.from_bytes(g_b, "big")
+                assert 1 < g < dh_prime - 1
+                assert 1 < g_a < dh_prime - 1
+                assert 1 < g_b < dh_prime - 1
+                assert 2 ** (2048 - 64) < g_a < dh_prime - 2 ** (2048 - 64)
+                assert 2 ** (2048 - 64) < g_b < dh_prime - 2 ** (2048 - 64)
+                log.debug("g_a and g_b validation: OK")
+
+                # https://core.telegram.org/mtproto/security_guidelines#checking-sha1-hash-values
+                answer = server_dh_inner_data.write()  # Call .write() to remove padding
+                assert answer_with_hash[:20] == sha1(answer).digest()
+                log.debug("SHA1 hash values check: OK")
+
+                # https://core.telegram.org/mtproto/security_guidelines#checking-nonce-server-nonce-and-new-nonce-fields
+                # 1st message
+                assert nonce == res_pq.nonce
+                # 2nd message
+                server_nonce = int.from_bytes(server_nonce, "little", signed=True)
+                assert nonce == server_dh_params.nonce
+                assert server_nonce == server_dh_params.server_nonce
+                # 3rd message
+                assert nonce == set_client_dh_params_answer.nonce
+                assert server_nonce == set_client_dh_params_answer.server_nonce
+                server_nonce = int.to_bytes(server_nonce, 16, "little", signed=True)
+                log.debug("Nonce fields check: OK")
+
+                # Step 9
+                server_salt = IGE.xor(new_nonce[:8], server_nonce[:8])
+
+                log.debug("Server salt: {}".format(int.from_bytes(server_salt, "little")))
+
+                log.info(
+                    "Done auth key exchange: {}".format(
+                        set_client_dh_params_answer.__class__.__name__
+                    )
+                )
+            except:  # TODO: Too broad exception clause
+                log.warning("Auth key creation failed. Let's try again.")
+                continue
+            else:
+                return auth_key
+            finally:
+                self.connection.close()