#!/usr/bin/python3

import getopt
import json
import os
import re
import subprocess
import sys
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(f'Failed running "{str_cmd}": {err}', 3)
    except subprocess.CalledProcessError as err:
        # For CalledProcessError, __str__ returns
        # "Command '%s' returned non-zero exit status %d" % (self.cmd, self.returncode)
        handle_error(f'Failed executing "{str_cmd}": {err} ({err.output})', 3)
    except Exception as e:
        handle_error(f'Unexpected error on running "{str_cmd}": {e}', 3)

    return response


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


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


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(f'"{" ".join(cmd)}" seems called incorrectly: {err}', 3)
    except KeyError as err:
        handle_error(f'Unexpected {clitool} JSON answer structure: "{err}"', 3)
    except:
        handle_error(f'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 = f'/c{cid}/s{slt}'
        else:
            cpath = f'/c{cid}/e{eid}/s{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(f'"{" ".join(cmd)}" seems called incorrectly: "{err}"', 3)
        except KeyError as err:
            handle_error(f'Unexpected {clitool} JSON answer structure: "{err}"', 3)

        qprint(f'Executing "{" ".join(cmd)}" - {run_status}')

    return True


def locate_devpath(tpath):
    tpath = os.readlink(tpath)
    qprint(f"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(f'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(f"Can't find raid drive by token (pci = {pci_addr}, DID = {str(did)})", 3)


def locate_using_enclosure(spath):
    global state
    path_prefix = os.path.join(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(os.path.join(path_prefix, entry, "locate"), 'w') as f:
                f.write(str(new_val))
            with open(os.path.join(path_prefix, entry, "locate"), 'r') as f:
                return str2iarr(f.readline()) == [new_val]

    return False


def devname_no_part(devname):
    return re.findall(r"([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):
            devname_path = os.path.join(blkdev_prefix, devname)
            with open(os.path.join(devname_path, "dev"), 'r') as f:
                maj, min_v = [int(x) for x in f.readline().rstrip().split(":")]
                if os.major(dev) != maj or os.minor(dev) != min_v:
                    continue

            qprint(f"Matching drive is {devname_path}")

            if locate_using_enclosure(os.path.join(blkdev_prefix, devname_no_part(devname))):
                qprint(f'New location state for {devname_path} was setted by enclosure')
                return

            blk_dev_slaves_prefix = os.path.join(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(os.path.join(blkdev_prefix, devname))
                return

            # Software RAID/LVM cases
            for slavename in slaves_list:
                qprint(f"Find {devname}'s slave {slavename}")
                locate_devpath(os.path.join(blk_dev_slaves_prefix, slavename))
            return

    except OSError as err:
        handle_error(f"OSError: {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:])
