Inquirer script to SSH into instances

sc-ec2-ssh.py

#!/usr/bin/env python

# %%SC_COPYRIGHT%%
# Author: Vinay Dandekar

import argparse
import logging
import signal
import subprocess
import sys
from operator import itemgetter

import inquirer
from inquirer.render.console import ConsoleRender
# noinspection PyProtectedMember
from inquirer.render.console._list import List

import sc
from sc.utils import boto_connection_factory as factory


class OtherColorList(List):
    def get_options(self):
        choices = self.question.choices

        for choice in choices:
            selected = choice == choices[self.current]

            if selected:
                color = self.terminal.black_on_yellow
                symbol = '>'
            else:
                color = self.terminal.normal
                symbol = ' '
            yield choice, symbol, color


class OtherListConsoleRender(ConsoleRender):
    def render_factory(self, question_type):
        if question_type != 'list':
            # noinspection PyUnresolvedReferences
            # pylint: disable=bad-super-call, no-member
            return super(ConsoleRender, self).render_factory(question_type)
        return OtherColorList


def __setup_log():
    """."""
    formatter = logging.Formatter(
        '%(asctime)s - %(name)s - %(levelname)s - %(message)s')
    ch = logging.StreamHandler()
    ch.setLevel(logging.INFO)
    ch.setFormatter(formatter)
    logger.addHandler(ch)
    sc.set_stream_logger('sc', logging.INFO)


def signal_handler():
    sys.exit(0)


def main(query_str):
    """
    :type query_str: str
    :param query_str: the query to filter on
    """
    ec2 = factory.get_ec2_resource()
    instances = ec2.instances.filter(
        Filters=[{'Name': 'tag:Name', 'Values': ['*' + query_str + '*']}])
    names = []
    for instance in instances:
        for tag in instance.tags:
            if tag['Key'] == 'Name':
                if instance.state['Name'] == 'running':
                    name = tag['Value']
                    names.append(
                        {
                            'name': name,
                            'instance': instance,
                            'ip_address': instance.private_ip_address
                        })

    if len(names) > 0:
        names.sort(key=itemgetter('name'))
        name_length = max(len(name['name']) for name in names) + 2
        choices = []
        for instance_name in names:
            choices.append(
                ('%s\t%s' % (instance_name['name'], instance_name['instance'].private_ip_address))
                .expandtabs(name_length) + '\t' + str(instance_name['instance'].launch_time))

        questions = [
            inquirer.List('instance',
                          message="Select an EC2 instance",
                          choices=choices
                          ),
        ]
        answers = inquirer.prompt(questions, render=OtherListConsoleRender())

        if answers:
            for name in names:
                if name['name'] in answers['instance']:
                    # noinspection PyBroadException
                    try:
                        print 'SSH to %s (%s)\n' % \
                              (name['name'], name['instance'].private_ip_address)
                        subprocess.check_call(["ssh",
                                               "-A",
                                               name['instance'].private_ip_address])
                    except:
                        pass
                    break
    else:
        print 'No instances found'


if __name__ == '__main__':
    logger = logging.getLogger(__file__)
    logger.setLevel(logging.INFO)
    __setup_log()

    signal.signal(signal.SIGINT, signal_handler)

    parser = argparse.ArgumentParser(
        description='list inventory for autoscaling groups')
    parser.add_argument('query', help="query to filter on")
    args = parser.parse_args()
    query = args.query

    if not query:
        query = ''

    main(query)