def lambda_handler()

in lambda/isolate-ec2-security-group.py [0:0]

• 46 lines of code
• 16 McCabe index (conditional complexity)

def lambda_handler(event, context):
    
    security_team_sg = None
    sec_groups = []
    vpc_id = None


    logger.info(event)
    for findings in event['detail']['findings']:
        for res in findings['Resources']:
            id_arn = (res['Id'])
            print(id_arn)
            instance_id = id_arn.split("/")[1]
            print(instance_id)

    instance_info = ec2client.describe_instances(
        InstanceIds=[instance_id])


    # Get all the Security Groups for the instance
    for res in instance_info['Reservations']:
        for ins in res['Instances']:
            for sg in ins['SecurityGroups']:
                sec_groups.append(sg['GroupId'])
                print(sec_groups)

                logger.info(
                    'Modifying Security Groups for Instance ID: %s', instance_id)
                logger.info('## Original Security Groups ##')
                logger.info(sec_groups)

    for res in instance_info['Reservations']:
        for ins in res['Instances']:
            vpc_id = ins['VpcId']
            print(vpc_id)

    try:
        all_security_groups = ec2client.describe_security_groups()
    except ClientError as e:
        logger.exception(e)

    
    for sg in all_security_groups['SecurityGroups']:
        if sg['GroupName'] == security_team_name and sg['VpcId'] == vpc_id:
            security_team_sg = sg['GroupId']

    # modify_instance_attribute will wipe out all the existing security groups with the new one
    logger.info('Modified instance to now have security group: %s',
                security_team_sg)

    if security_team_sg == None:
        security_team_sg = create_security_sg(vpc_id)  
    
    if security_team_sg == None:
        return 500
        
    try:
        response = ec2client.modify_instance_attribute(
            InstanceId=instance_id, Groups=[security_team_sg])
    except ClientError as e:
        logger.exception(e)
        return 500

    return response['ResponseMetadata']['HTTPStatusCode']