From e3b05cbc4efb9aaaab72c39a5c43b2125bee1f85 Mon Sep 17 00:00:00 2001
From: Dan <14043624+delivrance@users.noreply.github.com>
Date: Sat, 9 Dec 2017 01:10:09 +0100
Subject: [PATCH] Add MTProto 2.0 key derivation function

---
 pyrogram/crypto/__init__.py |  2 +-
 pyrogram/crypto/kdf.py      | 16 +++++++++++++++-
 2 files changed, 16 insertions(+), 2 deletions(-)

diff --git a/pyrogram/crypto/__init__.py b/pyrogram/crypto/__init__.py
index 71f177ed..3f87b6fc 100644
--- a/pyrogram/crypto/__init__.py
+++ b/pyrogram/crypto/__init__.py
@@ -17,6 +17,6 @@
 # along with Pyrogram.  If not, see <http://www.gnu.org/licenses/>.
 
 from .ige import IGE
-from .kdf import KDF
+from .kdf import KDF, KDF2
 from .prime import Prime
 from .rsa import RSA
diff --git a/pyrogram/crypto/kdf.py b/pyrogram/crypto/kdf.py
index 1b941ddd..abc6b976 100644
--- a/pyrogram/crypto/kdf.py
+++ b/pyrogram/crypto/kdf.py
@@ -16,7 +16,7 @@
 # You should have received a copy of the GNU Lesser General Public License
 # along with Pyrogram.  If not, see <http://www.gnu.org/licenses/>.
 
-from hashlib import sha1
+from hashlib import sha1, sha256
 
 
 class KDF:
@@ -33,3 +33,17 @@ class KDF:
         aes_iv = sha1_a[8:20] + sha1_b[:8] + sha1_c[16:20] + sha1_d[:8]
 
         return aes_key, aes_iv
+
+
+class KDF2:
+    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
+
+        sha256_a = sha256(msg_key + auth_key[x: x + 36]).digest()
+        sha256_b = sha256(auth_key[x + 40:x + 76] + msg_key).digest()  # 76 = 40 + 36
+
+        aes_key = sha256_a[:8] + sha256_b[8:24] + sha256_a[24:32]
+        aes_iv = sha256_b[:8] + sha256_a[8:24] + sha256_b[24:32]
+
+        return aes_key, aes_iv