Tag instances automatically

ec2-autotagger.py

#!/usr/bin/env python
# Author: Vinay Dandekar

"""
sc-ec2-autotagger
~~~~~~~~~~~~~

Tags the EC2 instance in a autoscaling group.

"""
import argparse
import logging
import string
from operator import attrgetter

from sc.utils import asg_utils
from sc.utils import boto_connection_factory as factory


def __setup_log(debug):
    """Sets up the logger.

    :type debug: bool
    :param debug: debug flag
    """
    formatter = logging.Formatter(
        '%(asctime)s - %(name)s - %(levelname)s - %(message)s')
    ch = logging.StreamHandler()
    if debug:
        logger.setLevel(logging.DEBUG)
    else:
        logger.setLevel(logging.INFO)
    ch.setFormatter(formatter)
    logger.addHandler(ch)


def _get_asg_data(instance_id):
    """Get the EC2 autoscaling group name tag, instances, and size based on the EC2 instance ID.

    :type instance_id: str
    :param instance_id: EC2 instance ID
    :rtype: (list, int, str)
    :return: autoscaling group instances, max # of instances in the autoscaling group, the tag prefix
    """
    asg = asg_utils.get_instance_asg(instance_id)
    max_size = asg['MaxSize']
    tag_prefix = ''
    for tag in asg['Tags']:
        if tag['Key'] == 'Name':
            tag_prefix = tag['Value']
            logger.info('Autoscaling group name: ' + tag_prefix)
            break
    return asg['Instances'], max_size, tag_prefix


def _tag_instance(instance_id, asg_instances, max_size, tag_prefix):
    """Tag the EC2 instance.

    :type instance_id: str
    :param instance_id: EC2 instance ID
    :type asg_instances: list
    :param asg_instances: autoscaling group instances
    :type max_size: int
    :param max_size: max # of instances in the autoscaling group
    :type tag_prefix: str
    :param tag_prefix: the autoscaling group tag prefix
    """
    ec2 = factory.get_ec2_resource()

    unused_suffix_nums = []
    for i in range(1, max_size + 1):
        unused_suffix_nums.append(i)

    instances = []
    for instance in asg_instances:
        instance = ec2.Instance(instance['InstanceId'])
        if instance.state['Name'] == 'running':
            instances.append(instance)

    tagged_instances = []
    for instance in instances:
        if instance.tags:
            for tag in instance.tags:
                if tag['Key'] == 'Name':
                    name = tag['Value']
                    str_number = string.replace(name, tag_prefix, '')
                    if str_number != '':
                        logger.info('Found instance with name tag: %s' % name)
                        str_number = string.replace(str_number, '-', '')
                        unused_suffix_nums.remove(int(str_number))
                        tagged_instances.append(instance)

    logger.debug('Tagged instances: %s' % str(tagged_instances))
    untagged_instances = [item for item in instances if item not in set(tagged_instances)]

    # Sort the untagged_instances by instance_id and secondary sort by launch_time
    untagged_instances = sorted(sorted(untagged_instances, key=attrgetter('instance_id')),
                                key=attrgetter('launch_time'))

    for instance in untagged_instances:
        logger.debug('Instance ID: %s, Launch time: %s' % (instance.instance_id, instance.launch_time))

    instance_tag = tag_prefix
    for index, instance in enumerate(untagged_instances):
        if instance.instance_id == instance_id:
            num = unused_suffix_nums[index]
            instance_tag += '-%02d' % num
            logger.info('Instance (%s) Name: %s' % (instance_id, instance_tag))
            instance.create_tags(Tags=[{'Key': 'Name', 'Value': instance_tag}])
            break


def main():
    """."""
    parser = argparse.ArgumentParser()
    parser.add_argument('--debug',
                        help="use for debug output",
                        action="store_true")
    args = parser.parse_args()
    debug = args.debug
    __setup_log(debug)

    # Get instance ID
    instance_id = factory.get_instance_id()

    asg_instances, max_size, tag_prefix = _get_asg_data(instance_id)
    _tag_instance(instance_id, asg_instances, max_size, tag_prefix)


if __name__ == '__main__':
    logger = logging.getLogger(__file__)
    main()