manage.py

#!/usr/bin/env python3

import argparse
import os
import json
import sys
import time
from datetime import datetime
import subprocess
import yaml

class bcolors:
    HEADER = '\033[95m'
    OKBLUE = '\033[94m'
    OKGREEN = '\033[92m'
    WARNING = '\033[93m'
    FAIL = '\033[91m'
    ENDC = '\033[0m'
    BOLD = '\033[1m'
    UNDERLINE = '\033[4m'

def print_color(text: str, color:bcolors = bcolors.BOLD):
    """
    Utility method to print colored text to stdout.

    :param text:        The text to print
    :param color:       The bcolors to print text in (defaults to bold)
    :return:
    """
    print(f"{color}{text}{bcolors.ENDC}")

def run_shell(cmd: str):
    """
    Runs a shell command and prints command to stdout before
    running so user can see what was run

    :param cmd:     The shell command to run
    :return:
    """
    print_color(f"** RUNNING: {cmd}")
    os.system(cmd)

# Subcommand options

def start(args):
    """
    Start a GKE Cluster with Zebrium's demo environment deployed.
    """
    print_color(f"Starting GKE cluster in project {args.project} with name {args.name} in zone {args.zone}", bcolors.OKBLUE)

    # Ensure GCloud SDK is up to date
    run_shell("gcloud components update")

    # Set GCloud project
    run_shell(f"gcloud config set project \"{args.project}\"")

    # Spinup cluster
    run_shell(f"gcloud container clusters create {args.name} --zone {args.zone} --cluster-version 1.14.10-gke.17 --machine-type n1-standard-2 --no-enable-autoupgrade")

    # Get kubectl credentials
    run_shell(f"gcloud container clusters get-credentials {args.name} --zone {args.zone}")

    print_color("\nGKE Cluster Running with following nodes:\n")
    run_shell(f"kubectl get nodes")

    # Deploy Zebrium Collector using Helm
    ze_deployment_name = "zebrium-k8s-demo"
    ze_logs_url = "https://zapi03.zebrium.com"
    ze_stats_url = "https://zapi03.zebrium.com/stats/api/v1/zstats"
    run_shell("sleep 90") # Wait 1.5 mins for cluster to finish setting up fully
    run_shell("kubectl create namespace zebrium")
    run_shell(f"helm install zlog-collector --namespace zebrium --set zebrium.deployment={ze_deployment_name},zebrium.collectorUrl={ze_logs_url},zebrium.authToken={args.key} --repo https://raw.githubusercontent.com/zebrium/ze-kubernetes-collector/master/charts zlog-collector")
    # Install Prometheus collector (EXPERIMENTAL)
    run_shell(f"helm install zstats-collector --namespace zebrium --set zebrium.deployment={ze_deployment_name},zebrium.collectorUrl={ze_stats_url},zebrium.authToken={args.key} --repo https://raw.githubusercontent.com/zebrium/ze-stats/master/charts zstats")

    # Deploy all demo apps
    run_shell("kubectl create -f ./deploy/sock-shop.yaml")
    run_shell("kubectl create -f ./deploy/random-log-counter.yaml")

    # Deploy kafka demo app
    run_shell("kubectl create namespace kafka")
    run_shell("helm repo add confluentinc https://confluentinc.github.io/cp-helm-charts/")
    run_shell("helm repo update")
    run_shell("helm install kafka-cluster --set cp-schema-registry.enabled=false,cp-kafka-rest.enabled=false,cp-kafka-connect.enabled=false,cp-control-center.enabled=false,cp-ksql-server.enabled=false confluentinc/cp-helm-charts --namespace=kafka")
    run_shell('kubectl annotate sts/kafka-cluster-cp-kafka litmuschaos.io/chaos="true" -n kafka')

    # Deploy Litmus ChaosOperator to run Experiments that create incidents
    run_shell("kubectl apply -f https://litmuschaos.github.io/litmus/litmus-operator-v1.6.1.yaml")

    # Install Litmus Experiments - TEMP Workaround to set experiment versions until Chaos Hub supports in URL
    run_shell("curl -sL https://github.com/litmuschaos/chaos-charts/archive/1.6.1.tar.gz -o litmus.tar.gz")
    run_shell("tar -zxvf litmus.tar.gz")
    run_shell("rm litmus.tar.gz")
    run_shell("find chaos-charts-1.6.1 -name experiments.yaml | grep generic | xargs kubectl apply -n sock-shop -f")
    run_shell("find chaos-charts-1.6.1 -name experiments.yaml | grep kafka | xargs kubectl apply -n kafka -f")
    #run_shell("kubectl create -f https://hub.litmuschaos.io/api/chaos?file=charts/generic/experiments.yaml -n sock-shop")
    #run_shell("kubectl create -f https://hub.litmuschaos.io/api/chaos?file=charts/kafka/experiments.yaml -n kafka")

    # Create the chaos serviceaccount with permissions needed to run the generic K8s experiments
    run_shell("kubectl create -f ./deploy/litmus-rbac.yaml")

    # Get ingress IP address
    run_shell("sleep 60")  # Wait 1 min for ingress to finish setting up
    print_color("\nIngress Details:\n", bcolors.UNDERLINE)
    run_shell("kubectl get ingress basic-ingress --namespace=sock-shop")

    try:
        ingress_ip = \
        json.loads(os.popen('kubectl get ingress basic-ingress --namespace=sock-shop -o json').read())["status"][
            "loadBalancer"]["ingress"][0]["ip"]
        print_color(f"\nYou can access the web application in a few minutes at: http://{ingress_ip}\n\n")
    except:
        print_color("Ingress still being setup. Use the following command to get the IP later:", bcolors.WARNING)
        print_color("\tkubectl get ingress basic-ingress --namespace=sock-shop", bcolors.WARNING)

    print_color("***************************************************************************************************", bcolors.WARNING)
    print_color(f"* {datetime.now().strftime('%Y-%m-%d %H:%M:%S')} Finished creating cluster.", bcolors.WARNING)
    print_color("* Please wait at least 15 minutes for environment to become fully initalised.")
    print_color("* The ingress to access the web application from your browser can take at least 5 minutes to create.", bcolors.WARNING)
    print_color("*", bcolors.WARNING)
    print_color("*", bcolors.WARNING)
    print_color("* IMPORTANT: To reliably detect Chaos experiment incidents you must reduce the Refractory Period for your account to 10 minutes.", bcolors.WARNING)
    print_color("* You can do this under your Advanced account settings at https://portal03.zebrium.com/Settings/advanced.",bcolors.WARNING)
    print_color("*", bcolors.WARNING)
    print_color("***************************************************************************************************\n\n", bcolors.WARNING)

