Use an Enum for the ConnectionMode and support specifying it

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

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

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:
+       'mode' must be any of the ConnectionMode enumeration.
-         'tcp_full', 'tcp_intermediate', 'tcp_abridged', 'tcp_obfuscated'
 
        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)

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:

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

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