#!/usr/bin/env python

# Copyright C 2002-2012 by Isilon Systems LLC. All rights not expressly
# granted herein are reserved.
# Created in the United States of America.
# 
# This digital media is protected by U.S. and international copyright and
# intellectual property laws. Unless otherwise specified, all information
# and screens appearing as part of this digital medium, including
# software, services, documents, text, images, icons, and logos design;
# the selection, assembly, arrangement, and design thereof; and the code
# that enables its presentation, are the sole property of Isilon Systems,
# Inc.
# 
# Except as otherwise required by applicable law, no part of any
# copyrighted material may be reproduced, modified, translated,
# distributed, published, or transmitted by any means, electronic or
# mechanical, for any purpose, without the prior written consent of Isilon
# Systems. 
#  
# Isilon and Isilon Systems are registered trademarks of Isilon Systems,
# Inc.  Legion is a trademark of Isilon Systems, Inc.  All other
# trademarks and registered trademarks are the property of their
# respective holders.

import os
import re
import sys
import syslog
import time
import getopt

import isi.app.lib.procs as procs
import isi.celog
import isi.celog.eventids
import isi.celog.query
import isi.config.arrcfg as arrcfg
import isi.gconfig as gconfig
import isi.hw.hal as hal
import isi.misc as misc

# Error return codes
ERROR_READ_ONLY_FAILED  = 3
ERROR_VALIDATION_FAILED = 4
ERROR_SHUTDOWN_FAILED   = 5

ERROR_MSG_MAP = {
    ERROR_READ_ONLY_FAILED  : "Read-only transition failed",
    ERROR_VALIDATION_FAILED : "Journal backup validation failed",
    ERROR_SHUTDOWN_FAILED   : "Final shutdown/reboot failed",
}

ADMIN_REBOOT_EVID = isi.celog.eventids.eventid_map["REBOOT_NODE_REBOOT"]
ADMIN_SHUTDOWN_EVID = isi.celog.eventids.eventid_map["REBOOT_NODE_SHUTDOWN"]
JRNL_BKUP_EVID = isi.celog.eventids.eventid_map["REBOOT_JRNL_BKUP_FAIL"]
RO_TRANS_EVID = isi.celog.eventids.eventid_map["REBOOT_RO_TRANS_FAILED"]
SHUTDOWN_ERR_EVID = isi.celog.eventids.eventid_map["REBOOT_SHUTDOWN_FAILED"]

SYSCTL_BIN = "/sbin/sysctl"
SYNC_BIN = "/bin/sync"
SYSCTL_COALESCER_ENABLED = "efs.bam.coalescer.enabled"
SYSCTL_COALESCER_STATS = "efs.bam.coalescer_stats"
COALESCER_STAT_SPACE_REGEX = r"^Space allocated: (\d+)"
COALESCER_FLUSH_WAIT_MAX = 10
REBOOT_FLAG_FILENAME = "/tmp/isi_rebooting"

devid = None
lnn = None
if arrcfg.is_local_configured():
    try:
        ac = arrcfg.config_load()
        # Leave devid and lnn None if we can't find valid values.
        if ac.local_devid > 0:
            devid = ac.local_devid
        if ac.local_lnn > 0:
            lnn = ac.local_lnn
    except arrcfg.Error:
        pass

def get_coalescer_space():
    cstats_raw = misc.sysctlbyname(SYSCTL_COALESCER_STATS)
    cstats = cstats_raw.strip().split("\n")
    for l in cstats:
        if re.match(COALESCER_STAT_SPACE_REGEX, l):
            return long(re.match(COALESCER_STAT_SPACE_REGEX, l).group(1))
    return None

class Shutdown(object):
    def __init__(self, verbose):
        self._verbose = verbose
        syslog.openlog(sys.argv[0])

    def exit(self, evid, exit_code, line, extra=None):
        # stderr gets the basic error 
        print >> sys.stderr, line
        # syslog gets the error and the extra information
        self.log(syslog.LOG_WARNING, line, extra)
        # CELOG gets the code to build the specifier
        self.send_event(evid, exit_code)
        sys.exit(exit_code)

    def send_event(self, evid, ex_code):
        if lnn is None:
            # We can't send reboot events if we don't know what our lnn is. 
            return

        try:
            msg_str = "%s" % ERROR_MSG_MAP[ex_code]
            spec = {}
            isi.celog.Event(event_id=evid, specifier = spec,
                     severity=isi.celog.EVENT_LEVEL_CRITICAL,
                     message=msg_str, send_event=True, devid=devid)
            self.log(syslog.LOG_INFO, "Sent %s event" %
                     repr(ERROR_MSG_MAP[ex_code]))
        except Exception, e:
            self.log(syslog.LOG_WARNING, "Could not send %s event (%s) "
                     % (repr(ERROR_MSG_MAP[ex_code]), e))

    def log(self, prio, line, extra=None):
        syslog.syslog(prio, line)
        if extra:
            syslog.syslog(syslog.LOG_WARNING, 'Failed command output:')
            for l in extra:
                syslog.syslog(syslog.LOG_WARNING, l)

    def out(self, line):
        if self._verbose:
            print '  %s' % line.rstrip()
            sys.stdout.flush()

    def disable_and_flush_coalescer(self):
        cenabled = misc.sysctlbyname(SYSCTL_COALESCER_ENABLED)
        if cenabled == 0L:
            return cenabled

        rc = os.system("%s %s=%d" %
                (SYSCTL_BIN, SYSCTL_COALESCER_ENABLED, 0))
        if not os.WIFEXITED(rc) or os.WEXITSTATUS(rc):
            self.log(syslog.LOG_WARNING,
                    "Failed to disable coalescer: status 0x%X" % rc)
            return cenabled

        rc = os.system(SYNC_BIN)
        if not os.WIFEXITED(rc) or os.WEXITSTATUS(rc):
            self.log(syslog.LOG_WARNING,
                    "Failed to flush coalescer: status 0x%X" % rc)
            return cenabled

        slept = 0
        cspace = get_coalescer_space()
        while cspace is not None and cspace > 0 and \
                slept < COALESCER_FLUSH_WAIT_MAX:
            time.sleep(2)
            slept = slept + 2
            cspace = get_coalescer_space()

        if cspace is not None and cspace > 0:
            self.log(syslog.LOG_WARNING,
                    "Coalescer flush incomplete after %d seconds; "
                    "proceeding with shutdown/reboot" % slept)

        return cenabled

    def enable_coalescer(self, cstate_orig):
        cenabled = misc.sysctlbyname(SYSCTL_COALESCER_ENABLED)
        if cenabled == cstate_orig:
            return

        rc = os.system("%s %s=%d" %
                (SYSCTL_BIN, SYSCTL_COALESCER_ENABLED, 1))
        if not os.WIFEXITED(rc) or os.WEXITSTATUS(rc):
            self.log(syslog.LOG_WARNING,
                    "Failed to re-enable coalescer: status 0x%X" % rc)

    def get_shutdown_type(self, argv):
        """ Determine the shutdown shutdown_type (reboot|halt) based
            on argv
            """
        reboot_re = re.compile('^\-[^-]*r.*$')
        shutdown_type = 'halt'
        for arg in argv:
            if reboot_re.match(arg):
                shutdown_type = 'reboot'
                break

        return shutdown_type

    def save_journal(self):
        """Switches the node to read-only mode via the system-shutdown
           flag, then validates the backup journal. Any failure will
           result in shutdown being aborted.
           """

        # No journal to save on nodes without NVRAM (e.g., accelerators)
        if not hal.hasNvram():
            return

        # Do not need to save here if the IFS is not mounted
        if not misc.ifs_mounted():
            return

        # If the node is unconfigured or reset, we don't need to save
        # the journal.
        if not arrcfg.is_local_configured():
            return

        print "Disabling and flushing coalescer"
        cenabled = self.disable_and_flush_coalescer()

        print 'Putting node into read-only mode for shutdown'

        # First switch the node to read-only. Fail shutdown on error.
        cmd = '/usr/bin/isi_hwtools/isi_read_only --set=system-shutdown -v'
        (error, output) = procs.proc_cmd_output(cmd, self.out)
        if error:
            self.enable_coalescer(cenabled)
            self.exit(RO_TRANS_EVID, ERROR_READ_ONLY_FAILED,
                'ERROR: Switch to read-only mode failed. Shutdown aborted',
                output)

        print 'Validating backup journal before executing shutdown'

        # Next, validate the journal save was successful. This is
        # probably unnecessary if we actually switched from read/write
        # to read-only with the previous command, but for the case
        # where the node was already read-only for some reason (and
        # thus has supposedly had a saved journal sitting around for
        # some amount of time), it's probably wise to double-check
        # before shutting down.
        cmd = '/usr/sbin/isi_save_journal --validate'
        (error, output) = procs.proc_cmd_output(cmd, self.out)
        if error:
            self.enable_coalescer(cenabled)
            self.exit(JRNL_BKUP_EVID, ERROR_VALIDATION_FAILED,
                    'ERROR: Validation failed for backup journal. '
                    'Shutdown aborted', output)

    def isi_flush(self):
        """ Runs isi_flush, because it takes more than 2 minutes """
        print 'Flushing caches before executing shutdown'

        cmd = '/usr/sbin/isi_flush'
        (error, output) = procs.proc_cmd_output(cmd, self.out)

    def send_shutdown_event(self, argv, critical=None):
        """Sends the 'node shut down/reboot by admin' event.
        """
        self.log(syslog.LOG_WARNING, "Trying to send shutdown event.")
        if lnn is None:
            # We can't send reboot events if we don't know what our lnn is. 
            self.log(syslog.LOG_WARNING, "Unknown LNN. Not sending event.")
            return

        reason = " for maintenance purposes"
        severity = isi.celog.EVENT_LEVEL_INFO
        if critical:
            reason = ": %s" % critical
            severity = isi.celog.EVENT_LEVEL_CRITICAL

        # Determine the shutdown shutdown_type (reboot|shutdown)
        shutdown_type = self.get_shutdown_type(argv)
        if shutdown_type == 'reboot':
            evid = ADMIN_REBOOT_EVID
            msg="Node {lnn} is being rebooted%s" % reason
            print 'Sending reboot event'
        elif shutdown_type == 'halt':
            evid = ADMIN_SHUTDOWN_EVID
            msg="Node {lnn} is being shutdown%s" % reason
            print 'Sending shutdown event'

        try:
            ev = isi.celog.Event(event_id=evid,
                    severity=severity,
                    message=msg,
                    send_event=True,
                    devid=devid)
        except Exception, e:
            self.log(syslog.LOG_WARNING, "Failed to send event (%s)" % e)
            print "Failed to send event (%s)" % e
            return

        # Check for the event
        max_time = 30
        print "\tWaiting for %d seconds for the event to appear " \
                "in the database..." % max_time
        if isi.celog.query.wait_for_event(ev.instance_id, max_time=max_time):
            print "\tFound the event..."
        else:
            msg = "\tNo event found in the database. Unsure if one was sent."
            self.log(syslog.LOG_WARNING, msg)
            print msg

    # Quick validation of shutdown args. This doesn't otherwise happen until
    # we get into shutdown_real.
    # See src/sbin/shutdown/shutdown_real.c::get_offset()
    def validate_shutdown_args(self, argv):

        optlist, nonopts = getopt.getopt(argv[1:], '-hknopr')
        try:
            timearg = nonopts[0]
        except IndexError:
            raise Exception("Time required.")

        if timearg == "now":
            return

        if timearg[0] == '+' and not timearg[1:].isdigit():
            raise Exception("Bad time format.")

        colon = timearg.find(':')
        if colon != -1:
            if len(timearg[colon:] != 3):
                raise Exception("Bad time format.")

            timearg.replace(':', '')

        if not timearg.isdigit():
            raise Exception("Bad time format.")

        if len(timearg) not in [4, 6, 8, 10]:
            raise Exception("Bad time format.")

    def shutdown(self):
        """Passes the arguments provided to this script (minus any leading
           --verbose flag we've already extracted) on to the renamed
           shutdown_real binary, which does the real work.
           """
        cmd = '/sbin/shutdown_real ' + " ".join(sys.argv[1:])
        (error, output) = procs.proc_cmd_output(cmd, self.out)
        if error:
            # It's pretty unlikely this will ever happen, as the shutdown_real
            # binary appears to only ever return 0. So barring seg faults or
            # signals, I wouldn't ever expect this code to actually run.
            # One exception is option parse error of shutdown_real.
            if not self._verbose:
                print >> sys.stderr, output
            self.exit(SHUTDOWN_ERR_EVID, ERROR_SHUTDOWN_FAILED,
                    'ERROR: Final shutdown failed', output)

    def is_kload_enabled(self):
        ctx = gconfig.open(gconfig.gcfg_nodeconfig).ctx_new()
        try:
            result = ctx.read('kload_enabled')
            return result
        except isi.error.isiSystemError:
            return False

    def run_kload(self):
        # Load kernel into memory
        os.system("/usr/sbin/kload_wrapper");


def main():
    """
    shutdown [optional args] [shutdown args]

    Optional args:
    --critical "MESSAGE"  Use Critical alerts, with specified MESSAGE
    --verbose   Verbose mode
    --dryrun    Test mode; do everything but the final shutdown
    --no_event  Don't send an event
    --unreal    Skip the final /sbin/shutdown_real
    --no_kload  Don't run kload
    --help      Print this string

    Shutdown args:
    (See man shutdown)
    """
    # Handle isilon-specific options that won't conflict with shutdown
    # arguments in the first positions
    ISI_OPTS = {
        '--critical': None,
        '--verbose': False,
        '--dryrun':  False,
        '--no_event':  False,
        '--no_reboot_flag':  False,
        '--help':  False,
        '--unreal':  False,
        '--no_kload': False,
    }
    while len(sys.argv) > 1:
        if ISI_OPTS.has_key(sys.argv[1]):
            if ISI_OPTS[sys.argv[1]] in [True, False]:
                ISI_OPTS[sys.argv[1]] = True
                sys.argv = [sys.argv[0]] + sys.argv[2:]
            else:
                if len(sys.argv) < 3:
                    print "Invalid args: Missing --critical MESSAGE string"
                    print main.__doc__
                    sys.exit(1)
                ISI_OPTS[sys.argv[1]] = sys.argv[2]
                sys.argv = [sys.argv[0]] + sys.argv[3:]
        else:
            break
    if ISI_OPTS['--help']:
        print main.__doc__
        sys.exit(0)

    shutdown = Shutdown(ISI_OPTS['--verbose'])

    try:
        shutdown.validate_shutdown_args(sys.argv)
    except Exception, e:
        print e
        print main.__doc__
        sys.exit(1)

    if not ISI_OPTS['--no_event']:
        shutdown.send_shutdown_event(sys.argv, ISI_OPTS['--critical'])
    if ISI_OPTS['--dryrun']:
        shutdown.log(syslog.LOG_WARNING, "FAKE SHUTDOWN CALLED")
        print "TEST: Not shutting down due to --dryrun option"
    else:
        # Only run if enabled on gconfig
        shutdown_type = shutdown.get_shutdown_type(sys.argv)
        if (shutdown_type == 'reboot' and
            not ISI_OPTS['--no_kload'] and shutdown.is_kload_enabled()):
            shutdown.run_kload()
        shutdown.save_journal()
        shutdown.isi_flush()

        if shutdown_type == 'reboot' and not ISI_OPTS['--no_reboot_flag']:
            #
            # Place file in /tmp to indicate we are rebooting.
            #
            # This file should not already exist.  If it does, something went
            # wrong and we should delete the file and not recreate it.
            #
            if os.path.exists(REBOOT_FLAG_FILENAME):
                os.unlink(REBOOT_FLAG_FILENAME)
            else:
                # if something goes wrong, we should just abort
                fd = None
                try:
                    fd = os.open(REBOOT_FLAG_FILENAME, 
                                 os.O_CREAT|os.O_RDWR|os.O_TRUNC)

                    os.write(fd, "%d\n" % int(time.time()))
                    os.close(fd)
                except:
                    if fd > 0:
                        os.close(fd)
                    if os.path.exists(REBOOT_FLAG_FILENAME):
                        os.unlink(REBOOT_FLAG_FILENAME)
                    pass
        if not ISI_OPTS['--unreal']:
            shutdown.shutdown()

if __name__ == '__main__':
    main()
