#!/usr/bin/python3

import tarfile
import zipfile
import tempfile
import io
import re

#
# PowerFlex DAX issues detection & remediation utilities V1.0.0
#
###############################################################################
# GENERAL UTILITIES START
###############################################################################
import logging
import argparse
import subprocess
import sys
import os
import glob
import shutil
import collections
import re

import json

logger = logging.getLogger(__name__)

###############################################################################
def verifyRequiredCommands(*args):
    for curCmd in args:
        if shutil.which(curCmd) == None:
            logger.error(
                "Required command {} is not available.".format(curCmd))
            return False
        else:
            logger.info("Required command {} is available.".format(curCmd))
            
    return True

###############################################################################
def setupLogging(logFileName):
    logging.basicConfig(level=logging.INFO,
                        handlers=(logging.StreamHandler(sys.stdout),
                                  logging.FileHandler(
                                      os.path.join(os.path.dirname(__file__),
                                                   logFileName))),
                        format="%(asctime)s - %(levelname)s - %(message)s")



###############################################################################
BadNvdimmStatus = collections.namedtuple(
    "BadNvdimmStatus", [ "badDimms", "badNamespaces" ])


###############################################################################
HostStatus = collections.namedtuple(
    "HostStatus", [ "hostId",
                    "hostDescr",
                    "badNvdimmStatus" ])

