mirror of
https://github.com/LonamiWebs/Telethon.git
synced 2025-06-20 03:56:38 +00:00
105 lines
3.9 KiB
Python
105 lines
3.9 KiB
Python
import io
|
|
import struct
|
|
|
|
from .mtprotostate import MTProtoState
|
|
from ..tl.core.messagecontainer import MessageContainer
|
|
|
|
|
|
class MTProtoLayer:
|
|
"""
|
|
This class is the message encryption layer between the methods defined
|
|
in the schema and the response objects. It also holds the necessary state
|
|
necessary for this encryption to happen.
|
|
|
|
The `connection` parameter is through which these messages will be sent
|
|
and received.
|
|
|
|
The `auth_key` must be a valid authorization key which will be used to
|
|
encrypt these messages. This class is not responsible for generating them.
|
|
"""
|
|
def __init__(self, connection, auth_key):
|
|
self._connection = connection
|
|
self._state = MTProtoState(auth_key)
|
|
|
|
def connect(self):
|
|
"""
|
|
Wrapper for ``self._connection.connect()``.
|
|
"""
|
|
return self._connection.connect()
|
|
|
|
def disconnect(self):
|
|
"""
|
|
Wrapper for ``self._connection.disconnect()``.
|
|
"""
|
|
self._connection.disconnect()
|
|
|
|
async def send(self, data_list):
|
|
"""
|
|
A list of serialized RPC queries as bytes must be given to be sent.
|
|
Nested lists imply an order is required for the messages in them.
|
|
Message containers will be used if there is more than one item.
|
|
|
|
Returns ``(container_id, msg_ids)``.
|
|
"""
|
|
data, container_id, msg_ids = self._pack_data_list(data_list)
|
|
await self._connection.send(self._state.encrypt_message_data(data))
|
|
return container_id, msg_ids
|
|
|
|
async def recv(self):
|
|
"""
|
|
Reads a single message from the network, decrypts it and returns it.
|
|
"""
|
|
body = await self._connection.recv()
|
|
return self._state.decrypt_message_data(body)
|
|
|
|
def _pack_data_list(self, data_list):
|
|
"""
|
|
A list of serialized RPC queries as bytes must be given to be packed.
|
|
Nested lists imply an order is required for the messages in them.
|
|
|
|
Returns ``(data, container_id, msg_ids)``.
|
|
"""
|
|
# TODO write_data_as_message raises on invalid messages, handle it
|
|
# TODO This method could be an iterator yielding messages while small
|
|
# respecting the ``MessageContainer.MAXIMUM_SIZE`` limit.
|
|
#
|
|
# Note that the simplest case is writing a single query data into
|
|
# a message, and returning the message data and ID. For efficiency
|
|
# purposes this method supports more than one message and automatically
|
|
# uses containers if deemed necessary.
|
|
#
|
|
# Technically the message and message container classes could be used
|
|
# to store and serialize the data. However, to keep the context local
|
|
# and relevant to the only place where such feature is actually used,
|
|
# this is not done.
|
|
msg_ids = []
|
|
buffer = io.BytesIO()
|
|
for data in data_list:
|
|
if not isinstance(data, list):
|
|
msg_ids.append(self._state.write_data_as_message(buffer, data))
|
|
else:
|
|
last_id = None
|
|
for d in data:
|
|
last_id = self._state.write_data_as_message(
|
|
buffer, d, after_id=last_id)
|
|
msg_ids.append(last_id)
|
|
|
|
if len(msg_ids) == 1:
|
|
container_id = None
|
|
else:
|
|
# Inlined code to pack several messages into a container
|
|
#
|
|
# TODO This part and encrypting data prepend a few bytes but
|
|
# force a potentially large payload to be appended, which
|
|
# may be expensive. Can we do better?
|
|
data = struct.pack(
|
|
'<Ii', MessageContainer.CONSTRUCTOR_ID, len(msg_ids)
|
|
) + buffer.getvalue()
|
|
buffer = io.BytesIO()
|
|
container_id = self._state.write_data_as_message(buffer, data)
|
|
|
|
return buffer.getvalue(), container_id, msg_ids
|
|
|
|
def __str__(self):
|
|
return str(self._connection)
|