[py] Handle Socks Proxy for Remote Connections. Fixes #10091

AutomatedTester · AutomatedTester · commit b25378793ee4 · 2021-11-29T12:36:05.000Z
diff --git a/py/BUILD.bazel b/py/BUILD.bazel
@@ -129,7 +129,7 @@ py_wheel(
     python_requires = "~=3.7",
     python_tag = "py3",
     requires = [
-        "urllib3[secure]~=1.26",
+        "urllib3[secure, socks]~=1.26",
         "trio~=0.17",
         "trio-websocket~=0.9",
     ],
diff --git a/py/requirements.txt b/py/requirements.txt
@@ -20,5 +20,5 @@ sortedcontainers
 sniffio
 trio~=0.17
 trio_websocket~=0.9
-urllib3~=1.26
+urllib3[secure, socks]~=1.26
 zipp
diff --git a/py/selenium/webdriver/remote/remote_connection.py b/py/selenium/webdriver/remote/remote_connection.py
@@ -136,8 +136,13 @@ def _get_connection_manager(self):
             pool_manager_init_args['cert_reqs'] = 'CERT_REQUIRED'
             pool_manager_init_args['ca_certs'] = self._ca_certs
 
-        return urllib3.PoolManager(**pool_manager_init_args) if not self._proxy_url else \
-            urllib3.ProxyManager(self._proxy_url, **pool_manager_init_args)
+        if self._proxy_url:
+            if self._proxy_url.startswith('sock'):
+                from urllib3.contrib.socks import SOCKSProxyManager
+                return SOCKSProxyManager(self._proxy_url, **pool_manager_init_args)
+            return urllib3.ProxyManager(self._proxy_url, **pool_manager_init_args)
+
+        return urllib3.PoolManager(**pool_manager_init_args)
 
     def __init__(self, remote_server_addr, keep_alive=False, resolve_ip=None, ignore_proxy=False):
         if resolve_ip:
diff --git a/py/test/unit/selenium/webdriver/remote/remote_connection_tests.py b/py/test/unit/selenium/webdriver/remote/remote_connection_tests.py
@@ -118,6 +118,13 @@ def test_ignore_proxy_env_vars(mock_proxy_settings):
     assert type(conn) == urllib3.PoolManager
 
 
+def test_get_socks_proxy_when_set(mock_socks_proxy_settings):
+    remote_connection = RemoteConnection("http://127.0.0.1:4444/wd/hub")
+    conn = remote_connection._get_connection_manager()
+    from urllib3.contrib.socks import SOCKSProxyManager
+    assert type(conn) == SOCKSProxyManager
+
+
 class MockResponse:
     code = 200
     headers = []
@@ -140,6 +147,16 @@ def mock_proxy_settings_missing(monkeypatch):
     monkeypatch.delenv("http_proxy", raising=False)
 
 
+@pytest.fixture(scope="function")
+def mock_socks_proxy_settings(monkeypatch):
+    http_proxy = 'socks5://http_proxy.com:8080'
+    https_proxy = 'socks5://https_proxy.com:8080'
+    monkeypatch.setenv("HTTPS_PROXY", https_proxy)
+    monkeypatch.setenv("HTTP_PROXY", http_proxy)
+    monkeypatch.setenv("https_proxy", https_proxy)
+    monkeypatch.setenv("http_proxy", http_proxy)
+
+
 @pytest.fixture(scope="function")
 def mock_proxy_settings(monkeypatch):
     http_proxy = 'http://http_proxy.com:8080'
