devutils

#!/usr/bin/env python3

# Supress warning
import sys
from devutil import conn_graph_helper
from devutil.task_runner import TaskRunner
from devutil.ssh_utils import SSHClient
from devutil.inv_helpers import HostManager
from tabulate import tabulate
import time
import subprocess
import json
import argparse
import warnings
warnings.filterwarnings("ignore")

sys.path.append("..")

from tests.common.plugins.pdu_controller.pdu_manager import pdu_manager_factory     # noqa E402
from tests.common.connections.console_host import ConsoleHost   # noqa E402

g_inv_mgr = None
g_task_runner = None
g_pdu_dict = {}
g_conn_graph_facts = {}


def run_cmd(cmd):
    '''
    @summary: Utility that runs a command in a subprocess
    @param cmd: Command to be run
    @return: stdout of the command run
    @return: stderr of the command run
    '''
    out = subprocess.Popen(cmd, stdout=subprocess.PIPE,
                           stderr=subprocess.STDOUT, shell=True)
    stdout, stderr = out.communicate()
    return out.returncode, stdout, stderr


def get_conn_graph_facts(hosts):

    global g_conn_graph_facts
    if g_conn_graph_facts:
        return g_conn_graph_facts
    hostnames = hosts.keys()
    g_conn_graph_facts = conn_graph_helper.get_conn_graph_facts(hostnames)
    return g_conn_graph_facts


def build_global_vars(concurrency, inventory):
    global g_task_runner, g_inv_mgr
    g_task_runner = TaskRunner(max_worker=concurrency)
    g_inv_mgr = HostManager(inventory)


def retrieve_hosts(group, limit):
    global g_inv_mgr
    return g_inv_mgr.get_host_list(group, limit)


def get_pdu_info_from_conn_graph(hostname):
    """
    Read pdu info from conn graph.
    Returns a tuple of success, err msg or result
    Result is a tuple of pdu_links, pdu_info and pdu_vars
    {
        "PSU1": {
            "A": {
                "peerdevice": "pdu-2",
                "peerport": "12",
                "feed": "A",
            },
        },
    }
    {
        "pdu-2": {
            "Hostname": "pdu-2",
            "Protocol": "snmp",
            "ManagementIp": "10.3.155.107",
            "HwSku": "Sentry",
            "Type": "Pdu",
        },
    }
    {
        "pdu-2": {
        ...
        }
    }
    The 3 variables can be directly used in pdu_manager_factory
    """
    global g_inv_mgr, g_conn_graph_facts
    device_pdu_links = g_conn_graph_facts['device_pdu_links']
    device_pdu_info = g_conn_graph_facts['device_pdu_info']
    if hostname not in device_pdu_links or hostname not in device_pdu_info:
        return (False, '{} does not exist in connection graph'.format(hostname))
    pdu_links = device_pdu_links[hostname]
    pdu_info = device_pdu_info[hostname]
    pdu_vars = {}
    for pdu_name in pdu_info.keys():
        pdu_vars[pdu_name] = g_inv_mgr.get_host_vars(pdu_name)
    return (True, (pdu_links, pdu_info, pdu_vars))


def get_pdu_info_from_inventory(hostname, attrs):
    """
    Read pdu info from inventory. This should be a fallback of get_pdu_info_from_conn_graph.
    Returns a tuple of success, err msg or result
    Result is a tuple of pdu_links, pdu_info and pdu_vars.
    """
    global g_inv_mgr, g_pdu_dict
    pdu_host = attrs['pdu_host'] if 'pdu_host' in attrs else None
    if not pdu_host:
        return (False, 'DUT has no PDU configuration')

    if pdu_host in g_pdu_dict:
        return (True, g_pdu_dict[pdu_host])

    hosts = retrieve_hosts('all', pdu_host)
    if not hosts:
        return (False, 'PDU not found in inventory')

    pdu_links = {}
    pdu_info = {}
    pdu_vars = {}
    index = 1
    for ph in pdu_host.split(','):
        if ph in hosts:
            host_vars = hosts[ph]
            pdu_links['PSU{}'.format(index)] = {
                'N/A': {
                    'peerdevice': ph,
                    'peerport': 'probing',
                    'feed': 'N/A',
                }
            }
            pdu_info[ph] = {
                'Hostname': ph,
                'Protocol': host_vars['protocol'],
                'ManagementIp': host_vars['ansible_host'],
                'Type': 'Pdu',
            }
            pdu_vars[ph] = host_vars
            index = index + 1
    g_pdu_dict[pdu_host] = (pdu_links, pdu_info, pdu_vars)

    return (True, g_pdu_dict[pdu_host])


