report-marketplace-amis-by-account.py

#!/usr/bin/env -S uv run
# /// script
# dependencies = [
#   "boto3",
# ]
# ///
"""
Report marketplace AMIs across all of the accounts in an Organization

Assumes AWS credentials for an account which has permission to query a Config
aggregator, Organizations, and EC2.
"""

import argparse
from functools import cache
import json

import boto3


AWS_CONFIG = boto3.client("config")
EC2 = boto3.client("ec2")
ORGANIZATIONS = boto3.client("organizations")


@cache
def get_account_name(account_id):
    r = ORGANIZATIONS.describe_account(AccountId=account_id)
    return r["Account"]["Name"]


@cache
def describe_image(image_id):
    r = EC2.describe_images(ImageIds=[image_id])

    if images := r["Images"]:
        return images[0]
    else:
        return None


def get_marketplace_amis(config_aggregator_name):
    ami_query = AWS_CONFIG.select_aggregate_resource_config(
        ConfigurationAggregatorName=config_aggregator_name,
        Expression="SELECT accountId, configuration.imageId, COUNT(*) WHERE resourceType = 'AWS::EC2::Instance' AND configuration.productCodes.productCodeType = 'marketplace' GROUP BY accountId, configuration.imageId",
    )

    amis_by_account = {}

    for raw_result in ami_query["Results"]:
        result = json.loads(raw_result)
        account_id = result["accountId"]
        image_id = result["configuration"]["imageId"]
        amis_by_account.setdefault(account_id, {})[image_id] = result["COUNT(*)"]

    return amis_by_account


def main(config_aggregator_name):
    amis_by_account = get_marketplace_amis(config_aggregator_name)

    print(
        "account_id",
        "account_name",
        "instance_count",
        "image_id",
        "image_name",
        sep="\t",
    )

    for account_id, images in sorted(amis_by_account.items(), key=lambda i: i[0]):
        account_name = get_account_name(account_id)

        images = sorted(images.items(), key=lambda i: (i[1], i[0]), reverse=True)

        for image_id, instance_count in images:
            image_info = describe_image(image_id)
            if image_info:
                image_name = image_info["Name"]
            else:
                image_name = ""
            print(
                account_id, account_name, instance_count, image_id, image_name, sep="\t"
            )


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description=__doc__.strip())
    parser.add_argument(
        "config_aggregator",
        help="Name of the AWS Config aggregator to query",
    )

    args = parser.parse_args()

    main(args.config_aggregator)