From 3eafe18d0b8bf535c140e845eae2e3a0be78701c Mon Sep 17 00:00:00 2001
From: Lonami Exo <totufals@hotmail.com>
Date: Sat, 6 Jan 2018 01:55:11 +0100
Subject: [PATCH] Implement MtProto 2.0 (closes #484, thanks @delivrance!)

Huge shoutout to @delivrance's pyrogram, specially this commit:
pyrogram/pyrogram/commit/42f9a2d6994baaf9ecad590d1ff4d175a8c56454
---
 telethon/extensions/binary_reader.py |  5 ++-
 telethon/helpers.py                  | 67 ++++++++++++++++++++++++++--
 telethon/network/mtproto_sender.py   | 44 +++---------------
 3 files changed, 75 insertions(+), 41 deletions(-)

diff --git a/telethon/extensions/binary_reader.py b/telethon/extensions/binary_reader.py
index 460bed96..1402083f 100644
--- a/telethon/extensions/binary_reader.py
+++ b/telethon/extensions/binary_reader.py
@@ -56,8 +56,11 @@ class BinaryReader:
         return int.from_bytes(
             self.read(bits // 8), byteorder='little', signed=signed)
 
-    def read(self, length):
+    def read(self, length=None):
         """Read the given amount of bytes."""
+        if length is None:
+            return self.reader.read()
+
         result = self.reader.read(length)
         if len(result) != length:
             raise BufferError(
diff --git a/telethon/helpers.py b/telethon/helpers.py
index 3c9af2cb..d97b8a9f 100644
--- a/telethon/helpers.py
+++ b/telethon/helpers.py
@@ -1,6 +1,11 @@
 """Various helpers not related to the Telegram API itself"""
-from hashlib import sha1, sha256
 import os
+import struct
+from hashlib import sha1, sha256
+
+from telethon.crypto import AES
+from telethon.extensions import BinaryReader
+
 
 # region Multiple utilities
 
@@ -21,9 +26,48 @@ def ensure_parent_dir_exists(file_path):
 # region Cryptographic related utils
 
 
+def pack_message(session, message):
+    """Packs a message following MtProto 2.0 guidelines"""
+    # See https://core.telegram.org/mtproto/description
+    data = struct.pack('<qq', session.salt, session.id) + bytes(message)
+    padding = os.urandom(-(len(data) + 12) % 16 + 12)
+
+    # Being substr(what, offset, length); x = 0 for client
+    # "msg_key_large = SHA256(substr(auth_key, 88+x, 32) + pt + padding)"
+    msg_key_large = sha256(
+        session.auth_key.key[88:88 + 32] + data + padding).digest()
+
+    # "msg_key = substr (msg_key_large, 8, 16)"
+    msg_key = msg_key_large[8:24]
+    aes_key, aes_iv = calc_key_2(session.auth_key.key, msg_key, True)
+
+    key_id = struct.pack('<Q', session.auth_key.key_id)
+    return key_id + msg_key + AES.encrypt_ige(data + padding, aes_key, aes_iv)
+
+
+def unpack_message(session, reader):
+    """Unpacks a message following MtProto 2.0 guidelines"""
+    # See https://core.telegram.org/mtproto/description
+    reader.read_long(signed=False)  # remote_auth_key_id
+
+    msg_key = reader.read(16)
+    aes_key, aes_iv = calc_key_2(session.auth_key.key, msg_key, False)
+    data = BinaryReader(AES.decrypt_ige(reader.read(), aes_key, aes_iv))
+
+    data.read_long()  # remote_salt
+    data.read_long()  # remote_session_id
+    remote_msg_id = data.read_long()
+    remote_sequence = data.read_int()
+    msg_len = data.read_int()
+    message = data.read(msg_len)
+
+    return message, remote_msg_id, remote_sequence
+
+
 def calc_key(shared_key, msg_key, client):
-    """Calculate the key based on Telegram guidelines,
-       specifying whether it's the client or not
+    """
+    Calculate the key based on Telegram guidelines,
+    specifying whether it's the client or not.
     """
     x = 0 if client else 8
 
@@ -40,6 +84,23 @@ def calc_key(shared_key, msg_key, client):
     return key, iv
 
 
+def calc_key_2(auth_key, msg_key, client):
+    """
+    Calculate the key based on Telegram guidelines
+    for MtProto 2, specifying whether it's the client or not.
+    """
+    # https://core.telegram.org/mtproto/description#defining-aes-key-and-initialization-vector
+    x = 0 if client else 8
+
+    sha256a = sha256(msg_key + auth_key[x: x + 36]).digest()
+    sha256b = sha256(auth_key[x + 40:x + 76] + msg_key).digest()
+
+    aes_key = sha256a[:8] + sha256b[8:24] + sha256a[24:32]
+    aes_iv = sha256b[:8] + sha256a[8:24] + sha256b[24:32]
+
+    return aes_key, aes_iv
+
+
 def calc_msg_key(data):
     """Calculates the message key from the given data"""
     return sha1(data).digest()[4:20]
diff --git a/telethon/network/mtproto_sender.py b/telethon/network/mtproto_sender.py
index 7e4d2f18..82a378ba 100644
--- a/telethon/network/mtproto_sender.py
+++ b/telethon/network/mtproto_sender.py
@@ -156,17 +156,7 @@ class MtProtoSender:
 
         :param message: the TLMessage to be sent.
         """
-        plain_text = \
-            struct.pack('<qq', self.session.salt, self.session.id) \
-            + bytes(message)
-
-        msg_key = utils.calc_msg_key(plain_text)
-        key_id = struct.pack('<Q', self.session.auth_key.key_id)
-        key, iv = utils.calc_key(self.session.auth_key.key, msg_key, True)
-        cipher_text = AES.encrypt_ige(plain_text, key, iv)
-
-        result = key_id + msg_key + cipher_text
-        self.connection.send(result)
+        self.connection.send(utils.pack_message(self.session, message))
 
     def _decode_msg(self, body):
         """
@@ -175,34 +165,14 @@ class MtProtoSender:
         :param body: the body to be decoded.
         :return: a tuple of (decoded message, remote message id, remote seq).
         """
-        message = None
-        remote_msg_id = None
-        remote_sequence = None
+        if len(body) < 8:
+            if body == b'l\xfe\xff\xff':
+                raise BrokenAuthKeyError()
+            else:
+                raise BufferError("Can't decode packet ({})".format(body))
 
         with BinaryReader(body) as reader:
-            if len(body) < 8:
-                if body == b'l\xfe\xff\xff':
-                    raise BrokenAuthKeyError()
-                else:
-                    raise BufferError("Can't decode packet ({})".format(body))
-
-            # TODO Check for both auth key ID and msg_key correctness
-            reader.read_long()  # remote_auth_key_id
-            msg_key = reader.read(16)
-
-            key, iv = utils.calc_key(self.session.auth_key.key, msg_key, False)
-            plain_text = AES.decrypt_ige(
-                reader.read(len(body) - reader.tell_position()), key, iv)
-
-            with BinaryReader(plain_text) as plain_text_reader:
-                plain_text_reader.read_long()  # remote_salt
-                plain_text_reader.read_long()  # remote_session_id
-                remote_msg_id = plain_text_reader.read_long()
-                remote_sequence = plain_text_reader.read_int()
-                msg_len = plain_text_reader.read_int()
-                message = plain_text_reader.read(msg_len)
-
-        return message, remote_msg_id, remote_sequence
+            return utils.unpack_message(self.session, reader)
 
     def _process_msg(self, msg_id, sequence, reader, state):
         """