Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
34 changes: 17 additions & 17 deletions SecretsManagerRDSPostgreSQLRotationMultiUser/lambda_function.py
Original file line number Diff line number Diff line change
Expand Up @@ -61,17 +61,17 @@ def lambda_handler(event, context):
# 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)
logger.error("Secret %s is not enabled for rotation" % arn.encode())
raise ValueError("Secret %s is not enabled for rotation" % arn.encode())
versions = metadata['VersionIdsToStages']
if token not in versions:
logger.error("Secret version %s has no stage for rotation of secret %s." % (token, arn))
logger.error("Secret version %s has no stage for rotation of secret %s." % (token.encode(), arn.encode()))
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))
logger.info("Secret version %s already set as AWSCURRENT for secret %s." % (token.encode(), arn.encode()))
return
elif "AWSPENDING" not in versions[token]:
logger.error("Secret version %s not set as AWSPENDING for rotation of secret %s." % (token, arn))
logger.error("Secret version %s not set as AWSPENDING for rotation of secret %s." % (token.encode(), arn.encode()))
raise ValueError("Secret version %s not set as AWSPENDING for rotation of secret %s." % (token, arn))

# Call the appropriate step
Expand All @@ -88,7 +88,7 @@ def lambda_handler(event, context):
finish_secret(service_client, arn, token)

else:
logger.error("lambda_handler: Invalid step parameter %s for secret %s" % (step, arn))
logger.error("lambda_handler: Invalid step parameter %s for secret %s" % (step.encode(), arn.encode()))
raise ValueError("Invalid step parameter %s for secret %s" % (step, arn))


Expand Down Expand Up @@ -117,15 +117,15 @@ def create_secret(service_client, arn, token):
# 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)
logger.info("createSecret: Successfully retrieved secret for %s." % arn.encode())
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'])
current_dict['password'] = get_random_password(service_client)

# 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))
logger.info("createSecret: Successfully put secret for ARN %s and version %s." % (arn.encode(), token.encode()))


def set_secret(service_client, arn, token):
Expand Down Expand Up @@ -158,7 +158,7 @@ def set_secret(service_client, arn, token):
conn = get_connection(pending_dict)
if conn:
conn.close()
logger.info("setSecret: AWSPENDING secret is already set as password in PostgreSQL DB for secret arn %s." % arn)
logger.info("setSecret: AWSPENDING secret is already set as password in PostgreSQL DB for secret arn %s." % arn.encode())
return

# Make sure the user from current and pending match
Expand All @@ -177,7 +177,7 @@ def set_secret(service_client, arn, token):
# This ensures that the credential we are rotating is valid to protect against a confused deputy attack
conn = get_connection(current_dict)
if not conn:
logger.error("setSecret: Unable to log into database using current credentials for secret %s" % arn)
logger.error("setSecret: Unable to log into database using current credentials for secret %s" % arn.encode())
raise ValueError("Unable to log into database using current credentials for secret %s" % arn)
conn.close()

Expand Down Expand Up @@ -220,7 +220,7 @@ def set_secret(service_client, arn, token):
cur.execute(alter_role + " WITH PASSWORD %s", (pending_dict['password'],))

conn.commit()
logger.info("setSecret: Successfully set password for %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.encode()))
finally:
conn.close()

Expand Down Expand Up @@ -258,11 +258,11 @@ def test_secret(service_client, arn, token):
finally:
conn.close()

logger.info("testSecret: Successfully signed into PostgreSQL DB with AWSPENDING secret in %s." % arn)
logger.info("testSecret: Successfully signed into PostgreSQL DB with AWSPENDING secret in %s." % arn.encode())
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)
logger.error("testSecret: Unable to log into database with pending secret of secret ARN %s" % arn.encode())
raise ValueError("Unable to log into database with pending secret of secret ARN %s" % arn.encode())


def finish_secret(service_client, arn, token):
Expand All @@ -288,14 +288,14 @@ def finish_secret(service_client, arn, token):
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))
logger.info("finishSecret: Version %s already marked as AWSCURRENT for %s" % (version, arn.encode()))
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." % (token, arn))
logger.info("finishSecret: Successfully set AWSCURRENT stage to version %s for secret %s." % (token.encode(), arn.encode()))


def get_connection(secret_dict):
Expand Down Expand Up @@ -457,7 +457,7 @@ def get_secret_dict(service_client, arn, stage, token=None, master_secret=False)
db_instance_info = fetch_instance_arn_from_system_tags(service_client, arn)
if len(db_instance_info) != 0:
secret_dict = get_connection_params_from_rds_api(secret_dict, db_instance_info)
logger.info("setSecret: Successfully fetched connection params for Master Secret %s from DescribeDBInstances API." % arn)
logger.info("setSecret: Successfully fetched connection params for Master Secret %s from DescribeDBInstances API." % arn.encode())

