Let all connection modes implement the new Connection

telethon/network/connection/tcpobfuscated.py

@@ -1,7 +1,7 @@
 import os
 
+from .connection import Connection
 from .tcpabridged import ConnectionTcpAbridged
-from .tcpfull import ConnectionTcpFull
 from ...crypto import AESModeCTR
 
 
@@ -11,16 +11,22 @@ class ConnectionTcpObfuscated(ConnectionTcpAbridged):
     every message with a randomly generated key using the
     AES-CTR mode so the packets are harder to discern.
     """
-    def __init__(self, *, loop, timeout, proxy=None):
-        super().__init__(loop=loop, timeout=timeout, proxy=proxy)
-        self._aes_encrypt, self._aes_decrypt = None, None
-        self.read = lambda s: self._aes_decrypt.encrypt(self.conn.read(s))
-        self.write = lambda d: self.conn.write(self._aes_encrypt.encrypt(d))
+    def __init__(self, ip, port, *, loop):
+        super().__init__(ip, port, loop=loop)
+        self._aes_encrypt = None
+        self._aes_decrypt = None
+
+    def _write(self, data):
+        self._writer.write(self._aes_encrypt.encrypt(data))
+
+    async def _read(self, n):
+        return self._aes_decrypt.encrypt(await self._reader.readexactly(n))
+
+    async def connect(self):
+        await Connection.connect(self)
 
-    async def connect(self, ip, port):
-        result = await ConnectionTcpFull.connect(self, ip, port)
         # Obfuscated messages secrets cannot start with any of these
-        keywords = (b'PVrG', b'GET ', b'POST', b'\xee' * 4)
+        keywords = (b'PVrG', b'GET ', b'POST', b'\xee\xee\xee\xee')
         while True:
             random = os.urandom(64)
             if (random[0] != b'\xef' and
@@ -28,11 +34,11 @@ class ConnectionTcpObfuscated(ConnectionTcpAbridged):
                     random[4:4] != b'\0\0\0\0'):
                 break
 
-        random = list(random)
+        random = bytearray(random)
         random[56] = random[57] = random[58] = random[59] = 0xef
         random_reversed = random[55:7:-1]  # Reversed (8, len=48)
 
-        # encryption has "continuous buffer" enabled
+        # Encryption has "continuous buffer" enabled
         encrypt_key = bytes(random[8:40])
         encrypt_iv = bytes(random[40:56])
         decrypt_key = bytes(random_reversed[:32])
@@ -42,5 +48,4 @@ class ConnectionTcpObfuscated(ConnectionTcpAbridged):
         self._aes_decrypt = AESModeCTR(decrypt_key, decrypt_iv)
 
         random[56:64] = self._aes_encrypt.encrypt(bytes(random))[56:64]
-        await self.conn.write(bytes(random))
-        return result
+        await self._send(random)