#!/usr/share/venvs/netq-agent/bin/python
#
# Copyright (c) 2017-2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# LicenseRef-NvidiaProprietary
#
# NVIDIA CORPORATION, its affiliates and licensors retain all intellectual
# property and proprietary rights in and to this material, related
# documentation and any modifications thereto. Any use, reproduction,
# disclosure or distribution of this material and related documentation
# without an express license agreement from NVIDIA CORPORATION or
# its affiliates is strictly prohibited
#

import subprocess
import os
import sys
import socket
from time import localtime, sleep
from netq_lib.common.utils import subprocess_wrapper, is_sonic, get_logger, is_cl3, is_cl4


DATE = '%d%d%d' % (localtime().tm_year, localtime().tm_mon, localtime().tm_mday)
TIME = '%d%d%d' % (localtime().tm_hour, localtime().tm_min, localtime().tm_sec)

SUPPORT_PATH = '/var/support/'
EXTENSION = '.txz' if not is_sonic() else '.tar.gz'
TAR_OPTIONS = '-cJvf' if not is_sonic() else '-czvf'
OPTAPREFIX = 'opta'
NETQPREFIX = 'netq'
BASENAME = 'support_%s_%s_%s' % (socket.gethostname(), DATE, TIME)
NETQ_AGENT_PY3_4_VENV_PATH = '/usr/share/venvs/netq-agent/lib/python3.4/site-packages/netq_agent'
NETQ_AGENT_PY3_6_VENV_PATH = '/usr/share/venvs/netq-agent/lib/python3.6/site-packages/netq_agent'
NETQ_AGENT_PY3_7_VENV_PATH = '/usr/share/venvs/netq-agent/lib/python3.7/site-packages/netq_agent'
NETQ_AGENT_PY3_8_VENV_PATH = '/usr/share/venvs/netq-agent/lib/python3.8/site-packages/netq_agent'
NETQ_CLI_CLIENT_PATH = '/tests/netq_agent/agent_cli_client.py'
NETQ_AGENT_VENV_PATH = ''
NETQ_CLIENT_CLI=''
NETQ_CMD_SCRIPT = 'netq_script.sh'
OPTA_FILE = '/etc/app-release'
COPY_DB_LOGS = 'db_logs'

RC_SUCCESS = 0
RC_FAIL = 1

CMD_TIMEOUT = 60
FAILED_COMMANDS_FILE = 'failed-commands.txt'

OPTA_FILES = {
    '/var/log/': ['syslog', 'auth', 'netq-agent', 'netqd', 'boot', 'kern', 'containers', 'dpkg', 'kubelet', 'apt'],
    '/etc/': ['os-release', 'app-release', 'resolv.conf', 'netq', 'kubernetes', 'hosts', 'netplan'],
    '/home/cumulus/': ['kube'],
    '/var/run/': ['netq'],
}

PODS_FILES ={
    'cassandra': 'var/log/cassandra'
}

CUMULUS_FILES = {
    '/var/log/': ['syslog', 'netq-agent', 'netqd', 'netq-model', 'boot', 'dpkg', 'auth'],
    '/etc/': ['os-release', 'netq'],
    '/var/run/': ['netq'],
}

OPTA_CMDS = {
    'kubectl-describe.txt': ['kubectl', 'describe', 'pods', '--all-namespaces'],
    'kubectl-getall.txt': ['kubectl', 'get', 'all', '--all-namespaces'],
    'kubectl-describe-node.txt': ['kubectl', 'describe', 'node'],
    'docker-images.txt': ['docker', 'image', 'ls'],
    'docker-info.txt': ['docker', 'info'],
    'docker-ps.txt': ['docker', 'ps', '-a'],
    'opta-check.txt': ['opta-check'],
    'netq-show-agents.txt': ['netq', 'show', 'agents'],
    'netq-show-lldp.txt': ['netq', 'show', 'lldp'],
    'netq-show-inv-os.txt': ['netq', 'show', 'inventory', 'os'],
    'netq-show-clag.txt': ['netq', 'show', 'clag'],
    'netq-show-bgp.txt': ['netq', 'show', 'bgp'],
    'netq-show-evpn.txt': ['netq', 'show', 'evpn'],
    'netq-show-ospf.txt': ['netq', 'show', 'ospf'],
    'netq-show-opta-health.txt': ['netq', 'show', 'opta-health'],
    'netq-app-admin.txt': "docker ps -aq -f label=io.kubernetes.container.name=netq-app-admin",
    'topology-json.txt': ['topology-json'],
    'dmesg.txt': ['dmesg'],
    'lspci.txt': ['lspci'],
    'ls-mnt-r.txt': ['ls', '-lR', '/mnt/'],
    'chrony.txt': ['chronyc', 'tracking'],
    'kubectl-describe-configmap.txt': ['kubectl', 'describe', 'configmap'],
    'kubectl-describe-ingress.txt': ['kubectl', 'describe', 'ingress'],
    'kubectl-describe-clusterrole.txt': ['kubectl', 'describe', 'clusterrole'],
    'kubectl-describe-role.txt': ['kubectl', 'describe', 'role'],
    'kubectl-describe-serviceaccount.txt': ['kubectl', 'describe', 'serviceaccount'],
    'kubectl-get-serviceaccount.txt': ['kubectl', 'get', 'serviceaccount'],
    'kubectl-get-role.txt': ['kubectl', 'get', 'role'],
    'kubectl-get-clusterrole.txt': ['kubectl', 'get', 'clusterrole'],
    'kubectl-get-crd.txt': ['kubectl', 'get', 'crd'],
    'kubectl-get-secret.txt': ['kubectl', 'get', 'secret'],
    'kubectl-get-configmap.txt': ['kubectl', 'get', 'configmap'],
    'kubectl-get-ingress.txt': ['kubectl', 'get', 'ingress'],
    'opta-info.txt': ['/usr/sbin/opta-info.py'],
    'cassandra-tables-disk-usage.txt': ['du','-h', '/mnt/cassandra'],
    'cassandra-describe-cluster.txt': ['bash', '/usr/sbin/netq-cassandra-information-dump.sh']
}

LINUX_CMDS = {
    'df-h.txt': ['df', '-h'],
    'lsmod.txt': ['lsmod'],
    'top-n1.txt': ['top', '-b', '-n1'],
    'ss-pan.txt': ['ss', '-pan'],
    'free-lm.txt': ['free', '-lm'],
    'ps-aux.txt': ['ps', 'aux'],
    'systemctl.txt': ['systemctl', '-l', '--type', 'service', '--all'],
    'ip-link-show.txt': ['ip', '-d', 'link', 'show'],
    'ip-addr-show.txt': ['ip', 'addr', 'show'],
    'ip-route-show-table-all.txt': ['ip', 'route', 'show', 'table', 'all'],
    'ip-neighbor-show.txt': ['ip', 'neighbor', 'show'],
    'brctl-show.txt': ['brctl', 'show'],
    'bridge-fdb-show.txt': ['bridge', 'fdb', 'show'],
    'bridge-vlan-show.txt': ['bridge', 'vlan', 'show'],
    'ntpq.txt': ['ntpq', '-pn'],
    'timedatectl.txt': ['timedatectl'],
    'netq-agent-cli.txt': ['python', NETQ_CLIENT_CLI],
    'proc_meminfo.txt': ['cat', '/proc/meminfo'],
    'proc_cpuinfo.txt': ['cat', '/proc/cpuinfo'],
    'journalctl_docker.txt': ['journalctl', '--no-pager', '-u', 'docker'],
    'journalctl_kubelet.txt': ['journalctl', '--no-pager', '-u', 'kubelet'],
    'env_cumulus.txt': ['sudo', 'su', '-l', 'cumulus', '-c', 'printenv'],
    'env_root.txt': ['sudo', 'su', '-l', 'root', '-c', 'printenv'],
}

