lw_add_kubernetes.py

#!/usr/bin/env python3

import csv
import sys
import os
import argparse


def csv_dictionary(filename):
    clusters = []
    try:
        with open(filename, 'r') as data:
            for line in csv.DictReader(data):
                if "NAME" in line and "RESOURCE GROUP" in line and "Subscription ID" in line:
                    clusters.append(line)
        if len(clusters) < 1:
            print("ERROR: The CSV file must container the following three headers: 'NAME', 'RESOURCE GROUP', and 'Subscription ID'")
            sys.exit(3)
    except FileNotFoundError:
        print("ERROR: {} file was not found.".format(filename))
        sys.exit(2)
    return clusters

def add_lw_kubernetes(subscription_id, resource_group, cluster_name, template):
    if args.debug:
        print("---\nSubscription ID: {0}, Resource Group: {1}, Cluster: {2}".format(subscription_id, resource_group, cluster_name))
        print("Create: {}".format(args.create))
        print("Execute: {}".format(args.execute))

    create_dir_terraform(cluster_name, template)
    execute_terraform(subscription_id, resource_group, cluster_name)


def execute_terraform(subscription_id, resource_group, cluster_name):
    try:
        if os.path.isfile('{}/terraform.tfstate'.format(cluster_name)):
            if args.debug: print("tfstate file exists for cluster '{}', skipping.".format(cluster_name))
            return
    except:
        print("Error checking for tfstate file in '{}'".format(cluster_name))
        sys.exit(7)

    if not args.execute:
        if args.debug: print("'{}' directory would have terraform executed.".format(cluster_name))
        return

    ret = os.system("az account set --subscription {}".format(subscription_id))
    if ret:
        print("Failed to set subscription '{}'".format(subscription_id))
        sys.exit(8)

    ret = os.system("az aks get-credentials --resource-group {0} --name {1}".format(resource_group, cluster_name))
    if ret:
        print("Failed to get credentials for resource group '{0}' and cluster '{1}'".format(resource_group, cluster_name))
        sys.exit(9)

    ret = os.system("cd {} && terraform init && terraform plan && terraform apply && cd ..".format(cluster_name))
    # ret = os.system("cd {} && terraform init && terraform plan && terraform apply --auto-approve && cd ..".format(cluster_name))
    if ret:
        print("Failed to execute terraform init/plan/apply in '{}'".format(cluster_name))
        sys.exit(10)

    print("Verifying daemonset deployed ---")
    os.system("kubectl get pods -l name=lacework -o=wide --all-namespaces")


def create_dir_terraform(cluster_name, template):
    if not os.path.isfile(template):
        print("Terraform file '{}' does not exist".format(template))
        sys.exit(4)
    if os.path.isdir(cluster_name) and not args.create:
        if args.debug: print("Directory '{}' already exists, would skip directory creation.".format(cluster_name))
        return
    if os.path.isdir(cluster_name) and args.create:
        if args.debug: print("Directory '{}' already exists, skipping directory creation.".format(cluster_name))
        return
    elif not args.create:
        return

    try:
        os.makedirs(cluster_name)
        print("Created directory for cluster: {}".format(cluster_name))
    except:
        print("Could not create directory '{}'. Do you have write access to current directory?".format(cluster_name))
    try:
        with open(template, 'r') as data:
            terraform = data.read()
    except:
        print("Error reading in terraform file '{}'. Do you have read permissions?".format(template))
        sys.exit(5)
    terraform = terraform.replace("[CLUSTER NAME]", cluster_name)
    try:
        with open('{}/main.tf'.format(cluster_name), 'w') as data:
            data.write(terraform)
    except:
        print("Failure creating {}/main.tf".format(cluster_name))
        sys.exit(6)


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description='You need to specify the --execute or --create flags to take any action.')
    parser.add_argument('-f', '--file', type=csv_dictionary, required=True, help='This must be a CSV file with these headers: NAME, RESOURCE GROUP, Subscription ID')
    parser.add_argument('-t', '--terraform', nargs='?', required=False, type=str, default="main.tf.template", help='This will default to main.tf.template if not specified. [CLUSTER NAME] placeholders are replaced in this file.')
    parser.add_argument('-c', '--create', action='store_true', help='Specify this to create a directory for each cluster and create customized main.tf for each directory')
    parser.add_argument('-e', '--execute', action='store_true', help='WARNING: This will execute terraform commands in supplied clusters directory as long as terraform.state does not exist')
    parser.add_argument('-d', '--debug', action='store_true')
    args = parser.parse_args()
    for item in args.file:
        add_lw_kubernetes(item['Subscription ID'], item['RESOURCE GROUP'], item['NAME'], args.terraform)