PYTHON-6016 Consolidate auth helper methods into auth_shared.py

pymongo/asynchronous/auth.py

@@ -21,35 +21,36 @@
 import hmac
 import socket
 from base64 import standard_b64decode, standard_b64encode
-from collections.abc import Coroutine, Mapping, MutableMapping
+from collections.abc import Coroutine, Mapping
 from typing import (
     TYPE_CHECKING,
     Any,
     Callable,
-    Optional,
-    cast,
 )
 from urllib.parse import quote
 
 from bson.binary import Binary
 from pymongo.asynchronous.auth_aws import _authenticate_aws
 from pymongo.asynchronous.auth_oidc import (
     _authenticate_oidc,
-    _get_authenticator,
+    _OIDCAuthenticator,
 )
 from pymongo.asynchronous.helpers import _getaddrinfo
 from pymongo.auth_shared import (
     MongoCredential,
     _authenticate_scram_start,
+    _OIDCContext,
     _parse_scram_response,
+    _password_digest,
+    _ScramContext,
+    _X509Context,
     _xor,
 )
 from pymongo.errors import ConfigurationError, OperationFailure
 from pymongo.saslprep import saslprep
 
 if TYPE_CHECKING:
     from pymongo.asynchronous.pool import AsyncConnection
-    from pymongo.hello import Hello
 
 HAVE_KERBEROS = True
 _USE_PRINCIPAL = False
@@ -153,21 +154,6 @@ async def _authenticate_scram(
             raise OperationFailure("SASL conversation failed to complete.")
 
 
-def _password_digest(username: str, password: str) -> str:
-    """Get a password digest to use for authentication."""
-    if not isinstance(password, str):
-        raise TypeError("password must be an instance of str")
-    if len(password) == 0:
-        raise ValueError("password can't be empty")
-    if not isinstance(username, str):
-        raise TypeError(f"username must be an instance of str, not {type(username)}")
-
-    md5hash = hashlib.md5()  # noqa: S324
-    data = f"{username}:mongo:{password}"
-    md5hash.update(data.encode("utf-8"))
-    return md5hash.hexdigest()
-
-
 def _auth_key(nonce: str, username: str, password: str) -> str:
     """Get an auth key to use for authentication."""
     digest = _password_digest(username, password)
@@ -370,71 +356,11 @@ async def _authenticate_default(credentials: MongoCredential, conn: AsyncConnect
 }
 
 
-class _AuthContext:
-    def __init__(self, credentials: MongoCredential, address: tuple[str, int]) -> None:
-        self.credentials = credentials
-        self.speculative_authenticate: Optional[Mapping[str, Any]] = None
-        self.address = address
-
-    @staticmethod
-    def from_credentials(
-        creds: MongoCredential, address: tuple[str, int]
-    ) -> Optional[_AuthContext]:
-        spec_cls = _SPECULATIVE_AUTH_MAP.get(creds.mechanism)
-        if spec_cls:
-            return cast(_AuthContext, spec_cls(creds, address))
-        return None
-
-    def speculate_command(self) -> Optional[MutableMapping[str, Any]]:
-        raise NotImplementedError
-
-    def parse_response(self, hello: Hello[Mapping[str, Any]]) -> None:
-        self.speculative_authenticate = hello.speculative_authenticate
-
-    def speculate_succeeded(self) -> bool:
-        return bool(self.speculative_authenticate)
-
-
-class _ScramContext(_AuthContext):
-    def __init__(
-        self, credentials: MongoCredential, address: tuple[str, int], mechanism: str
-    ) -> None:
-        super().__init__(credentials, address)
-        self.scram_data: Optional[tuple[bytes, bytes]] = None
-        self.mechanism = mechanism
-
-    def speculate_command(self) -> Optional[MutableMapping[str, Any]]:
-        nonce, first_bare, cmd = _authenticate_scram_start(self.credentials, self.mechanism)
-        # The 'db' field is included only on the speculative command.
-        cmd["db"] = self.credentials.source
-        # Save for later use.
-        self.scram_data = (nonce, first_bare)
-        return cmd
-
-
-class _X509Context(_AuthContext):
-    def speculate_command(self) -> MutableMapping[str, Any]:
-        cmd = {"authenticate": 1, "mechanism": "MONGODB-X509"}
-        if self.credentials.username is not None:
-            cmd["user"] = self.credentials.username
-        return cmd
-
-
-class _OIDCContext(_AuthContext):
-    def speculate_command(self) -> Optional[MutableMapping[str, Any]]:
-        authenticator = _get_authenticator(self.credentials, self.address)
-        cmd = authenticator.get_spec_auth_cmd()
-        if cmd is None:
-            return None
-        cmd["db"] = self.credentials.source
-        return cmd
-
-
 _SPECULATIVE_AUTH_MAP: Mapping[str, Any] = {
     "MONGODB-X509": _X509Context,
     "SCRAM-SHA-1": functools.partial(_ScramContext, mechanism="SCRAM-SHA-1"),
     "SCRAM-SHA-256": functools.partial(_ScramContext, mechanism="SCRAM-SHA-256"),
-    "MONGODB-OIDC": _OIDCContext,
+    "MONGODB-OIDC": functools.partial(_OIDCContext, authenticator_cls=_OIDCAuthenticator),
     "DEFAULT": functools.partial(_ScramContext, mechanism="SCRAM-SHA-256"),
 }
 

pymongo/asynchronous/auth_oidc.py

@@ -35,9 +35,10 @@
     OIDCCallbackContext,
     OIDCCallbackResult,
     OIDCIdPInfo,
+    _get_authenticator,
     _OIDCProperties,
 )
