#!/usr/bin/python

# xxx could use a port to python3..

import sys
import pypci
import pypcimap
import os
import time

def now():
    format = "%H:%M:%S(%Z)"
    return time.strftime(format, time.localtime())

def verbose_message(single):
    print now(), single

def modprobe(module, summary_file, args = ""):
    ret = os.system("/sbin/modprobe {} {}".format(module, args))
    if os.WEXITSTATUS(ret) == 0:
        summary_file.write("{}\n".format(module))
        return True
    else:
        return False

def main(argv):
    if len(argv) == 0:
        kernel = os.uname()[2]
    else:
        kernel = argv[0]

    if os.path.exists(kernel):
        path = kernel
    else:
        path = "/lib/modules/{}/modules.pcimap".format(kernel)

    blacklisted_modules = []
    blacklists = os.listdir("/etc/modprobe.d")
    for blacklist in blacklists:
        blf = "/etc/modprobe.d/{}".format(blacklist)
        if os.path.exists(blf):
            f = open(blf)
            for i in f.readlines():
                if i.startswith("blacklist"):
                    blacklisted_modules.append(i.split()[1])
    # unify the list
    blacklisted_modules = list(set(blacklisted_modules))

    pcimap = pypcimap.PCIMap(path)
    verbose_message("pl_hwinit: loading applicable modules")
    devices = pypci.get_devices()
    storage_devices = 0
    network_devices = 0
    missing = []
    with open('/tmp/loadedmodules', 'w') as loadedmodules:
        for slot in sorted(devices.keys()):
            dev = devices[slot]
            modules = pcimap.get(dev)
            base = (dev[4] & 0xff0000) >> 16
            if len(modules) == 0:
                if base == 0x01 or base == 0x02:
                    # storage or network device, in that order
                    missing.append((slot, dev))
            else:
                if base == 0x01:
                    storage_devices += 1
                elif base == 0x02:
                    network_devices += 1

                # FIXME: This needs improved logic in the case of multiple matching modules
                for module in modules:
                    if module not in blacklisted_modules:
                        verbose_message("pl_hwinit: found and loading module {} (%s)" % (module, slot))
                        modprobe(module, loadedmodules)

        if network_devices == 0:
            verbose_message("pl_hwinit: no supported network devices found!")
            verbose_message("pl_hwinit: the following devices were found, but have no driver:")
            lines = [ "{} x {}".format(slot, dev) for slot, dev in missing ]
            for line in lines:
                verbose_message("pl_hwinit: missing " + line)

        # XXX: could check for storage devices too, but older kernels have a lot of that built-in

        # sd_mod won't get loaded automatically
        verbose_message("pl_hwinit: loading sd_mod")
        modprobe("sd_mod", loadedmodules)

        # load usb_storage to support node conf files on flash disks
        verbose_message("pl_hwinit: loading usb_storage")
        modprobe("usb_storage", loadedmodules)

        verbose_message("pl_hwinit: loading floppy device driver")
        modprobe("floppy", loadedmodules, "floppy=0,allowed_drive_mask")

    # always wait a bit between loading the usb drivers, and checking /sys/
    # for usb devices (this isn't necessarily for waiting for mass storage files,
    # that is done below)
    verbose_message("pl_hwinit: waiting for usb system to initialize.")
    time.sleep(10)

    # sometimes, flash devices take a while to initialize. in fact, the kernel
    # intentionally waits 5 seconds for a device to 'settle'. some take even longer
    # to show up. if there are any mass storage devices on the system, try to
    # delay until they come online, up to a max delay of 30s.

    # the way this will be done is to look for files in /sys/devices that are named
    # 'bInterfaceClass', these will be a list of the usb devices on the system, and
    # their primary usb device interface class ids. The base directory these files
    # exist in will be the full path to the /sys/device entry for that device.
    # for each mass storage devices (they have an interface class value of 08),
    # we wait for a new symbolic link named 'driver' to exist in that directory,
    # indicating the kernel loaded a driver for that device.

    # usb interface class id for mass storage
    INTERFACE_CLASS_MASS_STORAGE = "08"

    # how long to wait in seconds before continuing on if devices
    # aren't available
    MAX_USB_WAIT_TIME = 30

    # low long in seconds to wait between checks
    PER_CHECK_USB_WAIT_TIME = 5


    # find out if the device identified by the /sys dir has a module
    # loaded for it. check for a symlink in the dir named driver.
    def does_device_dir_have_driver(device):
        return os.path.exists(os.path.join(device, "driver"))

    def filter_and_add(list, directory, files):
        if ("bInterfaceClass" in files and
            int(file(os.path.join(directory, "bInterfaceClass")).read(), 16) == INTERFACE_CLASS_MASS_STORAGE):
            list.append(directory)

    wait_dev_list = []
    os.path.walk("/sys/devices", filter_and_add, wait_dev_list)

    if len(wait_dev_list) > 0:
        verbose_message("pl_hwinit: found USB mass storage device(s). Attempting to wait")
        verbose_message("pl_hwinit: up to %d seconds for them to come online." % MAX_USB_WAIT_TIME)

        total_wait_time = 0
        success = False
        while total_wait_time < MAX_USB_WAIT_TIME:
            total_wait_time += PER_CHECK_USB_WAIT_TIME

            verbose_message("pl_hwinit: waiting {} seconds.".format(PER_CHECK_USB_WAIT_TIME))
            time.sleep(PER_CHECK_USB_WAIT_TIME)

            all_devices_online = True
            for device in wait_dev_list:
                if not does_device_dir_have_driver(device):
                    all_devices_online = False

            if all_devices_online:
                success = True
                verbose_message("pl_hwinit: looks like the devices are now online.")
                break
            else:
                verbose_message("pl_hwinit: not all devices online yet, waiting...")

        if success:
            verbose_message("pl_hwinit: Succesfully waited for USB mass storage devices")
            verbose_message("pl_hwinit: to come online.")
        else:
            verbose_message("pl_hwinit: One or more USB mass storage devices did not")
            verbose_message("pl_hwinit: initialize in time. Continuing anyway.")

if __name__ == "__main__":
    main(sys.argv[1:])