#!/usr/libexec/platform-python

from __future__ import print_function

import os
import re
import sys
import time
import socket
import datetime
import optparse
try:
    import ConfigParser
except ImportError:
    import configparser as ConfigParser
import xml.sax.saxutils

import gratia.common.Gratia as Gratia
import gratia.services.StorageElement as StorageElement
import gratia.services.StorageElementRecord as StorageElementRecord

def bootstrap_hadoop():
    if 'JAVA_HOME' not in os.environ:
        os.environ['JAVA_HOME'] = '/usr/java/default'

    global HDFS_CMD
    if get_version(None).startswith("Hadoop 0.20"):
        HDFS_CMD = "hadoop"
        HADOOP_DEFAULT_CONF_DIR = '/etc/hadoop-0.20/conf'
    else:
        # for hadoop 2
        HDFS_CMD = "hdfs"
        HADOOP_DEFAULT_CONF_DIR = '/etc/hadoop/conf'

    if 'HADOOP_CONF_DIR' not in os.environ:
        os.environ['HADOOP_CONF_DIR'] = HADOOP_DEFAULT_CONF_DIR

# Prevent us from sending in overly-large objects:
MAX_DATA_LEN = 50*1024

class Area(object):

    def __init__(self, name, path, prefix):
        self.name = name
        self.path = path
        self.prefix = prefix.split()

def configure_gratia(cp):
    global Gratia
    global StorageElement
    global StorageElementRecord
    try:
        probe_config = cp.get("Gratia", "ProbeConfig")
    except:
        probe_config = '/etc/gratia/hadoop-storage/ProbeConfig'
    if not os.path.exists(probe_config):
        raise Exception("ProbeConfig, %s, does not exist." % probe_config)
    Gratia.Initialize(probe_config)
    try:
        Gratia.Config.setSiteName(cp.get("Gratia", "SiteName"))
    except:
        if Gratia.Config.get_SiteName().lower().find('generic') >= 0:
            Gratia.Config.setSiteName(socket.getfqdn())
    try:
        Gratia.Config._ProbeConfiguration__CollectorHost = cp.get("Gratia",
            "Collector")
    except:
        pass
    try:
        Gratia.Config.setMeterName('hadoop-storage:%s' % socket.getfqdn())
    except:
        pass

def get_hostname(ipaddr):
    try:
        return socket.gethostbyaddr(ipaddr)[0]
    except:
        return ipaddr

divider_re = re.compile('^-+$')
attr_re = re.compile('(.*?): (.*)')
def collect_site_data(cp):
    cmd = '%s dfsadmin -report' % HDFS_CMD
    fd = os.popen(cmd)
    output = fd.read()
    if fd.close():
        raise Exception("Command failed: %s" % cmd)
    current_node = {}
    system_props = {}
    all_nodes = []
    summary_mode = 0
    for line in output.splitlines():
        if summary_mode == 0:
             m = divider_re.match(line)
             if m:
                 summary_mode += 1
                 continue
             m = attr_re.match(line)
             if not m:
                 continue
             attr, val = m.groups()
             system_props[attr] = val
        elif summary_mode == 1:
             if len(line) == 0 and current_node:
                 all_nodes.append(current_node)
                 current_node = {}
                 continue
             m = attr_re.match(line)
             if not m:
                 continue
             attr, val = m.groups()
             current_node[attr.strip()] = val.strip()

    send_system_props(system_props, cp)

    for node in all_nodes:
        send_node_props(node, cp)

time_format = '%a %b %d %H:%M:%S %Z %Y'
def send_node_props(node, cp):
    try:
        time_tuple = time.strptime(node['Last contact'], time_format)
    except:
        return
    contact = datetime.datetime(*time_tuple[:6])
    age = (datetime.datetime.now() - contact)
    if age.days * 86400 + age.seconds > 3600:
        status = 'Closed'
    elif node.get('Decommission Status').lower().find('progress') >= 0:
        status = 'Draining'
    else:
        status = 'Production'
    name = node.get('Name', 'Unknown').split(':')[0]
    name = get_hostname(name)
    se = get_se(cp)
    unique_id = '%s:Pool:%s' % (se, name)
    parent_id = '%s:SE:%s' % (se, se)
    timestamp = time.time()
    sa = StorageElement.StorageElement()
    sar = StorageElementRecord.StorageElementRecord()
    sa.UniqueID(unique_id)
    sa.Name(name)
    sa.SE(se)
    sa.SpaceType('Pool')
    sa.Implementation('Hadoop')
    sa.Version(get_version(cp))
    sa.Status(status)
    sa.ParentID(parent_id)
    sa.Timestamp(timestamp)
    sar.Timestamp(timestamp)
    sar.UniqueID(unique_id)
    sar.MeasurementType('raw')
    sar.StorageType('disk')
    sar.TotalSpace(node.get('Configured Capacity', '0').split()[0])
    sar.FreeSpace(node.get('DFS Remaining', '0').split('(')[0].strip())
    sar.UsedSpace(node.get('DFS Used', '0').split()[0])
    #print unique_id, parent_id
    Gratia.Send(sa)
    Gratia.Send(sar)

