Merge pull request #1126 from seriyps/mtproto-proxy-other-protocols

Implement different MTProto proxy protocols

telethon/client/telegrambaseclient.py

@@ -9,7 +9,7 @@ from datetime import datetime, timezone
 from .. import version, __name__ as __base_name__
 from ..crypto import rsa
 from ..extensions import markdown
-from ..network import MTProtoSender, ConnectionTcpFull, ConnectionTcpMTProxy
+from ..network import MTProtoSender, ConnectionTcpFull, TcpMTProxy
 from ..sessions import Session, SQLiteSession, MemorySession
 from ..tl import TLObject, functions, types
 from ..tl.alltlobjects import LAYER
@@ -64,7 +64,7 @@ class TelegramBaseClient(abc.ABC):
 
         proxy (`tuple` | `list` | `dict`, optional):
             An iterable consisting of the proxy info. If `connection` is
-            `ConnectionTcpMTProxy`, then it should contain MTProxy credentials:
+            one of `MTProxy`, then it should contain MTProxy credentials:
             ``('hostname', port, 'secret')``. Otherwise, it's meant to store
             function parameters for PySocks, like ``(type, 'hostname', port)``.
             See https://github.com/Anorov/PySocks#usage-1 for more.
@@ -249,8 +249,8 @@ class TelegramBaseClient(abc.ABC):
 
         assert isinstance(connection, type)
         self._connection = connection
-        init_proxy = None if connection is not ConnectionTcpMTProxy else \
+        init_proxy = None if not issubclass(connection, TcpMTProxy) else \
-            types.InputClientProxy(*ConnectionTcpMTProxy.address_info(proxy))
+            types.InputClientProxy(*connection.address_info(proxy))
 
         # Used on connection. Capture the variables in a lambda since
         # exporting clients need to create this InvokeWithLayerRequest.

telethon/network/__init__.py

@@ -7,5 +7,7 @@ from .authenticator import do_authentication
 from .mtprotosender import MTProtoSender
 from .connection import (
     ConnectionTcpFull, ConnectionTcpIntermediate, ConnectionTcpAbridged,
-    ConnectionTcpObfuscated, ConnectionTcpMTProxy, ConnectionHttp
+    ConnectionTcpObfuscated, ConnectionTcpMTProxyAbridged,
+    ConnectionTcpMTProxyIntermediate,
+    ConnectionTcpMTProxyRandomizedIntermediate, ConnectionHttp, TcpMTProxy
 )

telethon/network/connection/__init__.py

@@ -2,5 +2,10 @@ from .tcpfull import ConnectionTcpFull
 from .tcpintermediate import ConnectionTcpIntermediate
 from .tcpabridged import ConnectionTcpAbridged
 from .tcpobfuscated import ConnectionTcpObfuscated
-from .tcpmtproxy import ConnectionTcpMTProxy
+from .tcpmtproxy import (
+    TcpMTProxy,
+    ConnectionTcpMTProxyAbridged,
+    ConnectionTcpMTProxyIntermediate,
+    ConnectionTcpMTProxyRandomizedIntermediate
+)
 from .http import ConnectionHttp

telethon/network/connection/connection.py

@@ -18,6 +18,10 @@ class Connection(abc.ABC):
     ``ConnectionError``, which will raise when attempting to send if
     the client is disconnected (includes remote disconnections).
     """
+    # this static attribute should be redefined by `Connection` subclasses and
+    # should be one of `PacketCodec` implementations
+    packet_codec = None
+
     def __init__(self, ip, port, dc_id, *, loop, loggers, proxy=None):
         self._ip = ip
         self._port = port
@@ -30,6 +34,8 @@ class Connection(abc.ABC):
         self._connected = False
         self._send_task = None
         self._recv_task = None
+        self._codec = None
+        self._obfuscation = None  # TcpObfuscated and MTProxy
         self._send_queue = asyncio.Queue(1)
         self._recv_queue = asyncio.Queue(1)
 
@@ -76,6 +82,7 @@ class Connection(abc.ABC):
                 await asyncio.open_connection(sock=s, loop=self._loop)
 
         self._connected = True
+        self._codec = self.packet_codec(self)
         self._init_conn()
         await self._writer.drain()
 
@@ -182,27 +189,71 @@ class Connection(abc.ABC):
         data to Telegram to indicate which connection mode will
         be used.
         """
