#!/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 sys
import os
import signal
import logging
import getopt
import errno
import re
import isi.hw.imdd.devices as devices
import isi.hw.imdd.log as imddlog
import isi.hw.imdd.report as imddreport
import isi.hw.imdd.util as imddutil
import isi.hw.bay as dbay
import isi.app.daemon as daemon
import isi.sys.disk as sysdisk

IMDD_CONFIG = '/usr/bin/isi_imdd_config'

# Custom version of usable_part that allows root on non-boot drives
def factory_usable_part(part, imdd = None):
    # Most of the time it's just a matter of whether it's in chas 0
    result = (part.slice.drive.chasn == 0)

    if imdd != None:
        # If an imdd is specified, then in addition to being in chas 0,
        # either the drive is a boot disk or the imdd is root
        result = result and (sysdisk.is_bios_disk(part.device) or 
                             imdd.part == "root")
        # However, never allow conflicts
        if imdd.partitions.lookup(part.device) == None and imdd.conflicts(part) == 1:
            result = False
    return result

# Implementation of devices.ImddList.target_mir_cnts() for --factory
def factory_target_mir_cnts(imddsort, bios_parts, total_non_bios_disks):
    # Calculate the number of drives in chassis 0
    c0drives = range(1,dbay.getCounts(0)[0]+1)          # list of baynums in c0
    c0drives = map(lambda b:dbay.bay2unit(b), c0drives) # list of unitnums
    c0drives = filter(lambda u:u != -1, c0drives)       # remove unitnums == -1
    c0drives = len(c0drives)                            # # of unitnums

    # Set imdd0 to one mirror for each drive in chassis 0
    imddr = imddsort[0]
    imddr.count = min(c0drives, devices.imdd_maxpart)
    if imddr.count > total_non_bios_disks:
        imdd0_bios_parts = imddr.count - total_non_bios_disks
        if bios_parts >= imdd0_bios_parts:
            bios_parts -= imdd0_bios_parts
        else:
            imddutil.panic('%d bios drives have only %d partitions' % 
                (imdd0_bios_parts, bios_parts))

    # Add to the other imdds until they each have two mirrors
    for x in range(2):
        for imddr in imddsort[1:]:
            if bios_parts == 0: break
            if (imddr.count >= devices.imdd_maxpart): continue
            imddr.count += 1
            bios_parts -= 1

# Determine if an isi_reformat_node (aka isi_reset) is in progress.
def isi_reset_in_progress():
    try:
        if os.stat("/var/run/isi_resetting"):
            return True
    except:
        pass
    return False

