def enable_flow_logs()

in source/remediation_runbooks/scripts/EnableVPCFlowLogs.py [0:0]

• 59 lines of code
• 17 McCabe index (conditional complexity)

def enable_flow_logs(event, context):
    """
    remediates CloudTrail.2 by enabling SSE-KMS
    On success returns a string map
    On failure returns NoneType
    """
    max_retries = event.get('retries', 12) # max number of waits for actions to complete.
    wait_interval = event.get('wait', 5) # how many seconds between attempts

    boto_config_args = {
        'retries': {
            'mode': 'standard'
        }
    }

    boto_config = Config(**boto_config_args)

    if 'vpc' not in event or 'remediation_role' not in event or 'kms_key_arn' not in event:
        exit('Error: missing vpc from input')

    logs_client = connect_to_logs(boto_config)
    ec2_client = connect_to_ec2(boto_config)
    
    kms_key_arn = event['kms_key_arn'] # for logs encryption at rest
    
    # set dynamic variable for CW Log Group for VPC Flow Logs
    vpc_flow_loggroup = "VPCFlowLogs/" + event['vpc']        
    # create cloudwatch log group
    try:
        logs_client.create_log_group(
            logGroupName=vpc_flow_loggroup,
            kmsKeyId=kms_key_arn
        )
    except ClientError as client_error:
        exception_type = client_error.response['Error']['Code']

        if exception_type in ["ResourceAlreadyExistsException"]:
            print(f'CloudWatch Logs group {vpc_flow_loggroup} already exists')
        else:
            exit(f'ERROR CREATING LOGGROUP {vpc_flow_loggroup}: {str(exception_type)}')
            
    except Exception as e:
        exit(f'ERROR CREATING LOGGROUP {vpc_flow_loggroup}: {str(e)}')

    # wait for CWL creation to propagate
    wait_for_loggroup(logs_client, wait_interval, max_retries, vpc_flow_loggroup)

    # create VPC Flow Logging
    try:
        ec2_client.create_flow_logs(
            DryRun=False,
            DeliverLogsPermissionArn=event['remediation_role'],
            LogGroupName=vpc_flow_loggroup,
            ResourceIds=[event['vpc']],
            ResourceType='VPC',
            TrafficType='REJECT',
            LogDestinationType='cloud-watch-logs'
        )
    except ClientError as client_error:
        exception_type = client_error.response['Error']['Code']

        if exception_type in ["FlowLogAlreadyExists"]:
            return {
                "response": {
                    "message": f'VPC Flow Logs for {event["vpc"]} already enabled',
                    "status": "Success"
                }
            }
        else:
            exit(f'ERROR CREATING LOGGROUP {vpc_flow_loggroup}: {str(exception_type)}')
    except Exception as e:
        exit(f'create_flow_logs failed {str(e)}')

    # wait for Flow Log creation to propagate. Exits on timeout (no need to check results)
    wait_for_flowlogs(ec2_client, wait_interval, max_retries, vpc_flow_loggroup)

    # wait_for_flowlogs will exit if unsuccessful after max_retries * wait_interval (60 seconds by default)
    return {
        "response": {
            "message": f'VPC Flow Logs enabled for {event["vpc"]} to {vpc_flow_loggroup}',
            "status": "Success"
        }
    }