diff --git a/pyrogram/crypto/__init__.py b/pyrogram/crypto/__init__.py index 0c8298c3..71f177ed 100644 --- a/pyrogram/crypto/__init__.py +++ b/pyrogram/crypto/__init__.py @@ -17,6 +17,6 @@ # along with Pyrogram. If not, see . from .ige import IGE -from .kdf import KDF2 +from .kdf import KDF from .prime import Prime from .rsa import RSA diff --git a/pyrogram/crypto/kdf.py b/pyrogram/crypto/kdf.py index b37a6de4..2a7efa6a 100644 --- a/pyrogram/crypto/kdf.py +++ b/pyrogram/crypto/kdf.py @@ -19,23 +19,7 @@ from hashlib import sha256 -# class KDF: -# def __new__(cls, auth_key: bytes, msg_key: bytes, outgoing: bool) -> tuple: -# # https://core.telegram.org/mtproto/description#defining-aes-key-and-initialization-vector -# x = 0 if outgoing else 8 -# -# sha1_a = sha1(msg_key + auth_key[x:x + 32]).digest() -# sha1_b = sha1(auth_key[x + 32:x + 48] + msg_key + auth_key[x + 48:x + 64]).digest() -# sha1_c = sha1(auth_key[x + 64:x + 96] + msg_key).digest() -# sha1_d = sha1(msg_key + auth_key[x + 96:x + 128]).digest() -# -# aes_key = sha1_a[:8] + sha1_b[8:20] + sha1_c[4:16] -# aes_iv = sha1_a[8:20] + sha1_b[:8] + sha1_c[16:20] + sha1_d[:8] -# -# return aes_key, aes_iv - - -class KDF2: +class KDF: def __new__(cls, auth_key: bytes, msg_key: bytes, outgoing: bool) -> tuple: # https://core.telegram.org/mtproto/description#defining-aes-key-and-initialization-vector x = 0 if outgoing else 8 diff --git a/pyrogram/session/session.py b/pyrogram/session/session.py index 1c90bc7a..1cd21014 100644 --- a/pyrogram/session/session.py +++ b/pyrogram/session/session.py @@ -32,7 +32,7 @@ from pyrogram.api.all import layer from pyrogram.api.core import Message, Object, MsgContainer, Long, FutureSalt, Int from pyrogram.api.errors import Error from pyrogram.connection import Connection -from pyrogram.crypto import IGE, KDF2 +from pyrogram.crypto import IGE, KDF from .internals import MsgId, MsgFactory, DataCenter log = logging.getLogger(__name__) @@ -178,15 +178,7 @@ class Session: self.stop() self.start() - # def pack(self, message: Message) -> bytes: - # data = Long(self.current_salt.salt) + self.session_id + message.write() - # msg_key = sha1(data).digest()[-16:] - # aes_key, aes_iv = KDF(self.auth_key, msg_key, True) - # padding = urandom(-len(data) % 16) - # - # return self.auth_key_id + msg_key + IGE.encrypt(data + padding, aes_key, aes_iv) - - def pack2(self, message: Message): + def pack(self, message: Message): data = Long(self.current_salt.salt) + self.session_id + message.write() # MTProto 2.0 requires a minimum of 12 padding bytes. # I don't get why it says up to 1024 when what it actually needs after the @@ -197,39 +189,15 @@ class Session: # 88 = 88 + 0 (outgoing message) msg_key_large = sha256(self.auth_key[88: 88 + 32] + data + padding).digest() msg_key = msg_key_large[8:24] - aes_key, aes_iv = KDF2(self.auth_key, msg_key, True) + aes_key, aes_iv = KDF(self.auth_key, msg_key, True) return self.auth_key_id + msg_key + IGE.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(IGE.decrypt(b.read(), aes_key, aes_iv)) - # data.read(8) # Server salt - # - # # https://core.telegram.org/mtproto/security_guidelines#checking-session-id - # assert data.read(8) == self.session_id - # - # message = Message.read(data) - # - # # https://core.telegram.org/mtproto/security_guidelines#checking-sha1-hash-value-of-msg-key - # # https://core.telegram.org/mtproto/security_guidelines#checking-message-length - # # 32 = salt (8) + session_id (8) + msg_id (8) + seq_no (4) + length (4) - # assert msg_key == sha1(data.getvalue()[:32 + message.length]).digest()[-16:] - # - # # https://core.telegram.org/mtproto/security_guidelines#checking-msg-id - # # TODO: check for lower msg_ids - # assert message.msg_id % 2 != 0 - # - # return message - - def unpack2(self, b: BytesIO) -> Message: + 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 = KDF2(self.auth_key, msg_key, False) + aes_key, aes_iv = KDF(self.auth_key, msg_key, False) data = BytesIO(IGE.decrypt(b.read(), aes_key, aes_iv)) data.read(8) @@ -268,7 +236,7 @@ class Session: def unpack_dispatch_and_ack(self, packet: bytes): # TODO: A better dispatcher - data = self.unpack2(BytesIO(packet)) + data = self.unpack(BytesIO(packet)) messages = ( data.body.messages @@ -398,7 +366,7 @@ class Session: if wait_response: self.results[msg_id] = Result() - payload = self.pack2(message) + payload = self.pack(message) try: self.connection.send(payload)