110 lines
3.5 KiB
Python
110 lines
3.5 KiB
Python
# Copyright 2020-present MongoDB, Inc.
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
|
|
"""MONGODB-AWS Authentication helpers."""
|
|
|
|
try:
|
|
import pymongo_auth_aws
|
|
from pymongo_auth_aws import AwsCredential, AwsSaslContext, PyMongoAuthAwsError
|
|
|
|
_HAVE_MONGODB_AWS = True
|
|
except ImportError:
|
|
|
|
class AwsSaslContext: # type: ignore
|
|
def __init__(self, credentials):
|
|
pass
|
|
|
|
_HAVE_MONGODB_AWS = False
|
|
|
|
try:
|
|
from pymongo_auth_aws.auth import set_cached_credentials, set_use_cached_credentials
|
|
|
|
# Enable credential caching.
|
|
set_use_cached_credentials(True)
|
|
except ImportError:
|
|
|
|
def set_cached_credentials(creds):
|
|
pass
|
|
|
|
|
|
import bson
|
|
from bson.binary import Binary
|
|
from bson.son import SON
|
|
from pymongo.errors import ConfigurationError, OperationFailure
|
|
|
|
|
|
class _AwsSaslContext(AwsSaslContext): # type: ignore
|
|
# Dependency injection:
|
|
def binary_type(self):
|
|
"""Return the bson.binary.Binary type."""
|
|
return Binary
|
|
|
|
def bson_encode(self, doc):
|
|
"""Encode a dictionary to BSON."""
|
|
return bson.encode(doc)
|
|
|
|
def bson_decode(self, data):
|
|
"""Decode BSON to a dictionary."""
|
|
return bson.decode(data)
|
|
|
|
|
|
def _authenticate_aws(credentials, sock_info):
|
|
"""Authenticate using MONGODB-AWS."""
|
|
if not _HAVE_MONGODB_AWS:
|
|
raise ConfigurationError(
|
|
"MONGODB-AWS authentication requires pymongo-auth-aws: "
|
|
"install with: python -m pip install 'pymongo[aws]'"
|
|
)
|
|
|
|
if sock_info.max_wire_version < 9:
|
|
raise ConfigurationError("MONGODB-AWS authentication requires MongoDB version 4.4 or later")
|
|
|
|
try:
|
|
ctx = _AwsSaslContext(
|
|
AwsCredential(
|
|
credentials.username,
|
|
credentials.password,
|
|
credentials.mechanism_properties.aws_session_token,
|
|
)
|
|
)
|
|
client_payload = ctx.step(None)
|
|
client_first = SON(
|
|
[("saslStart", 1), ("mechanism", "MONGODB-AWS"), ("payload", client_payload)]
|
|
)
|
|
server_first = sock_info.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 = SON(
|
|
[
|
|
("saslContinue", 1),
|
|
("conversationId", server_first["conversationId"]),
|
|
("payload", client_payload),
|
|
]
|
|
)
|
|
res = sock_info.command("$external", cmd)
|
|
if res["done"]:
|
|
# SASL complete.
|
|
break
|
|
except PyMongoAuthAwsError as exc:
|
|
# Clear the cached credentials if we hit a failure in auth.
|
|
set_cached_credentials(None)
|
|
# Convert to OperationFailure and include pymongo-auth-aws version.
|
|
raise OperationFailure(f"{exc} (pymongo-auth-aws version {pymongo_auth_aws.__version__})")
|
|
except Exception:
|
|
# Clear the cached credentials if we hit a failure in auth.
|
|
set_cached_credentials(None)
|
|
raise
|