update_ssh_config.py

'''
Update your ssh keys for a given client.
e.g.
  python update_ssh_config.py qgc

Usage:
    update_ssh_config.py <client>

'''
from docopt import docopt
from boto import ec2
import os


TEMPLATE = '''Host {shortcut}
    IdentityFile ~/.ssh/{key}.pem
    Hostname {host}
    User ubuntu'''


def list_instances(access_key_id, secret_access_key):

    ec2conn = ec2.connect_to_region(
        'ap-southeast-2',
        aws_access_key_id=access_key_id,
        aws_secret_access_key=secret_access_key)

    all_instances = ec2conn.get_all_instances()

    return [i for r in all_instances for i in r.instances]


def make_config(host, stack, logical, key, arguments):

    client = arguments['<client>']

    shortcut = '{client}-{stack}-{logical}'.format(
        client=client,
        stack=stack,
        logical=logical)

    config = TEMPLATE.format(
        shortcut=shortcut,
        host=host,
        key=key)

    return config


def update_ssh_with_config(config):

    fn = os.path.expanduser('~/.ssh/config')

    if not os.path.exists(fn):
        f = open(fn, 'w')
        f.write()
        f.close()

    first_line = config.splitlines()[0].strip()

    matching_block = False
    new_file = []
    last_line = None
    with open(fn, 'r') as f:
        data = f.read()
        for line in data.splitlines():
            line = line.rstrip()

            if first_line == line:
                matching_block = True
            if line == '':
                matching_block = False

            if not matching_block:
                if line != last_line:
                    new_file.append(line)

            last_line = line

    if len(new_file) and new_file[-1].strip() != '':
        new_file.append('')

    # Add in the new config
    for line in config.splitlines():
        new_file.append(line)

    content = '\n'.join(new_file)

    # write the new file
    with open(fn, 'w') as f:
        f.write(content)


def update_configuration(arguments, access_key_id, secret_access_key):

    instances = list_instances(access_key_id, secret_access_key)

    for i in instances:

        host = i.public_dns_name
        key = i.key_name
        stack = i.tags['aws:cloudformation:stack-name']
        if '-' in stack:
            stack = stack.split('-')[1].lower()
        logical = i.tags['aws:cloudformation:logical-id'].lower()

        config = make_config(host, stack, logical, key, arguments)

        update_ssh_with_config(config)
        print config.splitlines()[0]


if __name__ == '__main__':
    arguments = docopt(__doc__, version='!.0')

    access_key_id = os.environ.get('AWS_ACCESS_KEY_ID', None)
    if not access_key_id:
        print('Please set AWS_ACCESS_KEY_ID, e.g.:')
        print('export AWS_ACCESS_KEY_ID=AKAIXXXXXXXX')
        print('export AWS_SECRET_ACCESS_KEY=Uxswre12...')
        exit(-1)

    secret_access_key = os.environ.get('AWS_SECRET_ACCESS_KEY', None)
    if not access_key_id:
        print('Please set AWS_SECRET_ACCESS_KEY, e.g.:')
        print('export AWS_SECRET_ACCESS_KEY=Uxswre12...')
        exit(-1)

    update_configuration(arguments, access_key_id, secret_access_key)