From 1dac866118d587b2fbe196f1833fcaa6c768519f Mon Sep 17 00:00:00 2001 From: Lonami Date: Fri, 26 Aug 2016 12:58:53 +0200 Subject: [PATCH] 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 | 21 +++ README.md | 9 + main.py | 3 + network/__init__.py | 0 network/mtproto_plain_sender.py | 42 +++++ network/mtproto_sender.py | 290 ++++++++++++++++++++++++++++++++ network/tcp_client.py | 20 +++ network/tcp_message.py | 61 +++++++ network/tcp_transport.py | 52 ++++++ requests/__init__.py | 0 requests/ack_request.py | 20 +++ requests/mtproto_request.py | 39 +++++ utils/__init__.py | 0 utils/binary_reader.py | 87 ++++++++++ utils/binary_writer.py | 93 ++++++++++ utils/helpers.py | 56 ++++++ 16 files changed, 793 insertions(+) create mode 100644 LICENSE create mode 100644 README.md create mode 100644 main.py create mode 100644 network/__init__.py create mode 100644 network/mtproto_plain_sender.py create mode 100644 network/mtproto_sender.py create mode 100644 network/tcp_client.py create mode 100644 network/tcp_message.py create mode 100644 network/tcp_transport.py create mode 100644 requests/__init__.py create mode 100644 requests/ack_request.py create mode 100644 requests/mtproto_request.py create mode 100644 utils/__init__.py create mode 100644 utils/binary_reader.py create mode 100644 utils/binary_writer.py create mode 100644 utils/helpers.py diff --git a/LICENSE b/LICENSE new file mode 100644 index 00000000..8aa26455 --- /dev/null +++ b/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. diff --git a/README.md b/README.md new file mode 100644 index 00000000..d2e87b25 --- /dev/null +++ b/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 ` 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)! diff --git a/main.py b/main.py new file mode 100644 index 00000000..06aba510 --- /dev/null +++ b/main.py @@ -0,0 +1,3 @@ + +if __name__ == '__main__': + print('Hello worldz! Wooho... This is harder than it looks!') diff --git a/network/__init__.py b/network/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/network/mtproto_plain_sender.py b/network/mtproto_plain_sender.py new file mode 100644 index 00000000..a2e81398 --- /dev/null +++ b/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 diff --git a/network/mtproto_sender.py b/network/mtproto_sender.py new file mode 100644 index 00000000..345fa2c4 --- /dev/null +++ b/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) diff --git a/network/tcp_client.py b/network/tcp_client.py new file mode 100644 index 00000000..ef1fe471 --- /dev/null +++ b/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) diff --git a/network/tcp_message.py b/network/tcp_message.py new file mode 100644 index 00000000..f08a8b45 --- /dev/null +++ b/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) diff --git a/network/tcp_transport.py b/network/tcp_transport.py new file mode 100644 index 00000000..aa33a131 --- /dev/null +++ b/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() diff --git a/requests/__init__.py b/requests/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/requests/ack_request.py b/requests/ack_request.py new file mode 100644 index 00000000..3f208f05 --- /dev/null +++ b/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 diff --git a/requests/mtproto_request.py b/requests/mtproto_request.py new file mode 100644 index 00000000..abdba99c --- /dev/null +++ b/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 diff --git a/utils/__init__.py b/utils/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/utils/binary_reader.py b/utils/binary_reader.py new file mode 100644 index 00000000..86f67310 --- /dev/null +++ b/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 diff --git a/utils/binary_writer.py b/utils/binary_writer.py new file mode 100644 index 00000000..5b86e552 --- /dev/null +++ b/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() diff --git a/utils/helpers.py b/utils/helpers.py new file mode 100644 index 00000000..39ef37ea --- /dev/null +++ b/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]