Re-implement enable_cloud_password using SRP

pyrogram/client/methods/password/enable_cloud_password.py

@@ -16,34 +16,18 @@
 # You should have received a copy of the GNU Lesser General Public License
 # along with Pyrogram.  If not, see <http://www.gnu.org/licenses/>.
 
-from hashlib import sha256, pbkdf2_hmac
 import os
 
 from pyrogram.api import functions, types
+from .utils import compute_hash, btoi, itob
 from ...ext import BaseClient
 
 
-def compute_hash(algo: types.PasswordKdfAlgoSHA256SHA256PBKDF2HMACSHA512iter100000SHA256ModPow, password: str):
-    hash1 = sha256(algo.salt1 + password.encode() + algo.salt1).digest()
-    hash2 = sha256(algo.salt2 + hash1 + algo.salt2).digest()
-    hash3 = pbkdf2_hmac("sha512", hash2, algo.salt1, 100000)
-
-    return sha256(algo.salt2 + hash3 + algo.salt2).digest()
-
-
-def btoi(b: bytes):
-    return int.from_bytes(b, "big")
-
-
-def itob(i: int):
-    return i.to_bytes(256, "big")
-
-
 class EnableCloudPassword(BaseClient):
-    def enable_cloud_password(self, password: str, hint: str = "", email: str = ""):
+    def enable_cloud_password(self, password: str, hint: str = "", email: str = None):
         """Use this method to enable the Two-Step Verification security feature (Cloud Password) on your account.
 
-        This password will be asked when you log in on a new device in addition to the SMS code.
+        This password will be asked when you log-in on a new device in addition to the SMS code.
 
         Args:
             password (``str``):
@@ -56,92 +40,31 @@ class EnableCloudPassword(BaseClient):
                 Recovery e-mail.
 
         Returns:
-            True on success, False otherwise.
+            True on success.
 
         Raises:
             :class:`Error <pyrogram.Error>` in case of a Telegram RPC error.
+            ``ValueError`` in case there is already a cloud password enabled.
         """
         r = self.send(functions.account.GetPassword())
-        print(r)
 
-        algo = r.new_algo
+        if r.has_password:
+            raise ValueError("There is already a cloud password enabled")
 
-        p_bytes = algo.p
-        p = btoi(algo.p)
+        r.new_algo.salt1 += os.urandom(32)
+        new_hash = btoi(compute_hash(r.new_algo, password))
+        new_hash = itob(pow(r.new_algo.g, new_hash, btoi(r.new_algo.p)))
 
-        g_bytes = itob(algo.g)
-        g = algo.g
-
-        B_bytes = r.srp_B or b""
-        B = btoi(B_bytes)
-
-        srp_id = r.srp_id or 0
-
-        x_bytes = compute_hash(algo, password)
-        x = btoi(x_bytes)
-
-        g_x = pow(g, x, p)
-
-        k_bytes = sha256(p_bytes + g_bytes).digest()
-        k = btoi(k_bytes)
-
-        kg_x = (k * g_x) % p
-
-        while True:
-            a_bytes = os.urandom(256)
-            a = btoi(a_bytes)
-
-            A = pow(g, a, p)
-            A_bytes = itob(A)
-
-            u = btoi(sha256(A_bytes + B_bytes).digest())
-
-            if u > 0:
-                break
-
-        g_b = (B - kg_x) % p
-
-        ux = u * x
-        a_ux = a + ux
-        S = pow(g_b, a_ux, p)
-        S_bytes = itob(S)
-
-        K_bytes = sha256(S_bytes).digest()
-        M1_bytes = sha256(
-            b"".join([bytes([int(i) ^ int(j)]) for (i, j) in zip(sha256(p_bytes).digest(), sha256(g_bytes).digest())])
-            + sha256(algo.salt1).digest()
-            + sha256(algo.salt2).digest()
-            + A_bytes
-            + B_bytes
-            + K_bytes
-        ).digest()
-
-        input_check_password = types.InputCheckPasswordSRP(srp_id, A_bytes, M1_bytes)
-
-        r2 = self.send(functions.account.UpdatePasswordSettings(
-            input_check_password, types.account.PasswordInputSettings(
-                new_algo=algo,
-                new_password_hash=b"",
-                hint=""
+        self.send(
+            functions.account.UpdatePasswordSettings(
+                password=types.InputCheckPasswordEmpty(),
+                new_settings=types.account.PasswordInputSettings(
+                    new_algo=r.new_algo,
+                    new_password_hash=new_hash,
+                    hint=hint,
+                    email=email
+                )
             )
-        ))
+        )
 
-        print(r2)
-
-        # if isinstance(r, types.account.NoPassword):
-        #     salt = r.new_salt + os.urandom(8)
-        #     password_hash = sha256(salt + password.encode() + salt).digest()
-        #
-        #     return self.send(
-        #         functions.account.UpdatePasswordSettings(
-        #             current_password_hash=salt,
-        #             new_settings=types.account.PasswordInputSettings(
-        #                 new_salt=salt,
-        #                 new_password_hash=password_hash,
-        #                 hint=hint,
-        #                 email=email
-        #             )
-        #         )
-        #     )
-        # else:
-        #     return False
+        return True