#!/usr/bin/python2

import os, sys, getopt, re, json, subprocess
from stat import *

storcli_cmd = "/opt/MegaRAID/storcli/storcli64"
perccli_cmd = "/opt/MegaRAID/perccli/perccli64"
tools = [storcli_cmd, perccli_cmd]

quiet = False
path = ''
state = None


def qprint(msg):
    global quiet
    if not quiet:
        print msg


def handle_error(msg, err):
    qprint(msg)
    sys.exit(err)


def usage():
    handle_error('Usage: [-q to be quiet] -p <path> -s <start|stop>', 2)


def run(cmd):
    response = None
    str_cmd = ' '.join(cmd)

    try:
        response = subprocess.check_output(cmd)
    except OSError as err:
        handle_error("Failed running \""+str_cmd+"\" : "+err[1], 3)
    except subprocess.CalledProcessError as err:
        # For CalledProcessError, __str__ returns
        # "Command '%s' returned non-zero exit status %d" % (self.cmd, self.returncode)
        handle_error("Failed executing \""+str_cmd+"\" : "+str(err)+" ("+err.output+")", 3)
    except:
        handle_error("Unexpected error on running \""+str_cmd+"\"", 3)

    return response


def test_clitool(clitool):
    return os.access(clitool, os.W_OK)


def str2iarr(str):
    return [int(x) for x in re.findall('\d+', str)]


def pdlist2eid_slt(pdlist, did):
    eid_slt_arr = []

    for pd in pdlist:
        list = [int(x) for x in pd['EID:Slt'].split(":") if x.isdigit()]
        eid_slt = (list[0], list[1]) if len(list) == 2 else (None, list[0])

        if (pd['DG'] == '-' and pd['DID'] == did) or (pd['DG'] == did):
            eid_slt_arr.append(eid_slt)

    return eid_slt_arr


def pciaddr2cid_eid_slt(clitool, pci_addr, did):
    cmd = [clitool, '/cAll', 'show', 'J']

    try:
        controller_json_info = json.loads(run(cmd))['Controllers']

        for cid, controller in enumerate(controller_json_info):
            if controller['Command Status']['Status'] != 'Success':
                continue

            data = controller["Response Data"]
            if not (str2iarr(data["PCI Address"]) == str2iarr(pci_addr)):
                continue

            eid_slt_arr = pdlist2eid_slt(data['PD LIST'], did)
            cid = controller["Command Status"]["Controller"]
            return cid, eid_slt_arr

    except ValueError as err:
        handle_error("\""+' '.join(cmd)+"\" seems called incorrectly : "+err.message, 3)
    except KeyError as err:
        handle_error("Unexpected "+clitool+" JSON answer structure : \""+err.message+"\"", 3)
    except:
        handle_error("Failed "+clitool+" JSON answer decoding", 3)

    return None, [(None, None)]


def set_locate(clitool, pci_addr, did):
    cid, eid_slt_arr = pciaddr2cid_eid_slt(clitool, pci_addr, did)

    if cid is None or eid_slt_arr == []:
        return False

    for eid, slt in eid_slt_arr:
        if eid is None:
            cpath = '/c'+str(cid)+'/s'+str(slt)
        else:
            cpath = '/c'+str(cid)+'/e'+str(eid)+"/s"+str(slt)

        global state
        cmd = [clitool, cpath, state, "locate", "J"]
        run_status = None
        try:
            run_status = json.loads(run(cmd))['Controllers'][0]['Command Status']['Status']
        except ValueError as err:
            handle_error("\"" + ' '.join(cmd) + "\" seems called incorrectly : " + err.message, 3)
        except KeyError as err:
            handle_error("Unexpected "+clitool+" JSON answer structure : \"" + err.message + "\"", 3)

        qprint("Executing \"" + ' '.join(cmd) + "\" - "+run_status)

    return True


