LonamiWebs/Telethon
1
1
Fork 0
mirror of https://github.com/LonamiWebs/Telethon.git synced 2024-09-21 03:08:49 +03:00
Code Issues Packages Projects Releases Wiki Activity

Create a proper Message class (msg_id, seqno, body; only .to_bytes())

Browse Source
This commit is contained in:
Lonami Exo 2017-09-27 21:01:20 +02:00
parent b0839a028e
commit bd3dd371a2
4 changed files with 67 additions and 72 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

99
telethon/network/mtproto_sender.py
View File

@ -1,5 +1,6 @@
import gzip
import logging
import struct
from threading import RLock
from .. import helpers as utils
@ -8,8 +9,8 @@ from ..errors import (
BadMessageError, InvalidChecksumError, BrokenAuthKeyError,
rpc_message_to_error
)
from ..extensions import BinaryReader, BinaryWriter
from ..tl import MessageContainer, GzipPacked
from ..extensions import BinaryReader
from ..tl import Message, MessageContainer, GzipPacked
from ..tl.all_tlobjects import tlobjects
from ..tl.types import MsgsAck
@ -29,8 +30,11 @@ class MtProtoSender:
self.connection = connection
self._logger = logging.getLogger(__name__)
self._need_confirmation = [] # Message IDs that need confirmation
self._pending_receive = [] # Requests sent waiting to be received
# Message IDs that need confirmation
self._need_confirmation = []
# Requests (as msg_id: Message) sent waiting to be received
self._pending_receive = {}
# Sending and receiving are independent, but two threads cannot
# send or receive at the same time no matter what.
@ -65,21 +69,22 @@ class MtProtoSender:
self._send_acknowledges()
# Finally send our packed request(s)
self._pending_receive.extend(requests)
if len(requests) == 1:
request = requests[0]
data = GzipPacked.gzip_if_smaller(request)
else:
request = MessageContainer(self.session, requests)
data = request.to_bytes()
messages = [Message(self.session, r) for r in requests]
self._pending_receive.update({m.msg_id: m for m in messages})
self._send_packet(data, request)
if len(messages) == 1:
message = messages[0]
else:
message = Message(self.session, MessageContainer(messages))
self._send_message(message)
def _send_acknowledges(self):
"""Sends a messages acknowledge for all those who _need_confirmation"""
if self._need_confirmation:
msgs_ack = MsgsAck(self._need_confirmation)
self._send_packet(msgs_ack.to_bytes(), msgs_ack)
self._send_message(
Message(self.session, MsgsAck(self._need_confirmation))
)
del self._need_confirmation[:]
def receive(self, update_state):
@ -114,36 +119,21 @@ class MtProtoSender:
# region Low level processing
def _send_packet(self, packet, request):
"""Sends the given packet bytes with the additional
information of the original request.
"""
request.request_msg_id = self.session.get_new_msg_id()
def _send_message(self, message):
"""Sends the given Message(TLObject) encrypted through the network"""
# First calculate plain_text to encrypt it
with BinaryWriter() as plain_writer:
plain_writer.write_long(self.session.salt, signed=False)
plain_writer.write_long(self.session.id, signed=False)
plain_writer.write_long(request.request_msg_id)
plain_writer.write_int(
self.session.generate_sequence(request.content_related))
plain_text = \
struct.pack('<QQ', self.session.salt, self.session.id) \
+ message.to_bytes()
plain_writer.write_int(len(packet))
plain_writer.write(packet)
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)
msg_key = utils.calc_msg_key(plain_writer.get_bytes())
key, iv = utils.calc_key(self.session.auth_key.key, msg_key, True)
cipher_text = AES.encrypt_ige(plain_writer.get_bytes(), key, iv)
# And then finally send the encrypted packet
with BinaryWriter() as cipher_writer:
cipher_writer.write_long(
self.session.auth_key.key_id, signed=False)
cipher_writer.write(msg_key)
cipher_writer.write(cipher_text)
with self._send_lock:
self.connection.send(cipher_writer.get_bytes())
result = key_id + msg_key + cipher_text
with self._send_lock:
self.connection.send(result)
def _decode_msg(self, body):
"""Decodes an received encrypted message body bytes"""
@ -211,13 +201,12 @@ class MtProtoSender:
# msgs_ack, it may handle the request we wanted
if code == 0x62d6b459:
ack = reader.tgread_object()
for r in self._pending_receive:
if r.request_msg_id in ack.msg_ids:
self._logger.debug('Ack found for the a request')
if self.logging_out:
self._logger.debug('Message ack confirmed a request')
self._pending_receive.remove(r)
# We only care about ack requests if we're logging out
if self.logging_out:
for msg_id in ack.msg_ids:
r = self._pop_request(msg_id)
if r:
self._logger.debug('Message ack confirmed', r)
r.confirm_received.set()
return True
@ -244,16 +233,16 @@ class MtProtoSender:
# region Message handling
def _pop_request(self, request_msg_id):
"""Pops a pending request from self._pending_receive, or
returns None if it's not found
def _pop_request(self, msg_id):
"""Pops a pending REQUEST from self._pending_receive, or
returns None if it's not found.
"""
for i in range(len(self._pending_receive)):
if self._pending_receive[i].request_msg_id == request_msg_id:
return self._pending_receive.pop(i)
message = self._pending_receive.pop(msg_id, None)
if message:
return message.request
def _clear_all_pending(self):
for r in self._pending_receive:
for r in self._pending_receive.values():
r.confirm_received.set()
self._pending_receive.clear()

1
telethon/tl/__init__.py
View File

@ -1,4 +1,5 @@
from .tlobject import TLObject
from .session import Session
from .gzip_packed import GzipPacked
from .message import Message
from .message_container import MessageContainer

17
telethon/tl/message.py Normal file
View File

@ -0,0 +1,17 @@
import struct
from . import TLObject, GzipPacked
class Message(TLObject):
"""https://core.telegram.org/mtproto/service_messages#simple-container"""
def __init__(self, session, request):
super().__init__()
del self.content_related
self.msg_id = session.get_new_msg_id()
self.seq_no = session.generate_sequence(request.content_related)
self.request = request
def to_bytes(self):
body = GzipPacked.gzip_if_smaller(self.request)
return struct.pack('<qii', self.msg_id, self.seq_no, len(body)) + body

22
telethon/tl/message_container.py
View File

@ -5,31 +5,19 @@ from ..extensions import BinaryWriter
class MessageContainer(TLObject):
constructor_id = 0x73f1f8dc
# TODO Currently it's a bit of a hack, since the container actually holds
# messages (message id, sequence number, request body), not requests.
# Probably create a proper "Message" class
def __init__(self, session, requests):
def __init__(self, messages):
super().__init__()
self.content_related = False
self.session = session
self.requests = requests
self.messages = messages
def to_bytes(self):
# TODO Change this to delete the on_send from this class
with BinaryWriter() as writer:
writer.write_int(MessageContainer.constructor_id, signed=False)
writer.write_int(len(self.requests))
for x in self.requests:
x.request_msg_id = self.session.get_new_msg_id()
writer.write_int(len(self.messages))
for m in self.messages:
writer.write(m.to_bytes())
writer.write_long(x.request_msg_id)
writer.write_int(
self.session.generate_sequence(x.content_related)
)
packet = GzipPacked.gzip_if_smaller(x)
writer.write_int(len(packet))
writer.write(packet)
return writer.get_bytes()
@staticmethod