From f72ddbdd5adc2e0153d173764aef4f2ad5ca164a Mon Sep 17 00:00:00 2001
From: Lonami Exo <totufals@hotmail.com>
Date: Thu, 7 Jun 2018 17:20:45 +0200
Subject: [PATCH] Implement retry and fail cases in authenticator

---
 telethon/crypto/auth_key.py       |  2 +-
 telethon/network/authenticator.py | 55 +++++++++++++++----------------
 telethon/network/mtprotosender.py |  2 +-
 3 files changed, 29 insertions(+), 30 deletions(-)

diff --git a/telethon/crypto/auth_key.py b/telethon/crypto/auth_key.py
index a6c0675b..31d1fc9c 100644
--- a/telethon/crypto/auth_key.py
+++ b/telethon/crypto/auth_key.py
@@ -37,4 +37,4 @@ class AuthKey:
         data = new_nonce + struct.pack('<BQ', number, self.aux_hash)
 
         # Calculates the message key from the given data
-        return sha1(data).digest()[4:20]
+        return int.from_bytes(sha1(data).digest()[4:20], 'little', signed=True)
diff --git a/telethon/network/authenticator.py b/telethon/network/authenticator.py
index 02b1128d..82e96a61 100644
--- a/telethon/network/authenticator.py
+++ b/telethon/network/authenticator.py
@@ -72,11 +72,7 @@ async def do_authentication(sender):
         encrypted_data=cipher_text
     ))
 
-    if isinstance(server_dh_params, ServerDHParamsFail):
-        raise SecurityError('Server DH params fail: TODO')
-
-    if not isinstance(server_dh_params, ServerDHParamsOk):
-        raise AssertionError(server_dh_params)
+    assert isinstance(server_dh_params, (ServerDHParamsOk, ServerDHParamsFail))
 
     if server_dh_params.nonce != res_pq.nonce:
         raise SecurityError('Invalid nonce from server')
@@ -84,6 +80,16 @@ async def do_authentication(sender):
     if server_dh_params.server_nonce != res_pq.server_nonce:
         raise SecurityError('Invalid server nonce from server')
 
+    if isinstance(server_dh_params, ServerDHParamsFail):
+        nnh = int.from_bytes(
+            sha1(new_nonce.to_bytes(32, 'little', signed=True)).digest()[4:20],
+            'little', signed=True
+        )
+        if server_dh_params.new_nonce_hash != nnh:
+            raise SecurityError('Invalid DH fail nonce from server')
+
+    assert isinstance(server_dh_params, ServerDHParamsOk)
+
     # Step 3 sending: Complete DH Exchange
     key, iv = utils.generate_key_data_from_nonce(
         res_pq.server_nonce, new_nonce
@@ -99,8 +105,7 @@ async def do_authentication(sender):
     with BinaryReader(plain_text_answer) as reader:
         reader.read(20)  # hash sum
         server_dh_inner = reader.tgread_object()
-        if not isinstance(server_dh_inner, ServerDHInnerData):
-            raise AssertionError(server_dh_inner)
+        assert isinstance(server_dh_inner, ServerDHInnerData)
 
     if server_dh_inner.nonce != res_pq.nonce:
         raise SecurityError('Invalid nonce in encrypted answer')
@@ -136,31 +141,25 @@ async def do_authentication(sender):
         encrypted_data=client_dh_encrypted,
     ))
 
-    if isinstance(dh_gen, DhGenOk):
-        if dh_gen.nonce != res_pq.nonce:
-            raise SecurityError('Invalid nonce from server')
+    nonce_types = (DhGenOk, DhGenRetry, DhGenFail)
+    assert isinstance(dh_gen, nonce_types)
+    name = dh_gen.__class__.__name__
+    if dh_gen.nonce != res_pq.nonce:
+        raise SecurityError('Invalid {} nonce from server'.format(name))
 
-        if dh_gen.server_nonce != res_pq.server_nonce:
-            raise SecurityError('Invalid server nonce from server')
+    if dh_gen.server_nonce != res_pq.server_nonce:
+        raise SecurityError('Invalid {} server nonce from server'.format(name))
 
-        auth_key = AuthKey(rsa.get_byte_array(gab))
-        new_nonce_hash = int.from_bytes(
-            auth_key.calc_new_nonce_hash(new_nonce, 1), 'little', signed=True
-        )
+    auth_key = AuthKey(rsa.get_byte_array(gab))
+    nonce_number = 1 + nonce_types.index(type(dh_gen))
+    new_nonce_hash = auth_key.calc_new_nonce_hash(new_nonce, nonce_number)
 
-        if dh_gen.new_nonce_hash1 != new_nonce_hash:
-            raise SecurityError('Invalid new nonce hash')
+    dh_hash = getattr(dh_gen, 'new_nonce_hash{}'.format(nonce_number))
+    if dh_hash != new_nonce_hash:
+        raise SecurityError('Invalid new nonce hash')
 
-        return auth_key, time_offset
-
-    elif isinstance(dh_gen, DhGenRetry):
-        raise NotImplementedError('DhGenRetry')
-
-    elif isinstance(dh_gen, DhGenFail):
-        raise NotImplementedError('DhGenFail')
-
-    else:
-        raise NotImplementedError('DH Gen unknown: {}'.format(dh_gen))
+    assert isinstance(dh_gen, DhGenOk)
+    return auth_key, time_offset
 
 
 def get_int(byte_array, signed=True):
diff --git a/telethon/network/mtprotosender.py b/telethon/network/mtprotosender.py
index ae89cb01..9b10817d 100644
--- a/telethon/network/mtprotosender.py
+++ b/telethon/network/mtprotosender.py
@@ -105,7 +105,7 @@ class MTProtoSender:
             await self._connection.connect(ip, port)
         self._user_connected = True
 
-        # TODO Handle SecurityError, AssertionError, NotImplementedError
+        # TODO Handle SecurityError, AssertionError
         if self.session.auth_key is None:
             plain = MTProtoPlainSender(self._connection)
             self.session.auth_key, self.session.time_offset =\