#!/usr/bin/env python
#-
# Copyright (c) 2013 iXsystems, Inc. All rights reserved.
#
# This file is a part of TrueNAS and may be copied and/or distributed
# without the express permission of iXsystems.


from subprocess import *
from time import time, sleep
import sys
import os
import hashlib
import libxml2
import ctypes
import logging
import logging.handlers

try:
    log = logging.getLogger('failover.fenced')
    handler = logging.handlers.SysLogHandler(address='/var/run/log')
    handler.setFormatter(logging.Formatter(fmt = '[%(name)s:%(lineno)s] %(message)s'))
    log.addHandler(handler)
    log.setLevel(logging.INFO)
except:
    pass

hostid = "00000000"

def pipeopen(command):
    try:
        log.debug("Executing: %s" % command)
    except:
        pass
    return Popen(command, stdin = PIPE, stdout = PIPE, stderr = PIPE, shell = True, close_fds = False)

# Successful -- do not try this pair again.
SUCCESS = 0
# Requeue immediately.
RETRY = 35

"""
Represents a disk on SCSI/SAS/FC bus and provide methods to register key, reserve device, etc.
"""
class Disk:
    """ Initialize an object with specified name (e.g. "da0").  """
    def __init__(self, name, resvtype = 1):
        self.diskname = name
        self.ourkey = "0"
        self.newkey = "0"
        self.remote_keys = set()
        self._works = []
        self.resvtype = resvtype

    """ Start a process in background with parameters set right now """
    def issuework_immediate(self, command, waiter = None, retries = 5):
        if waiter == None:
            waiter = self._waiter_generic
        proc = pipeopen(command)
        self._works.insert(0, (command, waiter, proc, retries))
    """ Queue work at tail of work queue """
    def queuework(self, command, waiter = None, retries = 5):
        if waiter == None:
            waiter = self._waiter_generic
        if len(self._works) == 0:
            self.issuework_immediate(command, waiter, retries)
        else:
            self._works.append((command, waiter, None, retries))
    """ Poke the object for an action.  Returns 0 when everything is done.  """
    def poke(self):
        if len(self._works) > 0:
            command, waiter, proc, retries = self._works.pop(0)
            rv = waiter(proc)
            if rv == SUCCESS:
                if len(self._works) == 0:
                    return 0
                else:
                    command, waiter, proc, retries = self._works.pop(0)
                    if proc == None and command != "":
                        proc = pipeopen(command)
                    self._works.insert(0, (command, waiter, proc, retries))
                    return 1
            elif rv == RETRY:
                assert (command != "")
                if retries > 0:
                    retries = retries - 1
                    # Retry, and requeue the waiter.
                    self.issuework_immediate(command, waiter, retries)
                    return 1
                else:
                    # Retry exhausted.
                    return -1
        else:
            """TODO: This should not happen.  Pretend successful for now."""
            """assert(False)"""
            return 0
    def poke_drain(self):
        result = self.poke()
        while result > 0:
            result = self.poke()
        return result

    """ Abort the currently running action. """
    def abort(self):
        if len(self._works) == 0:
            return
        command, waiter, proc, retries = self._works.pop(0)
        if proc != None:
            res = proc.communicate()
            retval = proc.wait()
        self._works = []

    """ Generic waiter: discard all output, returns SUCCESS when exit code is 0 or RETRY otherwise """
    def _waiter_generic(self, proc):
        if proc != None:
            res = proc.communicate()
            retval = proc.wait()
        else:
            retval = -1
        if retval == 0:
            return SUCCESS
        else:
            return RETRY

    """ Examine for keys """
    def command_fillkeys(self):
        return "/usr/local/bin/sg_persist -kn %s" % (self.diskname)
    def queue_fillkeys(self, retries = 5):
        self.queuework(self.command_fillkeys(), self._waiter_fillkeys, retries)
    def issue_fillkeys_immediate(self, retries = 5):
        self.issuework_immediate(self.command_fillkeys(), self._waiter_fillkeys, retries)
    def _waiter_fillkeys(self, proc):
        if proc != None:
            res = proc.communicate()[0].split()
            retval = proc.wait()
        else:
            retval = -1
        if retval == 0:
            keys = [ x[2:] for x in res if x.startswith("0x") ]
            self.remote_keys.clear()
            for k in keys:
                if k.startswith(hostid):
                    self.ourkey = k
                else:
                    self.remote_keys.add(k)
            return SUCCESS
        else:
            return RETRY

    def issue_updatekey(self, timestamp="00000000"):
        command = "/usr/local/bin/sg_persist -oGn --param-rk=%s --param-sark=%s%s %s" % (self.ourkey, hostid, timestamp, self.diskname)
        self.queuework(command)
        self.queue_fillkeys()

    def issue_forceupdatekey(self, timestamp="00000000"):
        command = "/usr/local/bin/sg_persist -oIn --param-sark=%s%s %s" % (hostid, timestamp, self.diskname)
        self.queuework(command)
        self.queue_fillkeys()

    def issue_clear(self):
        command = "/usr/local/bin/sg_persist -oIn --param-sark=%s55aa55aa %s" % (hostid, self.diskname)
        self.queuework(command)
        command = "/usr/local/bin/sg_persist -oCn --param-rk=%s55aa55aa %s" % (hostid, self.diskname)
        self.queuework(command)
        # Set our key to ffffffff
        self.issue_forceupdatekey("ffffffff")

    def issue_reserve(self):
        command = "/usr/local/bin/sg_persist -oRn --param-rk=%s --prout-type=%d %s" % (self.ourkey, self.resvtype, self.diskname)
        self.queuework(command)

    def issue_release(self):
        command = "/usr/local/bin/sg_persist -oLn --param-rk=%s --prout-type=%d %s" % (self.ourkey, self.resvtype, self.diskname)
        self.queuework(command)

    # Clear out all reservations.
    def fullreset(self):
        self.issue_clear()
        self.poke_drain()

class Byte(object):
    """
    Used in sysctl hack to return a byte array
    """

    def __init__(self, array, as_string=True):
        self.array = array
        self.as_string = as_string

    @property
    def value(self):
        array = ""
        for byte in self.array:
            if byte == 0 and self.as_string:
                break
            array += chr(byte)
        return array

def sysctl(name, type="CHAR"):
    libc = ctypes.CDLL('libc.so.7')
    size = ctypes.c_size_t()

    if type == "CHAR":
        if libc.sysctlbyname(str(name), None, ctypes.byref(size), None, 0) == 0:
            arg = (ctypes.c_ubyte * size.value)()
            buf = Byte(arg)
    elif type == "ULONG":
        buf = ctypes.c_ulong()
        size.value = ctypes.sizeof(buf)
        arg = ctypes.byref(buf)
    if libc.sysctlbyname(str(name), arg, ctypes.byref(size), None, 0) == 0:
        return buf.value
    
    return None

def fence_daemon():
    os.closerange(0, 1024)
    counter = 2
    while True:
        try:
            log.debug("Waken up.")
        except:
            pass
        if counter > 0xfffffffe:
            counter = 2
        else:
            counter = counter + 1
        stamp = "%08x" % counter
        try:
            log.debug("Stamp is now %s" % stamp)
        except:
            pass

        saved_key = alldisks[0].ourkey

        for disk in alldisks:
            disk.issue_updatekey(stamp)

        retries = 0
        retry_disks = alldisks
        while len(retry_disks) > 0 and retries < 10:
            retry_disks = [ disk for disk in retry_disks if disk.poke() > 0 ]
            retries = retries + 1

        if (saved_key == alldisks[0].ourkey):
            log.error("FATAL: someone cleared our registration!")
            log.error("FATAL: issuing an immediate panic.")
            proc1 = pipeopen("/usr/sbin/watchdog -t 1")
            proc2 = pipeopen("/sbin/sysctl debug.kdb.panic=1")
            proc3 = pipeopen("/sbin/shutdown -p now")
            proc1.communicate()
            proc2.communicate()
            proc3.communicate()
            # NOT REACHED
            sys.exit(255)

        sleep(5)

args = sys.argv[1:]
showconfig = 'showconfig' in args
force = 'force' in args
status = 'status' in args

if force and status:
    print "Can not force and status at the same time."
    sys.exit(255)

try:
    log.info("Entering fenced")
except:
    pass

# Get real hostid
hostid = '%08x' % (sysctl('kern.hostid', 'ULONG') | 1<<31)

# Certain Supermicro systems does not supply hostid.  Workaround by using a
# blacklist and derive the value from the license.
if hostid == 'fe4ac89c':
    hostid =  hashlib.md5(open('/data/truenas_license', 'r').read()).hexdigest()[:8]
    if hostid[0] == "0":
        hostid = "8%s" % (hostid[-7:])

try:
    log.info("Getting disk list starts")
except:
    pass

# Get a list of all 'da' disks.
geomxml = libxml2.parseDoc(sysctl('kern.geom.confxml'))

# We need to do special multipath handling.
mp_set = set()
approved_mp_prefixes = ["VIOLIN ",]
for prefix in approved_mp_prefixes:
    mp_set = mp_set.union(set([ x.content for x in geomxml.xpathEval('/mesh/class[name="DISK"]/geom/provider/config[(starts-with(descr,"%s"))]/../name' % (prefix)) ]))

# Excluded multipath devices.
excluded_mp_set = set()
excluded_mp_prefixes = ["3PAR", "HP "]
for prefix in excluded_mp_prefixes:
    excluded_mp_set = excluded_mp_set.union(set([ x.content for x in geomxml.xpathEval('/mesh/class[name="DISK"]/geom/provider/config[(starts-with(descr,"%s"))]/../name' % (prefix)) ]))

alldisks_set = set([ x for x in sysctl('kern.disks').split(' ') if x.startswith('da') ])
if showconfig:
    print "alldisks_set=%s" % alldisks_set
    print "mp_set=%s" % (mp_set)
    print "alldisks_set - mp_set - excluded_mp_set =%s" % ((alldisks_set - mp_set) - excluded_mp_set)

