#!/usr/bin/env python

# Clean up tagged/untagged ECR images by age
# You need to set AWS_DEFAULT_REGION environment variable if you're deleting from an ECR
# repository outside your configured default region
#
# python cleanup_ecr.py --registry-id 123456789000 --repository-name ecr-foobar \
#                       --untagged-age 1 --tagged-age 30 \
#                       --tag-prefix 3.8 --dry-run

import argparse
import datetime
import logging
import operator
import sys
import boto3


logging.basicConfig(format='[%(asctime)-15s] [%(module)s] %(levelname)s %(message)s', level='INFO')
logger = logging.getLogger(__name__)

client = boto3.client('ecr')
paginator = client.get_paginator('describe_images')

UNTAGGED_AGE_DAYS, TAGGED_AGE_DAYS = 3, 30  # default image ages
BATCH_DELETE_SIZE = 100  # maximum allowed


def get_image_details(registry_id, repository_name):
    '''Return a 2-tuple (untagged, tagged) of image details sorted by push date'''
    logger.info('Getting UNTAGGED image details')
    untagged_images = paginator.paginate(
        registryId=registry_id, repositoryName=repository_name, filter={'tagStatus': 'UNTAGGED'},
    ).build_full_result()['imageDetails']

    logger.info('Getting TAGGED image details')
    tagged_images = paginator.paginate(
        registryId=registry_id, repositoryName=repository_name, filter={'tagStatus': 'TAGGED'},
    ).build_full_result()['imageDetails']

    return (
        sorted(untagged_images, key=operator.itemgetter('imagePushedAt')),
        sorted(tagged_images, key=operator.itemgetter('imagePushedAt')),
    )


def get_image_digests_to_delete(untagged, tagged, untagged_age_days, tagged_age_days, tag_prefix):
    '''Return a list of image digests for images to delete'''
    digests = []
    now = datetime.datetime.now().astimezone()  # ECR returns tz-aware as well
    for image in untagged:
        image_age = now - image['imagePushedAt']
        if image_age.days > untagged_age_days:
            logger.info('UNTAGGED image age %s', image_age)
            digests.append(image['imageDigest'])

    for image in tagged:
        image_age = now - image['imagePushedAt']
        if image_age.days > tagged_age_days:
            if tag_prefix is None:
                logger.info('TAGGED image age %s - %s', image_age, image['imageTags'])
                digests.append(image['imageDigest'])
            elif len(image['imageTags']) == 1 and any(tag_prefix in s for s in image['imageTags']):
                logger.info('TAGGED image age %s - %s', image_age, image['imageTags'])
                digests.append(image['imageDigest'])
    return digests


def delete_images(registry_id, repository_name, digests):
    '''Delete ECR images for given image digests'''
    while digests:
        client.batch_delete_image(
            registryId=registry_id, repositoryName=repository_name,
            imageIds=[{'imageDigest': image_digest} for image_digest in digests[:BATCH_DELETE_SIZE]],
        )
        digests = digests[BATCH_DELETE_SIZE:]


if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='Delete ECR images by age')
    parser.add_argument('--registry-id', type=str, required=True, metavar='REGISTRY_ID')
    parser.add_argument('--repository-name', type=str, required=True, metavar='REPOSITORY_NAME')
    parser.add_argument(
        '--untagged-age', type=int, required=False, default=UNTAGGED_AGE_DAYS, metavar='DAYS',
        help='age threshold for UNTAGGED images in days (default: {})'.format(UNTAGGED_AGE_DAYS),
    )
    parser.add_argument(
        '--tagged-age', type=int, required=False, default=TAGGED_AGE_DAYS, metavar='DAYS',
        help='age threshold for TAGGED images in days (default: {})'.format(TAGGED_AGE_DAYS),
    )
    parser.add_argument(
        '--tag-prefix', type=str, required=False, default=None, metavar='TAG_PREFIX',
        help='prefix of the image tag (default: None)',
    )

    parser.add_argument('--dry-run', action='store_true')
    args = parser.parse_args()

    untagged, tagged = get_image_details(args.registry_id, args.repository_name)
    digests = get_image_digests_to_delete(untagged, tagged, args.untagged_age, args.tagged_age, args.tag_prefix)
    if not digests:
        logger.info('No images to delete')
        sys.exit(0)

    logger.info('Retrieved %d image digests for deletion', len(digests))
    if args.dry_run:
        logger.info('Exiting before deletion, --dry-run is set')
    else:
        logger.info('Deleting %d images from registry=%s repository=%s', len(digests), args.registry_id, args.repository_name)
        delete_images(args.registry_id, args.repository_name, digests)