-from pymongo.errors import ConfigurationError, OperationFailure
+from pymongo.errors import OperationFailure
 from pymongo.helpers_shared import _AUTHENTICATION_FAILURE_CODE
 from pymongo.lock import Lock, _async_create_lock
 
@@ -48,35 +49,6 @@
 _IS_SYNC = False
 
 
-def _get_authenticator(
-    credentials: MongoCredential, address: tuple[str, int]
-) -> _OIDCAuthenticator:
-    # Extract values.
-    principal_name = credentials.username
-    properties = credentials.mechanism_properties
-
-    # Validate that the address is allowed.
-    if properties.human_callback is not None:
-        found = False
-        allowed_hosts = properties.allowed_hosts
-        for patt in allowed_hosts:
-            if patt == address[0]:
-                found = True
-            elif patt.startswith("*.") and address[0].endswith(patt[1:]):
-                found = True
-        if not found:
-            raise ConfigurationError(
-                f"Refusing to connect to {address[0]}, which is not in authOIDCAllowedHosts: {allowed_hosts}"
-            )
-
-    if credentials.cache.data:
-        return credentials.cache.data
-
-    # Get or create the cache data.
-    credentials.cache.data = _OIDCAuthenticator(username=principal_name, properties=properties)
-    return credentials.cache.data
-
-
 @dataclass
 class _OIDCAuthenticator:
     username: str
@@ -300,7 +272,7 @@ async def _authenticate_oidc(
     credentials: MongoCredential, conn: AsyncConnection, reauthenticate: bool
 ) -> Optional[Mapping[str, Any]]:
     """Authenticate using MONGODB-OIDC."""
-    authenticator = _get_authenticator(credentials, conn.address)
+    authenticator = _get_authenticator(credentials, conn.address, _OIDCAuthenticator)
     if reauthenticate:
         return await authenticator.reauthenticate(conn)
     else:

pymongo/asynchronous/pool.py

@@ -34,6 +34,7 @@
 from pymongo._telemetry import _CmapTelemetry
 from pymongo.asynchronous.command_runner import run_command
 from pymongo.asynchronous.helpers import _handle_reauth
