Rename AES class methods and add ctr256_encrypt

delivrance · delivrance · commit f0c801be8c75 · 2018-05-18T14:15:35.000+02:00
diff --git a/pyrogram/crypto/aes.py b/pyrogram/crypto/aes.py
@@ -38,31 +38,32 @@
 # TODO: Ugly IFs
 class AES:
     @classmethod
-    def ige_encrypt(cls, data: bytes, key: bytes, iv: bytes) -> bytes:
+    def ige256_encrypt(cls, data: bytes, key: bytes, iv: bytes) -> bytes:
         if is_fast:
-            return tgcrypto.ige_encrypt(data, key, iv)
+            return tgcrypto.ige256_encrypt(data, key, iv)
         else:
             return cls.ige(data, key, iv, True)
 
     @classmethod
-    def ige_decrypt(cls, data: bytes, key: bytes, iv: bytes) -> bytes:
+    def ige256_decrypt(cls, data: bytes, key: bytes, iv: bytes) -> bytes:
         if is_fast:
-            return tgcrypto.ige_decrypt(data, key, iv)
+            return tgcrypto.ige256_decrypt(data, key, iv)
         else:
             return cls.ige(data, key, iv, False)
 
     @staticmethod
-    def ctr_decrypt(data: bytes, key: bytes, iv: bytes, offset: int) -> bytes:
-        replace = int.to_bytes(offset // 16, 4, "big")
-        iv = iv[:-4] + replace
-
+    def ctr256_encrypt(data: bytes, key: bytes, iv: bytes, state: bytes) -> bytes:
         if is_fast:
-            return tgcrypto.ctr_decrypt(data, key, iv)
+            return tgcrypto.ctr256_decrypt(data, key, iv, state)
         else:
             ctr = pyaes.AESModeOfOperationCTR(key)
             ctr._counter._counter = list(iv)
             return ctr.decrypt(data)
 
+    @staticmethod
+    def ctr256_decrypt(data: bytes, key: bytes, iv: bytes, state: bytes) -> bytes:
+        return AES.ctr256_encrypt(data, key, iv, state)
+
     @staticmethod
     def xor(a: bytes, b: bytes) -> bytes:
         return int.to_bytes(
diff --git a/pyrogram/session/auth.py b/pyrogram/session/auth.py
@@ -163,7 +163,7 @@ def create(self):
 
                 server_nonce = int.from_bytes(server_nonce, "little", signed=True)
 
-                answer_with_hash = AES.ige_decrypt(encrypted_answer, tmp_aes_key, tmp_aes_iv)
+                answer_with_hash = AES.ige256_decrypt(encrypted_answer, tmp_aes_key, tmp_aes_iv)
                 answer = answer_with_hash[20:]
 
                 server_dh_inner_data = Object.read(BytesIO(answer))
@@ -192,7 +192,7 @@ def create(self):
                 sha = sha1(data).digest()
                 padding = urandom(- (len(data) + len(sha)) % 16)
                 data_with_hash = sha + data + padding
-                encrypted_data = AES.ige_encrypt(data_with_hash, tmp_aes_key, tmp_aes_iv)
+                encrypted_data = AES.ige256_encrypt(data_with_hash, tmp_aes_key, tmp_aes_iv)
 
                 log.debug("Send set_client_DH_params")
                 set_client_dh_params_answer = self.send(
diff --git a/pyrogram/session/session.py b/pyrogram/session/session.py
@@ -222,14 +222,14 @@ def pack(self, message: Message):
         msg_key = msg_key_large[8:24]
         aes_key, aes_iv = KDF(self.auth_key, msg_key, True)
 
-        return self.auth_key_id + msg_key + AES.ige_encrypt(data + padding, aes_key, aes_iv)
+        return self.auth_key_id + msg_key + AES.ige256_encrypt(data + padding, aes_key, aes_iv)
 
     def unpack(self, b: BytesIO) -> Message:
         assert b.read(8) == self.auth_key_id, b.getvalue()
 
         msg_key = b.read(16)
         aes_key, aes_iv = KDF(self.auth_key, msg_key, False)
-        data = BytesIO(AES.ige_decrypt(b.read(), aes_key, aes_iv))
+        data = BytesIO(AES.ige256_decrypt(b.read(), aes_key, aes_iv))
         data.read(8)
 
         # https://core.telegram.org/mtproto/security_guidelines#checking-session-id
