#!/usr/bin/env python
#
# Copyright (c) 2017-2022 Dell Inc. or its subsidiaries. All Rights Reserved.
#
"""
This script disruptively applies firmware updates from the Drive Support
Package (DSP) on nodes that are not in a cluster. This script
does not install the contents of the DSP on the node and/or cluster.
"""
from __future__ import print_function

import datetime
import optparse
import os
import re
import string
import subprocess
import sys
import time
from builtins import input
from builtins import map
from builtins import zip

LOG = None


def log(msg):
    """Try to log to the local log file"""
    global LOG
    ts = datetime.datetime.now().strftime("%Y-%m-%dT%H:%M:%S")
    msg = "%s %s\n" % (ts, msg.strip("\n"))
    if LOG:
        LOG.write(msg)
        LOG.flush()
    else:
        print(msg.strip("\n"))


def log_print(msg):
    """Log and print to stdout"""
    log(msg)
    print(msg)


def cmd_err_log(msg, cmd, rc, out, err):
    """Log detail of failed subprocess commands"""
    log_print("ERROR: %s" % msg)
    log("    command: %s" % " ".join(cmd))
    log("    returncode: %s" % rc)
    log("    stdout:")
    for line in out.split("\n"):
        log("        %s" % line)
    log("    stderr:")
    for line in err.split("\n"):
        log("        %s" % line)


def cleanup_exit(rc):
    """Close log and exit"""
    global LOG
    if LOG:
        LOG.write("***closing log file***\n")
        LOG.flush()
        LOG.close()
    sys.exit(rc)


def run_cmd(cmd):
    """Run command and return tuple of returncode, stdout, stderr"""
    proc = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
    (output, error) = proc.communicate()
    err = error.decode("utf-8")
    out = output.decode("utf-8")
    rc = proc.returncode
    return (rc, out, err)


MIRRORCTL_PLATFORMS = ["110-365", "110-385"]


def platform_check():
    """Check that the platform is Gen5/6 without dependencies
    on OneFS libs that may change between releases or not work
    correctly in a degraded state.
    """
    supported_platforms = ["S1400FP", "S2600CP"] + MIRRORCTL_PLATFORMS
    cmd = ["/usr/local/sbin/dmidecode", "-t", "baseboard"]
    (rc, out, err) = run_cmd(cmd)
    if not out:
        cmd_err_log("Platform check failed to complete", cmd, rc, out, err)
        cleanup_exit(1)
    for platform in supported_platforms:
        if platform in out:
            return platform
    log_print("ERROR: This platform is not supported by this script")
    cleanup_exit(1)


def option_parser_setup():
    """Set up option parser to be used in main"""
    parser = optparse.OptionParser()
    parser.add_option(
        "-a",
        "--assess",
        action="store_true",
        dest="assess",
        default=False,
        help="show available updates (default)",
    )
    parser.add_option(
        "-u",
        "--update",
        action="store_true",
        dest="update",
        default=False,
        help="update drive firmware",
    )
    parser.add_option(
        "-s",
        "--serial",
        action="store_true",
        dest="serial",
        default=False,
        help="update drives one at a time (default)",
    )
    parser.add_option(
        "-p",
        "--parallel",
        action="store_true",
        dest="parallel",
        default=False,
        help="update drives in parallel",
    )
    parser.add_option(
        "-r",
        "--reboot",
        action="store_true",
        dest="reboot",
        default=False,
        help="reboot after updates (default)",
    )
    parser.add_option(
        "-n",
        "--no-reboot",
        action="store_true",
        dest="no_reboot",
        default=False,
        help="don't reboot after updates",
    )
    parser.add_option(
        "-o",
        "--power-off",
        action="store_true",
        dest="power_off",
        default=False,
        help="power off instead of reboot",
    )
    parser.add_option(
        "-f",
        "--force",
        action="store_true",
        dest="force",
        default=False,
        help="skip user prompts",
    )
    parser.add_option(
        "-c",
        "--continue-on-error",
        action="store_true",
        dest="continue_on_error",
        default=False,
        help=optparse.SUPPRESS_HELP,
    )
    parser.add_option(
        "-i",
        "--ignore-version-check",
        action="store_true",
        dest="ignore_version_check",
        default=False,
        help=optparse.SUPPRESS_HELP,
    )
    parser.add_option(
        "-m",
        "--skip-mirror-rebalances",
        action="store_true",
        dest="skip_mirror_rebalances",
        default=False,
        help=optparse.SUPPRESS_HELP,
    )
    return parser