+        if self._codec.tag:
+            self._writer.write(self._codec.tag)
 
-    @abc.abstractmethod
     def _send(self, data):
-        """
+        self._writer.write(self._codec.encode_packet(data))
-        This method should be implemented differently under each
-        connection mode and serialize the data into the packet
-        the way it should be sent through `self._writer`.
-        """
-        raise NotImplementedError
 
-    @abc.abstractmethod
     async def _recv(self):
-        """
+        return await self._codec.read_packet(self._reader)
-        This method should be implemented differently under each
-        connection mode and deserialize the data from the packet
-        the way it should be read from `self._reader`.
-        """
-        raise NotImplementedError
 
     def __str__(self):
         return '{}:{}/{}'.format(
             self._ip, self._port,
             self.__class__.__name__.replace('Connection', '')
         )
+
+
+class ObfuscatedConnection(Connection):
+    """
+    Base class for "obfuscated" connections ("obfuscated2", "mtproto proxy")
+    """
+    """
+    This attribute should be redefined by subclasses
+    """
+    obfuscated_io = None
+
+    def _init_conn(self):
+        self._obfuscation = self.obfuscated_io(self)
+        self._writer.write(self._obfuscation.header)
+
+    def _send(self, data):
+        self._obfuscation.write(self._codec.encode_packet(data))
+
+    async def _recv(self):
+        return await self._codec.read_packet(self._obfuscation)
+
+
+class PacketCodec(abc.ABC):
+    """
+    Base class for packet codecs
+    """
+
+    """
+    This attribute should be re-defined by subclass to define if some
+    "magic bytes" should be sent to server right after conection is made to
+    signal which protocol will be used
+    """
+    tag = None
+
+    def __init__(self, connection):
+        """
+        Codec is created when connection is just made.
+        """
+        pass
+
+    @abc.abstractmethod
+    def encode_packet(self, data):
+        """
+        Encodes single packet and returns encoded bytes.
+        """
+        raise NotImplementedError
+
+    @abc.abstractmethod
+    async def read_packet(self, reader):
+        """
+        Reads single packet from `reader` object that should have
+        `readexactly(n)` method.
+        """
+        raise NotImplementedError

telethon/network/connection/http.py

@@ -1,34 +1,43 @@
 import asyncio
 
-from .connection import Connection
+from .connection import Connection, PacketCodec
 
 
 SSL_PORT = 443
 
 
-class ConnectionHttp(Connection):
+class HttpPacketCodec(PacketCodec):
-    async def connect(self, timeout=None, ssl=None):
+    tag = None
-        await super().connect(timeout=timeout, ssl=self._port == SSL_PORT)
+    obfuscate_tag = None
 
-    def _send(self, message):
+    def __init__(self, connection):
-        self._writer.write(
+        self._ip = connection._ip
-            'POST /api HTTP/1.1\r\n'
+        self._port = connection._port
-            'Host: {}:{}\r\n'
-            'Content-Type: application/x-www-form-urlencoded\r\n'
-            'Connection: keep-alive\r\n'
-            'Keep-Alive: timeout=100000, max=10000000\r\n'
-            'Content-Length: {}\r\n\r\n'
-            .format(self._ip, self._port, len(message))
-            .encode('ascii') + message
-        )
 
-    async def _recv(self):
+    def encode_packet(self, data):
+        return ('POST /api HTTP/1.1\r\n'
+                'Host: {}:{}\r\n'
+                'Content-Type: application/x-www-form-urlencoded\r\n'
+                'Connection: keep-alive\r\n'
+                'Keep-Alive: timeout=100000, max=10000000\r\n'
+                'Content-Length: {}\r\n\r\n'
+                .format(self._ip, self._port, len(data))
+                .encode('ascii') + data)
+
+    async def read_packet(self, reader):
         while True:
-            line = await self._reader.readline()
+            line = await reader.readline()
             if not line or line[-1] != b'\n':
                 raise asyncio.IncompleteReadError(line, None)
 
             if line.lower().startswith(b'content-length: '):
