ec2-control

#!/usr/bin/env python3

# NB Python 3

# need boto3, attr, and tabulate python libraries.
# all installable via pip; your linux distribution
# may also supply them via its package installer

# usage -- you should have already configured the AWS CLI
# before using this.   Mostly it's a wrapper around some common
# operations.

import boto3
import attr
import itertools
import argparse
import sys
import re
from botocore.exceptions import ClientError
from tabulate import tabulate


@attr.s
class Instance:
    id = attr.ib(default="unknown")
    name = attr.ib(default="unknown")
    hostname = attr.ib(default='unknown')
    keyname = attr.ib(default='')
    state = attr.ib(default='unknown')
    tags = attr.ib(default={})

client = boto3.client('ec2')


def hydrate_instance(inst_data):
    iid = inst_data['InstanceId']
    tags = dict((x.get("Key", '-'), x.get('Value'))
                for x in inst_data.get("Tags", []))
    name = tags.get("Name", '--unknown--')
    state = inst_data.get("State", {}).get("Name", '--unknown--')
    hostname = inst_data.get('PublicDnsName', 'n/a')
    keyname = inst_data.get('KeyName', '<unknown')
    return Instance(iid, name, hostname, keyname, state, tags)


def instance_to_fields(i):
    return [i.id, i.name, i.hostname, i.state]


def get_instances():
    response = client.describe_instances()
    reservations = response['Reservations']
    instances = itertools.chain.from_iterable(
        (x['Instances'] for x in reservations))
    return (hydrate_instance(x) for x in instances)


def list_instances(requested=()):
    req_ids = [x.id for x in requested]
    instances = (instance_to_fields(x)
                 for x in get_instances() if len(req_ids) == 0 or x.id in req_ids)

    return tabulate(instances, headers=("ID", "Name", "Hostname", "State"))


def connect(instances=()):
    def keyname_to_file(kn):
        if kn is None:
            return ""
        return "~/.aws/keypairs/" + kn.lower().replace(" ", "") + ".pem"

    return ["ssh -p 22 -i {key} ec2-user@{hostname}".format(key=keyname_to_file(i.keyname), hostname=i.hostname) for i in instances if i.state == 'running']


def start_instances(inst_ids):
    try:
        client.start_instances(InstanceIds=inst_ids, DryRun=True)
    except ClientError as e:
        if 'DryRunOperation' not in str(e):
            raise
    try:
        response = client.start_instances(InstanceIds=inst_ids, DryRun=False)
        print(response)
    except ClientError as e:
        print(e)


def stop_instances(inst_ids):
    try:
        client.stop_instances(InstanceIds=inst_ids, DryRun=True)
    except ClientError as e:
        if 'DryRunOperation' not in str(e):
            raise
    try:
        response = client.stop_instances(InstanceIds=inst_ids, DryRun=False)
        print(response)
    except ClientError as e:
        print(e)


def find_by_name(names=()):
    instances = get_instances()
    regexes = []
    for name in names:
        if name[0] == '~':
            regexes.append(re.compile("".join(name[1:])))
        else:
            regexes.append(re.compile("^" + name + "$"))

    def finder(x):
        for rg in regexes:
            if rg.search(x):
                return True
        return False

    return [x for x in instances if finder(x.name)]


# utility internal methods

def _to_ids(instances=()):
    return [x.id for x in instances]


def _falsey(x):
    return x is None or len(x) == 0


def filter_instances(ids=(), names=()):
    if _falsey(ids) and _falsey(names):
        return get_instances()
    if ids is not None and len(ids) != 0:
        return (x for x in get_instances() if x.id in ids)
    return find_by_name(names)

if __name__ == '__main__':
    parser = argparse.ArgumentParser(description="Control EC2 instances")
    parser.add_argument(
        'command', type=str, nargs='?',
        default='list', help='Command to be executed')
    parser.add_argument('--names', nargs='*', type=str,
                        help='Names (tags) of instances to be operated on')
    parser.add_argument('--instances', metavar='id', type=str, nargs='*',
                        help='instance IDs to be operated on')
    args = parser.parse_args(sys.argv[1:])
    found_instances = filter_instances(args.instances, args.names)
    if len(found_instances) == 0:
        print("No instances found matching criteria")
        sys.exit(0)
    if args.command is None or args.command == 'list':
        print(list_instances(found_instances))
    elif args.command == 'start':
        start_instances(_to_ids(found_instances))
    elif args.command == 'stop':
        stop_instances(_to_ids(found_instances))
    elif args.command == 'connect':
        for cstr in connect(found_instances):
            print(cstr)
    else:
        parser.print_help()