Merge replace PySocks with python-socks for Python >= 3.6 (#1627)

optional-requirements.txt

@@ -1,4 +1,5 @@
 cryptg
 pysocks
+python-socks[asyncio]
 hachoir
 pillow

readthedocs/basic/signing-in.rst

@@ -117,7 +117,12 @@ Signing In behind a Proxy
 =========================
 
 If you need to use a proxy to access Telegram,
-you will need to  `install PySocks`__ and then change:
+you will need to either:
+
+* For Python >= 3.6 : `install python-socks[asyncio]`__
+* For Python <= 3.5 : `install PySocks`__
+
+and then change
 
 .. code-block:: python
 
@@ -127,13 +132,47 @@ with
 
 .. code-block:: python
 
-    TelegramClient('anon', api_id, api_hash, proxy=(socks.SOCKS5, '127.0.0.1', 4444))
+    TelegramClient('anon', api_id, api_hash, proxy=("socks5", '127.0.0.1', 4444))
 
-(of course, replacing the IP and port with the IP and port of the proxy).
+(of course, replacing the protocol, IP and port with the protocol, IP and port of the proxy).
 
-The ``proxy=`` argument should be a tuple, a list or a dict,
+The ``proxy=`` argument should be a dict (or tuple, for backwards compatibility),
 consisting of parameters described `in PySocks usage`__.
 
+The allowed values for the argument ``proxy_type`` are:
+
+* For Python <= 3.5:
+    * ``socks.SOCKS5`` or ``'socks5'``
+    * ``socks.SOCKS4`` or ``'socks4'``
+    * ``socks.HTTP`` or ``'http'``
+
+* For Python >= 3.6:
+    * All of the above
+    * ``python_socks.SOCKS5``
+    * ``python_socks.SOCKS4``
+    * ``python_socks.HTTP``
+
+
+Example:
+
+.. code-block:: python
+
+    proxy = {
+        'proxy_type': 'socks5', # (mandatory) protocol to use (see above)
+        'addr': '1.1.1.1',      # (mandatory) proxy IP address
+        'port': 5555,           # (mandatory) proxy port number
+        'username': 'foo',      # (optional) username if the proxy requires auth
+        'password': 'bar',      # (optional) password if the proxy requires auth
+        'rdns': True            # (optional) whether to use remote or local resolve, default remote
+    }
+
+For backwards compatibility with ``PySocks`` the following format
+is possible (but discouraged):
+
+.. code-block:: python
+    proxy = (socks.SOCKS5, '1.1.1.1', 5555, True, 'foo', 'bar')
+
+.. __: https://github.com/romis2012/python-socks#installation
 .. __: https://github.com/Anorov/PySocks#installation
 .. __: https://github.com/Anorov/PySocks#usage-1
 

telethon/client/telegrambaseclient.py

@@ -110,8 +110,8 @@ class TelegramBaseClient(abc.ABC):
             function parameters for PySocks, like ``(type, 'hostname', port)``.
             See https://github.com/Anorov/PySocks#usage-1 for more.
 
-        local_addr (`str`, optional):
+        local_addr (`str` | `tuple`, optional):
-            Local host address used to bind the socket to locally.
+            Local host address (and port, optionally) used to bind the socket to locally.
             You only need to use this if you have multiple network cards and
             want to use a specific one.
 
@@ -220,7 +220,7 @@ class TelegramBaseClient(abc.ABC):
             connection: 'typing.Type[Connection]' = ConnectionTcpFull,
             use_ipv6: bool = False,
             proxy: typing.Union[tuple, dict] = None,
-            local_addr=None,
+            local_addr: typing.Union[str, tuple] = None,
             timeout: int = 10,
             request_retries: int = 5,
             connection_retries: int = 5,

telethon/network/connection/connection.py

@@ -45,54 +45,228 @@ class Connection(abc.ABC):
         self._send_queue = asyncio.Queue(1)
         self._recv_queue = asyncio.Queue(1)
 
-    async def _connect(self, timeout=None, ssl=None):
+    @staticmethod
-        if not self._proxy:
+    def _wrap_socket_ssl(sock):
-            if self._local_addr is not None:
-                local_addr = (self._local_addr, None)
-            else:
-                local_addr = None
 
-            self._reader, self._writer = await asyncio.wait_for(
-                asyncio.open_connection(self._ip, self._port, ssl=ssl, local_addr=local_addr),
-                timeout=timeout
-            )
-        else:
-            import socks
-            if ':' in self._ip:
-                mode, address = socket.AF_INET6, (self._ip, self._port, 0, 0)
-            else:
-                mode, address = socket.AF_INET, (self._ip, self._port)
-
-            s = socks.socksocket(mode, socket.SOCK_STREAM)
-            if isinstance(self._proxy, dict):
-                s.set_proxy(**self._proxy)
-            else:
-                s.set_proxy(*self._proxy)
-
-            s.settimeout(timeout)
-            if self._local_addr is not None:
-                s.bind((self._local_addr, None))
-            await asyncio.wait_for(
-                asyncio.get_event_loop().sock_connect(s, address),
-                timeout=timeout
-            )
-            if ssl:
         if ssl_mod is None:
             raise RuntimeError(
                 'Cannot use proxy that requires SSL '
                 'without the SSL module being available'
             )
 
-                s = ssl_mod.wrap_socket(
+        return ssl_mod.wrap_socket(
-                    s,
+            sock,
             do_handshake_on_connect=True,
             ssl_version=ssl_mod.PROTOCOL_SSLv23,
-                    ciphers='ADH-AES256-SHA'
+            ciphers='ADH-AES256-SHA')
+
+    @staticmethod
+    def _parse_proxy(library, proxy_type, addr, port, rdns=True, username=None, password=None):
+
+        if isinstance(proxy_type, str):
+            proxy_type = proxy_type.lower()
+
+        if library == "python-socks":
+
+            from python_socks import ProxyType
+
+            # We do the check for numerical values here
+            # to be backwards compatible with PySocks proxy format,
+            # (since socks.SOCKS5 == 2, socks.SOCKS4 == 1, socks.HTTP == 3)
+
+            if proxy_type == ProxyType.SOCKS5 or proxy_type == 2 or proxy_type == "socks5":
+                protocol = ProxyType.SOCKS5
+            elif proxy_type == ProxyType.SOCKS4 or proxy_type == 1 or proxy_type == "socks4":
+                protocol = ProxyType.SOCKS4
+            elif proxy_type == ProxyType.HTTP or proxy_type == 3 or proxy_type == "http":
+                protocol = ProxyType.HTTP
+            else:
+                raise ValueError("Unknown proxy protocol type: {}".format(proxy_type))
+
+            # NOTE: We return a tuple compatible with the
+            # signature (order) of python_socks `Proxy.create()`
+            # I.E.: (proxy_type, host, port, username, password, rdns)
+
+            return protocol, addr, port, username, password, rdns
+
+        elif library == "pysocks":
+
+            from socks import SOCKS5, SOCKS4, HTTP
+
+            if proxy_type == 2 or proxy_type == "socks5":
+                protocol = SOCKS5
+            elif proxy_type == 1 or proxy_type == "socks4":
+                protocol = SOCKS4
+            elif proxy_type == 3 or proxy_type == "http":
+                protocol = HTTP
+            else:
+                raise ValueError("Unknown proxy protocol type: {}".format(proxy_type))
+
+            # NOTE: We return a tuple compatible with the
+            # signature of `PySocks` `socksocket.set_proxy()`
+            # I.E.: (proxy_type, addr, port, rdns, username, password)
+
+            return protocol, addr, port, rdns, username, password
+
+        else:
+            raise ValueError("Unknown proxy library: {}".format(library))
+
+    async def _proxy_connect(self, timeout=None, local_addr=None):
+
+        # Use `python-socks` library for newer Python >= 3.6
+        # Use `PySocks` library for older Python <= 3.5
+        if sys.version_info >= (3, 6):
+
+            import python_socks
+
+            # python_socks internal errors are not inherited from
+            # builtin IOError (just from Exception). Instead of adding those
+            # in exceptions clauses everywhere through the code, we
+            # rather monkey-patch them in place.
+
+            python_socks._errors.ProxyError = ConnectionError
+            python_socks._errors.ProxyConnectionError = ConnectionError
+            python_socks._errors.ProxyTimeoutError = ConnectionError
+
+            from python_socks.async_.asyncio import Proxy
+
+            # We expect a dict/tuple in the format of `PySocks` (for compatibility)
+            # I.E.: (proxy_type, addr, port, rdns, username, password).
+
+            if isinstance(self._proxy, (tuple, list)):
+                parsed = self._parse_proxy("python-socks", *self._proxy)
+            elif isinstance(self._proxy, dict):
+                parsed = self._parse_proxy("python-socks", **self._proxy)
+            else:
+                raise TypeError("Proxy of unknown format!", type(self._proxy))
+
+            proxy = Proxy.create(*parsed)
+
+            # WARNING: If `local_addr` is set we use manual socket creation, because,
+            # unfortunately, `Proxy.connect()` does not expose `local_addr`
+            # argument, so if we want to bind socket locally, we need to manually
+            # create, bind and connect socket, and then pass to `Proxy.connect()` method.
+
+            if local_addr is None:
+                sock = await proxy.connect(
+                    dest_host=self._ip,
+                    dest_port=self._port,
+                    timeout=timeout
+                )
+            else:
+
+                # Here we start manual setup of the socket.
+                # The `address` represents the proxy ip and proxy port,
+                # not the destination one (!), because the socket
+                # connects to the proxy server, not destination server.
+                # IPv family is also checked on proxy address.
+
+                if ':' in proxy.proxy_host:
+                    mode, address = socket.AF_INET6, (proxy.proxy_host, proxy.proxy_port, 0, 0)
+                else:
+                    mode, address = socket.AF_INET, (proxy.proxy_host, proxy.proxy_port)
+
+                # Create a non-blocking socket and bind it (if local address is specified).
+                sock = socket.socket(mode, socket.SOCK_STREAM)
+                sock.setblocking(False)
+                sock.bind(local_addr)
+
+                # Actual TCP connection is performed here.
+                await asyncio.wait_for(
+                    asyncio.get_event_loop().sock_connect(sock=sock, address=address),
+                    timeout=timeout
                 )
 
-            s.setblocking(False)
+                # As our socket is already created and connected,
+                # this call sets the destination host/port and
+                # starts protocol negotiations with the proxy server.
 
-            self._reader, self._writer = await asyncio.open_connection(sock=s)
+                sock = await proxy.connect(
+                    dest_host=self._ip,
+                    dest_port=self._port,
+                    timeout=timeout,
+                    _socket=sock
+                )
+
+        else:
+
+            import socks
+
+            # We expect a dict/tuple in the format of `PySocks`.
+            # I.E.: (proxy_type, addr, port, rdns, username, password).
+
+            if isinstance(self._proxy, (tuple, list)):
+                parsed = self._parse_proxy("pysocks", *self._proxy)
+            elif isinstance(self._proxy, dict):
+                parsed = self._parse_proxy("pysocks", **self._proxy)
+            else:
+                raise TypeError("Proxy of unknown format!", type(self._proxy))
+
+            # Here `address` represents destination address (not proxy), because of
+            # the `PySocks` implementation of the connection routine.
+            # IPv family is checked on proxy address, not destination address.
+
+            if ':' in parsed[1]:
+                mode, address = socket.AF_INET6, (self._ip, self._port, 0, 0)
+            else:
+                mode, address = socket.AF_INET, (self._ip, self._port)
+
+            # Setup socket, proxy, timeout and bind it (if necessary).
+            sock = socks.socksocket(mode, socket.SOCK_STREAM)
+            sock.set_proxy(*parsed)
+            sock.settimeout(timeout)
+
+            if local_addr is not None:
+                sock.bind(local_addr)
+
+            # Actual TCP connection and negotiation performed here.
+            await asyncio.wait_for(
+                asyncio.get_event_loop().sock_connect(sock=sock, address=address),
+                timeout=timeout
+            )
+
+            sock.setblocking(False)
+
+        return sock
+
+    async def _connect(self, timeout=None, ssl=None):
+
+        if self._local_addr is not None:
+
+            # NOTE: If port is not specified, we use 0 port
+            # to notify the OS that port should be chosen randomly
+            # from the available ones.
+
+            if isinstance(self._local_addr, tuple) and len(self._local_addr) == 2:
+                local_addr = self._local_addr
+            elif isinstance(self._local_addr, str):
+                local_addr = (self._local_addr, 0)
+            else:
+                raise ValueError("Unknown local address format: {}".format(self._local_addr))
+        else:
+            local_addr = None
+
+        if not self._proxy:
+            self._reader, self._writer = await asyncio.wait_for(
+                asyncio.open_connection(
+                    host=self._ip,
+                    port=self._port,
+                    ssl=ssl,
+                    local_addr=local_addr
+                ), timeout=timeout)
+        else:
+
+            # Proxy setup, connection and negotiation is performed here.
+            sock = await self._proxy_connect(
+                timeout=timeout,
+                local_addr=local_addr
+            )
+
+            # Wrap socket in SSL context (if provided)
+            if ssl:
+                sock = self._wrap_socket_ssl(sock)
+
+            self._reader, self._writer = await asyncio.open_connection(sock=sock)
 
         self._codec = self.packet_codec(self)
         self._init_conn()