Add SSL rotation support for PostgreSQL

Author: JoeJesse

SecretsManagerRDSPostgreSQLRotationMultiUser/lambda_function.py

@@ -18,8 +18,8 @@ def lambda_handler(event, context):
     This handler uses the master-user rotation scheme to rotate an RDS PostgreSQL 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.
+    simply creates a new secret with the AWSPREVIOUS user credentials, 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:
     {
@@ -165,8 +165,10 @@ def set_secret(service_client, arn, token):
 
     # Make sure the user from current and pending match
     if get_alt_username(current_dict['username']) != pending_dict['username']:
-        logger.error("setSecret: Attempting to modify user %s other than current user clone %s" % (pending_dict['username'], get_alt_username(current_dict['username'])))
-        raise ValueError("Attempting to modify user %s other than current user clone %s" % (pending_dict['username'], get_alt_username(current_dict['username'])))
+        logger.error(
+            "setSecret: Attempting to modify user %s other than current user or clone %s" % (pending_dict['username'], current_dict['username']))
+        raise ValueError(
+            "Attempting to modify user %s other than current user or clone %s" % (pending_dict['username'], current_dict['username']))
 
     # Make sure the host from current and pending match
     if current_dict['host'] != pending_dict['host']:
@@ -216,7 +218,7 @@ def set_secret(service_client, arn, token):
                 cur.execute(alter_role + " WITH PASSWORD %s", (pending_dict['password'],))
 
             conn.commit()
-            logger.info("setSecret: Successfully created user %s in PostgreSQL DB for secret arn %s." % (pending_dict['username'], arn))
+            logger.info("setSecret: Successfully set password for %s in PostgreSQL DB for secret arn %s." % (pending_dict['username'], arn))
     finally:
         conn.close()
 
@@ -297,8 +299,9 @@ def finish_secret(service_client, arn, token):
 def get_connection(secret_dict):
     """Gets a connection to PostgreSQL DB 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
+    This helper function uses connectivity information from the secret dictionary to initiate
+    connection attempt(s) to the database. Will attempt a fallback, non-SSL connection when
+    initial connection fails using SSL and fall_back is True.
 
     Args:
         secret_dict (dict): The Secret Dictionary
@@ -314,9 +317,87 @@ def get_connection(secret_dict):
     port = int(secret_dict['port']) if 'port' in secret_dict else 5432
     dbname = secret_dict['dbname'] if 'dbname' in secret_dict else "postgres"
 
+    # Get SSL connectivity configuration
+    use_ssl, fall_back = get_ssl_config(secret_dict)
+
+    # if an 'ssl' key is not found or does not contain a valid value, attempt an SSL connection and fall back to non-SSL on failure
+    conn = connect_and_authenticate(secret_dict, port, dbname, use_ssl)
+    if conn or not fall_back:
+        return conn
+    else:
+        return connect_and_authenticate(secret_dict, port, dbname, False)
+
+
+def get_ssl_config(secret_dict):
+    """Gets the desired SSL and fall back behavior using a secret dictionary
+
+    This helper function uses the existance and value the 'ssl' key in a secret dictionary
+    to determine desired SSL connectivity configuration. Its behavior is as follows:
+        - 'ssl' key DNE or invalid type/value: return True, True
+        - 'ssl' key is bool: return secret_dict['ssl'], False
+        - 'ssl' key equals "true" ignoring case: return True, False
+        - 'ssl' key equals "false" ignoring case: return False, False
+
+    Args:
+        secret_dict (dict): The Secret Dictionary
+
+    Returns:
+        Tuple(use_ssl, fall_back): SSL configuration
+            - use_ssl (bool): Flag indicating if an SSL connection should be attempted
+            - fall_back (bool): Flag indicating if non-SSL connection should be attempted if SSL connection fails
+
+    """
+    # Default to True for SSL and fall_back mode if 'ssl' key DNE
+    if 'ssl' not in secret_dict:
+        return True, True
+
+    # Handle type bool
+    if isinstance(secret_dict['ssl'], bool):
+        return secret_dict['ssl'], False
+
+    # Handle type string
+    if isinstance(secret_dict['ssl'], str):
+        ssl = secret_dict['ssl'].lower()
+        if ssl == "true":
+            return True, False
+        elif ssl == "false":
+            return False, False
+        else:
+            # Invalid string value, default to True for both SSL and fall_back mode
+            return True, True
+
+    # Invalid type, default to True for both SSL and fall_back mode
+    return True, True
+
+
+def connect_and_authenticate(secret_dict, port, dbname, use_ssl):
+    """Attempt to connect and authenticate to a PostgreSQL instance
+
+    This helper function tries to connect to the database using connectivity info passed in.
+    If successful, it returns the connection, else None
+
+    Args:
+        - secret_dict (dict): The Secret Dictionary
+        - port (int): The databse port to connect to
+        - dbname (str): Name of the database
+        - use_ssl (bool): Flag indicating whether connection should use SSL/TLS
+
+    Returns:
+        Connection: The pymongo.database.Database object if successful. None otherwise
+
+    Raises:
+        KeyError: If the secret json does not contain the expected keys
+
+    """
     # Try to obtain a connection to the db
     try:
-        conn = pgdb.connect(host=secret_dict['host'], user=secret_dict['username'], password=secret_dict['password'], database=dbname, port=port, connect_timeout=5)
+        if use_ssl:
+            # Setting sslmode='verify-full' will verify the server's certificate and check the server's host name
+            conn = pgdb.connect(host=secret_dict['host'], user=secret_dict['username'], password=secret_dict['password'], database=dbname, port=port,
+                                connect_timeout=5, sslrootcert='/etc/pki/tls/cert.pem', sslmode='verify-full')
+        else:
+            conn = pgdb.connect(host=secret_dict['host'], user=secret_dict['username'], password=secret_dict['password'], database=dbname, port=port,
+                                connect_timeout=5, sslmode='disable')
         return conn
     except pg.InternalError:
         return None
@@ -392,6 +473,7 @@ def get_alt_username(current_username):
             raise ValueError("Unable to clone user, username length with _clone appended would exceed 63 characters")
         return new_username
 
+
 def is_rds_replica_database(replica_dict, master_dict):
     """Validates that the database of a secret is a replica of the database of the master secret
 
@@ -429,4 +511,4 @@ def is_rds_replica_database(replica_dict, master_dict):
 
     # DB Instance identifiers are unique - can only be one result
     current_instance = instances[0]
-    return master_instance_id == current_instance.get('ReadReplicaSourceDBInstanceIdentifier')
+    return master_instance_id == current_instance.get('ReadReplicaSourceDBInstanceIdentifier')

SecretsManagerRDSPostgreSQLRotationSingleUser/lambda_function.py

@@ -172,18 +172,23 @@ def set_secret(service_client, arn, token):
 
     # Now try the current password
     conn = get_connection(current_dict)
-    if not conn:
-        if previous_dict:
-            # If both current and pending do not work, try previous
-            conn = get_connection(previous_dict)
-
-            # Make sure the user/host from previous and pending match
-            if previous_dict['username'] != pending_dict['username']:
-                logger.error("setSecret: Attempting to modify user %s other than previous valid user %s" % (pending_dict['username'], previous_dict['username']))
-                raise ValueError("Attempting to modify user %s other than previous valid user %s" % (pending_dict['username'], previous_dict['username']))
-            if previous_dict['host'] != pending_dict['host']:
-                logger.error("setSecret: Attempting to modify user for host %s other than previous valid host %s" % (pending_dict['host'], previous_dict['host']))
-                raise ValueError("Attempting to modify user for host %s other than current previous valid %s" % (pending_dict['host'], previous_dict['host']))
+
+    # If both current and pending do not work, try previous
+    if not conn and previous_dict:
+        # Update previous_dict to leverage current SSL settings
+        previous_dict.pop('ssl', None)
+        if 'ssl' in current_dict:
+            previous_dict['ssl'] = current_dict['ssl']
+
+        conn = get_connection(previous_dict)
+
+        # Make sure the user/host from previous and pending match
+        if previous_dict['username'] != pending_dict['username']:
+            logger.error("setSecret: Attempting to modify user %s other than previous valid user %s" % (pending_dict['username'], previous_dict['username']))
+            raise ValueError("Attempting to modify user %s other than previous valid user %s" % (pending_dict['username'], previous_dict['username']))
+        if previous_dict['host'] != pending_dict['host']:
+            logger.error("setSecret: Attempting to modify user for host %s other than previous valid host %s" % (pending_dict['host'], previous_dict['host']))
+            raise ValueError("Attempting to modify user for host %s other than current previous valid %s" % (pending_dict['host'], previous_dict['host']))
 
     # If we still don't have a connection, raise a ValueError
     if not conn:
@@ -278,8 +283,9 @@ def finish_secret(service_client, arn, token):
 def get_connection(secret_dict):
     """Gets a connection to PostgreSQL DB 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
+    This helper function uses connectivity information from the secret dictionary to initiate
+    connection attempt(s) to the database. Will attempt a fallback, non-SSL connection when
+    initial connection fails using SSL and fall_back is True.
 
     Args:
         secret_dict (dict): The Secret Dictionary
@@ -295,9 +301,87 @@ def get_connection(secret_dict):
     port = int(secret_dict['port']) if 'port' in secret_dict else 5432
     dbname = secret_dict['dbname'] if 'dbname' in secret_dict else "postgres"
 
+    # Get SSL connectivity configuration
+    use_ssl, fall_back = get_ssl_config(secret_dict)
+
+    # if an 'ssl' key is not found or does not contain a valid value, attempt an SSL connection and fall back to non-SSL on failure
+    conn = connect_and_authenticate(secret_dict, port, dbname, use_ssl)
+    if conn or not fall_back:
+        return conn
+    else:
+        return connect_and_authenticate(secret_dict, port, dbname, False)
+
+
+def get_ssl_config(secret_dict):
+    """Gets the desired SSL and fall back behavior using a secret dictionary
+
+    This helper function uses the existance and value the 'ssl' key in a secret dictionary
+    to determine desired SSL connectivity configuration. Its behavior is as follows:
+        - 'ssl' key DNE or invalid type/value: return True, True
+        - 'ssl' key is bool: return secret_dict['ssl'], False
+        - 'ssl' key equals "true" ignoring case: return True, False
+        - 'ssl' key equals "false" ignoring case: return False, False
+
+    Args:
+        secret_dict (dict): The Secret Dictionary
+
+    Returns:
+        Tuple(use_ssl, fall_back): SSL configuration
+            - use_ssl (bool): Flag indicating if an SSL connection should be attempted
+            - fall_back (bool): Flag indicating if non-SSL connection should be attempted if SSL connection fails
+
+    """
+    # Default to True for SSL and fall_back mode if 'ssl' key DNE
+    if 'ssl' not in secret_dict:
+        return True, True
+
+    # Handle type bool
+    if isinstance(secret_dict['ssl'], bool):
+        return secret_dict['ssl'], False
+
+    # Handle type string
+    if isinstance(secret_dict['ssl'], str):
+        ssl = secret_dict['ssl'].lower()
+        if ssl == "true":
+            return True, False
+        elif ssl == "false":
+            return False, False
+        else:
+            # Invalid string value, default to True for both SSL and fall_back mode
+            return True, True
+
+    # Invalid type, default to True for both SSL and fall_back mode
+    return True, True
+
+
+def connect_and_authenticate(secret_dict, port, dbname, use_ssl):
+    """Attempt to connect and authenticate to a PostgreSQL instance
+
+    This helper function tries to connect to the database using connectivity info passed in.
+    If successful, it returns the connection, else None
+
+    Args:
+        - secret_dict (dict): The Secret Dictionary
+        - port (int): The databse port to connect to
+        - dbname (str): Name of the database
+        - use_ssl (bool): Flag indicating whether connection should use SSL/TLS
+
+    Returns:
+        Connection: The pymongo.database.Database object if successful. None otherwise
+
+    Raises:
+        KeyError: If the secret json does not contain the expected keys
+
+    """
     # Try to obtain a connection to the db
     try:
-        conn = pgdb.connect(host=secret_dict['host'], user=secret_dict['username'], password=secret_dict['password'], database=dbname, port=port, connect_timeout=5)
+        if use_ssl:
+            # Setting sslmode='verify-full' will verify the server's certificate and check the server's host name
+            conn = pgdb.connect(host=secret_dict['host'], user=secret_dict['username'], password=secret_dict['password'], database=dbname, port=port,
+                                connect_timeout=5, sslrootcert='/etc/pki/tls/cert.pem', sslmode='verify-full')
+        else:
+            conn = pgdb.connect(host=secret_dict['host'], user=secret_dict['username'], password=secret_dict['password'], database=dbname, port=port,
+                                connect_timeout=5, sslmode='disable')
         return conn
     except pg.InternalError:
         return None