Rework class hierarchy, try to DRY more

telethon/network/connection/tcpintermediate.py

@@ -2,31 +2,12 @@ import struct
 import random
 import os
 
-from .connection import Connection
+from .connection import Connection, PacketCodec
 
 
-class ConnectionTcpIntermediate(Connection):
-    """
-    Intermediate mode between `ConnectionTcpFull` and `ConnectionTcpAbridged`.
-    Always sends 4 extra bytes for the packet length.
-    """
-    def __init__(self, *args, **kwargs):
-        super().__init__(*args, **kwargs)
-        self._codec = IntermediatePacket()
-
-    def _init_conn(self):
-        self._writer.write(self._codec.tag)
-
-    def _send(self, data):
-        self._writer.write(self._codec.encode_packet(data))
-
-    async def _recv(self):
-        return await self.codec.read_packet(self._reader)
-
-
-class IntermediatePacket:
+class IntermediatePacketCodec(PacketCodec):
     tag = b'\xee\xee\xee\xee'
-    mtproto_proxy_tag = tag
+    obfuscate_tag = tag
 
     def encode_packet(self, data):
         return struct.pack('<i', len(data)) + data
@@ -36,12 +17,13 @@ class IntermediatePacket:
         return await reader.readexactly(length)
 
 
-class RandomizedIntermediatePacket(IntermediatePacket):
+class RandomizedIntermediatePacketCodec(IntermediatePacketCodec):
     """
     Data packets are aligned to 4bytes. This codec adds random bytes of size
     from 0 to 3 bytes, which are ignored by decoder.
     """
-    mtproto_proxy_tag = b'\xdd\xdd\xdd\xdd'
+    tag = None
+    obfuscate_tag = b'\xdd\xdd\xdd\xdd'
 
     def encode_packet(self, data):
         pad_size = random.randint(0, 3)
@@ -54,3 +36,11 @@ class RandomizedIntermediatePacket(IntermediatePacket):
         if pad_size > 0:
             return packet_with_padding[:-pad_size]
         return packet_with_padding
+
+
+class ConnectionTcpIntermediate(Connection):
+    """
+    Intermediate mode between `ConnectionTcpFull` and `ConnectionTcpAbridged`.
+    Always sends 4 extra bytes for the packet length.
+    """
+    packet_codec = IntermediatePacketCodec