#!/usr/bin/python

import os
import re
import sys
import time
import socket
import datetime
import optparse
import 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, 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()

