Make TCPAbridgedO async

delivrance · delivrance · commit 78a768f52ccc · 2018-10-08T20:17:04.000+02:00
diff --git a/pyrogram/connection/transport/tcp/tcp_abridged_o.py b/pyrogram/connection/transport/tcp/tcp_abridged_o.py
@@ -34,8 +34,8 @@ def __init__(self, ipv6: bool, proxy: dict):
         self.encrypt = None
         self.decrypt = None
 
-    def connect(self, address: tuple):
-        super().connect(address)
+    async def connect(self, address: tuple):
+        await super().connect(address)
 
         while True:
             nonce = bytearray(os.urandom(64))
@@ -53,12 +53,12 @@ def connect(self, address: tuple):
 
         nonce[56:64] = AES.ctr256_encrypt(nonce, *self.encrypt)[56:64]
 
-        super().sendall(nonce)
+        await super().send(nonce)
 
-    def sendall(self, data: bytes, *args):
+    async def send(self, data: bytes, *args):
         length = len(data) // 4
 
-        super().sendall(
+        await super().send(
             AES.ctr256_encrypt(
                 (bytes([length])
                  if length <= 126
@@ -68,23 +68,23 @@ def sendall(self, data: bytes, *args):
             )
         )
 
-    def recvall(self, length: int = 0) -> bytes or None:
-        length = super().recvall(1)
+    async def recv(self, length: int = 0) -> bytes or None:
+        length = await super().recv(1)
 
         if length is None:
             return None
 
         length = AES.ctr256_decrypt(length, *self.decrypt)
 
         if length == b"\x7f":
-            length = super().recvall(3)
+            length = await super().recv(3)
 
             if length is None:
                 return None
 
             length = AES.ctr256_decrypt(length, *self.decrypt)
 
-        data = super().recvall(int.from_bytes(length, "little") * 4)
+        data = await super().recv(int.from_bytes(length, "little") * 4)
 
         if data is None:
             return None
