From 4682dac91e9c9ff9c202337c9337b8dc5755b34f Mon Sep 17 00:00:00 2001 From: Steven Silvester Date: Tue, 31 Dec 2024 16:29:48 -0600 Subject: [PATCH 1/4] PYTHON-3096 Finish implementation and tests for GSSAPI options (#1985) --- .../scripts/run-enterprise-auth-tests.sh | 3 +- doc/contributors.rst | 1 + pymongo/asynchronous/auth.py | 14 ++-- pymongo/asynchronous/pool.py | 22 ++++++- pymongo/asynchronous/topology.py | 1 + pymongo/client_options.py | 10 +++ pymongo/common.py | 4 ++ pymongo/pool_options.py | 9 ++- pymongo/synchronous/auth.py | 14 ++-- pymongo/synchronous/pool.py | 22 ++++++- pymongo/synchronous/topology.py | 1 + pyproject.toml | 1 + requirements/socks.txt | 1 + test/asynchronous/test_auth.py | 66 +++++++++++++++++-- test/test_auth.py | 66 +++++++++++++++++-- 15 files changed, 213 insertions(+), 22 deletions(-) create mode 100644 requirements/socks.txt diff --git a/.evergreen/scripts/run-enterprise-auth-tests.sh b/.evergreen/scripts/run-enterprise-auth-tests.sh index 11f8db22e1..7f936b1955 100755 --- a/.evergreen/scripts/run-enterprise-auth-tests.sh +++ b/.evergreen/scripts/run-enterprise-auth-tests.sh @@ -1,7 +1,8 @@ #!/bin/bash +set -eu # Disable xtrace for security reasons (just in case it was accidentally set). set +x # Use the default python to bootstrap secrets. -PYTHON_BINARY="" bash "${DRIVERS_TOOLS}"/.evergreen/auth_aws/setup_secrets.sh drivers/enterprise_auth +bash "${DRIVERS_TOOLS}"/.evergreen/secrets_handling/setup-secrets.sh drivers/enterprise_auth TEST_ENTERPRISE_AUTH=1 AUTH=auth bash "${PROJECT_DIRECTORY}"/.evergreen/hatch.sh test:test-eg diff --git a/doc/contributors.rst b/doc/contributors.rst index 4a7f5424b1..fbdad001db 100644 --- a/doc/contributors.rst +++ b/doc/contributors.rst @@ -103,3 +103,4 @@ The following is a list of people who have contributed to - Terry Patterson - Romain Morotti - Navjot Singh (navjots18) +- Yuval Zaif (zaif-yuval) diff --git a/pymongo/asynchronous/auth.py b/pymongo/asynchronous/auth.py index fc563ec48f..48ce4bbd39 100644 --- a/pymongo/asynchronous/auth.py +++ b/pymongo/asynchronous/auth.py @@ -177,13 +177,20 @@ def _auth_key(nonce: str, username: str, password: str) -> str: return md5hash.hexdigest() -def _canonicalize_hostname(hostname: str) -> str: +def _canonicalize_hostname(hostname: str, option: str | bool) -> str: """Canonicalize hostname following MIT-krb5 behavior.""" # https://github.com/krb5/krb5/blob/d406afa363554097ac48646a29249c04f498c88e/src/util/k5test.py#L505-L520 + if option in [False, "none"]: + return hostname + af, socktype, proto, canonname, sockaddr = socket.getaddrinfo( hostname, None, 0, 0, socket.IPPROTO_TCP, socket.AI_CANONNAME )[0] + # For forward just to resolve the cname as dns.lookup() will not return it. + if option == "forward": + return canonname.lower() + try: name = socket.getnameinfo(sockaddr, socket.NI_NAMEREQD) except socket.gaierror: @@ -205,9 +212,8 @@ async def _authenticate_gssapi(credentials: MongoCredential, conn: AsyncConnecti props = credentials.mechanism_properties # Starting here and continuing through the while loop below - establish # the security context. See RFC 4752, Section 3.1, first paragraph. - host = conn.address[0] - if props.canonicalize_host_name: - host = _canonicalize_hostname(host) + host = props.service_host or conn.address[0] + host = _canonicalize_hostname(host, props.canonicalize_host_name) service = props.service_name + "@" + host if props.service_realm is not None: service = service + "@" + props.service_realm diff --git a/pymongo/asynchronous/pool.py b/pymongo/asynchronous/pool.py index 5dc5675a0a..0bac28ed42 100644 --- a/pymongo/asynchronous/pool.py +++ b/pymongo/asynchronous/pool.py @@ -121,6 +121,12 @@ def _set_non_inheritable_non_atomic(fd: int) -> None: # noqa: ARG001 """Dummy function for platforms that don't provide fcntl.""" +try: + from python_socks import ProxyType + from python_socks.sync import Proxy +except ImportError: + Proxy = ProxyType = None + _IS_SYNC = False _MAX_TCP_KEEPIDLE = 120 @@ -838,7 +844,21 @@ 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..3d931bd522 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/auth.py b/pymongo/synchronous/auth.py index 7b370843c5..0e51ff8b7f 100644 --- a/pymongo/synchronous/auth.py +++ b/pymongo/synchronous/auth.py @@ -174,13 +174,20 @@ def _auth_key(nonce: str, username: str, password: str) -> str: return md5hash.hexdigest() -def _canonicalize_hostname(hostname: str) -> str: +def _canonicalize_hostname(hostname: str, option: str | bool) -> str: """Canonicalize hostname following MIT-krb5 behavior.""" # https://github.com/krb5/krb5/blob/d406afa363554097ac48646a29249c04f498c88e/src/util/k5test.py#L505-L520 + if option in [False, "none"]: + return hostname + af, socktype, proto, canonname, sockaddr = socket.getaddrinfo( hostname, None, 0, 0, socket.IPPROTO_TCP, socket.AI_CANONNAME )[0] + # For forward just to resolve the cname as dns.lookup() will not return it. + if option == "forward": + return canonname.lower() + try: name = socket.getnameinfo(sockaddr, socket.NI_NAMEREQD) except socket.gaierror: @@ -202,9 +209,8 @@ def _authenticate_gssapi(credentials: MongoCredential, conn: Connection) -> None props = credentials.mechanism_properties # Starting here and continuing through the while loop below - establish # the security context. See RFC 4752, Section 3.1, first paragraph. - host = conn.address[0] - if props.canonicalize_host_name: - host = _canonicalize_hostname(host) + host = props.service_host or conn.address[0] + host = _canonicalize_hostname(host, props.canonicalize_host_name) service = props.service_name + "@" + host if props.service_realm is not None: service = service + "@" + props.service_realm diff --git a/pymongo/synchronous/pool.py b/pymongo/synchronous/pool.py index 1a155c82d7..7418a1aea2 100644 --- a/pymongo/synchronous/pool.py +++ b/pymongo/synchronous/pool.py @@ -121,6 +121,12 @@ def _set_non_inheritable_non_atomic(fd: int) -> None: # noqa: ARG001 """Dummy function for platforms that don't provide fcntl.""" +try: + from python_socks import ProxyType + from python_socks.sync import Proxy +except ImportError: + Proxy = ProxyType = None + _IS_SYNC = True _MAX_TCP_KEEPIDLE = 120 @@ -836,7 +842,21 @@ 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] diff --git a/test/asynchronous/test_auth.py b/test/asynchronous/test_auth.py index 4f26200fb0..08dc4d7247 100644 --- a/test/asynchronous/test_auth.py +++ b/test/asynchronous/test_auth.py @@ -35,7 +35,7 @@ import pytest from pymongo import AsyncMongoClient, monitoring -from pymongo.asynchronous.auth import HAVE_KERBEROS +from pymongo.asynchronous.auth import HAVE_KERBEROS, _canonicalize_hostname from pymongo.auth_shared import _build_credentials_tuple from pymongo.errors import OperationFailure from pymongo.hello import HelloCompat @@ -96,10 +96,11 @@ def setUpClass(cls): cls.service_realm_required = ( GSSAPI_SERVICE_REALM is not None and GSSAPI_SERVICE_REALM not in GSSAPI_PRINCIPAL ) - mech_properties = f"SERVICE_NAME:{GSSAPI_SERVICE_NAME}" - mech_properties += f",CANONICALIZE_HOST_NAME:{GSSAPI_CANONICALIZE}" + mech_properties = dict( + SERVICE_NAME=GSSAPI_SERVICE_NAME, CANONICALIZE_HOST_NAME=GSSAPI_CANONICALIZE + ) if GSSAPI_SERVICE_REALM is not None: - mech_properties += f",SERVICE_REALM:{GSSAPI_SERVICE_REALM}" + mech_properties["SERVICE_REALM"] = GSSAPI_SERVICE_REALM cls.mech_properties = mech_properties async def test_credentials_hashing(self): @@ -167,7 +168,10 @@ async def test_gssapi_simple(self): await client[GSSAPI_DB].collection.find_one() # Log in using URI, with authMechanismProperties. - mech_uri = uri + f"&authMechanismProperties={self.mech_properties}" + mech_properties_str = "" + for key, value in self.mech_properties.items(): + mech_properties_str += f"{key}:{value}," + mech_uri = uri + f"&authMechanismProperties={mech_properties_str[:-1]}" client = self.simple_client(mech_uri) await client[GSSAPI_DB].collection.find_one() @@ -268,6 +272,58 @@ async def test_gssapi_threaded(self): thread.join() self.assertTrue(thread.success) + async def test_gssapi_canonicalize_host_name(self): + # Test the low level method. + assert GSSAPI_HOST is not None + result = _canonicalize_hostname(GSSAPI_HOST, "forward") + if "compute-1.amazonaws.com" not in result: + self.assertEqual(result, GSSAPI_HOST) + result = _canonicalize_hostname(GSSAPI_HOST, "forwardAndReverse") + self.assertEqual(result, GSSAPI_HOST) + + # Use the equivalent named CANONICALIZE_HOST_NAME. + props = self.mech_properties.copy() + if props["CANONICALIZE_HOST_NAME"] == "true": + props["CANONICALIZE_HOST_NAME"] = "forwardAndReverse" + else: + props["CANONICALIZE_HOST_NAME"] = "none" + client = self.simple_client( + GSSAPI_HOST, + GSSAPI_PORT, + username=GSSAPI_PRINCIPAL, + password=GSSAPI_PASS, + authMechanism="GSSAPI", + authMechanismProperties=props, + ) + await client.server_info() + + async def test_gssapi_host_name(self): + props = self.mech_properties + props["SERVICE_HOST"] = "example.com" + + # Authenticate with authMechanismProperties. + client = self.simple_client( + GSSAPI_HOST, + GSSAPI_PORT, + username=GSSAPI_PRINCIPAL, + password=GSSAPI_PASS, + authMechanism="GSSAPI", + authMechanismProperties=self.mech_properties, + ) + with self.assertRaises(OperationFailure): + await client.server_info() + + props["SERVICE_HOST"] = GSSAPI_HOST + client = self.simple_client( + GSSAPI_HOST, + GSSAPI_PORT, + username=GSSAPI_PRINCIPAL, + password=GSSAPI_PASS, + authMechanism="GSSAPI", + authMechanismProperties=self.mech_properties, + ) + await client.server_info() + class TestSASLPlain(AsyncPyMongoTestCase): @classmethod diff --git a/test/test_auth.py b/test/test_auth.py index 70c061b747..345d16121b 100644 --- a/test/test_auth.py +++ b/test/test_auth.py @@ -40,7 +40,7 @@ from pymongo.hello import HelloCompat from pymongo.read_preferences import ReadPreference from pymongo.saslprep import HAVE_STRINGPREP -from pymongo.synchronous.auth import HAVE_KERBEROS +from pymongo.synchronous.auth import HAVE_KERBEROS, _canonicalize_hostname _IS_SYNC = True @@ -96,10 +96,11 @@ def setUpClass(cls): cls.service_realm_required = ( GSSAPI_SERVICE_REALM is not None and GSSAPI_SERVICE_REALM not in GSSAPI_PRINCIPAL ) - mech_properties = f"SERVICE_NAME:{GSSAPI_SERVICE_NAME}" - mech_properties += f",CANONICALIZE_HOST_NAME:{GSSAPI_CANONICALIZE}" + mech_properties = dict( + SERVICE_NAME=GSSAPI_SERVICE_NAME, CANONICALIZE_HOST_NAME=GSSAPI_CANONICALIZE + ) if GSSAPI_SERVICE_REALM is not None: - mech_properties += f",SERVICE_REALM:{GSSAPI_SERVICE_REALM}" + mech_properties["SERVICE_REALM"] = GSSAPI_SERVICE_REALM cls.mech_properties = mech_properties def test_credentials_hashing(self): @@ -167,7 +168,10 @@ def test_gssapi_simple(self): client[GSSAPI_DB].collection.find_one() # Log in using URI, with authMechanismProperties. - mech_uri = uri + f"&authMechanismProperties={self.mech_properties}" + mech_properties_str = "" + for key, value in self.mech_properties.items(): + mech_properties_str += f"{key}:{value}," + mech_uri = uri + f"&authMechanismProperties={mech_properties_str[:-1]}" client = self.simple_client(mech_uri) client[GSSAPI_DB].collection.find_one() @@ -268,6 +272,58 @@ def test_gssapi_threaded(self): thread.join() self.assertTrue(thread.success) + def test_gssapi_canonicalize_host_name(self): + # Test the low level method. + assert GSSAPI_HOST is not None + result = _canonicalize_hostname(GSSAPI_HOST, "forward") + if "compute-1.amazonaws.com" not in result: + self.assertEqual(result, GSSAPI_HOST) + result = _canonicalize_hostname(GSSAPI_HOST, "forwardAndReverse") + self.assertEqual(result, GSSAPI_HOST) + + # Use the equivalent named CANONICALIZE_HOST_NAME. + props = self.mech_properties.copy() + if props["CANONICALIZE_HOST_NAME"] == "true": + props["CANONICALIZE_HOST_NAME"] = "forwardAndReverse" + else: + props["CANONICALIZE_HOST_NAME"] = "none" + client = self.simple_client( + GSSAPI_HOST, + GSSAPI_PORT, + username=GSSAPI_PRINCIPAL, + password=GSSAPI_PASS, + authMechanism="GSSAPI", + authMechanismProperties=props, + ) + client.server_info() + + def test_gssapi_host_name(self): + props = self.mech_properties + props["SERVICE_HOST"] = "example.com" + + # Authenticate with authMechanismProperties. + client = self.simple_client( + GSSAPI_HOST, + GSSAPI_PORT, + username=GSSAPI_PRINCIPAL, + password=GSSAPI_PASS, + authMechanism="GSSAPI", + authMechanismProperties=self.mech_properties, + ) + with self.assertRaises(OperationFailure): + client.server_info() + + props["SERVICE_HOST"] = GSSAPI_HOST + client = self.simple_client( + GSSAPI_HOST, + GSSAPI_PORT, + username=GSSAPI_PRINCIPAL, + password=GSSAPI_PASS, + authMechanism="GSSAPI", + authMechanismProperties=self.mech_properties, + ) + client.server_info() + class TestSASLPlain(PyMongoTestCase): @classmethod From 480318597be7eda1bd70437c64669841fbf1c97b Mon Sep 17 00:00:00 2001 From: Steven Silvester Date: Tue, 31 Dec 2024 16:29:48 -0600 Subject: [PATCH 2/4] PYTHON-3096 Finish implementation and tests for GSSAPI options (#1985) --- doc/contributors.rst | 1 + pymongo/asynchronous/pool.py | 22 +++++++++++++++++++++- pymongo/asynchronous/topology.py | 1 + pymongo/client_options.py | 10 ++++++++++ pymongo/common.py | 4 ++++ pymongo/pool_options.py | 9 ++++++++- pymongo/synchronous/pool.py | 22 +++++++++++++++++++++- pymongo/synchronous/topology.py | 1 + pyproject.toml | 1 + requirements/socks.txt | 1 + 10 files changed, 69 insertions(+), 3 deletions(-) create mode 100644 requirements/socks.txt diff --git a/doc/contributors.rst b/doc/contributors.rst index 4a7f5424b1..fbdad001db 100644 --- a/doc/contributors.rst +++ b/doc/contributors.rst @@ -103,3 +103,4 @@ The following is a list of people who have contributed to - Terry Patterson - Romain Morotti - Navjot Singh (navjots18) +- Yuval Zaif (zaif-yuval) diff --git a/pymongo/asynchronous/pool.py b/pymongo/asynchronous/pool.py index 5dc5675a0a..0bac28ed42 100644 --- a/pymongo/asynchronous/pool.py +++ b/pymongo/asynchronous/pool.py @@ -121,6 +121,12 @@ def _set_non_inheritable_non_atomic(fd: int) -> None: # noqa: ARG001 """Dummy function for platforms that don't provide fcntl.""" +try: + from python_socks import ProxyType + from python_socks.sync import Proxy +except ImportError: + Proxy = ProxyType = None + _IS_SYNC = False _MAX_TCP_KEEPIDLE = 120 @@ -838,7 +844,21 @@ 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..3d931bd522 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..7418a1aea2 100644 --- a/pymongo/synchronous/pool.py +++ b/pymongo/synchronous/pool.py @@ -121,6 +121,12 @@ def _set_non_inheritable_non_atomic(fd: int) -> None: # noqa: ARG001 """Dummy function for platforms that don't provide fcntl.""" +try: + from python_socks import ProxyType + from python_socks.sync import Proxy +except ImportError: + Proxy = ProxyType = None + _IS_SYNC = True _MAX_TCP_KEEPIDLE = 120 @@ -836,7 +842,21 @@ 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] From 7392aaaefb5ecab2d5e04f91bc18d493e389574c Mon Sep 17 00:00:00 2001 From: Yuval Zaif Date: Mon, 6 Jan 2025 22:39:53 +0200 Subject: [PATCH 3/4] Fixed PR notes --- pymongo/asynchronous/pool.py | 2 +- pymongo/client_options.py | 6 +++--- pymongo/synchronous/pool.py | 2 +- pyproject.toml | 2 +- requirements/{socks.txt => proxy.txt} | 0 5 files changed, 6 insertions(+), 6 deletions(-) rename requirements/{socks.txt => proxy.txt} (100%) diff --git a/pymongo/asynchronous/pool.py b/pymongo/asynchronous/pool.py index 0bac28ed42..02ebd693d7 100644 --- a/pymongo/asynchronous/pool.py +++ b/pymongo/asynchronous/pool.py @@ -848,7 +848,7 @@ def _create_connection(address: _Address, options: PoolOptions) -> socket.socket 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]`" + "This can be done by re-installing pymongo with `pip install pymongo[proxy]`" ) proxy_host = proxy["host"] proxy_port = proxy["port"] or 1080 diff --git a/pymongo/client_options.py b/pymongo/client_options.py index 3d931bd522..5df24c6e28 100644 --- a/pymongo/client_options.py +++ b/pymongo/client_options.py @@ -171,14 +171,14 @@ def _parse_pool_options( load_balanced = options.get("loadbalanced") max_connecting = options.get("maxconnecting", common.MAX_CONNECTING) if proxy_host := options.get("proxyHost"): - proxy = { + proxy_options = { "host": proxy_host, "port": options.get("proxyPort"), "username": options.get("proxyUserName"), "password": options.get("proxyPassword"), } else: - proxy = None + proxy_options = None return PoolOptions( max_pool_size, min_pool_size, @@ -197,7 +197,7 @@ def _parse_pool_options( load_balanced=load_balanced, credentials=credentials, is_sync=is_sync, - proxy=proxy, + proxy=proxy_options, ) diff --git a/pymongo/synchronous/pool.py b/pymongo/synchronous/pool.py index 7418a1aea2..6fecbc3cbf 100644 --- a/pymongo/synchronous/pool.py +++ b/pymongo/synchronous/pool.py @@ -846,7 +846,7 @@ def _create_connection(address: _Address, options: PoolOptions) -> socket.socket 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]`" + "This can be done by re-installing pymongo with `pip install pymongo[proxy]`" ) proxy_host = proxy["host"] proxy_port = proxy["port"] or 1080 diff --git a/pyproject.toml b/pyproject.toml index 54ed0676e4..7472f81363 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -67,7 +67,7 @@ ocsp = ["requirements/ocsp.txt"] snappy = ["requirements/snappy.txt"] test = ["requirements/test.txt"] zstd = ["requirements/zstd.txt"] -socks = ["requirements/socks.txt"] +proxy = ["requirements/proxy.txt"] [tool.pytest.ini_options] minversion = "7" diff --git a/requirements/socks.txt b/requirements/proxy.txt similarity index 100% rename from requirements/socks.txt rename to requirements/proxy.txt From 0f81749ca9eff1268c5e21cc2eb94375f538fb58 Mon Sep 17 00:00:00 2001 From: Yuval Zaif Date: Mon, 6 Jan 2025 22:44:16 +0200 Subject: [PATCH 4/4] Added proxy to readme --- README.md | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) diff --git a/README.md b/README.md index bd0755620e..fa27ad64d9 100644 --- a/README.md +++ b/README.md @@ -145,11 +145,18 @@ Client-Side Field Level Encryption requires ```bash python -m pip install "pymongo[encryption]" ``` + +Proxy support requires +[python-socks] +```bash +python -m pip install "pymongo[proxy]" +``` + You can install all dependencies automatically with the following command: ```bash -python -m pip install "pymongo[gssapi,aws,ocsp,snappy,zstd,encryption]" +python -m pip install "pymongo[gssapi,aws,ocsp,snappy,zstd,encryption,proxy]" ``` Additional dependencies are: