Move MTProto related methods into a separate module

pyrogram/crypto/__init__.py

@@ -18,5 +18,6 @@
 
 from .aes import AES
 from .kdf import KDF
+from .mtproto import MTProto
 from .prime import Prime
 from .rsa import RSA

pyrogram/crypto/mtproto.py

@@ -0,0 +1,65 @@
+# Pyrogram - Telegram MTProto API Client Library for Python
+# Copyright (C) 2017-2018 Dan Tès <https://github.com/delivrance>
+#
+# This file is part of Pyrogram.
+#
+# Pyrogram is free software: you can redistribute it and/or modify
+# it under the terms of the GNU Lesser General Public License as published
+# by the Free Software Foundation, either version 3 of the License, or
+# (at your option) any later version.
+#
+# Pyrogram is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+# GNU Lesser General Public License for more details.
+#
+# 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
+from io import BytesIO
+from os import urandom
+
+from pyrogram.api.core import Message, Long
+from . import AES, KDF
+
+
+class MTProto:
+    INITIAL_SALT = 0x616e67656c696361
+
+    @staticmethod
+    def pack(message: Message, salt: int, session_id: bytes, auth_key: bytes, auth_key_id: bytes) -> bytes:
+        data = Long(salt) + session_id + message.write()
+        padding = urandom(-(len(data) + 12) % 16 + 12)
+
+        # 88 = 88 + 0 (outgoing message)
+        msg_key_large = sha256(auth_key[88: 88 + 32] + data + padding).digest()
+        msg_key = msg_key_large[8:24]
+        aes_key, aes_iv = KDF(auth_key, msg_key, True)
+
+        return auth_key_id + msg_key + AES.ige256_encrypt(data + padding, aes_key, aes_iv)
+
+    @staticmethod
+    def unpack(b: BytesIO, salt: int, session_id: bytes, auth_key: bytes, auth_key_id: bytes) -> Message:
+        assert b.read(8) == auth_key_id, b.getvalue()
+
+        msg_key = b.read(16)
+        aes_key, aes_iv = KDF(auth_key, msg_key, False)
+        data = BytesIO(AES.ige256_decrypt(b.read(), aes_key, aes_iv))
+
+        assert data.read(8) == Long(salt) or Long(salt) == Long(MTProto.INITIAL_SALT)
+
+        # https://core.telegram.org/mtproto/security_guidelines#checking-session-id
+        assert data.read(8) == session_id
+
+        message = Message.read(data)
+
+        # https://core.telegram.org/mtproto/security_guidelines#checking-sha256-hash-value-of-msg-key
+        # https://core.telegram.org/mtproto/security_guidelines#checking-message-length
+        # 96 = 88 + 8 (incoming message)
+        assert msg_key == sha256(auth_key[96:96 + 32] + data.getvalue()).digest()[8:24]
+
+        # https://core.telegram.org/mtproto/security_guidelines#checking-msg-id
+        assert message.msg_id % 2 != 0
+
+        return message