-                await self._reader.readexactly(2)
+                await reader.readexactly(2)
                 length = int(line[16:-2])
-                return await self._reader.readexactly(length)
+                return await reader.readexactly(length)
+
+
+class ConnectionHttp(Connection):
+    packet_codec = HttpPacketCodec
+
+    async def connect(self, timeout=None, ssl=None):
+        await super().connect(timeout=timeout, ssl=self._port == SSL_PORT)

telethon/network/connection/tcpabridged.py

@@ -1,6 +1,27 @@
 import struct
 
-from .connection import Connection
+from .connection import Connection, PacketCodec
+
+
+class AbridgedPacketCodec(PacketCodec):
+    tag = b'\xef'
+    obfuscate_tag = b'\xef\xef\xef\xef'
+
+    def encode_packet(self, data):
+        length = len(data) >> 2
+        if length < 127:
+            length = struct.pack('B', length)
+        else:
+            length = b'\x7f' + int.to_bytes(length, 3, 'little')
+        return length + data
+
+    async def read_packet(self, reader):
+        length = struct.unpack('<B', await reader.readexactly(1))[0]
+        if length >= 127:
+            length = struct.unpack(
+                '<i', await reader.readexactly(3) + b'\0')[0]
+
+        return await reader.readexactly(length << 2)
 
 
 class ConnectionTcpAbridged(Connection):
@@ -9,34 +30,4 @@ class ConnectionTcpAbridged(Connection):
     only require 1 byte if the packet length is less than
     508 bytes (127 << 2, which is very common).
     """
-    def _init_conn(self):
+    packet_codec = AbridgedPacketCodec
-        self._writer.write(b'\xef')
-
-    def _write(self, data):
-        """
-        Define wrapper write methods for `TcpObfuscated` to override.
-        """
-        self._writer.write(data)
-
-    async def _read(self, n):
-        """
-        Define wrapper read methods for `TcpObfuscated` to override.
-        """
-        return await self._reader.readexactly(n)
-
-    def _send(self, data):
-        length = len(data) >> 2
-        if length < 127:
-            length = struct.pack('B', length)
-        else:
-            length = b'\x7f' + int.to_bytes(length, 3, 'little')
-
-        self._write(length + data)
-
-    async def _recv(self):
-        length = struct.unpack('<B', await self._read(1))[0]
-        if length >= 127:
-            length = struct.unpack(
-                '<i', await self._read(3) + b'\0')[0]
-
-        return await self._read(length << 2)

telethon/network/connection/tcpfull.py

@@ -1,36 +1,29 @@
 import struct
 from zlib import crc32
 
-from .connection import Connection
+from .connection import Connection, PacketCodec
 from ...errors import InvalidChecksumError
 
 
-class ConnectionTcpFull(Connection):
+class FullPacketCodec(PacketCodec):
-    """
+    tag = None
-    Default Telegram mode. Sends 12 additional bytes and
-    needs to calculate the CRC value of the packet itself.
-    """
-    def __init__(self, ip, port, dc_id, *, loop, loggers, proxy=None):
-        super().__init__(
-            ip, port, dc_id, loop=loop, loggers=loggers, proxy=proxy)
-        self._send_counter = 0
 
-    def _init_conn(self):
+    def __init__(self, _conn):
         self._send_counter = 0  # Important or Telegram won't reply
 
-    def _send(self, data):
+    def encode_packet(self, data):
         # https://core.telegram.org/mtproto#tcp-transport
         # total length, sequence number, packet and checksum (CRC32)
         length = len(data) + 12
         data = struct.pack('<ii', length, self._send_counter) + data
         crc = struct.pack('<I', crc32(data))
         self._send_counter += 1
-        self._writer.write(data + crc)
+        return data + crc
 
-    async def _recv(self):
+    async def read_packet(self, reader):
-        packet_len_seq = await self._reader.readexactly(8)  # 4 and 4
+        packet_len_seq = await reader.readexactly(8)  # 4 and 4
         packet_len, seq = struct.unpack('<ii', packet_len_seq)
-        body = await self._reader.readexactly(packet_len - 8)
+        body = await reader.readexactly(packet_len - 8)
         checksum = struct.unpack('<I', body[-4:])[0]
         body = body[:-4]
 
@@ -39,3 +32,11 @@ class ConnectionTcpFull(Connection):
             raise InvalidChecksumError(checksum, valid_checksum)
 
         return body
+
+
+class ConnectionTcpFull(Connection):
+    """
+    Default Telegram mode. Sends 12 additional bytes and
+    needs to calculate the CRC value of the packet itself.
+    """
+    packet_codec = FullPacketCodec

telethon/network/connection/tcpintermediate.py

@@ -1,6 +1,41 @@
 import struct
+import random
+import os
 
-from .connection import Connection
+from .connection import Connection, PacketCodec
+
+
+class IntermediatePacketCodec(PacketCodec):
+    tag = b'\xee\xee\xee\xee'
+    obfuscate_tag = tag
+
+    def encode_packet(self, data):
+        return struct.pack('<i', len(data)) + data
+
+    async def read_packet(self, reader):
+        length = struct.unpack('<i', await reader.readexactly(4))[0]
+        return await reader.readexactly(length)
+
+
+class RandomizedIntermediatePacketCodec(IntermediatePacketCodec):
+    """
+    Data packets are aligned to 4bytes. This codec adds random bytes of size
+    from 0 to 3 bytes, which are ignored by decoder.
+    """
+    tag = None
+    obfuscate_tag = b'\xdd\xdd\xdd\xdd'
+
+    def encode_packet(self, data):
+        pad_size = random.randint(0, 3)
+        padding = os.urandom(pad_size)
+        return super().encode_packet(data + padding)
+
+    async def read_packet(self, reader):
+        packet_with_padding = await super().read_packet(reader)
+        pad_size = len(packet_with_padding) % 4
+        if pad_size > 0:
+            return packet_with_padding[:-pad_size]
+        return packet_with_padding
 
 
 class ConnectionTcpIntermediate(Connection):
@@ -8,12 +43,4 @@ class ConnectionTcpIntermediate(Connection):
     Intermediate mode between `ConnectionTcpFull` and `ConnectionTcpAbridged`.
     Always sends 4 extra bytes for the packet length.
     """
