From d2ac7e5b0a760f1dddcefc4178a24805bc4efb30 Mon Sep 17 00:00:00 2001
From: Lonami Exo <totufals@hotmail.com>
Date: Mon, 24 Dec 2018 14:16:50 +0100
Subject: [PATCH] Actually perform all checks in 2fa

---
 telethon/crypto/factorization.py  |  3 ++-
 telethon/network/authenticator.py |  2 +-
 telethon/password.py              | 12 +++++++++---
 3 files changed, 12 insertions(+), 5 deletions(-)

diff --git a/telethon/crypto/factorization.py b/telethon/crypto/factorization.py
index 75df4dd9..d06d270f 100644
--- a/telethon/crypto/factorization.py
+++ b/telethon/crypto/factorization.py
@@ -47,7 +47,8 @@ class Factorization:
                 if g > 1:
                     break
 
-        return g, pq // g
+        p, q = g, pq // g
+        return (p, q) if p < q else (q, p)
 
     @staticmethod
     def gcd(a, b):
diff --git a/telethon/network/authenticator.py b/telethon/network/authenticator.py
index 68006abc..bdf0f664 100644
--- a/telethon/network/authenticator.py
+++ b/telethon/network/authenticator.py
@@ -38,7 +38,7 @@ async def do_authentication(sender):
 
     # Step 2 sending: DH Exchange
     p, q = Factorization.factorize(pq)
-    p, q = rsa.get_byte_array(min(p, q)), rsa.get_byte_array(max(p, q))
+    p, q = rsa.get_byte_array(p), rsa.get_byte_array(q)
     new_nonce = int.from_bytes(os.urandom(32), 'little', signed=True)
 
     pq_inner_data = bytes(PQInnerData(
diff --git a/telethon/password.py b/telethon/password.py
index 36050216..17d6438d 100644
--- a/telethon/password.py
+++ b/telethon/password.py
@@ -1,6 +1,7 @@
 import hashlib
 import os
 
+from .crypto import factorization
 from .tl import types
 
 
@@ -10,7 +11,9 @@ def check_prime_and_good_check(prime: int, g: int):
         raise ValueError('bad prime count {}, expected {}'
                          .format(prime.bit_length(), good_prime_bits_count))
 
-    # TODO if not is_prime(prime) raise
+    if factorization.Factorization.factorize(prime)[0] != 1:
+        raise ValueError('given "prime" is not prime')
+
     if g == 2:
         if prime % 8 != 7:
             raise ValueError('bad g {}, mod8 {}'.format(g, prime % 8))
@@ -32,8 +35,10 @@ def check_prime_and_good_check(prime: int, g: int):
         raise ValueError('bad g {}'.format(g))
 
     prime_sub1_div2 = (prime - 1) // 2
-    # TODO if not is_prime(prime_sub1_div2) raise
-    # It's good
+    if factorization.Factorization.factorize(prime_sub1_div2)[0] != 1:
+        raise ValueError('(prime - 1) // 2 is not prime')
+
+    # Else it's good
 
 
 def check_prime_and_good(prime_bytes: bytes, g: int):
@@ -112,6 +117,7 @@ def compute_hash(algo: types.PasswordKdfAlgoSHA256SHA256PBKDF2HMACSHA512iter1000
     return sha256(algo.salt2, hash3, algo.salt2)
 
 
+# https://github.com/telegramdesktop/tdesktop/blob/18b74b90451a7db2379a9d753c9cbaf8734b4d5d/Telegram/SourceFiles/core/core_cloud_password.cpp
 def compute_check(request: types.account.Password, password: str):
     algo = request.current_algo
     if not isinstance(algo, types.PasswordKdfAlgoSHA256SHA256PBKDF2HMACSHA512iter100000SHA256ModPow):