3232from pyrogram .api .core import Message , Object , MsgContainer , Long , FutureSalt , Int
3333from pyrogram .api .errors import Error
3434from pyrogram .connection import Connection
35- from pyrogram .crypto import IGE , KDF2
35+ from pyrogram .crypto import IGE , KDF
3636from .internals import MsgId , MsgFactory , DataCenter
3737
3838log = logging .getLogger (__name__ )
@@ -178,15 +178,7 @@ def restart(self):
178178 self .stop ()
179179 self .start ()
180180
181- # def pack(self, message: Message) -> bytes:
182- # data = Long(self.current_salt.salt) + self.session_id + message.write()
183- # msg_key = sha1(data).digest()[-16:]
184- # aes_key, aes_iv = KDF(self.auth_key, msg_key, True)
185- # padding = urandom(-len(data) % 16)
186- #
187- # return self.auth_key_id + msg_key + IGE.encrypt(data + padding, aes_key, aes_iv)
188-
189- def pack2 (self , message : Message ):
181+ def pack (self , message : Message ):
190182 data = Long (self .current_salt .salt ) + self .session_id + message .write ()
191183 # MTProto 2.0 requires a minimum of 12 padding bytes.
192184 # I don't get why it says up to 1024 when what it actually needs after the
@@ -197,39 +189,15 @@ def pack2(self, message: Message):
197189 # 88 = 88 + 0 (outgoing message)
198190 msg_key_large = sha256 (self .auth_key [88 : 88 + 32 ] + data + padding ).digest ()
199191 msg_key = msg_key_large [8 :24 ]
200- aes_key , aes_iv = KDF2 (self .auth_key , msg_key , True )
192+ aes_key , aes_iv = KDF (self .auth_key , msg_key , True )
201193
202194 return self .auth_key_id + msg_key + IGE .encrypt (data + padding , aes_key , aes_iv )
203195
204- # def unpack(self, b: BytesIO) -> Message:
205- # assert b.read(8) == self.auth_key_id, b.getvalue()
206- #
207- # msg_key = b.read(16)
208- # aes_key, aes_iv = KDF(self.auth_key, msg_key, False)
209- # data = BytesIO(IGE.decrypt(b.read(), aes_key, aes_iv))
210- # data.read(8) # Server salt
211- #
212- # # https://core.telegram.org/mtproto/security_guidelines#checking-session-id
213- # assert data.read(8) == self.session_id
214- #
215- # message = Message.read(data)
216- #
217- # # https://core.telegram.org/mtproto/security_guidelines#checking-sha1-hash-value-of-msg-key
218- # # https://core.telegram.org/mtproto/security_guidelines#checking-message-length
219- # # 32 = salt (8) + session_id (8) + msg_id (8) + seq_no (4) + length (4)
220- # assert msg_key == sha1(data.getvalue()[:32 + message.length]).digest()[-16:]
221- #
222- # # https://core.telegram.org/mtproto/security_guidelines#checking-msg-id
223- # # TODO: check for lower msg_ids
224- # assert message.msg_id % 2 != 0
225- #
226- # return message
227-
228- def unpack2 (self , b : BytesIO ) -> Message :
196+ def unpack (self , b : BytesIO ) -> Message :
229197 assert b .read (8 ) == self .auth_key_id , b .getvalue ()
230198
231199 msg_key = b .read (16 )
232- aes_key , aes_iv = KDF2 (self .auth_key , msg_key , False )
200+ aes_key , aes_iv = KDF (self .auth_key , msg_key , False )
233201 data = BytesIO (IGE .decrypt (b .read (), aes_key , aes_iv ))
234202 data .read (8 )
235203
@@ -268,7 +236,7 @@ def worker(self):
268236
269237 def unpack_dispatch_and_ack (self , packet : bytes ):
270238 # TODO: A better dispatcher
271- data = self .unpack2 (BytesIO (packet ))
239+ data = self .unpack (BytesIO (packet ))
272240
273241 messages = (
274242 data .body .messages
@@ -398,7 +366,7 @@ def _send(self, data: Object, wait_response: bool = True):
398366 if wait_response :
399367 self .results [msg_id ] = Result ()
400368
401- payload = self .pack2 (message )
369+ payload = self .pack (message )
402370
403371 try :
404372 self .connection .send (payload )
0 commit comments