ecs_tasks_def_cleanup.py

from __future__ import print_function
import os
import argparse
import boto3
import time
# Requries 'requests' to be bundled into zip.
#import requests

REGION = None
DRYRUN = None
DEFS_TO_KEEP = None


def initialize(regionname):
    global REGION
    global DRYRUN
    global DEFS_TO_KEEP
    global ecs_client

    ecs_client = boto3.client('ecs', region_name=regionname)

    REGION = os.environ.get('REGION', "None")
    DRYRUN = os.environ.get('DRYRUN', "false").lower()
    if DRYRUN == "false":
        DRYRUN = False
    else:
        DRYRUN = True
    DEFS_TO_KEEP = int(os.environ.get('DEFS_TO_KEEP', 10))


def handler(event, context):
    initialize(REGION)
    #
    # Requries 'requests' to be bundled into zip.
    #
    # if REGION == "None":
    #     partitions = requests.get("https://raw.githubusercontent.com/boto/botocore/develop/botocore/data/endpoints.json").json()['partitions']
    #     for partition in partitions:
    #         if partition['partition'] == "aws":
    #             for endpoint in partition['services']['ecs']['endpoints']:
    #                 discover_tasks_definitions(endpoint)
    # else:
    discover_tasks_definitions(REGION)


def discover_tasks_definitions(regionname):
    print("Discovering tasks definitions in "+regionname)

    for family in list_task_definition_families():
        tasksdef = []
        describetasksdef_paginator = ecs_client.get_paginator('list_task_definitions')
        for response_listtasksdefpaginator in describetasksdef_paginator.paginate(familyPrefix=family):
            for task_def in response_listtasksdefpaginator['taskDefinitionArns']:
                # if task_def not in active_tasksdef:
                tasksdef.append(task_def)
        if len(tasksdef) >DEFS_TO_KEEP:
            del tasksdef[-DEFS_TO_KEEP:]
        else:
            tasksdef = []

        print("Task definitions to remove in family %s: " % family)
        for task in tasksdef:
            deregister_task_def(task)


def deregister_task_def(arn):
    if not DRYRUN:
        print("Deregistering: %s" % arn)
        response = ecs_client.deregister_task_definition(
            taskDefinition=arn
        )
        time.sleep(0.3)
    else:
        print("[DRYRUN] Deregistering: %s" % arn)
        time.sleep(0.4)
    pass


def list_task_definition_families():
    task_definition_families = []
    list_task_def_fam_pages = ecs_client.get_paginator('list_task_definition_families')
    for page in list_task_def_fam_pages.paginate():
        task_definition_families.extend(page['families'])
    return task_definition_families


# Below is the test harness
if __name__ == '__main__':
    request = {"None": "None"}
    parser = argparse.ArgumentParser(description='Deregister old ECS Task definitions')
    parser.add_argument('-dryrun', help='Prints the tasks definitions to be deactivated without deactivating them', default='true', action='store', dest='dryrun')
    parser.add_argument('-defstokeep', help='Number of task definitions to keep', default='10', action='store', dest='defstokeep')
    parser.add_argument('-region', help='ECR/ECS region', default=None, action='store', dest='region')

    args = parser.parse_args()
    if args.region:
        os.environ["REGION"] = args.region
    else:
        os.environ["REGION"] = "None"
    os.environ["DRYRUN"] = args.dryrun.lower()
    os.environ["DEFS_TO_KEEP"] = args.defstokeep
    handler(request, None)