def send_system_props(props, cp):
    cmd = "hadoop fs -count -q /"
    fd = os.popen(cmd)
    output = fd.read()
    if fd.close():
        raise Exception("Command failed: %s" % cmd)
    line = output.splitlines()[0]
    info = line.split()
    if len(info) != 8:
        raise Exception("Invalid output from %s." % cmd)
    quota, remaining_quota, space_quota, remaining_space_quota, dir_count, \
        file_count, size, file_name = info
    sa = StorageElement.StorageElement()
    sar = StorageElementRecord.StorageElementRecord()
    sar2 = StorageElementRecord.StorageElementRecord()
    timestamp = time.time()
    se = get_se(cp)
    name = get_se(cp)
    unique_id = '%s:SE:%s' % (se, se)
    parent_id = '%s:SE:%s' % (se, se)
    sa.UniqueID(unique_id)
    sa.Name(name)
    sa.SE(se)
    sa.SpaceType('SE')
    sa.Implementation('Hadoop')
    sa.Version(get_version(cp))
    sa.Status('Production')
    sa.ParentID(parent_id)
    sa.Timestamp(timestamp)
    sar.UniqueID(unique_id)
    sar2.UniqueID(unique_id)
    sar.MeasurementType('raw')
    sar2.MeasurementType('logical')
    sar.StorageType('disk')
    sar2.StorageType('disk')
    sar.TotalSpace(props.get('Present Capacity', '0').split()[0])
    sar2.TotalSpace(size)
    sar.FreeSpace(props.get('DFS Remaining', '0').split()[0])
    sar2.FreeSpace(0)
    sar.UsedSpace(props.get('DFS Used', '0').split()[0])
    sar2.UsedSpace(size)
    sar.Timestamp(timestamp)
    sar2.Timestamp(timestamp)
    sar.FileCount(file_count)
    sar2.FileCount(file_count)
    if quota != 'none':
        sar.FileCountLimit(quota)
        sar2.FileCountLimit(quota)
    if space_quota != 'none':
        sar2.TotalSpace(space_quota)
        sar2.FreeSpace(remaining_space_quota)
    #print unique_id, parent_id
    Gratia.Send(sa)
    Gratia.Send(sar)
    Gratia.Send(sar2)

def _get_se(cp):
    try:
        return cp.get('Gratia', 'SiteName')
    except:
        pass
    try:
        return Gratia.Config.get_SiteName()
    except:
        pass
    try:
        return socket.getfqdn()
    except:
        return 'Unknown'

_my_se = None
def get_se(cp):
    global _my_se
    if _my_se:
        return _my_se
    _my_se = _get_se(cp)
    return _my_se

url_re = re.compile('.*?://.*?(:\d*)?/*(/.*)')
def trim_name(file_name, area):
    m = url_re.match(file_name)
    if m:
        file_name = m.groups()[1]
    for prefix in area.prefix:
        if file_name.startswith(prefix):
            return file_name[len(prefix):]
    return file_name

_my_version = None
def get_version(cp):
    global _my_version
    if _my_version:
        return _my_version
    fd = os.popen("hadoop version")
    output = fd.read()
    if fd.close():
        raise Exception("Unable to determine version with 'hadoop version'")
    _my_version = output
    return _my_version

def get_areas(cp):
    se = get_se(cp)
    space_type = 'Area'
    areas = []
    for section in cp.sections():
        if not section.lower().startswith('area'):
            continue
        try:
            name = cp.get(section, 'name')
            path = cp.get(section, 'path')
        except:
            continue
        try:
            trim = cp.get(section, 'trim')
        except:
            trim = ''
        sa = StorageElement.StorageElement()
        unique_id = '%s:%s:%s' % (se, space_type, name)
        parent_id = '%s:SE:%s' % (se, se)
        sa.UniqueID(unique_id)
        sa.Name(name)
        sa.SE(se)
        sa.SpaceType(space_type)
        sa.Implementation('Hadoop')
        sa.Version(get_version(cp))
        sa.Status('Production')
        sa.ParentID(parent_id)
        sa.Timestamp(time.time())
        #print unique_id, parent_id
        Gratia.Send(sa)
        areas.append(Area(name, path, trim))
    return areas

