REVIEW PLEASE

jstack.py

#!/usr/bin/env python3
'''
    JSTack Tool for Kubernetes environments

    Libraries:
        * connect_to_cluster()
        * list_pod_namespace()
        * collect_thread_dumps()
        * upload_file()
        * restart_application()

    Features:
        * Connect to K8s Clusters
        * Retrieve Pods(podname, application, namespace)
        * Collect Thread Dumps from pods
        * Upload a file to AWS S3 bucket
        * Restart applications non-statefulset on K8s

'''
import argparse
import datetime
import logging
import subprocess
import sys

from os import remove, environ
from kubernetes import client, config
from kubernetes.stream import stream
from kubernetes.client.rest import ApiException
from botocore.exceptions import ClientError

import boto3


def connect_to_cluster():
    '''
    Function to creates a stream connection to
    Kubernetes Cluster

    param:
        None

    return:
        Returns a stream connection object
    '''
    try:
        config.load_kube_config()
        con_v1 = client.CoreV1Api()
    except ApiException as error:
        return f"Please check the credentials of K8s cluster! \n" \
               f"For example: $ export KUBECONFIG=~/kubernetes-prod.conf\n" \
               f"ERROR: {error}"

    return con_v1


def list_pod_namespace(connection, apps):
    '''
    Function to retrieve a Podname and Namespace from an application

    param:
        connection: The stream connection object to
                    communicates with the K8s Cluster
        apps: List of applications

    return:
        Returns a dict {'podname': {'application': 'namespace'}}
    '''
    list_apps = list(apps.split(","))
    pod_dict = {}

    for app in list_apps:
        string_query = "app={0}".format(app.strip())
        list_of_pods = connection.list_pod_for_all_namespaces(label_selector=string_query,
                                                              watch=False)

        for i in list_of_pods.items:
            pod_dict[i.metadata.name] = {f"{app.strip()}": f"{i.metadata.namespace}"}

    return pod_dict


def collect_thread_dumps(connection, apps_dict):
    '''
    That function collects the thread dumps of a
    list of applications received

    param:
        connection: The stream connection object to
                    communicates with the K8s Cluster
        apps_dict: Dictionary generates by list_pod_namespace()

    return:
        Files names generated
    '''
    files = []

    for pod in apps_dict.keys():
        podname = pod

        for data in apps_dict[pod].items():
            namespace = data[1]

            pid_command = ['/bin/bash', '-c', 'pidof java']

            pid = stream(connection.connect_get_namespaced_pod_exec,
                         name=podname,
                         namespace=namespace,
                         command=pid_command,
                         stderr=True, stdin=False,
                         stdout=True, tty=False)

            jstack_command = ['jstack', '-l', '{0}'.format(pid)]

            jstack = stream(connection.connect_get_namespaced_pod_exec,
                            name=podname,
                            namespace=namespace,
                            command=jstack_command,
                            stderr=True, stdin=False,
                            stdout=True, tty=False)

            file_jstack = open("{0}-{1}.txt".format(namespace, podname), "w")
            file_jstack.write(jstack)
            file_jstack.close()
            files.append("{0}-{1}.txt".format(namespace, podname))

    return files


def upload_file(file_name, bucket, object_name=None, delete=False):
    '''
    Function to upload a file to an AWS S3 bucket

    param:
        file_name: File to upload to AWS S3 bucket
        bucket: Bucket name that will be used in upload process
        object_name: S3 object name. If not specified then file_name
                     is used plus the path pattern defined to each
                     application
        delete: Delete the file uploaded to AWS S3 bucket. if not
                specified will not delete if specified True will
                delete the file

    return:
        True if file was uploaded or False with error in case of a
        client issue
    '''
    if object_name is None:
        app = file_name.split('-')[1]
        folder = f"{datetime.datetime.now():%Y-%m-%d-%H}"
        path = f"jstack/{app}/{folder}"
        object_name = f"{path}/{file_name}"

    try:
        s3_client = boto3.client('s3')
        s3_client.upload_file(file_name, bucket, object_name)

        if delete is True:
            remove(file_name)

        return True

    except ClientError as error:
        return False, error


def restart_application(apps_dict):
    '''
    Function to restart the applications of Kubernetes

    param:
        apps_dict: Dictionary generates by list_pod_namespace()

    return:
        Returns a list with applications in what namespace was
        rebooted
    '''
    restarts = []

    for pod in apps_dict.keys():

        for data in apps_dict[pod].items():
            application = data[0]
            namespace = data[1]

            if application is not None or namespace is not None:
                try:
                    subprocess.run(["kubectl", "rollout",
                                    "restart", f"deployment/{application}",
                                    "-n", f"{namespace}"],
                                   stdout=subprocess.DEVNULL, check=True)

                    restarts.append(f"{application} in {namespace}")

                except subprocess.CalledProcessError as error:
                    return error

    return restarts


if __name__ == "__main__":
    BASE = "[%(asctime)s] [%(levelname)s] [%(name)s] "\
           "[%(funcName)s():%(lineno)s] [PID:%(process)d TID:%(thread)d] "\
           " %(message)s"
    TIME_CONTROL = "%d/%m/%Y %H:%M:%S"

    HANDLER = logging.StreamHandler(sys.stdout)
    HANDLER.setFormatter(logging.Formatter(BASE, TIME_CONTROL))

    LOG = logging.getLogger(__name__)
    LOG.addHandler(HANDLER)
    LOG.setLevel(logging.INFO)


    PARSER = argparse.ArgumentParser()

    PARSER.add_argument('-r', '--restart',
                        action='store_true',
                        help="restart the applications")

    PARSER.add_argument('-a', '--apps',
                        action='store',
                        required=True,
                        type=str,
                        help="list of applications to collect thread dumps." \
                             "e.g ./jstack -a \"cms, external\"")

    ARGS = PARSER.parse_args()

    MANDATORY_ENV_VARS = ["AWS_ACCESS_KEY_ID",
                          "AWS_SECRET_ACCESS_KEY",
                          "KUBECONFIG"]

    LOG.info('Checking credentials...')

    for var in MANDATORY_ENV_VARS:
        if var not in environ:
            raise EnvironmentError(f'Failed because {var} is not set.')

    if ARGS.apps:
        LOG.info('Connecting to K8s Cluster...')
        CONN = connect_to_cluster()
        LOG.info('Retrieving PODS...')
        PODS_DICT = list_pod_namespace(connection=CONN, apps=ARGS.apps)
        LOG.info('Collecting Thread Dumps...')
        GET_FILES = collect_thread_dumps(connection=CONN, apps_dict=PODS_DICT)

        LOG.info('Uploading files to AWS S3 bucket...')

        for handle_file in GET_FILES:
            upload = upload_file(file_name=handle_file,
                                 bucket='betpawa-dumps-gcheap',
                                 object_name=None,
                                 delete=True)

            LOG.info(' %s', handle_file)

    if ARGS.restart:
        LOG.info('Restarting Applications...')
        REST = restart_application(apps_dict=PODS_DICT)

        for retrieve_pod in REST:
            LOG.info(' restarted %s namespace', retrieve_pod)