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:
-         '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)

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