def locate_devpath(prefix, devname):
    tpath = os.readlink(prefix+devname)
    qprint("Matching device mapped as " + tpath)

    tpath_parts = tpath.split('/')

    host_idx = -1
    devices_idx = -1
    target_idx = -1
    for idx, part in enumerate(tpath_parts):
        if part.find("host") >= 0:
            host_idx = idx
        elif part.find("devices") >= 0:
            devices_idx = idx
        elif part.find("target") >= 0:
            target_idx = idx

    # ../../devices/pci0000:00/0000:00:02.0/0000:02:00.0/host0/target0:0:16/0:0:16:0/block/sdh/sdh1
    if host_idx == -1 or devices_idx == -1 or target_idx == -1:
        handle_error("The device by path \"" + path + "\" is not like RAID drive..", 3)

    pci_addr = tpath_parts[host_idx - 1]
    did = str2iarr(tpath_parts[target_idx])[2]

    tools_found = 0
    for tool in tools:
        if not test_clitool(tool):
            continue
        tools_found += 1
        if set_locate(tool, pci_addr, did):
            return

    if tools_found == 0:
        handle_error("No RAID controller tools found. If you have controllers then install appropriate tool, please.", 3)
        return

    handle_error("Can't find raid drive by token (pci = " + pci_addr + ", DID = " + str(did) + ")", 3)


def locate_using_enclosure(spath):
    global state
    path_prefix = spath+"/device/"

    if not os.path.exists(path_prefix):
        return False

    for entry in os.listdir(path_prefix):
        if entry.find("enclosure_device") >= 0:
            new_val = 1 if (state == "start") else 0
            with open(path_prefix+entry+"/locate", 'w') as f:
                f.write(str(new_val))
            with open(path_prefix+entry+"/locate", 'r') as f:
                return str2iarr(f.readline()) == [new_val]

    return False


def devname_no_part(devname):
    return re.findall("([a-z]+)", devname)[0]


def locate_by_path():
    global path

    try:
        st = os.stat(path)
        dev = st.st_rdev if S_ISBLK(st.st_mode) else st.st_dev

        blkdev_prefix = "/sys/class/block/"
        for devname in os.listdir(blkdev_prefix):
            with open(blkdev_prefix+devname+"/dev", 'r') as f:
                maj, min = [int(x) for x in f.readline().rstrip().split(":")]
                if os.major(dev) != maj or os.minor(dev) != min:
                    continue

            qprint("Matching drive is "+blkdev_prefix+devname)

            if locate_using_enclosure(blkdev_prefix+devname_no_part(devname)):
                qprint("New location state for "+blkdev_prefix+devname+" was setted by enclosure")
                return

            blk_dev_slaves_prefix = blkdev_prefix+devname+"/slaves/"
            slaves_list = os.listdir(blk_dev_slaves_prefix) if os.path.exists(blk_dev_slaves_prefix) else []

            if not slaves_list:
                # True RAID case
                locate_devpath(blkdev_prefix, devname)
                return

            # Software RAID/LVM cases
            for slavename in slaves_list:
                qprint("Find "+devname+"\'s slave "+slavename)
                locate_devpath(blk_dev_slaves_prefix, slavename)
            return

    except OSError as err:
        handle_error("OSError : {0}".format(err), 3)
    except IOError as err:
        handle_error("IOError : {0}".format(err), 3)

    qprint("There is no device in /sys/class/block matching device id obtained by provided path")


def main(argv):
    global path
    global state
    global quiet

    try:
        opts, args = getopt.getopt(argv, "qhp:s:", ['path=', "status="])
    except getopt.GetoptError:
        usage()
    for opt, arg in opts:
        if opt == '-h':
            usage()
        elif opt in ("-p", "--path"):
            path = arg
        elif opt in ("-s", "--status"):
            if arg == "start":
                state = arg
            elif arg == "stop":
                state = arg
        elif opt == '-q':
            quiet = True
        else:
            usage()

    if path == '' or state is None:
        usage()

    locate_by_path()


if __name__ == "__main__":
    main(sys.argv[1:])