if len(alldisks_set) == 0:
    try:
        log.error("No disks available for doing fencing operation, exiting")
    except:
        pass
    sys.exit(3)

# Create a list consisting all disks.
#
# For multipath disks, we use type 7 reservation (write exclusive to all registrants);
# For non-multipath disks, we use type 1 reservation (write exclusive).
#
alldisks = [ Disk(x, 1) for x in ((alldisks_set - mp_set) - excluded_mp_set) ]
alldisks = alldisks + [ Disk(x, 7) for x in mp_set ]

retries = 0
fails = 0
retry_disks = alldisks
alldisks_done = []
while retries < 5 and len(retry_disks) > 0:
    new_retry_disks = []
    for disk in retry_disks:
        retval = disk.poke()
        if retval == 0:
            alldisks_done.append(disk)
        elif retval == 1:
            new_retry_disks.append(disk)
        else:
            print "WARNING: permanently disabling %d" % (disk.diskname)
            fails = fails + 1
    retry_disks = new_retry_disks
    retries = retries + 1

if showconfig:
    sys.exit(3)

alldisks = alldisks_done

try:
    log.debug("Getting disk list done.")
except:
    pass

if not force:
    # Observe if there is other node is alive.  We must quit if that's the case.
    for disk in alldisks:
        disk.queue_fillkeys()
    
    observed_keys_first = set()
    
    # Check disk status
    
    retries = 0
    retry_disks = alldisks
    while retries < 5 and len(retry_disks) > 0:
        new_retry_disks = []
        for disk in retry_disks:
            retval = disk.poke()
            if retval == 0:
                observed_keys_first = observed_keys_first.union(disk.remote_keys)
            else:
                new_retry_disks.append(disk)
        retry_disks = new_retry_disks
    
    if len(observed_keys_first) > 0:
        probes = 0
        observed_keys_new = set()
        while probes < 3:
            sleep(7)
            for disk in alldisks:
                disk.queue_fillkeys()
            retries = 0
            retry_disks = alldisks
            while retries < 25 and len(retry_disks) > 0:
                new_retry_disks = []
                for disk in retry_disks:
                    retval = disk.poke()
                    if retval == 0:
                        observed_keys_new = observed_keys_new.union(disk.remote_keys)
                    else:
                        new_retry_disks.append(disk)
                retry_disks = new_retry_disks
            if len(retry_disks) == 0 and not observed_keys_new.issubset(observed_keys_first):
                # The other node is still alive.
                print "FATAL: The other node is still alive!"
                try:
                    log.error("The other node is still alive!")
                except:
                    pass
                # Defeated!
                sys.exit(2)
            probes = probes + 1

    if status:
        print "The other node is not alive!"
        sys.exit(0)

# We reached here because the other node did not respond within a reasonable
# timeframe, or we are forced to do so.
#
# Clear existing registration and reservation

for disk in alldisks:
    disk.issue_clear()

retry_disks = alldisks
registered_disks = []
while len(retry_disks) > 0:
    new_retry_disks = []
    for disk in retry_disks:
        retval = disk.poke()
        if retval == 0:
            registered_disks.append(disk)
        elif retval == 1:
            new_retry_disks.append(disk)
        else:
            pass # This should not happen
    retry_disks = new_retry_disks

# Stamp keys to drives.
curtime = int(time())
if curtime == 0xffffffff or curtime == 0:
    curtime = 1
timestamp = '%08x' % curtime

# Register our key.

for disk in alldisks:
    disk.issue_updatekey(timestamp)

retry_disks = alldisks
registered_disks = []
while len(retry_disks) > 0:
    new_retry_disks = []
    for disk in retry_disks:
        retval = disk.poke()
        if retval == 0:
            registered_disks.append(disk)
        elif retval == 1:
            new_retry_disks.append(disk)
        else:
            pass # This should not happen
    retry_disks = new_retry_disks

# Clear existing registration and reservation

if len(registered_disks) < len(alldisks):
    if len(registered_disks) * 3 > len(alldisks) * 2:
        try:
            log.info("Proceeding with less drives.")
        except:
            pass
        alldisks = registered_disks
    else:
        try:
            log.error("Too many fails, quitting")
        except:
            pass
        # We have a problem that this script can not handle.
        sys.exit(1)

# Now that all registrations are cleared.  Proceed with reserve.
for disk in alldisks:
    disk.issue_reserve()

retry_disks = alldisks
reserved_disks = []
failed_disks = []
while len(retry_disks) > 0:
    new_retry_disks = []
    for disk in retry_disks:
        retval = disk.poke()
        if retval == 0:
            registered_disks.append(disk)
        elif retval == 1:
            new_retry_disks.append(disk)
        else:
            failed_disks.append(disk)
    retry_disks = new_retry_disks

if len(failed_disks) > 0:
    # Bail out.  We do not want to take the chance.
    sys.exit(1)

# Now we have all disks owned, start fenced.
try:
    log.info("Forking as daemon")
except:
    pass

if os.fork() != 0:
    sys.exit(0)
fence_daemon()
# NOTREACHED
sys.exit(0)
