#!/usr/bin/python2
#
# Copyright (C) 2013 by Red Hat, Inc.  All rights reserved.
#
# This program is free software; you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation; either version 2 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.
#
# Author(s): Brian C. Lane <bcl@brianlane.com>
#
"""
Driver Update Disk UI

/tmp/dd_modules is a copy of /proc/modules at startup time
/tmp/dd_args is a parsed list of the inst.dd= cmdline args, and may include
             'dd' or 'inst.dd' if it was specified without arguments
/tmp/dd_args_ks is the same format, but skips processing existing OEMDRV devices.

Pass a path and it will install the driver rpms from the path before checking
for new OEMDRV devices.

Repositories for installed drivers are copied into /run/install/DD-X where X
starts at 1 and increments for each repository.

Selected driver package names are saved in /run/install/dd_packages

Anaconda uses the repository and package list to install the same set of drivers
to the target system.
"""
import logging
from logging.handlers import SysLogHandler
import sys
import os
import subprocess
import time
import glob
import readline # pylint:disable=unused-import

log = logging.getLogger("DD")


class RunCmdError(Exception):
    """ Raised when run_cmd gets a non-zero returncode
    """
    pass


def run_cmd(cmd):
    """ Run a command, collect stdout and the returncode. stderr is ignored.

        :param cmd: command and arguments to run
        :type cmd:  list of strings
        :returns:   exit code and stdout from the command
        :rtype:     (int, string)
        :raises:    OSError if the cmd doesn't exist, RunCmdError if the rc != 0
    """
    try:
        with open("/dev/null", "w") as fd_null:
            log.debug(" ".join(cmd))
            proc = subprocess.Popen(cmd,
                                    stdout=subprocess.PIPE,
                                    stderr=fd_null)
            out = proc.communicate()[0]
            if out:
                for line in out.splitlines():
                    log.debug(line)
    except OSError as e:
        log.error("Error running %s: %s", cmd[0], e.strerror)
        raise
    if proc.returncode:
        log.debug("%s returned %s", cmd[0], proc.returncode)
        raise RunCmdError()
    return (proc.returncode, out)


def oemdrv_list():
    """ Get a list of devices labeled as OEMDRV

        :returns: list of devices
        :rtype:   list
    """
    try:
        outlines = run_cmd(["blkid", "-t", "LABEL=OEMDRV", "-o", "device"])[1]
    except (OSError, RunCmdError):
        # Nothing with that label
        return []
    else:
        return outlines.splitlines()


def get_dd_args():
    """ Get the dd arguments from /tmp/dd_args or /tmp/dd_args_ks

        :returns: List of arguments
        :rtype:   list of strings
    """
    net_protocols = ["http", "https", "ftp", "nfs", "nfs4"]
    args = []
    for dd_args_file in ["/tmp/dd_args", "/tmp/dd_args_ks"]:
        if not os.path.exists(dd_args_file):
            continue
        try:
            dd_args = open(dd_args_file, "r").readline().split()
        except IOError:
            return []

        # skip dd args that need networking
        args.extend(filter(lambda x: x.split(":")[0].lower() not in net_protocols, dd_args))
    return args


def is_interactive():
    """ Determine if the user requested interactive driver selection

        :returns: True if 'dd' or 'inst.dd' included in /tmp/dd_args False if not
        :rtype:   bool
    """
    dd_args = get_dd_args()
    if "dd" in dd_args or "inst.dd" in dd_args:
        return True
    else:
        return False


def umount(device):
    """ Unmount the device

        :param device: Device or mountpoint to unmount
        :type device:  string
        :returns:      None
    """
    if not device:
        return

    try:
        run_cmd(["umount", device])
    except (OSError, RunCmdError):
        pass


def mount_device(device, mnt="/media/DD/"):
    """ Mount a device and check to see if it really is a driver disk

        :param device: path to device to mount
        :type device:  string
        :param mnt:    path to mount the device on
        :type mnt:     string
        :returns:      True if it is a DD, False if not
        :rtype:        bool

        It is unmounted if it is not a DD and left mounted if it is.
    """
    try:
        run_cmd(["mount", device, mnt])
    except (OSError, RunCmdError):
        return False
    return True


