From 768bf1760a45eb94ccd4a93eff504f2e04f54beb Mon Sep 17 00:00:00 2001 From: Yuval Zaif Date: Wed, 1 Jan 2025 13:55:31 +0200 Subject: [PATCH] Add support in socks5 proxy --- pymongo/asynchronous/pool.py | 25 ++++++++++++++++++++++++- pymongo/asynchronous/topology.py | 1 + pymongo/client_options.py | 10 ++++++++++ pymongo/common.py | 4 ++++ pymongo/pool_options.py | 9 ++++++++- pymongo/synchronous/pool.py | 20 +++++++++++++++++++- pymongo/synchronous/topology.py | 1 + pyproject.toml | 1 + requirements/socks.txt | 1 + 9 files changed, 69 insertions(+), 3 deletions(-) create mode 100644 requirements/socks.txt diff --git a/pymongo/asynchronous/pool.py b/pymongo/asynchronous/pool.py index 5dc5675a0a..1c9e198522 100644 --- a/pymongo/asynchronous/pool.py +++ b/pymongo/asynchronous/pool.py @@ -120,6 +120,11 @@ def _set_non_inheritable_non_atomic(fd: int) -> None: def _set_non_inheritable_non_atomic(fd: int) -> None: # noqa: ARG001 """Dummy function for platforms that don't provide fcntl.""" +try: + from python_socks.sync import Proxy + from python_socks import ProxyType +except ImportError: + Proxy = ProxyType = None _IS_SYNC = False @@ -838,7 +843,25 @@ def _create_connection(address: _Address, options: PoolOptions) -> socket.socket sock.settimeout(timeout) sock.setsockopt(socket.SOL_SOCKET, socket.SO_KEEPALIVE, True) _set_keepalive_times(sock) - sock.connect(sa) + if proxy := options.proxy: + if Proxy is None: + raise RuntimeError( + "In order to use SOCKS5 proxy, python_socks must be installed. " + "This can be done by re-installing pymongo with `pip install pymongo[socks]`" + ) + proxy_host = proxy['host'] + proxy_port = proxy['port'] or 1080 + sock.connect((proxy_host, proxy_port)) + proxy = Proxy( + ProxyType.SOCKS5, + proxy_host, + proxy_port, + proxy['username'], + proxy['password'] + ) + proxy.connect(sa[0], dest_port=sa[1], _socket=sock) + else: + sock.connect(sa) return sock except OSError as e: err = e diff --git a/pymongo/asynchronous/topology.py b/pymongo/asynchronous/topology.py index 6d67710a7e..0883c92f5e 100644 --- a/pymongo/asynchronous/topology.py +++ b/pymongo/asynchronous/topology.py @@ -967,6 +967,7 @@ def _create_pool_for_monitor(self, address: _Address) -> Pool: driver=options.driver, pause_enabled=False, server_api=options.server_api, + proxy=options.proxy, ) return self._settings.pool_class( diff --git a/pymongo/client_options.py b/pymongo/client_options.py index 9b9b88a736..cd30ac81b1 100644 --- a/pymongo/client_options.py +++ b/pymongo/client_options.py @@ -170,6 +170,15 @@ def _parse_pool_options( ssl_context, tls_allow_invalid_hostnames = _parse_ssl_options(options) load_balanced = options.get("loadbalanced") max_connecting = options.get("maxconnecting", common.MAX_CONNECTING) + if proxy_host := options.get("proxyHost"): + proxy = { + "host": proxy_host, + "port": options.get("proxyPort"), + "username": options.get("proxyUserName"), + "password": options.get("proxyPassword"), + } + else: + proxy = None return PoolOptions( max_pool_size, min_pool_size, @@ -188,6 +197,7 @@ def _parse_pool_options( load_balanced=load_balanced, credentials=credentials, is_sync=is_sync, + proxy=proxy, ) diff --git a/pymongo/common.py b/pymongo/common.py index 5661de011c..a0139e01c3 100644 --- a/pymongo/common.py +++ b/pymongo/common.py @@ -729,6 +729,10 @@ def validate_server_monitoring_mode(option: str, value: str) -> str: "srvmaxhosts": validate_non_negative_integer, "timeoutms": validate_timeoutms, "servermonitoringmode": validate_server_monitoring_mode, + "proxyhost": validate_string, + "proxyport": validate_positive_integer_or_none, + "proxyusername": validate_string_or_none, + "proxypassword": validate_string_or_none, } # Dictionary where keys are the names of URI options specific to pymongo, diff --git a/pymongo/pool_options.py b/pymongo/pool_options.py index 038dbb3b5d..c17330caaa 100644 --- a/pymongo/pool_options.py +++ b/pymongo/pool_options.py @@ -312,6 +312,7 @@ class PoolOptions: "__server_api", "__load_balanced", "__credentials", + "__proxy", ) def __init__( @@ -334,6 +335,7 @@ def __init__( load_balanced: Optional[bool] = None, credentials: Optional[MongoCredential] = None, is_sync: Optional[bool] = True, + proxy: Optional[dict] = None, ): self.__max_pool_size = max_pool_size self.__min_pool_size = min_pool_size @@ -353,7 +355,7 @@ def __init__( self.__load_balanced = load_balanced self.__credentials = credentials self.__metadata = copy.deepcopy(_METADATA) - + self.__proxy = copy.deepcopy(proxy) if appname: self.__metadata["application"] = {"name": appname} @@ -522,3 +524,8 @@ def server_api(self) -> Optional[ServerApi]: def load_balanced(self) -> Optional[bool]: """True if this Pool is configured in load balanced mode.""" return self.__load_balanced + + @property + def proxy(self) -> Optional[dict]: + """Proxy settings, if configured""" + return self.__proxy diff --git a/pymongo/synchronous/pool.py b/pymongo/synchronous/pool.py index 1a155c82d7..a7ce0def1a 100644 --- a/pymongo/synchronous/pool.py +++ b/pymongo/synchronous/pool.py @@ -836,7 +836,25 @@ def _create_connection(address: _Address, options: PoolOptions) -> socket.socket sock.settimeout(timeout) sock.setsockopt(socket.SOL_SOCKET, socket.SO_KEEPALIVE, True) _set_keepalive_times(sock) - sock.connect(sa) + if proxy := options.proxy: + if Proxy is None: + raise RuntimeError( + "In order to use SOCKS5 proxy, python_socks must be installed. " + "This can be done by re-installing pymongo with `pip install pymongo[socks]`" + ) + proxy_host = proxy['host'] + proxy_port = proxy['port'] or 1080 + sock.connect((proxy_host, proxy_port)) + proxy = Proxy( + ProxyType.SOCKS5, + proxy_host, + proxy_port, + proxy['username'], + proxy['password'] + ) + proxy.connect(sa[0], dest_port=sa[1], _socket=sock) + else: + sock.connect(sa) return sock except OSError as e: err = e diff --git a/pymongo/synchronous/topology.py b/pymongo/synchronous/topology.py index b03269ae43..f3b805d10b 100644 --- a/pymongo/synchronous/topology.py +++ b/pymongo/synchronous/topology.py @@ -965,6 +965,7 @@ def _create_pool_for_monitor(self, address: _Address) -> Pool: driver=options.driver, pause_enabled=False, server_api=options.server_api, + proxy=options.proxy, ) return self._settings.pool_class( diff --git a/pyproject.toml b/pyproject.toml index 9a29a777fc..54ed0676e4 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -67,6 +67,7 @@ ocsp = ["requirements/ocsp.txt"] snappy = ["requirements/snappy.txt"] test = ["requirements/test.txt"] zstd = ["requirements/zstd.txt"] +socks = ["requirements/socks.txt"] [tool.pytest.ini_options] minversion = "7" diff --git a/requirements/socks.txt b/requirements/socks.txt new file mode 100644 index 0000000000..eb6c7a304c --- /dev/null +++ b/requirements/socks.txt @@ -0,0 +1 @@ +python-socks[asyncio]