Implement MtProto 2.0 (closes #484, thanks @delivrance!)

Huge shoutout to @delivrance's pyrogram, specially this commit:
pyrogram/pyrogram/commit/42f9a2d6994baaf9ecad590d1ff4d175a8c56454
This commit is contained in:
Lonami Exo 2018-01-06 01:55:11 +01:00
parent c039ba3e16
commit 3eafe18d0b
3 changed files with 75 additions and 41 deletions

View File

@ -56,8 +56,11 @@ class BinaryReader:
return int.from_bytes( return int.from_bytes(
self.read(bits // 8), byteorder='little', signed=signed) self.read(bits // 8), byteorder='little', signed=signed)
def read(self, length): def read(self, length=None):
"""Read the given amount of bytes.""" """Read the given amount of bytes."""
if length is None:
return self.reader.read()
result = self.reader.read(length) result = self.reader.read(length)
if len(result) != length: if len(result) != length:
raise BufferError( raise BufferError(

View File

@ -1,6 +1,11 @@
"""Various helpers not related to the Telegram API itself""" """Various helpers not related to the Telegram API itself"""
from hashlib import sha1, sha256
import os import os
import struct
from hashlib import sha1, sha256
from telethon.crypto import AES
from telethon.extensions import BinaryReader
# region Multiple utilities # region Multiple utilities
@ -21,9 +26,48 @@ def ensure_parent_dir_exists(file_path):
# region Cryptographic related utils # 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): 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 x = 0 if client else 8
@ -40,6 +84,23 @@ def calc_key(shared_key, msg_key, client):
return key, iv 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): def calc_msg_key(data):
"""Calculates the message key from the given data""" """Calculates the message key from the given data"""
return sha1(data).digest()[4:20] return sha1(data).digest()[4:20]

View File

@ -156,17 +156,7 @@ class MtProtoSender:
:param message: the TLMessage to be sent. :param message: the TLMessage to be sent.
""" """
plain_text = \ self.connection.send(utils.pack_message(self.session, message))
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)
def _decode_msg(self, body): def _decode_msg(self, body):
""" """
@ -175,34 +165,14 @@ class MtProtoSender:
:param body: the body to be decoded. :param body: the body to be decoded.
:return: a tuple of (decoded message, remote message id, remote seq). :return: a tuple of (decoded message, remote message id, remote seq).
""" """
message = None if len(body) < 8:
remote_msg_id = None if body == b'l\xfe\xff\xff':
remote_sequence = None raise BrokenAuthKeyError()
else:
raise BufferError("Can't decode packet ({})".format(body))
with BinaryReader(body) as reader: with BinaryReader(body) as reader:
if len(body) < 8: return utils.unpack_message(self.session, reader)
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
def _process_msg(self, msg_id, sequence, reader, state): def _process_msg(self, msg_id, sequence, reader, state):
""" """