Quick attempt at building an AMI lineage

lineage.py

#!/usr/bin/env python3

import boto3
from botocore.exceptions import ClientError
import argparse
from tabulate import tabulate  # Import tabulate


def get_ami_lineage(image_id, region, lineage=None):
    """
    Recursively retrieves the lineage of AMIs starting from the given image_id,
    including the image name and owner account ID.

    Parameters:
        image_id (str): The AMI ID to start from.
        region (str): The region where the AMI is located.
        lineage (list): Accumulator for storing the lineage.

    Returns:
        list: A list of dictionaries containing ImageId, Region, Name, and OwnerId.
    """
    if lineage is None:
        lineage = []

    # Create an EC2 client for the specified region
    client = boto3.client("ec2", region_name=region)

    try:
        # Describe the image to get its details
        response = client.describe_images(ImageIds=[image_id])
        images = response.get("Images", [])

        if images:
            image = images[0]
            source_image_id = image.get("SourceImageId")
            source_image_region = image.get("SourceImageRegion")
            image_name = image.get("Name", "Unnamed AMI")
            owner_id = image.get("OwnerId", "Unknown")

            # Append the current image details to the lineage
            lineage.append(
                {
                    "ImageId": image_id,
                    "Region": region,
                    "Name": image_name,
                    "OwnerId": owner_id,
                }
            )

            # If there's a source image, recurse into it
            if source_image_id and source_image_region:
                return get_ami_lineage(source_image_id, source_image_region, lineage)
        else:
            print(f"No details found for image {image_id} in region {region}.")
    except ClientError as e:
        print(f"Error retrieving image {image_id} in region {region}: {e}")

    # Return the accumulated lineage
    return lineage


if __name__ == "__main__":
    # Set up command-line argument parsing
    parser = argparse.ArgumentParser(description="Traverse the lineage of an AWS AMI.")
    parser.add_argument("-i", "--image-id", required=True, help="The starting AMI ID.")
    parser.add_argument(
        "-r", "--region", required=True, help="The region of the starting AMI."
    )

    args = parser.parse_args()

    starting_image_id = args.image_id
    starting_region = args.region

    # Get the lineage
    lineage = get_ami_lineage(starting_image_id, starting_region)

    # Reverse the lineage to show the oldest ancestor first
    lineage = list(reversed(lineage))

    # Prepare data for tabulate
    headers = ["ImageId", "Region", "Name", "OwnerId"]
    table_data = [
        [ami["ImageId"], ami["Region"], ami["Name"], ami["OwnerId"]] for ami in lineage
    ]

    print("\nAMI Lineage:")
    print(tabulate(table_data, headers=headers, tablefmt="grid"))