✨ Crypto-MTProto: Improved pack and unpack function

pyrogram/crypto/mtproto.py

@@ -43,29 +43,34 @@ def kdf(auth_key: bytes, msg_key: bytes, outgoing: bool) -> tuple:
     return aes_key, aes_iv
 
 
-def pack(message: Message, salt: int, session_id: bytes, auth_key: bytes, auth_key_id: bytes) -> bytes:
+def pack(message: Message, salt: int, session_id: bytes, auth_key: bytes, auth_key_id: bytes, outgoing=True) -> bytes:
     data = Long(salt) + session_id + message.write()
     padding = urandom(-(len(data) + 12) % 16 + 12)
 
-    # 88 = 88 + 0 (outgoing message)
+    # 88 = 88 + 0 (outgoing message [Client])
-    msg_key_large = sha256(auth_key[88: 88 + 32] + data + padding).digest()
+    # 96 = 88 + 8 (incoming message [Server])
+    x = 0 if outgoing else 8
+    msg_key_large = sha256(auth_key[88 + x: 88 + x + 32] + data + padding).digest()
     msg_key = msg_key_large[8:24]
-    aes_key, aes_iv = kdf(auth_key, msg_key, True)
+    aes_key, aes_iv = kdf(auth_key, msg_key, outgoing)
 
     return auth_key_id + msg_key + aes.ige256_encrypt(data + padding, aes_key, aes_iv)
 
 
+# noinspection DuplicatedCode
 def unpack(
-    b: BytesIO,
+        b: BytesIO,
-    session_id: bytes,
+        session_id: bytes,
-    auth_key: bytes,
+        auth_key: bytes,
-    auth_key_id: bytes,
+        auth_key_id: bytes,
-    stored_msg_ids: List[int]
+        stored_msg_ids: List[int],
+        outgoing=False,
+        skip_security=False
 ) -> Message:
     SecurityCheckMismatch.check(b.read(8) == auth_key_id)
 
     msg_key = b.read(16)
-    aes_key, aes_iv = kdf(auth_key, msg_key, False)
+    aes_key, aes_iv = kdf(auth_key, msg_key, outgoing)
     data = BytesIO(aes.ige256_decrypt(b.read(), aes_key, aes_iv))
     data.read(8)  # Salt
 
@@ -87,8 +92,10 @@ def unpack(
         raise ValueError(f"The server sent an unknown constructor: {hex(e.args[0])}\n{left}")
 
     # https://core.telegram.org/mtproto/security_guidelines#checking-sha256-hash-value-of-msg-key
-    # 96 = 88 + 8 (incoming message)
+    # 88 = 88 + 0 (outgoing message [Client])
-    SecurityCheckMismatch.check(msg_key == sha256(auth_key[96:96 + 32] + data.getvalue()).digest()[8:24])
+    # 96 = 88 + 8 (incoming message [Server])
+    x = 0 if outgoing else 8
+    SecurityCheckMismatch.check(msg_key == sha256(auth_key[88 + x:88 + x + 32] + data.getvalue()).digest()[8:24])
 
     # https://core.telegram.org/mtproto/security_guidelines#checking-message-length
     data.seek(32)  # Get to the payload, skip salt (8) + session_id (8) + msg_id (8) + seq_no (4) + length (4)
@@ -98,30 +105,31 @@ def unpack(
     SecurityCheckMismatch.check(len(payload) % 4 == 0)
 
     # https://core.telegram.org/mtproto/security_guidelines#checking-msg-id
-    SecurityCheckMismatch.check(message.msg_id % 2 != 0)
+    if outgoing is False:
+        SecurityCheckMismatch.check(message.msg_id % 2 != 0)
 
-    if len(stored_msg_ids) > STORED_MSG_IDS_MAX_SIZE:
+        if len(stored_msg_ids) > STORED_MSG_IDS_MAX_SIZE:
-        del stored_msg_ids[:STORED_MSG_IDS_MAX_SIZE // 2]
+            del stored_msg_ids[:STORED_MSG_IDS_MAX_SIZE // 2]
 
-    if stored_msg_ids:
+        if stored_msg_ids:
-        # Ignored message: msg_id is lower than all of the stored values
+            # Ignored message: msg_id is lower than all the stored values
-        if message.msg_id < stored_msg_ids[0]:
+            if message.msg_id < stored_msg_ids[0]:
-            raise SecurityCheckMismatch
+                raise SecurityCheckMismatch
 
-        # Ignored message: msg_id is equal to any of the stored values
+            # Ignored message: msg_id is equal to any of the stored values
-        if message.msg_id in stored_msg_ids:
+            if message.msg_id in stored_msg_ids:
-            raise SecurityCheckMismatch
+                raise SecurityCheckMismatch
 
-        time_diff = (message.msg_id - MsgId()) / 2 ** 32
+            time_diff = (message.msg_id - MsgId()) / 2 ** 32
 
-        # Ignored message: msg_id belongs over 30 seconds in the future
+            # Ignored message: msg_id belongs over 30 seconds in the future
-        if time_diff > 30:
+            if time_diff > 30 and not skip_security:
-            raise SecurityCheckMismatch
+                raise SecurityCheckMismatch
 
-        # Ignored message: msg_id belongs over 300 seconds in the past
+            # Ignored message: msg_id belongs over 300 seconds in the past
-        if time_diff < -300:
+            if time_diff < -300:
-            raise SecurityCheckMismatch
+                raise SecurityCheckMismatch
 
-    bisect.insort(stored_msg_ids, message.msg_id)
+        bisect.insort(stored_msg_ids, message.msg_id)
 
     return message