Automate VPC Flow logs with boto3

Hello, I want to share how can I create a script to enable VPC flow logs for AWS accounts. 
First, I need to install boto3

pip install boto3

To successfully finish the task I need to complete 6 steps. 
Get VPC’s list, get or create log group, role arn, policy and enable flow logs.
As usual, I start from import and boto3 client initialization:

import boto3
import json
from botocore.exceptions import ClientError

ec2 = boto3.resource("ec2")
ec2_client = boto3.client("ec2")
log_client = boto3.client("logs")
iam_client = boto3.client("iam")

Then I will start from the first function. In this function. Here I am getting all VPC’s and return their ID’s.

def get_all_vpcs():
    return [vpc.id for vpc in list(ec2.vpcs.all())]

Next one is to get a log group name.  In this function, I need to check if the group already exist and use it if exist. But if the group not exist I need to create a new one.

def get_flow_log_group(pattern="flow-logs-group", retention_perion=14):
    # Check if log group exist. Use it.
    for log_group in log_client.describe_log_groups()["logGroups"]:
        if pattern in log_group["logGroupName"]:
            print(
                f"Found existing group {log_group['logGroupName']}")
            return log_group["logGroupName"]
    else:
        # Create new log group, if no existing groups
        print("Log group not found. Creating new one!")
        log_client.create_log_group(
            logGroupName=f"{pattern}-{log_client.meta.region_name}"
        )
        log_group_name = log_client.describe_log_groups(
            logGroupNamePrefix=pattern
        )["logGroups"][0]["logGroupName"]
        log_client.put_retention_policy(
            logGroupName=log_group_name,
            retentionInDays=retention_perion
        )
        print(f"New log group is created - {log_group_name}")
        return log_group_name

And the third method is to get a role ARN if it exists. Otherwise, I need to check if FlowLogs policy exists, and if not exist create a new one and then create a new role. I split policy and role into separate functions. And also the policy required a policy document and role trust document. To handle this I am using the json module.

policy_trust_document = json.dumps({
    "Version": "2012-10-17",
    "Statement": [
        {
               "Sid": "",
            "Effect": "Allow",
            "Principal": {
                "Service": "vpc-flow-logs.amazonaws.com"
            },
            "Action": "sts:AssumeRole"
        }
    ]
})

flow_logs_policy_document = json.dumps({
    "Version": "2012-10-17",
    "Statement": [
        {
            "Action": [
                "logs:CreateLogGroup",
                "logs:CreateLogStream",
                "logs:PutLogEvents",
                "logs:DescribeLogGroups",
                "logs:DescribeLogStreams"
            ],
            "Effect": "Allow",
            "Resource": "*"
        }
    ]
})

def get_flow_logs_policy(pattern="flow-logs-policy"):
    for policy in iam_client.list_policies(Scope="Local")["Policies"]:
        if pattern in policy["PolicyName"]:
            print(f"Found existing policy {policy['Arn']}.")
            return policy["Arn"]
    else:
        print("Policy not found. Creating new one!")
        flow_logs_policy = iam_client.create_policy(
            PolicyName=pattern,
            PolicyDocument=flow_logs_policy_document,
            Description="Policy for Flow Logs"
        )
        print(f"New policy is created - {flow_logs_policy['Policy']['Arn']}")
        return flow_logs_policy["Policy"]["Arn"]


def get_flow_log_role_arn(pattern="flow-logs-role"):
    for role in iam_client.list_roles()["Roles"]:
        if pattern in role["RoleName"]:
            print(f"Found existing role {role['Arn']}.")
            return role["Arn"]
    else:
        print("Role not found. Creating new one!")
        new_role = iam_client.create_role(
            RoleName="flow-logs-role",
            AssumeRolePolicyDocument=policy_trust_document,
            PermissionsBoundary=get_flow_logs_policy()

        )
        print(f"New role is created - {new_role['Role']['Arn']}")
        return new_role["Role"]["Arn"]

The final function is to enable flow logs. In this function will be one API call to create flow logs with all required parameters. And catch the error if FlowLogAlreadyExists, if exist print message if no create.

def enable_flow_logs(vpc_id, flow_logs_role_arn, log_group):
    try:
        print(
            f"Trying to enable flow logs for {vpc_id},"
            f" using {log_group} log group and role {flow_logs_role_arn}")
        ec2_client.create_flow_logs(
            DeliverLogsPermissionArn=flow_logs_role_arn,
            LogGroupName=log_group,
            ResourceIds=[vpc_id],
            ResourceType="VPC",
            TrafficType="ALL",
            LogDestinationType="cloud-watch-logs"
        )
    except ClientError as e:
        if e.response['Error']['Code'] == "FlowLogAlreadyExists":
            print(f"Flow logs is already enabled for {vpc_id}\n")
    else:
        print(f"Flow logs is successfully enabled for {vpc_id}\n")

What left is to execute all these functions. At first, get all required parameters, role arn, log group VPC’s. Then loop through VPC and enable flow logs.

if __name__ == "__main__":
    role_arn = get_flow_log_role_arn()
    log_group = get_flow_log_group()
    vpcs = get_all_vpcs()

    for vpc in vpcs:
        enable_flow_logs(vpc, role_arn, log_group)

Save the file, execute and you should see the similar output:

python.exe -u "e:\vpc_flow_logs.py"
Role not found. Creating new one!
Policy not found. Creating new one!
New policy is created - arn:aws:iam::1234:policy/flow-logs-policy
New role is created - arn:aws:iam::1234:role/flow-logs-role
Log group not found. Creating new one!
New log group is created - flow-logs-group-us-east-1
Trying to enable flow logs for vpc-e5543380, using flow-logs-group-us-east-1 log group and role arn:aws:iam::1234:role/flow-logs-role
Flow logs is successfully enabled for vpc-e5543380

The full code can be found in my GitHub