Automating AWS Image Management using Serverless, Python, and NoSQL
Tags: nosql, python, serverless, aws
Published: 2023-10-21 08:03:31

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:

```python
# 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.
```python

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
![AMI_Table](https://dev-to-uploads.s3.amazonaws.com/uploads/articles/iwiz3yrb0171q0gxgtzl.png)

### 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.

```python
# 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 `InstanceId` using the `ec2.create_image` function.
- The `current_time` is used to generate a unique name for the new AMI.
- The script updates the `image_id` global variable with the newly created AMI's ID.
- It stores the `image_id` value in the DynamoDB table to maintain the state.

```python
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.

```python
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.

```python
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.

```python
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](https://github.com/network-charles/AWS-Sys-Ops/tree/main) to find the full codes implementation.