def option_conflict(options, parser):
    """Show end user conflicting options and help text"""
    log_print("ERROR: Conflicting options: %s\n" % ", ".join(options))
    parser.print_help()
    cleanup_exit(1)


def find_drives():
    """Find drives recognized by the system"""
    cmd = ["/usr/bin/isi_drivenum", "-bmfd"]
    (rc, out, err) = run_cmd(cmd)
    if not out:
        cmd_err_log("No drives found", cmd, rc, out, err)
        cleanup_exit(1)
    drive_re = re.compile(
        r"^(?P<bay>Bay\s+\S+)\s+(?P<model>(\S+\s+)*\S+)\s+"
        r"FW:(?P<fw>\S+)\s+/dev/(?P<dev>da\d+)$"
    )
    drives = []
    for line in out.split("\n"):
        match = drive_re.match(line)
        if match:
            drive = {
                "bay": match.group("bay"),
                "model": match.group("model"),
                "fw": match.group("fw"),
                "dev": match.group("dev"),
            }
            drives.append(drive)
    return drives


MODEL_RE = re.compile(r"model\s*=\s*(\"|\')(?P<val>(\S+\s+)*\S+)(\"|\')")
MANUFACTURER_RE = re.compile(r"manufacturer\s*=\s*(\"|\')(?P<val>\w+)(\"|\')")
ALGORITHM_RE = re.compile(r"algorithm\s*=\s*(?P<val>\w+)")
VERSION_RE = re.compile(r"version\s*=\s*(\"|\')(?P<val>\S+)(\"|\')")
FILENAME_RE = re.compile(r"filename\s*=\s*(\"|\')(?P<val>\S+)(\"|\')")
CHECKSUM_RE = re.compile(r"checksum\s*=\s*(\"|\')(?P<val>\w+)(\"|\')")
VALID_VERSIONS_RE = re.compile(r"valid_versions\s*=\s*\[\s*(?P<val>.+)\s*\]")


def parse_config(config, drive, ignore_version_check):
    """Parse the config for a drive and return update information if found"""
    update = {}
    model = None
    manufacturer = None
    algorithm = None
    fwupdate_versions = {}
    with open(config) as fin:
        for line in fin:
            match = MODEL_RE.search(line)
            if match:
                model = match.group("val")
                continue
            match = MANUFACTURER_RE.search(line)
            if match:
                manufacturer = match.group("val")
                continue
            match = ALGORITHM_RE.search(line)
            if match:
                algorithm = match.group("val")
                continue
            if "fwupdate_versions" in line:
                version = line.split(".")[1].split(".")[0]
                if version not in fwupdate_versions:
                    fwupdate_versions[version] = []
                fwupdate_versions[version].append(line)
    for group in fwupdate_versions:
        version = None
        filename = None
        checksum = None
        vv_str = None
        valid_versions = []
        for line in fwupdate_versions[group]:
            match = VERSION_RE.search(line)
            if match:
                version = match.group("val")
                continue
            match = FILENAME_RE.search(line)
            if match:
                filename = match.group("val")
                continue
            match = CHECKSUM_RE.search(line)
            if match:
                checksum = match.group("val")
                continue
            match = VALID_VERSIONS_RE.search(line)
            if match:
                vv_str = match.group("val")
                continue
        if vv_str:
            for val_ver in vv_str.split(","):
                valid_versions.append(val_ver.strip(" '\""))
        if (
            drive.get("model") == model
            and drive.get("fw") != version
            and manufacturer
            and algorithm
            and version
            and filename
            and checksum
        ):
            update["manufacturer"] = manufacturer
            update["algorithm"] = algorithm
            update["desired"] = version
            update["filename"] = filename
            update["checksum"] = checksum
            if ignore_version_check:
                return update
            if drive.get("fw") in valid_versions:
                return update
    return None


UPDATE_CACHE = {}
FW_CACHE = {}


