1# -*- coding: utf-8 -*-
2"""
3This module contains provisional support for SOCKS proxies from within
4urllib3. This module supports SOCKS4, SOCKS4A (an extension of SOCKS4), and
5SOCKS5. To enable its functionality, either install PySocks or install this
6module with the ``socks`` extra.
7
8The SOCKS implementation supports the full range of urllib3 features. It also
9supports the following SOCKS features:
10
11- SOCKS4A (``proxy_url='socks4a://...``)
12- SOCKS4 (``proxy_url='socks4://...``)
13- SOCKS5 with remote DNS (``proxy_url='socks5h://...``)
14- SOCKS5 with local DNS (``proxy_url='socks5://...``)
15- Usernames and passwords for the SOCKS proxy
16
17.. note::
18   It is recommended to use ``socks5h://`` or ``socks4a://`` schemes in
19   your ``proxy_url`` to ensure that DNS resolution is done from the remote
20   server instead of client-side when connecting to a domain name.
21
22SOCKS4 supports IPv4 and domain names with the SOCKS4A extension. SOCKS5
23supports IPv4, IPv6, and domain names.
24
25When connecting to a SOCKS4 proxy the ``username`` portion of the ``proxy_url``
26will be sent as the ``userid`` section of the SOCKS request:
27
28.. code-block:: python
29
30    proxy_url="socks4a://<userid>@proxy-host"
31
32When connecting to a SOCKS5 proxy the ``username`` and ``password`` portion
33of the ``proxy_url`` will be sent as the username/password to authenticate
34with the proxy:
35
36.. code-block:: python
37
38    proxy_url="socks5h://<username>:<password>@proxy-host"
39
40"""
41from __future__ import absolute_import
42
43try:
44    import socks
45except ImportError:
46    import warnings
47
48    from ..exceptions import DependencyWarning
49
50    warnings.warn(
51        (
52            "SOCKS support in urllib3 requires the installation of optional "
53            "dependencies: specifically, PySocks.  For more information, see "
54            "https://urllib3.readthedocs.io/en/latest/contrib.html#socks-proxies"
55        ),
56        DependencyWarning,
57    )
58    raise
59
60from socket import error as SocketError
61from socket import timeout as SocketTimeout
62
63from ..connection import HTTPConnection, HTTPSConnection
64from ..connectionpool import HTTPConnectionPool, HTTPSConnectionPool
65from ..exceptions import ConnectTimeoutError, NewConnectionError
66from ..poolmanager import PoolManager
67from ..util.url import parse_url
68
69try:
70    import ssl
71except ImportError:
72    ssl = None
73
74
75class SOCKSConnection(HTTPConnection):
76    """
77    A plain-text HTTP connection that connects via a SOCKS proxy.
78    """
79
80    def __init__(self, *args, **kwargs):
81        self._socks_options = kwargs.pop("_socks_options")
82        super(SOCKSConnection, self).__init__(*args, **kwargs)
83
84    def _new_conn(self):
85        """
86        Establish a new connection via the SOCKS proxy.
87        """
88        extra_kw = {}
89        if self.source_address:
90            extra_kw["source_address"] = self.source_address
91
92        if self.socket_options:
93            extra_kw["socket_options"] = self.socket_options
94
95        try:
96            conn = socks.create_connection(
97                (self.host, self.port),
98                proxy_type=self._socks_options["socks_version"],
99                proxy_addr=self._socks_options["proxy_host"],
100                proxy_port=self._socks_options["proxy_port"],
101                proxy_username=self._socks_options["username"],
102                proxy_password=self._socks_options["password"],
103                proxy_rdns=self._socks_options["rdns"],
104                timeout=self.timeout,
105                **extra_kw
106            )
107
108        except SocketTimeout:
109            raise ConnectTimeoutError(
110                self,
111                "Connection to %s timed out. (connect timeout=%s)"
112                % (self.host, self.timeout),
113            )
114
115        except socks.ProxyError as e:
116            # This is fragile as hell, but it seems to be the only way to raise
117            # useful errors here.
118            if e.socket_err:
119                error = e.socket_err
120                if isinstance(error, SocketTimeout):
121                    raise ConnectTimeoutError(
122                        self,
123                        "Connection to %s timed out. (connect timeout=%s)"
124                        % (self.host, self.timeout),
125                    )
126                else:
127                    raise NewConnectionError(
128                        self, "Failed to establish a new connection: %s" % error
129                    )
130            else:
131                raise NewConnectionError(
132                    self, "Failed to establish a new connection: %s" % e
133                )
134
135        except SocketError as e:  # Defensive: PySocks should catch all these.
136            raise NewConnectionError(
137                self, "Failed to establish a new connection: %s" % e
138            )
139
140        return conn
141
142
143# We don't need to duplicate the Verified/Unverified distinction from
144# urllib3/connection.py here because the HTTPSConnection will already have been
145# correctly set to either the Verified or Unverified form by that module. This
146# means the SOCKSHTTPSConnection will automatically be the correct type.
147class SOCKSHTTPSConnection(SOCKSConnection, HTTPSConnection):
148    pass
149
150
151class SOCKSHTTPConnectionPool(HTTPConnectionPool):
152    ConnectionCls = SOCKSConnection
153
154
155class SOCKSHTTPSConnectionPool(HTTPSConnectionPool):
156    ConnectionCls = SOCKSHTTPSConnection
157
158
159class SOCKSProxyManager(PoolManager):
160    """
161    A version of the urllib3 ProxyManager that routes connections via the
162    defined SOCKS proxy.
163    """
164
165    pool_classes_by_scheme = {
166        "http": SOCKSHTTPConnectionPool,
167        "https": SOCKSHTTPSConnectionPool,
168    }
169
170    def __init__(
171        self,
172        proxy_url,
173        username=None,
174        password=None,
175        num_pools=10,
176        headers=None,
177        **connection_pool_kw
178    ):
179        parsed = parse_url(proxy_url)
180
181        if username is None and password is None and parsed.auth is not None:
182            split = parsed.auth.split(":")
183            if len(split) == 2:
184                username, password = split
185        if parsed.scheme == "socks5":
186            socks_version = socks.PROXY_TYPE_SOCKS5
187            rdns = False
188        elif parsed.scheme == "socks5h":
189            socks_version = socks.PROXY_TYPE_SOCKS5
190            rdns = True
191        elif parsed.scheme == "socks4":
192            socks_version = socks.PROXY_TYPE_SOCKS4
193            rdns = False
194        elif parsed.scheme == "socks4a":
195            socks_version = socks.PROXY_TYPE_SOCKS4
196            rdns = True
197        else:
198            raise ValueError("Unable to determine SOCKS version from %s" % proxy_url)
199
200        self.proxy_url = proxy_url
201
202        socks_options = {
203            "socks_version": socks_version,
204            "proxy_host": parsed.host,
205            "proxy_port": parsed.port,
206            "username": username,
207            "password": password,
208            "rdns": rdns,
209        }
210        connection_pool_kw["_socks_options"] = socks_options
211
212        super(SOCKSProxyManager, self).__init__(
213            num_pools, headers, **connection_pool_kw
214        )
215
216        self.pool_classes_by_scheme = SOCKSProxyManager.pool_classes_by_scheme
217