DEBIAN_UBUNTU_FILES = {
    '/var/log/': ['syslog', 'netq-agent', 'netqd', 'netq-model'],
    '/etc/': ['os-release', 'netq'],
    '/var/run/': ['netq'],
}
DEBIAN_UBUNTU_CMDS = {
    'dpkg.txt': ['dpkg', '-l'],
}
CENTOS_RHEL_FILES = {
    '/var/log/': ['messages', 'netq-agent', 'netqd', 'netq-model'],
    '/etc/': ['os-release', 'netq'],
    '/var/run/': ['netq'],
}
CENTOS_RHEL_CMDS = {
    'rpm-qa.txt': ['rpm', '-qa']
}

SONIC_FILES = {
    '/var/log/': ['syslog', 'netq-agent', 'netqd', 'netq-model', 'boot', 'dpkg', 'auth'],
    '/etc/': ['os-release', 'netq'],
    '/etc/sonic/': ['sonic_version.yml'],
    '/var/run/': ['netq'],
}

SONIC_CMDS = {
    'config-db-json.txt': ['/usr/local/bin/sonic-cfggen', '-d', '--print-data'],
}

def is_opta():
    return os.path.isfile(OPTA_FILE)

def is_cumulus():
    return (is_cl3() or is_cl4())

def platform():
    sysname, nodename, release, version, machine = os.uname()

    if os.path.isdir(NETQ_AGENT_PY3_4_VENV_PATH):
        NETQ_AGENT_VENV_PATH = NETQ_AGENT_PY3_4_VENV_PATH
    elif os.path.isdir(NETQ_AGENT_PY3_6_VENV_PATH):
        NETQ_AGENT_VENV_PATH = NETQ_AGENT_PY3_6_VENV_PATH
    elif os.path.isdir(NETQ_AGENT_PY3_7_VENV_PATH):
        NETQ_AGENT_VENV_PATH = NETQ_AGENT_PY3_7_VENV_PATH
    elif os.path.isdir(NETQ_AGENT_PY3_8_VENV_PATH):
        NETQ_AGENT_VENV_PATH = NETQ_AGENT_PY3_8_VENV_PATH
    else:
        NETQ_AGENT_VENV_PATH = ''
        major = sys.version_info.major
        minor = sys.version_info.minor
        print('skipping model dump - python version {}.{} - not supported'.format(major, minor))

    if len(NETQ_AGENT_VENV_PATH):
        NETQ_CLIENT_CLI = '%s%s' % (NETQ_AGENT_VENV_PATH, NETQ_CLI_CLIENT_PATH)
        LINUX_CMDS.update({'netq-agent-cli.txt': ['python3', NETQ_CLIENT_CLI]})

    if is_opta():
        OPTA_CMDS.update(LINUX_CMDS)
        OPTA_CMDS.update(DEBIAN_UBUNTU_CMDS)
        return OPTA_FILES, OPTA_CMDS
    elif 'Cumulus' in version:
        LINUX_CMDS.update(DEBIAN_UBUNTU_CMDS)
        return CUMULUS_FILES, LINUX_CMDS
    elif 'Ubuntu' in version or 'Debian' in version:
        DEBIAN_UBUNTU_CMDS.update(LINUX_CMDS)
        if is_sonic():
            DEBIAN_UBUNTU_CMDS.update(SONIC_CMDS)
            return SONIC_FILES, DEBIAN_UBUNTU_CMDS
        return DEBIAN_UBUNTU_FILES, DEBIAN_UBUNTU_CMDS
    elif 'el7' in release:
        try:
            # brctl is not installed by default
            subprocess_wrapper(['yum', 'install', 'bridge-utils.x86_64', '-y'])
        except subprocess.CalledProcessError as ex:
            print('ERROR: {}'.format(ex.cmd))
        CENTOS_RHEL_CMDS.update(LINUX_CMDS)
        return CENTOS_RHEL_FILES, CENTOS_RHEL_CMDS

    else: # rhel
        try:
            # brctl is not installed by default
            subprocess_wrapper(['yum', 'install', 'bridge-utils.x86_64', '-y'])
        except subprocess.CalledProcessError as ex:
            print('ERROR: {}'.format(ex.cmd))
        CENTOS_RHEL_CMDS.update(LINUX_CMDS)
        return CENTOS_RHEL_FILES, CENTOS_RHEL_CMDS


def copy_files(file_dict, support_path):
    cp_files = []
    for path in file_dict:
        for file in os.listdir(path):
            for file_pattern in file_dict[path]:
                if file_pattern in file:
                    cp_files.append('%s%s' % (path, file))

    try:
        for file in cp_files:
            subprocess_wrapper(['cp', '-r', '-L', '--parents', file, support_path + '/Logs'])
    except subprocess.CalledProcessError as ex:
        print('ERROR: failed to {}'.format(ex.cmd))
        return RC_FAIL
    else:
        return RC_SUCCESS

def copy_pods_files(support_path):
    try:
        for pod_name in PODS_FILES:
            copy_pod_files(pod_name, PODS_FILES[pod_name], support_path)
    except Exception as ex:
        print('ERROR: failed to {}'.format(ex.cmd))
        return RC_FAIL
    else:
        return RC_SUCCESS

def copy_pod_files(pod_name, directory_path, support_path):
    full_pod_names = get_full_pod_name(pod_name)
    for full_pod_name in full_pod_names:
        kubectl_exec_cassandra_command = f'kubectl cp {full_pod_name}:{directory_path} {support_path}/Logs/{full_pod_name}'
        shell_process = subprocess.Popen(kubectl_exec_cassandra_command, stdin=subprocess.PIPE,
                stdout=subprocess.PIPE, stderr=subprocess.PIPE, shell=True)

def get_full_pod_name(pod_name):
    kubectl_get_pod_command = f"kubectl get pods | grep {pod_name}"
    output = subprocess.check_output(kubectl_get_pod_command, shell=True)
    lines = output.decode("utf-8").split("\n")
    full_name_pods = []
    for line in lines:
        if pod_name in line:
            full_name_pods.append(line.split()[0])
    return full_name_pods

def cmd_outputs(cmd_dict, support_path):
    devnull = open(os.devnull, 'w')
    failed_cmds = ''
    try:
        for file in cmd_dict:
            if file == 'systemctl.txt':
                output = subprocess_wrapper(cmd_dict[file]).split()
                services = set()
                for word in output:
                    if 'service' in word:
                        services.add(word)

                properties = 'Names,ActiveState,UnitFilePreset,MainPID,ActiveEnterTimestamp'
                run_cmd = ['systemctl', 'show', '-p', properties] + list(services)
            elif file == 'netq-app-admin.txt':
                output = subprocess_wrapper(cmd_dict[file].split())
                with open(NETQ_CMD_SCRIPT, 'w') as fd:
                    os.chmod("netq_script.sh", 0o744)
                    fd.write("#!/bin/sh\n")
                    for container in output.splitlines():
                        fd.write("echo %%%%% Logs of {} %%%%%\n".format(container))
                        fd.write("docker logs {}\n".format(container))
                    run_cmd = './' + NETQ_CMD_SCRIPT
            elif type(cmd_dict[file]) == str:
                # call this in a temp bash script
                with open(NETQ_CMD_SCRIPT, 'w') as fd:
                    os.chmod("netq_script.sh", 0o744)
                    fd.write("#!/bin/sh\n")
                    fd.write(cmd_dict[file] + "\n")
                    run_cmd = './' + NETQ_CMD_SCRIPT
            else:
                run_cmd = cmd_dict[file]

            with open(file, 'w') as fp:
                try:
                    subprocess.check_call(run_cmd, stdout=fp, stderr=devnull, timeout=CMD_TIMEOUT)
                    subprocess_wrapper(['mv', file, support_path + '/Support'])
                except Exception as ex:
                    reattempt_status = RC_FAIL
                    if is_opta() and ' '.join(run_cmd) == "netq show agents":
                        with open(file, 'r') as fp1:
                            line = fp1.readline()
                            if "Login Failed" in line or "Access key is not found" in line:
                                print("Access key is not found. Please check the access key entered or generate a fresh access_key,secret_key pair and add it to the CLI configuration")
                                print("Proceeding with opta-support generation without netq show outputs")
                    subprocess_wrapper(['rm', '-f', file])
                    if reattempt_status != RC_SUCCESS:
                        failed_cmds += "Command: %s, Error Output: %s\n" % (' '.join(run_cmd), str(ex))

        if failed_cmds:
            fd_failed_cmds = open(FAILED_COMMANDS_FILE, 'a')
            fd_failed_cmds.write(failed_cmds)
            fd_failed_cmds.close()
            subprocess_wrapper(['mv', FAILED_COMMANDS_FILE, support_path + '/Support'])

            try:
                os.remove(NETQ_CMD_SCRIPT)
            except Exception:
                pass

    except subprocess.CalledProcessError as ex:
        print('ERROR: failed to {}'.format(ex.cmd))
        return RC_FAIL
    else:
        return RC_SUCCESS


