Added and updated documentation

network/mtproto_plain_sender.py

@@ -1,11 +1,12 @@
-
+# This file is based on TLSharp
+# https://github.com/sochix/TLSharp/blob/master/TLSharp.Core/Network/MtProtoPlainSender.cs
 import time
 from utils.binary_writer import BinaryWriter
 from utils.binary_reader import BinaryReader
 
 
 class MtProtoPlainSender:
-
+    """MTProto Mobile Protocol plain sender (https://core.telegram.org/mtproto/description#unencrypted-messages)"""
     def __init__(self, transport):
         self._sequence = 0
         self._time_offset = 0
@@ -13,6 +14,7 @@ class MtProtoPlainSender:
         self._transport = transport
 
     def send(self, data):
+        """Sends a plain packet (auth_key_id = 0) containing the given message body (data)"""
         with BinaryWriter() as writer:
             writer.write_long(0)
             writer.write_int(self.get_new_msg_id())
@@ -23,18 +25,21 @@ class MtProtoPlainSender:
             self._transport.send(packet)
 
     def receive(self):
+        """Receives a plain packet, returning the body of the response"""
         result = self._transport.receive()
         with BinaryReader(result.body) as reader:
             auth_key_id = reader.read_long()
-            message_id = reader.read_long()
+            msg_id = reader.read_long()
             message_length = reader.read_int()
 
             response = reader.read(message_length)
             return response
 
     def get_new_msg_id(self):
-        new_msg_id = int(self._time_offset + time.time() * 1000)  # multiply by 1000 to get milliseconds
+        """Generates a new message ID based on the current time (in ms) since epoch"""
+        new_msg_id = int(self._time_offset + time.time() * 1000)  # Multiply by 1000 to get milliseconds
 
+        # Ensure that we always return a message ID which is higher than the previous one
         if self._last_msg_id >= new_msg_id:
             new_msg_id = self._last_msg_id + 4
 

network/mtproto_sender.py

@@ -1,3 +1,5 @@
+# This file is based on TLSharp
+# https://github.com/sochix/TLSharp/blob/master/TLSharp.Core/Network/MtProtoSender.cs
 import re
 import zlib
 import pyaes
@@ -5,51 +7,68 @@ from time import sleep
 
 from utils.binary_writer import BinaryWriter
 from utils.binary_reader import BinaryReader
-from requests.ack_request import AckRequest
+from tl.types.msgs_ack import MsgsAck
 import utils.helpers as helpers
 
 
 class MtProtoSender:
+    """MTProto Mobile Protocol sender (https://core.telegram.org/mtproto/description)"""
     def __init__(self, transport, session):
-        self._transport = transport
-        self._session = session
-
-        self.need_confirmation = []
-
-    def change_transport(self, transport):
-        self._transport = transport
+        self.transport = transport
+        self.session = session
+        self.need_confirmation = []  # Message IDs that need confirmation
 
     def generate_sequence(self, confirmed):
+        """Generates the next sequence number, based on whether it was confirmed yet or not"""
         if confirmed:
-            result = self._session.sequence * 2 + 1
-            self._session.sequence += 1
+            result = self.session.sequence * 2 + 1
+            self.session.sequence += 1
             return result
         else:
-            return self._session.sequence * 2
+            return self.session.sequence * 2
 
-    # TODO async?
+    # region Send and receive
+
+    # TODO In TLSharp, this was async. Should this be?
     def send(self, request):
-        if self.need_confirmation:
-            ack_request = AckRequest(self.need_confirmation)
+        """Sends the specified MTProtoRequest, previously sending any message which needed confirmation"""
 
+        # First check if any message needs confirmation, if this is the case, send an "AckRequest"
+        if self.need_confirmation:
+            msgs_ack = MsgsAck(self.need_confirmation)
             with BinaryWriter() as writer:
-                ack_request.on_send(writer)
-                self.send_packet(writer.get_bytes(), ack_request)
+                msgs_ack.on_send(writer)
+                self.send_packet(writer.get_bytes(), msgs_ack)
                 del self.need_confirmation[:]
 
+        # Then send our packed request
         with BinaryWriter() as writer:
             request.on_send(writer)
             self.send_packet(writer.get_bytes(), request)
 
-        self._session.save()
+        # And update the saved session
+        self.session.save()
+
+    def receive(self, request):
+        """Receives the specified MTProtoRequest ("fills in it" the received data)"""
+        while not request.confirm_received:
+            message, remote_msg_id, remote_sequence = self.decode_msg(self.transport.receive().body)
+
+            with BinaryReader(message) as reader:
+                self.process_msg(remote_msg_id, remote_sequence, reader, request)
+
+    # endregion
+
+    # region Low level processing
 
     def send_packet(self, packet, request):
