diff --git a/pyrogram/connection/transport/tcp/tcp.py b/pyrogram/connection/transport/tcp/tcp.py index 5df8aacb..f006cadd 100644 --- a/pyrogram/connection/transport/tcp/tcp.py +++ b/pyrogram/connection/transport/tcp/tcp.py @@ -16,6 +16,7 @@ # You should have received a copy of the GNU Lesser General Public License # along with Pyrogram. If not, see . +import asyncio import logging import socket @@ -32,14 +33,18 @@ except ImportError as e: log = logging.getLogger(__name__) -class TCP(socks.socksocket): +class TCP: def __init__(self, proxy: dict): - super().__init__() - self.settimeout(10) + self.proxy = proxy + + self.socket = socks.socksocket() + self.reader = None # type: asyncio.StreamReader + self.writer = None # type: asyncio.StreamWriter + self.proxy_enabled = proxy.get("enabled", False) if proxy and self.proxy_enabled: - self.set_proxy( + self.socket.set_proxy( proxy_type=socks.SOCKS5, addr=proxy.get("hostname", None), port=proxy.get("port", None), @@ -52,26 +57,81 @@ class TCP(socks.socksocket): proxy.get("port", None) )) + async def connect(self, address: tuple): + self.socket.connect(address) + self.reader, self.writer = await asyncio.open_connection(sock=self.socket) + def close(self): try: - self.shutdown(socket.SHUT_RDWR) - except OSError: - pass - finally: - super().close() + self.writer.close() + except AttributeError: + try: + self.socket.shutdown(socket.SHUT_RDWR) + except OSError: + pass + finally: + self.socket.close() - def recvall(self, length: int) -> bytes or None: + async def send(self, data: bytes): + self.writer.write(data) + await self.writer.drain() + + async def recv(self, length: int = 0): data = b"" while len(data) < length: try: - packet = super().recv(length - len(data)) + chunk = await self.reader.read(length - len(data)) except OSError: return None else: - if packet: - data += packet + if chunk: + data += chunk else: return None return data + +# class TCP(socks.socksocket): +# def __init__(self, proxy: dict): +# super().__init__() +# self.settimeout(10) +# self.proxy_enabled = proxy.get("enabled", False) +# +# if proxy and self.proxy_enabled: +# self.set_proxy( +# proxy_type=socks.SOCKS5, +# addr=proxy.get("hostname", None), +# port=proxy.get("port", None), +# username=proxy.get("username", None), +# password=proxy.get("password", None) +# ) +# +# log.info("Using proxy {}:{}".format( +# proxy.get("hostname", None), +# proxy.get("port", None) +# )) +# +# def close(self): +# try: +# self.shutdown(socket.SHUT_RDWR) +# except OSError: +# pass +# finally: +# super().close() +# +# def recvall(self, length: int) -> bytes or None: +# data = b"" +# +# while len(data) < length: +# try: +# packet = super().recv(length - len(data)) +# except OSError: +# return None +# else: +# if packet: +# data += packet +# else: +# return None +# +# return data