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 . 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 . -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