diff --git a/telethon/extensions/__init__.py b/telethon/extensions/__init__.py index 46de4467..06e2f087 100644 --- a/telethon/extensions/__init__.py +++ b/telethon/extensions/__init__.py @@ -7,4 +7,3 @@ strings, bytes, etc.) from .binary_writer import BinaryWriter from .binary_reader import BinaryReader from .tcp_client import TcpClient -from .tcp_client_obfuscated import TcpClientObfuscated diff --git a/telethon/extensions/tcp_client_obfuscated.py b/telethon/extensions/tcp_client_obfuscated.py deleted file mode 100644 index 302c388e..00000000 --- a/telethon/extensions/tcp_client_obfuscated.py +++ /dev/null @@ -1,171 +0,0 @@ -# Python rough implementation of a C# TCP client -import socket -import time -import os -from datetime import datetime, timedelta -from io import BytesIO, BufferedWriter -from threading import Event, Lock -import errno - -from ..crypto import AESModeCTR -from ..errors import ReadCancelledError - - -# Obfuscated messages secrets cannot start with any of these -OBFUSCATED_ANTI_KEYWORDS = (b'PVrG', b'GET ', b'POST', b'\xee' * 4) - - -class TcpClientObfuscated: - # TODO Avoid duplicating so much code - transport for TCPO - - def __init__(self, proxy=None): - self.connected = False - self._proxy = proxy - self._recreate_socket() - - # Support for multi-threading advantages and safety - self.cancelled = Event() # Has the read operation been cancelled? - self.delay = 0.1 # Read delay when there was no data available - self._lock = Lock() - - self.aes_encrypt = None - self.aes_decrypt = None - - def _recreate_socket(self): - if self._proxy is None: - self._socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM) - else: - import socks - self._socket = socks.socksocket(socket.AF_INET, socket.SOCK_STREAM) - if type(self._proxy) is dict: - self._socket.set_proxy(**self._proxy) - else: # tuple, list, etc. - self._socket.set_proxy(*self._proxy) - - def connect(self, ip, port, timeout): - """Connects to the specified IP and port number. - 'timeout' must be given in seconds - """ - if not self.connected: - self._socket.settimeout(timeout) - self._socket.connect((ip, port)) - self._socket.setblocking(False) - self.connected = True - - # TCP Obfuscated bits - while True: - random = os.urandom(64) - if (random[0] != b'\xef' and - random[:4] not in OBFUSCATED_ANTI_KEYWORDS and - random[4:4] != b'\0\0\0\0'): - # Invalid random generated - break - - random = list(random) - random[56] = random[57] = random[58] = random[59] = 0xef - random_reversed = random[55:7:-1] # Reversed (8, len=48) - - # encryption has "continuous buffer" enabled - encrypt_key = bytes(random[8:40]) - encrypt_iv = bytes(random[40:56]) - decrypt_key = bytes(random_reversed[:32]) - decrypt_iv = bytes(random_reversed[32:48]) - - self.aes_encrypt = AESModeCTR(encrypt_key, encrypt_iv) - self.aes_decrypt = AESModeCTR(decrypt_key, decrypt_iv) - - random[56:64] = self.aes_encrypt.encrypt(bytes(random))[56:64] - self._socket.sendall(bytes(random)) - - def close(self): - """Closes the connection""" - if self.connected: - try: - self._socket.shutdown(socket.SHUT_RDWR) - self._socket.close() - except OSError as e: - if e.errno != errno.ENOTCONN: - raise - - self.connected = False - self._recreate_socket() - - def write(self, data): - """Writes (sends) the specified bytes to the connected peer""" - data = self.aes_encrypt.encrypt(data) - - # Ensure that only one thread can send data at once - with self._lock: - try: - view = memoryview(data) - total_sent, total = 0, len(data) - while total_sent < total: - try: - sent = self._socket.send(view[total_sent:]) - if sent == 0: - raise ConnectionResetError( - 'The server has closed the connection.') - total_sent += sent - - except BlockingIOError: - time.sleep(self.delay) - except BrokenPipeError: - self.close() - raise - - def read(self, size, timeout=timedelta(seconds=5)): - """Reads (receives) a whole block of 'size bytes - from the connected peer. - - A timeout can be specified, which will cancel the operation if - no data has been read in the specified time. If data was read - and it's waiting for more, the timeout will NOT cancel the - operation. Set to None for no timeout - """ - - # 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.clear() - - # Set the starting time so we can - # calculate whether the timeout should fire - start_time = datetime.now() if timeout is not None else None - - with BufferedWriter(BytesIO(), buffer_size=size) as buffer: - bytes_left = size - while bytes_left != 0: - # Only do cancel if no data was read yet - # Otherwise, carry on reading and finish - if self.cancelled.is_set() and bytes_left == size: - raise ReadCancelledError() - - try: - partial = self._socket.recv(bytes_left) - if len(partial) == 0: - self.close() - raise ConnectionResetError( - 'The server has closed the connection.') - - buffer.write(partial) - bytes_left -= len(partial) - - except BlockingIOError as error: - # No data available yet, sleep a bit - time.sleep(self.delay) - - # Check if the timeout finished - if timeout is not None: - time_passed = datetime.now() - start_time - if time_passed > timeout: - raise TimeoutError( - 'The read operation exceeded the timeout.') from error - - # If everything went fine, return the read bytes - buffer.flush() - return self.aes_decrypt.encrypt(buffer.raw.getvalue()) - - def cancel_read(self): - """Cancels the read operation IF it hasn't yet - started, raising a ReadCancelledError""" - self.cancelled.set() diff --git a/telethon/network/connection.py b/telethon/network/connection.py index b7be09ce..8594ddd9 100644 --- a/telethon/network/connection.py +++ b/telethon/network/connection.py @@ -1,33 +1,51 @@ +import os from datetime import timedelta from zlib import crc32 +from ..crypto import AESModeCTR from ..extensions import BinaryWriter, TcpClient from ..errors import InvalidChecksumError class Connection: - def __init__(self, ip, port, mode='tcp_abridged', + """Represents an abstract connection (TCP, TCP abridged...). + 'mode' may be any of 'tcp_full', 'tcp_abridged'. + + Note that '.send()' and '.recv()' refer to messages, which + will be packed accordingly, whereas '.write()' and '.read()' + work on plain bytes, with no further additions. + """ + + def __init__(self, ip, port, mode='tcp_obfuscated', proxy=None, timeout=timedelta(seconds=5)): - """Represents an abstract connection (TCP, TCP abridged...). - 'mode' may be any of 'tcp_full', 'tcp_abridged' - """ self.ip = ip self.port = port self._mode = mode self.timeout = timeout - self._send_counter = 0 - # TODO Rename "TcpClient" as some sort of generic socket + self._send_counter = 0 + self._aes_encrypt, self._aes_decrypt = None, None + + # TODO Rename "TcpClient" as some sort of generic socket? self.conn = TcpClient(proxy=proxy) + # Sending messages if mode == 'tcp_full': setattr(self, 'send', self._send_tcp_full) setattr(self, 'recv', self._recv_tcp_full) - elif mode == 'tcp_abridged': + elif mode in ('tcp_abridged', 'tcp_obfuscated'): setattr(self, 'send', self._send_abridged) setattr(self, 'recv', self._recv_abridged) + # Writing and reading from the socket + if mode == 'tcp_obfuscated': + setattr(self, 'write', self._write_obfuscated) + setattr(self, 'read', self._read_obfuscated) + else: + setattr(self, 'write', self._write_plain) + setattr(self, 'read', self._read_plain) + def connect(self): self._send_counter = 0 self.conn.connect(self.ip, self.port, @@ -35,6 +53,35 @@ class Connection: if self._mode == 'tcp_abridged': self.conn.write(int.to_bytes(239, 1, 'little')) + elif self._mode == 'tcp_obfuscated': + self._setup_obfuscation() + + def _setup_obfuscation(self): + # Obfuscated messages secrets cannot start with any of these + keywords = (b'PVrG', b'GET ', b'POST', b'\xee' * 4) + while True: + random = os.urandom(64) + if (random[0] != b'\xef' and + random[:4] not in keywords and + random[4:4] != b'\0\0\0\0'): + # Invalid random generated + break + + random = list(random) + random[56] = random[57] = random[58] = random[59] = 0xef + random_reversed = random[55:7:-1] # Reversed (8, len=48) + + # encryption has "continuous buffer" enabled + encrypt_key = bytes(random[8:40]) + encrypt_iv = bytes(random[40:56]) + decrypt_key = bytes(random_reversed[:32]) + decrypt_iv = bytes(random_reversed[32:48]) + + self._aes_encrypt = AESModeCTR(encrypt_key, encrypt_iv) + self._aes_decrypt = AESModeCTR(decrypt_key, decrypt_iv) + + random[56:64] = self._aes_encrypt.encrypt(bytes(random))[56:64] + self.conn.write(bytes(random)) def is_connected(self): return self.conn.connected @@ -51,7 +98,7 @@ class Connection: """Gets the client read delay""" return self.conn.delay - # region Receive implementations + # region Receive message implementations def recv(self, **kwargs): """Receives and unpacks a message""" @@ -60,14 +107,14 @@ class Connection: raise ValueError('Invalid connection mode specified: ' + self._mode) def _recv_tcp_full(self, **kwargs): - packet_length_bytes = self.conn.read(4, self.timeout) + packet_length_bytes = self.read(4) packet_length = int.from_bytes(packet_length_bytes, 'little') - seq_bytes = self.conn.read(4, self.timeout) + seq_bytes = self.read(4) seq = int.from_bytes(seq_bytes, 'little') - body = self.conn.read(packet_length - 12, self.timeout) - checksum = int.from_bytes(self.conn.read(4, self.timeout), 'little') + body = self.read(packet_length - 12) + checksum = int.from_bytes(self.read(4), 'little') valid_checksum = crc32(packet_length_bytes + seq_bytes + body) if checksum != valid_checksum: @@ -76,15 +123,15 @@ class Connection: return body def _recv_abridged(self, **kwargs): - length = int.from_bytes(self.conn.read(1, self.timeout), 'little') + length = int.from_bytes(self.read(1), 'little') if length >= 127: - length = int.from_bytes(self.conn.read(3, self.timeout) + b'\0', 'little') + length = int.from_bytes(self.read(3) + b'\0', 'little') - return self.conn.read(length << 2) + return self.read(length << 2) # endregion - # region Send implementations + # region Send message implementations def send(self, message): """Encapsulates and sends the given message""" @@ -100,7 +147,7 @@ class Connection: writer.write(message) writer.write_int(crc32(writer.get_bytes()), signed=False) self._send_counter += 1 - self.conn.write(writer.get_bytes()) + self.write(writer.get_bytes()) def _send_abridged(self, message): with BinaryWriter() as writer: @@ -111,6 +158,34 @@ class Connection: writer.write_byte(127) writer.write(int.to_bytes(length, 3, 'little')) writer.write(message) - self.conn.write(writer.get_bytes()) + self.write(writer.get_bytes()) + + # endregion + + # region Read implementations + + def read(self, length): + raise ValueError('Invalid connection mode specified: ' + self._mode) + + def _read_plain(self, length): + return self.conn.read(length, timeout=self.timeout) + + def _read_obfuscated(self, length): + return self._aes_decrypt.encrypt( + self.conn.read(length, timeout=self.timeout) + ) + + # endregion + + # region Write implementations + + def write(self, data): + raise ValueError('Invalid connection mode specified: ' + self._mode) + + def _write_plain(self, data): + self.conn.write(data) + + def _write_obfuscated(self, data): + self.conn.write(self._aes_encrypt.encrypt(data)) # endregion