###############################################################################
class NdctlAnalyzer(object):
    def __init__(self):
        self.__goodHostsList = []
        self.__badHostsList = []
        self.__unkHostsList = []

    def addUnknownHost(self, hostId, hostDescr, badNvdimmStatus=None):
        self.__unkHostsList.append(
            HostStatus(hostId = hostId,
                       hostDescr = hostDescr,
                       badNvdimmStatus = badNvdimmStatus))

    def processNdctlFileByPath(self, hostId, hostDescr, ndctlFilePath):
        try:
            with open(ndctlFilePath, "r") as ndctlFile:
                self.processNdctlFile(hostId, hostDescr, ndctlFile)
                
        except Exception as ex:
            logger.error(
                "Caught exception {} while processing file {}.".format(
                    ex, ndctlFilePath))

    def processNdctlFile(self, hostId, hostDescr, ndctlFile):
        try:
            ndctlJson = json.load(ndctlFile)

            badNvdimmStatus = NdctlAnalyzer.sGetBadNvdimmStatusFromNdctlJson(
                ndctlJson)

            if badNvdimmStatus is None:
                # No bad NVDIMM status means that the host is good.
                self.__goodHostsList.append(
                    HostStatus(hostId = hostId,
                               hostDescr = hostDescr,
                               badNvdimmStatus = None))
            elif NdctlAnalyzer.sIsBadStatusUnknown(badNvdimmStatus):
                # Maybe it's just "unknown" flags due to /dev/nmem permissions
                # errors
                self.__unkHostsList.append(
                    HostStatus(hostId = hostId,
                               hostDescr = hostDescr,
                               badNvdimmStatus = badNvdimmStatus))
                
            else:
                # Definitely bad.
                self.__badHostsList.append(
                    HostStatus(hostId = hostId,
                               hostDescr = hostDescr,
                               badNvdimmStatus = badNvdimmStatus))
            
        except Exception as ex:
            logger.error(
                "Caught exception {} while processing host {} {}!".format(
                    ex,
                    hostId,
                    hostDescr))
            self.__unkHostsList.append(
                HostStatus(hostId = hostId,
                           hostDescr = hostDescr,
                           badNvdimmStatus = None))
    @staticmethod
    def sGetBadNvdimmStatusFromNdctlJson(ndctlJson):
        badDimmsByName = {}
        badNamespacesByName = {}

        # 1. Walk over dimms and collect bad ones
        # 2. Walk over regions
        #    2a. Collect a union of bad flags for the region, based on its
        #        dimms.
        #    2b. Collect all namespaces in the region, associating the union
        #        of bad flags with each.
    
        dimmsJsonArray = ndctlJson[0].get("dimms")

        logger.info("ndctl JSON contains {} dimms.".format(len(dimmsJsonArray)))

        for curDimmJson in dimmsJsonArray:
            (curDimmName,
             curDimmBadStatusSet) = NdctlAnalyzer.sGetDimmJsonNodeBadStatus(
                 curDimmJson)

            if len(curDimmBadStatusSet) > 0:
                logger.error("DIMM {} is NOT HEALTHY {}.".format(
                    curDimmName, curDimmBadStatusSet))
                badDimmsByName[curDimmName] = curDimmBadStatusSet
            
            else:
                logger.info("DIMM {} is OK".format(curDimmName))

        regionsJsonArray = ndctlJson[0].get("regions")
        for curRegionJson in regionsJsonArray:
            curRegionBadFlags = set()
            curRegionName = curRegionJson["dev"]

            if curRegionJson.get("badblock_count") is not None:
                curRegionBadFlags.add("badblock_count")

            for curMappingJson in curRegionJson["mappings"]:
                curDimmName = curMappingJson["dimm"]
                curDimmBadFlags = badDimmsByName.get(curDimmName)

                if curDimmBadFlags is not None:
                    curRegionBadFlags.update(curDimmBadFlags)
                    logger.error(
                        "Region {} has mapping to an unhealthy dimm {}".format(
                            curRegionName, curDimmName))

            if len(curRegionBadFlags) > 0:
                # This region is bad, so let's add all its namespaces
                # to the bad dictionary.
                for curNsJson in curRegionJson["namespaces"]:
                    curBadNsName = curNsJson["dev"]
                    curBadNsData = {}

                    curBadNsData["region"] = curRegionName
                    curBadNsData["mode"] = curNsJson["mode"]
                    curBadNsData["size"] = curNsJson["size"]
                    curBadNsData["align"] = curNsJson.get("align")
                    curBadNsData["region_align"] = curRegionJson["align"]
                    curBadNsData["flags"] = list(curRegionBadFlags)
                    curBadNsData["devices"] = []

                    if curNsJson["mode"] == "devdax":
                        # map is "dev" or "mem"
                        curBadNsData["map"] = curNsJson["map"]
                        
                        curNsDaxregionJson = curNsJson.get("daxregion")

                        if curNsDaxregionJson is None:
                            logger.warning(
                                ("Unable to find \"daxregion\" " +
                                 "object in {}").format(curBadNsName))
                            continue

                        curNsDaxregionDevicesJson = curNsDaxregionJson["devices"]

                        if curNsDaxregionDevicesJson is None:
                            logger.warning(
                                ("No \"devices\" section in " +
                                 "\"daxregion\" namespace {}").format(
                                     curBadNsName))
                            continue

                        for curDaxDevJson in curNsDaxregionDevicesJson:
                        
                            curDaxDevName = curDaxDevJson.get("chardev")

                            if curDaxDevName is None:
                                logger.warning(
                                    "Unable to find name for device in {}".format(
                                        curBadNsName))
                                continue

                            logger.info("Adding DAX device {} to bad list.".format(
                                curDaxDevName))
                        
                            curBadNsData["devices"].append(curDaxDevName)

                    elif curNsJson["size"] == 0:
                        logger.info("Skipping dummy namespace {}".format(
                            curBadNsName))
                        continue

                    logger.info(
                        "Adding namespace {} (mode {}) to the bad list.".format(
                        curBadNsName,
                        curBadNsData["mode"]))

                    badNamespacesByName[curBadNsName] = curBadNsData

        # Return the status tuple if we have something
        if len(badDimmsByName) + len(badNamespacesByName) > 0:
            return BadNvdimmStatus( badDimms = badDimmsByName,
                                    badNamespaces = badNamespacesByName )
        else:
            return None

    @staticmethod
    def sGetDimmJsonNodeBadStatus(dimmJson):
        badStatusSet = set()
    
        dimmJsonName = dimmJson.get("dev")
        dimmJsonHealth = dimmJson.get("health")
        dimmJsonHealthState = dimmJsonHealth.get("health_state")

        if dimmJsonHealthState != "ok":
            badStatusSet.add("health_state_{}".format(dimmJsonHealthState))
        
        if dimmJson.get("flag_failed_arm") == True:
            badStatusSet.add("flag_failed_arm")

        if dimmJson.get("flag_failed_map") == True:
            badStatusSet.add("flag_failed_map")
        
        if dimmJson.get("flag_failed_save") == True:
            badStatusSet.add("flag_failed_save")

        if dimmJson.get("flag_failed_restore") == True:
            badStatusSet.add("flag_failed_restore")

        if dimmJson.get("flag_failed_flush") == True:
            badStatusSet.add("flag_failed_flush")

        if dimmJson.get("flag_smart_event") == True:
            badStatusSet.add("flag_smart_event")

        return (dimmJsonName, badStatusSet)

    def writeHealthSummary(self, logsDir, hostIdPrefix = ""):
        logger.info("")
        logger.info("--------------------")
        logger.info("NVDIMM CHECK SUMMARY")
        logger.info("--------------------")
        logger.info("")        

        # Print good hosts
        outFileName = "{}/hosts_ok".format(logsDir)
        outFile = open(outFileName, "w")

        if len(self.__goodHostsList) > 0:
            logger.info("Hosts without errors")
            logger.info("====================")
            
            for curHost in self.__goodHostsList:
                logger.info("{} {} {}:".format(hostIdPrefix,
                                               curHost.hostId,
                                               curHost.hostDescr))
                outFile.write("{}\n".format(curHost.hostId))

            logger.info("Wrote {}".format(outFileName))
            logger.info("")
        outFile.close()

        # Print unknown hosts
        outFileName = "{}/hosts_unknown".format(logsDir)
        outFile = open(outFileName, "w")

        if len(self.__unkHostsList) > 0:
            logger.info("Hosts with UNKNOWN status")
            logger.info("=========================")
            
            for curHost in self.__unkHostsList:
                logger.info("{} {} {}:".format(hostIdPrefix,
                                               curHost.hostId,
                                               curHost.hostDescr))
                if curHost.badNvdimmStatus:
                    NdctlAnalyzer.sLogBadNvdimmStatus(curHost.badNvdimmStatus)
                outFile.write("{}\n".format(curHost.hostId))

            logger.info("Wrote {}".format(outFileName))
            logger.info("")

        outFile.close()

        # Print bad hosts
        outFileName = "{}/hosts_with_errors".format(logsDir)
        outFile = open(outFileName, "w")        
    
        if len(self.__badHostsList) == 0:
            logger.info("There are no hosts with NVDIMM errors.")
        else:
            logger.info("HOSTS WITH ERRORS")
            logger.info("=================")
            
            for curHost in self.__badHostsList:
                logger.info("{} {} {}:".format(hostIdPrefix,
                                               curHost.hostId,
                                               curHost.hostDescr))
                NdctlAnalyzer.sLogBadNvdimmStatus(curHost.badNvdimmStatus)
                outFile.write("{}\n".format(curHost.hostId))

            logger.info("Wrote {}".format(outFileName))
            logger.info("")

        outFile.close()

    @staticmethod
    def sLogBadNvdimmStatus(badNvdimmStatus):
        for (curBadDimmName, curBadDimmFlags) in \
                badNvdimmStatus.badDimms.items():
            logger.info("DIMM: {} ({})".format(
                curBadDimmName, curBadDimmFlags))

        for (curBadNsName, curBadNsData) in \
                badNvdimmStatus.badNamespaces.items():
            logger.info("NAMESPACE: {} ({}) {} ({})".format(
                curBadNsName,
                curBadNsData["mode"],
                curBadNsData["devices"],
                curBadNsData["flags"]))

    @staticmethod
    def sIsBadStatusCritical(badNvdimmStatus):
        bCritical = False
    
        logger.info("Checking for critical flags in:")
        NdctlAnalyzer.sLogBadNvdimmStatus(badNvdimmStatus)

        for (curDimmName, curDimmFlags) in badNvdimmStatus.badDimms.items():
            if (("flag_failed_restore" in curDimmFlags) or
                ("flag_failed_save" in curDimmFlags)):
                logger.error("Critical flags found in DIMM {} ({})".format(
                    curDimmName, curDimmFlags))
                bCritical = True

        for (curBadNsName, curBadNsData) in \
                badNvdimmStatus.badNamespaces.items():
            if "badblock_count" in curBadNsData["flags"]:
                logger.error("Bad blocks are present in namespace {}".format(
                    curBadNsName))
                bCritical = True

        return bCritical

    @staticmethod
    def sIsBadStatusUnknown(badNvdimmStatus):
        for (curDimmName, curDimmFlags) in badNvdimmStatus.badDimms.items():
            if (len(curDimmFlags) != 1 or
                "health_state_unknown" not in curDimmFlags):
                return False

        return True

