diff --git a/pyrogram/crypto/mtproto.py b/pyrogram/crypto/mtproto.py index c893e3e5..40e433c1 100644 --- a/pyrogram/crypto/mtproto.py +++ b/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) - msg_key_large = sha256(auth_key[88: 88 + 32] + data + padding).digest() + # 88 = 88 + 0 (outgoing message [Client]) + # 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, - session_id: bytes, - auth_key: bytes, - auth_key_id: bytes, - stored_msg_ids: List[int] + b: BytesIO, + session_id: bytes, + auth_key: bytes, + auth_key_id: bytes, + 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) - SecurityCheckMismatch.check(msg_key == sha256(auth_key[96:96 + 32] + data.getvalue()).digest()[8:24]) + # 88 = 88 + 0 (outgoing message [Client]) + # 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: - del stored_msg_ids[:STORED_MSG_IDS_MAX_SIZE // 2] + if len(stored_msg_ids) > STORED_MSG_IDS_MAX_SIZE: + del stored_msg_ids[:STORED_MSG_IDS_MAX_SIZE // 2] - if stored_msg_ids: - # Ignored message: msg_id is lower than all of the stored values - if message.msg_id < stored_msg_ids[0]: - raise SecurityCheckMismatch + if stored_msg_ids: + # Ignored message: msg_id is lower than all the stored values + if message.msg_id < stored_msg_ids[0]: + raise SecurityCheckMismatch - # Ignored message: msg_id is equal to any of the stored values - if message.msg_id in stored_msg_ids: - raise SecurityCheckMismatch + # Ignored message: msg_id is equal to any of the stored values + if message.msg_id in stored_msg_ids: + 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 - if time_diff > 30: - raise SecurityCheckMismatch + # Ignored message: msg_id belongs over 30 seconds in the future + if time_diff > 30 and not skip_security: + raise SecurityCheckMismatch - # Ignored message: msg_id belongs over 300 seconds in the past - if time_diff < -300: - raise SecurityCheckMismatch + # Ignored message: msg_id belongs over 300 seconds in the past + if time_diff < -300: + raise SecurityCheckMismatch - bisect.insort(stored_msg_ids, message.msg_id) + bisect.insort(stored_msg_ids, message.msg_id) return message