Accommodate TgCrypto

pyrogram/crypto/ige.py

@@ -16,7 +16,8 @@
 # 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 pyaes import AES
+# from pyaes import AES
+import tgcrypto
 
 BLOCK_SIZE = 16
 
@@ -26,11 +27,13 @@ BLOCK_SIZE = 16
 class IGE:
     @classmethod
     def encrypt(cls, data: bytes, key: bytes, iv: bytes) -> bytes:
-        return cls.ige(data, key, iv, True)
+        return tgcrypto.ige_encrypt(data, key, iv)
+        # return cls.ige(data, key, iv, True)
 
     @classmethod
     def decrypt(cls, data: bytes, key: bytes, iv: bytes) -> bytes:
-        return cls.ige(data, key, iv, False)
+        return tgcrypto.ige_decrypt(data, key, iv)
+        # return cls.ige(data, key, iv, False)
 
     @staticmethod
     def xor(a: bytes, b: bytes) -> bytes:
@@ -40,22 +43,22 @@ class IGE:
             "big",
         )
 
-    @classmethod
-    def ige(cls, data: bytes, key: bytes, iv: bytes, encrypt: bool) -> bytes:
-        cipher = AES(key)
-
-        iv_1 = iv[:BLOCK_SIZE]
-        iv_2 = iv[BLOCK_SIZE:]
-
-        data = [data[i: i + BLOCK_SIZE] for i in range(0, len(data), BLOCK_SIZE)]
-
-        if encrypt:
-            for i, chunk in enumerate(data):
-                iv_1 = data[i] = cls.xor(cipher.encrypt(cls.xor(chunk, iv_1)), iv_2)
-                iv_2 = chunk
-        else:
-            for i, chunk in enumerate(data):
-                iv_2 = data[i] = cls.xor(cipher.decrypt(cls.xor(chunk, iv_2)), iv_1)
-                iv_1 = chunk
-
-        return b"".join(data)
+    # @classmethod
+    # def ige(cls, data: bytes, key: bytes, iv: bytes, encrypt: bool) -> bytes:
+    #     cipher = AES(key)
+    #
+    #     iv_1 = iv[:BLOCK_SIZE]
+    #     iv_2 = iv[BLOCK_SIZE:]
+    #
+    #     data = [data[i: i + BLOCK_SIZE] for i in range(0, len(data), BLOCK_SIZE)]
+    #
+    #     if encrypt:
+    #         for i, chunk in enumerate(data):
+    #             iv_1 = data[i] = cls.xor(cipher.encrypt(cls.xor(chunk, iv_1)), iv_2)
+    #             iv_2 = chunk
+    #     else:
+    #         for i, chunk in enumerate(data):
+    #             iv_2 = data[i] = cls.xor(cipher.decrypt(cls.xor(chunk, iv_2)), iv_1)
+    #             iv_1 = chunk
+    #
+    #     return b"".join(data)

setup.py

@@ -67,6 +67,6 @@ setup(
     ],
     packages=find_packages(),
     zip_safe=False,
-    install_requires=["pyaes", "pysocks"],
+    install_requires=["pyaes", "pysocks", "tgcrypto"],
     include_package_data=True,
 )