# For non-RDS-made Master Secrets that are missing `host`, this will error below when checking for required connection params.

Expand Down
30 changes: 15 additions & 15 deletions SecretsManagerRDSPostgreSQLRotationSingleUser/lambda_function.py
Original file line number Diff line number Diff line change
Expand Up @@ -56,17 +56,17 @@ def lambda_handler(event, context):
# 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)
logger.error("Secret %s is not enabled for rotation" % arn.encode())
raise ValueError("Secret %s is not enabled for rotation" % arn.encode())
versions = metadata['VersionIdsToStages']
if token not in versions:
logger.error("Secret version %s has no stage for rotation of secret %s." % (token, arn))
logger.error("Secret version %s has no stage for rotation of secret %s." % (token.encode(), arn.encode()))
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))
logger.info("Secret version %s already set as AWSCURRENT for secret %s." % (token.encode(), arn.encode()))
return
elif "AWSPENDING" not in versions[token]:
logger.error("Secret version %s not set as AWSPENDING for rotation of secret %s." % (token, arn))
logger.error("Secret version %s not set as AWSPENDING for rotation of secret %s." % (token.encode(), arn.encode()))
raise ValueError("Secret version %s not set as AWSPENDING for rotation of secret %s." % (token, arn))

# Call the appropriate step
Expand All @@ -83,7 +83,7 @@ def lambda_handler(event, context):
finish_secret(service_client, arn, token)

else:
logger.error("lambda_handler: Invalid step parameter %s for secret %s" % (step, arn))
logger.error("lambda_handler: Invalid step parameter %s for secret %s" % (step.encode(), arn.encode()))
raise ValueError("Invalid step parameter %s for secret %s" % (step, arn))


Expand Down Expand Up @@ -112,13 +112,13 @@ def create_secret(service_client, arn, token):
# 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)
logger.info("createSecret: Successfully retrieved secret for %s." % arn.encode())
except service_client.exceptions.ResourceNotFoundException:
# Generate a random password
current_dict['password'] = get_random_password(service_client)
# 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))
logger.info("createSecret: Successfully put secret for ARN %s and version %s." % (arn.encode(), token.encode()))


def set_secret(service_client, arn, token):
Expand Down Expand Up @@ -154,7 +154,7 @@ def set_secret(service_client, arn, token):
conn = get_connection(pending_dict)
if conn:
conn.close()
logger.info("setSecret: AWSPENDING secret is already set as password in PostgreSQL DB for secret arn %s." % arn)
logger.info("setSecret: AWSPENDING secret is already set as password in PostgreSQL DB for secret arn %s." % arn.encode())
return

# Make sure the user from current and pending match
Expand Down Expand Up @@ -189,7 +189,7 @@ def set_secret(service_client, arn, token):

# If we still don't have a connection, raise a ValueError
if not conn:
logger.error("setSecret: Unable to log into database with previous, current, or pending secret of secret arn %s" % arn)
logger.error("setSecret: Unable to log into database with previous, current, or pending secret of secret arn %s" % arn.encode())
raise ValueError("Unable to log into database with previous, current, or pending secret of secret arn %s" % arn)

# Now set the password to the pending password
Expand All @@ -202,7 +202,7 @@ def set_secret(service_client, arn, token):
alter_role = "ALTER USER %s" % escaped_username
cur.execute(alter_role + " WITH PASSWORD %s", (pending_dict['password'],))
conn.commit()
logger.info("setSecret: Successfully set password for user %s in PostgreSQL DB for secret arn %s." % (pending_dict['username'], arn))
logger.info("setSecret: Successfully set password for user %s in PostgreSQL DB for secret arn %s." % (pending_dict['username'], arn.encode()))
finally:
conn.close()

Expand Down Expand Up @@ -240,10 +240,10 @@ def test_secret(service_client, arn, token):
finally:
conn.close()

logger.info("testSecret: Successfully signed into PostgreSQL DB with AWSPENDING secret in %s." % arn)
logger.info("testSecret: Successfully signed into PostgreSQL DB with AWSPENDING secret in %s." % arn.encode())
return
else:
logger.error("testSecret: Unable to log into database with pending secret of secret ARN %s" % arn)
logger.error("testSecret: Unable to log into database with pending secret of secret ARN %s" % arn.encode())
raise ValueError("Unable to log into database with pending secret of secret ARN %s" % arn)


Expand All @@ -267,14 +267,14 @@ def finish_secret(service_client, arn, token):
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))
logger.info("finishSecret: Version %s already marked as AWSCURRENT for %s" % (version, arn.encode()))
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." % (token, arn))
logger.info("finishSecret: Successfully set AWSCURRENT stage to version %s for secret %s." % (token.encode(), arn.encode()))


def get_connection(secret_dict):
Expand Down