jpayne@7: """
jpayne@7: This module contains provisional support for SOCKS proxies from within
jpayne@7: urllib3. This module supports SOCKS4, SOCKS4A (an extension of SOCKS4), and
jpayne@7: SOCKS5. To enable its functionality, either install PySocks or install this
jpayne@7: module with the ``socks`` extra.
jpayne@7: 
jpayne@7: The SOCKS implementation supports the full range of urllib3 features. It also
jpayne@7: supports the following SOCKS features:
jpayne@7: 
jpayne@7: - SOCKS4A (``proxy_url='socks4a://...``)
jpayne@7: - SOCKS4 (``proxy_url='socks4://...``)
jpayne@7: - SOCKS5 with remote DNS (``proxy_url='socks5h://...``)
jpayne@7: - SOCKS5 with local DNS (``proxy_url='socks5://...``)
jpayne@7: - Usernames and passwords for the SOCKS proxy
jpayne@7: 
jpayne@7: .. note::
jpayne@7:    It is recommended to use ``socks5h://`` or ``socks4a://`` schemes in
jpayne@7:    your ``proxy_url`` to ensure that DNS resolution is done from the remote
jpayne@7:    server instead of client-side when connecting to a domain name.
jpayne@7: 
jpayne@7: SOCKS4 supports IPv4 and domain names with the SOCKS4A extension. SOCKS5
jpayne@7: supports IPv4, IPv6, and domain names.
jpayne@7: 
jpayne@7: When connecting to a SOCKS4 proxy the ``username`` portion of the ``proxy_url``
jpayne@7: will be sent as the ``userid`` section of the SOCKS request:
jpayne@7: 
jpayne@7: .. code-block:: python
jpayne@7: 
jpayne@7:     proxy_url="socks4a://<userid>@proxy-host"
jpayne@7: 
jpayne@7: When connecting to a SOCKS5 proxy the ``username`` and ``password`` portion
jpayne@7: of the ``proxy_url`` will be sent as the username/password to authenticate
jpayne@7: with the proxy:
jpayne@7: 
jpayne@7: .. code-block:: python
jpayne@7: 
jpayne@7:     proxy_url="socks5h://<username>:<password>@proxy-host"
jpayne@7: 
jpayne@7: """
jpayne@7: 
jpayne@7: from __future__ import annotations
jpayne@7: 
jpayne@7: try:
jpayne@7:     import socks  # type: ignore[import-not-found]
jpayne@7: except ImportError:
jpayne@7:     import warnings
jpayne@7: 
jpayne@7:     from ..exceptions import DependencyWarning
jpayne@7: 
jpayne@7:     warnings.warn(
jpayne@7:         (
jpayne@7:             "SOCKS support in urllib3 requires the installation of optional "
jpayne@7:             "dependencies: specifically, PySocks.  For more information, see "
jpayne@7:             "https://urllib3.readthedocs.io/en/latest/advanced-usage.html#socks-proxies"
jpayne@7:         ),
jpayne@7:         DependencyWarning,
jpayne@7:     )
jpayne@7:     raise
jpayne@7: 
jpayne@7: import typing
jpayne@7: from socket import timeout as SocketTimeout
jpayne@7: 
jpayne@7: from ..connection import HTTPConnection, HTTPSConnection
jpayne@7: from ..connectionpool import HTTPConnectionPool, HTTPSConnectionPool
jpayne@7: from ..exceptions import ConnectTimeoutError, NewConnectionError
jpayne@7: from ..poolmanager import PoolManager
jpayne@7: from ..util.url import parse_url
jpayne@7: 
jpayne@7: try:
jpayne@7:     import ssl
jpayne@7: except ImportError:
jpayne@7:     ssl = None  # type: ignore[assignment]
jpayne@7: 
jpayne@7: from typing import TypedDict
jpayne@7: 
jpayne@7: 
jpayne@7: class _TYPE_SOCKS_OPTIONS(TypedDict):
jpayne@7:     socks_version: int
jpayne@7:     proxy_host: str | None
jpayne@7:     proxy_port: str | None
jpayne@7:     username: str | None
jpayne@7:     password: str | None
jpayne@7:     rdns: bool
jpayne@7: 
jpayne@7: 
jpayne@7: class SOCKSConnection(HTTPConnection):
jpayne@7:     """
jpayne@7:     A plain-text HTTP connection that connects via a SOCKS proxy.
jpayne@7:     """
jpayne@7: 
jpayne@7:     def __init__(
jpayne@7:         self,
jpayne@7:         _socks_options: _TYPE_SOCKS_OPTIONS,
jpayne@7:         *args: typing.Any,
jpayne@7:         **kwargs: typing.Any,
jpayne@7:     ) -> None:
jpayne@7:         self._socks_options = _socks_options
jpayne@7:         super().__init__(*args, **kwargs)
jpayne@7: 
jpayne@7:     def _new_conn(self) -> socks.socksocket:
jpayne@7:         """
jpayne@7:         Establish a new connection via the SOCKS proxy.
jpayne@7:         """
jpayne@7:         extra_kw: dict[str, typing.Any] = {}
jpayne@7:         if self.source_address:
jpayne@7:             extra_kw["source_address"] = self.source_address
jpayne@7: 
jpayne@7:         if self.socket_options:
jpayne@7:             extra_kw["socket_options"] = self.socket_options
jpayne@7: 
jpayne@7:         try:
jpayne@7:             conn = socks.create_connection(
jpayne@7:                 (self.host, self.port),
jpayne@7:                 proxy_type=self._socks_options["socks_version"],
jpayne@7:                 proxy_addr=self._socks_options["proxy_host"],
jpayne@7:                 proxy_port=self._socks_options["proxy_port"],
jpayne@7:                 proxy_username=self._socks_options["username"],
jpayne@7:                 proxy_password=self._socks_options["password"],
jpayne@7:                 proxy_rdns=self._socks_options["rdns"],
jpayne@7:                 timeout=self.timeout,
jpayne@7:                 **extra_kw,
jpayne@7:             )
jpayne@7: 
jpayne@7:         except SocketTimeout as e:
jpayne@7:             raise ConnectTimeoutError(
jpayne@7:                 self,
jpayne@7:                 f"Connection to {self.host} timed out. (connect timeout={self.timeout})",
jpayne@7:             ) from e
jpayne@7: 
jpayne@7:         except socks.ProxyError as e:
jpayne@7:             # This is fragile as hell, but it seems to be the only way to raise
jpayne@7:             # useful errors here.
jpayne@7:             if e.socket_err:
jpayne@7:                 error = e.socket_err
jpayne@7:                 if isinstance(error, SocketTimeout):
jpayne@7:                     raise ConnectTimeoutError(
jpayne@7:                         self,
jpayne@7:                         f"Connection to {self.host} timed out. (connect timeout={self.timeout})",
jpayne@7:                     ) from e
jpayne@7:                 else:
jpayne@7:                     # Adding `from e` messes with coverage somehow, so it's omitted.
jpayne@7:                     # See #2386.
jpayne@7:                     raise NewConnectionError(
jpayne@7:                         self, f"Failed to establish a new connection: {error}"
jpayne@7:                     )
jpayne@7:             else:
jpayne@7:                 raise NewConnectionError(
jpayne@7:                     self, f"Failed to establish a new connection: {e}"
jpayne@7:                 ) from e
jpayne@7: 
jpayne@7:         except OSError as e:  # Defensive: PySocks should catch all these.
jpayne@7:             raise NewConnectionError(
jpayne@7:                 self, f"Failed to establish a new connection: {e}"
jpayne@7:             ) from e
jpayne@7: 
jpayne@7:         return conn
jpayne@7: 
jpayne@7: 
jpayne@7: # We don't need to duplicate the Verified/Unverified distinction from
jpayne@7: # urllib3/connection.py here because the HTTPSConnection will already have been
jpayne@7: # correctly set to either the Verified or Unverified form by that module. This
jpayne@7: # means the SOCKSHTTPSConnection will automatically be the correct type.
jpayne@7: class SOCKSHTTPSConnection(SOCKSConnection, HTTPSConnection):
jpayne@7:     pass
jpayne@7: 
jpayne@7: 
jpayne@7: class SOCKSHTTPConnectionPool(HTTPConnectionPool):
jpayne@7:     ConnectionCls = SOCKSConnection
jpayne@7: 
jpayne@7: 
jpayne@7: class SOCKSHTTPSConnectionPool(HTTPSConnectionPool):
jpayne@7:     ConnectionCls = SOCKSHTTPSConnection
jpayne@7: 
jpayne@7: 
jpayne@7: class SOCKSProxyManager(PoolManager):
jpayne@7:     """
jpayne@7:     A version of the urllib3 ProxyManager that routes connections via the
jpayne@7:     defined SOCKS proxy.
jpayne@7:     """
jpayne@7: 
jpayne@7:     pool_classes_by_scheme = {
jpayne@7:         "http": SOCKSHTTPConnectionPool,
jpayne@7:         "https": SOCKSHTTPSConnectionPool,
jpayne@7:     }
jpayne@7: 
jpayne@7:     def __init__(
jpayne@7:         self,
jpayne@7:         proxy_url: str,
jpayne@7:         username: str | None = None,
jpayne@7:         password: str | None = None,
jpayne@7:         num_pools: int = 10,
jpayne@7:         headers: typing.Mapping[str, str] | None = None,
jpayne@7:         **connection_pool_kw: typing.Any,
jpayne@7:     ):
jpayne@7:         parsed = parse_url(proxy_url)
jpayne@7: 
jpayne@7:         if username is None and password is None and parsed.auth is not None:
jpayne@7:             split = parsed.auth.split(":")
jpayne@7:             if len(split) == 2:
jpayne@7:                 username, password = split
jpayne@7:         if parsed.scheme == "socks5":
jpayne@7:             socks_version = socks.PROXY_TYPE_SOCKS5
jpayne@7:             rdns = False
jpayne@7:         elif parsed.scheme == "socks5h":
jpayne@7:             socks_version = socks.PROXY_TYPE_SOCKS5
jpayne@7:             rdns = True
jpayne@7:         elif parsed.scheme == "socks4":
jpayne@7:             socks_version = socks.PROXY_TYPE_SOCKS4
jpayne@7:             rdns = False
jpayne@7:         elif parsed.scheme == "socks4a":
jpayne@7:             socks_version = socks.PROXY_TYPE_SOCKS4
jpayne@7:             rdns = True
jpayne@7:         else:
jpayne@7:             raise ValueError(f"Unable to determine SOCKS version from {proxy_url}")
jpayne@7: 
jpayne@7:         self.proxy_url = proxy_url
jpayne@7: 
jpayne@7:         socks_options = {
jpayne@7:             "socks_version": socks_version,
jpayne@7:             "proxy_host": parsed.host,
jpayne@7:             "proxy_port": parsed.port,
jpayne@7:             "username": username,
jpayne@7:             "password": password,
jpayne@7:             "rdns": rdns,
jpayne@7:         }
jpayne@7:         connection_pool_kw["_socks_options"] = socks_options
jpayne@7: 
jpayne@7:         super().__init__(num_pools, headers, **connection_pool_kw)
jpayne@7: 
jpayne@7:         self.pool_classes_by_scheme = SOCKSProxyManager.pool_classes_by_scheme