-    def _init_conn(self):
+    packet_codec = IntermediatePacketCodec
-        self._writer.write(b'\xee\xee\xee\xee')
-
-    def _send(self, data):
-        self._writer.write(struct.pack('<i', len(data)) + data)
-
-    async def _recv(self):
-        return await self._reader.readexactly(
-            struct.unpack('<i', await self._reader.readexactly(4))[0])

telethon/network/connection/tcpmtproxy.py

@@ -1,20 +1,97 @@
 import hashlib
+import os
 
-from .tcpobfuscated import ConnectionTcpObfuscated
+from .connection import ObfuscatedConnection
+from .tcpabridged import AbridgedPacketCodec
+from .tcpintermediate import (
+    IntermediatePacketCodec,
+    RandomizedIntermediatePacketCodec
+)
+
+from ...crypto import AESModeCTR
 
 
-class ConnectionTcpMTProxy(ConnectionTcpObfuscated):
+class MTProxyIO:
     """
-    Wrapper around the "obfuscated2" mode that modifies it a little and allows
+    It's very similar to tcpobfuscated.ObfuscatedIO, but the way
-    user to connect to the Telegram proxy servers commonly known as MTProxy.
+    encryption keys, protocol tag and dc_id are encoded is different.
+    """
+    header = None
+
+    def __init__(self, connection):
+        self._reader = connection._reader
+        self._writer = connection._writer
+
+        (self.header,
+         self._encrypt,
+         self._decrypt) = self.init_header(
+             connection._secret, connection._dc_id, connection.packet_codec)
+
+    def init_header(self, secret, dc_id, packet_codec):
+        # Validate
+        is_dd = (len(secret) == 17) and (secret[0] == 0xDD)
+        is_rand_codec = issubclass(
+            packet_codec, RandomizedIntermediatePacketCodec)
+        if is_dd and not is_rand_codec:
+            raise ValueError(
+                "Only RandomizedIntermediate can be used with dd-secrets")
+        secret = secret[:-1] if is_dd else secret
+        if len(secret) != 16:
+            raise ValueError(
+                "MTProxy secret must be a hex-string representing 16 bytes")
+
+        # Obfuscated messages secrets cannot start with any of these
+        keywords = (b'PVrG', b'GET ', b'POST', b'\xee\xee\xee\xee')
+        while True:
+            random = os.urandom(64)
+            if (random[0] != 0xef and
+                    random[:4] not in keywords and
+                    random[4:4] != b'\0\0\0\0'):
+                break
+
+        random = bytearray(random)
+        random_reversed = random[55:7:-1]  # Reversed (8, len=48)
+
+        # Encryption has "continuous buffer" enabled
+        encrypt_key = hashlib.sha256(
+            bytes(random[8:40]) + secret).digest()
+        encrypt_iv = bytes(random[40:56])
+        decrypt_key = hashlib.sha256(
+            bytes(random_reversed[:32]) + secret).digest()
+        decrypt_iv = bytes(random_reversed[32:48])
+
+        encryptor = AESModeCTR(encrypt_key, encrypt_iv)
+        decryptor = AESModeCTR(decrypt_key, decrypt_iv)
+
+        random[56:60] = packet_codec.obfuscate_tag
+
+        dc_id_bytes = dc_id.to_bytes(2, "little", signed=True)
+        random = random[:60] + dc_id_bytes + random[62:]
+        random[56:64] = encryptor.encrypt(bytes(random))[56:64]
+        return (random, encryptor, decryptor)
+
+    async def readexactly(self, n):
+        return self._decrypt.encrypt(await self._reader.readexactly(n))
+
+    def write(self, data):
+        self._writer.write(self._encrypt.encrypt(data))
+
+
+class TcpMTProxy(ObfuscatedConnection):
+    """
+    Connector which allows user to connect to the Telegram via proxy servers
+    commonly known as MTProxy.
     Implemented very ugly due to the leaky abstractions in Telethon networking
     classes that should be refactored later (TODO).
 
     .. warning::
 
-        The support for MTProtoProxies class is **EXPERIMENTAL** and prone to
+        The support for TcpMTProxy classes is **EXPERIMENTAL** and prone to
         be changed. You shouldn't be using this class yet.
     """
