Add support for tcp_obfuscated on the Connection class (#112)

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

telethon/extensions/tcp_client_obfuscated.py

@@ -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()

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)
+        body = self.read(packet_length - 12)
-        checksum = int.from_bytes(self.conn.read(4, self.timeout), 'little')
+        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