Function _authenticate_aws

pymongo/asynchronous/auth_aws.py:32–100

Authenticate using MONGODB-AWS.

(credentials: MongoCredential, conn: AsyncConnection)

async def _authenticate_aws(credentials: MongoCredential, conn: AsyncConnection) -> None:
    """Authenticate using MONGODB-AWS."""
    try:
        import pymongo_auth_aws  # type:ignore[import]
    except ImportError as e:
        raise ConfigurationError(
            "MONGODB-AWS authentication requires pymongo-auth-aws: "
            "install with: python -m pip install 'pymongo[aws]'"
        ) from e
    # Delayed import.
    from pymongo_auth_aws.auth import (  # type:ignore[import]
        set_cached_credentials,
        set_use_cached_credentials,
    )

    set_use_cached_credentials(True)

    if conn.max_wire_version < 9:
        raise ConfigurationError("MONGODB-AWS authentication requires MongoDB version 4.4 or later")

    class AwsSaslContext(pymongo_auth_aws.AwsSaslContext):  # type: ignore
        # Dependency injection:
        def binary_type(self) -> Type[Binary]:
            """Return the bson.binary.Binary type."""
            return Binary

        def bson_encode(self, doc: Mapping[str, Any]) -> bytes:
            """Encode a dictionary to BSON."""
            return bson.encode(doc)

        def bson_decode(self, data: _ReadableBuffer) -> Mapping[str, Any]:
            """Decode BSON to a dictionary."""
            return bson.decode(data)

    try:
        ctx = AwsSaslContext(
            pymongo_auth_aws.AwsCredential(
                credentials.username,
                credentials.password,
                credentials.mechanism_properties.aws_session_token,
            )
        )
        client_payload = ctx.step(None)
        client_first = {"saslStart": 1, "mechanism": "MONGODB-AWS", "payload": client_payload}
        server_first = await conn.command("$external", client_first)
        res = server_first
        # Limit how many times we loop to catch protocol / library issues
        for _ in range(10):
            client_payload = ctx.step(res["payload"])
            cmd = {
                "saslContinue": 1,
                "conversationId": server_first["conversationId"],
                "payload": client_payload,
            }
            res = await conn.command("$external", cmd)
            if res["done"]:
                # SASL complete.
                break