#!/usr/bin/env python

"""
wmcore-SysStat

Hardware monitoring utility for launched sensor daemons.
"""
import sys
import os
import getopt
import subprocess
import time
import signal
import string
from datetime import datetime
from WMCore.Agent.Daemon.Details import Details
from WMCore.Configuration import loadConfigurationFile

########################
### Helper functions ###
########################

def error(msg):
    """
    General error handler.
    """
    
    print("ERROR:", msg) 
    sys.exit(1)
    
def warning(msg):
    """
    General warning handler.
    """
    
    print("WARNING:", msg)

def extractList(targetType, arg):
    """
    Makes python list of comma sepeated string list.
    """
    
    global ITEM_TYPE
    
    targetList = arg.split(",")
    listItems = []
    for item in targetList:
        if len(item.strip()) == 0:
            continue
        if checkItem(targetType, item) == False:
            warning("Ignoring target %s %s" % (ITEM_TYPE[targetType], item))
            continue
        listItems.append(item)    
    return listItems
    
def checkOutput(arg):
    """
    Checks string if it contains the name of supported output format.
    """
    
    global OUT_XML, OUT_TUI, OUT_JSON, OUT_CSV
    
    if arg == "XML":
        return OUT_XML
    elif arg == "TUI":
        return OUT_TUI
    elif arg == "JSON":
        return OUT_JSON
    elif arg == "CSV":
        return OUT_CSV
    else:
        return None

def help(short=True):
    """
    Returns short or long help information.
    """
    
    short_help = \
    """
Usage: wmcore-SysStat --config=<WMCoreConfig.py> <command> {<specified_list>} [optional]
       Where:
       command: --start, --shutdown, --status, --restart, --help, --list=
       specified_list: --wmcomponents=, --services=, --resources=, --disks=
       optional: --ouput=
    """
    
    long_help = \
    """
Help: '--config=' option is obligatory, you must always provide WMCore
      configuration file. It must be the first option passed to
      wmcore-SysStat.
      
      This utility is case-sensitive!
      
      Only one command can be provided for wmcore-SysStat:
        --start - start sensors for all/specified items;
        --shutdown - stop sensors for all/specified items;
        --status - give status of sensors for all/specified items;
        --restart - restarts sensors for all/specified items;
        --help - prints this help information.
        --list=<list> - prints all available targets.
          Possible values: WMComponent, Service, Resource, Disk
        
      You can specify comma separeted lists for different targets:
        --wmcomponents=<list> - for WMComponents;
        --services=<list> - for Services;
        --resources=<list> - for general machine Resources;
        --disks=<list> - for machine Disks;
        
      You can speficy output format for '--status' and '--list=' 
      commands by providing '--ouput=' option. Possible values:
        TUI - Text User Interface (best for terminal);
        XML - Best for third-party applications;
        JSON - Best for web-based applications;
        CSV - Best for machine and human;
        
      <list> - comma separated list, eg. 'GridFTP,mySQL'
      
      You can turn on debug mode on wmcore-sensord, by using '--debug'
      option on wmcore-SysStat.
    """
    
    if short == True:
        return short_help
    else:
        return short_help + long_help
        
def checkItem(targetType, item):
    """
    Checks if items belongs to specific family.
    """
    
    global ITEM_WMCOMPONENT, ITEM_SERVICE, ITEM_RESOURCE, ITEM_DISK
    global config, servicesList, resourcesList
    
    if targetType == ITEM_WMCOMPONENT:
        wmComponentsList = config.listComponents_()
        if item not in wmComponentsList:
            return False
    elif targetType == ITEM_SERVICE:
        if item not in servicesList:
            return False
    elif targetType == ITEM_RESOURCE:
        if item not in resourcesList:
            return False
    elif targetType == ITEM_DISK:
        if item not in disks():
            return False
            
    return True

def disks():
    """
    Retuns list of available disks on machine.
    """
    
    disks = os.popen('ls -1 /dev/?d[a-z]').readlines()
    disks = [x.split('/')[2].rstrip() for x in disks]
    return disks

def loadConfiguration():
    """
    Loads WMCore configuration file.
    """
    
    global config, configFile
    
    try:
        if config == None:
            config = loadConfigurationFile(configFile)
            
            if "HWMon" not in dir(config):
                error("No 'HWMon' section found in WMCore configuration file! Please, check your configuration file.")
    except ImportError as ex:
        error("Can no load configuration file! Please, check configuration file (" + os.path.abspath(configFile) + ")")

def prepareList(targetType, arg):
    """
    Parse target list.
    """
    
    global ITEM_WMCOMPONENT, ITEM_SERVICE, ITEM_RESOURCE, ITEM_DISK
    global loadDefaults, config
    
    loadDefaults = False
    
    if targetType == ITEM_WMCOMPONENT:
        loadConfiguration() # must load configuratio to confirm WMComponents
        
    if len(arg) != 0:
        listItems = extractList(targetType, arg)
        if len(listItems) != 0:
            specifiedList[targetType] = listItems
    else:
        if targetType == ITEM_WMCOMPONENT:
            specifiedList[ITEM_WMCOMPONENT] = config.listComponents_()
        elif targetType == ITEM_SERVICE:
            specifiedList[ITEM_SERVICE] = list(servicesList)
        elif targetType == ITEM_RESOURCE:
            specifiedList[ITEM_RESOURCE] = resourcesList
        elif targetType == ITEM_DISK:
            specifiedList[ITEM_DISK] = disks()

def isWMComponentRunning(wmcomponent):
    """
    Checks if WMComponent is running.
    """
    
    global config
    
    daemonXml = os.path.abspath("%s/Components/%s/Daemon.xml" % (config.General.workDir, wmcomponent))
    
    if not os.path.isfile(daemonXml):
        return False, 0
        
    daemon = Details(daemonXml)
    if daemon.isAlive():
        return True, int(daemon['ProcessID'])
    else:
        return False, 0
        
def isServiceRunning(service):
    """
    Checks if Service is running.
    """
    
    global servicesList
    
    if service not in servicesList:
        return False, 0
    
    ret = os.popen("ps -C %s wwho pid" % servicesList[service]).read()
    if ret:
        return True, int(ret.strip()) 
    else:
        return False, 0

def isSensorDaemonRunning(target):
    """
    Checks if wmcore-sensord is running for target.
    """
    
    global ITEM_WMCOMPONENT, ITEM_SERVICE, ITEM_RESOURCE, ITEM_DISK
    global specifiedList
    
    patterns = {
        ITEM_WMCOMPONENT : "\\-\\-wmcomponent=%s", 
        ITEM_SERVICE     : "\\-\\-service=%s",
        ITEM_RESOURCE    : "\\-\\-resource=%s", 
        ITEM_DISK        : "\\-\\-disk=%s"
    }
    
    pattern = None
    targetTypes = (ITEM_WMCOMPONENT, ITEM_SERVICE, ITEM_RESOURCE, ITEM_DISK)
    for ttype in targetTypes:
        if ttype in specifiedList and target in specifiedList[ttype]:
            pattern = patterns[ttype] % target
            break
    
    if pattern == None:
        return False, 0
           
    ret = os.popen("ps -C python wwho pid,cmd | grep -i 'wmcore-sensord' | grep -i '%s'" % pattern).read()
    if ret:
        return True, int(ret.split()[0])
    else:
        return False, 0
    
def isSensorRunning(target):
    """
    Checks if sar/iostat is running for target.
    """
    
    ret = None
    if ITEM_DISK in specifiedList and target in specifiedList[ITEM_DISK]:
        ret = os.popen("ps -C iostat wwho pid,cmd | grep '%s'" % target).read()  
    elif ITEM_RESOURCE in specifiedList and target in specifiedList[ITEM_RESOURCE]:
        patterns = {'CPU': '\\-u', 'MEM': '\\-r', 'SWAP': '\\-W', 'LOAD': '\\-q'}
        ret = os.popen("ps -C sar wwho pid,cmd | grep '%s'" % patterns[target]).read()
    elif ITEM_WMCOMPONENT in specifiedList and target in specifiedList[ITEM_WMCOMPONENT]:
        stat = isWMComponentRunning(target)
        if stat[0] == True:
            ret = os.popen("ps -C sar wwho pid,cmd | grep '%s'" % stat[1]).read()
    elif ITEM_SERVICE in specifiedList and target in specifiedList[ITEM_SERVICE]:
        stat = isServiceRunning(target)
        if stat[0] == True:
            ret = os.popen("ps -C sar wwho pid,cmd | grep '%s'" % stat[1]).read()
    
    if ret:
        return True, int(ret.split()[0])
    else:   
        return False, 0

def kill(pid, signal):
    """
    Sends singal to specified process.
    """
    
    try:
        os.kill(int(pid), signal)
    except OSError as ex:
        err = str(err)
        if err.find("No such process") >= 0:
            pass
        else:
            raise

def formatOutput(caller, data):
    """
    Prints '--status' information in '--output=' specified fromat.
    """
    
    global CMD_TYPE, OUT_TYPE
    global output
    
    def status_XML(data):
        msg = "<?xml version=\"1.0\"?>\n"
        msg += "<targets>\n"
        for ttype in data:
            for target in data[ttype]:
                msg += "<target type=\"%s\" name=\"%s\">\n" % (ttype, target)
                for status in  data[ttype][target]:
                    msg += "<status type=\"%s\" status=\"%s\" pid=\"%s\"/>\n" % \
                           (status, data[ttype][target][status][0], data[ttype][target][status][1])
                msg += "</target>\n"
        msg += "</targets>"
        return msg
    
    # TODO: Other formats
    def status_TUI(data):
        def showPid(format, pid):
            if pid != 0:
                return str(pid)
            else:
                return format[2]
                
        def showStatus(format, status, ttype = "NONE"):
            if ttype in ["DISK", "RESOURCE"]:
                return format[2]
            else:
                return format[status]
                
        status = ["\033[41m OFF  \033[m", "\033[42m  ON  \033[m", "\033[43m  NO  \033[m"]
        
        msg = "\033[30m\033[47m%-11s | %-40s | %-5s | %-6s | %-5s | %-6s | %-5s | %-6s\033[m\n" %\
              ("Target Type", "Target Name", "Target", "-->Pid", "Daemon", "-->Pid", "Sensor", "-->Pid")
        for ttype in data:
            for target in data[ttype]:
                msg += "%-11s | %-40s | %-5s | %-6s | %-5s | %-6s | %-5s | %-6s\n" %\
                       (ttype, target, showStatus(status, data[ttype][target]['TARGET'][0], ttype), showPid(status, data[ttype][target]['TARGET'][1]), \
                       showStatus(status, data[ttype][target]['DAEMON'][0]), showPid(status, data[ttype][target]['DAEMON'][1]), \
                       showStatus(status, data[ttype][target]['SENSOR'][0]), showPid(status, data[ttype][target]['SENSOR'][1]))
                    
        return msg
    
    def list_XML(data):
        msg = "<?xml version=\"1.0\"?>\n"
        msg += "<targets>\n"
        for ttype in data:
            for target in data[ttype]:
                msg += "<target type=\"%s\">%s</target>\n" % (ttype, target)
        msg += "</targets>"
        return msg
        
    def list_TUI(data):
        msg = ""
        msg += "\033[30m\033[47m%-11s | %-40s \033[m\n" % ("Target Type", "Target")
        for ttype in data:
            for target in data[ttype]:
                msg += "%-11s | %-56s\n" % (ttype, target)
        return msg
    
    template = CMD_TYPE[caller].lower() + "_" + OUT_TYPE[output]
    if template in locals():
        return locals()[template](data)
    else:
        return "Can not find template '%s'" % template

def start():
    """
    Starts sensor daemons for specified targets.
    """
    
    global ITEM_WMCOMPONENT, ITEM_SERVICE, ITEM_RESOURCE, ITEM_DISK, ITEM_TYPE
    global specifiedList, configFile, debugMode
    
    patterns = {
        ITEM_WMCOMPONENT : "--wmcomponent=%s", 
        ITEM_SERVICE     : "--service=%s",
        ITEM_RESOURCE    : "--resource=%s", 
        ITEM_DISK        : "--disk=%s"
    }
    
    status = ["\033[41m   OFF   \033[m", "\033[42m   ON    \033[m", "\033[43m WARNING \033[m"]
    
    for ttype in specifiedList:
        for target in specifiedList[ttype]:
            statSensorDaemon = isSensorDaemonRunning(target)
            if statSensorDaemon[0]:
                print("%-7s | Sensor daemon (pid: %s) is already running for %s %s." % (status[2], str(statSensorDaemon[1]), ITEM_TYPE[ttype], target))
                continue
            
            #statSensor = isSensorRunning(target)
            #if statSensor[0]:
            #    print "%-7s | Sensor (pid: %s) is already running for %s %s." % (status[2], str(statSensor[1]), ITEM_TYPE[ttype], target)
            #    continue
              
            if ttype == ITEM_WMCOMPONENT and not isWMComponentRunning(target)[0]:
                print("%-7s | %s %s is not running." % (status[2], ITEM_TYPE[ttype], target))
                continue
                
            if ttype == ITEM_SERVICE and not isServiceRunning(target)[0]:
                print("%-7s | %s %s is not running." % (status[2], ITEM_TYPE[ttype], target))
                continue
            
            cmd = ["wmcore-sensord"]
            cmd.append("--config=%s" % configFile)
            cmd.append(patterns[ttype] % target)
            if debugMode == True:
                cmd.append("--debug")
            sensord = subprocess.Popen(cmd)
            time.sleep(1)
            sensord.poll()
            
            pid = 0
            if ttype == ITEM_WMCOMPONENT:
                pid = isWMComponentRunning(target)[1]
            elif ttype == ITEM_SERVICE:
                pid = isServiceRunning(target)[1]
            
            if sensord.returncode == None:
                if pid != 0:
                    print("%-7s | Sensor daemon (pid: %s) for %s %s (pid: %s) started." \
                          % (status[1], str(sensord.pid), ITEM_TYPE[ttype], target, pid))
                else:
                    print("%-7s | Sensor daemon (pid: %s) for %s %s started." \
                          % (status[1], str(sensord.pid), ITEM_TYPE[ttype], target))
            else:                
                if pid != 0:
                    print("%-7s | Sensor daemon for %s %s (pid: %s) did not start." % (status[0], ITEM_TYPE[ttype], target, str(pid)))
                else:
                    print("%-7s | Sensor daemon for %s %s did not start." % (status[0], ITEM_TYPE[ttype], target))
    
def shutdown():
    """
    Shutdowns sensor daemons for specified targets.
    """
    
    global ITEM_TYPE
    
    for ttype in specifiedList:
        for target in specifiedList[ttype]:
            stat = isSensorDaemonRunning(target)
            if stat[0]:
                print(">>> Shutting down sensor daemon (pid: %s) for %s %s" % (stat[1], ITEM_TYPE[ttype], target))
                kill(stat[1], signal.SIGTERM)
    
def restart():
    """
    Restarts sensor daemons for specified targets.
    """
    
    print(">>> Shutting down specified sensor daemons.")
    shutdown()
    print(">>> Please wait 10 sec.")
    startTime = datetime.today()
    while (datetime.today() - startTime).seconds <= 10:
        time.sleep(1)
    print(">>> Starting up specified sensor daemons.")
    start()

def status():
    """
    Gives sensors information.
    """
    
    def showPid(status, pid):    
        if pid != 0:
            return str(pid)
        else:
            return status[2]
    
    global ITEM_WMCOMPONENT, ITEM_SERVICE, ITEM_TYPE
    
    status = ["\033[41m OFF  \033[m", "\033[42m  ON  \033[m", "\033[43m  NO  \033[m"]
    
    if len(specifiedList) == 0:
        print("No information for specified targets.")
        return
    
    data = {}
    
    for ttype in specifiedList:
        data[ITEM_TYPE[ttype]] = {}
        for target in specifiedList[ttype]:
            statSensorDaemon = isSensorDaemonRunning(target)
            statSensor = isSensorRunning(target)
            stat = None
            
            data[ITEM_TYPE[ttype]][target] = {}
            
            if ttype == ITEM_WMCOMPONENT:
                stat = isWMComponentRunning(target)

            elif ttype == ITEM_SERVICE:
                stat = isServiceRunning(target)
            else:
                stat = [False, 0]
                
            data[ITEM_TYPE[ttype]][target]["TARGET"] = [stat[0], stat[1]]
            data[ITEM_TYPE[ttype]][target]["DAEMON"] = [statSensorDaemon[0], statSensorDaemon[1]]
            data[ITEM_TYPE[ttype]][target]["SENSOR"] = [statSensor[0], statSensor[1]]
            
    print(formatOutput(CMD_STATUS, data))

def prepareAvailableList(arg):
    """
    Prepares available WMComponents, Resources, Services, Disks list.
    """
    
    global ITEM_WMCOMPONENT, ITEM_SERVICE, ITEM_RESOURCE, ITEM_DISK, ITEM_TYPE
    global config, servicesList, resourcesList
    
    valid = ['WMComponent', 'Service', 'Resource', 'Disk']
    
    listAvail = {}
    
    loadConfiguration()
    if len(arg) == 0:
        listAvail[ITEM_WMCOMPONENT] = config.listComponents_()
        listAvail[ITEM_SERVICE] = list(servicesList)
        listAvail[ITEM_RESOURCE] = resourcesList
        listAvail[ITEM_DISK] = disks()
    else:    
        targetTypeList = arg.split(",")
        for ttype in targetTypeList:
            if len(ttype.strip()) == 0:
                continue
            if ttype in valid:
                targetType = ITEM_TYPE.index(ttype.upper())
                if targetType == ITEM_WMCOMPONENT:
                    listAvail[ITEM_WMCOMPONENT] = config.listComponents_()
                elif targetType == ITEM_SERVICE:
                    listAvail[ITEM_SERVICE] = list(servicesList)
                elif targetType == ITEM_RESOURCE:
                    listAvail[ITEM_RESOURCE] = resourcesList
                elif targetType == ITEM_DISK:
                    listAvail[ITEM_DISK] = disks()
            else:
                warning("Ignoring target type: %s." % ttype)
            
    return listAvail

def available():
    """
    Returns available WMComponents, Resources, Services, Disk string.
    """
    
    global ITEM_TYPE
    global availableList
    
    data = {}
    if len(availableList) == 0:
        return  "No targets for specified target types."
    
    for ttype in availableList:
        data[ITEM_TYPE[ttype]] = availableList[ttype]
    
    return formatOutput(CMD_LIST, data)

##########################
### Global definitions ###
##########################

servicesList = {
    'GridFTP' : 'globus-gridftp-server',
    'mySQL'   : 'mysqld'
}

resourcesList = ['CPU', 'MEM', 'SWAP', 'LOAD']

ITEM_TYPE = ['WMCOMPONENT', 'SERVICE', 'RESOURCE', 'DISK']

CMD_TYPE = ['START', 'SHUTDOWN', 'STATUS', 'RESTART', 'HELP', 'LIST']

OUT_TYPE = ['XML', 'TUI', 'JSON', 'CSV']

ITEM_WMCOMPONENT = 0
ITEM_SERVICE     = 1
ITEM_RESOURCE    = 2
ITEM_DISK        = 3

CMD_START    = 0
CMD_SHUTDOWN = 1
CMD_STATUS   = 2
CMD_RESTART  = 3
CMD_HELP     = 4
CMD_LIST     = 5

OUT_XML  = 0
OUT_TUI  = 1
OUT_JSON = 2
OUT_CSV  = 3

#################
### Main part ###
#################

valid = ['config=', 'wmcomponents=', 'services=', 'resources=', 'disks=', 'list=',
         'start', 'shutdown', 'restart', 'status', 'output=', 'help', 'debug']

try:
    opts, args = getopt.getopt(sys.argv[1:], "", valid)
except getopt.GetoptError as ex:
    print(error(str(ex) + help()))

configFile    = None
config        = None
command       = None
output        = None
specifiedList = {}
loadDefaults  = True
debugMode     = False
availableList = {}

if len(opts) == 0:
    print(help())
    sys.exit(0)

if opts[0][0] != "--config":
    error("First passed option must be '--config='!")

for opt, arg in opts:
    if opt == "--config":
        configFile = arg
    elif opt == "--start":
        if command != None:
            error("Command specified twise!" + help())
        command = CMD_START
    elif opt == "--shutdown":
        if command != None:
            error("Command specified twise!" + help())
        command = CMD_SHUTDOWN
    elif opt == "--restart":
        if command != None:
            error("Command specified twise!" + help())
        command = CMD_RESTART
    elif opt == "--status":
        if command != None:
            error("Command specified twise!" + help())
        command = CMD_STATUS
    elif opt == "--help":
        if command != None:
            error("Command specified twise!" + help())
        command = CMD_HELP
    elif opt == "--output":
        if output != None:
            error("'--output=' specified twise!" + help())
        output = checkOutput(arg)
        if output == None:
            error("Not support output format! Following output format are posible: 'TUI', 'XML', 'CSV', 'JSON'.")
    elif opt == "--debug":
        if debugMode == True:
            error("'--debug' was specified twise!" + help())
        debugMode = True
    elif opt == "--wmcomponents":
        if ITEM_WMCOMPONENT in specifiedList:
            error("'--wmcomponents=' specified twise!" + help())
        prepareList(ITEM_WMCOMPONENT, arg)
    elif opt == "--services":
        if ITEM_SERVICE in specifiedList:
            error("'--services=' specified twise!" + help())
        prepareList(ITEM_SERVICE, arg)
    elif opt == "--resources":
        if ITEM_RESOURCE in specifiedList:
            error("'--resources=' specified twise!" + help())
        prepareList(ITEM_RESOURCE, arg)
    elif opt == "--disks":
        if ITEM_DISK in specifiedList:
            error("'--disks=' specified twise!" + help())
        prepareList(ITEM_DISK, arg)
    elif opt == "--list":
        if len(availableList) != 0:
            error("'--list=' specified twise!" + help())
        availableList = prepareAvailableList(arg)
        command = CMD_LIST

# Checking curcial variables

if configFile == None:
    error("No configuration file set! Configuration file must be passed via '--config=' option.")
    
if command == None:
    error("No commmand specified! You must specify one of the following commands:\n\
          '--start', '--shutdown', '--status', '--restart', '--help'")

# Loading configuration

loadConfiguration()

# Load defaults if necessary
  
if output == None:
    output = OUT_TUI

if loadDefaults == True:
    specifiedList[ITEM_WMCOMPONENT] = config.listComponents_()
    specifiedList[ITEM_SERVICE] = list(servicesList)
    specifiedList[ITEM_RESOURCE] = resourcesList
    specifiedList[ITEM_DISK] = disks()

# Execute command

if command == CMD_START:
    start()
    sys.exit(0)
elif command == CMD_SHUTDOWN:
    shutdown()
    sys.exit(0)
elif command == CMD_STATUS:
    status()
    sys.exit(0)
elif command == CMD_RESTART:
    restart()
    sys.exit(0)
elif command == CMD_HELP:
    print(help(short=False))
    sys.exit(0)
elif command == CMD_LIST:
    print(available())
    sys.exit(0)