#!/usr/bin/python

import sys
import pypci
import pypcimap
import os
import time

loadedmodules = None

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

def modprobe(module, args = ""):
    ret = os.system("/sbin/modprobe %s %s" % (module, args))
    if os.WEXITSTATUS(ret) == 0:
        globals()['loadedmodules'].write("%s\n" % 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/%s/modules.pcimap" % kernel

    blacklisted_modules = []
    blacklists = os.listdir("/etc/modprobe.d")
    for blacklist in blacklists:
        blf = "/etc/modprobe.d/%s" % 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)
    print now(),"pl_hwinit: loading applicable modules"
    devices = pypci.get_devices()
    storage_devices = 0
    network_devices = 0
    missing = []
    globals()['loadedmodules'] = file('/tmp/loadedmodules', 'w')
    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:
                    print now(),"pl_hwinit: found and loading module %s (%s)" % (module, slot)
                    modprobe(module)

    if network_devices == 0:
        print now(),"pl_hwinit: no supported network devices found!"
        print now(),"pl_hwinit: the following devices were found, but have no driver:"
        print now(),"pl_hwinit: ", "\npl_hwinit: ".join(missing)

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

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

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

    print now(),"pl_hwinit: loading floppy device driver"
    modprobe("floppy","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)
    print now(),"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:
        print now(),"pl_hwinit: found USB mass storage device(s). Attempting to wait"
        print now(),"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

            print now(),"pl_hwinit: waiting %d seconds." % 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
                print now(),"pl_hwinit: looks like the devices are now online."
                break
            else:
                print now(),"pl_hwinit: not all devices online yet, waiting..."

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

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