def copy_repo(dd_path, dest_prefix):
    """ Copy the current arch's repository to a unique destination

        :param dd_path:     Path to the driver repo directory
        :type dd_path:      string
        :param dest_prefix: Destination directory prefix, a number is added
        :type dest_prefix:  string
        :returns:           None

        The destination directory names are in the order that the drivers
        were loaded, starting from 1
    """
    suffix = 1
    while os.path.exists(dest_prefix+str(suffix)):
        suffix += 1
    dest = dest_prefix+str(suffix)
    os.makedirs(dest)
    try:
        run_cmd(["cp", "-ar", dd_path, dest])
    except (OSError, RunCmdError):
        pass


def copy_file(src, dest):
    """ Copy a file

        :param src:  Source file
        :type src:   string
        :param dest: Destination file
        :type dest:  string
        :returns:    None
    """
    try:
        run_cmd(["cp", "-a", src, dest])
    except (OSError, RunCmdError):
        pass


def move_file(src, dest):
    """ Move a file

        :param src:  Source file
        :type src:   string
        :param dest: Destination file
        :type dest:  string
        :returns:    None
    """
    try:
        run_cmd(["mv", "-f", src, dest])
    except (OSError, RunCmdError):
        pass


def find_dd(mnt="/media/DD"):
    """ Find all suitable DD repositories under a path

        :param mnt: Top of the directory tree to search
        :type mnt:  string
        :returns:   list of DD repositories
        :rtype:     list
    """
    dd_repos = []
    arch = os.uname()[4]
    for root, dirs, files in os.walk(mnt, followlinks=True):
        if "rhdd3" in files and "rpms" in dirs and \
          os.path.exists(root+"/rpms/"+arch):
            dd_repos.append(root+"/rpms/"+arch)
    log.debug("Found repos - %s", " ".join(dd_repos))
    return dd_repos


def get_module_set(fname):
    """ Read a module list and return a set of the names

        :param fname: Full path to filename
        :type fname:  string
        :returns:     set of the module names
    """
    modules = set()
    if os.path.exists(fname):
        with open(fname, "r") as f:
            for line in f:
                mod_args = line.strip().split()
                if mod_args:
                    modules.update([mod_args[0]])
    return modules

def to_modname(modfile):
    return os.path.basename(modfile)[:-3].replace('-','_')

def reload_modules(newfiles):
    """ Reload new module versions from /lib/modules/<kernel>/updates/
    """
    try:
        run_cmd(["depmod", "-a"])
    except (OSError, RunCmdError):
        pass

    # Make a list of modules added since startup
    startup_modules = get_module_set("/tmp/dd_modules")
    current_modules = get_module_set("/proc/modules")
    new_modules = current_modules.difference(startup_modules)
    log.debug("new_modules = %s", " ".join(new_modules))

    # And a list of modules contained in the disk we just extracted
    dd_modules = set(to_modname(f) for f in newfiles if f.endswith(".ko"))
    # TODO: what modules do we unload when there's new firmware?
    log.debug("dd_modules = %s", " ".join(dd_modules))
    new_modules.update(dd_modules)

    # I think we can just iterate once using modprobe -r to remove unused deps
    for module in new_modules:
        try:
            run_cmd(["modprobe", "-r", module])
        except (OSError, RunCmdError):
            pass

    time.sleep(2)

    # Reload the modules, using the new versions from /lib/modules/<kernel>/updates/
    try:
        run_cmd(["udevadm", "trigger"])
    except (OSError, RunCmdError):
        pass


class Driver(object):
    def __init__(self):
        self.source = ""
        self.name = ""
        self.flags = ""
        self.description = []
        self.selected = False

    @property
    def args(self):
        return ["--%s" % a for a in self.flags.split()]

    @property
    def rpm(self):
        return self.source


def fake_drivers(num):
    """ Generate a number of fake drivers for testing
    """
    drivers = []
    for i in range(0, num):
        d = Driver()
        d.source = "driver-%d" % i
        d.flags = "modules"
        drivers.append(d)
    return drivers


def dd_list(dd_path, kernel_ver=None, anaconda_ver=None):
    """ Build a list of the drivers in the directory

        :param dd_path: Path to the driver repo
        :type dd_path:  string
        :returns:       list of drivers
        :rtype:         Driver object

        By default none of the drivers are selected
    """
    if not kernel_ver:
        kernel_ver = os.uname()[2]
    if not anaconda_ver:
        anaconda_ver = "19.0"

    try:
        outlines = run_cmd(["dd_list", "-k", kernel_ver, "-a", anaconda_ver, "-d", dd_path])[1]
    except (OSError, RunCmdError):
        return []

    # Output format is:
    #   source rpm\n
    #   name\n
    #   flags\n
    #   description (multi-line)\n
    #   ---\n
    drivers = []
    new_driver = Driver()
    line_idx = 0
    for line in outlines.splitlines():
        log.debug(line)
        if line == "---":
            drivers.append(new_driver)
            new_driver = Driver()
            line_idx = 0
        elif line_idx == 0:
            new_driver.source = line
            line_idx += 1
        elif line_idx == 1:
            new_driver.name = line
            line_idx += 1
        elif line_idx == 2:
            new_driver.flags = line
            line_idx += 1
        elif line_idx == 3:
            new_driver.description.append(line)

    return drivers


def dd_extract(driver, dest_path="/updates/", kernel_ver=None):
    """ Extract a driver rpm to a destination path

        :param driver:    Driver to extract
        :type driver:     Driver object
        :param dest_path: Top directory of the destination path
        :type dest_path:  string
        :returns:         list of paths to extracted firmware and modules

        This extracts the driver's files into 'dest_path' (which defaults
        to /updates/ so that the normal live updates handling will overlay
        any binary or library updates onto the initrd automatically.
    """
    if not kernel_ver:
        kernel_ver = os.uname()[2]

    cmd = ["dd_extract", "-k", kernel_ver]
    cmd += driver.args
    cmd += ["--rpm", driver.rpm, "--directory", dest_path]
    log.info("Extracting files from %s", driver.rpm)

    # make sure the to be used directory exists
    if not os.path.isdir(dest_path):
        os.makedirs(dest_path)

    try:
        run_cmd(cmd)
    except (OSError, RunCmdError):
        log.error("dd_extract failed, skipped %s", driver.rpm)
        return

    # Create the destination directories
    initrd_updates = "/lib/modules/" + os.uname()[2] + "/updates/"
    ko_updates = dest_path + initrd_updates
    initrd_firmware = "/lib/firmware/updates/"
    firmware_updates = dest_path + initrd_firmware
    for d in (initrd_updates, ko_updates, initrd_firmware, firmware_updates):
        if not os.path.exists(d):
            os.makedirs(d)

    filelist = []

    # Copy *.ko files over to /updates/lib/modules/<kernel>/updates/
    for root, _dirs, files in os.walk(dest_path+"/lib/modules/"):
        if root.endswith("/updates") and os.path.isdir(root):
            continue
        for f in (f for f in files if f.endswith(".ko")):
            src = root+"/"+f
            filelist.append(src)
            copy_file(src, ko_updates)
            move_file(src, initrd_updates)

    # Copy the firmware updates
    for root, _dirs, files in os.walk(dest_path+"/lib/firmware/"):
        if root.endswith("/updates") and os.path.isdir(root):
            continue
        for f in (f for f in files):
            src = root+"/"+f
            filelist.append(src)
            copy_file(src, firmware_updates)
            move_file(src, initrd_firmware)

    # Tell our caller about the newly-extracted stuff
    return filelist


# an arbitrary value to signal refreshing the menu contents
DoRefresh = True