+    packet_codec = None
+    obfuscated_io = MTProxyIO
+
     @staticmethod
     def address_info(proxy_info):
         if proxy_info is None:
@@ -22,22 +99,29 @@ class ConnectionTcpMTProxy(ConnectionTcpObfuscated):
         return proxy_info[:2]
 
     def __init__(self, ip, port, dc_id, *, loop, loggers, proxy=None):
+        # connect to proxy's host and port instead of telegram's ones
         proxy_host, proxy_port = self.address_info(proxy)
+        self._secret = bytes.fromhex(proxy[2])
         super().__init__(
             proxy_host, proxy_port, dc_id, loop=loop, loggers=loggers)
 
-        # TODO: Implement the dd-secret secure mode (adds noise to fool DPI)
-        self._secret = bytes.fromhex(proxy[2])
-        if len(self._secret) != 16:
-            raise ValueError(
-                "MTProxy secure mode is not implemented for now"
-                if len(self._secret) == 17 and self._secret[0] == 0xDD else
-                "MTProxy secret must be a hex-string representing 16 bytes"
-            )
 
-    def _compose_key(self, data):
+class ConnectionTcpMTProxyAbridged(TcpMTProxy):
-        return hashlib.sha256(data + self._secret).digest()
+    """
+    Connect to proxy using abridged protocol
+    """
+    packet_codec = AbridgedPacketCodec
 
-    def _compose_tail(self, data):
+
-        dc_id_bytes = self._dc_id.to_bytes(2, "little", signed=True)
+class ConnectionTcpMTProxyIntermediate(TcpMTProxy):
-        return super()._compose_tail(data[:60] + dc_id_bytes + data[62:])
+    """
+    Connect to proxy using intermediate protocol
+    """
+    packet_codec = IntermediatePacketCodec
+
+
+class ConnectionTcpMTProxyRandomizedIntermediate(TcpMTProxy):
+    """
+    Connect to proxy using randomized intermediate protocol (dd-secrets)
+    """
+    packet_codec = RandomizedIntermediatePacketCodec

