Merge pull request #15 from willtong1234/develop

Support DocumentDB and MySQL 8

Author: willtong1234

SecretsManagerMongoDBRotationMultiUser/lambda_function.py

@@ -29,6 +29,7 @@ def lambda_handler(event, context):
         '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>
+        'ssl': <optional: if not specified, defaults to false. This must be true if being used for DocumentDB rotations where the cluster has TLS enabled>
     }
 
     Args:
@@ -282,10 +283,13 @@ def get_connection(secret_dict):
     """
     port = int(secret_dict['port']) if 'port' in secret_dict else 27017
     dbname = secret_dict['dbname'] if 'dbname' in secret_dict else "admin"
+    ssl = False
+    if 'ssl' in secret_dict:
+        ssl = (secret_dict['ssl'].lower() == "true") if type(secret_dict['ssl']) is str else bool(secret_dict['ssl'])
 
     # Try to obtain a connection to the db
     try:
-        client = MongoClient(host=secret_dict['host'], port=port, connectTimeoutMS=5000, serverSelectionTimeoutMS=5000)
+        client = MongoClient(host=secret_dict['host'], port=port, connectTimeoutMS=5000, serverSelectionTimeoutMS=5000, ssl=ssl)
         db = client[dbname]
         db.authenticate(secret_dict['username'], secret_dict['password'])
         return db

SecretsManagerMongoDBRotationSingleUser/lambda_function.py

@@ -26,6 +26,7 @@ def lambda_handler(event, context):
         'password': <required: password>,
         'dbname': <optional: database name>,
         'port': <optional: if not specified, default port 27017 will be used>
+        'ssl': <optional: if not specified, defaults to false. This must be true if being used for DocumentDB rotations where the cluster has TLS enabled>
     }
 
     Args:
@@ -261,10 +262,13 @@ def get_connection(secret_dict):
     # Parse and validate the secret JSON string
     port = int(secret_dict['port']) if 'port' in secret_dict else 27017
     dbname = secret_dict['dbname'] if 'dbname' in secret_dict else "admin"
+    ssl = False
+    if 'ssl' in secret_dict:
+        ssl = (secret_dict['ssl'].lower() == "true") if type(secret_dict['ssl']) is str else bool(secret_dict['ssl'])
 
     # Try to obtain a connection to the db
     try:
-        client = MongoClient(host=secret_dict['host'], port=port, connectTimeoutMS=5000, serverSelectionTimeoutMS=5000)
+        client = MongoClient(host=secret_dict['host'], port=port, connectTimeoutMS=5000, serverSelectionTimeoutMS=5000, ssl=ssl)
         db = client[dbname]
         db.authenticate(secret_dict['username'], secret_dict['password'])
         return db

SecretsManagerRDSMySQLRotationMultiUser/lambda_function.py

@@ -182,14 +182,24 @@ def set_secret(service_client, arn, token):
     # Now set the password to the pending password
     try:
         with conn.cursor() as cur:
-            # List the grants on the current user and add them to the pending user.
-            # This also creates the user if it does not already exist
+            cur.execute("SELECT User FROM mysql.user WHERE User = %s", pending_dict['username'])
+            # Create the user if it does not exist
+            if cur.rowcount == 0:
+                cur.execute("CREATE USER %s IDENTIFIED BY %s", (pending_dict['username'], pending_dict['password']))
+
+            # Copy grants to the new user
             cur.execute("SHOW GRANTS FOR %s", current_dict['username'])
             for row in cur.fetchall():
                 grant = row[0].split(' TO ')
                 new_grant = "%s TO '%s'" % (grant[0], pending_dict['username'])
                 new_grant_escaped = new_grant.replace('%','%%') # % is a special character in Python format strings.
-                cur.execute(new_grant_escaped + " IDENTIFIED BY %s", pending_dict['password'])
+                cur.execute(new_grant_escaped)
+
+            # Set the password for the user and commit
+            cur.execute("SELECT VERSION()")
+            ver = cur.fetchone()
+            password_option = get_password_option(ver[0])
+            cur.execute("SET PASSWORD FOR %s = " + password_option, (pending_dict['username'], pending_dict['password']))
             conn.commit()
             logger.info("setSecret: Successfully set password for %s in MySQL DB for secret arn %s." % (pending_dict['username'], arn))
     finally:
@@ -363,3 +373,22 @@ def get_alt_username(current_username):
         if len(new_username) > 16:
             raise ValueError("Unable to clone user, username length with _clone appended would exceed 16 characters")
         return new_username
+
+
+def get_password_option(version):
+    """Gets the password option template string to use for the SET PASSWORD sql query
+
+    This helper function takes in the mysql version and returns the appropriate password option template string that can
+    be used in the SET PASSWORD query for that mysql version.
+
+    Args:
+        version (string): The mysql database version
+
+    Returns:
+        PasswordOption: The password option string
+
+    """
+    if version.startswith("8"):
+        return "%s"
+    else:
+        return "PASSWORD(%s)"

SecretsManagerRDSMySQLRotationSingleUser/lambda_function.py

@@ -168,7 +168,10 @@ def set_secret(service_client, arn, token):
     # Now set the password to the pending password
     try:
         with conn.cursor() as cur:
-            cur.execute("SET PASSWORD = PASSWORD(%s)", pending_dict['password'])
+            cur.execute("SELECT VERSION()")
+            ver = cur.fetchone()
+            password_option = get_password_option(ver[0])
+            cur.execute("SET PASSWORD = " + password_option, pending_dict['password'])
             conn.commit()
             logger.info("setSecret: Successfully set password for user %s in MySQL DB for secret arn %s." % (pending_dict['username'], arn))
     finally:
@@ -315,3 +318,22 @@ def get_secret_dict(service_client, arn, stage, token=None):
 
     # Parse and return the secret JSON string
     return secret_dict
+
+
+def get_password_option(version):
+    """Gets the password option template string to use for the SET PASSWORD sql query
+
+    This helper function takes in the mysql version and returns the appropriate password option template string that can
+    be used in the SET PASSWORD query for that mysql version.
+
+    Args:
+        version (string): The mysql database version
+
+    Returns:
+        PasswordOption: The password option string
+
+    """
+    if version.startswith("8"):
+        return "%s"
+    else:
+        return "PASSWORD(%s)"