if __name__ == '__main__':
    prog_name = os.path.basename(sys.argv[0])
    if len(sys.argv) > 1:
        copy_db_logs = COPY_DB_LOGS in sys.argv

    if os.getuid():
        print('Must run {} as root'.format(prog_name))
        sys.exit(1)

    if OPTAPREFIX in prog_name and not is_opta():
        print('Must run {} on a NetQ Appliance. Please run netq-support instead'.format(prog_name))
        sys.exit(0)

    if NETQPREFIX in prog_name and is_opta():
        print('Must run opta-support on a NetQ Appliance')
        sys.exit(0)

    print('Generating opta-support archive. Process takes few minutes to complete...')
    file_dict, cmd_dict = platform()

    if is_opta():
        SUPPORT_NAME = OPTAPREFIX + '_' + BASENAME
    else:
        SUPPORT_NAME = NETQPREFIX + '_' + BASENAME

    SUPPORT_FILE = '%s%s' % (SUPPORT_NAME, EXTENSION)
    SUPPORT_PATH_FILE = '%s%s' % (SUPPORT_PATH, SUPPORT_FILE)
    TMP_SUPPORT_PATH = '/tmp/%s' % (SUPPORT_NAME)

    subprocess_wrapper(['mkdir', TMP_SUPPORT_PATH])
    subprocess_wrapper(['mkdir', TMP_SUPPORT_PATH + '/Logs'])
    subprocess_wrapper(['mkdir', TMP_SUPPORT_PATH + '/Support'])

    if is_cumulus() == True:
        CUMULUS_SUPPORT_NAME = '/cl-support'
        CL_SUPPORT_PATH = TMP_SUPPORT_PATH + (CUMULUS_SUPPORT_NAME)
        subprocess_wrapper(['mkdir', CL_SUPPORT_PATH])
        print('Collecting cl-support...')
        try:
            subprocess_wrapper(['cl-support', '-S', CL_SUPPORT_PATH, '-p', \
                'netq-cl', '-r', '\"triggered from netq-support\"', '-T', '300' , '-j'])
            print('Collecting netq-support...')
        except subprocess.CalledProcessError as ex:
            print('cl-support collection did not succeed. collecting netq-support alone')
            pass

    if not os.path.isdir(SUPPORT_PATH):
        subprocess_wrapper(['mkdir', SUPPORT_PATH])

    cmd_outputs(cmd_dict, TMP_SUPPORT_PATH)
    sleep(3)
    copy_files(file_dict, TMP_SUPPORT_PATH)
    if 'copy_db_logs' in locals() and copy_db_logs:
        copy_pods_files(TMP_SUPPORT_PATH)

    try:
        subprocess_wrapper(['/bin/tar', '--directory=/tmp', TAR_OPTIONS, SUPPORT_PATH_FILE, SUPPORT_NAME])
        subprocess_wrapper(['rm', '-rf', TMP_SUPPORT_PATH])
        print('Please send {} to Nvidia support.'.format(SUPPORT_PATH_FILE))
    except subprocess.CalledProcessError as ex:
        print('ERROR: could not create support file ({})'.format(ex.cmd))
        pass