telethon/network/connection/tcpobfuscated.py

@@ -1,30 +1,23 @@
 import os
 
-from .connection import Connection
+from .tcpabridged import AbridgedPacketCodec
-from .tcpabridged import ConnectionTcpAbridged
+from .connection import ObfuscatedConnection
+
 from ...crypto import AESModeCTR
 
 
-class ConnectionTcpObfuscated(ConnectionTcpAbridged):
+class ObfuscatedIO:
-    """
+    header = None
-    Mode that Telegram defines as "obfuscated2". Encodes the packet
-    just like `ConnectionTcpAbridged`, but encrypts every message with
-    a randomly generated key using the AES-CTR mode so the packets are
-    harder to discern.
-    """
-    def __init__(self, ip, port, dc_id, *, loop, loggers, proxy=None):
-        super().__init__(
-            ip, port, dc_id, loop=loop, loggers=loggers, proxy=proxy)
-        self._aes_encrypt = None
-        self._aes_decrypt = None
 
-    def _write(self, data):
+    def __init__(self, connection):
-        self._writer.write(self._aes_encrypt.encrypt(data))
+        self._reader = connection._reader
+        self._writer = connection._writer
 
-    async def _read(self, n):
+        (self.header,
-        return self._aes_decrypt.encrypt(await self._reader.readexactly(n))
+         self._encrypt,
+         self._decrypt) = self.init_header(connection.packet_codec)
 
-    def _init_conn(self):
+    def init_header(self, packet_codec):
         # Obfuscated messages secrets cannot start with any of these
         keywords = (b'PVrG', b'GET ', b'POST', b'\xee\xee\xee\xee')
         while True:
@@ -38,23 +31,31 @@ class ConnectionTcpObfuscated(ConnectionTcpAbridged):
         random_reversed = random[55:7:-1]  # Reversed (8, len=48)
 
         # Encryption has "continuous buffer" enabled
-        encrypt_key = self._compose_key(bytes(random[8:40]))
+        encrypt_key = bytes(random[8:40])
         encrypt_iv = bytes(random[40:56])
-        decrypt_key = self._compose_key(bytes(random_reversed[:32]))
+        decrypt_key = bytes(random_reversed[:32])
         decrypt_iv = bytes(random_reversed[32:48])
 
-        self._aes_encrypt = AESModeCTR(encrypt_key, encrypt_iv)
+        encryptor = AESModeCTR(encrypt_key, encrypt_iv)
-        self._aes_decrypt = AESModeCTR(decrypt_key, decrypt_iv)
+        decryptor = AESModeCTR(decrypt_key, decrypt_iv)
 
-        random[56:60] = b'\xef\xef\xef\xef'
+        random[56:60] = packet_codec.obfuscate_tag
-        random[56:64] = self._compose_tail(bytes(random))
+        random[56:64] = encryptor.encrypt(bytes(random))[56:64]
+        return (random, encryptor, decryptor)
 
-        self._writer.write(random)
+    async def readexactly(self, n):
+        return self._decrypt.encrypt(await self._reader.readexactly(n))
 
-    # Next functions provide the variable parts of the connection handshake.
+    def write(self, data):
-    # This is necessary to modify obfuscated2 the way that MTProxy requires.
+        self._writer.write(self._encrypt.encrypt(data))
-    def _compose_key(self, data):
-        return data
 
-    def _compose_tail(self, data):
+
-        return self._aes_encrypt.encrypt(data)[56:64]
+class ConnectionTcpObfuscated(ObfuscatedConnection):
+    """
+    Mode that Telegram defines as "obfuscated2". Encodes the packet
+    just like `ConnectionTcpAbridged`, but encrypts every message with
+    a randomly generated key using the AES-CTR mode so the packets are
+    harder to discern.
+    """
+    obfuscated_io = ObfuscatedIO
+    packet_codec = AbridgedPacketCodec