class IsiImddRebalance :

    def __init__(self, log_file = "/var/log/isi_drive_d.log") :
        self.query_only = False
        self.log_file = log_file
        self.rebalance = True
        self.factory = False
        self.read_only = False
        self.action = None
        self.verbose = False
        self.pidfile = None
        self.device = None
            
    def usage(self):
        print >> sys.stderr, "     usage: %s [-qhvme] [-d device] " \
              "[-f device] [-u device] " % sys.argv[0]
        print >> sys.stderr, ""
        print >> sys.stderr, "        -q: query only without any change."
        print >> sys.stderr, "        -m: mark device without rebalance."
        print >> sys.stderr, "        -r: read-only file system; do not write any files."
        print >> sys.stderr, "            (defeats concurrency protection)"
        print >> sys.stderr, "        -h: help."
        print >> sys.stderr, "        -v: verbose."
        print >> sys.stderr, " -d device: remove from imdd but do not mark dead."
        print >> sys.stderr, " -f device: fail the given device."
        print >> sys.stderr, " -u device: unfail the given device."
        print >> sys.stderr, "        -e: unfail all the devices."
        print >> sys.stderr, "        -x: rebalance for shipping."
        sys.exit(1)
        
    def remove_drive_from_imdd(self):
        """
        Removes the given device's imdd partitions from their
        respective imdds, if any.
        
        Returns 0 on success, non-zero on failure
        """
        ret = 0
        lines = imddutil.dd_popen("%s -i" % IMDD_CONFIG)
        if (lines == None or len(lines) < 2):
            return 1

        for line in lines[1:]:
            imdd_set = line.split()
            for part in imdd_set[1:]:
                if (not part.startswith(self.device + 's1')):
                    continue
                imddreport.report("status", "Removing %s from %s"
                           % (part, imdd_set[0]))
                result = imddutil.dd_popen("%s -x %s %s"
                        % (IMDD_CONFIG, imdd_set[0], part))
                if (result == None):
                    ret = 1
        return ret

    def imdd_partitions(self):
        lines = imddutil.dd_popen("%s -i" % IMDD_CONFIG)
        if (lines == None or len(lines) < 2):
            return None

        ret = []
        for parts in lines[1:]:
            parts = parts.split()
            ret.extend(parts[1:])
        return ret
        
    def do_down(self):
        return self.remove_drive_from_imdd()
    
    def do_fail(self):
        err = self.remove_drive_from_imdd()
        if err == 0:
            err = self.do_relabel("idead")
        return err
    
    def do_unfail(self):
        return self.do_relabel("none")
    
    def do_enable(self):
        lines = imddutil.dd_popen("/sbin/sysctl kern.disks")
        if (lines and len(lines) == 1):
            for disk in lines[0].split()[1:]:
                drive = "/dev/" + disk
                if (devices.imdd_drive_status(drive) !=
                    devices.IMDD_DRIVE_VALID):
                    continue
                parts = devices.drive2parts(drive)
                if parts == None:
                    continue
                for part in parts:
                    (p, flags, tstamp, blocks) = devices.imddvars(part)
                    if p == "idead":
                        label, imddre = devices.imdd_part_label("none", blocks = blocks)
                        devices.relabel(part, "", blocks, label, imddre)
        return 0
    
    def contains(self, imdd_parts, part):
        for d in imdd_parts:
            if d == part:
                return True
        return False

        
    def do_relabel(self, flag):
        imdd_parts = self.imdd_partitions()
        if imdd_parts == None:
            return 1
        
        err = 0
        parts = devices.drive2parts(self.device)
        if parts == None:
            return 1
        
        for part in parts:
            imddreport.report("status", 'Setting imddflags for '
                       '%s to %s' % (part, flag))

            if self.contains(imdd_parts, part):
                err = errno.EINVAL
            else:
                (p, flags, tstamp, blocks) = devices.imddvars(part)
                label, imddre = devices.imdd_part_label(flag, blocks = blocks)
                devices.relabel(part, "", blocks, label, imddre)
        return err
    
    def process_args(self) :
        try:
            opts, args = getopt.getopt(sys.argv[1:], 'qvrmef:d:u:x')
        except:
            self.usage()

        if (not len(args) == 0):
            self.usage()

        for opt in opts:
            if opt[0] == '-h':
                self.usage()
            if opt[0] == '-v':
                if self.verbose:
                    imddreport.report_hash["rebal_dbg"] = 1
                self.verbose = True
            elif opt[0] == '-q':
                self.query_only = True
                imddreport.report("status", "** Query mode: will"
                           " not add new partitions **")
            elif opt[0] == '-r':
                self.read_only = True
                self.log_file = None
            elif opt[0] == '-m':
                self.rebalance = False
            elif opt[0] == '-e':
                self.action = IsiImddRebalance.do_enable
            elif opt[0] == '-d':
                self.action = IsiImddRebalance.do_down
                self.device = opt[1]
            elif opt[0] == '-f':
                self.action = IsiImddRebalance.do_fail
                self.device = opt[1]
            elif opt[0] == '-u':
                self.action = IsiImddRebalance.do_unfail
                self.device = opt[1]
            elif opt[0] == '-x':
                self.factory = True

        if self.query_only and self.action is not None:
            self.usage()

        if not self.rebalance and self.action is None:
            self.usage()


    def start_logging(self):
        if self.log_file is None:
            hdlr = None
        else:
            hdlr = logging.FileHandler(self.log_file)
            formatter = logging.Formatter('%(asctime)s %(process)d %(message)s')
            hdlr.setFormatter(formatter)
        if not self.factory or imddreport.report_hash["rebal_dbg"] == 1:
            imddreport.report_hash["popen"] = 1
        imddlog.log_start(self.verbose, hdlr)
        
    def do_rebalance(self) :
        if not self.read_only:
            # only allow one instance to run at a time
            pid = str(os.getpid())
            self.pidfile = daemon.PidFile()
            self.pidfile.write(pid)

        def handle_sigint(signal, stack_frame):
            imddutil.dd_popen("pkill -INT -g 0 isi_imdd_config")
            if self.pidfile:
                del(self.pidfile)
            os._exit(1)

        signal.signal(signal.SIGINT, handle_sigint)
        
        # Quit if isi_reformat_node (aka isi_reset) is in progress and
        # we're not doing a rebalance for shipping
        if not self.factory and isi_reset_in_progress():
            imddreport.report("status",
                    "Abort due to isi_reformat_node in progress")
            del(self.pidfile)
            return

        if self.factory:
            # Probe all disks
            devices.disks_file_2 = ""

        if not self.query_only:
            # Disable dumps in case we panic mid-rebalance and the kernel
            # is still set to dump to a device we reassign elsewhere. We'll
            # pick a new dump partition later, either in probe or rebalance.
            devices.dumpoff()

        # populates globals such as the mount_list, uninit_drives, ... 
        devices.probe()

        # init imdd objects
        devices.attach()
        
        if (not self.query_only) :
            devices.device_lists_print()
            if self.factory:
                # Clear excess partitions
                devices.imdds.uninit_partitions.absorb(devices.imdds.dead_partitions)
                devices.imdds.uninit_partitions.clear()
                devices.sdumpimdds.uninit_partitions.absorb(devices.sdumpimdds.dead_partitions)
                devices.sdumpimdds.uninit_partitions.clear()

                # Customize rebalance
                devices.imdds.target_mir_cnts = factory_target_mir_cnts
                devices.imdds.usable_part = factory_usable_part

            devices.rebalance()

            if self.factory:
                imddreport.report("status", "Resetting imdds")
                for imdd in devices.imdds.list:
                    part = imdd.part
                    if imdd.part == 'root':
                        # Fix timestamp for root at 0
                        imdd.part = 'root tstamp_disable'
                        imdd.label(0)
                    else:
                        # For all others, mark as none--but don't clear() them
                        # since some are still mounted.
                        # XXX use tstamp_disable to prevent Bug16587 warnings
                        imdd.part = 'none tstamp_disable'
                        imdd.update()
                    # Preserve the original 'part' name...just in case we
                    # ever try a lookup() later on
                    imdd.part = part
        else :
            imddreport.report("status", "Skipped repair.")
        imddreport.report("status", "Done.")
        del(self.pidfile)

    def normalize_device(self):
        if self.device == None:
            return True
        reg_dev = re.compile('^[a-z]+\d+$')
        if reg_dev.match(self.device):
            self.device = '/dev/' + self.device
            
        reg_dev = re.compile('^/dev/[a-z]+\d+$')
        if reg_dev.match(self.device):
            return True
        return False
        
    def main(self):
        self.process_args()
        self.start_logging()
        pid = os.getpid()
        os.setpgid(0, pid)
        if self.action is not None:
            if isi_reset_in_progress():
                imddreport.report("status",
                        "Abort due to isi_reformat_node in progress")
                return 0
            if not self.normalize_device():
                return errno.EINVAL
            ret = self.action(self)
            if ret:
                return ret
        if self.rebalance:
            self.do_rebalance()
        return 0

if __name__ == "__main__" :
    ret = IsiImddRebalance().main()
    sys.exit(ret)
