diff --git a/pyrogram/crypto/__init__.py b/pyrogram/crypto/__init__.py index 08ed44f0..3112729d 100644 --- a/pyrogram/crypto/__init__.py +++ b/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 diff --git a/pyrogram/crypto/mtproto.py b/pyrogram/crypto/mtproto.py new file mode 100644 index 00000000..d42caf7c --- /dev/null +++ b/pyrogram/crypto/mtproto.py @@ -0,0 +1,65 @@ +# Pyrogram - Telegram MTProto API Client Library for Python +# Copyright (C) 2017-2018 Dan Tès +# +# 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 . + +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