def get_pdu_info(dut_hostname, attrs):
    success, results = get_pdu_info_from_conn_graph(dut_hostname)
    if success:
        return (success, results)
    return get_pdu_info_from_inventory(dut_hostname, attrs)


def get_console_info_from_conn_graph(hostname):
    """
    Read console info from conn_graph_facts.
    """
    console_info = {}
    if hostname in g_conn_graph_facts['device_console_info'] and g_conn_graph_facts['device_console_info'][hostname]:
        console_info['console_type'] = g_conn_graph_facts['device_console_info'][hostname]['Protocol']
        console_info['console_host'] = g_conn_graph_facts['device_console_info'][hostname]['ManagementIp']
        console_info['console_port'] = g_conn_graph_facts['device_console_link'][hostname]['ConsolePort']['peerport']
    return console_info


def get_console_info_from_inventory(attrs):
    """
    Read console info from inventory file. This should be a fallback of get_console_info_from_conn_graph.
    """
    console_info = {}
    keys = ['console_type', 'console_host', 'console_port']
    for k in keys:
        if k in attrs:
            console_info[k] = attrs[k]
    return console_info


def get_console_info(hostname, attrs):
    console_info = get_console_info_from_conn_graph(hostname)
    if not console_info:
        console_info = get_console_info_from_inventory(attrs)
    if not console_info:
        print("Failed to get console info for {}".format(hostname))
    return console_info


def show_data_output(header, data, json_output=False, output_file=None):
    if json_output:
        output_content = json.dumps(sorted(data, key=lambda x: x['Host']), indent=4)
    else:
        output_content = tabulate(sorted(data, key=lambda x: x[0]), headers=header, tablefmt='grid')

    if output_file:
        try:
            with open(output_file, 'w') as f:
                f.write(output_content)
            print("Output saved to {}".format(output_file))
        except Exception as e:
            print("Failed to write to output file: {}".format(e))
    else:
        print(output_content)


def action_list(parameters):
    hosts = parameters['hosts']
    header = ['Host', 'Ansible_host']
    data = []
    if parameters['json']:
        for name, vars in hosts.items():
            data.append(dict(zip(header, (name, vars['ansible_host']))))
    else:
        for name, vars in hosts.items():
            data.append(
                (name, vars['ansible_host'] if 'ansible_host' in vars else 'not_available'))
    show_data_output(header, data, parameters['json'], parameters['out'])


def action_ping(parameters):
    hosts = parameters['hosts']
    header = ['Host', 'Hostname', 'Ping result']
    data = []
    for name, vars in hosts.items():
        if 'ansible_host' in vars:
            cmd = 'timeout 1 ping -q -c 1 -w 1 {}'.format(vars['ansible_host'])
            g_task_runner.submit_task(
                name + '|' + vars['ansible_host'], run_cmd, cmd=cmd)
    if parameters['json']:
        for name, result in g_task_runner.task_results():
            data.append(
                dict(zip(header, (name.split('|')[0], name.split('|')[1],
                                  'Success' if result['result'][0] == 0 else "Fail"))))
    else:
        for name, result in g_task_runner.task_results():
            data.append((name.split('|')[0], name.split('|')[
                        1], 'Success' if result['result'][0] == 0 else "Fail"))

    if parameters['ipv6']:
        for name, vars in hosts.items():
            if 'ansible_hostv6' in vars.keys():
                cmd = 'timeout 1 ping -6 -q -c 1 -w 1 {}'.format(
                    vars['ansible_hostv6'])
                g_task_runner.submit_task(
                    name + '|' + vars['ansible_hostv6'], run_cmd, cmd=cmd)
        if parameters['json']:
            for name, result in g_task_runner.task_results():
                data.append(
                    dict(zip(header, (name.split('|')[0], name.split('|')[1],
                                      'Success' if result['result'][0] == 0 else "Fail"))))
        else:
            for name, result in g_task_runner.task_results():
                data.append((name.split('|')[0], name.split('|')[
                            1], 'Success' if result['result'][0] == 0 else "Fail"))

    show_data_output(header, data, parameters['json'], parameters['out'])


def action_ssh(parameters):
    hosts = parameters['hosts']
    for _, vars in hosts.items():
        client = SSHClient()
        client.connect(hostname=vars['ansible_host'], username=vars[
                       'creds']['username'], passwords=vars['creds']['password'])
        client.posix_shell()