-        request.message_id = self._session.get_new_msg_id()
+        """Sends the given packet bytes with the additional information of the original request"""
+        request.msg_id = self.session.get_new_msg_id()
 
+        # First calculate the ciphered bit
         with BinaryWriter() as writer:
-            # TODO Is there any difference with unsigned long and long?
-            writer.write_long(self._session.salt, signed=False)
-            writer.write_long(self._session.id, signed=False)
+            writer.write_long(self.session.salt, signed=False)
+            writer.write_long(self.session.id, signed=False)
             writer.write_long(request.msg_id)
             writer.write_int(self.generate_sequence(request.confirmed))
             writer.write_int(len(packet))
@@ -57,21 +76,22 @@ class MtProtoSender:
 
             msg_key = helpers.calc_msg_key(writer.get_bytes())
 
-            key, iv = helpers.calc_key(self._session.auth_key.data, msg_key, True)
+            key, iv = helpers.calc_key(self.session.auth_key.data, msg_key, True)
             aes = pyaes.AESModeOfOperationCFB(key, iv, 16)
             cipher_text = aes.encrypt(writer.get_bytes())
 
+        # And then finally send the packet
         with BinaryWriter() as writer:
-            # TODO is it unsigned long?
-            writer.write_long(self._session.auth_key.id, signed=False)
+            writer.write_long(self.session.auth_key.id, signed=False)
             writer.write(msg_key)
             writer.write(cipher_text)
 
-            self._transport.send(writer.get_bytes())
+            self.transport.send(writer.get_bytes())
 
     def decode_msg(self, body):
+        """Decodes an received encrypted message body bytes"""
         message = None
-        remote_message_id = None
+        remote_msg_id = None
         remote_sequence = None
 
         with BinaryReader(body) as reader:
@@ -82,72 +102,70 @@ class MtProtoSender:
             remote_auth_key_id = reader.read_long()
             msg_key = reader.read(16)
 
-            key, iv = helpers.calc_key(self._session.auth_key.data, msg_key, False)
+            key, iv = helpers.calc_key(self.session.auth_key.data, msg_key, False)
             aes = pyaes.AESModeOfOperationCFB(key, iv, 16)
             plain_text = aes.decrypt(reader.read(len(body) - reader.tell_position()))
 
             with BinaryReader(plain_text) as plain_text_reader:
                 remote_salt = plain_text_reader.read_long()
                 remote_session_id = plain_text_reader.read_long()
-                remote_message_id = plain_text_reader.read_long()
+                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_message_id, remote_sequence
+        return message, remote_msg_id, remote_sequence
 
-    def receive(self, mtproto_request):
-        while not mtproto_request.confirm_received:
-            message, remote_message_id, remote_sequence = self.decode_msg(self._transport.receive().body)
-
-            with BinaryReader(message) as reader:
-                self.process_msg(remote_message_id, remote_sequence, reader, mtproto_request)
-
-    def process_msg(self, message_id, sequence, reader, mtproto_request):
+    def process_msg(self, msg_id, sequence, reader, request):
+        """Processes and handles a Telegram message"""
         # TODO Check salt, session_id and sequence_number
-        self.need_confirmation.append(message_id)
+        self.need_confirmation.append(msg_id)
 
         code = reader.read_int(signed=False)
         reader.seek(-4)
 
         if code == 0x73f1f8dc:  # Container
-            return self.handle_container(message_id, sequence, reader, mtproto_request)
+            return self.handle_container(msg_id, sequence, reader, request)
         if code == 0x7abe77ec:  # Ping
-            return self.handle_ping(message_id, sequence, reader)
+            return self.handle_ping(msg_id, sequence, reader)
         if code == 0x347773c5:  # pong
-            return self.handle_pong(message_id, sequence, reader)
+            return self.handle_pong(msg_id, sequence, reader)
         if code == 0xae500895:  # future_salts
-            return self.handle_future_salts(message_id, sequence, reader)
+            return self.handle_future_salts(msg_id, sequence, reader)
         if code == 0x9ec20908:  # new_session_created
-            return self.handle_new_session_created(message_id, sequence, reader)
+            return self.handle_new_session_created(msg_id, sequence, reader)
         if code == 0x62d6b459:  # msgs_ack
-            return self.handle_msgs_ack(message_id, sequence, reader)
+            return self.handle_msgs_ack(msg_id, sequence, reader)
         if code == 0xedab447b:  # bad_server_salt
-            return self.handle_bad_server_salt(message_id, sequence, reader, mtproto_request)
+            return self.handle_bad_server_salt(msg_id, sequence, reader, request)
         if code == 0xa7eff811:  # bad_msg_notification
