Skip to content

Instantly share code, notes, and snippets.

@asimzeeshan
Forked from rcoh/sshupdate.py
Created August 16, 2022 11:07
Show Gist options
  • Save asimzeeshan/3bf90e61185a53d8cdf688aad657547a to your computer and use it in GitHub Desktop.
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.
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