ECS task definition rollback script

ecs-rollback.py

#!/usr/bin/env python
# -*- coding: utf-8 -*-

import re
import json
import argparse
from datetime import datetime
from subprocess import Popen, PIPE


def execute(command):
    child = Popen(command, stdout=PIPE, stderr=PIPE, shell=True)
    if child.wait():
        raise Exception('\n'.join(child.communicate()).rstrip())
    else:
        return '\n'.join(child.communicate()).rstrip()


def _get_td_image_tag(task_definition):
    image = re.search(
        '/([\w-]*(/)?([\w-]*)?):',
        task_definition['image']).groups()[0]
    tag = re.search(':(\w*)$', task_definition['image']).groups()[0]
    return image, tag


def get_task_defs(service):
    """ Retrieve task definitions """
    srv = execute('aws ecs describe-services '
        '--region {region} --cluster {cluster} --services {service}'
        .format(region=region, cluster=cluster, service=service))
    try:
        task = json.loads(srv)['services'][0]['taskDefinition']
    except Exception, e:
        raise Exception('Service {0} not found!\nError: {1}'
                .format(service, e))
    task = task.split('/')[-1]
    td = execute('aws ecs describe-task-definition '
            '--region {region} --task-definition {td}'
            .format(region=region, td=task))
    return json.loads(td)['taskDefinition']['containerDefinitions']


def get_ecr_tags(service):
    """ Get all tags for all ECR images for given service """
    cd = get_task_defs(service)
    images = {
        re.search(
            '/([\w-]*(/)?([\w-]*)?):', i['image'])
            .groups()[0]: re.search(
                '([a-z0-9-]*)\.amazonaws', i['image'])
                .groups()[0] for i in cd}
    tags = {}
    for i in images:
        old_image = execute('aws ecr describe-images '
            '--region {region} --repository-name {repo}'
            .format(region=images[i], repo=i))
        revs = {int(i['imagePushedAt']):i['imageTags']
                for i in json.loads(old_image)['imageDetails']}
        tags[i] = revs
    return tags


def get_previous_tag(service):
    """ Retrieve previous tag from ECR images for a given service """
    prevs = {}
    images = get_ecr_tags(service)
    for i in images:
        prevs[i] = images[i][sorted(images[i])[-2]][0]
    return prevs


def rollback_tag(service, tag=None):
    prev_tags = get_previous_tag(service)
    # Override tags if provided
    if tag:
        prev_tags = {i:tag for i in prev_tags}
    task_defs = get_task_defs(service)
    new_task_defs = []
    for i in task_defs:
        image, curr_tag = _get_td_image_tag(i)
        curr_tag = re.search(':(\w*)$', i['image']).groups()[0]
        if curr_tag == prev_tags[image]:
            raise Exception('Service {0} is already using tag {1}'
                    .format(service, image))
        new_image = re.sub(
            ':\w*$', ':{0}'.format(prev_tags[image]), i['image'])
        i['image'] = new_image
        new_task_defs.append(i)
    out = execute('aws ecs register-task-definition '
            '--region {region} '
            '--family {service} '
            '--container-definitions \'{tasks}\''
            .format(
                region=region,
                service=service,
                tasks=json.dumps(new_task_defs)))
    new_td = json.loads(out)['taskDefinition']['taskDefinitionArn']
    new_srv = execute('aws ecs update-service '
            '--region {region} '
            '--cluster {cluster} '
            '--service {service} '
            '--task-definition {td}'
            .format(
                region=region,
                cluster=cluster,
                service=service,
                td=new_td))
    return new_srv


def set_arguments():
    parser.add_argument(
            '-r',
            '--region',
            help='AWS region where service is, default: eu-west-1',
            type=str,
            default='eu-west-1')
    parser.add_argument(
            '-c',
            '--cluster',
            help='ECS cluster name, default: aws-eu-ecs-prod-apps',
            type=str,
            default='aws-eu-ecs-prod-apps')
    subparsers = parser.add_subparsers(dest='sub')
    parent = argparse.ArgumentParser(add_help=False)
    rb = subparsers.add_parser(
            'rollback', parents=[parent],
            help='rollback ECS service')
    rb.add_argument(
            '-t',
            '--tag',
            help='''ECR image tag for task definition(s),
            if not provided previous will be used''',
            type=str)
    rb.add_argument(
            'service',
            help='ECS service name')
    li = subparsers.add_parser(
            'images', parents=[parent],
            help='list ECS service ECR images')
    li.add_argument(
            'service',
            help='ECS service name')
    logs = subparsers.add_parser(
            'logs', parents=[parent],
            help='show logs for ECS service')
    logs.add_argument(
            '-l',
            '--length',
            type=int,
            default=10,
            help='Lines of log to output, default: 10')
    logs.add_argument(
            'service',
            help='ECS service name')
    srv = subparsers.add_parser(
            'services', parents=[parent],
            help='list ECS services')


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    set_arguments()
    args = parser.parse_args()
    region = args.region
    cluster = args.cluster
    if args.sub == 'rollback':
        if args.tag:
            tag = args.tag
        else:
            tag = 'previous'
        print('Rolling service {0} to {1} tag'
                .format(args.service, tag))
        rollback_tag(args.service, args.tag)

    elif args.sub == 'images':
        # Retrieve map of deployed tags
        task_defs = get_task_defs(args.service)
        curr_tags = {}
        for td in task_defs:
            i, t = _get_td_image_tag(td)
            curr_tags[i] = t
        images = get_ecr_tags(args.service)
        for img in images:
            print('{0} {1}'.format('-' * 5, img))
            for date in sorted(images[img], reverse=True):
                for tag in sorted(images[img][date], reverse=True):
                    if tag == curr_tags[img]:
                        current = '\tDeployed'
                    else:
                        current = ''
                    print('Tag: {t}\tPushed: {d}{current}'
                        .format(
                            t=tag,
                            current=current,
                            d=datetime.fromtimestamp(
                                date).strftime('%Y-%m-%d %H:%M:%S')))

    elif args.sub == 'logs':
        srv = execute(
            'aws ecs describe-services '
            '--region {region} '
            '--cluster {cluster} '
            '--services {service}'
            .format(
                region=region,
                cluster=cluster,
                service=args.service))
        srv = json.loads(srv)
        try:
            logs = srv['services'][0]['events']
        except Exception, e:
            raise Exception('Service {0} not found!\nError: {1}'
                    .format(args.service, e))
        for log in logs[:args.length]:
            print('{0}  {1}'.format(
                datetime.fromtimestamp(int(log['createdAt']))
                    .strftime('%Y-%m-%d %H:%M:%S'),
                log['message']))

    elif args.sub == 'services':
        srv = execute('aws ecs list-services '
                '--region {region} '
                '--cluster {cluster} '
                .format(
                    region=region,
                    cluster=cluster))
        for service in sorted(json.loads(srv)['serviceArns']):
            print(service.split(':service/')[-1])