def find_updates(drives, ignore_version_check, script_dir, opts):
    """Find available updates for drives in the system"""
    global UPDATE_CACHE
    global FW_CACHE
    updates = []
    for drive in drives:
        if (drive.get("model"), drive.get("fw")) in UPDATE_CACHE:
            update = UPDATE_CACHE[(drive.get("model"), drive.get("fw"))]
            if update:
                drive["update"] = update
                updates.append(drive)
            continue
        model_file = string.replace(drive.get("model"), " ", "_")
        model_file = string.replace(model_file, "-", "_")
        model_file = "%s.gc" % model_file
        model_file = os.path.join(script_dir, "config", "models", model_file)
        if not os.path.isfile(model_file):
            log_print("No config file found for %s" % drive.get("model"))
            continue
        update = parse_config(model_file, drive, ignore_version_check)
        if update:
            if update.get("filename") in FW_CACHE:
                update["fw_path"] = FW_CACHE[update.get("filename")]
                drive["update"] = update
                updates.append(drive)
            else:
                fw_path = os.path.join(
                    script_dir,
                    "firmware",
                    update.get("manufacturer"),
                    update.get("filename"),
                )
                if not os.path.exists(fw_path):
                    log_print("ERROR: File not found: %s" % fw_path)
                    if not opts.continue_on_error:
                        cleanup_exit(1)
                cmd = ["/sbin/md5", "-q", fw_path]
                (rc, out, err) = run_cmd(cmd)
                if not out:
                    cmd_err_log("MD5 command failed", cmd, rc, out, err)
                    if not opts.continue_on_error:
                        cleanup_exit(1)
                if update.get("checksum").lower() not in out.lower():
                    log_print("ERROR: Checksum mismatch on %s" % fw_path)
                    if not opts.continue_on_error:
                        cleanup_exit(1)
                if fw_path.endswith(".gz") and os.uname()[2].startswith("v7"):
                    new_fw_path = fw_path[:-3]
                    if not os.path.exists(new_fw_path):
                        cmd = ["/usr/bin/gunzip", "-fk", fw_path]
                        (rc, out, err) = run_cmd(cmd)
                        if rc:
                            cmd_err_log("gunzip command failed", cmd, rc, out, err)
                            if not opts.continue_on_error:
                                cleanup_exit(1)
                    fw_path = new_fw_path
                FW_CACHE[update.get("filename")] = fw_path
                update["fw_path"] = fw_path
                drive["update"] = update
                updates.append(drive)
        UPDATE_CACHE[(drive.get("model"), drive.get("fw"))] = update
    return updates


def kill_busy_umount_ifs(opts):
    """Run isi_kill_busy and umount /ifs"""
    cmd = ["/usr/sbin/isi_kill_busy"]
    (rc, out, err) = run_cmd(cmd)
    if rc:
        cmd_err_log("Failed to stop running processes", cmd, rc, out, err)
        if not opts.continue_on_error:
            cleanup_exit(1)
    cmd = ["/sbin/umount", "-f", "/ifs"]
    (rc, out, err) = run_cmd(cmd)
    if rc:
        cmd_err_log("Failed to unmount /ifs", cmd, rc, out, err)
        if not opts.continue_on_error:
            cleanup_exit(1)


def kill_sasphymon(opts):
    cmd = ["/bin/pkill", "-9", "isi_sasphymon"]
    (rc, out, err) = run_cmd(cmd)
    if rc > 1:
        cmd_err_log("Failed to kill isi_sasphymon", cmd, rc, out, err)
        if not opts.continue_on_error:
            cleanup_exit(1)


def start_key_manager():
    """Check if isi_km_d is running, start it if not"""
    cmd = ["/bin/pgrep", "isi_km_d"]
    (rc, out, err) = run_cmd(cmd)
    if rc:
        cmd = ["/bin/pgrep", "rpcbind"]
        (rc, out, err) = run_cmd(cmd)
        if rc:
            cmd = ["/usr/sbin/rpcbind"]
            (rc, out, err) = run_cmd(cmd)
            if rc:
                cmd_err_log(
                    "Failed to start rpcbind, SED updates will fail", cmd, rc, out, err
                )
        cmd = ["/usr/bin/isi_km_d", "start"]
        (rc, out, err) = run_cmd(cmd)
        if "successfully started" not in out:
            cmd_err_log(
                "Failed to start key manager, SED updates will fail", cmd, rc, out, err
            )
        time.sleep(5)
        cmd = ["/bin/pgrep", "isi_km_d"]
        (rc, out, err) = run_cmd(cmd)
        if rc:
            cmd_err_log("Key manager died, SED updates will fail", cmd, rc, out, err)


GMIRROR_CACHE = None


def in_mirror(dev):
    """Return True if dev is in a mirror"""
    global GMIRROR_CACHE
    if not GMIRROR_CACHE:
        cmd = ["/sbin/gmirror", "status"]
        (rc, out, err) = run_cmd(cmd)
        if not out:
            cmd_err_log("Gmirror command failed", cmd, rc, out, err)
            cleanup_exit(1)
        GMIRROR_CACHE = out
    if "%sp" % dev in GMIRROR_CACHE:
        return True
    return False


BAY_CACHE = {}


def mirror_rebalance(dev):
    """Rebalance boot mirror off of a drive to allow a safe update"""
    global GMIRROR_CACHE
    global BAY_CACHE
    if in_mirror(dev):
        if not BAY_CACHE:
            cmd = ["/usr/bin/isi_drivenum", "-bmfdV"]
            (rc, out, err) = run_cmd(cmd)
            if not out:
                cmd_err_log("No drives found", cmd, rc, out, err)
                cleanup_exit(1)
            bay_re = re.compile(
                r"^Bay\s+(?P<bay>\d+)\s+(\S+\s+)*\S+\s+"
                r"FW:\S+\s+/dev/(?P<dev>da\d+)$"
            )
            for line in out.split("\n"):
                match = bay_re.match(line)
                if match:
                    BAY_CACHE[match.group("dev")] = match.group("bay")
        if dev not in BAY_CACHE:
            log_print("ERROR: Not able to find bay number for %s" % dev)
            cleanup_exit(1)
        cmd = ["/usr/bin/isi_mirrorctl", "rebalance", "--exclude-bay", BAY_CACHE[dev]]
        (rc, out, err) = run_cmd(cmd)
        if rc:
            cmd_err_log("Rebalance failed", cmd, rc, out, err)
            cleanup_exit(1)
        GMIRROR_CACHE = None


def do_updates(updates, opts, use_mirrorctl, script_dir):
    """Perform drive firmware updates"""
    parallel = []
    serial = []
    if opts.parallel:
        if use_mirrorctl and not opts.skip_mirror_rebalances:
            for drive in updates:
                if in_mirror(drive.get("dev")):
                    serial.append(drive)
                else:
                    parallel.append(drive)
        else:
            parallel = updates
    else:
        serial = updates
    processes = []
    for drive in parallel:
        log_print("Updating %s" % drive.get("bay"))
        cmd = [
            "/usr/bin/isi_hwtools/isi_hdfw_update",
            "-a",
            drive.get("update").get("algorithm"),
            "-d",
            drive.get("dev"),
            "-f",
            drive.get("update").get("fw_path"),
        ]
        proc = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
        processes.append(proc)
    while processes:
        for proc in processes:
            if proc.poll() is not None:
                (output, error) = proc.communicate()
                out = output.decode("utf-8")
                err = error.decode("utf-8")
                rc = proc.returncode
                if rc:
                    cmd_err_log(
                        "Parallel update command failed", ["Unknown"], rc, out, err
                    )
                processes.remove(proc)
        time.sleep(1)
    for drive in serial:
        log_print("Updating %s" % drive.get("bay"))
        if use_mirrorctl and not opts.skip_mirror_rebalances:
            mirror_rebalance(drive.get("dev"))
        cmd = [
            "/usr/bin/isi_hwtools/isi_hdfw_update",
            "-a",
            drive.get("update").get("algorithm"),
            "-d",
            drive.get("dev"),
            "-f",
            drive.get("update").get("fw_path"),
        ]
        (rc, out, err) = run_cmd(cmd)
        if rc:
            cmd_err_log(
                "Update command failed on %s" % drive.get("bay"), cmd, rc, out, err
            )
            if not opts.continue_on_error:
                cleanup_exit(1)


