diff --git a/telethon/__init__.py b/telethon/__init__.py index b3971245..c8593168 100644 --- a/telethon/__init__.py +++ b/telethon/__init__.py @@ -1,3 +1,4 @@ from .telegram_bare_client import TelegramBareClient from .telegram_client import TelegramClient +from .network import ConnectionMode from . import tl diff --git a/telethon/network/__init__.py b/telethon/network/__init__.py index 8ef2643b..77bd4406 100644 --- a/telethon/network/__init__.py +++ b/telethon/network/__init__.py @@ -1,4 +1,4 @@ from .mtproto_plain_sender import MtProtoPlainSender from .authenticator import do_authentication from .mtproto_sender import MtProtoSender -from .connection import Connection +from .connection import Connection, ConnectionMode diff --git a/telethon/network/connection.py b/telethon/network/connection.py index 2116464a..77ef349d 100644 --- a/telethon/network/connection.py +++ b/telethon/network/connection.py @@ -1,23 +1,46 @@ import os from datetime import timedelta from zlib import crc32 +from enum import Enum from ..crypto import AESModeCTR from ..extensions import BinaryWriter, TcpClient from ..errors import InvalidChecksumError +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' may be any of: - 'tcp_full', 'tcp_intermediate', 'tcp_abridged', 'tcp_obfuscated' + '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, ip, port, mode='tcp_intermediate', + def __init__(self, ip, port, mode=ConnectionMode.TCP_FULL, proxy=None, timeout=timedelta(seconds=5)): self.ip = ip self.port = port @@ -30,20 +53,20 @@ class Connection: self.conn = TcpClient(proxy=proxy, timeout=timeout) # Sending messages - if mode == 'tcp_full': + if mode == ConnectionMode.TCP_FULL: setattr(self, 'send', self._send_tcp_full) setattr(self, 'recv', self._recv_tcp_full) - elif mode == 'tcp_intermediate': + elif mode == ConnectionMode.TCP_INTERMEDIATE: setattr(self, 'send', self._send_intermediate) setattr(self, 'recv', self._recv_intermediate) - elif mode in ('tcp_abridged', 'tcp_obfuscated'): + 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 == 'tcp_obfuscated': + if mode == ConnectionMode.TCP_OBFUSCATED: setattr(self, 'write', self._write_obfuscated) setattr(self, 'read', self._read_obfuscated) else: @@ -54,11 +77,11 @@ class Connection: self._send_counter = 0 self.conn.connect(self.ip, self.port) - if self._mode == 'tcp_abridged': + if self._mode == ConnectionMode.TCP_ABRIDGED: self.conn.write(b'\xef') - elif self._mode == 'tcp_intermediate': + elif self._mode == ConnectionMode.TCP_INTERMEDIATE: self.conn.write(b'\xee\xee\xee\xee') - elif self._mode == 'tcp_obfuscated': + elif self._mode == ConnectionMode.TCP_OBFUSCATED: self._setup_obfuscation() def _setup_obfuscation(self): @@ -103,7 +126,7 @@ class Connection: def recv(self): """Receives and unpacks a message""" # Default implementation is just an error - raise ValueError('Invalid connection mode specified: ' + self._mode) + raise ValueError('Invalid connection mode specified: ' + str(self._mode)) def _recv_tcp_full(self): packet_length_bytes = self.read(4) @@ -138,7 +161,7 @@ class Connection: def send(self, message): """Encapsulates and sends the given message""" # Default implementation is just an error - raise ValueError('Invalid connection mode specified: ' + self._mode) + raise ValueError('Invalid connection mode specified: ' + str(self._mode)) def _send_tcp_full(self, message): # https://core.telegram.org/mtproto#tcp-transport @@ -174,7 +197,7 @@ class Connection: # region Read implementations def read(self, length): - raise ValueError('Invalid connection mode specified: ' + self._mode) + raise ValueError('Invalid connection mode specified: ' + str(self._mode)) def _read_plain(self, length): return self.conn.read(length) @@ -189,7 +212,7 @@ class Connection: # region Write implementations def write(self, data): - raise ValueError('Invalid connection mode specified: ' + self._mode) + raise ValueError('Invalid connection mode specified: ' + str(self._mode)) def _write_plain(self, data): self.conn.write(data) diff --git a/telethon/telegram_bare_client.py b/telethon/telegram_bare_client.py index 46c878ca..d7a031c6 100644 --- a/telethon/telegram_bare_client.py +++ b/telethon/telegram_bare_client.py @@ -10,7 +10,7 @@ from . import helpers as utils from .errors import ( RPCError, FloodWaitError, FileMigrateError, TypeNotFoundError ) -from .network import authenticator, MtProtoSender, Connection +from .network import authenticator, MtProtoSender, Connection, ConnectionMode from .utils import get_appropriated_part_size from .crypto import rsa, CdnDecrypter @@ -66,6 +66,7 @@ class TelegramBareClient: # region Initialization def __init__(self, session, api_id, api_hash, + connection_mode=ConnectionMode.TCP_FULL, proxy=None, timeout=timedelta(seconds=5)): """Initializes the Telegram client with the specified API ID and Hash. Session must always be a Session instance, and an optional proxy @@ -74,6 +75,7 @@ class TelegramBareClient: self.session = session self.api_id = int(api_id) self.api_hash = api_hash + self._connection_mode = connection_mode self.proxy = proxy self._timeout = timeout self._logger = logging.getLogger(__name__) @@ -125,7 +127,7 @@ class TelegramBareClient: connection = Connection( self.session.server_address, self.session.port, - proxy=self.proxy, timeout=self._timeout + mode=self._connection_mode, proxy=self.proxy, timeout=self._timeout ) try: diff --git a/telethon/telegram_client.py b/telethon/telegram_client.py index 28e86dda..1ca46500 100644 --- a/telethon/telegram_client.py +++ b/telethon/telegram_client.py @@ -5,6 +5,7 @@ from threading import Event, RLock, Thread from time import sleep, time from . import TelegramBareClient +from .network import ConnectionMode # Import some externalized utilities to work with the Telegram types and more from . import helpers as utils @@ -62,7 +63,9 @@ class TelegramClient(TelegramBareClient): # region Initialization - def __init__(self, session, api_id, api_hash, proxy=None, + def __init__(self, session, api_id, api_hash, + connection_mode=ConnectionMode.TCP_FULL, + proxy=None, timeout=timedelta(seconds=5), **kwargs): """Initializes the Telegram client with the specified API ID and Hash. @@ -72,6 +75,10 @@ class TelegramClient(TelegramBareClient): would probably not work). Pass 'None' for it to be a temporary session - remember to '.log_out()'! + The 'connection_mode' should be any value under ConnectionMode. + This will only affect how messages are sent over the network + and how much processing is required before sending them. + If more named arguments are provided as **kwargs, they will be used to update the Session instance. Most common settings are: device_model = platform.node() @@ -93,7 +100,10 @@ class TelegramClient(TelegramBareClient): raise ValueError( 'The given session must be a str or a Session instance.') - super().__init__(session, api_id, api_hash, proxy, timeout=timeout) + super().__init__( + session, api_id, api_hash, + connection_mode=connection_mode, proxy=proxy, timeout=timeout + ) # Safety across multiple threads (for the updates thread) self._lock = RLock() diff --git a/telethon_examples/interactive_telegram_client.py b/telethon_examples/interactive_telegram_client.py index 37aa46a2..ae858cb7 100644 --- a/telethon_examples/interactive_telegram_client.py +++ b/telethon_examples/interactive_telegram_client.py @@ -1,7 +1,7 @@ import os from getpass import getpass -from telethon import TelegramClient +from telethon import TelegramClient, ConnectionMode from telethon.errors import SessionPasswordNeededError from telethon.tl.types import UpdateShortChatMessage, UpdateShortMessage from telethon.utils import get_display_name @@ -49,7 +49,10 @@ class InteractiveTelegramClient(TelegramClient): print_title('Initialization') print('Initializing interactive example...') - super().__init__(session_user_id, api_id, api_hash, proxy) + super().__init__( + session_user_id, api_id, api_hash, + connection_mode=ConnectionMode.TCP_ABRIDGED, proxy=proxy + ) # Store all the found media in memory here, # so it can be downloaded if the user wants