###############################################################################
def getNvdimmBadStatusWithNdctl():
    try:
        ndctlOutput = subprocess.check_output("ndctl list -vvv", shell = True)

        ndctlJson = json.loads(ndctlOutput)

        badNvdimmStatus = NdctlAnalyzer.sGetBadNvdimmStatusFromNdctlJson(
            ndctlJson)

    except Exception as ex:
        logger.error("Caught exception {} during ndctl run.".format(ex))
        raise

    return badNvdimmStatus

###############################################################################
def destroyNamespace(nsName):
    try:
        ndctlOutput = subprocess.check_output(
            "ndctl destroy-namespace {} -f".format(nsName),
            shell = True)

    except Exception as ex:
        logger.error("Failed to destroy namespace {}".format(nsName))
        return False    

    return True

###############################################################################
class UnexpectedData(RuntimeError):
    pass

###############################################################################
# GENERAL UTILITIES END
###############################################################################


###############################################################################


# Read-only file-like wrapper class, which discards get_info command output header if it exists
class GetInfoCommandTextIOWrapper(io.TextIOBase):
    def __init__(self, buffer, encoding=None, errors=None, newline=None):
        super().__init__()
        self._closed = False
        # Normalize source as a text stream
        if isinstance(buffer, (io.RawIOBase, io.BufferedIOBase, io.BytesIO)):
            self._text = io.TextIOWrapper(
                buffer, encoding=encoding, errors=errors, newline=newline
            )
            self._owns_text = True
        elif isinstance(buffer, io.TextIOBase):
            self._text = buffer
            self._owns_text = False
        else:
            raise TypeError("buffer must be a binary or text file-like object")
        self._checked = False
        # Maintain state between reading functions
        self._read_iter = iter(self)
        self._line_buffer = ""

    def readable(self):
        return True

    def writable(self):
        return False

    def seekable(self):
        return False

    @property
    def closed(self):
        return self._closed

    def close(self):
        if not self._closed:
            if self._owns_text:
                self._text.close()
            self._closed = True

    # Iterate over lines
    # Discard the get_info command output header, if it appears at the beginning of the file
    def __iter__(self):
        line_iter = iter(self._text)
        try:
            while True:
                if self._line_buffer:
                    line = self._line_buffer
                    self._line_buffer = ""
                else:
                    line = next(line_iter)
                # Check the first line (_checked is True otherwise)
                if not self._checked:
                    if re.match(r"\[COMMAND:.*\]\s*$", line):
                        # Skip lines matching 3.x get_info header pattern
                        line = next(line_iter)
                        while re.match(r"\[[A-Z]+:.*\]\s*$", line):
                            line = next(line_iter)
                self._checked = True
                yield line
        except StopIteration:
            pass

    # Read and return one line from the stream. If size is specified, at most size bytes will be read
    def readline(self, size=-1):
        if self._line_buffer == "":
            try:
                self._line_buffer = next(self._read_iter)
            except StopIteration:
                self._line_buffer = ""
        slice_size = (
            min(size, len(self._line_buffer)) if size >= 0 else len(self._line_buffer)
        )
        return_slice = self._line_buffer[:slice_size]
        self._line_buffer = self._line_buffer[slice_size:]
        return return_slice

    # Read and return up to size bytes. If the argument is omitted, None, or negative read as much as possible
    def read(self, size=-1):
        if size is None:
            size = -1
        buffer = []
        while size != 0:
            if self._line_buffer == "":
                try:
                    self._line_buffer = next(self._read_iter)
                except StopIteration:
                    break
            slice_size = (
                min(size, len(self._line_buffer))
                if size >= 0
                else len(self._line_buffer)
            )
            return_slice = self._line_buffer[:slice_size]
            self._line_buffer = self._line_buffer[slice_size:]
            buffer.append(return_slice)
            if size > 0:
                size -= slice_size
        return "".join(buffer)


###############################################################################
def parseArgs():
    try:
        prog = __file__
    except NameError:
        prog = sys.argv[0]
    parser = argparse.ArgumentParser(prog=prog)
    parser.add_argument(
        "path",
        nargs="*",
        default=[os.path.curdir],
        help="path of get_info bundle files, directories containing get_info bundles, or zip files containing get_info bundles",
    )
    parser.add_argument(
        "-t",
        "--tmpdir",
        metavar="DIR",
        help="when extracting zip files, create temporary directory under directory DIR",
    )
    return parser.parse_args()

###############################################################################
def processGetInfos(args):
    tmpZipDir = None

    analyzer = NdctlAnalyzer()
    
    try:
        for path in args.path:
            if zipfile.is_zipfile(path):
                tmpZipDir = tempfile.TemporaryDirectory(
                    dir=args.tmpdir, prefix="check_ndctl_tmp-"
                )
                unzipGetInfos(path, tmpZipDir.name)
                tarsPath = tmpZipDir.name

            else:
                tarsPath = path

            processTarFiles(analyzer, tarsPath)

        analyzer.writeHealthSummary(os.getcwd(), "SDS")

    except Exception as ex:
        logger.error("Caught exception {}".format(ex))

###############################################################################
def processTarFiles(analyzer, tarsPath):
    
    if os.path.isdir(tarsPath):
        tarFileIter = (
            os.path.join(path, filename)
            for (path, dirs, files) in os.walk(tarsPath)
            for filename in files
            if os.path.splitext(filename)[1] in [".tgz", ".tar", ".txz"]
            or filename.endswith(".tar.gz")
        )
    else:
        tarFileIter = [tarsPath]

    for curTarFileName in tarFileIter:
        logger.info("Processing tar file {}".format(curTarFileName))

        try:
            with tarfile.open(curTarFileName, "r:*") as tar:
                processSingleTarFile(analyzer, tar, curTarFileName)

                logger.info("Finished processing tar file {}.".format(
                    curTarFileName))

        except tarfile.ReadError as ex:
            logger.error(
                "Caught exception {} while processing file {}.".format(
                    ex,
                    curTarFileName))


###############################################################################
# Extract SDS ID from rep_tgt file
# Expect binary input
def parseRepTgt(repTgtFile):
    match = re.search(b"tgtId=(?P<tgtId>[0-9a-fA-F]+)", repTgtFile.read())
    return match.group("tgtId").decode("UTF-8") if match else ""


###############################################################################
# Extract a comma-separated list of IPv4 addresses from the output of "ip address list"
# Omit loopback addresses
# Expect text input
def parseIpAddrFile(ipAddrFile):
    matches = re.finditer(
        r"inet\s+(?P<ipAddr>[0-9]+\.[0-9]+\.[0-9]+\.[0-9]+)/",
        ipAddrFile.read(),
    )
    return (
        ",".join(
            ipAddr
            for ipAddr in (match.group("ipAddr") for match in matches)
            if not ipAddr.startswith("127.0.")
        )
        if matches
        else ""
    )


###############################################################################
def processSingleTarFile(analyzer, tar, tarName):
    ndctlStatusFiles = [
        member
        for member in tar.getmembers()
        if member.name.endswith("/ndctl_list_-vvv.txt")
        or member.name.endswith("/ndctl_list_vvv.txt")
    ]

    if len(ndctlStatusFiles) > 1:
        logger.warning("More than one ndctl status files found: {}".format(
            ndctlStatusFiles))

    repTgtFiles = [
        member
        for member in tar.getmembers()
        if member.name.endswith("/sds/cfg/rep_tgt.txt")
    ]

    sdsId = ""
    if len(repTgtFiles) > 0:
        if len(repTgtFiles) > 1:
            logger.warning(
                "More than one rep_tgt.txt files found: {}, using first one".format(
                    repTgtFiles
                )
            )
        try:
            sdsId = parseRepTgt(tar.extractfile(repTgtFiles[0]))
        except:
            logger.exception("exception when getting SDS ID")
    else:
        logger.warning("Couldn't find a rep_tgt.txt file.")

    ipAddrFiles = [
        member
        for member in tar.getmembers()
        if member.name.endswith("/server/ip_-s_addr.txt")
        or member.name.endswith("/server/ip_addr_s.txt")
    ]

    ipAddrs = ""
    if len(ipAddrFiles) > 0:
        if len(ipAddrFiles) > 1:
            logger.warning(
                "More than one ip address files found: {}, using first one".format(
                    ipAddrFiles
                )
            )
        try:
            ipAddrs = parseIpAddrFile(
                GetInfoCommandTextIOWrapper(tar.extractfile(ipAddrFiles[0]))
            )
        except:
            logger.exception("exception when getting IP addresses")
    else:
        logger.warning("Couldn't find a ip address file.")

    hostId = sdsId if sdsId else ipAddrs if ipAddrs else tarName
    hostDescr = []
    hostDescr.append("get_info {}".format(tarName))
    if ipAddrs:
        hostDescr.append("IPs {}".format(ipAddrs))
    hostDescr = " ".join(hostDescr)

    # Only now do we check whether we have an NDCTL status file because
    # we want descriptive hostId and hostDescr to put into the UNKNOWN
    # list.
    if len(ndctlStatusFiles) == 0:
        logger.error("Couldn't find an ndctl query file.")
        analyzer.addUnknownHost(hostId, hostDescr)
        return

    analyzer.processNdctlFile(
        hostId,
        hostDescr,
        GetInfoCommandTextIOWrapper(tar.extractfile(ndctlStatusFiles[0])),
    )

###############################################################################
def unzipGetInfos(zipFilePath, outDir):
    logger.info("Unzipping {} to {}".format(zipFilePath, outDir))
    with zipfile.ZipFile(zipFilePath, "r") as zip_file:
        zip_file.extractall(path=outDir)
    
###############################################################################
def main():
    setupLogging("check_nvdimm_health_in_getinfos.log")
    
    args = parseArgs()

    logger.info('Started')

    processGetInfos(args)
    
    logger.info('Finished')
    
###############################################################################
if __name__ == '__main__':
    main()