def main():
    """Main function"""

    # Attempt to set up local log file
    global LOG
    log_path = "%s.log" % os.path.abspath(os.path.realpath(__file__))
    try:
        LOG = open(log_path, "a")
        LOG.write("***opened log file***\n")
        LOG.flush()
    except:  # noqa: E722 # pylint: disable=bare-except
        # Explicitly catch all exceptions to be backwards compatible with
        # old upgrade code.
        LOG = None
        log_print("ERROR: Logging to stdout only. Failed to open %s." % log_path)

    # Platform check
    platform = platform_check()

    # Parse and log options
    parser = option_parser_setup()
    opts, args = parser.parse_args()
    log("Running with opts=%s" % opts)

    # Log OneFS version
    cmd = ["/usr/bin/uname", "-a"]
    (rc, out, err) = run_cmd(cmd)
    log("Running on %s" % out.strip())

    # Check for conflicting options
    if (opts.assess, opts.update).count(True) > 1:
        option_conflict(("assess", "update"), parser)
    if (opts.serial, opts.parallel).count(True) > 1:
        option_conflict(("serial", "parallel"), parser)
    if (opts.reboot, opts.no_reboot, opts.power_off).count(True) > 1:
        option_conflict(("reboot", "no-reboot", "power-off"), parser)

    # Set default options
    if (opts.assess, opts.update).count(True) == 0:
        opts.assess = True
    if opts.update and (opts.serial, opts.parallel).count(True) == 0:
        opts.serial = True
    if opts.update and (opts.reboot, opts.no_reboot, opts.power_off).count(True) == 0:
        opts.reboot = True

    # Find drives and available updates
    drives = find_drives()
    script_dir = os.path.dirname(os.path.abspath(os.path.realpath(__file__)))
    updates = find_updates(drives, opts.ignore_version_check, script_dir, opts)
    if len(updates) < 1:
        log_print("\nNo updates available\n")
        cleanup_exit(0)

    # Display available updates
    log_print("\nAvailable updates:\n")
    headers = ["Location", "Firmware", "Desired", "Model"]
    update_list = [headers]
    for drive in updates:
        update_list.append(
            [
                drive.get("bay"),
                drive.get("fw"),
                drive.get("update").get("desired"),
                drive.get("model"),
            ]
        )
    widths = [max(list(map(len, col))) for col in zip(*update_list)]
    update_list.remove(headers)
    log_print("  ".join((h.ljust(w) for h, w in zip(headers, widths))))
    log_print("--".join(("-" * w for w in widths)))
    for drive in update_list:
        log_print("  ".join((d.ljust(w) for d, w in zip(drive, widths))))
    log_print("")

    # Execute actions specified by user
    if opts.assess:
        cleanup_exit(0)
    if opts.update:
        if not opts.force:
            yesno = input("Proceed with updates? (y/N): ")
            if not yesno.lower().startswith("y"):
                log_print("Updates canceled")
                cleanup_exit(0)
        if os.path.ismount("/ifs"):
            log_print(
                "WARNING: /ifs is mounted. In order to proceed the node "
                "will be disconnected from the cluster. All clients "
                "currently connected to this node will be disconnected."
            )
            yesno = input("Proceed with updates? (y/N): ")
            if not yesno.lower().startswith("y"):
                log_print("Updates canceled")
                cleanup_exit(0)
            kill_busy_umount_ifs(opts)
        kill_sasphymon(opts)
        start_key_manager()
        # If var migration exists, make sure it's not migrated
        if (
            os.path.exists("/usr/bin/isi_var_migrate")
            and not opts.skip_mirror_rebalances
        ):
            try:
                import isi.sys.bootdisk as bootdisk

                bootdisk.undo_var_migration()
            except (ImportError, AttributeError) as e:
                log_print("ERROR: %s" % e)
                if not opts.continue_on_error:
                    cleanup_exit(1)
        do_updates(updates, opts, platform in MIRRORCTL_PLATFORMS, script_dir)
        log_print("Updates complete")
        if opts.reboot:
            log_print("Rebooting node")
            cmd = ["/sbin/shutdown", "-r", "now"]
            (rc, out, err) = run_cmd(cmd)
            if rc:
                cmd_err_log("Failed to reboot node", cmd, rc, out, err)
        elif opts.power_off:
            log_print("Shutting down node")
            cmd = ["/sbin/shutdown", "-p", "now"]
            (rc, out, err) = run_cmd(cmd)
            if rc:
                cmd_err_log("Failed to shut down node", cmd, rc, out, err)

    cleanup_exit(0)


if __name__ == "__main__":
    main()