def collect_area_data(area, cp):
    se = get_se(cp)
    space_type_dir = 'Directory'
    space_type_quota = 'Quota'
    measurement_type = 'logical'
    storage_type = 'disk'
    timestamp = time.time()
    cmd = "hadoop fs -count -q '%s'" % area.path
    fd = os.popen(cmd)
    output = fd.read()
    if fd.close():
        raise Exception("Command failed: %s" % cmd)
    for line in output.splitlines():
        info = line.split()
        if len(info) != 8:
            continue
        quota, remaining_quota, space_quota, remaining_space_quota, dir_count, \
        file_count, size, file_name = info
        is_quota = False
        if quota != 'none' or space_quota != 'none':
            is_quota = True
        sa = StorageElement.StorageElement()
        sar = StorageElementRecord.StorageElementRecord()
        if is_quota:
            space_type = space_type_quota
        else:
            space_type = space_type_dir
        space_name = trim_name(file_name, area)
        space_unique_id = '%s:%s:%s' % (se, space_type, space_name)
        parent_id = '%s:%s:%s' % (se, 'Area', area.name)
        sa.Name(space_name)
        sa.SE(se)
        sa.UniqueID(space_unique_id)
        sa.SpaceType(space_type)
        sa.Implementation("Hadoop")
        sa.Version(get_version(cp))
        sa.Status("Production")
        sa.ParentID(parent_id)
        sar.UniqueID(space_unique_id)
        sar.MeasurementType(measurement_type)
        sar.StorageType(storage_type)
        sar.TotalSpace(size)
        sar.FreeSpace(0)
        sar.FileCountLimit(0)
        sar.FileCount(file_count)
        sar.UsedSpace(size)
        sar.Timestamp(timestamp)
        sa.Timestamp(timestamp)
        if is_quota:
            if quota != 'none':
                sar.FileCountLimit(quota)
            if space_quota != 'none':
                sar.TotalSpace(space_quota)
                sar.FreeSpace(remaining_space_quota)
        #print space_unique_id, parent_id
        Gratia.Send(sa)
        Gratia.Send(sar)

def collect_fsck_data(cp):
    collect_data("%s fsck / | grep -v '^[\.].*$' | grep -v '^$'" % HDFS_CMD,
        "FSCK Data", cp)

def collect_data(cmd, name, cp):
    time_start = time.time()
    fd = os.popen(cmd)
    output = fd.read()
    result = fd.close()
    time_end = time.time()
    if result:
        output = "WARNING: Non-zero exit code\n" + output
    else:
        result = 0
    output = xml.sax.saxutils.escape(output)
    if len(output) > MAX_DATA_LEN:
        output = "...Snipped output due to size...\n" + output[-MAX_DATA_LEN:]
    ur = Gratia.UsageRecord()
    ur.StartTime(time_start)
    ur.ResourceType("Storage-CustomInfo")
    ur.AdditionalInfo("CustomInfo", "CommandOutput")
    ur.GenericAddToList("CommandOutput", output)
    ur.Status(str(result))
    ur.WallDuration(time_end - time_start, "Was entered in seconds")
    ur.EndTime(time_end)
    ur.JobName(name, description=cmd)
    se = get_se(cp)
    ur.ProjectName('%s:SE:%s' % (se, se))
    Gratia.Send(ur)

def collect_custom_info(cp):
    for section in cp.sections():
        if section.lower().startswith("probe"):
            collect_custom_info_internal(cp, section)

custom_ctr = 0
def collect_custom_info_internal(cp, section):
    global ctr
    custom_ctr += 1
    try:
        name = cp.get(section, "Name")
    except:
        name = "Custom Probe %i" % custom_ctr
    try:
        cmd = cp.get(section, "Command")
    except:
        return
    collect_data(cmd, name, cp)

def configure():
    parser = optparse.OptionParser()
    parser.add_option("-c", "--config", dest="config", help="Config file to " \
        "use.", default="/etc/gratia/hadoop-storage/storage.cfg")
    options, args = parser.parse_args()
    config = options.config
    if not os.path.exists(config):
        raise Exception("Config file %s does not exist." % config)
    try:
        open(config, 'r').read()
    except:
        raise Exception("Config file %s exists, but an error occurred when " \
            "trying to read it." % config)
    cp = ConfigParser.ConfigParser()
    cp.read(config)
    return cp

def main():
    bootstrap_hadoop()
    try:
        cp = configure()
    except Exception as e:
        print(e)
        sys.exit(1)
    configure_gratia(cp)
    collect_site_data(cp)
    interesting_areas = get_areas(cp)
    for area in interesting_areas:
        collect_area_data(area, cp)
    collect_custom_info(cp)
    collect_fsck_data(cp)

if __name__ == '__main__':
    main()

