Make the Connection a proper ABC (#509)

telethon/network/connection/__init__.py

@@ -0,0 +1,4 @@
+from .tcpfull import ConnectionTcpFull
+from .tcpabridged import ConnectionTcpAbridged
+from .tcpobfuscated import ConnectionTcpObfuscated
+from .tcpintermediate import ConnectionTcpIntermediate

telethon/network/connection/common.py

@@ -0,0 +1,57 @@
+"""
+This module holds the abstract `Connection` class.
+"""
+import abc
+from datetime import timedelta
+
+
+class Connection(abc.ABC):
+    """
+    Represents an abstract connection for Telegram.
+
+    Subclasses should implement the actual protocol
+    being used when encoding/decoding messages.
+    """
+    def __init__(self, proxy=None, timeout=timedelta(seconds=5)):
+        """
+        Initializes a new connection.
+
+        :param proxy: whether to use a proxy or not.
+        :param timeout: timeout to be used for all operations.
+        """
+        self._proxy = proxy
+        self._timeout = timeout
+
+    @abc.abstractmethod
+    def get_timeout(self):
+        """Returns the timeout used by the connection."""
+        raise NotImplementedError
+
+    @abc.abstractmethod
+    def is_connected(self):
+        """
+        Determines whether the connection is alive or not.
+
+        :return: true if it's connected.
+        """
+        raise NotImplementedError
+
+    @abc.abstractmethod
+    def close(self):
+        """Closes the connection."""
+        raise NotImplementedError
+
+    @abc.abstractmethod
+    def clone(self):
+        """Creates a copy of this Connection."""
+        raise NotImplementedError
+
+    @abc.abstractmethod
+    def recv(self):
+        """Receives and unpacks a message"""
+        raise NotImplementedError
+
+    @abc.abstractmethod
+    def send(self, message):
+        """Encapsulates and sends the given message"""
+        raise NotImplementedError

telethon/network/connection/tcpabridged.py

@@ -0,0 +1,34 @@
+import struct
+
+from .tcpfull import ConnectionTcpFull
+
+
+class ConnectionTcpAbridged(ConnectionTcpFull):
+    """
+    This is the mode with the lowest overhead, as it will
+    only require 1 byte if the packet length is less than
+    508 bytes (127 << 2, which is very common).
+    """
+    def connect(self, ip, port):
+        result = super().connect(ip, port)
+        self.conn.write(b'\xef')
+        return result
+
+    def clone(self):
+        return ConnectionTcpAbridged(self._proxy, self._timeout)
+
+    def recv(self):
+        length = struct.unpack('<B', self.read(1))[0]
+        if length >= 127:
+            length = struct.unpack('<i', self.read(3) + b'\0')[0]
+
+        return self.read(length << 2)
+
+    def send(self, message):
+        length = len(message) >> 2
+        if length < 127:
+            length = struct.pack('B', length)
+        else:
+            length = b'\x7f' + int.to_bytes(length, 3, 'little')
+
+        self.write(length + message)

telethon/network/connection/tcpfull.py

@@ -0,0 +1,65 @@
+import errno
+import struct
+from datetime import timedelta
+from zlib import crc32
+
+from .common import Connection
+from ...errors import InvalidChecksumError
+from ...extensions import TcpClient
+
+
+class ConnectionTcpFull(Connection):
+    """
+    Default Telegram mode. Sends 12 additional bytes and
+    needs to calculate the CRC value of the packet itself.
+    """
+    def __init__(self, proxy=None, timeout=timedelta(seconds=5)):
+        super().__init__(proxy, timeout)
+        self._send_counter = 0
+        self.conn = TcpClient(proxy=self._proxy, timeout=self._timeout)
+        self.read = self.conn.read
+        self.write = self.conn.write
+
+    def connect(self, ip, port):
+        try:
+            self.conn.connect(ip, port)
+        except OSError as e:
+            if e.errno == errno.EISCONN:
+                return  # Already connected, no need to re-set everything up
+            else:
+                raise
+
+        self._send_counter = 0
+
+    def get_timeout(self):
+        return self.conn.timeout
+
+    def is_connected(self):
+        return self.conn.connected
+
+    def close(self):
+        self.conn.close()
+
+    def clone(self):
+        return ConnectionTcpFull(self._proxy, self._timeout)
+
+    def recv(self):
+        packet_len_seq = self.read(8)  # 4 and 4
+        packet_len, seq = struct.unpack('<ii', packet_len_seq)
+        body = self.read(packet_len - 12)
+        checksum = struct.unpack('<I', self.read(4))[0]
+
+        valid_checksum = crc32(packet_len_seq + body)
+        if checksum != valid_checksum:
+            raise InvalidChecksumError(checksum, valid_checksum)
+
+        return body
+
+    def send(self, message):
+        # https://core.telegram.org/mtproto#tcp-transport
+        # total length, sequence number, packet and checksum (CRC32)
+        length = len(message) + 12
+        data = struct.pack('<ii', length, self._send_counter) + message
+        crc = struct.pack('<I', crc32(data))
+        self._send_counter += 1
+        self.write(data + crc)

telethon/network/connection/tcpintermediate.py

@@ -0,0 +1,23 @@
+import struct
+
+from .tcpfull import ConnectionTcpFull
+
+
+class ConnectionTcpIntermediate(ConnectionTcpFull):
+    """
+    Intermediate mode between `ConnectionTcpFull` and `ConnectionTcpAbridged`.
+    Always sends 4 extra bytes for the packet length.
+    """
+    def connect(self, ip, port):
+        result = super().connect(ip, port)
+        self.conn.write(b'\xee\xee\xee\xee')
+        return result
+
+    def clone(self):
+        return ConnectionTcpIntermediate(self._proxy, self._timeout)
+
+    def recv(self):
+        return self.read(struct.unpack('<i', self.read(4))[0])
+
+    def send(self, message):
+        self.write(struct.pack('<i', len(message)) + message)

telethon/network/connection/tcpobfuscated.py

@@ -0,0 +1,50 @@
+import os
+from datetime import timedelta
+
+from .tcpfull import ConnectionTcpFull
+from .tcpabridged import ConnectionTcpAbridged
+from ...crypto import AESModeCTR
+
+
+class ConnectionTcpObfuscated(ConnectionTcpAbridged):
+    """
+    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, proxy=None, timeout=timedelta(seconds=5)):
+        super().__init__(proxy, timeout)
+        self._aes_encrypt, self._aes_decrypt = None, None
+        self.read = lambda s: self._aes_decrypt.encrypt(self.conn.read(s))
+        self.write = lambda d: self.conn.write(self._aes_encrypt.encrypt(d))
+
+    def connect(self, ip, port):
+        result = ConnectionTcpFull.connect(self, ip, port)
+        # 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'):
+                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))
+        return result
+
+    def clone(self):
+        return ConnectionTcpObfuscated(self._proxy, self._timeout)