+from pymongo.auth_shared import _AuthContext
 from pymongo.client_session_shared import _validate_session_write_concern
 from pymongo.common import (
     MAX_BSON_SIZE,
@@ -89,7 +90,6 @@
 
     from bson import CodecOptions
     from bson.objectid import ObjectId
-    from pymongo.asynchronous.auth import _AuthContext
     from pymongo.asynchronous.client_session import AsyncClientSession
     from pymongo.asynchronous.mongo_client import AsyncMongoClient, _ClientCheckout
     from pymongo.compression_support import (
@@ -267,7 +267,9 @@ async def _hello(
                 cmd["saslSupportedMechs"] = creds.source + "." + creds.username
             from pymongo.asynchronous import auth
 
-            auth_ctx = auth._AuthContext.from_credentials(creds, self.address)
+            auth_ctx = _AuthContext.from_credentials(
+                creds, self.address, auth._SPECULATIVE_AUTH_MAP
+            )
             if auth_ctx:
                 speculative_authenticate = auth_ctx.speculate_command()
                 if speculative_authenticate is not None:

pymongo/auth_oidc_shared.py

@@ -20,11 +20,15 @@
 import abc
 import os
 from dataclasses import dataclass, field
-from typing import Optional
+from typing import TYPE_CHECKING, Any, Callable, Optional
 from urllib.parse import quote
 
 from pymongo._azure_helpers import _get_azure_response
 from pymongo._gcp_helpers import _get_gcp_response
+from pymongo.errors import ConfigurationError
+
+if TYPE_CHECKING:
+    from pymongo.auth_shared import MongoCredential
 
 
 @dataclass
@@ -131,3 +135,32 @@ def _get_k8s_token() -> str:
             fname = os.environ[key]
     with open(fname) as fid:
         return fid.read()
+
+
+def _get_authenticator(
+    credentials: MongoCredential, address: tuple[str, int], authenticator_cls: Callable[..., Any]
+) -> Any:
+    # Extract values.
+    principal_name = credentials.username
+    properties = credentials.mechanism_properties
+
+    # Validate that the address is allowed.
+    if properties.human_callback is not None:
+        found = False
+        allowed_hosts = properties.allowed_hosts
+        for patt in allowed_hosts:
+            if patt == address[0]:
+                found = True
+            elif patt.startswith("*.") and address[0].endswith(patt[1:]):
+                found = True
+        if not found:
+            raise ConfigurationError(
+                f"Refusing to connect to {address[0]}, which is not in authOIDCAllowedHosts: {allowed_hosts}"
+            )
+
+    if credentials.cache.data:
+        return credentials.cache.data
+
+    # Get or create the cache data.
+    credentials.cache.data = authenticator_cls(username=principal_name, properties=properties)
+    return credentials.cache.data

pymongo/auth_shared.py

@@ -17,15 +17,17 @@
 
 from __future__ import annotations
 
+import hashlib
 import os
 import typing
 from base64 import standard_b64encode
 from collections import namedtuple
-from collections.abc import Mapping
-from typing import Any, Optional
+from collections.abc import Mapping, MutableMapping
+from typing import TYPE_CHECKING, Any, Callable, Optional, cast
 
 from bson import Binary
 from pymongo.auth_oidc_shared import (
+    _get_authenticator,
     _OIDCAzureCallback,
     _OIDCGCPCallback,
     _OIDCK8SCallback,
@@ -34,6 +36,9 @@
 )
 from pymongo.errors import ConfigurationError
 
+if TYPE_CHECKING:
+    from pymongo.hello import Hello
+
 MECHANISMS = frozenset(
     [
         "GSSAPI",
@@ -255,3 +260,87 @@ def _authenticate_scram_start(
         "options": {"skipEmptyExchange": True},
     }
     return nonce, first_bare, cmd
+
+
+def _password_digest(username: str, password: str) -> str:
+    """Get a password digest to use for authentication."""
+    if not isinstance(password, str):
+        raise TypeError("password must be an instance of str")
+    if len(password) == 0:
+        raise ValueError("password can't be empty")
+    if not isinstance(username, str):
+        raise TypeError(f"username must be an instance of str, not {type(username)}")
+
+    md5hash = hashlib.md5()  # noqa: S324
+    data = f"{username}:mongo:{password}"
+    md5hash.update(data.encode("utf-8"))
+    return md5hash.hexdigest()
+
+
+class _AuthContext:
+    def __init__(self, credentials: MongoCredential, address: tuple[str, int]) -> None:
+        self.credentials = credentials
+        self.speculative_authenticate: Optional[Mapping[str, Any]] = None
+        self.address = address
+
+    @staticmethod
+    def from_credentials(
+        creds: MongoCredential, address: tuple[str, int], spec_auth_map: Mapping[str, Any]
+    ) -> Optional[_AuthContext]:
+        spec_cls = spec_auth_map.get(creds.mechanism)
+        if spec_cls:
+            return cast(_AuthContext, spec_cls(creds, address))
+        return None
+
+    def speculate_command(self) -> Optional[MutableMapping[str, Any]]:
+        raise NotImplementedError
+
+    def parse_response(self, hello: Hello[Mapping[str, Any]]) -> None:
+        self.speculative_authenticate = hello.speculative_authenticate
+
+    def speculate_succeeded(self) -> bool:
+        return bool(self.speculative_authenticate)
+
+
+class _ScramContext(_AuthContext):
+    def __init__(
+        self, credentials: MongoCredential, address: tuple[str, int], mechanism: str
+    ) -> None:
+        super().__init__(credentials, address)
+        self.scram_data: Optional[tuple[bytes, bytes]] = None
+        self.mechanism = mechanism
+
+    def speculate_command(self) -> Optional[MutableMapping[str, Any]]:
+        nonce, first_bare, cmd = _authenticate_scram_start(self.credentials, self.mechanism)
+        # The 'db' field is included only on the speculative command.
+        cmd["db"] = self.credentials.source
+        # Save for later use.
+        self.scram_data = (nonce, first_bare)
+        return cmd
+
+
+class _X509Context(_AuthContext):
+    def speculate_command(self) -> MutableMapping[str, Any]:
+        cmd = {"authenticate": 1, "mechanism": "MONGODB-X509"}
+        if self.credentials.username is not None:
+            cmd["user"] = self.credentials.username
+        return cmd
+
+
+class _OIDCContext(_AuthContext):
+    def __init__(
+        self,
+        credentials: MongoCredential,
+        address: tuple[str, int],
+        authenticator_cls: Callable[..., Any],
+    ) -> None:
+        super().__init__(credentials, address)
+        self.authenticator_cls = authenticator_cls
+
+    def speculate_command(self) -> Optional[MutableMapping[str, Any]]:
+        authenticator = _get_authenticator(self.credentials, self.address, self.authenticator_cls)
+        cmd = authenticator.get_spec_auth_cmd()
+        if cmd is None:
+            return None
+        cmd["db"] = self.credentials.source
+        return cmd