#!/usr/bin/python

import argparse
import boto3
import copy
import inspect
import json
import os
import sys
import time
import urllib2
import yaml
from botocore.exceptions import WaiterError
from botocore.utils import ClientError
from os import listdir
from os.path import isfile
from os.path import join


class SecGroupUpdater:
    COMMON_FOLDER = 'common'
    GROUPS_FOLDER = 'groups'

    CFN_STACK_NAME = 'navadmins-secgroups-'

    ACCOUNT_ID = '{{ ACCOUNT_ID }}'
    ACCOUNT_SHORTNAME = '{{ ACCOUNT_SHORTNAME }}'
    CFN_ROLE_ARN = 'arn:aws:iam::{{ ACCOUNT_ID }}:role/{{ ACCOUNT_SHORTNAME }}-navadmins-cloudformation'

    ACCOUNT_STRUCTURE = {
        '466131539206': 'navprod',
        '235580133017': 'navprodc1',
        '743593264885': 'navnonprod',
        '890655436785': 'navprodadmin',
        '784869473894': 'navprodstandby',
        '472973657150': 'navsandbox',
        '873015214695': 'navtest',
        '888995848466': 'splunkops',
        '972242968777': 'datalake-nonprod',
        '107413163702': 'datalake-prod',
    }

    BLANK_CFN = {
        'AWSTemplateFormatVersion': '2010-09-09',
        'Description': 'all of the security groups for ',
        'Outputs': {},
        'Resources': {}
    }

    ALL = 'all'

    NO_UPDATES = 'An error occurred (ValidationError) when calling the UpdateStack operation: No updates are to be performed.'

    def __init__(self, change_set):
        """Creates a cloudformation stack for each vpc in this account containing all of the security groups in this directory"""

        boto3.setup_default_session()
        if args.region is not None:
            boto3.setup_default_session(region_name=args.region)
        else:
            boto3.setup_default_session(region_name=self.get_region())

        account_id = self.get_account_id()
        self.cfn_role_arn = self.CFN_ROLE_ARN.replace(self.ACCOUNT_ID, account_id)

        if account_id in self.ACCOUNT_STRUCTURE:
            account_shortname = self.ACCOUNT_STRUCTURE[account_id]
        else:
            print '\nAccount shortname not found, please enter the shortname for this account:\n'
            account_shortname = raw_input()

        self.cfn_role_arn = self.cfn_role_arn.replace(self.ACCOUNT_SHORTNAME, account_shortname)

        alias_response = boto3.client('iam').list_account_aliases()
        if 'AccountAliases' not in alias_response or len(alias_response['AccountAliases']) < 1:
            print 'Unable to find account alias for this account, this script requires an alias'
            sys.exit(-1)
        self.account_alias = alias_response['AccountAliases'][0]

        self.vpcs = []
        vpc_response = boto3.client('ec2').describe_vpcs()
        if 'Vpcs' in vpc_response:
            self.vpcs = vpc_response['Vpcs']
        if len(self.vpcs) <= 0:
            print 'Unable to load vpcs, exiting'
            sys.exit(-2)
        self.create_vpc_names(self.vpcs)

        path = inspect.getfile(inspect.currentframe())[0:inspect.getfile(inspect.currentframe()).rfind('/')+1]

        stack_names = []
        waiter_list = {}

        failed = False

        for vpc in self.vpcs:
            self.common_folder = path + self.GROUPS_FOLDER + '/' + self.COMMON_FOLDER
            self.account_folder = path + self.GROUPS_FOLDER + '/' + self.account_alias + '/' + self.COMMON_FOLDER
            self.vpc_folder = path + self.GROUPS_FOLDER + '/' + self.account_alias + '/' + vpc['VpcName']

            self.template = copy.deepcopy(self.BLANK_CFN)
            self.template['Description'] = self.template['Description'] + vpc['VpcName']

            if self.populate_template(vpc):
                failed = True
                continue

            stack_name = self.CFN_STACK_NAME + vpc['VpcId'][4:]
            stack_names.append(stack_name)
            self.update_cfn_stack(stack_name, change_set, waiter_list)

        if failed:
            sys.exit(-4)

        self.wait_for_stack_updates(waiter_list)

    @staticmethod
    def get_region():
        region = urllib2.urlopen('http://169.254.169.254/latest/meta-data/placement/availability-zone').read()
        return region[0:len(region)-1]

    @staticmethod
    def get_account_id():
        """Gets the current account's ID"""

        return boto3.client('sts').get_caller_identity().get('Account')

    @staticmethod
    def stack_exists(stack_name):
        """Check if a cfn stack exists"""

        try:
            boto3.client('cloudformation').describe_stacks(StackName=stack_name)
            return True
        except ClientError:
            return False

    @staticmethod
    def json_default(obj):
        """JSON serializer for datetimes"""
        return str(obj)

    @staticmethod
    def create_vpc_names(vpcs):
        for vpc in vpcs:
            for vpc_tag in vpc['Tags']:
                if vpc_tag['Key'] == 'Name':
                    vpc['VpcName'] = vpc_tag['Value']

    @staticmethod
    def check_for_duplicates(dict1, dict2):
        for entry in dict1:
            if entry in dict2:
                return True
        return False

    def populate_template(self, vpc):
        """Creates a cloudformation template using the yml files from the three applicable folders"""

        failed = False
        for folder in [self.common_folder, self.account_folder, self.vpc_folder]:
            try:
                files = [f for f in listdir(folder) if isfile(join(folder, f)) and f.endswith('.yml')]
                for yaml_file in files:
                    print 'Loading ' + folder + '/' + yaml_file
                    with open(folder + '/' + yaml_file) as open_file:
                        input_file = yaml.safe_load(open_file)
                        if input_file is not None:
                            if self.generate_groups(input_file, vpc):
                                failed = True
                            self.generate_rules(input_file, vpc)
                            self.generate_outputs(input_file, vpc)

                        else:
                            print 'WARNING: empty file ' + folder + '/' + yaml_file

            except OSError as e:
                print 'WARNING: Unable to load folder: ' + str(e)

        return failed

    def generate_groups(self, cfn_resources, vpc):
        """Creates all of the security groups for this vpc"""

        failed = False
        if 'Resources' in cfn_resources and self.find_vpc(cfn_resources, vpc):
            if self.check_for_duplicates(self.template['Resources'], cfn_resources['Resources']):
                print '\nDuplicate Security Group CFN name found:\n' + json.dumps(self.template, indent=4) + '\n' + json.dumps(cfn_resources, indent=4)
                sys.exit(-3)

            for resource_name in cfn_resources['Resources']:
                resource = cfn_resources['Resources'][resource_name]

                if resource['Type'] == 'AWS::EC2::SecurityGroup':
                    if ('Properties' in resource and
                            'VpcId' in resource['Properties'] and
                            (resource['Properties']['VpcId'] == self.ALL or
                                (len(resource['Properties']['VpcId'].split(',')) > 1 and
                                    vpc['VpcId'] in resource['Properties']['VpcId']))):

                        generated_resource = self.generate_group_for_vpc(resource, vpc)
                        self.template['Resources'][resource_name] = generated_resource

                    elif len(resource['Properties']['VpcId'].split(',')) <= 1 and vpc['VpcId'] in resource['Properties']['VpcId']:
                        self.template['Resources'][resource_name] = resource

                    if resource_name in self.template['Resources'] and 'Properties' in self.template['Resources'][resource_name]:
                        for key, value in self.template['Resources'][resource_name]['Properties'].items():
                            if key == 'SecurityGroupEgress' or key == 'SecurityGroupIngress':
                                for rule in value:
                                    if 'Comments' not in rule or len(rule['Comments']) <= 0:
                                        print 'No Comments field found in ' + resource_name + ' at ' + str(rule)
                                        failed = True
                                    else:
                                        rule.pop('Comments')
        return failed

    @staticmethod
    def generate_group_for_vpc(resource, vpc):
        """Sets the appropriate Name for a security group based on its vpc"""

        generated_resource = copy.deepcopy(resource)
        generated_resource['Properties']['VpcId'] = vpc['VpcId']

        if 'Tags' in generated_resource['Properties']:
            for resource_tag in generated_resource['Properties']['Tags']:
                if resource_tag['Key'] == 'Name':
                    if 'Tags' in vpc:
                        for vpc_tag in vpc['Tags']:
                            if vpc_tag['Key'] == 'Name':
                                resource_tag['Value'] = vpc_tag['Value'] + '-' + resource_tag['Value']

        return generated_resource

    def generate_rules(self, cfn_resources, vpc):
        """Creates all of the rules that apply to the security groups for this vpc"""

        if 'Resources' in cfn_resources and self.find_vpc(cfn_resources, vpc):
            for resource_name in cfn_resources['Resources']:
                resource = cfn_resources['Resources'][resource_name]

                if resource['Type'] == 'AWS::EC2::SecurityGroupIngress' or resource['Type'] == 'AWS::EC2::SecurityGroupEgress':
                    self.template['Resources'][resource_name] = resource

    def generate_outputs(self, cfn_resources, vpc):
        """Creates all of the CloudFormation Outputs for one vpc"""

        if 'Outputs' in cfn_resources and self.find_vpc(cfn_resources, vpc):
            if self.check_for_duplicates(self.template['Outputs'], cfn_resources['Outputs']):
                print '\nDuplicate Output found:\n' + json.dumps(self.template, indent=4) + '\n' + json.dumps(cfn_resources, indent=4)
                sys.exit(-5)

            for resource_name in cfn_resources['Outputs']:
                self.template['Outputs'][resource_name] = cfn_resources['Outputs'][resource_name]

    def find_vpc(self, cfn_resources, vpc):
        for resource_name in cfn_resources['Resources']:
            resource = cfn_resources['Resources'][resource_name]
            if ('Properties' in resource and
                    'VpcId' in resource['Properties'] and
                    (resource['Properties']['VpcId'] == self.ALL or
                     (len(resource['Properties']['VpcId'].split(',')) >= 1 and
                      vpc['VpcId'] in resource['Properties']['VpcId']))):
                return True
        return False

    def update_cfn_stack(self, stack_name, change_set, waiter_list):
        """Update or create a cfn stack containing the security groups for one vpc"""

        print '\nCreating change set for stack ' + stack_name if change_set else '\nUpdating cloudformation stack ' + stack_name

        print 'Uploading template to S3.'

        s3_url = self.upload_template_and_get_url(stack_name, json.dumps(self.template, indent=4, sort_keys=True))

        if s3_url is not None:
            kwargs = {
                'StackName': stack_name,
                'TemplateURL': s3_url,
                'RoleARN': self.cfn_role_arn
            }

            if self.stack_exists(stack_name):
                if not change_set:
                    call_method = boto3.client('cloudformation').update_stack
                else:
                    call_method = boto3.client('cloudformation').create_change_set
                    change_set_name = '-'.join([stack_name, str(time.time())])
                    change_set_name = change_set_name[:change_set_name.index('.')]
                    kwargs['ChangeSetName'] = change_set_name

                self.call_cloudformation(call_method=call_method,
                                         kwargs=kwargs,
                                         stack_name=stack_name,
                                         waiter_list=waiter_list)

            elif not change_set:
                self.call_cloudformation(call_method=boto3.client('cloudformation').create_stack,
                                         kwargs=kwargs,
                                         stack_name=stack_name,
                                         waiter_list=waiter_list)

            else:
                print '\nNo change set required, this would create a new stack'

        else:
            print 'ERROR: something went wrong with S3 upload'

    @staticmethod
    def upload_template_and_get_url(stack_name, template_body):
        """Uploads a template to S3 (used if a template is too large to upload directly)"""
        filename = stack_name + '.template'
        with open(filename, 'w') as open_file:
            open_file.write(template_body)

        upload_bucket = None

        response = boto3.client('s3').list_buckets()
        for bucket in response['Buckets']:
            if bucket['Name'].startswith('cf-templates-'):
                upload_bucket = bucket['Name']

        s3_response = None
        if upload_bucket is not None:
            boto3.client('s3').upload_file(filename, upload_bucket, filename)
            s3_response = boto3.client('s3').generate_presigned_url(ClientMethod='get_object',
                                                                    Params={
                                                                        'Bucket': upload_bucket,
                                                                        'Key': filename
                                                                    })
        else:
            print 'Template is too large to upload directly to Cloudformation, and unable to determine S3 bucket to upload to'

        os.remove(filename)

        return s3_response

    def call_cloudformation(self, call_method, kwargs, stack_name, waiter_list):
        try:
            response = call_method(**kwargs)

            print json.dumps(response, indent=4, default=self.json_default)

            if 'stack_update_complete' not in waiter_list:
                waiter_list['stack_update_complete'] = []
            waiter_list['stack_update_complete'].append(stack_name)

        except ClientError as e:
            if str(e) != self.NO_UPDATES:
                print json.dumps(self.template, indent=4, sort_keys=True)
                raise e

            else:
                print e

    @staticmethod
    def wait_for_stack_updates(waiter_list):
        for waiter_key in waiter_list:
            for stack_name in waiter_list[waiter_key]:
                try:
                    print '\nWaiting for ' + waiter_key + ' on ' + stack_name
                    waiter = boto3.client('cloudformation').get_waiter(waiter_key)
                    waiter.wait(StackName=stack_name)
                    print 'Done waiting on ' + waiter_key
                except WaiterError as e:
                    raise e


if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='Splits "=" separated param file into files named key with value')
    parser.add_argument('-c', '--change-set', required=False, action='store_true', help='Creates a change set instead of directly updating a stack')
    parser.add_argument('-r', '--region', required=False, help='Select AWS region to apply changes to')
    args = parser.parse_args()

    sec_group_updater = SecGroupUpdater(args.change_set)
© 2018 GitHub, Inc.
Help
Support
API
Training
Shop
Blog
About
GitHub Enterprise 2.14.2