Totally refactored source files location

Now it *should* be easier to turn Telethon
into a pip package

telethon/network/__init__.py

@@ -0,0 +1,5 @@
+from .mtproto_plain_sender import MtProtoPlainSender
+from .tcp_client import TcpClient
+from .authenticator import do_authentication
+from .mtproto_sender import MtProtoSender
+from .tcp_transport import TcpTransport

telethon/network/authenticator.py

@@ -0,0 +1,209 @@
+import os
+import time
+import telethon.helpers as utils
+from telethon.utils import BinaryWriter, BinaryReader
+from telethon.crypto import AES, AuthKey, Factorizator, RSA
+from telethon.network import MtProtoPlainSender
+
+
+def do_authentication(transport):
+    """Executes the authentication process with the Telegram servers.
+    If no error is rose, returns both the authorization key and the time offset"""
+    sender = MtProtoPlainSender(transport)
+
+    # Step 1 sending: PQ Request
+    nonce = os.urandom(16)
+    with BinaryWriter() as writer:
+        writer.write_int(0x60469778, signed=False)  # Constructor number
+        writer.write(nonce)
+        sender.send(writer.get_bytes())
+
+    # Step 1 response: PQ Request
+    pq, pq_bytes, server_nonce, fingerprints = None, None, None, []
+    with BinaryReader(sender.receive()) as reader:
+        response_code = reader.read_int(signed=False)
+        if response_code != 0x05162463:
+            raise AssertionError('Invalid response code: {}'.format(hex(response_code)))
+
+        nonce_from_server = reader.read(16)
+        if nonce_from_server != nonce:
+            raise AssertionError('Invalid nonce from server')
+
+        server_nonce = reader.read(16)
+
+        pq_bytes = reader.tgread_bytes()
+        pq = get_int(pq_bytes)
+
+        vector_id = reader.read_int()
+        if vector_id != 0x1cb5c415:
+            raise AssertionError('Invalid vector constructor ID: {}'.format(hex(response_code)))
+
+        fingerprints = []
+        fingerprint_count = reader.read_int()
+        for _ in range(fingerprint_count):
+            fingerprints.append(reader.read(8))
+
+    # Step 2 sending: DH Exchange
+    new_nonce = os.urandom(32)
+    p, q = Factorizator.factorize(pq)
+    with BinaryWriter() as pq_inner_data_writer:
+        pq_inner_data_writer.write_int(0x83c95aec, signed=False)  # PQ Inner Data
+        pq_inner_data_writer.tgwrite_bytes(get_byte_array(pq, signed=False))
+        pq_inner_data_writer.tgwrite_bytes(get_byte_array(min(p, q), signed=False))
+        pq_inner_data_writer.tgwrite_bytes(get_byte_array(max(p, q), signed=False))
+        pq_inner_data_writer.write(nonce)
+        pq_inner_data_writer.write(server_nonce)
+        pq_inner_data_writer.write(new_nonce)
+
+        cipher_text, target_fingerprint = None, None
+        for fingerprint in fingerprints:
+            cipher_text = RSA.encrypt(get_fingerprint_text(fingerprint), pq_inner_data_writer.get_bytes())
+
+            if cipher_text is not None:
+                target_fingerprint = fingerprint
+                break
+
+        if cipher_text is None:
+            raise AssertionError('Could not find a valid key for fingerprints: {}'
+                                 .format(', '.join([get_fingerprint_text(f) for f in fingerprints])))
+
+        with BinaryWriter() as req_dh_params_writer:
+            req_dh_params_writer.write_int(0xd712e4be, signed=False)  # Req DH Params
+            req_dh_params_writer.write(nonce)
+            req_dh_params_writer.write(server_nonce)
+            req_dh_params_writer.tgwrite_bytes(get_byte_array(min(p, q), signed=False))
+            req_dh_params_writer.tgwrite_bytes(get_byte_array(max(p, q), signed=False))
+            req_dh_params_writer.write(target_fingerprint)
+            req_dh_params_writer.tgwrite_bytes(cipher_text)
+
+            req_dh_params_bytes = req_dh_params_writer.get_bytes()
+            sender.send(req_dh_params_bytes)
+
+    # Step 2 response: DH Exchange
+    encrypted_answer = None
+    with BinaryReader(sender.receive()) as reader:
+        response_code = reader.read_int(signed=False)
+
+        if response_code == 0x79cb045d:
+            raise AssertionError('Server DH params fail: TODO')
+
+        if response_code != 0xd0e8075c:
+            raise AssertionError('Invalid response code: {}'.format(hex(response_code)))
+
+        nonce_from_server = reader.read(16)
+        if nonce_from_server != nonce:
+            raise NotImplementedError('Invalid nonce from server')
+
+        server_nonce_from_server = reader.read(16)
+        if server_nonce_from_server != server_nonce:
+            raise NotImplementedError('Invalid server nonce from server')
+
+        encrypted_answer = reader.tgread_bytes()
+
+    # Step 3 sending: Complete DH Exchange
+    key, iv = utils.generate_key_data_from_nonces(server_nonce, new_nonce)
+    plain_text_answer = AES.decrypt_ige(encrypted_answer, key, iv)
+
+    g, dh_prime, ga, time_offset = None, None, None, None
+    with BinaryReader(plain_text_answer) as dh_inner_data_reader:
+        hashsum = dh_inner_data_reader.read(20)
+        code = dh_inner_data_reader.read_int(signed=False)
+        if code != 0xb5890dba:
+            raise AssertionError('Invalid DH Inner Data code: {}'.format(code))
+
+        nonce_from_server1 = dh_inner_data_reader.read(16)
+        if nonce_from_server1 != nonce:
+            raise AssertionError('Invalid nonce in encrypted answer')
+
+        server_nonce_from_server1 = dh_inner_data_reader.read(16)
+        if server_nonce_from_server1 != server_nonce:
+            raise AssertionError('Invalid server nonce in encrypted answer')
+
+        g = dh_inner_data_reader.read_int()
+        dh_prime = get_int(dh_inner_data_reader.tgread_bytes(), signed=False)
+        ga = get_int(dh_inner_data_reader.tgread_bytes(), signed=False)
+
+        server_time = dh_inner_data_reader.read_int()
+        time_offset = server_time - int(time.time())
+
+    b = get_int(os.urandom(2048), signed=False)
+    gb = pow(g, b, dh_prime)
+    gab = pow(ga, b, dh_prime)
+
+    # Prepare client DH Inner Data
+    with BinaryWriter() as client_dh_inner_data_writer:
+        client_dh_inner_data_writer.write_int(0x6643b654, signed=False)  # Client DH Inner Data
+        client_dh_inner_data_writer.write(nonce)
+        client_dh_inner_data_writer.write(server_nonce)
+        client_dh_inner_data_writer.write_long(0)  # TODO retry_id
+        client_dh_inner_data_writer.tgwrite_bytes(get_byte_array(gb, signed=False))
+
+        with BinaryWriter() as client_dh_inner_data_with_hash_writer:
+            client_dh_inner_data_with_hash_writer.write(utils.sha1(client_dh_inner_data_writer.get_bytes()))
+            client_dh_inner_data_with_hash_writer.write(client_dh_inner_data_writer.get_bytes())
+            client_dh_inner_data_bytes = client_dh_inner_data_with_hash_writer.get_bytes()
+
+    # Encryption
+    client_dh_inner_data_encrypted_bytes = AES.encrypt_ige(client_dh_inner_data_bytes, key, iv)
+
+    # Prepare Set client DH params
+    with BinaryWriter() as set_client_dh_params_writer:
+        set_client_dh_params_writer.write_int(0xf5045f1f, signed=False)
+        set_client_dh_params_writer.write(nonce)
+        set_client_dh_params_writer.write(server_nonce)
+        set_client_dh_params_writer.tgwrite_bytes(client_dh_inner_data_encrypted_bytes)
+
+        set_client_dh_params_bytes = set_client_dh_params_writer.get_bytes()
+        sender.send(set_client_dh_params_bytes)
+
+    # Step 3 response: Complete DH Exchange
+    with BinaryReader(sender.receive()) as reader:
+        code = reader.read_int(signed=False)
+        if code == 0x3bcbf734:  # DH Gen OK
+            nonce_from_server = reader.read(16)
+            if nonce_from_server != nonce:
+                raise NotImplementedError('Invalid nonce from server')
+
+            server_nonce_from_server = reader.read(16)
+            if server_nonce_from_server != server_nonce:
+                raise NotImplementedError('Invalid server nonce from server')
+
+            new_nonce_hash1 = reader.read(16)
+            auth_key = AuthKey(get_byte_array(gab, signed=False))
+
+            new_nonce_hash_calculated = auth_key.calc_new_nonce_hash(new_nonce, 1)
+            if new_nonce_hash1 != new_nonce_hash_calculated:
+                raise AssertionError('Invalid new nonce hash')
+
+            return auth_key, time_offset
+
+        elif code == 0x46dc1fb9:  # DH Gen Retry
+            raise NotImplementedError('dh_gen_retry')
+
+        elif code == 0xa69dae02:  # DH Gen Fail
+            raise NotImplementedError('dh_gen_fail')
+
+        else:
+            raise AssertionError('DH Gen unknown: {}'.format(hex(code)))
+
+
+def get_fingerprint_text(fingerprint):
+    """Gets a fingerprint text in 01-23-45-67-89-AB-CD-EF format (no hyphens)"""
+    return ''.join(hex(b)[2:].rjust(2, '0').upper() for b in fingerprint)
+
+
+# The following methods operate in big endian (unlike most of Telegram API) because:
+# > "...pq is a representation of a natural number (in binary *big endian* format)..."
+# > "...current value of dh_prime equals (in *big-endian* byte order)..."
+# Reference: https://core.telegram.org/mtproto/auth_key
+def get_byte_array(integer, signed):
+    """Gets the arbitrary-length byte array corresponding to the given integer"""
+    bits = integer.bit_length()
+    byte_length = (bits + 8 - 1) // 8  # 8 bits per byte
+    return int.to_bytes(integer, length=byte_length, byteorder='big', signed=signed)
+
+
+def get_int(byte_array, signed=True):
+    """Gets the specified integer from its byte array. This should be used by the authenticator,
+       who requires the data to be in big endian"""
+    return int.from_bytes(byte_array, byteorder='big', signed=signed)

