ec2-sshconfig-generate.py

#!/usr/bin/env python

import os
import sys
import getopt

from sys import stderr

import boto.ec2


_ec2 = None

def short_usage():
    print >>stderr, """Usage: ec2-host [-k KEY] [-s SECRET] [--region REGION] [--user USERNAME] [--port SSHPORT]
    ec2-sshconfig-generate --user deploy --port 2222 --region us-west-1 >> ~/.ssh/config
Try `ec2-sshconfig-generate --help' for more information."""


def full_usage():
    print >>stderr, """Usage: ec2-sshconfig-generate [-k KEY] [-s SECRET] [--region REGION] [--username USER] [--port PORT] 
Prints server host name.

      --help                 display this help and exit
  -k, --aws-key KEY          Amazon EC2 Key, defaults to ENV[AWS_ACCESS_KEY_ID]
  -s, --aws-secret SECRET    Amazon EC2 Secret, defaults to ENV[AWS_SECRET_ACCESS_KEY]
  --region REGION            Amazon EC2 Region, defaults to us-east-1 or ENV[AWS_EC2_REGION]
  --user USERNAME            Amazon EC2 Instance UserName, defaults to ubuntu
  --port SSHPORT             Amazon EC2 Instance SSH Port, defaults to 22 """


def generate_host_map(connection):
    host_map = {}
    for reservation in connection.get_all_instances():
        for instance in reservation.instances:
            if instance.public_dns_name:
                host_map[instance.tags.get('Name')] = instance.public_dns_name
    return host_map

def create_ssh_config(connection, default_port=None, default_user=None):
    host_map = generate_host_map(connection)

    host_file = []
    host_file.append('# Add this to your ~/.ssh/config file')

    for host, dns_name in host_map.iteritems():
        host_file.append('Host %s' % (host))
        host_file.append('\tHostName %s' % (dns_name))
        if default_user:
            host_file.append('\tUser %s' % (default_user))
        if default_port:
            host_file.append('\tPort %s' % (default_port))
        host_file.append('')

    return '\n'.join(host_file)


def main(argv):
    global user, port, _ec2
    try:
        opts, args = getopt.getopt(argv, 'h:k:s',
                                         ["help", "aws-key=", "aws-secret=", "region=", "port=", "user="])
    except getopt.GetoptError, err:
        print >>sys.stderr, err
        short_usage()
        sys.exit(2)

    aws_key = os.environ.get("AWS_ACCESS_KEY_ID")
    aws_secret = os.environ.get("AWS_SECRET_ACCESS_KEY")
    region = os.environ.get("AWS_EC2_REGION")
    user = 'ubuntu'
    port = 22

    for opt, arg in opts:
        if opt in ("-h", "--help"):
            full_usage()
            sys.exit()
        elif opt in("-k", "--aws-key"):
            aws_key = arg
        elif opt in("-s", "--aws-secret"):
            aws_secret = arg
        elif opt in ("--region"):
            region = arg
        elif opt in ("--port"):
            port = arg
        elif opt in ("--user"):
            user = arg

    if not aws_key or not aws_secret:
        if not aws_key:
            print >>sys.stderr,\
            "AWS_ACCESS_KEY_ID not set in environment and not",\
            "specified by --aws-key KEY or -k KEY"

        if not aws_secret:
            print >>sys.stderr,\
            "AWS_SECRET_ACCESS_KEY not set in environment and not",\
            "specified by --aws-secret SECRET or -s SECRET"

        short_usage()
        sys.exit(2)

    region = region and boto.ec2.get_region(region,
                                            aws_access_key_id=aws_key,
                                            aws_secret_access_key=aws_secret)

    _ec2 = boto.ec2.connection.EC2Connection(aws_key, aws_secret, region=region)

    host_file_content = create_ssh_config(_ec2, port , user)
    print(host_file_content)

if __name__ == "__main__":
    main(sys.argv[1:])