def stop(args):
    """
    Shutdown the GKE Cluster with Zebrium's demo environment deployed.
    """
    print_color(f"Stopping GKE cluster in project {args.project} with name {args.name} in zone {args.zone}", bcolors.OKBLUE)

    # Set GCloud project
    run_shell(f"gcloud config set project \"{args.project}\"")

    # Stop cluster
    run_shell(f"gcloud container clusters delete {args.name} --zone {args.zone}")

class ExperimentResult(object):
    """
    Holds Experiment Result
    """

    def __init__(self, name:str, status:str, startTime:datetime):
        self.name = name
        self.status = status
        self.startTime = startTime

def run_experiment(experiment: str, delay: int = 0):
    """
    Run a specific experiment

    :param experiment:  The name of the experiment as defined in the YAML, i.e. container-kill
    :param ramp_time:   The number of seconds to delay experiment after setup to avoid confusing setup events with experiment events in Zebrium
    :return:            ExperimentResult object with results of experiment
    """
    print_color("***************************************************************************************************", bcolors.OKBLUE)
    print_color(f"* {datetime.now().strftime('%Y-%m-%d %H:%M:%S')} Experiment: {experiment}", bcolors.OKBLUE)
    print_color("***************************************************************************************************", bcolors.OKBLUE)

    experiment_file = experiment + ".yaml"

    # Set namespace to check
    with open(f"./litmus/{experiment_file}") as f:
        spec = yaml.load(f, Loader=yaml.FullLoader)
        result_name = spec['metadata']['name']
        namespace = spec['metadata']['namespace']

        # Create temp file with updated RAMP_TIME
        if (delay > 0):
            spec['spec']['experiments'][0]['spec']['components']['env'].append({'name': 'RAMP_TIME', 'value': str(delay)})
        with open(r"temp.yaml", 'w') as temp:
            yaml.dump(spec, temp)

    print_color(f"Running Litmus ChaosEngine Experiment {experiment_file} in namespace {namespace} with delay {delay} seconds...")
    print_color(f"Deploying {experiment_file}...")
    run_shell(f"kubectl delete chaosengine {result_name} -n {namespace}")
    run_shell(f"kubectl create -f temp.yaml -n {namespace}")

    # Check status of experiment execution
    startTime = datetime.now()
    print_color(f"{startTime.strftime('%Y-%m-%d %H:%M:%S')} Running experiment...")
    expStatusCmd = "kubectl get chaosengine " + result_name + " -o jsonpath='{.status.experiments[0].status}' -n " + namespace
    run_shell(expStatusCmd)
    logs_cmd = f"kubectl logs --since=10s -l name={experiment} -n {namespace}"
    print(f"\n{bcolors.OKGREEN}//** Experiment Logs ({logs_cmd}) **//\n\n")
    try:
        while subprocess.check_output(expStatusCmd, shell=True).decode('unicode-escape') != "Execution Successful":
            os.system(logs_cmd)
            os.system("sleep 10")

        print(f"\n\n//** End of Experiment Logs **//{bcolors.ENDC}\n")

        # View experiment results
        run_shell(f"kubectl describe chaosresult {result_name}-{experiment} -n {namespace}")

        # Delete temp file
        run_shell('rm temp.yaml')
    except:
        print_color("User has cancelled script execution.", bcolors.FAIL)
        sys.exit(2)

    # Store Experiment Result
    status = subprocess.check_output("kubectl get chaosresult " + result_name + "-" + experiment + " -n " + namespace + " -o jsonpath='{.spec.experimentstatus.verdict}'", shell=True).decode('unicode-escape')
    return ExperimentResult(experiment, status, startTime)

