ssm-login.py

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()