From cbcb1c78c4498c1ac3a4b9fc84734e4b9881a0f9 Mon Sep 17 00:00:00 2001 From: Dan <14043624+delivrance@users.noreply.github.com> Date: Sat, 9 Dec 2017 15:26:33 +0100 Subject: [PATCH] Don't throw errors if auth key creation fails; try again instead --- pyrogram/session/auth.py | 271 ++++++++++++++++++++------------------- 1 file changed, 138 insertions(+), 133 deletions(-) diff --git a/pyrogram/session/auth.py b/pyrogram/session/auth.py index bf3cd05c..f90c64df 100644 --- a/pyrogram/session/auth.py +++ b/pyrogram/session/auth.py @@ -31,9 +31,6 @@ from .internals import MsgId, DataCenter 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 @@ class Auth: 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() + # 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)) - # 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)) + self.connection.connect() - # 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 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 4 - server_nonce = res_pq.server_nonce - new_nonce = int.from_bytes(urandom(32), "little", signed=True) + # 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)) - data = types.PQInnerData( - res_pq.pq, - int.to_bytes(p, 4, "big"), - int.to_bytes(q, 4, "big"), - nonce, - server_nonce, - new_nonce, - ).write() + # Step 4 + server_nonce = res_pq.server_nonce + new_nonce = int.from_bytes(urandom(32), "little", signed=True) - 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]) + data = types.PQInnerData( + res_pq.pq, + int.to_bytes(p, 4, "big"), + int.to_bytes(q, 4, "big"), + nonce, + server_nonce, + new_nonce, + ).write() - log.debug("Done encrypt data with RSA") + 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]) - # 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 - ) - ) + log.debug("Done encrypt data with RSA") - encrypted_answer = server_dh_params.encrypted_answer + # 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 + ) + ) - server_nonce = int.to_bytes(server_nonce, 16, "little", signed=True) - new_nonce = int.to_bytes(new_nonce, 32, "little", signed=True) + encrypted_answer = server_dh_params.encrypted_answer - tmp_aes_key = ( - sha1(new_nonce + server_nonce).digest() - + sha1(server_nonce + new_nonce).digest()[:12] - ) + server_nonce = int.to_bytes(server_nonce, 16, "little", signed=True) + new_nonce = int.to_bytes(new_nonce, 32, "little", signed=True) - tmp_aes_iv = ( - sha1(server_nonce + new_nonce).digest()[12:] - + sha1(new_nonce + new_nonce).digest() + new_nonce[:4] - ) + tmp_aes_key = ( + sha1(new_nonce + server_nonce).digest() + + sha1(server_nonce + new_nonce).digest()[:12] + ) - server_nonce = int.from_bytes(server_nonce, "little", signed=True) + tmp_aes_iv = ( + sha1(server_nonce + new_nonce).digest()[12:] + + sha1(new_nonce + new_nonce).digest() + new_nonce[:4] + ) - answer_with_hash = IGE.decrypt(encrypted_answer, tmp_aes_key, tmp_aes_iv) - answer = answer_with_hash[20:] + server_nonce = int.from_bytes(server_nonce, "little", signed=True) - server_dh_inner_data = Object.read(BytesIO(answer)) + answer_with_hash = IGE.decrypt(encrypted_answer, tmp_aes_key, tmp_aes_iv) + answer = answer_with_hash[20:] - log.debug("Done decrypting answer") + server_dh_inner_data = Object.read(BytesIO(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("Done decrypting answer") - log.debug("Delta time: {}".format(round(delta_time, 3))) + dh_prime = int.from_bytes(server_dh_inner_data.dh_prime, "big") + delta_time = server_dh_inner_data.server_time - time.time() - # 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") + log.debug("Delta time: {}".format(round(delta_time, 3))) - retry_id = 0 + # 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") - data = types.ClientDhInnerData( - nonce, - server_nonce, - retry_id, - g_b - ).write() + retry_id = 0 - 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) + data = types.ClientDhInnerData( + nonce, + server_nonce, + retry_id, + g_b + ).write() - log.debug("Send set_client_DH_params") - set_client_dh_params_answer = self.send( - functions.SetClientDhParams( - nonce, - server_nonce, - encrypted_data - ) - ) + 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) - # TODO: Handle "auth_key_aux_hash" if the previous step fails + log.debug("Send set_client_DH_params") + set_client_dh_params_answer = self.send( + functions.SetClientDhParams( + nonce, + server_nonce, + encrypted_data + ) + ) - # 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 "auth_key_aux_hash" if the previous step fails - # TODO: Handle errors + # 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) - ####################### - # Security checks - ####################### + # TODO: Handle errors - assert dh_prime == self.CURRENT_DH_PRIME - log.debug("DH parameters check: OK") + ####################### + # Security checks + ####################### - # 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") + assert dh_prime == self.CURRENT_DH_PRIME + log.debug("DH parameters check: 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#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-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") + # 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") - # Step 9 - server_salt = IGE.xor(new_nonce[:8], server_nonce[:8]) + # 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") - log.debug("Server salt: {}".format(int.from_bytes(server_salt, "little"))) + # Step 9 + server_salt = IGE.xor(new_nonce[:8], server_nonce[:8]) - log.info( - "Done auth key exchange: {}".format( - set_client_dh_params_answer.__class__.__name__ - ) - ) + log.debug("Server salt: {}".format(int.from_bytes(server_salt, "little"))) - self.connection.close() - - return auth_key + 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()