Use an Enum for the ConnectionMode and support specifying it

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)