Initial release

The initial release contains the most basic implementation of TLSharp core.
This is also fully untested, since no test can be done until more work is done.

LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) [year] [fullname]
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

README.md

@@ -0,0 +1,9 @@
+# Telethon
+**Telethon** is Telegram client implementation in Python. This project is **completely based** on [TLSharp](https://github.com/sochix/TLSharp), so please, also have a look to the original project!
+
+### Requirements
+This project requires the following Python modules, which can be installed by issuing `sudo -H pip install <module>` on a Linux terminal:
+- `pyaes` ([GitHub](https://github.com/ricmoo/pyaes), [package index](https://pypi.python.org/pypi/pyaes))
+
+### We need your help!
+As of now, the project is fully **untested** and with many pending things to do. If you know both Python and C#, please don't think it twice and help us (me)!

main.py

@@ -0,0 +1,3 @@
+
+if __name__ == '__main__':
+    print('Hello worldz! Wooho... This is harder than it looks!')

network/mtproto_plain_sender.py

@@ -0,0 +1,42 @@
+
+import time
+from utils.binary_writer import BinaryWriter
+from utils.binary_reader import BinaryReader
+
+
+class MtProtoPlainSender:
+
+    def __init__(self, transport):
+        self._sequence = 0
+        self._time_offset = 0
+        self._last_msg_id = 0
+        self._transport = transport
+
+    def send(self, data):
+        with BinaryWriter() as writer:
+            writer.write_long(0)
+            writer.write_int(self.get_new_msg_id())
+            writer.write_int(len(data))
+            writer.write(data)
+
+            packet = writer.get_bytes()
+            self._transport.send(packet)
+
+    def receive(self):
+        result = self._transport.receive()
+        with BinaryReader(result.body) as reader:
+            auth_key_id = reader.read_long()
+            message_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
+
+        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

network/mtproto_sender.py

@@ -0,0 +1,290 @@
+import re
+import zlib
+import pyaes
+from time import sleep
+
+from utils.binary_writer import BinaryWriter
+from utils.binary_reader import BinaryReader
+from requests.ack_request import AckRequest
+import utils.helpers as helpers
+
+
+class MtProtoSender:
+    def __init__(self, transport, session):
+        self._transport = transport
+        self._session = session
+
+        self.need_confirmation = []
+
+    def change_transport(self, transport):
+        self._transport = transport
+
+    def generate_sequence(self, confirmed):
+        if confirmed:
+            result = self._session.sequence * 2 + 1
+            self._session.sequence += 1
+            return result
+        else:
+            return self._session.sequence * 2
+
+    # TODO async?
+    def send(self, request):
+        if self.need_confirmation:
+            ack_request = AckRequest(self.need_confirmation)
+
+            with BinaryWriter() as writer:
+                ack_request.on_send(writer)
+                self.send_packet(writer.get_bytes(), ack_request)
+                del self.need_confirmation[:]
+
+        with BinaryWriter() as writer:
+            request.on_send(writer)
+            self.send_packet(writer.get_bytes(), request)
+
+        self._session.save()
+
+    def send_packet(self, packet, request):
+        request.message_id = self._session.get_new_msg_id()
+
+        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(request.message_id)
+            writer.write_int(self.generate_sequence(request.confirmed))
+            writer.write_int(len(packet))
+            writer.write(packet)
+
+            msg_key = helpers.calc_msg_key(writer.get_bytes())
+
+            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())
+
+        with BinaryWriter() as writer:
+            # TODO is it unsigned long?
+            writer.write_long(self._session.auth_key.id, signed=False)
+            writer.write(msg_key)
+            writer.write(cipher_text)
+
+            self._transport.send(writer.get_bytes())
+
+    def decode_msg(self, body):
+        message = None
+        remote_message_id = None
+        remote_sequence = None
+
+        with BinaryReader(body) as reader:
+            if len(body) < 8:
+                raise BufferError("Can't decode packet")
+
+            # 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 = 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_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
+
+    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):
+        # TODO Check salt, session_id and sequence_number
+        self.need_confirmation.append(message_id)
+
+        code = reader.read_int(signed=False)
+        reader.seek(-4)
+
+        if code == 0x73f1f8dc:  # Container
+            return self.handle_container(message_id, sequence, reader, mtproto_request)
+        if code == 0x7abe77ec:  # Ping
+            return self.handle_ping(message_id, sequence, reader)
+        if code == 0x347773c5:  # pong
+            return self.handle_pong(message_id, sequence, reader)
+        if code == 0xae500895:  # future_salts
+            return self.handle_future_salts(message_id, sequence, reader)
+        if code == 0x9ec20908:  # new_session_created
+            return self.handle_new_session_created(message_id, sequence, reader)
+        if code == 0x62d6b459:  # msgs_ack
+            return self.handle_msgs_ack(message_id, sequence, reader)
+        if code == 0xedab447b:  # bad_server_salt
+            return self.handle_bad_server_salt(message_id, sequence, reader, mtproto_request)
+        if code == 0xa7eff811:  # bad_msg_notification
+            return self.handle_bad_msg_notification(message_id, sequence, reader)
+        if code == 0x276d3ec6:  # msg_detailed_info
+            return self.hangle_msg_detailed_info(message_id, sequence, reader)
+        if code == 0xf35c6d01:  # rpc_result
+            return self.handle_rpc_result(message_id, sequence, reader, mtproto_request)
+        if code == 0x3072cfa1:  # gzip_packed
+            return self.handle_gzip_packed(message_id, sequence, reader, mtproto_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)
+
+        # TODO Log unknown message code
+        return False
+
+    def handle_update(self, message_id, sequence, reader):
+        return False
+
+    def handle_container(self, message_id, sequence, reader, mtproto_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()
+            try:
+                if not self.process_msg(inner_msg_id, sequence, reader, mtproto_request):
+                    reader.set_position(begin_position + inner_length)
+
+            except:
+                reader.set_position(begin_position + inner_length)
+
+        return False
+
+    def handle_ping(self, message_id, sequence, reader):
+        return False
+
+    def handle_pong(self, message_id, sequence, reader):
+        return False
+
+    def handle_future_salts(self, message_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):
+        return False
+
+    def handle_msgs_ack(self, message_id, sequence, reader):
+        return False
+
+    def handle_bad_server_salt(self, message_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
+
+        # Resend
+        self.send(mtproto_request)
+
+        return True
+
+    def handle_bad_msg_notification(self, message_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()
+
+        if error_code == 16:
+            raise RuntimeError("msg_id too low (most likely, client time is wrong it would be worthwhile to "
+                               "synchronize it using msg_id notifications and re-send the original message "
+                               "with the “correct” msg_id or wrap it in a container with a new msg_id if the "
+                               "original message had waited too long on the client to be transmitted)")
+        if error_code == 17:
+            raise RuntimeError("msg_id too high (similar to the previous case, the client time has to be "
+                               "synchronized, and the message re-sent with the correct msg_id)")
+        if error_code == 18:
+            raise RuntimeError("Incorrect two lower order msg_id bits (the server expects client message msg_id "
+                               "to be divisible by 4)")
+        if error_code == 19:
+            raise RuntimeError("Container msg_id is the same as msg_id of a previously received message "
+                               "(this must never happen)")
+        if error_code == 20:
+            raise RuntimeError("Message too old, and it cannot be verified whether the server has received a "
+                               "message with this msg_id or not")
+        if error_code == 32:
+            raise RuntimeError("msg_seqno too low (the server has already received a message with a lower "
+                               "msg_id but with either a higher or an equal and odd seqno)")
+        if error_code == 33:
+            raise RuntimeError("msg_seqno too high (similarly, there is a message with a higher msg_id but with "
+                               "either a lower or an equal and odd seqno)")
+        if error_code == 34:
+            raise RuntimeError("An even msg_seqno expected (irrelevant message), but odd received")
+        if error_code == 35:
+            raise RuntimeError("Odd msg_seqno expected (relevant message), but even received")
+        if error_code == 48:
+            raise RuntimeError("Incorrect server salt (in this case, the bad_server_salt response is received with "
+                               "the correct salt, and the message is to be re-sent with it)")
+        if error_code == 64:
+            raise RuntimeError("Invalid container")
+
+        raise NotImplementedError('This should never happen!')
+
+    def hangle_msg_detailed_info(self, message_id, sequence, reader):
+        return False
+
+    def handle_rpc_result(self, message_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:
+            mtproto_request.confirm_received = True
+
+        inner_code = reader.read_int(signed=False)
+        if inner_code == 0x2144ca19:  # RPC Error
+            error_code = reader.read_int()
+            error_msg = reader.tgread_string()
+
+            if error_msg.startswith('FLOOD_WAIT_'):
+                seconds = int(re.search(r'\d+', error_msg).group(0))
+                print('Should wait {}s. Sleeping until then.')
+                sleep(seconds)
+
+            elif error_msg.startswith('PHONE_MIGRATE_'):
+                dc_index = int(re.search(r'\d+', error_msg).group(0))
+                raise ConnectionError('Your phone number registered to {} dc. Please update settings. '
+                                      'See https://github.com/sochix/TLSharp#i-get-an-error-migrate_x '
+                                      'for details.'.format(dc_index))
+            else:
+                raise ValueError(error_msg)
+
+        elif inner_code == 0x3072cfa1:  # GZip packed
+            try:
+                packed_data = reader.tgread_bytes()
+                unpacked_data = zlib.decompress(packed_data)
+
+                with BinaryReader(unpacked_data) as compressed_reader:
+                    mtproto_request.on_response(compressed_reader)
+
+            except:
+                pass
+
+        else:
+            reader.seek(-4)
+            mtproto_request.on_response(reader)
+
+    def handle_gzip_packed(self, message_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)

network/tcp_client.py

@@ -0,0 +1,20 @@
+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):
+        self.socket.connect((ip, port))
+
+    def close(self):
+        self.socket.close()
+
+    def write(self, data):
+        self.socket.send(data)
+
+    def read(self, buffer_size):
+        self.socket.recv(buffer_size)

network/tcp_message.py

@@ -0,0 +1,61 @@
+
+from zlib import crc32
+
+from utils.binary_writer import BinaryWriter
+from utils.binary_reader import BinaryReader
+
+
+class TcpMessage:
+
+    def __init__(self, seq_number, body):
+        """
+        :param seq_number: Sequence number
+        :param body: Message body byte array
+        """
+        if body is None:
+            raise ValueError('body cannot be None')
+
+        self.sequence_number = seq_number
+        self.body = body
+
+    def encode(self):
+        with BinaryWriter() as writer:
+            ''' https://core.telegram.org/mtproto#tcp-transport
+
+                4 length bytes are added at the front
+                (to include the length, the sequence number, and CRC32; always divisible by 4)
+                and 4 bytes with the packet sequence number within this TCP connection
+                (the first packet sent is numbered 0, the next one 1, etc.),
+                and 4 CRC32 bytes at the end (length, sequence number, and payload together).
+            '''
+            writer.write_int(len(self.body) + 12)
+            writer.write_int(self.sequence_number)
+            writer.write(self.body)
+            writer.flush()  # Flush so we can get the buffer in the CRC
+
+            crc = crc32(writer.get_bytes()[0:8 + len(self.body)])
+            writer.write_int(crc, signed=False)
+
+            return writer.get_bytes()
+
+    def decode(self, body):
+        if body is None:
+            raise ValueError('body cannot be None')
+
+        if len(body) < 12:
+            raise ValueError('Wrong size of input packet')
+
+        with BinaryReader(body) as reader:
+            packet_len = int.from_bytes(reader.read(4), byteorder='big')
+            if packet_len < 12:
+                raise ValueError('Invalid packet length: {}'.format(packet_len))
+
+            seq = reader.read_int()
+            packet = reader.read(packet_len - 12)
+            checksum = reader.read_int()
+
+            valid_checksum = crc32(body[:packet_len - 4])
+            if checksum != valid_checksum:
+                raise ValueError('Invalid checksum, skip')
+
+            return TcpMessage(seq, packet)

network/tcp_transport.py

@@ -0,0 +1,52 @@
+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
+
+        self._tcp_client.connect(ip_address, port)
+
+    def send(self, packet):
+        """
+        :param packet: Bytes array representing the packet to be sent
+        """
+        if not self._tcp_client.connected:
+            raise ConnectionError('Client not connected to server.')
+
+        tcp_message = TcpMessage(self._send_counter, packet)
+
+        # TODO async? and receive too, of course
+        self._tcp_client.write(tcp_message.encode())
+
+        self._send_counter += 1
+
+    def receive(self):
+        # First read everything
+        packet_length_bytes = self._tcp_client.read(4)
+        packet_length = int.from_bytes(packet_length_bytes, byteorder='big')
+
+        seq_bytes = self._tcp_client.read(4)
+        seq = int.from_bytes(seq_bytes, byteorder='big')
+
+        body = self._tcp_client.read(packet_length - 12)
+
+        checksum = int.from_bytes(self._tcp_client.read(4), byteorder='big')
+
+        # Then perform the checks
+        rv = packet_length_bytes + seq_bytes + body
+        valid_checksum = crc32(rv) & 0xFFFFFFFF  # Ensure it's unsigned (http://stackoverflow.com/a/30092291/4759433)
+
+        if checksum != valid_checksum:
+            raise ValueError('Invalid checksum, skip')
+
+        return TcpMessage(seq, body)
+
+    def dispose(self):
+        if self._tcp_client.connected:
+            self._tcp_client.close()

requests/ack_request.py

@@ -0,0 +1,20 @@
+from requests.mtproto_request import MTProtoRequest
+
+
+class AckRequest(MTProtoRequest):
+    def __init__(self, msgs):
+        super().__init__()
+        self.msgs = msgs
+
+    def on_send(self, writer):
+        writer.write_int(0x62d6b459)  # msgs_ack
+        writer.write_int(0x1cb5c415)  # vector
+        writer.write_int(len(self.msgs))
+        for msg_id in self.msgs:
+            writer.write_int(msg_id, signed=False)
+
+    def on_response(self, reader):
+        pass
+
+    def on_exception(self, exception):
+        pass

requests/mtproto_request.py

@@ -0,0 +1,39 @@
+from datetime import datetime, timedelta
+
+
+class MTProtoRequest:
+    def __init__(self):
+        self.sent = False
+
+        self.msg_id = 0  # Long
+        self.sequence = 0
+
+        self.dirty = False
+        self.send_time = None
+        self.confirm_received = False
+
+        # These should be overrode
+        self.confirmed = False
+        self.responded = False
+
+    # These should not be overrode
+    def on_send_success(self):
+        self.send_time = datetime.now()
+        self.sent = True
+
+    def on_confirm(self):
+        self.confirm_received = True
+
+    def need_resend(self):
+        return self.dirty or (self.confirmed and not self.confirm_received and
+                              datetime.now() - self.send_time > timedelta(seconds=3))
+
+    # These should be overrode
+    def on_send(self, writer):
+        pass
+
+    def on_response(self, reader):
+        pass
+
+    def on_exception(self, exception):
+        pass

utils/binary_reader.py

@@ -0,0 +1,87 @@
+from io import BytesIO, BufferedReader
+import os
+
+
+class BinaryReader:
+    """
+    Small utility class to read binary data.
+    Also creates a "Memory Stream" if necessary
+    """
+    def __init__(self, data=None, stream=None):
+        if data:
+            self.stream = BytesIO(data)
+        elif stream:
+            self.stream = stream
+        else:
+            raise ValueError("Either bytes or a stream must be provided")
+
+        self.reader = BufferedReader(self.stream)
+
+    # region Reading
+
+    def read_int(self, signed=True):
+        return int.from_bytes(self.reader.read(4), signed=signed, byteorder='big')
+
+    def read_long(self, signed=True):
+        return int.from_bytes(self.reader.read(8), signed=signed, byteorder='big')
+
+    def read(self, length):
+        return self.reader.read(length)
+
+    def get_bytes(self):
+        return self.stream.getbuffer()
+
+    # endregion
+
+    # region Telegram custom reading
+
+    def tgread_bytes(self):
+        first_byte = self.read(1)
+        if first_byte == 254:
+            length = self.read(1) | (self.read(1) << 8) | (self.read(1) << 16)
+            padding = length % 4
+        else:
+            length = first_byte
+            padding = (length + 1) % 4
+
+        data = self.read(length)
+        if padding > 0:
+            padding = 4 - padding
+            self.read(padding)
+
+        return data
+
+    def tgread_string(self):
+        return str(self.tgread_bytes(), encoding='utf-8')
+
+    # endregion
+
+    def close(self):
+        self.reader.close()
+        # TODO Do I need to close the underlying stream?
+
+    # region Position related
+
+    def tell_position(self):
+        """Tells the current position on the stream"""
+        return self.reader.tell()
+
+    def set_position(self, position):
+        """Sets the current position on the stream"""
+        self.reader.seek(position)
+
+    def seek(self, offset):
+        """Seeks the stream position given an offset from the current position. May be negative"""
+        self.reader.seek(offset, os.SEEK_CUR)
+
+    # endregion
+
+    # region with block
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        self.close()
+
+    # endregion

utils/binary_writer.py

@@ -0,0 +1,93 @@
+from io import BytesIO, BufferedWriter
+from struct import pack
+
+
+class BinaryWriter:
+    """
+    Small utility class to write binary data.
+    Also creates a "Memory Stream" if necessary
+    """
+
+    def __init__(self, stream=None):
+        if not stream:
+            stream = BytesIO()
+
+        self.stream = stream
+        self.writer = BufferedWriter(self.stream)
+
+    # region Writing
+
+    def write_byte(self, byte):
+        self.writer.write(pack('B', byte))
+
+    def write_int(self, integer, signed=True):
+        if not signed:
+            integer &= 0xFFFFFFFF  # Ensure it's unsigned (see http://stackoverflow.com/a/30092291/4759433)
+        self.writer.write(pack('I', integer))
+
+    def write_long(self, long, signed=True):
+        if not signed:
+            long &= 0xFFFFFFFFFFFFFFFF
+        self.writer.write(pack('Q', long))
+
+    def write(self, data):
+        self.writer.write(data)
+
+    # endregion
+
+    # region Telegram custom writing
+
+    def tgwrite_bytes(self, data):
+
+        if len(data) < 254:
+            padding = (len(data) + 1) % 4
+            if padding != 0:
+                padding = 4 - padding
+
+            self.write(bytes([len(data)]))
+            self.write(data)
+
+        else:
+            padding = len(data) % 4
+            if padding != 0:
+                padding = 4 - padding
+
+            # TODO ensure that _this_ is right (it appears to be)
+            self.write(bytes([254]))
+            self.write(bytes([len(data) % 256]))
+            self.write(bytes([(len(data) >> 8) % 256]))
+            self.write(bytes([(len(data) >> 16) % 256]))
+            self.write(data)
+
+            """ Original:
+                    binaryWriter.Write((byte)254);
+                    binaryWriter.Write((byte)(bytes.Length));
+                    binaryWriter.Write((byte)(bytes.Length >> 8));
+                    binaryWriter.Write((byte)(bytes.Length >> 16));
+            """
+
+        self.write(bytes(padding))
+
+    def tgwrite_string(self, string):
+        return self.tgwrite_bytes(string.encode('utf-8'))
+
+    # endregion
+
+    def flush(self):
+        self.writer.flush()
+
+    def close(self):
+        self.writer.close()
+        # TODO Do I need to close the underlying stream?
+
+    def get_bytes(self, flush=True):
+        if flush:
+            self.writer.flush()
+        self.stream.getbuffer()
+
+    # with block
+    def __enter__(self):
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        self.close()

utils/helpers.py

@@ -0,0 +1,56 @@
+import random
+from utils.binary_writer import BinaryWriter
+from hashlib import sha1
+
+
+def generate_random_long(signed=True):
+    result = random.getrandbits(64)
+    if not signed:
+        result &= 0xFFFFFFFFFFFFFFFF  # Ensure it's unsigned
+
+    return result
+
+
+def generate_random_bytes(count):
+    with BinaryWriter() as writer:
+        for _ in range(count):
+            writer.write(random.getrandbits(8))
+
+    return writer.get_bytes()
+
+
+def calc_key(shared_key, msg_key, client):
+    x = 0 if client else 8
+
+    buffer = [0] * 48
+    buffer[0:16] = msg_key
+    buffer[16:48] = shared_key[x:x + 32]
+    sha1a = sha1(buffer)
+
+    buffer[0:16] = shared_key[x + 32:x + 48]
+    buffer[16:32] = msg_key
+    buffer[32:48] = shared_key[x + 48:x + 64]
+    sha1b = sha1(buffer)
+
+    buffer[0:32] = shared_key[x + 64:x + 96]
+    buffer[32:48] = msg_key
+    sha1c = sha1(buffer)
+
+    buffer[0:16] = msg_key
+    buffer[16:48] = shared_key[x + 96:x + 128]
+    sha1d = sha1(buffer)
+
+    key = sha1a[0:8] + sha1b[8:20] + sha1c[4:16]
+    iv = sha1a[8:20] + sha1b[0:8] + sha1c[16:20] + sha1d[0:8]
+
+    return key, iv
+
+
+def calc_msg_key(data):
+    return sha1(data)[4:20]
+
+
+def calc_msg_key_offset(data, offset, limit):
+    # TODO untested, may not be offset like this
+    # In the original code it was as parameters for the sha function, not slicing the array
+    return sha1(data[offset:offset + limit])[4:20]