ecr.py

#!/usr/bin/env python
'''ECR Commands'''

import itertools
import os
import sys
import boto3
import humanize
import shellish
from . import base
from . import cmd
from . import IMAGES
from . import REGISTRIES
from . import REPOSITORIES
from concurrent import futures

AWS_REGIONS = [
    'us-east-2',
    'us-west-2',
]
AWS_REGION = os.environ.get('AWS_REGION', 'us-east-2')


class ECRCommand(base.HostingCommand):
    '''Manage ECR Images and Tags'''
    name = 'ecr'

    def setup_args(self, parser):
        self.add_subcommand(PushCommand)
        self.add_subcommand(CatalogCommand)
        self.add_subcommand(ListImagesCommand)
        self.add_subcommand(RemoveTagCommand)


class PushCommand(base.HostingCommand):
    '''Push Current Images to AWS ECR'''

    name = 'push'

    def setup_args(self, paser):
        self.add_argument('--image-tag', help='Override Image Tag')
        self.add_argument('--verbose', '-v',
                          action='store_true',
                          help='Turn on verbose output')

    def run(self, args):
        git_sha = os.environ.get('GIT_SHORT_REV')
        image_tags = [('g-%s' % git_sha)]
        if args.image_tag:
            image_tags.append(args.image_tag)

        for region in AWS_REGIONS:
            for tag in image_tags:
                push_odyssey_images(tag, region, args.verbose)


class CatalogCommand(base.HostingCommand):
    '''Display List of Images and Tags in ECR'''

    name = 'catalog'

    def setup_args(self, parser):
        self.add_argument('--region', help='Only list images from this region')

    def run(self, args):
        display_image_catalog(args.region)


class ListImagesCommand(base.HostingCommand):
    '''Display List of Images and Tags in ECR'''

    name = 'list'

    def setup_args(self, parser):
        self.add_argument('--region', help='Only list images from this region')

    def run(self, args):
        display_image_catalog(args.region)


class RemoveTagCommand(base.HostingCommand):
    '''Remove Docker tag from images in ECR'''

    name = 'rm'

    def setup_args(self, parser):
        self.add_argument('tag_name', help='Name of tag to delete')
        self.add_argument('--verbose', '-v', action='store_true',
                          help='Print verbose output')

    def confirm_action(self, message):
        confirm = input(message)
        if confirm != 'yes':
            raise SystemExit('Aborted')

    def run(self, args):
        self.confirm_action('Are you sure you want to remove image tag "%s" '
                            '[yes|NO]: ' % args.tag_name)

        remove_image_tag(args.tag_name, verbose=args.verbose)


def display_image_catalog(region=None):
    regions = AWS_REGIONS
    if region in AWS_REGIONS:
        regions = [region]
    for region in regions:
        print_image_catalog_table(region)


def print_image_catalog_table(region):
    catalog = build_catalog_table(region)
    for repository in catalog.keys():
        table_config = {
            'headers': [
                'Image Tags',
                'Pushed At',
                'Digest',
                'Size',
            ],
            'columns': [
                0.25,
                0.1,
                None,
                0.1,
            ],
            'title': 'Images from %s/%s' % (region, repository)
        }
        t = shellish.Table(**table_config)
        rows = []
        for image in catalog[repository]:
            rows.append((
                ','.join(image['imageTags']),
                humanize.naturaldate(image['imagePushedAt']),
                image['imageDigest'],
                humanize.naturalsize(image['imageSizeInBytes']),
            ))
        t.print(rows)
        t.print_footer('')
        t.close()


def build_catalog_table(region):
    '''build table of images for each repository'''
    table = {}
    for repository in REPOSITORIES:
        table[repository] = list(get_ecr_images(region, repository))

    return table


def get_ecr_images(region, repository_name):
    def _get_ecr_images_(region, repository_name):
        client = boto3.client('ecr', region_name=region)
        paginator = client.get_paginator('describe_images')
        pages = paginator.paginate(
            repositoryName=repository_name,
            filter={
                'tagStatus': 'TAGGED'
            }
        )
        for page in pages:
            yield from page['imageDetails']

    return itertools.chain(_get_ecr_images_(region, repository_name))


def push_odyssey_images(tag, region, verbose):
    registry = REGISTRIES[region]
    cmd('$(aws --region %s ecr get-login --no-include-email)' % region,
        verbose=verbose)
    with futures.ThreadPoolExecutor(max_workers=10) as pool:
        push_futures = []
        for image, remote_image in IMAGES:
            def fn(_image=image, _remote_image=remote_image):
                ecr_tag = '%s/%s:%s' % (registry,
                                        _remote_image,
                                        tag)
                cmd('docker tag', _image, ecr_tag, verbose=verbose)
                try:
                    cmd('docker push', ecr_tag, verbose=verbose)
                finally:
                    cmd('docker rmi', ecr_tag, verbose=verbose)
            push_futures.append(pool.submit(fn))
        for f in futures.as_completed(push_futures):
            try:
                f.result()  # check for exceptions
            except Exception as e:
                raise SystemExit(str(e))


def remove_image_tag(tag_name, **kwargs):
    def __remove_tag_from_ecr__(tag_name, region_name):
        client = boto3.client('ecr', region_name=region_name)
        for repository in REPOSITORIES:
            response = client.batch_delete_image(
                repositoryName=repository,
                imageIds=[
                    {
                        'imageTag': tag_name,
                    }
                ]
            )
            for image_id in response['imageIds']:
                shellish.vtmlprint(
                    '<red>Deleted image: %s</red>' % (
                        image_id['imageTag']))
            for failure in response['failures']:
                shellish.vtmlprint(
                    '<red>Unable to delete: "%s": %s</red>' % (
                        failure['imageId']['imageTag'],
                        failure['failureReason']))

    verbose = kwargs.get('verbose', False)
    with futures.ThreadPoolExecutor(max_workers=10) as pool:
        push_futures = []

        for region in AWS_REGIONS:
            push_futures.append(pool.submit(__remove_tag_from_ecr__,
                                            tag_name,
                                            region))
        for f in futures.as_completed(push_futures):
            try:
                f.result()
            except Exception as ex:
                if verbose:
                    print(str(ex), file=sys.stderr)