telethon/network/mtproto_plain_sender.py

@@ -0,0 +1,49 @@
+import time
+import random
+from telethon.utils import BinaryWriter, 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
+        self._last_msg_id = 0
+        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_long(self.get_new_msg_id())
+            writer.write_int(len(data))
+            writer.write(data)
+
+            packet = writer.get_bytes()
+            self._transport.send(packet)
+
+    def receive(self):
+        """Receives a plain packet, returning the body of the response"""
+        seq, body = self._transport.receive()
+        with BinaryReader(body) as reader:
+            auth_key_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):
+        """Generates a new message ID based on the current time (in ms) since epoch"""
+        # See https://core.telegram.org/mtproto/description#message-identifier-msg-id
+        ms_time = int(time.time() * 1000)
+        new_msg_id = (((ms_time // 1000) << 32) |  # "must approximately equal unixtime*2^32"
+                      ((ms_time % 1000) << 22) |  # "approximate moment in time the message was created"
+                      random.randint(0, 524288) << 2)  # "message identifiers are divisible by 4"
+
+        # 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
+
+        self._last_msg_id = new_msg_id
+        return new_msg_id

telethon/network/mtproto_sender.py

@@ -0,0 +1,331 @@
+import gzip
+from telethon.errors import *
+from time import sleep
+from threading import Thread, Lock
+
+import telethon.helpers as utils
+from telethon.crypto import AES
+from telethon.utils import BinaryWriter, BinaryReader
+from telethon.tl.types import MsgsAck
+from telethon.tl.all_tlobjects import tlobjects
+
+
+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 = []  # Message IDs that need confirmation
+        self.on_update_handlers = []
+
+        # Store a Lock instance to make this class safely multi-threaded
+        self.lock = Lock()
+
+        self.updates_thread = Thread(target=self.updates_thread_method, name='Updates thread')
+        self.updates_thread_running = False
+        self.updates_thread_receiving = False
+
+    def disconnect(self):
+        """Disconnects and **stops all the running threads** if any"""
+        self.set_listen_for_updates(enabled=False)
+        self.transport.close()
+
+    def add_update_handler(self, handler):
+        """Adds an update handler (a method with one argument, the received
+           TLObject) that is fired when there are updates available"""
+
+        first_handler = not self.on_update_handlers
+        self.on_update_handlers.append(handler)
+
+        # If this is the first added handler,
+        # we must start the thread to receive updates
+        if first_handler:
+            self.set_listen_for_updates(enabled=True)
+
+    def remove_update_handler(self, handler):
+        self.on_update_handlers.remove(handler)
+
+        # If there are no more update handlers, stop the thread
+        if not self.on_update_handlers:
+            self.set_listen_for_updates(False)
+
+    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
+            return result
+        else:
+            return self.session.sequence * 2
+
+    # region Send and receive
+
+    def send(self, request, resend=False):
+        """Sends the specified MTProtoRequest, previously sending any message
+           which needed confirmation. This also pauses the updates thread"""
+
+        # Only cancel the receive *if* it was the
+        # updates thread who was receiving. We do
+        # not want to cancel other pending requests!
+        if self.updates_thread_receiving:
+            self.transport.cancel_receive()
+
+        # Now only us can be using this method if we're not resending
+        if not resend:
+            self.lock.acquire()
+
+        # If any message needs confirmation send an AckRequest first
+        if self.need_confirmation:
+            msgs_ack = MsgsAck(self.need_confirmation)
+            with BinaryWriter() as writer:
+                msgs_ack.on_send(writer)
+                self.send_packet(writer.get_bytes(), msgs_ack)
+
+            del self.need_confirmation[:]
+
+        # Finally send our packed request
+        with BinaryWriter() as writer:
+            request.on_send(writer)
+            self.send_packet(writer.get_bytes(), request)
+
+        # And update the saved session
+        self.session.save()
+        # Don't resume the updates thread yet,
+        # since every send() is preceded by a receive()
+
+    def receive(self, request):
+        """Receives the specified MTProtoRequest ("fills in it"
+           the received data). This also restores the updates thread"""
+
+        try:
+            # Don't stop trying to receive until we get the request we wanted
+            while not request.confirm_received:
+                seq, body = self.transport.receive()
+                message, remote_msg_id, remote_sequence = self.decode_msg(body)
+
+                with BinaryReader(message) as reader:
+                    self.process_msg(remote_msg_id, remote_sequence, reader, request)
+
+        finally:
+            # Once we are done trying to get our request,
+            # restore the updates thread and release the lock
+            self.lock.release()
+
+    # endregion
+
+    # region Low level processing
+
+    def send_packet(self, packet, request):
+        """Sends the given packet bytes with the additional
+           information of the original request. This does NOT lock the threads!"""
+        request.msg_id = self.session.get_new_msg_id()
+
+        # 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.msg_id)
+            plain_writer.write_int(self.generate_sequence(request.confirmed))
+            plain_writer.write_int(len(packet))
+            plain_writer.write(packet)
+
+            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)
+            self.transport.send(cipher_writer.get_bytes())
+
+    def decode_msg(self, body):
+        """Decodes an received encrypted message body bytes"""
+        message = None
+        remote_msg_id = None
+        remote_sequence = None
+
+        with BinaryReader(body) as reader:
+            if len(body) < 8:
+                raise BufferError("Can't decode packet ({})".format(body))
+
+            # TODO Check for both auth key ID and msg_key correctness
+            remote_auth_key_id = reader.read_long()
+            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:
+                remote_salt = plain_text_reader.read_long()
+                remote_session_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_msg_id, remote_sequence
+
+    def process_msg(self, msg_id, sequence, reader, request=None):
+        """Processes and handles a Telegram message"""
+
+        # TODO Check salt, session_id and sequence_number
+        self.need_confirmation.append(msg_id)
+
+        code = reader.read_int(signed=False)
+        reader.seek(-4)
+
+        # The following codes are "parsed manually"
+        if code == 0xf35c6d01:  # rpc_result, (response of an RPC call, i.e., we sent a request)
+            return self.handle_rpc_result(msg_id, sequence, reader, request)
+
+        if code == 0x73f1f8dc:  # msg_container
+            return self.handle_container(msg_id, sequence, reader, request)
+        if code == 0x3072cfa1:  # gzip_packed
+            return self.handle_gzip_packed(msg_id, sequence, reader, request)
+        if code == 0xedab447b:  # bad_server_salt
+            return self.handle_bad_server_salt(msg_id, sequence, reader, request)
+        if code == 0xa7eff811:  # bad_msg_notification
+            return self.handle_bad_msg_notification(msg_id, sequence, reader)
+
+        # If the code is not parsed manually, then it was parsed by the code generator!
+        # In this case, we will simply treat the incoming TLObject as an Update,
+        # if we can first find a matching TLObject
+        if code in tlobjects.keys():
+            return self.handle_update(msg_id, sequence, reader)
+
+        print('Unknown message: {}'.format(hex(code)))
+        return False
+
+    # endregion
+
+    # region Message handling
+
+    def handle_update(self, msg_id, sequence, reader):
+        tlobject = reader.tgread_object()
+        for handler in self.on_update_handlers:
+            handler(tlobject)
+
+        return False
+
+    def handle_container(self, msg_id, sequence, reader, request):
+        code = reader.read_int(signed=False)
+        size = reader.read_int()
+        for _ in range(size):
+            inner_msg_id = reader.read_long(signed=False)
+            inner_sequence = reader.read_int()
+            inner_length = reader.read_int()
+            begin_position = reader.tell_position()
+
+            if not self.process_msg(inner_msg_id, sequence, reader, request):
+                reader.set_position(begin_position + inner_length)
+
+        return False
+
+    def handle_bad_server_salt(self, msg_id, sequence, reader, 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
+
+        if request is None:
+            raise ValueError('Tried to handle a bad server salt with no request specified')
+
+        # Resend
+        self.send(request, resend=True)
+
+        return True
+
+    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()
+
+        error_code = reader.read_int()
+        raise BadMessageError(error_code)
+
+    def handle_rpc_result(self, msg_id, sequence, reader, request):
+        if not request:
+            raise ValueError('RPC results should only happen after a request was sent')
+
+        code = reader.read_int(signed=False)
+        request_id = reader.read_long(signed=False)
+        inner_code = reader.read_int(signed=False)
+
+        if request_id == request.msg_id:
+            request.confirm_received = True
+
+        if inner_code == 0x2144ca19:  # RPC Error
+            error = RPCError(code=reader.read_int(), message=reader.tgread_string())
+            if error.must_resend:
+                request.confirm_received = False
+
+            if error.message.startswith('FLOOD_WAIT_'):
+                print('Should wait {}s. Sleeping until then.'.format(error.additional_data))
+                sleep(error.additional_data)
+
+            elif error.message.startswith('PHONE_MIGRATE_'):
+                raise InvalidDCError(error.additional_data)
+
+            else:
+                raise error
+        else:
+            if inner_code == 0x3072cfa1:  # GZip packed
+                unpacked_data = gzip.decompress(reader.tgread_bytes())
+                with BinaryReader(unpacked_data) as compressed_reader:
+                    request.on_response(compressed_reader)
+            else:
+                reader.seek(-4)
+                request.on_response(reader)
+
+    def handle_gzip_packed(self, msg_id, sequence, reader, request):
+        code = reader.read_int(signed=False)
+        packed_data = reader.tgread_bytes()
+        unpacked_data = gzip.decompress(packed_data)
+
+        with BinaryReader(unpacked_data) as compressed_reader:
+            return self.process_msg(msg_id, sequence, compressed_reader, request)
+
+    # endregion
+
+    def set_listen_for_updates(self, enabled):
+        if enabled:
+            if not self.updates_thread_running:
+                self.updates_thread_running = True
+                self.updates_thread_receiving = False
+
+                self.updates_thread.start()
+        else:
+            self.updates_thread_running = False
+            if self.updates_thread_receiving:
+                self.transport.cancel_receive()
+
+    def updates_thread_method(self):
+        """This method will run until specified and listen for incoming updates"""
+        while self.updates_thread_running:
+            with self.lock:
+                try:
+                    self.updates_thread_receiving = True
+                    seq, body = self.transport.receive()
+                    message, remote_msg_id, remote_sequence = self.decode_msg(body)
+
+                    with BinaryReader(message) as reader:
+                        self.process_msg(remote_msg_id, remote_sequence, reader)
+
+                except ReadCancelledError:
+                    pass
+
+            self.updates_thread_receiving = False
+
+            # If we are here, it is because the read was cancelled
+            # Sleep a bit just to give enough time for the other thread
+            # to acquire the lock. No need to sleep if we're not running anymore
+            if self.updates_thread_running:
+                sleep(0.1)

telethon/network/tcp_client.py

@@ -0,0 +1,75 @@
+# Python rough implementation of a C# TCP client
+import socket
+import time
+from threading import Lock
+
+from telethon.errors import ReadCancelledError
+from telethon.utils import BinaryWriter
+
+
+class TcpClient:
+    def __init__(self):
+        self.connected = False
+        self.socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
+
+        # Support for multi-threading advantages and safety
+        self.cancelled = False  # Has the read operation been cancelled?
+        self.delay = 0.1  # Read delay when there was no data available
+        self.lock = Lock()
+
+    def connect(self, ip, port):
+        """Connects to the specified IP and port number"""
+        self.socket.connect((ip, port))
+        self.connected = True
+
+    def close(self):
+        """Closes the connection"""
+        self.socket.close()
+        self.connected = False
+        self.socket.setblocking(True)
+
+    def write(self, data):
+        """Writes (sends) the specified bytes to the connected peer"""
+
+        # Ensure that only one thread can send data at once
+        with self.lock:
+            # Set blocking so it doesn't error
+            self.socket.setblocking(True)
+            self.socket.sendall(data)
+
+    def read(self, buffer_size):
+        """Reads (receives) the specified bytes from the connected peer"""
+
+        # Ensure that only one thread can receive data at once
+        with self.lock:
+            # Ensure it is not cancelled at first, so we can enter the loop
+            self.cancelled = False
+
+            # Set non-blocking so it can be cancelled
+            self.socket.setblocking(False)
+
+            with BinaryWriter() as writer:
+                while writer.written_count < buffer_size:
+                    # Only do cancel if no data was read yet
+                    # Otherwise, carry on reading and finish
+                    if self.cancelled and writer.written_count == 0:
+                        raise ReadCancelledError()
+
+                    try:
+                        # When receiving from the socket, we may not receive all the data at once
+                        # This is why we need to keep checking to make sure that we receive it all
+                        left_count = buffer_size - writer.written_count
+                        partial = self.socket.recv(left_count)
+                        writer.write(partial)
+
+                    except BlockingIOError:
+                        # There was no data available for us to read. Sleep a bit
+                        time.sleep(self.delay)
+
+                # If everything went fine, return the read bytes
+                return writer.get_bytes()
+
+    def cancel_read(self):
+        """Cancels the read operation IF it hasn't yet
+           started, raising a ReadCancelledError"""
+        self.cancelled = True

telethon/network/tcp_transport.py

@@ -0,0 +1,67 @@
+from binascii import crc32
+from telethon.network import TcpClient
+from telethon.errors import *
+from telethon.utils import BinaryWriter
+
+
+class TcpTransport:
+    def __init__(self, ip_address, port):
+        self.tcp_client = TcpClient()
+        self.send_counter = 0
+
+        self.tcp_client.connect(ip_address, port)
+
+    # Original reference: https://core.telegram.org/mtproto#tcp-transport
+    # The packets are encoded as: total length, sequence number, packet and checksum (CRC32)
+    def send(self, packet):
+        """Sends the given packet (bytes array) to the connected peer"""
+        if not self.tcp_client.connected:
+            raise ConnectionError('Client not connected to server.')
+
+        with BinaryWriter() as writer:
+            writer.write_int(len(packet) + 12)  # 12 = size_of (integer) * 3
+            writer.write_int(self.send_counter)
+            writer.write(packet)
+
+            crc = crc32(writer.get_bytes())
+            writer.write_int(crc, signed=False)
+
+            self.tcp_client.write(writer.get_bytes())
+            self.send_counter += 1
+
+    def receive(self):
+        """Receives a TCP message (tuple(sequence number, body)) from the connected peer"""
+
+        # First read everything we need
+        packet_length_bytes = self.tcp_client.read(4)
+        packet_length = int.from_bytes(packet_length_bytes, byteorder='little')
+
+        seq_bytes = self.tcp_client.read(4)
+        seq = int.from_bytes(seq_bytes, byteorder='little')
+
+        body = self.tcp_client.read(packet_length - 12)
+
+        checksum = int.from_bytes(self.tcp_client.read(4), byteorder='little', signed=False)
+
+        # Then perform the checks
+        rv = packet_length_bytes + seq_bytes + body
+        valid_checksum = crc32(rv)
+
+        if checksum != valid_checksum:
+            raise InvalidChecksumError(checksum, valid_checksum)
+
+        # If we passed the tests, we can then return a valid TCP message
+        return seq, body
+
+    def close(self):
+        if self.tcp_client.connected:
+            self.tcp_client.close()
+
+    def cancel_receive(self):
+        """Cancels (stops) trying to receive from the
+        remote peer and raises a ReadCancelledError"""
+        self.tcp_client.cancel_read()
+
+    def get_client_delay(self):
+        """Gets the client read delay"""
+        return self.tcp_client.delay