From 4d8cd202f1dac9bc86fce39c501ea9f339dfd065 Mon Sep 17 00:00:00 2001 From: lds56 Date: Sat, 4 Apr 2015 15:18:45 +0800 Subject: [PATCH] Add auth error handling --- .gitignore | 6 ++++- TL.py | 4 ++- mtproto.py | 59 ++++++++++++++++++++++++++++-------------- testing.py | 5 +++- tests/auth_key_test.py | 21 +++++++++++++++ 5 files changed, 73 insertions(+), 22 deletions(-) create mode 100644 tests/auth_key_test.py diff --git a/.gitignore b/.gitignore index a8163489..33170b57 100644 --- a/.gitignore +++ b/.gitignore @@ -59,4 +59,8 @@ target/ # testing apps credentials -rsa.pub \ No newline at end of file +rsa.pub + +# emacs auto-saving files +\#*# +.#*# \ No newline at end of file diff --git a/TL.py b/TL.py index 090c0639..30294ddb 100644 --- a/TL.py +++ b/TL.py @@ -152,11 +152,13 @@ def deserialize(bytes_io, type_=None, subtype=None): except KeyError: # Unknown type raise Exception("Could not extract type: %s" % type_) + base_boxed_types = ["Vector t", "Int", "Long", "Double", "String", "Int128", "Int256"] if tl_elem.type in base_boxed_types: x = deserialize(bytes_io, type_=tl_elem.predicate, subtype=subtype) else: # other types x = {} + x[u'name'] = tl_elem.predicate for arg in tl_elem.params: x[arg['name']] = deserialize(bytes_io, type_=arg['type'], subtype=arg['subtype']) - return x \ No newline at end of file + return x diff --git a/mtproto.py b/mtproto.py index ff9e40b7..cf7f6eac 100644 --- a/mtproto.py +++ b/mtproto.py @@ -52,6 +52,9 @@ class Session: self.session_id = os.urandom(8) self.auth_key = auth_key self.auth_key_id = SHA.new(self.auth_key).digest()[-8:] if self.auth_key else None + self.sock.settimeout(5.0) + self.MAX_RETRY = 5; + self.AUTH_MAX_RETRY = 5; def __del__(self): # closing socket when session object is deleted @@ -129,9 +132,14 @@ class Session: return data def method_call(self, method, **kwargs): - self.send_message(TL.serialize_method(method, **kwargs)) - server_answer = self.recv_message() - return TL.deserialize(io.BytesIO(server_answer)) + for i in range(1, self.MAX_RETRY): + try: + self.send_message(TL.serialize_method(method, **kwargs)) + server_answer = self.recv_message() + except socket.timeout: + print("Retry call method") + continue + return TL.deserialize(io.BytesIO(server_answer)) def create_auth_key(self): @@ -223,29 +231,42 @@ class Session: data_with_sha_padded = data_with_sha + os.urandom(-len(data_with_sha) % 16) encrypted_data = crypt.ige_encrypt(data_with_sha_padded, tmp_aes_key, tmp_aes_iv) - Set_client_DH_params_answer = self.method_call('set_client_DH_params', + for i in range(1, self.AUTH_MAX_RETRY): # retry when dh_gen_retry or dh_gen_fail + Set_client_DH_params_answer = self.method_call('set_client_DH_params', nonce=nonce, server_nonce=server_nonce, encrypted_data=encrypted_data) - auth_key = pow(g_a, b, dh_prime) - auth_key_str = long_to_bytes(auth_key) - auth_key_sha = SHA.new(auth_key_str).digest() - auth_key_aux_hash = auth_key_sha[:8] + # print Set_client_DH_params_answer + auth_key = pow(g_a, b, dh_prime) + auth_key_str = long_to_bytes(auth_key) + auth_key_sha = SHA.new(auth_key_str).digest() + auth_key_aux_hash = auth_key_sha[:8] - new_nonce_hash1 = SHA.new(new_nonce+b'\x01'+auth_key_aux_hash).digest()[-16:] - new_nonce_hash2 = SHA.new(new_nonce+b'\x02'+auth_key_aux_hash).digest()[-16:] - new_nonce_hash3 = SHA.new(new_nonce+b'\x03'+auth_key_aux_hash).digest()[-16:] + new_nonce_hash1 = SHA.new(new_nonce+b'\x01'+auth_key_aux_hash).digest()[-16:] + new_nonce_hash2 = SHA.new(new_nonce+b'\x02'+auth_key_aux_hash).digest()[-16:] + new_nonce_hash3 = SHA.new(new_nonce+b'\x03'+auth_key_aux_hash).digest()[-16:] - assert Set_client_DH_params_answer['nonce'] == nonce - assert Set_client_DH_params_answer['server_nonce'] == server_nonce - assert Set_client_DH_params_answer['new_nonce_hash1'] == new_nonce_hash1 - print("Diffie Hellman key exchange processed successfully") + assert Set_client_DH_params_answer['nonce'] == nonce + assert Set_client_DH_params_answer['server_nonce'] == server_nonce - self.server_salt = strxor(new_nonce[0:8], server_nonce[0:8]) - self.auth_key = auth_key_str - self.auth_key_id = auth_key_sha[-8:] - print("Auth key generated") + if Set_client_DH_params_answer['name'] == 'dh_gen_ok': + assert Set_client_DH_params_answer['new_nonce_hash1'] == new_nonce_hash1 + print("Diffie Hellman key exchange processed successfully") + + self.server_salt = strxor(new_nonce[0:8], server_nonce[0:8]) + self.auth_key = auth_key_str + self.auth_key_id = auth_key_sha[-8:] + print("Auth key generated") + return "Auth Ok" + elif Set_client_DH_params_answer['name'] == 'dh_gen_retry': + assert Set_client_DH_params_answer['new_nonce_hash2'] == new_nonce_hash2 + print ("Retry Auth") + elif Set_client_DH_params_answer['name'] == 'dh_gen_fail': + assert Set_client_DH_params_answer['new_nonce_hash3'] == new_nonce_hash3 + print("Auth Failed") + raise Exception("Auth Failed") + else: raise Exception("Response Error") def aes_calculate(self, msg_key, direction="to server"): x = 0 if direction == "to server" else 8 diff --git a/testing.py b/testing.py index 707dfcea..a73a4011 100644 --- a/testing.py +++ b/testing.py @@ -1,5 +1,6 @@ # -*- coding: utf-8 -*- import os +import time import io import struct # Deal with py2 and py3 differences @@ -20,7 +21,9 @@ ip = config.get('App data', 'ip_address') port = config.getint('App data', 'port') Session = mtproto.Session(ip, port) + +# print("start create") Session.create_auth_key() future_salts = Session.method_call('get_future_salts', num=3) -print(future_salts) \ No newline at end of file +print(future_salts) diff --git a/tests/auth_key_test.py b/tests/auth_key_test.py new file mode 100644 index 00000000..969b3ea6 --- /dev/null +++ b/tests/auth_key_test.py @@ -0,0 +1,21 @@ +import unittest +import os +import io + +try: + import configparser +except ImportError: + import ConfigParser as configparser + +import mtproto + +class TestAuthKeyCase(unittest.TestCase): + def setUp(self): + + + def tearDown(self): + pass + + def test_socket_timeout(self): + + # def t