AWS provides managed SSM documents for specific actions which you can add to your maintenance window tasks. However, there are cases where you need a custom action, such as this one.
This blog post will discuss a Python script that leverages AWS serverless and NoSQL database services to automate the creation and deletion of an AMI whose instance had just been patched using SSM. We'll explore how the script manages state sharing between functions and stores state information in a DynamoDB table, all while ensuring that your AMIs are maintained seamlessly.
Understanding State Sharing Generally
State sharing in Python allows multiple functions to access and modify shared data, enabling them to communicate and coordinate using a common state. Here's a concise Python example:
# Shared global variable
shared_counter = 0
# Function to increment
def increment_counter():
global shared_counter
shared_counter += 1
# Function to get the value
def get_counter_value():
return shared_counter
# Example usage
increment_counter()
print(get_counter_value()) # Output: 1
increment_counter()
print(get_counter_value()) # Output: 2
In this code, shared_counter is a global variable, shared between functions for data sharing and coordination.
Understanding State Sharing in my Script
State sharing in the main script is achieved through a global variable called image_id. This variable is accessible and modifiable by multiple functions, allowing them to share and update the same state information, such as the currently active AMI. This facilitates the automation of AMI creation and deletion while maintaining a consistent state across functions.
The Script
The script starts by initializing AWS clients with Boto3. It creates clients for EC2 (ec2) and DynamoDB (dynamo) to handle AWS services. A DynamoDB table named AMI_Table is defined with a partition key structure for item identification.
import boto3
import datetime
# Initialize the AWS clients
ec2 = boto3.client('ec2')
dynamo = boto3.client('dynamodb')
table_name = 'AMI_Table'
key = {
'PK': {
'S': 'Partition_Key'
}
}
The table format
Storing State in DynamoDB
State management is crucial in any automation process, and DynamoDB serves as an excellent choice for this task. In this script, state information, specifically the image_id, is stored in the DynamoDB table.
The script retrieves the image_id from the DynamoDB table, which acts as the identifier for the currently active AMI. This is the key to maintaining state information across multiple function invocations.
# Get ami_id value from dynamodb AMI_Table
get_ami_id = dynamo.get_item(TableName=table_name, Key=key)
# Pass the value to a new variable
image_id = get_ami_id['Item']['AMI_ID']['S']
AMI Creation
The script defines a function create_ami(InstanceId) for creating a new AMI. It takes an EC2 InstanceId as input, and here's how it works:
- It creates a new AMI based on the provided
InstanceIdusing theec2.create_imagefunction. - The
current_timeis used to generate a unique name for the new AMI. - The script updates the
image_idglobal variable with the newly created AMI's ID. - It stores the
image_idvalue in the DynamoDB table to maintain the state.
def create_ami(InstanceId):
# Declare image_id as a global variable
global image_id
# Get the current time to use as the AMI name
current_time = datetime.datetime.now().strftime('%Y-%m- %d-%H-%M-%S')
if InstanceId:
create_image = ec2.create_image(
InstanceId=InstanceId,
# Use the current time as the image name
Name=f'{current_time}_AMI'
)
# Update the global image_id variable
image_id = create_image['ImageId']
# Store the image_id value as a state in a dynamodb ami_table
update_ami_id = dynamo.update_item(
TableName=table_name,
Key=key,
# (:) in :image_id is used to define a placeholder for
# attribute name or a value that will be set.
UpdateExpression="SET AMI_ID = :image_id",
# Use :image_id as a placeholder, its actual value is desired_state.
ExpressionAttributeValues={
':image_id': {
'S': image_id
}
}
)
# Return the value so when the function is passed to a variable, a value exists.
return image_id
else:
return "InstanceId not found."
AMI Deletion
Another function, delete_ami(image_id), is responsible for deleting the existing AMI. It follows these steps:
- Deregister the current AMI using
ec2.deregister_image. - Retrieve the snapshot associated with the AMI and delete it.
def delete_ami(image_id):
# Deregister the AMI
deregister_ami = ec2.deregister_image(ImageId=image_id)
# Get the snapshot ID associated with the AMI
ami_info = ec2.describe_images(ImageIds=[image_id])
snapshot_id = ami_info['Images'][0]['BlockDeviceMappings'][0]['Ebs']['SnapshotId']
# Delete the associated snapshot
delete_snapshot = ec2.delete_snapshot(SnapshotId=snapshot_id)
Updating State
The update_state(desired_state) function is used to update the state in the DynamoDB table. This function is essential for transitioning between creating and deleting AMIs.
def update_state(desired_state):
get_ami_state = dynamo.get_item(TableName=table_name, Key=key)
# Update the value of AMI_state
dynamo.update_item(
TableName=table_name,
Key=key,
# (:) in :new_state is used to define a placeholder for
# attribute name or a value that will be set.
UpdateExpression='SET AMI_State = :new_state',
# Use :new_state as a placeholder, its actual value is desired_state.
ExpressionAttributeValues={':new_state': {'S': desired_state}}
)
Lambda Function
Finally, the script provides a Lambda function called lambda_handler(event, context) that orchestrates the entire process. It reads the current state from DynamoDB, which will be 'Create_AMI' then it executes the code block for it. At the end of it, it updates the current NoSQL state to 'Delete_AMI' so during further execution, only the code block having this state value runs.
The script handles cases where the InstanceId is missing and ensures that the state is updated correctly.
def lambda_handler(event, context):
try:
# Read the dynamodb table
db_read = dynamo.get_item(TableName=table_name, Key=key)
# Get the AMI State from the dynamodb table
ami_state = db_read['Item']['AMI_State']['S']
if ami_state == 'Create_AMI':
# Extract the InstanceId from the event dictionary
InstanceId = event.get('InstanceId')
# Check if InstanceId value exists
if InstanceId:
# Call the create_ami function and pass the InstanceId as a string.
create_ami(InstanceId)
# Update the value of AMI_state
update_state('Delete_AMI')
return f'Created new AMI: {image_id}'
# When InstanceId value does not exist
else:
return 'InstanceId not found in event data.'
elif ami_state == 'Delete_AMI':
# Check if image_id value exists
if image_id:
# Store the image_id before deletion
existing_ami_id = image_id
# Delete the existing AMI
delete_ami(existing_ami_id)
# Extract the InstanceId from the event dictionary
InstanceId = event.get('InstanceId')
# Call the create_ami function
new_ami_id = create_ami(InstanceId)
return f'Deleted existing AMI: {existing_ami_id} and created a new AMI: {new_ami_id}'
else:
return 'image_id is not set, cannot delete AMI'
else:
return 'Invalid state'
except Exception as e:
return str(e)
Conclusion
By using Python, AWS services like EC2, Lambda, DynamoDB, and careful state management, this script automates the creation and deletion of AMIs seamlessly. This can be integrated into a larger infrastructure to ensure that your EC2 instances are always based on the latest and most reliable images. If you want to add several images, you can simply modify the NoSQL AMI_ID attribute in the AMI_Table item to use a list type for its values, and then tweak the rest of the code.
Using NoSQL databases like DynamoDB is a powerful choice for storing state information in such automation scripts. This script is a clear example of how state sharing between functions and external systems can be effectively managed to automate complex workflows.
Check out Lab 4 in this repo to find the full codes implementation.
Top comments (0)