import os
import subprocess
import boto.ec2
import itertools

AWS_KEY = os.getenv('AWS_ACCESS_KEY')
AWS_SECRET = os.getenv('AWS_SECRET_KEY')

# Tweak as necessary for your infrastructure
regions = ['us-east-1', 'us-west-1', 'eu-west-1']
def get_ec2_connection(region):
    conn = boto.ec2.connect_to_region(region,
                                      aws_access_key_id=AWS_KEY,
                                      aws_secret_access_key=AWS_SECRET)
    return conn

def get_instances_in_region(region="ALL"):
    ''' Get BOTO instances in a region '''
    if region == "ALL":
        instances = []
        for region in regions:
            instances.extend(get_instances_in_region(region=region))
        return instances
    else:
        conn = get_ec2_connection(region)
        try:
            reservations = conn.get_all_instances()
            instances = itertools.chain(*map(lambda r: r.instances, reservations))
        except boto.exception.EC2ResponseError:
            raise

        #Filter out stopped/terminated ones
        instances = filter(lambda i: i.public_dns_name is not None and len(i.public_dns_name) > 0, instances)
        return instances

HOSTS_HEADER = "### AUTOGENERATED SSH CONF DO NOT MODIFY ###\n"
HOSTS_FOOTER = "\n### END AUTOGENERATED SSH CONF ###"
SSH_CONF = os.path.expanduser('~/.ssh/config')

def update_ssh_config():
    # Ensure the SSH_CONF exists
    print "Updating SSH config. Storing in ", SSH_CONF 
    subprocess.call(['touch', SSH_CONF])


    nodes = get_node_names()
    print "Adding configuration for %d nodes" % len(nodes)
    entries = [to_entry(tag_name, dns_name) for (tag_name, dns_name) in nodes]

    new_hosts = '\n\n'.join(entries)


    with file(SSH_CONF, 'r') as f:
        current_contents = f.read()

    with file(SSH_CONF, 'w') as f:
        head_index = current_contents.find(HOSTS_HEADER)
        foot_index = current_contents.find(HOSTS_FOOTER)
        if head_index != -1 and foot_index != -1: 
            # Splice out 
            current_contents = current_contents[:head_index] + current_contents[foot_index + len(HOSTS_FOOTER):]

        current_contents += HOSTS_HEADER
        current_contents += new_hosts
        current_contents += HOSTS_FOOTER

        f.write(current_contents)

    with file(SSH_CONF, 'r') as f:
        assert(len(f.read()) == len(current_contents))

def get_node_names():
    instances = get_instances_in_region("ALL")  
    return [(safe_name(ins.tags['Name']), ins.public_dns_name) for ins in instances if 'Name' in ins.tags]

def safe_name(name):
    ''' If your instance names are especially bizarre, this may need to be refined'''
    return name.replace(' ', '_')

def to_entry(tag_name, dns_name):
    return "Host %s\n\tHostName %s\n\tUser ubuntu" % (tag_name, dns_name)


if __name__ == "__main__":
    update_ssh_config()