def test(args):
    """
    Run Litmus ChaosEngine Experiments inside Zebrium's demo environment.
    Each experiment is defined under its own yaml file under the /litmus directory. You can run
    a specific experiment by specifying a test name that matches one of the yaml file names in the directory
    but by default all '*' experiments will be run with 20 minute wait period between each experiment
    to ensure Zebrium doesn't cluster the incidents together into one incident
    """
    experiments = sorted(os.listdir('./litmus'))
    experiment_results = []

    if args.test == '*':
        # Run all experiments in /litmus directory with wait time between them
        print_color(f"Running all Litmus ChaosEngine Experiments with {args.wait} mins wait time between each one...")
        lstindex = len(experiments)
        for experiment_file in experiments:
            result = run_experiment(experiment_file.replace('.yaml', ''), args.delay)
            experiment_results.append(result)
            print_color(f"{datetime.now().strftime('%Y-%m-%d %H:%M:%S')} Waiting {args.wait} mins before running next experiment...", bcolors.WARNING)
            lstindex -= 1
            if lstindex != 0:
                time.sleep(args.wait * 60)
    else:
        # Check experiment exists
        experiment_file = args.test + ".yaml"
        if experiment_file in experiments:
            result = run_experiment(args.test, args.delay)
            experiment_results.append(result)
        else:
            print_color(f"ERROR: {experiment_file} not found in ./litmus directory. Please check the name and try again.", bcolors.FAIL)
            sys.exit(2)

    # Print out experiment result summary
    print_color("***************************************************************************************************", bcolors.OKBLUE)
    print_color("* Experiments Result Summary", bcolors.OKBLUE)
    print_color("***************************************************************************************************\n", bcolors.OKBLUE)
    headers = ["#", "Start Time", "Experiment", "Status"]
    row_format = "{:>25}" * (len(headers) + 1)
    print_color(row_format.format("", *headers), bcolors.OKBLUE)
    i = 1
    for result in experiment_results:
        print_color(row_format.format("", str(i), result.startTime.strftime('%Y-%m-%d %H:%M:%S'), result.name, result.status), bcolors.OKBLUE)
        i += 1
    print("\n")

def list(args):
    """
    List all available Litmus Chaos Experiments available in this repository
    """
    experiments = sorted(os.listdir('./litmus'))
    print_color("Available Litmus Chaos Experiments:\n\n")
    i = 1
    for experiment_file in experiments:
        print_color(f"\t{i}. {experiment_file.replace('.yaml', '')}")
        i += 1

if __name__ == "__main__":

    # Add command line arguments
    parser = argparse.ArgumentParser(description='Spin up Zebrium Demo Environment on Kubernetes.')
    subparsers = parser.add_subparsers()

    # Start command
    parser_start = subparsers.add_parser("start", help="Start a GKE Cluster with Zebrium's demo environment deployed.")
    parser_start.add_argument("-p", "--project", type=str,
                        help="Set GCloud Project to spin GKE cluster up in")
    parser_start.add_argument("-z", "--zone", type=str, default="us-central1-a",
                        help="Set GCloud Zone to spin GKE cluster up in")
    parser_start.add_argument("-n", "--name", type=str, default="zebrium-k8s-demo",
                        help="Set GKE cluster name")
    parser_start.add_argument("-k", "--key", type=str,
                        help="Set Zebrium collector key for demo account")
    parser_start.set_defaults(func=start)

    # Test command
    parser_test = subparsers.add_parser("test", help="Run Litmus ChaosEngine Experiments inside Zebrium's demo environment.")
    parser_test.add_argument("-t", "--test", type=str, default="*",
                             help="Name of test to run based on yaml file name under /litmus folder. '*' runs all of them with wait time between each experiement.")
    parser_test.add_argument("-w", "--wait", type=int, default=11,
                             help="Number of minutes to wait between experiments. Defaults to 11 mins to avoid Zebrium clustering incidents together.")
    parser_test.add_argument("-d", "--delay", type=int, default=660,
                             help="Delay time in seconds between setting up experiment and running it. Defaults to 660 seconds.")
    parser_test.set_defaults(func=test)

    # List Tests Command
    parser_list = subparsers.add_parser("list", help="List all available Litmus ChaosEngine Experiments available to run.")
    parser_list.set_defaults(func=list)

    # Stop command
    parser_stop = subparsers.add_parser("stop", help="Shutdown the GKE Cluster with Zebrium's demo environment deployed.")
    parser_stop.add_argument("-p", "--project", type=str,
                        help="Set GCloud Project to spin GKE cluster up in")
    parser_stop.add_argument("-z", "--zone", type=str, default="us-central1-a",
                        help="Set GCloud Zone to spin GKE cluster up in")
    parser_stop.add_argument("-n", "--name", type=str, default="zebrium-k8s-demo",
                        help="Set GKE cluster name")
    parser_stop.set_defaults(func=stop)

    args = parser.parse_args()
    args.func(args)