-            return self.handle_bad_msg_notification(message_id, sequence, reader)
+            return self.handle_bad_msg_notification(msg_id, sequence, reader)
         if code == 0x276d3ec6:  # msg_detailed_info
-            return self.hangle_msg_detailed_info(message_id, sequence, reader)
+            return self.hangle_msg_detailed_info(msg_id, sequence, reader)
         if code == 0xf35c6d01:  # rpc_result
-            return self.handle_rpc_result(message_id, sequence, reader, mtproto_request)
+            return self.handle_rpc_result(msg_id, sequence, reader, request)
         if code == 0x3072cfa1:  # gzip_packed
-            return self.handle_gzip_packed(message_id, sequence, reader, mtproto_request)
+            return self.handle_gzip_packed(msg_id, sequence, reader, request)
 
         if (code == 0xe317af7e or
-                    code == 0xd3f45784 or
-                    code == 0x2b2fbd4e or
-                    code == 0x78d4dec1 or
-                    code == 0x725b04c3 or
-                    code == 0x74ae4240):
-            return self.handle_update(message_id, sequence, reader)
+            code == 0xd3f45784 or
+            code == 0x2b2fbd4e or
+            code == 0x78d4dec1 or
+            code == 0x725b04c3 or
+                code == 0x74ae4240):
+            return self.handle_update(msg_id, sequence, reader)
 
-        # TODO Log unknown message code
+        print('Unknown message: {}'.format(hex(msg_id)))
         return False
 
-    def handle_update(self, message_id, sequence, reader):
+    # endregion
+
+    # region Message handling
+
+    def handle_update(self, msg_id, sequence, reader):
         return False
 
-    def handle_container(self, message_id, sequence, reader, mtproto_request):
+    def handle_container(self, msg_id, sequence, reader, request):
         code = reader.read_int(signed=False)
         size = reader.read_int()
         for _ in range(size):
@@ -156,7 +174,7 @@ class MtProtoSender:
             inner_length = reader.read_int()
             begin_position = reader.tell_position()
             try:
-                if not self.process_msg(inner_msg_id, sequence, reader, mtproto_request):
+                if not self.process_msg(inner_msg_id, sequence, reader, request):
                     reader.set_position(begin_position + inner_length)
 
             except:
@@ -164,40 +182,40 @@ class MtProtoSender:
 
         return False
 
-    def handle_ping(self, message_id, sequence, reader):
+    def handle_ping(self, msg_id, sequence, reader):
         return False
 
-    def handle_pong(self, message_id, sequence, reader):
+    def handle_pong(self, msg_id, sequence, reader):
         return False
 
-    def handle_future_salts(self, message_id, sequence, reader):
+    def handle_future_salts(self, msg_id, sequence, reader):
         code = reader.read_int(signed=False)
         request_id = reader.read_long(signed=False)
         reader.seek(-12)
 
         raise NotImplementedError("Handle future server salts function isn't implemented.")
 
-    def handle_new_session_created(self, message_id, sequence, reader):
+    def handle_new_session_created(self, msg_id, sequence, reader):
         return False
 
-    def handle_msgs_ack(self, message_id, sequence, reader):
+    def handle_msgs_ack(self, msg_id, sequence, reader):
         return False
 
-    def handle_bad_server_salt(self, message_id, sequence, reader, mtproto_request):
+    def handle_bad_server_salt(self, msg_id, sequence, reader, mtproto_request):
         code = reader.read_int(signed=False)
         bad_msg_id = reader.read_long(signed=False)
         bad_msg_seq_no = reader.read_int()
         error_code = reader.read_int()
         new_salt = reader.read_long(signed=False)
 
-        self._session.salt = new_salt
+        self.session.salt = new_salt
 
         # Resend
         self.send(mtproto_request)
 
         return True
 
-    def handle_bad_msg_notification(self, message_id, sequence, reader):
+    def handle_bad_msg_notification(self, msg_id, sequence, reader):
         code = reader.read_int(signed=False)
         request_id = reader.read_long(signed=False)
         request_sequence = reader.read_int()
@@ -238,14 +256,14 @@ class MtProtoSender:
 
         raise NotImplementedError('This should never happen!')
 
-    def hangle_msg_detailed_info(self, message_id, sequence, reader):
+    def hangle_msg_detailed_info(self, msg_id, sequence, reader):
         return False
 
-    def handle_rpc_result(self, message_id, sequence, reader, mtproto_request):
+    def handle_rpc_result(self, msg_id, sequence, reader, mtproto_request):
         code = reader.read_int(signed=False)
         request_id = reader.read_long(signed=False)
 
