From ba4b7ce881eed3d2ce2a327e4b4eb15190f066ff Mon Sep 17 00:00:00 2001 From: Lonami Exo Date: Thu, 10 May 2018 14:22:19 +0200 Subject: [PATCH] Make the Connection a proper ABC (#509) --- telethon/__init__.py | 2 +- telethon/network/__init__.py | 5 +- telethon/network/connection.py | 316 ------------------ telethon/network/connection/__init__.py | 4 + telethon/network/connection/common.py | 57 ++++ telethon/network/connection/tcpabridged.py | 34 ++ telethon/network/connection/tcpfull.py | 65 ++++ .../network/connection/tcpintermediate.py | 23 ++ telethon/network/connection/tcpobfuscated.py | 50 +++ telethon/telegram_bare_client.py | 12 +- telethon/telegram_client.py | 17 +- 11 files changed, 254 insertions(+), 331 deletions(-) delete mode 100644 telethon/network/connection.py create mode 100644 telethon/network/connection/__init__.py create mode 100644 telethon/network/connection/common.py create mode 100644 telethon/network/connection/tcpabridged.py create mode 100644 telethon/network/connection/tcpfull.py create mode 100644 telethon/network/connection/tcpintermediate.py create mode 100644 telethon/network/connection/tcpobfuscated.py diff --git a/telethon/__init__.py b/telethon/__init__.py index 2f984bf1..cfeba49e 100644 --- a/telethon/__init__.py +++ b/telethon/__init__.py @@ -1,7 +1,7 @@ import logging from .telegram_bare_client import TelegramBareClient from .telegram_client import TelegramClient -from .network import ConnectionMode +from .network import connection from . import tl, version diff --git a/telethon/network/__init__.py b/telethon/network/__init__.py index d2538924..756df771 100644 --- a/telethon/network/__init__.py +++ b/telethon/network/__init__.py @@ -5,4 +5,7 @@ with Telegram's servers and the protocol used (TCP full, abridged, etc.). from .mtproto_plain_sender import MtProtoPlainSender from .authenticator import do_authentication from .mtproto_sender import MtProtoSender -from .connection import Connection, ConnectionMode +from .connection import ( + ConnectionTcpFull, ConnectionTcpAbridged, ConnectionTcpObfuscated, + ConnectionTcpIntermediate +) diff --git a/telethon/network/connection.py b/telethon/network/connection.py deleted file mode 100644 index 45afaefc..00000000 --- a/telethon/network/connection.py +++ /dev/null @@ -1,316 +0,0 @@ -""" -This module holds both the Connection class and the ConnectionMode enum, -which specifies the protocol to be used by the Connection. -""" -import logging -import os -import struct -from datetime import timedelta -from zlib import crc32 -from enum import Enum - -import errno - -from ..crypto import AESModeCTR -from ..extensions import TcpClient -from ..errors import InvalidChecksumError - -__log__ = logging.getLogger(__name__) - - -class ConnectionMode(Enum): - """Represents which mode should be used to stabilise a connection. - - TCP_FULL: Default Telegram mode. Sends 12 additional bytes and - needs to calculate the CRC value of the packet itself. - - TCP_INTERMEDIATE: Intermediate mode between TCP_FULL and TCP_ABRIDGED. - Always sends 4 extra bytes for the packet length. - - TCP_ABRIDGED: 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). - - TCP_OBFUSCATED: Encodes the packet just like TCP_ABRIDGED, but encrypts - every message with a randomly generated key using the - AES-CTR mode so the packets are harder to discern. - """ - TCP_FULL = 1 - TCP_INTERMEDIATE = 2 - TCP_ABRIDGED = 3 - TCP_OBFUSCATED = 4 - - -class Connection: - """ - Represents an abstract connection (TCP, TCP abridged...). - 'mode' must be any of the ConnectionMode enumeration. - - 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, mode=ConnectionMode.TCP_FULL, - proxy=None, timeout=timedelta(seconds=5)): - """ - Initializes a new connection. - - :param mode: the ConnectionMode to be used. - :param proxy: whether to use a proxy or not. - :param timeout: timeout to be used for all operations. - """ - self._mode = mode - 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, timeout=timeout) - - # Sending messages - if mode == ConnectionMode.TCP_FULL: - setattr(self, 'send', self._send_tcp_full) - setattr(self, 'recv', self._recv_tcp_full) - - elif mode == ConnectionMode.TCP_INTERMEDIATE: - setattr(self, 'send', self._send_intermediate) - setattr(self, 'recv', self._recv_intermediate) - - elif mode in (ConnectionMode.TCP_ABRIDGED, - ConnectionMode.TCP_OBFUSCATED): - setattr(self, 'send', self._send_abridged) - setattr(self, 'recv', self._recv_abridged) - - # Writing and reading from the socket - if mode == ConnectionMode.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, ip, port): - """ - Estabilishes a connection to IP:port. - - :param ip: the IP to connect to. - :param port: the port to connect to. - """ - 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 - if self._mode == ConnectionMode.TCP_ABRIDGED: - self.conn.write(b'\xef') - elif self._mode == ConnectionMode.TCP_INTERMEDIATE: - self.conn.write(b'\xee\xee\xee\xee') - elif self._mode == ConnectionMode.TCP_OBFUSCATED: - self._setup_obfuscation() - - def get_timeout(self): - """Returns the timeout used by the connection.""" - return self.conn.timeout - - def _setup_obfuscation(self): - """ - Sets up the obfuscated protocol. - """ - # 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): - """ - Determines whether the connection is alive or not. - - :return: true if it's connected. - """ - return self.conn.connected - - def close(self): - """Closes the connection.""" - self.conn.close() - - def clone(self): - """Creates a copy of this Connection.""" - return Connection( - mode=self._mode, proxy=self.conn.proxy, timeout=self.conn.timeout - ) - - # region Receive message implementations - - def recv(self): - """Receives and unpacks a message""" - # Default implementation is just an error - raise ValueError('Invalid connection mode specified: ' + str(self._mode)) - - def _recv_tcp_full(self): - """ - Receives a message from the network, - internally encoded using the TCP full protocol. - - May raise InvalidChecksumError if the received data doesn't - match its valid checksum. - - :return: the read message payload. - """ - packet_len_seq = self.read(8) # 4 and 4 - packet_len, seq = struct.unpack('= 127: - length = struct.unpack('> 2 - if length < 127: - length = struct.pack('B', length) - else: - length = b'\x7f' + int.to_bytes(length, 3, 'little') - - self.write(length + message) - - # endregion - - # region Read implementations - - def read(self, length): - raise ValueError('Invalid connection mode specified: ' + str(self._mode)) - - def _read_plain(self, length): - """ - Reads data from the socket connection. - - :param length: how many bytes should be read. - :return: a byte sequence with len(data) == length - """ - return self.conn.read(length) - - def _read_obfuscated(self, length): - """ - Reads data and decrypts from the socket connection. - - :param length: how many bytes should be read. - :return: the decrypted byte sequence with len(data) == length - """ - return self._aes_decrypt.encrypt( - self.conn.read(length) - ) - - # endregion - - # region Write implementations - - def write(self, data): - raise ValueError('Invalid connection mode specified: ' + str(self._mode)) - - def _write_plain(self, data): - """ - Writes the given data through the socket connection. - - :param data: the data in bytes to be written. - """ - self.conn.write(data) - - def _write_obfuscated(self, data): - """ - Writes the given data through the socket connection, - using the obfuscated mode (AES encryption is applied on top). - - :param data: the data in bytes to be written. - """ - self.conn.write(self._aes_encrypt.encrypt(data)) - - # endregion diff --git a/telethon/network/connection/__init__.py b/telethon/network/connection/__init__.py new file mode 100644 index 00000000..0c7a07d0 --- /dev/null +++ b/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 diff --git a/telethon/network/connection/common.py b/telethon/network/connection/common.py new file mode 100644 index 00000000..504a0c12 --- /dev/null +++ b/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 diff --git a/telethon/network/connection/tcpabridged.py b/telethon/network/connection/tcpabridged.py new file mode 100644 index 00000000..7b5a7848 --- /dev/null +++ b/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('= 127: + length = struct.unpack('> 2 + if length < 127: + length = struct.pack('B', length) + else: + length = b'\x7f' + int.to_bytes(length, 3, 'little') + + self.write(length + message) diff --git a/telethon/network/connection/tcpfull.py b/telethon/network/connection/tcpfull.py new file mode 100644 index 00000000..37d130ec --- /dev/null +++ b/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('