|
| 1 | +import os |
| 2 | +import sys |
| 3 | + |
| 4 | +from six.moves import configparser |
| 5 | + |
| 6 | +import boto3 |
| 7 | +import click |
| 8 | +from jinja2 import Environment, FileSystemLoader |
| 9 | +from workflow import Workflow3 |
| 10 | +from workflow import ( |
| 11 | + MATCH_ALL, |
| 12 | + MATCH_ALLCHARS, |
| 13 | + MATCH_ATOM, |
| 14 | + MATCH_CAPITALS, |
| 15 | + MATCH_INITIALS, |
| 16 | + MATCH_INITIALS_CONTAIN, |
| 17 | + MATCH_INITIALS_STARTSWITH, |
| 18 | + MATCH_STARTSWITH, |
| 19 | + MATCH_SUBSTRING, |
| 20 | +) |
| 21 | + |
| 22 | +pass_wf = click.make_pass_decorator(Workflow3) |
| 23 | +log = None |
| 24 | + |
| 25 | + |
| 26 | +def get_ec2_instances(): |
| 27 | + client = boto3.client('ec2') |
| 28 | + next_token = {} |
| 29 | + instances = [] |
| 30 | + while True: |
| 31 | + log.debug('calling describe_instances') |
| 32 | + response = client.describe_instances(MaxResults=1000, **next_token) |
| 33 | + for reservation in response['Reservations']: |
| 34 | + for instance in reservation['Instances']: |
| 35 | + if 'Tags' in instance: |
| 36 | + for tag in instance['Tags']: |
| 37 | + instance['Tag:%s' % tag['Key']] = tag['Value'] |
| 38 | + instances.append(instance) |
| 39 | + if 'NextToken' in response: |
| 40 | + next_token['NextToken'] = response['NextToken'] |
| 41 | + else: |
| 42 | + break |
| 43 | + return instances |
| 44 | + |
| 45 | + |
| 46 | +def dedupe(lofd, key): |
| 47 | + seen = set() |
| 48 | + res = [] |
| 49 | + for d in lofd: |
| 50 | + if d[key] not in seen: |
| 51 | + seen.add(d[key]) |
| 52 | + res.append(d) |
| 53 | + return res |
| 54 | + |
| 55 | + |
| 56 | +def find_ec2(wf, profile, quicklook, query): |
| 57 | + instances = wf.cached_data('%s-ec2' % profile, get_ec2_instances, max_age=600) |
| 58 | + |
| 59 | + if query: |
| 60 | + matched = wf.filter(query, instances, key=lambda i: unicode(i['InstanceId']), match_on=MATCH_STARTSWITH) |
| 61 | + matched += wf.filter(query, instances, key=lambda i: i.get('Tag:Name', u'')) |
| 62 | + # TODO: parse query for facet:query to perform faceted search |
| 63 | + # matched += wf.filter(query, instances, key=lambda i: i.get('Tag:Application', u'')) |
| 64 | + # matched += wf.filter(query, instances, key=lambda i: i.get('Tag:Role', u'')) |
| 65 | + # matched += wf.filter(query, instances, key=lambda i: i.get('Tag:Vertical', u'')) |
| 66 | + # matched += wf.filter(query, instances, key=lambda i: i.get('Tag:Vpc', u'')) |
| 67 | + # matched += wf.filter(query, instances, key=lambda i: i.get('Tag:Environment', u'')) |
| 68 | + else: |
| 69 | + matched = instances |
| 70 | + |
| 71 | + matched = dedupe(matched, 'InstanceId') |
| 72 | + |
| 73 | + for instance in matched: |
| 74 | + if 'Tag:Name' in instance: |
| 75 | + title = '%s (%s)' % (instance['Tag:Name'], instance['InstanceId']) |
| 76 | + else: |
| 77 | + title = instance['InstanceId'] |
| 78 | + uid = '%s-ec2-%s' % (profile, instance['InstanceId']) |
| 79 | + valid = instance['State']['Name'] == 'running' |
| 80 | + quicklookurl = None |
| 81 | + # quicklookurl = quicklook(uid, 'ec2', { |
| 82 | + # 'title': title, |
| 83 | + # 'uid': uid, |
| 84 | + # 'instance': instance, |
| 85 | + # }) |
| 86 | + |
| 87 | + item = wf.add_item( |
| 88 | + title, |
| 89 | + subtitle='private ip', |
| 90 | + arg=instance.get('PrivateIpAddress', 'N/A'), |
| 91 | + valid=valid and 'PrivateIpAddress' in instance, |
| 92 | + uid=uid, |
| 93 | + icon='icons/ec2.eps', |
| 94 | + type='default', |
| 95 | + quicklookurl=quicklookurl |
| 96 | + ) |
| 97 | + item.setvar('title', title) |
| 98 | + item.add_modifier( |
| 99 | + "shift", |
| 100 | + subtitle='public ip', |
| 101 | + arg=instance.get('PublicIpAddress', 'N/A'), |
| 102 | + valid=valid and 'PublicIpAddress' in instance, |
| 103 | + ) |
| 104 | + item.add_modifier( |
| 105 | + "cmd", |
| 106 | + subtitle='open in console', |
| 107 | + arg='https://us-west-2.console.aws.amazon.com/ec2/v2/home?region=us-west-2#Instances:search=%s;sort=instanceState' % instance['InstanceId'], |
| 108 | + valid=True, |
| 109 | + ) |
| 110 | + |
| 111 | + |
| 112 | +@click.group() |
| 113 | +@pass_wf |
| 114 | +def cli(wf): |
| 115 | + pass |
| 116 | + |
| 117 | +@cli.command() |
| 118 | +@pass_wf |
| 119 | +def list_profiles(wf): |
| 120 | + parser = configparser.ConfigParser() |
| 121 | + parser.read(os.path.expanduser('~/.aws/credentials')) |
| 122 | + for profile in parser.sections(): |
| 123 | + wf.add_item( |
| 124 | + title=profile, |
| 125 | + valid=True, |
| 126 | + arg=profile, |
| 127 | + ) |
| 128 | + wf.send_feedback() |
| 129 | + |
| 130 | +@cli.command() |
| 131 | +@click.argument('profile') |
| 132 | +@pass_wf |
| 133 | +def set_profile(wf, profile): |
| 134 | + wf.settings['profile'] = profile |
| 135 | + |
| 136 | +@cli.command() |
| 137 | +@click.argument('query') |
| 138 | +@pass_wf |
| 139 | +def search(wf, query): |
| 140 | + profile = wf.settings['profile'] |
| 141 | + if profile is None: |
| 142 | + raise Exception('no profile selected') |
| 143 | + |
| 144 | + os.environ['AWS_PROFILE'] = profile |
| 145 | + |
| 146 | + def _quicklook_closure(): |
| 147 | + env = Environment(loader=FileSystemLoader(os.path.join(wf.workflowdir, 'quicklook'))) |
| 148 | + def build_quicklook(uid, template, context): |
| 149 | + filename = wf.cachefile('%s.html' % uid) |
| 150 | + template = env.get_template('%s.html.j2' % template) |
| 151 | + with open(filename, 'w') as f: |
| 152 | + f.write(template.render(**context)) |
| 153 | + return filename |
| 154 | + return build_quicklook |
| 155 | + quicklook = _quicklook_closure() |
| 156 | + |
| 157 | + find_ec2(wf, profile, quicklook, query) |
| 158 | + |
| 159 | + wf.send_feedback() |
| 160 | + |
| 161 | +if __name__ == '__main__': |
| 162 | + wf = Workflow3() |
| 163 | + log = wf.logger |
| 164 | + wf.run(lambda wf: cli(obj=wf, auto_envvar_prefix='WF')) |
0 commit comments