#!/usr/bin/env python3
from pyVmomi import vim
import pyVim.connect
import sys
import atexit
import argparse
import re
import siteconf

class Main:
    def __init__(self):
        self.ok = []
        self.warning = []
        self.critical = []
        self.site_conf = siteconf.read()
        self.get_args()

    def get_args(self):
        parser = argparse.ArgumentParser(
            description='check multipath redundancy')

        parser.add_argument('-s', '--site',
                            required=True,
                            action='store',
                            help='name of the site to connect to')

        parser.add_argument('-H', '--host',
                            required=True,
                            action='store',
                            help='host to check')

        self.args = parser.parse_args()

    def run(self):
        self.connect()
        host = self.get_obj([vim.HostSystem], self.args.host)
        if not host:
            print("host not found")
            exit(1)
        self.pathcheck(host)
        self.status_exit()

    def map_uuid_to_datastore(self, host):
        canonical_datastore = {}
        for ds in host.datastore:
            try:
                vmfs = ds.info.vmfs
            except:
                continue # not vmfs
            for ext in vmfs.extent:
                canonical_datastore[ext.diskName] = {
                    'name': vmfs.name,
                    'local': vmfs.local
                    }

        self.uuid_datastore = {}
        for lun in host.config.storageDevice.scsiLun:
            ds = canonical_datastore.get(lun.canonicalName)
            if not ds:
                continue # not a datastore
            #print(f"uuid {lun.uuid} name {ds['name']} local {ds['local']}")
            self.uuid_datastore[lun.uuid] = ds

    def pathcheck(self, host):
        self.map_uuid_to_datastore(host)
        storage_system = host.configManager.storageSystem
        multipath_info = storage_system.storageDeviceInfo.multipathInfo
        if multipath_info is None:
            self.critical.append("no multipath info")
            return
        for lun in multipath_info.lun:
            ds = self.uuid_datastore.get(lun.id)
            if not ds:
                #print(f"lun {lun.id} is not a datastore")
                continue # not a datastore

            active, standby, dead, other = [], [], [], []
            for path in lun.path:
                #print(f"{path.name} {path.pathState} {path.adapter} {path.transport.__class__.__name__ if path.transport else '?'}")
                if path.pathState == 'active':
                    active.append(path)
                elif path.pathState == 'standby':
                    standby.append(path)
                elif path.pathState == 'dead':
                    dead.append(path)
                else:
                    other.append(path)

            if ds['local']:
                minpath = 1
                prefix = f"{ds['name']} (local)"
            elif re.search(r'-local\d?$', ds['name']):
                minpath = 1
                prefix = f"{ds['name']} (local-by-name)"
            else:
                minpath = 2
                prefix = f"{ds['name']}"
            paths = f"{len(active)}+{len(standby)} active/standby"
            if len(dead) > 0:
                paths += f", {len(dead)}"
            if len(other) > 0:
                for path in other:
                    paths += f", 1 {path.pathState}"
            paths += " paths"
            if len(active) + len(standby) < minpath:
                self.critical.append(f"{prefix} NOT REDUNDANT: {paths}")
            elif len(dead) > 0:
                self.critical.append(f"{prefix} DEAD PATH: {paths}")
            elif len(other) > 0:
                self.warning.append(f"{prefix} OTHER STATE: {paths}")
            else:
                self.ok.append(f"{prefix}: {paths}")

    def connect(self):
        if not self.args.site in self.site_conf:
            print(f"site not found: {self.args.site}")
            exit(1)
        host = self.site_conf[self.args.site]['vcenter']
        user = self.site_conf[self.args.site]['vcenter_user']
        passwd = self.site_conf[self.args.site]['vcenter_passwd']

        si = pyVim.connect.SmartConnect(
            host = host,
            user = user,
            pwd = passwd,
            disableSslCertValidation = True
            )
        atexit.register(pyVim.connect.Disconnect, si)
        self.content = si.RetrieveContent()

    def get_obj(self, vimtype, name):
        """
        Return an object by name, if name is None the
        first found object is returned
        """
        obj = None
        container = self.content.viewManager.CreateContainerView(
            self.content.rootFolder, vimtype, True)
        for c in container.view:
            if name:
                if c.name == name:
                    obj = c
                    break
            else:
                obj = c
                break

        return obj

    def status_exit(self):
        message = []
        exitcode = 0
        if len(self.critical) > 0:
            exitcode = 2
            message.append("CRITICAL: "+ " / ".join(self.critical))
        if len(self.warning) > 0:
            if exitcode == 0:
                exitcode = 1
            message.append("WARNING: "+ " / ".join(self.warning))
        if exitcode > 0:
            print(" / ".join(message))
            sys.exit(exitcode)

        print("OK: "+ " / ".join(self.ok))
        sys.exit(exitcode)


if __name__ == "__main__":
    try:
        Main().run()
    except Exception as e:
        print(f"UNKNOWN: {e}")
        print(type(e))
        sys.exit(3)

# vim: set tabstop=4 shiftwidth=4 expandtab smarttab: