#! /usr/bin/env python

# Copyright (C) 2004 by Isilon Systems, Inc. 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.

# Description: Set on-disk imdd information for given partitions or
# imdd partitions of given disks by loading the imdd disklabel and
# changing the imdd info to "none" (or "idead", if specified).
#
# Run as "isi_imdd_forcepart [none|idead] [disk|part [disk|part ...]]"
# e.g. "isi_imdd_forcepart twed2"
# e.g. "isi_imdd_forcepart idead twed2 twed3s1a /dev/twed4s1d ad0"
#
# See isi_imdd_remove for removing it from runtime imdd sets, which
# should be done _before_ running isi_imdd_forcepart.
#
# This script is mostly useful for clearing disks that are about to
# be moved to another node, or have just been moved from another
# node... so that they have no chance of 'infecting' the second node
# with imdd info from the first node.

import os
import re
import sys
import tempfile
import getopt

from isi.app.lib.procs import get_cmd_output

def print_usage_and_exit(app_name, how):
    print 'usage: %s (none|idead) [disk|part [disk|part ...]]' % app_name
    sys.exit(how)

#
# set_part_flags
#
# Updates the imddflags in the disklabel for the given
# partition to the given state
#
def set_part_flags(part, state, rezero):
    # Check the device
    dev = '/dev/%s' % part
    if not os.access(dev, os.F_OK):
        print "ERROR: '%s' does not exist" % dev
        return -1

    print '%s => %s' % (state, dev)

    # Create a temporary file
    (file, filename) = tempfile.mkstemp('', 'forcepart-')

    # Figure out the rezero command
    if rezero:
        rezerocmd = "| sed -e 's~   63   ~    0   ~'"
    else:
        rezerocmd = ""

    # Read and modify the disklabel
    cmd = ("/sbin/disklabel -Ar %s | sed -e 's~\(imddflags:\).*~\\1 %s~' %s > %s"
           % (dev, state, rezerocmd, filename))
    (error, output) = get_cmd_output(cmd)
    if error:
        print "ERROR: disklabel -Ar '%s' failed" % dev
        # on failure, the temp file may or may not have gotten created.
        if os.path.exists(filename):
            os.remove(filename)
        return -1

    # Write the modified disklabel
    cmd = ('/sbin/disklabel -ArR %s %s' % (dev, filename))
    (error, output) = get_cmd_output(cmd)
    if error:
        print "ERROR: disklabel -ArR '%s' failed" % dev
        os.remove(filename)
        return -1

    # Delete the tempfile
    os.remove(filename)
    return 0
    

def main(args):
    app_name = os.path.basename(args[0])
    if len(args) <= 1:
        print_usage_and_exit(app_name, -1)

    # regex for stripping '/dev/' prefixes
    strip_reg = re.compile('^/dev/')

    # known device prefixes, and regexes for recognizing
    # strict device names of them (i.e. names for the raw
    # devices, not for any slices or partitions thereof)
    known_devs = ['twed', 'ad']
    known_dev_regs = []
    for dev in known_devs:
        known_dev_regs.append(re.compile('^%s\d+$' % dev))

    # imdd partitions to modify when passed a raw device name
    imdd_parts = ['s1a', 's1d']

    # process args
    state = 'none'
    parts = []
    rezero = False
    types = ['none', 'idead']
    unsup_types = types + ['var', 'var,lkg', 'root', 'root,lkg', 'crash', 'tmp']
    i = 0
    try:
        (opts, params) = getopt.gnu_getopt(args[1:], 'h', 
                                           ['help', 'unsup', 'rezero'])
    except getopt.GetoptError, e:
        print 'ERROR: %s' % e.msg
        print_usage_and_exit(app_name, 1)
    # Handle options
    for (opt, optarg) in opts:
        if opt == '-h' or opt == '--help':
            print_usage_and_exit(app_name, 0)
        elif opt == '--unsup':
            types = unsup_types
        elif opt == '--rezero':
            rezero = True
        else:
            print 'ERROR: Unexpected option: %s' % opt
            print_usage_and_exit(app_name, -1)
    # Handle non-options
    i = 0
    for param in params:
        if i == 0:
            if param in types:
                state = re.sub(',', ' ', param)
            else:
                print 'ERROR: Unknown partition type: %s' % param
                print_usage_and_exit(app_name, -1)
        else:
            # Strip any /dev/ prefix
            param = strip_reg.sub('', param)

            # See if we've found a known raw device. If so,
            # add its two imdd partitions. If not, just add
            # whatever it is we've found and assume it's a
            # slice or partition.
            found = 0
            for reg in known_dev_regs:
                if reg.match(param):
                    found = 1
                    break
            if found:
                for imdd_part in imdd_parts:
                    parts.append('%s%s' % (param, imdd_part))
            else:
                parts.append(param)
        i += 1
        
    if len(parts) == 0:
        print_usage_and_exit(app_name, -1)
    else:
        # do the dirty work
        result = 0
        for part in parts:
            if set_part_flags(part, state, rezero) < 0:
                result = -1

    sys.exit(result)        

#
# Start here
#
if __name__ == '__main__':
    main(sys.argv)