-        if request_id == mtproto_request.message_id:
+        if request_id == mtproto_request.msg_id:
             mtproto_request.confirm_received = True
 
         inner_code = reader.read_int(signed=False)
@@ -273,7 +291,6 @@ class MtProtoSender:
 
                 with BinaryReader(unpacked_data) as compressed_reader:
                     mtproto_request.on_response(compressed_reader)
-
             except:
                 pass
 
@@ -281,10 +298,13 @@ class MtProtoSender:
             reader.seek(-4)
             mtproto_request.on_response(reader)
 
-    def handle_gzip_packed(self, message_id, sequence, reader, mtproto_request):
+    def handle_gzip_packed(self, msg_id, sequence, reader, mtproto_request):
         code = reader.read_int(signed=False)
         packed_data = reader.tgread_bytes()
         unpacked_data = zlib.decompress(packed_data)
 
         with BinaryReader(unpacked_data) as compressed_reader:
-            self.process_msg(message_id, sequence, compressed_reader, mtproto_request)
+            self.process_msg(msg_id, sequence, compressed_reader, mtproto_request)
+
+    # endregion
+    pass

network/tcp_client.py

@@ -1,20 +1,24 @@
+# Python rough implementation of a C# TCP client
 import socket
 
 
 class TcpClient:
-
     def __init__(self):
         self.connected = False
         self.socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
 
     def connect(self, ip, port):
+        """Connects to the specified IP and port number"""
         self.socket.connect((ip, port))
 
     def close(self):
+        """Closes the connection"""
         self.socket.close()
 
     def write(self, data):
+        """Writes (sends) the specified bytes to the connected peer"""
         self.socket.send(data)
 
     def read(self, buffer_size):
+        """Reads (receives) the specified bytes from the connected peer"""
         self.socket.recv(buffer_size)

network/tcp_message.py

@@ -1,4 +1,5 @@
-
+# This file is based on TLSharp
+# https://github.com/sochix/TLSharp/blob/master/TLSharp.Core/Network/TcpMessage.cs
 from zlib import crc32
 
 from utils.binary_writer import BinaryWriter
@@ -6,7 +7,6 @@ from utils.binary_reader import BinaryReader
 
 
 class TcpMessage:
-
     def __init__(self, seq_number, body):
         """
         :param seq_number: Sequence number
@@ -19,6 +19,7 @@ class TcpMessage:
         self.body = body
 
     def encode(self):
+        """Returns the bytes of the this message encoded, following Telegram's guidelines"""
         with BinaryWriter() as writer:
             ''' https://core.telegram.org/mtproto#tcp-transport
 
@@ -38,7 +39,9 @@ class TcpMessage:
 
             return writer.get_bytes()
 
-    def decode(self, body):
+    @staticmethod
+    def decode(body):
+        """Returns a TcpMessage from the given encoded bytes, decoding them previously"""
         if body is None:
             raise ValueError('body cannot be None')
 

network/tcp_transport.py

@@ -1,11 +1,11 @@
+# This file is based on TLSharp
+# https://github.com/sochix/TLSharp/blob/master/TLSharp.Core/Network/TcpTransport.cs
 from zlib import crc32
-
 from network.tcp_message import TcpMessage
 from network.tcp_client import TcpClient
 
 
 class TcpTransport:
-
     def __init__(self, ip_address, port):
         self._tcp_client = TcpClient()
         self._send_counter = 0
@@ -13,20 +13,20 @@ class TcpTransport:
         self._tcp_client.connect(ip_address, port)
 
     def send(self, packet):
-        """
-        :param packet: Bytes array representing the packet to be sent
-        """
+        """Sends the given packet (bytes array) to the connected peer"""
         if not self._tcp_client.connected:
             raise ConnectionError('Client not connected to server.')
 
+        # Get a TcpMessage which contains the given packet
         tcp_message = TcpMessage(self._send_counter, packet)
 
-        # TODO async? and receive too, of course
+        # TODO In TLSharp, this is async; Should both send and receive be here too?
         self._tcp_client.write(tcp_message.encode())
-
         self._send_counter += 1
 
     def receive(self):
+        """Receives a TcpMessage from the connected peer"""
+
         # First read everything
         packet_length_bytes = self._tcp_client.read(4)
         packet_length = int.from_bytes(packet_length_bytes, byteorder='big')
@@ -45,6 +45,7 @@ class TcpTransport:
         if checksum != valid_checksum:
             raise ValueError('Invalid checksum, skip')
 
+        # If we passed the tests, we can then return a valid TcpMessage
         return TcpMessage(seq, body)
 
     def dispose(self):