Merge pull request #1 from aws-samples/develop

Initial commit of rotation function templates

Author: willtong1234

SecretsManagerMongoDBRotationMultiUser/lambda_function.py

@@ -0,0 +1,359 @@
+# Copyright 2018 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+# SPDX-License-Identifier: MIT-0
+
+import boto3
+import json
+import logging
+import os
+from pymongo import MongoClient, errors
+
+logger = logging.getLogger()
+logger.setLevel(logging.INFO)
+
+
+def lambda_handler(event, context):
+    """Secrets Manager MongoDB Handler
+
+    This handler uses the master-user rotation scheme to rotate an MongoDB user credential. During the first rotation, this
+    scheme logs into the database as the master user, creates a new user (appending _clone to the username), and grants the
+    new user all of the permissions from the user being rotated. Once the secret is in this state, every subsequent rotation
+    simply creates a new secret with the AWSPREVIOUS user credentials, adds any missing permissions that are in the current
+    secret, changes that user's password, and then marks the latest secret as AWSCURRENT.
+
+    The Secret SecretString is expected to be a JSON string with the following format:
+    {
+        'engine': <required: must be set to 'mongo'>,
+        'host': <required: instance host name>,
+        'username': <required: username>,
+        'password': <required: password>,
+        'dbname': <optional: database name>,
+        'port': <optional: if not specified, default port 27017 will be used>,
+        'masterarn': <required: the arn of the master secret which will be used to create users/change passwords>
+    }
+
+    Args:
+        event (dict): Lambda dictionary of event parameters. These keys must include the following:
+            - SecretId: The secret ARN or identifier
+            - ClientRequestToken: The ClientRequestToken of the secret version
+            - Step: The rotation step (one of createSecret, setSecret, testSecret, or finishSecret)
+
+        context (LambdaContext): The Lambda runtime information
+
+    Raises:
+        ResourceNotFoundException: If the secret with the specified arn and stage does not exist
+
+        ValueError: If the secret is not properly configured for rotation
+
+        KeyError: If the secret json does not contain the expected keys
+
+    """
+    arn = event['SecretId']
+    token = event['ClientRequestToken']
+    step = event['Step']
+
+    # Setup the client
+    service_client = boto3.client('secretsmanager', endpoint_url=os.environ['SECRETS_MANAGER_ENDPOINT'])
+
+    # Make sure the version is staged correctly
+    metadata = service_client.describe_secret(SecretId=arn)
+    if "RotationEnabled" in metadata and not metadata['RotationEnabled']:
+        logger.error("Secret %s is not enabled for rotation" % arn)
+        raise ValueError("Secret %s is not enabled for rotation" % arn)
+    versions = metadata['VersionIdsToStages']
+    if token not in versions:
+        logger.error("Secret version %s has no stage for rotation of secret %s." % (token, arn))
+        raise ValueError("Secret version %s has no stage for rotation of secret %s." % (token, arn))
+    if "AWSCURRENT" in versions[token]:
+        logger.info("Secret version %s already set as AWSCURRENT for secret %s." % (token, arn))
+        return
+    elif "AWSPENDING" not in versions[token]:
+        logger.error("Secret version %s not set as AWSPENDING for rotation of secret %s." % (token, arn))
+        raise ValueError("Secret version %s not set as AWSPENDING for rotation of secret %s." % (token, arn))
+
+    # Call the appropriate step
+    if step == "createSecret":
+        create_secret(service_client, arn, token)
+
+    elif step == "setSecret":
+        set_secret(service_client, arn, token)
+
+    elif step == "testSecret":
+        test_secret(service_client, arn, token)
+
+    elif step == "finishSecret":
+        finish_secret(service_client, arn, token)
+
+    else:
+        logger.error("lambda_handler: Invalid step parameter %s for secret %s" % (step, arn))
+        raise ValueError("Invalid step parameter %s for secret %s" % (step, arn))
+
+
+def create_secret(service_client, arn, token):
+    """Generate a new secret
+
+    This method first checks for the existence of a secret for the passed in token. If one does not exist, it will generate a
+    new secret and save it using the passed in token.
+
+    Args:
+        service_client (client): The secrets manager service client
+
+        arn (string): The secret ARN or other identifier
+
+        token (string): The ClientRequestToken associated with the secret version
+
+    Raises:
+        ValueError: If the current secret is not valid JSON
+
+        KeyError: If the secret json does not contain the expected keys
+
+    """
+    # Make sure the current secret exists
+    current_dict = get_secret_dict(service_client, arn, "AWSCURRENT")
+
+    # Now try to get the secret version, if that fails, put a new secret
+    try:
+        get_secret_dict(service_client, arn, "AWSPENDING", token)
+        logger.info("createSecret: Successfully retrieved secret for %s." % arn)
+    except service_client.exceptions.ResourceNotFoundException:
+        # Get the alternate username swapping between the original user and the user with _clone appended to it
+        current_dict['username'] = get_alt_username(current_dict['username'])
+
+        # Generate a random password
+        passwd = service_client.get_random_password(ExcludeCharacters='/@"\'\\')
+        current_dict['password'] = passwd['RandomPassword']
+
+        # Put the secret
+        service_client.put_secret_value(SecretId=arn, ClientRequestToken=token, SecretString=json.dumps(current_dict), VersionStages=['AWSPENDING'])
+        logger.info("createSecret: Successfully put secret for ARN %s and version %s." % (arn, token))
+
+
+def set_secret(service_client, arn, token):
+    """Set the pending secret in the database
+
+    This method tries to login to the database with the AWSPENDING secret and returns on success. If that fails, it
+    tries to login with the master credentials from the masterarn in the current secret. If this succeeds, it adds all
+    grants for AWSCURRENT user to the AWSPENDING user, creating the user and/or setting the password in the process.
+    Else, it throws a ValueError.
+
+    Args:
+        service_client (client): The secrets manager service client
+
+        arn (string): The secret ARN or other identifier
+
+        token (string): The ClientRequestToken associated with the secret version
+
+    Raises:
+        ResourceNotFoundException: If the secret with the specified arn and stage does not exist
+
+        ValueError: If the secret is not valid JSON or master credentials could not be used to login to DB
+
+        KeyError: If the secret json does not contain the expected keys
+
+    """
+    # First try to login with the pending secret, if it succeeds, return
+    pending_dict = get_secret_dict(service_client, arn, "AWSPENDING", token)
+    conn = get_connection(pending_dict)
+    if conn:
+        conn.logout()
+        logger.info("setSecret: AWSPENDING secret is already set as password in MongoDB for secret arn %s." % arn)
+        return
+
+    # Before we do anything with the secret, make sure the AWSCURRENT secret is valid by logging in to the db
+    # This ensures that the credential we are rotating is valid to protect against a confused deputy attack
+    current_dict = get_secret_dict(service_client, arn, "AWSCURRENT")
+    conn = get_connection(current_dict)
+    if not conn:
+        logger.error("setSecret: Unable to log into database using current credentials for secret %s" % arn)
+        raise ValueError("Unable to log into database using current credentials for secret %s" % arn)
+    conn.logout()
+
+    # Now get the master arn from the current secret
+    master_arn = current_dict['masterarn']
+    master_dict = get_secret_dict(service_client, master_arn, "AWSCURRENT")
+    if current_dict['host'] != master_dict['host']:
+        logger.warn("setSecret: Master database host %s is not the same host as current %s" % (master_dict['host'], current_dict['host']))
+
+    # Now log into the database with the master credentials
+    conn = get_connection(master_dict)
+    if not conn:
+        logger.error("setSecret: Unable to log into database using credentials in master secret %s" % master_arn)
+        raise ValueError("Unable to log into database using credentials in master secret %s" % master_arn)
+
+    # Now set the password to the pending password
+    try:
+        user_info = conn.command('usersInfo', pending_dict['username'])
+        if user_info['users']:
+            conn.command("updateUser", pending_dict['username'], pwd=pending_dict['password'])
+        else:
+            user_info = conn.command('usersInfo', current_dict['username'])
+            conn.command("createUser", pending_dict['username'], pwd=pending_dict['password'], roles=user_info['users'][0]['roles'])
+        logger.info("setSecret: Successfully set password for %s in MongoDB for secret arn %s." % (pending_dict['username'], arn))
+    finally:
+        conn.logout()
+
+
+def test_secret(service_client, arn, token):
+    """Test the pending secret against the database
+
+    This method tries to log into the database with the secrets staged with AWSPENDING and runs
+    a permissions check to ensure the user has the correct permissions.
+
+    Args:
+        service_client (client): The secrets manager service client
+
+        arn (string): The secret ARN or other identifier
+
+        token (string): The ClientRequestToken associated with the secret version
+
+    Raises:
+        ResourceNotFoundException: If the secret with the specified arn and stage does not exist
+
+        ValueError: If the secret is not valid JSON or pending credentials could not be used to login to the database
+
+        KeyError: If the secret json does not contain the expected keys
+
+    """
+    # Try to login with the pending secret, if it succeeds, return
+    pending_dict = get_secret_dict(service_client, arn, "AWSPENDING", token)
+    conn = get_connection(pending_dict)
+    if conn:
+        # This is where the lambda will validate the user's permissions. Modify the below lines to
+        # tailor these validations to your needs
+        try:
+            conn.command('usersInfo', pending_dict['username'])
+        finally:
+            conn.logout()
+
+        logger.info("testSecret: Successfully signed into MongoDB with AWSPENDING secret in %s." % arn)
+        return
+    else:
+        logger.error("testSecret: Unable to log into database with pending secret of secret ARN %s" % arn)
+        raise ValueError("Unable to log into database with pending secret of secret ARN %s" % arn)
+
+
+def finish_secret(service_client, arn, token):
+    """Finish the rotation by marking the pending secret as current
+
+    This method moves the secret from the AWSPENDING stage to the AWSCURRENT stage.
+
+    Args:
+        service_client (client): The secrets manager service client
+
+        arn (string): The secret ARN or other identifier
+
+        token (string): The ClientRequestToken associated with the secret version
+
+    Raises:
+        ResourceNotFoundException: If the secret with the specified arn and stage does not exist
+
+    """
+    # First describe the secret to get the current version
+    metadata = service_client.describe_secret(SecretId=arn)
+    current_version = None
+    for version in metadata["VersionIdsToStages"]:
+        if "AWSCURRENT" in metadata["VersionIdsToStages"][version]:
+            if version == token:
+                # The correct version is already marked as current, return
+                logger.info("finishSecret: Version %s already marked as AWSCURRENT for %s" % (version, arn))
+                return
+            current_version = version
+            break
+
+    # Finalize by staging the secret version current
+    service_client.update_secret_version_stage(SecretId=arn, VersionStage="AWSCURRENT", MoveToVersionId=token, RemoveFromVersionId=current_version)
+    logger.info("finishSecret: Successfully set AWSCURRENT stage to version %s for secret %s." % (version, arn))
+
+
+def get_connection(secret_dict):
+    """Gets a connection to MongoDB from a secret dictionary
+
+    This helper function tries to connect to the database grabbing connection info
+    from the secret dictionary. If successful, it returns the connection, else None
+
+    Args:
+        secret_dict (dict): The Secret Dictionary
+
+    Returns:
+        Connection: The pymongo.database.Database object if successful. None otherwise
+
+    Raises:
+        KeyError: If the secret json does not contain the expected keys
+
+    """
+    port = int(secret_dict['port']) if 'port' in secret_dict else 27017
+    dbname = secret_dict['dbname'] if 'dbname' in secret_dict else "admin"
+
+    # Try to obtain a connection to the db
+    try:
+        client = MongoClient(host=secret_dict['host'], port=port, connectTimeoutMS=5000, serverSelectionTimeoutMS=5000)
+        db = client[dbname]
+        db.authenticate(secret_dict['username'], secret_dict['password'])
+        return db
+    except errors.PyMongoError:
+        return None
+
+
+def get_secret_dict(service_client, arn, stage, token=None):
+    """Gets the secret dictionary corresponding for the secret arn, stage, and token
+
+    This helper function gets credentials for the arn and stage passed in and returns the dictionary by parsing the JSON string
+
+    Args:
+        service_client (client): The secrets manager service client
+
+        arn (string): The secret ARN or other identifier
+
+        token (string): The ClientRequestToken associated with the secret version, or None if no validation is desired
+
+        stage (string): The stage identifying the secret version
+
+    Returns:
+        SecretDictionary: Secret dictionary
+
+    Raises:
+        ResourceNotFoundException: If the secret with the specified arn and stage does not exist
+
+        ValueError: If the secret is not valid JSON
+
+    """
+    required_fields = ['host', 'username', 'password']
+
+    # Only do VersionId validation against the stage if a token is passed in
+    if token:
+        secret = service_client.get_secret_value(SecretId=arn, VersionId=token, VersionStage=stage)
+    else:
+        secret = service_client.get_secret_value(SecretId=arn, VersionStage=stage)
+    plaintext = secret['SecretString']
+    secret_dict = json.loads(plaintext)
+
+    # Run validations against the secret
+    if 'engine' not in secret_dict or secret_dict['engine'] != 'mongo':
+        raise KeyError("Database engine must be set to 'mongo' in order to use this rotation lambda")
+    for field in required_fields:
+        if field not in secret_dict:
+            raise KeyError("%s key is missing from secret JSON" % field)
+
+    # Parse and return the secret JSON string
+    return secret_dict
+
+
+def get_alt_username(current_username):
+    """Gets the alternate username for the current_username passed in
+
+    This helper function gets the username for the alternate user based on the passed in current username.
+
+    Args:
+        current_username (client): The current username
+
+    Returns:
+        AlternateUsername: Alternate username
+
+    Raises:
+        ValueError: If the new username length would exceed the maximum allowed
+
+    """
+    clone_suffix = "_clone"
+    if current_username.endswith(clone_suffix):
+        return current_username[:(len(clone_suffix) * -1)]
+    else:
+        return current_username + clone_suffix