def action_console(parameters):
    hosts = parameters['hosts']
    for hostname, vars in hosts.items():
        console_info = get_console_info(hostname, vars)
        if not console_info:
            continue
        console_host = ConsoleHost(console_type=console_info['console_type'],
                                   console_host=console_info['console_host'],
                                   console_port=console_info['console_port'],
                                   sonic_username=vars['creds']['username'],
                                   sonic_password=vars['creds']['password'],
                                   console_username=vars['creds']['console_user'][console_info['console_type']],
                                   console_password=vars['creds']['console_password'][console_info['console_type']])
        console_host.posix_shell()


def pdu_action_on_dut(host, attrs, action):
    ret = {'Host': host, 'PSU status': [], 'PDU status': [], 'Summary': [], 'Action': action}
    succeed, result = get_pdu_info(host, attrs)
    if not succeed:
        ret['Summary'] = result
        return ret

    pdu_links, pdu_info, pdu_vars = result
    pduman = pdu_manager_factory(host, pdu_links, pdu_info, pdu_vars)

    if not pduman:
        ret['Summary'].append(
            'Failed to communicate with PDU controller {}'.format(pdu_links.keys()))
        return ret

    if action == 'off':
        pduman.turn_off_outlet()
    elif action == 'on':
        pduman.turn_on_outlet()
    elif action != 'status':
        ret['Summary'].append('Unsupported action {}.'.format(action))
        return ret

    outlet_status = pduman.get_outlet_status()
    psu_status = {}
    for outlet in outlet_status:
        ret['PDU status'].append(outlet)
        psu_name = outlet['psu_name']
        status = psu_status.get(psu_name, {
                'power_on': False,
                'output_watts': 0,
        })
        status['power_on'] = status['power_on'] or outlet['outlet_on']
        if 'output_watts' not in outlet:
            status['output_watts'] = 'N/A'
        else:
            raw_outlet_watts = outlet['output_watts']
            if not raw_outlet_watts:
                outlet_watts = 0
            else:
                try:
                    outlet_watts = int(outlet['output_watts'])
                except ValueError:
                    outlet_watts = 0
            status['output_watts'] = status['output_watts'] + outlet_watts
        psu_status[psu_name] = status
    ret['PSU status'] = psu_status

    return ret


def action_pdu(parameters, action):
    hosts = parameters['hosts']
    data = []
    header = ['Host', 'Action', 'PSU status', 'PDU status', 'Summary']
    for host, attrs in hosts.items():
        g_task_runner.submit_task(
            host, pdu_action_on_dut, host=host, attrs=attrs, action=action)

    for _, ret in g_task_runner.task_results():
        status = ret['result']
        if parameters['json']:
            data.append(status)
        else:
            data.append([status[x] for x in header])

    return header, data


def action_pdu_status(parameters):
    header, data = action_pdu(parameters, 'status')
    show_data_output(header, data, parameters['json'], parameters['out'])


def action_pdu_failures(parameters):
    header, data = action_pdu(parameters, 'status')
    failures = []
    for entry in data:
        if parameters['json']:
            if any(x['output_watts'] == '0' for x in entry['PDU status']):
                failures.append(entry)
        else:
            if any(x['output_watts'] == '0' for x in entry[2]):
                failures.append(entry)
    show_data_output(header, failures, parameters['json'], parameters['out'])


def action_pdu_off(parameters):
    header, data = action_pdu(parameters, 'off')
    show_data_output(header, data, parameters['json'], parameters['out'])


def action_pdu_on(parameters):
    header, data = action_pdu(parameters, 'on')
    show_data_output(header, data, parameters['json'], parameters['out'])


def action_pdu_reboot(parameters):
    header, data = action_pdu(parameters, 'off')
    # sleep 1 second to ensure there is gap between power off and on
    time.sleep(1)
    _, data_on = action_pdu(parameters, 'on')

    data = data + data_on
    show_data_output(header, data, parameters['json'], parameters['out'])


def action_dispatcher(parameters):
    # Actions that can run simultaneously in different hosts
    parallel_actions = ['ssh_run_command']
    action = parameters['action'].__name__
    if action in parallel_actions:
        parallel_run(parameters)
    else:
        parameters['action'](parameters)


def parallel_run(parameters):
    action = parameters['action']
    cmd = parameters['cmd']
    hosts = parameters['hosts']
    if action.__name__ == 'ssh_run_command':
        for hostname, vars in hosts.items():
            args = {'hostname': vars['ansible_host'],
                    'username': vars['creds']['username'],
                    'passwords': vars['creds']['password'],
                    'cmd': cmd
                    }
            g_task_runner.submit_task(hostname, ssh_run_command, **args)
        for name, result in g_task_runner.task_results():
            print("task result for {} ===============>\n{}".format(
                name, str(result['result'][1])))


def ssh_run_command(hostname, username, passwords, cmd):
    client = SSHClient()
    client.connect(hostname=hostname, username=username, passwords=passwords)
    return client.run_command(cmd)


def validate_args(args):
    if args.action == 'run' and args.cmd == '':
        print("command is missing for run action")
        return False
    return True


def setup_logging(level):
    if level == 'close':
        return
    import logging
    if level == 'debug':
        logging.basicConfig(stream=sys.stdout, level=logging.DEBUG)
    elif level == 'info':
        logging.basicConfig(stream=sys.stdout, level=logging.INFO)
    elif level == 'warn':
        logging.basicConfig(stream=sys.stdout, level=logging.WARN)
    elif level == 'error':
        logging.basicConfig(stream=sys.stdout, level=logging.ERROR)


def main():
    parser = argparse.ArgumentParser(description='Device utilities')
    parser.add_argument('-6', '--ipv6', help='Include IPv6', action='store_true',
                        required=False, default=False)
    parser.add_argument('-a', '--action',
                        help='Action towards host(s): list, ping, run, ssh, console, pdu_status, '
                             'pdu_on, pdu_off, pdu_reboot, pdu_failures, default list',
                        type=str, required=False, default='list',
                        choices=['list', 'ping', 'ssh', 'console', 'run', 'pdu_status',
                                 'pdu_on', 'pdu_off', 'pdu_reboot', 'pdu_failures'])
    parser.add_argument('--cmd', help="Command to run on all hosts",
                        type=str, required=False)
    parser.add_argument('-g', '--group', help='Groups: all, sonic, ptf, pdu, default all',
                        type=str, required=False, default='all')
    parser.add_argument('-i', '--inventory', help='Categories: lab, etc, default lab',
                        type=str, required=False, default='lab')
    parser.add_argument('-l', '--limit', help='Host: limit to a single dut host name, default all',
                        type=str, required=False)
    parser.add_argument('--log-level', help='Log level: print logs to STDOUT (if not set to close)',
                        type=str, required=False, default='warn',
                        choices=['debug', 'info', 'warn', 'error', 'close'])
    parser.add_argument('-u', '--user', help='User: user account to login to host with, default admin',
                        type=str, required=False, default='admin')
    parser.add_argument(
        '-c', '--concurrency', help='Concurrency: the max concurrency for tasks that can run simultaneously, default 1',
        type=int, required=False, default=1)
    parser.add_argument('-j', '--json', help='json output', action='store_true',
                        required=False, default=False)
    parser.add_argument('-o', '--out', help='specify output file',
                        required=False, default=None)

    args = parser.parse_args()
    if not validate_args(args):
        return
    setup_logging(args.log_level)
    build_global_vars(args.concurrency, args.inventory)
    # Add limit argument check for pdu_reboot, pdu_on, pdu_off actions
    # If no limit argument for these actions, will not execute the process for all devices
    if args.action in ['pdu_reboot', 'pdu_on', 'pdu_off']:
        if not args.limit or args.limit == '':
            print("devutils: error: argument -l/--limit is required for pdu_reboot, pdu_on, "
                  "pdu_off actions to avoid impacting all devices accidentally.")
            print(
                "If you are sure to do pdu action to all devices, please use -l/--limit all.")
            return
    hosts = retrieve_hosts(args.group, args.limit)
    if not hosts:
        print('No matching hosts')
        return
    get_conn_graph_facts(hosts)
    actions = {'list': action_list,
               'ping': action_ping,
               'ssh': action_ssh,
               'console': action_console,
               'run': ssh_run_command,
               'pdu_status': action_pdu_status,
               'pdu_off': action_pdu_off,
               'pdu_on': action_pdu_on,
               'pdu_reboot': action_pdu_reboot,
               'pdu_failures': action_pdu_failures,
               }
    parameters = {'hosts': hosts,
                  'limit': args.limit,
                  'action': actions[args.action],
                  'user': args.user,
                  'ipv6': args.ipv6,
                  'cmd': args.cmd,
                  'json': args.json,
                  'out': args.out,
                  }
    action_dispatcher(parameters)


if __name__ == '__main__':
    main()