PYTHON-2138 Use pymongo-auth-aws for MONGODB-AWS support

Author: ShaneHarvey

.evergreen/run-mongodb-aws-ecs-test.sh

@@ -35,9 +35,9 @@ authtest () {
 
     $VIRTUALENV -p $PYTHON --system-site-packages --never-download venvaws
     . venvaws/bin/activate
-    pip install requests botocore
 
     cd src
+    pip install '.[aws]'
     python test/auth_aws/test_auth_aws.py
     cd -
     deactivate

.evergreen/run-mongodb-aws-test.sh

@@ -55,8 +55,7 @@ authtest () {
     else
       . venvaws/bin/activate
     fi
-    pip install requests botocore
-
+    pip install '.[aws]'
     python test/auth_aws/test_auth_aws.py
     deactivate
     rm -rf venvaws

README.rst

@@ -100,9 +100,8 @@ dependency can be installed automatically along with PyMongo::
 
   $ python -m pip install pymongo[gssapi]
 
-MONGODB-AWS authentication requires `botocore
-<https://pypi.org/project/botocore/>`_ and `requests
-<https://pypi.org/project/requests/>`_::
+MONGODB-AWS authentication requires `pymongo-auth-aws
+<https://pypi.org/project/pymongo-auth-aws/>`_::
 
   $ python -m pip install pymongo[aws]
 

doc/installation.rst

@@ -56,9 +56,8 @@ dependency can be installed automatically along with PyMongo::
 
   $ python -m pip install pymongo[gssapi]
 
-:ref:`MONGODB-AWS` authentication requires `botocore
-<https://pypi.org/project/botocore/>`_ and `requests
-<https://pypi.org/project/requests/>`_::
+:ref:`MONGODB-AWS` authentication requires `pymongo-auth-aws
+<https://pypi.org/project/pymongo-auth-aws/>`_::
 
   $ python -m pip install pymongo[aws]
 

pymongo/auth.py

@@ -43,7 +43,7 @@
 from bson.binary import Binary
 from bson.py3compat import string_type, _unicode, PY3
 from bson.son import SON
-from pymongo.auth_aws import _HAVE_MONGODB_AWS, _auth_aws, _AWSCredential
+from pymongo.auth_aws import _authenticate_aws
 from pymongo.errors import ConfigurationError, OperationFailure
 from pymongo.saslprep import saslprep
 
@@ -540,20 +540,6 @@ def _authenticate_x509(credentials, sock_info):
     sock_info.command('$external', cmd)
 
 
-def _authenticate_aws(credentials, sock_info):
-    """Authenticate using MONGODB-AWS.
-    """
-    if not _HAVE_MONGODB_AWS:
-        raise ConfigurationError(
-            "MONGODB-AWS authentication requires botocore and requests: "
-            "install these libraries with: "
-            "python -m pip install 'pymongo[aws]'")
-
-    _auth_aws(_AWSCredential(
-        credentials.username, credentials.password,
-        credentials.mechanism_properties.aws_session_token), sock_info)
-
-
 def _authenticate_mongo_cr(credentials, sock_info):
     """Authenticate using MONGODB-CR.
     """

pymongo/auth_aws.py

@@ -14,188 +14,69 @@
 
 """MONGODB-AWS Authentication helpers."""
 
-import os
-
 try:
-
-    from botocore.auth import SigV4Auth
-    from botocore.awsrequest import AWSRequest
-    from botocore.credentials import Credentials
-
-    import requests
-
+    import pymongo_auth_aws
+    from pymongo_auth_aws import (AwsCredential,
+                                  AwsSaslContext,
+                                  PyMongoAuthAwsError)
     _HAVE_MONGODB_AWS = True
 except ImportError:
     _HAVE_MONGODB_AWS = False
 
 import bson
-
-
-from base64 import standard_b64encode
-from collections import namedtuple
-
 from bson.binary import Binary
 from bson.son import SON
 from pymongo.errors import ConfigurationError, OperationFailure
 
 
-_AWS_REL_URI = 'http://169.254.170.2/'
-_AWS_EC2_URI = 'http://169.254.169.254/'
-_AWS_EC2_PATH = 'latest/meta-data/iam/security-credentials/'
-_AWS_HTTP_TIMEOUT = 10
-
-
-_AWSCredential = namedtuple('_AWSCredential',
-                            ['username', 'password', 'token'])
-"""MONGODB-AWS credentials."""
-
-
-def _aws_temp_credentials():
-    """Construct temporary MONGODB-AWS credentials."""
-    access_key = os.environ.get('AWS_ACCESS_KEY_ID')
-    secret_key = os.environ.get('AWS_SECRET_ACCESS_KEY')
-    if access_key and secret_key:
-        return _AWSCredential(
-            access_key, secret_key, os.environ.get('AWS_SESSION_TOKEN'))
-    # If the environment variable
-    # AWS_CONTAINER_CREDENTIALS_RELATIVE_URI is set then drivers MUST
-    # assume that it was set by an AWS ECS agent and use the URI
-    # http://169.254.170.2/$AWS_CONTAINER_CREDENTIALS_RELATIVE_URI to
-    # obtain temporary credentials.
-    relative_uri = os.environ.get('AWS_CONTAINER_CREDENTIALS_RELATIVE_URI')
-    if relative_uri is not None:
-        try:
-            res = requests.get(_AWS_REL_URI+relative_uri,
-                               timeout=_AWS_HTTP_TIMEOUT)
-            res_json = res.json()
-        except (ValueError, requests.exceptions.RequestException):
-            raise OperationFailure(
-                'temporary MONGODB-AWS credentials could not be obtained')
-    else:
-        # If the environment variable AWS_CONTAINER_CREDENTIALS_RELATIVE_URI is
-        # not set drivers MUST assume we are on an EC2 instance and use the
-        # endpoint
-        # http://169.254.169.254/latest/meta-data/iam/security-credentials
-        # /<role-name>
-        # whereas role-name can be obtained from querying the URI
-        # http://169.254.169.254/latest/meta-data/iam/security-credentials/.
-        try:
-            # Get token
-            headers = {'X-aws-ec2-metadata-token-ttl-seconds': "30"}
-            res = requests.post(_AWS_EC2_URI+'latest/api/token',
-                                headers=headers, timeout=_AWS_HTTP_TIMEOUT)
-            token = res.content
-            # Get role name
-            headers = {'X-aws-ec2-metadata-token': token}
-            res = requests.get(_AWS_EC2_URI+_AWS_EC2_PATH, headers=headers,
-                               timeout=_AWS_HTTP_TIMEOUT)
-            role = res.text
-            # Get temp creds
-            res = requests.get(_AWS_EC2_URI+_AWS_EC2_PATH+role,
-                               headers=headers, timeout=_AWS_HTTP_TIMEOUT)
-            res_json = res.json()
-        except (ValueError, requests.exceptions.RequestException):
-            raise OperationFailure(
-                'temporary MONGODB-AWS credentials could not be obtained')
-
-    try:
-        temp_user = res_json['AccessKeyId']
-        temp_password = res_json['SecretAccessKey']
-        token = res_json['Token']
-    except KeyError:
-        # If temporary credentials cannot be obtained then drivers MUST
-        # fail authentication and raise an error.
-        raise OperationFailure(
-            'temporary MONGODB-AWS credentials could not be obtained')
-
-    return _AWSCredential(temp_user, temp_password, token)
-
-
-_AWS4_HMAC_SHA256 = 'AWS4-HMAC-SHA256'
-_AWS_SERVICE = 'sts'
-
+class _AwsSaslContext(AwsSaslContext):
+    # Dependency injection:
+    def binary_type(self):
+        """Return the bson.binary.Binary type."""
+        return Binary
 
-def _get_region(sts_host):
-    """"""
-    parts = sts_host.split('.')
-    if len(parts) == 1 or sts_host == 'sts.amazonaws.com':
-        return 'us-east-1'  # Default
+    def bson_encode(self, doc):
+        """Encode a dictionary to BSON."""
+        return bson.encode(doc)
 
-    if len(parts) > 2 or not all(parts):
-        raise OperationFailure("Server returned an invalid sts host")
+    def bson_decode(self, data):
+        """Decode BSON to a dictionary."""
+        return bson.decode(data)
 
-    return parts[1]
-
-
-def _aws_auth_header(credentials, server_nonce, sts_host):
-    """Signature Version 4 Signing Process to construct the authorization header
-    """
-    region = _get_region(sts_host)
 
-    request_parameters = 'Action=GetCallerIdentity&Version=2011-06-15'
-    encoded_nonce = standard_b64encode(server_nonce).decode('utf8')
-    request_headers = {
-        'Content-Type': 'application/x-www-form-urlencoded',
-        'Content-Length': str(len(request_parameters)),
-        'Host': sts_host,
-        'X-MongoDB-Server-Nonce': encoded_nonce,
-        'X-MongoDB-GS2-CB-Flag': 'n',
-    }
-    request = AWSRequest(method="POST", url="/", data=request_parameters,
-                         headers=request_headers)
-    boto_creds = Credentials(credentials.username, credentials.password,
-                             token=credentials.token)
-    auth = SigV4Auth(boto_creds, "sts", region)
-    auth.add_auth(request)
-    final = {
-        'a': request.headers['Authorization'],
-        'd': request.headers['X-Amz-Date']
-    }
-    if credentials.token:
-        final['t'] = credentials.token
-    return final
-
-
-def _auth_aws(credentials, sock_info):
+def _authenticate_aws(credentials, sock_info):
     """Authenticate using MONGODB-AWS.
     """
     if not _HAVE_MONGODB_AWS:
         raise ConfigurationError(
-            "MONGODB-AWS authentication requires botocore and requests: "
-            "install these libraries with: "
-            "python -m pip install 'pymongo[aws]'")
+            "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")
 
-    # If a username and password are not provided, drivers MUST query
-    # a link-local AWS address for temporary credentials.
-    if credentials.username is None:
-        credentials = _aws_temp_credentials()
-
-    # Client first.
-    client_nonce = os.urandom(32)
-    payload = {'r': Binary(client_nonce), 'p': 110}
-    client_first = SON([('saslStart', 1),
-                        ('mechanism', 'MONGODB-AWS'),
-                        ('payload', Binary(bson.encode(payload)))])
-    server_first = sock_info.command('$external', client_first)
-
-    server_payload = bson.decode(server_first['payload'])
-    server_nonce = server_payload['s']
-    if len(server_nonce) != 64 or not server_nonce.startswith(client_nonce):
-        raise OperationFailure("Server returned an invalid nonce.")
-    sts_host = server_payload['h']
-    if len(sts_host) < 1 or len(sts_host) > 255 or '..' in sts_host:
-        # Drivers must also validate that the host is greater than 0 and less
-        # than or equal to 255 bytes per RFC 1035.
-        raise OperationFailure("Server returned an invalid sts host.")
-
-    payload = _aws_auth_header(credentials, server_nonce, sts_host)
-    client_second = SON([('saslContinue', 1),
-                         ('conversationId', server_first['conversationId']),
-                         ('payload', Binary(bson.encode(payload)))])
-    res = sock_info.command('$external', client_second)
-    if not res['done']:
-        raise OperationFailure('MONGODB-AWS conversation failed to complete.')
+    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:
+        # Convert to OperationFailure and include pymongo-auth-aws version.
+        raise OperationFailure('%s (pymongo-auth-aws version %s)' % (
+            exc, pymongo_auth_aws.__version__))

setup.py

@@ -329,7 +329,7 @@ def build_extension(self, ext):
     'snappy': ['python-snappy'],
     'tls': [],
     'zstd': ['zstandard'],
-    'aws': ['requests<3.0.0', 'botocore'],
+    'aws': ['pymongo-auth-aws<2.0.0'],
 }
 
 # https://jira.mongodb.org/browse/PYTHON-2117