Generate a list of EC2 instances and prompt to SSH into one

manage_ec2.py

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

'''
This script will list all hosts in an ec2 region and prompt you to connect
to them.

It expects the file .boto to exist in your home directory with contents
as follows:

[Credentials]
aws_access_key_id = <AWS_ACCESS_KEY_ID>
aws_secret_access_key = <AWS_SECRET_ACCESS_KEY>
'''

import boto.ec2
import os
import sys
from itertools import imap
from subprocess import call
import yaml
import datetime
import argparse
import calendar


def parse_args():
    '''
    Get arguments from the command like with argparse
    '''
    parser = argparse.ArgumentParser(
        description='Generate a list of ec2 hosts and select one to connect to'
        )
    parser.add_argument(
        '--login_name', '-l', required=False, help='Login Name')
    parser.add_argument(
        '--port', '-p', required=False, help='SSH Port')
    parser.add_argument(
        '--identity_file', '-i', required=False, help='SSH Identity File')
    parser.add_argument(
        '--force_download', '-F', required=False,
        help='Force Download', action='store_true')
    parser.add_argument(
        '--file', '-f', required=False, help='Config File',
        default=os.path.expanduser('~') + '/.ec2_addr_cache')
    parser.add_argument(
        '--region', '-r', required=False, help='EC2 Region',
        default='us-west-2')
    parser.add_argument(
        '--perpage', '-P', required=False,
        help='Hosts to display per page (default 25, 0 to display all)',
        default=25, type=int)
    args = parser.parse_args()
    return args


def generate_hosts(filename, args):
    '''
    Generate a list of hoosts in the ec2 region selected and save
    them to a yaml file
    '''
    hosts = {}
    conn = boto.ec2.connect_to_region(args.region)
    reservations = conn.get_all_reservations()
    for reservation in reservations:
        instances = reservation.instances
        for instance in instances:
            if instance.state == 'running':
                if 'Name' in instance.tags:
                    hosts.update(
                        {
                            instance.tags['Name']:
                            {
                                'ip': instance.private_ip_address,
                                'key_name': instance.key_name
                            }
                        }
                    )
    with open(filename, 'w') as outfile:
        outfile.write(yaml.dump(hosts, default_flow_style=False))
    return hosts


def read_hosts(filename):
    '''
    Load saved ec2 host file
    '''
    hostsfile = open(filename, 'r')
    hosts = yaml.load(hostsfile.read())
    return hosts


def hosts_dict(args):
    '''
    Obtain the hosts dict by either reading it if it is current from
    yaml or by generating it via an ec2 api call
    '''
    ec2_addr_cache_file = args.file
    if os.path.isfile(ec2_addr_cache_file) and not args.force_download:
        file_age = int(os.path.getmtime(ec2_addr_cache_file))
        now = int(
            calendar.timegm(
                datetime.datetime.now().timetuple()))
        if (file_age + 86400) < now:
            hosts = generate_hosts(ec2_addr_cache_file, args)
        else:
            hosts = read_hosts(ec2_addr_cache_file)
    else:
        hosts = generate_hosts(ec2_addr_cache_file, args)

    return hosts


def parse_input(args, startindex, perpage, hosts, hostlist):
    '''
    Parse the entered text and take action upon it

    q = quit
    n = next page
    p = previous page
    <number> = ssh to server <number>
    '''
    selected_index = raw_input(
        "Please enter the host to connect to or (n)ext, (p)revious, (q)uit: ")
    if isinstance(selected_index, basestring) and selected_index == 'q':
        sys.exit(0)
    elif isinstance(selected_index, basestring) and selected_index == 'n':
        args.force_download = False
        list_servers(args, startindex + perpage, perpage)
    elif isinstance(selected_index, basestring) and selected_index == 'p':
        args.force_download = False
        list_servers(args, startindex - perpage, perpage)
    else:
        try:
            selected_index = int(selected_index) - 1
            # Indexes start at 0, so subtract one from here and length tests.
            if not 0 <= selected_index <= len(hostlist) - 1:
                raise TypeError('Range not valid!')
        except (TypeError, ValueError):
            print "Invalid entry, try again!"
            args.force_download = False
            list_servers(args, startindex, perpage)
        else:
            host_ip = hosts[hostlist[selected_index]]['ip']
            ssh = ['ssh']
            if 'port' in args and args.port is not None:
                ssh.extend(['-p', str(args.port)])
            if 'login_name' in args and args.login_name is not None:
                ssh.extend(['-l', args.login_name])
            if 'identity_file' in args and args.identity_file is not None:
                ssh.extend(['-i', args.identity_file])
            ssh.extend([host_ip])
            call(ssh)


def list_servers(args, startindex, perpage):
    '''
    List the servers stored in the ec2 region file or regenerate the file if
    older than 24 hours or forced via the -F True flag
    '''
    hosts = hosts_dict(args)
    length = max(imap(len, hosts))
    hostcount = len(hosts)
    if perpage <= 0:
        perpage = hostcount
    countwidth = len(str(hostcount))
    hostlist = sorted(hosts, key=lambda s: s.lower())
    if startindex < 0:
        startindex = 0
    elif startindex >= hostcount:
        startindex = hostcount - perpage
    for index in range(startindex, startindex + perpage):
        if index >= hostcount:
            break
        host = hostlist[index]
        # Start the list with 1 instead of 0 for humans
        printindex = index + 1
        print "%s) %s | %s | %s" % (
            str(printindex).rjust(countwidth),
            host.ljust(length), hosts[host]['ip'].ljust(15),
            hosts[host]['key_name'])
    print "Displaying hosts %i to %i of %i" % (
        startindex + 1, printindex, hostcount)
    parse_input(args, startindex, perpage, hosts, hostlist)

def main():
    '''
    Main run function
    '''
    args = parse_args()
    list_servers(args, 0, args.perpage)


if __name__ == '__main__':
    sys.exit(main())