#!/usr/bin/python -u

import ConfigParser
import isi.app.lib.procs as procs
import isi.hw.bay as hwbay
import optparse
import os
import signal
import sys

# (model, fwversion) : filename
firmware_updates = {}

# the config file format is an ini file with the disk model as the
# section name, the firmware version and the firmware image filename
# as the key=value options.  
#
# A firmware image filename of "None" means no updates should be done
# for that firmware version.  A firmware version of "*" is a fallback
# that will match any firmware version unless there is a string match
# for the actual firmware version.  You can force an update (upgrade
# or downgrade) to a particular firmware version with something like
# the following:
#
# [STEC    MACH8 IOPS]
# *=/var/fw/disk/STEC/F2608.bin
# 2608=None
#

def load_config(filename = "/var/fw/disk/fwupdate.ini"):
    global options
    if options.inifile:
        filename = options.inifile
    config = ConfigParser.SafeConfigParser()
    # use option strings as-is, rather than transforming to lowercase
    config.optionxform = str
    config.read(filename)
    for model in config.sections():
        for version in config.options(model):
            filename = config.get(model, version)
            if options.debug:
                print "model '%s' fw '%s' file '%s'" % \
                    (model, version, filename)
            if filename in [ "None", "none", ""]:
                firmware_updates[(model, version)] = None
            elif not os.path.exists(filename):
                print "model '%s' fw '%s' missing file '%s'" % \
                    (model, version, filename)
            else:
                firmware_updates[(model, version)] = filename
    return firmware_updates        

def get_sysctl(name):
    cmd = "/sbin/sysctl " + name
    status, lines = procs.get_cmd_output(cmd)
    if (status != 0):
        print "%s failed with %s" % (cmd, procs.status_str(status))
        return None
    return lines

def known_disks():
    disks = []
    lines = get_sysctl("kern.disks")
    if (lines != None and len(lines) == 1):
        for disk in lines[0].split()[1:]:
            disks.append(disk)
    return disks

def check(disks):
    global options
    baydisks = []
    for disk in disks:
        bay = hwbay.dev2bay(disk)
        if bay > 0:
            baydisks.append((bay, disk))
        else:
            print "bogus disk name: %s" % disk
    # process disks in bay order
    baydisks.sort()
    if options.debug:
        print "baydisks: %s" % baydisks
    result = []
    for (bay, disk) in baydisks:
        # this returns a dict with: model, firmware, serial, bay
        info = hwbay.bay2mfginfo(bay)
        if info:
            model = info['model']
            firmware = info['firmware']
            serial = info['serial']
            if options.debug:
                print "info: %s" % info
            fwfile = firmware_updates.get((model,firmware),
                                          firmware_updates.get((model,"*"),
                                                               None))
            if fwfile:
                if not options.quiet:
                    print "bay %d disk %s needs a firmware update" % (bay, disk)
                if options.verbose:
                    print "\tmodel:    %s\n\tserial:   %s\n\tfirmware: %s\n\tfile:     %s" % \
                        (model, serial, firmware, fwfile)
                result.append((bay, disk, fwfile))
            else:
                if options.verbose:
                    print "bay %d disk %s does not need a firmware update" % \
                        (bay, disk)
    return result

# update all disks in parallel
def parallel_update(worklist):
    global options
    if not options.quiet:
        print "beginning parallel disk firmware update"
    signal.signal(signal.SIGHUP, signal.SIG_IGN)
    for attempt in range(3):
        if attempt > 0:
            print "retrying %d disk firmware updates" % len(worklist)
        cmds = []
        retry = []
        for (bay, disk, fw) in worklist:
            cmds.append("/usr/bin/isi_hwtools/isi_hdfw_update -d /dev/%s -f \"%s\"" % \
                            (disk, fw))
        if options.debug:
            print "cmds: %s" % cmds
        statuslist, outputlist = procs.get_cmds_output(cmds)
        for (bay, disk, fw, status, output) in zip([bay for (bay,_,_) in worklist],
                                                   [disk for (_,disk,_) in worklist],
                                                   [fw for (_,_,fw) in worklist],
                                                   statuslist, 
                                                   outputlist):
            if status != 0:
                print "bay %d disk %s firmware update failed with %s" % \
                    (bay, disk, procs.status_str(status))
                for line in output:
                    print "%s" % line
            	# queue it up for a possible retry
                retry.append((bay, disk, fw))
            elif not options.quiet:
                print "bay %d disk %s updated with %s" % (bay, disk, fw)
                if options.verbose:
                    for line in output:
                        print "%s" % line
        if retry:
            retry.sort()
            worklist = retry
        else:
            return 0 # success
    if retry:
        return 1 # ran out of retries with disks remaining, fail
    else:
        return 0 # success

# update each disk one at a time
def sequential_update(worklist):
    global options
    if not options.quiet:
        print "beginning sequential disk firmware update"
    signal.signal(signal.SIGHUP, signal.SIG_IGN)
    for (bay, disk, fw) in worklist:
        cmd = "/usr/bin/isi_hwtools/isi_hdfw_update -d /dev/%s -f \"%s\"" % \
            (disk, fw)
        if options.debug:
            print "cmd: %s" % cmd
        for attempt in range(3):
            if attempt > 0:
                print "retrying disk firmware update for bay %d disk %s" % \
                    (bay, disk)
            status, output = procs.get_cmd_output(cmd)
            if status != 0:
                print "bay %d disk %s firmware update failed with %s" % \
                    (bay, disk, procs.status_str(status))
                for line in output:
                    print "%s" % line
            else:
                if not options.quiet:
                    print "bay %d disk %s updated with fw %s" % (bay, disk, fw)
                    if options.verbose:
                        for line in output:
                            print "%s" % line
                break
        if status != 0:
            return 1 # fail
    return 0 # success

def safety_check(worklist):
    # never update disk firmware when /ifs is mounted
    if os.path.ismount("/ifs"):
        print "\nError: disk firmware cannot be updated when /ifs is mounted"
        print "Use the isi_disk_firmware_reboot command instead."
        sys.exit(1)

    # on systems using imdd, don't update disk firmware if imdd0 or
    # imdd1 contain anything other than ramdisks, since the OS doesn't 
    # like losing filesystems that are currently in use.
    for imdd in ["imdd0", "imdd1"]:
        status, lines = procs.get_cmd_output("isi_imdd_config -g %s | grep %s" % (imdd, imdd))
        if status == 0:
            devices = lines[0].split()[1:]
            for dev in devices:
                if not dev.startswith("/dev/md"):
                    print "\nError: disk firmware cannot be updated when imdd is using disks."
                    print "Use the isi_disk_firmware_reboot command instead."
                    sys.exit(1)

    # XXX: if we ever add the ability to update flash boot disk
    # firmware, we'll need to add a gmirror safety check as well.

if __name__ == "__main__":
    global options
    parser = optparse.OptionParser()
    parser.add_option('-c', '--config',  action="store", type="string", 
                      dest="inifile", help='fw update configuration file')
    parser.add_option('-d', '--debug',  action="store_true", dest="debug",
                      help='enable debug output')
    parser.add_option('-p', '--parallel', action="store_true", dest="parallel", 
                      help='do firmware updates for all disks in parallel')
    parser.add_option('-q', '--quiet',  action="store_true", dest="quiet",
                      help='quiet output')
    parser.add_option('-s', '--sequential', action="store_true", 
                      dest="sequential",
                      help='do firmware updates one disk at a time')
    parser.add_option('-v', '--verbose',  action="store_true", dest="verbose",
                      help='enable verbose output')
    options, args = parser.parse_args()

    if not load_config():
        print "no disk firmware update package installed."
        sys.exit(1)

    if args:
        # if there are arguments, assume they're disk device names.
        if options.debug:
            print "args: %s" % args
        disks = []
        for arg in args:
            if arg.startswith("/dev/"):
                disks.append(arg[5:])
            else:
                disks.append(arg)
    else:
        # if there are no arguments, operate on all known disk devices.
        disks = known_disks()

    if options.debug:
        print "disks: %s" % disks

    worklist = check(disks)
    if options.debug:
        print "worklist: %s" % worklist

    if not worklist:
        print "no disks need firmware updates."
        sys.exit(1)
    else:
        print "%d disks need firmware updates." % len(worklist)

    if options.sequential:
        safety_check(worklist)
        result = sequential_update(worklist)
        sys.exit(result)
    elif options.parallel:
        safety_check(worklist)
        result = parallel_update(worklist)
        sys.exit(result)
    else:
        if worklist:
            sys.exit(0) # updates are needed
        else:
            sys.exit(1) # no updates are needed

