Skip to content

Instantly share code, notes, and snippets.

@shomah4a
Last active May 4, 2021 02:33
Show Gist options
  • Save shomah4a/a5e2d2025368e649b13e7876d361d076 to your computer and use it in GitHub Desktop.
Save shomah4a/a5e2d2025368e649b13e7876d361d076 to your computer and use it in GitHub Desktop.
ssm-login.py
#!/usr/bin/env python3
import argparse
import os
import subprocess
import sys
import boto3
def make_parser():
p = argparse.ArgumentParser()
p.add_argument('-f', '--use-first-one', action='store_true', dest='use_first', default=False)
p.add_argument('-p', '--profile', type=str, default=None)
p.add_argument('-r', '--region', type=str, default='ap-northeast-1')
p.add_argument('instance_name', help='instance name or instance id (partial matching)')
return p
def main():
arg = make_parser().parse_args(sys.argv[1:])
if arg.profile is None:
session = boto3.Session(region_name=arg.region)
else:
session = boto3.Session(profile_name=arg.profile, region_name=arg.region)
ec2 = session.resource('ec2')
instances = find_instances(ec2, arg.instance_name)
if not instances:
print('instance not found')
return
if len(instances) == 1 or arg.use_first:
print('login to', name_fmt(*instances[0]))
start_session(arg.profile, arg.region, instances[0][1].id)
return
print('multiple instances found:', ', '.join([name_fmt(*x) for x in instances]))
def name_fmt(name, instance):
if name:
return '%s(%s)' % (instance.id, name)
else:
return instance.id
def find_instances(ec2, name):
result = []
instances = ec2.instances.all()
instances = sorted(instances, key = lambda x: x.launch_time, reverse=True)
for instance in instances:
if instance.state['Name'] != 'running':
continue
if name in instance.id:
result.append((None, instance))
if instance.tags is None:
continue
for tag in instance.tags:
if tag['Key'] != 'Name':
continue
if name in tag['Value']:
result.append((tag['Value'], instance))
return result
def start_session(profile, region, id):
if profile is None:
cmd = ['aws', 'ssm', 'start-session', '--target', id]
else:
cmd = ['aws', '--profile', profile, 'ssm', 'start-session', '--target', id]
env = os.environ.copy()
env['AWS_PROFILE'] = profile
env['AWS_DEFAULT_REGION'] = region
p = subprocess.Popen(cmd, env=env)
while True:
try:
p.wait()
break
except KeyboardInterrupt as e:
pass
if __name__ == '__main__':
main()
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment