Create a gist now

Instantly share code, notes, and snippets.

What would you like to do?
Generate a list of EC2 instances and prompt to SSH into one
#!/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())
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment