Add extra security checks during authkey gen

telethon/network/authenticator.py

@@ -126,19 +126,42 @@ async def do_authentication(sender):
         raise SecurityError('Step 3 Invalid server nonce in encrypted answer')
 
     dh_prime = get_int(server_dh_inner.dh_prime, signed=False)
+    g = server_dh_inner.g
     g_a = get_int(server_dh_inner.g_a, signed=False)
     time_offset = server_dh_inner.server_time - int(time.time())
 
     b = get_int(os.urandom(256), signed=False)
-    gb = pow(server_dh_inner.g, b, dh_prime)
+    g_b = pow(g, b, dh_prime)
     gab = pow(g_a, b, dh_prime)
 
+    # IMPORTANT: Apart from the conditions on the Diffie-Hellman prime
+    # dh_prime and generator g, both sides are to check that g, g_a and
+    # g_b are greater than 1 and less than dh_prime - 1. We recommend
+    # checking that g_a and g_b are between 2^{2048-64} and
+    # dh_prime - 2^{2048-64} as well.
+    # (https://core.telegram.org/mtproto/auth_key#dh-key-exchange-complete)
+    if not (1 < g < (dh_prime - 1)):
+        raise SecurityError('g_a is not within (1, dh_prime - 1)')
+
+    if not (1 < g_a < (dh_prime - 1)):
+        raise SecurityError('g_a is not within (1, dh_prime - 1)')
+
+    if not (1 < g_b < (dh_prime - 1)):
+        raise SecurityError('g_b is not within (1, dh_prime - 1)')
+
+    safety_range = 2 ** (2048 - 64)
+    if not (safety_range <= g_a <= (dh_prime - safety_range)):
+        raise SecurityError('g_a is not within (2^{2048-64}, dh_prime - 2^{2048-64})')
+
+    if not (safety_range <= g_b <= (dh_prime - safety_range)):
+        raise SecurityError('g_b is not within (2^{2048-64}, dh_prime - 2^{2048-64})')
+
     # Prepare client DH Inner Data
     client_dh_inner = bytes(ClientDHInnerData(
         nonce=res_pq.nonce,
         server_nonce=res_pq.server_nonce,
         retry_id=0,  # TODO Actual retry ID
-        g_b=rsa.get_byte_array(gb)
+        g_b=rsa.get_byte_array(g_b)
     ))
 
     client_dh_inner_hashed = sha1(client_dh_inner).digest() + client_dh_inner