lambda_function

#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""Lambda function to tag EC2 instances

Copyright (C) 2017 Peter Pakos <peter.pakos@wandisco.com>

Version 1.0.1

This program is free software: you can redistribute it and/or modify
it under the terms of the GNU General Public License as published by
the Free Software Foundation, either version 3 of the License, or
(at your option) any later version.

This program is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
GNU General Public License for more details.

You should have received a copy of the GNU General Public License
along with this program.  If not, see <http://www.gnu.org/licenses/>.
"""

from __future__ import print_function
import os
import sys
import gzip
import json
import urllib
import boto3
import botocore.exceptions


def lambda_handler(event, context):
    # print("Received event: " + json.dumps(event, indent=2))

    bucket = None
    key = None
    session = None
    if __name__ == '__main__':
        if len(sys.argv) == 3:
            profile = sys.argv[1]
            bucket = 'wandisco-%s-auto-tag' % profile
            key = sys.argv[2]
            try:
                session = boto3.Session(profile_name=profile)
            except botocore.exceptions.ProfileNotFound as err:
                print(err)
                exit(1)
        else:
            print('usage: %s AWS_PROFILE EVENT_PATH' % (sys.argv[0]))
            exit(1)
    else:
        session = boto3.Session()
        bucket = event['Records'][0]['s3']['bucket']['name']
        key = urllib.unquote_plus(event['Records'][0]['s3']['object']['key']).decode('utf8')

    s3 = session.client('s3')
    print('Loading events from %s...' % key)
    tmp_file = '/tmp/' + os.path.basename(key)
    try:
        s3.download_file(bucket, key, tmp_file)
    except Exception as err:
        print(err)
        exit(1)

    with gzip.open(tmp_file) as f:
        file_content = f.read()

    events = json.loads(file_content)['Records']
    asg_uid = {}
    for event in events:
        instance_ids = []

        if 'role' in event['eventName'].lower() or 'assume' in event['eventName'].lower():
            print(json.dumps(event, indent=2))
        if event['eventName'] == 'RequestSpotInstances':
            print('Processing event %s...' % event['eventName'])
            region = event['awsRegion']
            user = event['requestParameters']['launchSpecification']['keyName']
            ec2 = session.client('ec2', region_name=region)
            try:
                for item in event['responseElements']['spotInstanceRequestSet']['items']:
                    instance_ids.append(item['spotInstanceRequestId'])
            except TypeError as err:
                print(err)
                continue

            try:
                print('Tagging in progress... Last_user: %s (SSH Key Name), Request ID: %s' %
                      (user, ', '.join(instance_ids)))
                response = ec2.create_tags(Resources=instance_ids, Tags=[{'Key': 'Last_user', 'Value': user}])
            except botocore.exceptions.ClientError as err:
                print(err)
                continue
            print('HTTP response: %s' % response['ResponseMetadata']['HTTPStatusCode'])

            try:
                print('Requesting details about spot instance requests: %s' % ', '.join(instance_ids))
                response = ec2.describe_spot_instance_requests(SpotInstanceRequestIds=instance_ids)
            except botocore.exceptions.ClientError as err:
                print(err)
                continue
            print('HTTP response: %s' % response['ResponseMetadata']['HTTPStatusCode'])

            instance_ids = []
            for request in response['SpotInstanceRequests']:
                instance_ids.append(request['InstanceId'])
            try:
                print('Tagging in progress... Last_user: %s (SSH Key Name), Instances: %s' %
                      (user, ', '.join(instance_ids)))
                response = ec2.create_tags(Resources=instance_ids, Tags=[{'Key': 'Last_user', 'Value': user}])
            except botocore.exceptions.ClientError as err:
                print(err)
                continue
            print('HTTP response: %s' % response['ResponseMetadata']['HTTPStatusCode'])

        elif event['eventName'] == 'RunJobFlow':
            print('Processing event %s...' % event['eventName'])
            region = event['awsRegion']
            user = event['userIdentity']['userName']
            cluster_id = event['responseElements']['jobFlowId']
            emr = session.client('emr', region_name=region)
            instances = emr.list_instances(ClusterId=cluster_id)['Instances']
            for instance in instances:
                instance_ids.append(instance['Ec2InstanceId'])
            ec2 = session.client('ec2', region_name=region)
            try:
                print('Tagging in progress... Last_user: %s, Instances: %s' % (user, ', '.join(instance_ids)))
                response = ec2.create_tags(Resources=instance_ids, Tags=[{'Key': 'Last_user', 'Value': user}])
            except botocore.exceptions.ClientError as err:
                print(err)
                continue
            print('HTTP response: %s' % response['ResponseMetadata']['HTTPStatusCode'])

        elif event['eventName'] in ['CreateAutoScalingGroup', 'UpdateAutoScalingGroup']:
            print('Processing event %s...' % event['eventName'])
            user = event['userIdentity']['userName']
            as_group = event['requestParameters']['autoScalingGroupName']
            region = event['awsRegion']
            autoscaling = session.client('autoscaling', region_name=region)
            if as_group not in asg_uid:
                asg_uid[as_group] = user
            try:
                print('Tagging in progress... Last_user: %s, autoScalingGroupName: %s' % (user, as_group))
                response = autoscaling.create_or_update_tags(Tags=[{
                    'ResourceId': as_group,
                    'ResourceType': 'auto-scaling-group',
                    'Key': 'Last_user',
                    'Value': user,
                    'PropagateAtLaunch': True
                }])
            except botocore.exceptions.ClientError as err:
                print(err)
                continue
            print('HTTP response: %s' % response['ResponseMetadata']['HTTPStatusCode'])
            response = autoscaling.describe_auto_scaling_groups(AutoScalingGroupNames=[as_group])
            for instance in response['AutoScalingGroups'][0]['Instances']:
                instance_ids.append(instance['InstanceId'])
            ec2 = session.client('ec2', region_name=region)
            try:
                print('Tagging in progress... Last_user: %s, Instances: %s' % (user, ', '.join(instance_ids)))
                response = ec2.create_tags(Resources=instance_ids, Tags=[{'Key': 'Last_user', 'Value': user}])
            except botocore.exceptions.ClientError as err:
                print(err)
            print('HTTP response: %s' % response['ResponseMetadata']['HTTPStatusCode'])

        elif event['eventName'] in ['RunInstances', 'StartInstances', 'StopInstances']:
            print('Processing event %s...' % event['eventName'])
            region = event['awsRegion']
            if event['userAgent'] in ['autoscaling.amazonaws.com', 'elasticmapreduce.amazonaws.com']:
                print('Instance created by %s, skipping...' % event['userAgent'])
                continue
            if event['userIdentity']['type'] == 'Root':
                user = bucket.split('-')[1]
            else:
                user = event.get('userIdentity', {}).get('userName', 'unknown')
            ec2 = session.client('ec2', region_name=region)
            try:
                for item in event['responseElements']['instancesSet']['items']:
                    instance_ids.append(item['instanceId'])
            except TypeError as err:
                print(err)
                continue
            try:
                print('Tagging in progress... Last_user: %s, Instances: %s' % (user, ', '.join(instance_ids)))
                response = ec2.create_tags(Resources=instance_ids, Tags=[{'Key': 'Last_user', 'Value': user}])
            except botocore.exceptions.ClientError as err:
                print(err)
                continue
            print('HTTP response: %s' % response['ResponseMetadata']['HTTPStatusCode'])

        elif event['eventName'] == 'CreateTags' and event['userAgent'] == 'autoscaling.amazonaws.com':
            print('Processing event %s (%s)...' % (event['eventName'], event['userAgent']))
            region = event['awsRegion']
            user = None
            ec2 = session.client('ec2', region_name=region)
            for item in event['requestParameters']['tagSet']['items']:
                if item['key'] == 'aws:autoscaling:groupName' and item['value'] in asg_uid:
                    user = asg_uid[item['value']]
                    continue
            if not user:
                print('Unable to find user, skipping...')
                continue
            try:
                for item in event['requestParameters']['resourcesSet']['items']:
                    instance_ids.append(item['resourceId'])
            except TypeError as err:
                print(err)
                continue
            try:
                print('Tagging in progress... Last_user: %s, Instances: %s' % (user, ', '.join(instance_ids)))
                response = ec2.create_tags(Resources=instance_ids, Tags=[{'Key': 'Last_user', 'Value': user}])
            except botocore.exceptions.ClientError as err:
                print(err)
                continue
            print('HTTP response: %s' % response['ResponseMetadata']['HTTPStatusCode'])
            print(region, instance_ids)

        else:
            print('Skipping event %s...' % event['eventName'])


if __name__ == '__main__':
    lambda_handler('', '')