-
-
Save asimzeeshan/3bf90e61185a53d8cdf688aad657547a to your computer and use it in GitHub Desktop.
Automatically update your SSH config to allow ssh-ing into EC2 instances by their name. Works on Linux and Mac OS X.
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
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() |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment