如何修改AWS Secrets Manager自动轮换Lambda函数仅替换Password字段

我有一个AWS Secrets Manager自动轮换Lambda函数，当前它会用随机密码覆盖整个Secret String，但我需要仅替换其中名为"Password"的键值对。我的Secret String格式为：

{"Password":"*********","Username":"********"}

原始Lambda函数代码 #

import boto3

def lambda_handler(event, context):
    arn = event['SecretId']
    token = event['ClientRequestToken']
    step = event['Step']

    # Setup the client
    service_client = boto3.client('secretsmanager')

    # Make sure the version is staged correctly
    metadata = service_client.describe_secret(SecretId=arn)
    if not metadata['RotationEnabled']:
        raise ValueError("Secret %s is not enabled for rotation" % arn)
    
    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:
        raise ValueError("Invalid step parameter")


def create_secret(service_client, arn, token):

    service_client.get_secret_value(SecretId=arn, VersionStage="AWSCURRENT")

    # Now try to get the secret version, if that fails, put a new secret
    try:
        service_client.get_secret_value(SecretId=arn, VersionId=token, VersionStage="AWSPENDING")
    except service_client.exceptions.ResourceNotFoundException:
        # Generate a random password
        passwd = service_client.get_random_password(ExcludeCharacters='/@"\'\\')
        # Put the secret
        service_client.put_secret_value(SecretId=arn, ClientRequestToken=token, SecretString=passwd['RandomPassword'], VersionStages=['AWSPENDING'])


def set_secret(service_client, arn, token):
    print("No database user credentials to update...")


def test_secret(service_client, arn, token):
    print("No need to testing against any service...")


def finish_secret(service_client, arn, token):
   
    # First describe the secret to get the current version
    metadata = service_client.describe_secret(SecretId=arn)

    for version in metadata["VersionIdsToStages"]:
        if "AWSCURRENT" in metadata["VersionIdsToStages"][version]:
            if version == token:
                # The correct version is already marked as current, return
                return

            # Finalize by staging the secret version current
            service_client.update_secret_version_stage(SecretId=arn, VersionStage="AWSCURRENT", MoveToVersionId=token, RemoveFromVersionId=version)
            break

我尝试过但未成功的代码 #

import string
import random
import boto3

def lambda_handler(event, context):
    # Set the name of the secret
    secret_name = event['SecretId']
    
    # Connect to AWS Secrets Manager
    session = boto3.session.Session()
    client = session.client(service_name='secretsmanager')

    # Retrieve the current secret value
    get_secret_value_response = client.get_secret_value(SecretId=secret_name)
    current_secret_value = get_secret_value_response['SecretString']
    
    # Define the complexity of the new secret value
    password_length = 16
    uppercase_letters = string.ascii_uppercase
    lowercase_letters = string.ascii_lowercase
    digits = string.digits
    special_chars = "!@#$%^&*()_+-="
    
    # Parse the secret value to extract the part that needs to be rotated
    old_value = current_secret_value.split('=')[1]
    
    # Generate a new secret value with the defined complexity
    new_value_characters = []
    new_value_characters.extend(random.choices(uppercase_letters, k=4))
    new_value_characters.extend(random.choices(lowercase_letters, k=4))
    new_value_characters.extend(random.choices(digits, k=4))
    new_value_characters.extend(random.choices(special_chars, k=4))
    new_value_characters.extend(random.choices(uppercase_letters + lowercase_letters + digits + special_chars, k=password_length-16))
    new_value = ''.join(random.sample(new_value_characters, k=password_length))
    
    # Replace the rotated part of the secret value with the new value
    new_secret_value = current_secret_value.replace(old_value, new_value)
    
    # Update the secret value for the specified secret key in AWS Secrets Manager
    put_secret_value_response = client.put_secret_value(SecretId=secret_name, SecretString=new_secret_value)
    
    # Return the new secret value
    return {"new_secret_value": new_secret_value}

解决方案：修改create_secret函数仅替换Password字段 #

要实现仅更新Secret String中的Password键，需要解析当前的JSON格式Secret String，替换Password值后再重新存入。修改后的完整代码如下：

import boto3
import json

def lambda_handler(event, context):
    arn = event['SecretId']
    token = event['ClientRequestToken']
    step = event['Step']

    # Setup the client
    service_client = boto3.client('secretsmanager')

    # Make sure the version is staged correctly
    metadata = service_client.describe_secret(SecretId=arn)
    if not metadata['RotationEnabled']:
        raise ValueError("Secret %s is not enabled for rotation" % arn)
    
    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:
        raise ValueError("Invalid step parameter")


def create_secret(service_client, arn, token):
    # 获取当前版本的Secret String并解析为字典
    current_secret_response = service_client.get_secret_value(SecretId=arn, VersionStage="AWSCURRENT")
    secret_dict = json.loads(current_secret_response['SecretString'])

    # 检查AWSPENDING版本是否存在，不存在则创建
    try:
        service_client.get_secret_value(SecretId=arn, VersionId=token, VersionStage="AWSPENDING")
    except service_client.exceptions.ResourceNotFoundException:
        # 生成符合要求的随机密码
        passwd_response = service_client.get_random_password(ExcludeCharacters='/@"\'\\')
        new_password = passwd_response['RandomPassword']
        # 仅替换Password字段
        secret_dict['Password'] = new_password
        # 转换回JSON字符串
        new_secret_string = json.dumps(secret_dict)
        # 存入AWSPENDING版本
        service_client.put_secret_value(
            SecretId=arn,
            ClientRequestToken=token,
            SecretString=new_secret_string,
            VersionStages=['AWSPENDING']
        )


def set_secret(service_client, arn, token):
    print("No database user credentials to update...")


def test_secret(service_client, arn, token):
    print("No need to testing against any service...")


def finish_secret(service_client, arn, token):
    # 获取当前Secret的版本信息
    metadata = service_client.describe_secret(SecretId=arn)

    for version in metadata["VersionIdsToStages"]:
        if "AWSCURRENT" in metadata["VersionIdsToStages"][version]:
            if version == token:
                # 目标版本已经是当前版本，无需操作
                return

            # 将AWSPENDING版本标记为AWSCURRENT，移除旧版本的标记
            service_client.update_secret_version_stage(SecretId=arn, VersionStage="AWSCURRENT", MoveToVersionId=token, RemoveFromVersionId=version)
            break

修改说明 #

1. 新增json模块导入，用于解析和生成JSON格式的Secret String
2. 在create_secret函数中：
   • 获取当前AWSCURRENT版本的Secret内容，解析为字典格式
   • 生成新密码后，仅替换字典中的Password键值
   • 将修改后的字典转换回JSON字符串，存入AWSPENDING版本
3. 保留原有的轮换流程逻辑，若后续需要对接数据库更新密码，可在set_secret函数中添加对应操作

内容的提问来源于stack exchange，提问作者metalslug2004