def selection_menu(items, title, info_func, multi_choice=True, refresh=False):
    """ Display menu and let user select one or more choices.

        :param items: list of items
        :type items: list of objects (with the 'selected' property/attribute if
                     multi_choice=True is used)
        :param title: title for the menu
        :type title: str
        :param info_func: function providing info about items
        :type info_func: item -> str
        :param multi_choice: whether it is a multiple choice menu or not
        :type multi_choice: bool
        :returns: the selected item in case of multi_choice=False and user did
                  selection, None otherwise
    """

    page_length = 20
    page = 1
    num_pages = len(items) / page_length
    if len(items) % page_length > 0:
        num_pages += 1

    if multi_choice:
        choice_format = "[%s]"
    else:
        choice_format = ""
    format_str = "%3d) " + choice_format + " %s"

    while True:
        # show a page of items
        print("\nPage %d of %d" % (page, num_pages))
        print(title)
        if page * page_length <= len(items):
            num_items = page_length
        else:
            num_items = len(items) % page_length
        for i in range(0, num_items):
            item_idx = ((page-1) * page_length) + i
            if multi_choice:
                if items[item_idx].selected:
                    selected = "x"
                else:
                    selected = " "
                args = (i+1, selected, info_func(items[item_idx]))
            else:
                args = (i+1, info_func(items[item_idx]))
            print(format_str % args)

        # Select an item to toggle, continue or change pages
        opts = ["# to select",
                "'n'-next page",
                "'p'-previous page",
                "'c'-continue"]
        if multi_choice:
            opts[0] = "# to toggle selection"
        if refresh:
            opts.insert(1,"'r'-refresh")
        idx = raw_input(''.join(['\n',
                                ", ".join(opts[:-1]),
                                " or ", opts[-1], ": "]))
        if idx.isdigit() and not (int(idx) < 1 or int(idx) > num_items):
            item_idx = ((page-1) * page_length) + int(idx) - 1
            if multi_choice:
                items[item_idx].selected = not items[item_idx].selected
            else:
                # single choice only, we can return now
                return items[item_idx]
        elif idx.lower() == 'n':
            if page < num_pages:
                page += 1
            else:
                print("Last page")
        elif idx.lower() == 'p':
            if page > 1:
                page -= 1
            else:
                print("First page")
        elif idx.lower() == 'r' and refresh:
            return DoRefresh
        elif idx.lower() == 'c':
            return
        else:
            print("Invalid selection")

def select_drivers(drivers):
    """ Display pages of drivers to be loaded.

        :param drivers: Drivers to be selected by the user
        :type drivers:  list of Driver objects
        :returns:       None
    """
    if not drivers:
        return

    selection_menu(drivers, "Select drivers to install",
                   lambda driver: driver.source)

def process_dd(dd_path):
    """ Handle installing modules, firmware, enhancements from the dd repo

        :param dd_path: Path to the driver repository
        :type dd_path:  string
        :returns:       None
    """
    drivers = dd_list(dd_path)
    log.debug("drivers = %s", " ".join([d.rpm for d in drivers]))

    # If interactive mode or rhdd3.rules pass flag to deselect by default?
    if os.path.exists(dd_path+"/rhdd3.rules") or is_interactive():
        select_drivers(drivers)
        if not any((d.selected for d in drivers)):
            return
    else:
        map(lambda d: setattr(d, "selected", True), drivers)

    # Copy the repository for Anaconda to use during install
    copy_repo(dd_path, "/updates/run/install/DD-")

    extracted = []

    for driver in filter(lambda d: d.selected, drivers):
        extracted += dd_extract(driver, "/updates/")

        # Write the package names for all modules and firmware for Anaconda
        if "modules" in driver.flags or "firmwares" in driver.flags:
            with open("/run/install/dd_packages", "a") as f:
                f.write("%s\n" % driver.name)

    reload_modules(extracted)


def select_dd(device):
    """ Mount a device and check it for Driver Update repos

        :param device: Path to the device to mount and check
        :type device:  string
        :returns:      None
    """
    mnt = "/media/DD/"
    if not os.path.isdir(mnt):
        os.makedirs(mnt)
    if not mount_device(device, mnt):
        return

    dd_repos = find_dd(mnt)
    for repo in dd_repos:
        log.info("Processing DD repo %s on %s", repo, device)
        process_dd(repo)

    # TODO - does this need to be done before module reload?
    umount(device)


def network_driver(dd_path):
    """ Handle network driver download, then scan for new OEMDRV devices.

        :param dd_path: Path to the downloaded driver rpms
        :type dd_path:  string
        :returns:       None
    """
    skip_dds = set(oemdrv_list())

    log.info("Processing Network Drivers from %s", dd_path)
    isos = glob.glob(os.path.join(dd_path, "*.iso"))
    for iso in isos:
        select_dd(iso)

    process_dd(dd_path)

    # TODO: May need to add new drivers to /tmp/dd_modules to prevent them from being unloaded

    # Scan for new OEMDRV devices and ignore dd_args
    dd_scan(skip_dds, scan_dd_args=False, skip_device_menu=True)

class DeviceInfo(object):
    def __init__(self, **kwargs):
        self.device = kwargs.get("device", None)
        self.label = kwargs.get("label", None)
        self.uuid = kwargs.get("uuid", None)
        self.fs_type = kwargs.get("fs_type", None)

    def __str__(self):
        return "%-10s %-20s %-15s %s" % (self.device or "", self.fs_type or "",
                                        self.label or "", self.uuid or "")

def parse_blkid(line):
    """ Parse a line of output from blkid

        :param line: line of output from blkid
        :param type: string
        :returns: {} or dict of NAME=VALUE pairs including "device"
        :rtype: dict

        blkid output cannot be trusted. labels may be missing or in a different
        order so we parse what we get and return a dict with their values.
    """
    import shlex

    device = {"device":None, "label":None, "uuid":None, "fs_type":None}
    fields = shlex.split(line)
    if len(fields) < 2 or not fields[0].startswith("/dev/"):
        return {}

    # device is in [0] and the remainder are NAME=VALUE with possible spaces
    # Use the sda1 part of device "/dev/sda1:"
    device['device'] = fields[0][5:-1]
    for f in fields[1:]:
        if "=" in f:
            (key, val) = f.split("=", 1)
            if key == "TYPE":
                key = "fs_type"
            device[key.lower()] = val
    return device

def select_iso():
    """ Let user select device and DD ISO on it.

        :returns: path to the selected ISO file and mountpoint to be unmounted
                  or (None, None) if no ISO file is selected
        :rtype: (str, str)
    """
    header = "      %-10s %-20s %-15s %s" % ("DEVICE", "TYPE", "LABEL", "UUID")

    iso_dev = DoRefresh
    while iso_dev is DoRefresh:
        try:
            _ret, out = run_cmd(["blkid"])
        except (OSError, RunCmdError):
            return (None, None)

        devices = []
        for line in out.splitlines():
            dev = parse_blkid(line)
            if dev:
                devices.append(DeviceInfo(**dev))

        iso_dev = selection_menu(devices,
                                 "Driver disk device selection\n" + header,
                                 str, multi_choice=False, refresh=True)

    if not iso_dev:
        return (None, None)

    mnt = "/media/DD-search"
    if not os.path.isdir(mnt):
        os.makedirs(mnt)
    if not mount_device("/dev/" + iso_dev.device, mnt):
        print("===Cannot mount the chosen device!===\n")
        return select_iso()

    # is this device a Driver Update Disc?
    if find_dd(mnt):
        umount(mnt) # BLUH. unmount it first so select_dd can mount it OK
        return ("/dev/" + iso_dev.device, None)

    # maybe it's a device containing multiple DUDs - let the user pick one
    isos = list()
    for dir_path, _dirs, files in os.walk(mnt):
        # trim the mount point path
        rel_dir = dir_path[len(mnt):]

        # and the starting "/" (if any)
        if rel_dir.startswith("/"):
            rel_dir = rel_dir[1:]

        isos += (os.path.join(rel_dir, iso_file)
                 for iso_file in files if iso_file.endswith(".iso"))

    if not isos:
        print("===No ISO files found on %s!===\n" % iso_dev.device)
        umount(mnt)
        return select_iso()
    else:
        # mount writes out some mounting information, add blank line
        print

    # let user choose the ISO file
    dd_iso = selection_menu(isos, "Choose driver disk ISO file",
                            lambda iso_file: iso_file,
                            multi_choice=False)

    if not dd_iso:
        return (None, None)

    return (os.path.join(mnt, dd_iso), "/media/DD-search")

def dd_scan(skip_dds=None, scan_dd_args=True, skip_device_menu=False):
    """ Scan the system for OEMDRV devices and and specified by dd=/dev/<device>

        :param skip_dds:     devices to skip when checking for OEMDRV label
        :type skip_dds:      set()
        :param scan_dd_args: Scan devices passed in /tmp/dd_args or dd_args_ks
        :type scan_dd_args:  bool
        :returns:            None
    """
    dd_todo = set(oemdrv_list())

    if skip_dds is None:
        skip_dds = set()

    if skip_dds:
        dd_todo.difference_update(skip_dds)
        if dd_todo:
            log.info("Found new OEMDRV device(s) - %s", ", ".join(dd_todo))

    if scan_dd_args:
        # Add the user specified devices
        dd_devs = get_dd_args()
        dd_devs = [dev for dev in dd_devs if dev not in ("dd", "inst.dd")]
        dd_todo.update(dd_devs)
        log.info("Checking devices %s", ", ".join(dd_todo))

    # Process each Driver Disk, checking for new disks after each one
    dd_finished = dd_load(dd_todo, skip_dds=skip_dds)
    skip_dds.update(dd_finished)

    # Skip interactive selection of an iso if OEMDRV was found
    if skip_dds or skip_device_menu or not is_interactive():
        return

    # Handle interactive driver selection
    mount_point = None
    while True:
        iso, mount_point = select_iso()
        if iso:
            if iso in skip_dds:
                skip_dds.remove(iso)
            dd_load(set([iso]), skip_dds=skip_dds)
            # NOTE: we intentionally do not add the newly-loaded device to
            # skip_dds - the user might (e.g.) swap DVDs and use /dev/sr0 twice
            umount(mount_point)
        else:
            break

def dd_load(dd_todo, skip_dds=None):
    """ Process each Driver Disk, checking for new disks after each one.
        Return the set of devices that loaded stuff from.

        :param dd_todo:     devices to load drivers from
        :type dd_todo:      set
        :param skip_dds:    devices to skip when checking for OEMDRV label
        :type skip_dds:     set
        :returns:           set of devices that have been loaded
    """
    if skip_dds is None:
        skip_dds = set()

    dd_finished = set()
    while dd_todo:
        device = dd_todo.pop()
        if device in skip_dds:
            continue
        log.info("Checking device %s", device)
        select_dd(device)
        dd_finished.add(device)
        new_oemdrv = set(oemdrv_list()).difference(dd_finished, dd_todo)
        if new_oemdrv:
            log.info("Found new OEMDRV device(s) - %s", ", ".join(new_oemdrv))
        dd_todo.update(new_oemdrv)
    return dd_finished

if __name__ == '__main__':
    log.setLevel(logging.DEBUG)
    handler = SysLogHandler(address="/dev/log")
    log.addHandler(handler)
    handler = logging.StreamHandler()
    handler.setLevel(logging.INFO)
    formatter = logging.Formatter("DD: %(message)s")
    handler.setFormatter(formatter)
    log.addHandler(handler)

    if len(sys.argv) > 1:
        # Network driver source
        network_driver(sys.argv[1])
    elif os.path.exists("/tmp/DD-net/"):
        network_driver("/tmp/DD-net/")
    elif os.path.exists("/tmp/dd_args_ks"):
        # Kickstart driverdisk command, skip existing OEMDRV devices and
        # process cmdline dd entries. This will process any OEMDRV that
        # appear after loading the other drivers.
        skip_devices = set(oemdrv_list())
        dd_scan(skip_devices, skip_device_menu=True)
    else:
        # Process /tmp/dd_args and OEMDRV devices
        # Show device selection menu when inst.dd passed and no OEMDRV devices
        dd_scan()

    sys.exit(0)