#!/usr/bin/env python
# User configurable stuff
# Change the second value of get method or define environment variables:
# [jboss7_*]
#   env.jboss_username a
#   env.jboss_password b
#   env.jboss_realm ManagementRealm
#   env.jboss_manager "http://127.0.0.1:9990/management"
#   env.datasource xa-data-source
#   env.url_timeout 5
#

import urllib2 as ul
import sys
import os
import re
import json

# Auth stuff
user = os.environ.get('jboss_username', 'a')
password = os.environ.get('jboss_password', 'b')
realm = os.environ.get('jboss_realm', 'ManagementRealm')

# Manager location & timeout
manager = os.environ.get('jboss_manager', 'http://127.0.0.1:9990/management')
url_timeout = int(os.environ.get('url_timeout', 5))

# Datasource
datasource = os.environ.get('datasource', 'xa-data-source')

# JBoss data types -> Python data types
null = undefined = None
true = True
false = False

# Timeout for urlopen
required_version = (2, 6)
current_version = sys.version_info[:2]

options = {
    "ds": "/subsystem/datasources?include-runtime=true&recursive=true",
    "web": "/subsystem/web?include-runtime=true&recursive=true",
    "mem": "/core-service/platform-mbean/type/memory?include-runtime=true&recursive=true"
}

ds_stats = [
    'CreatedCount',
    'TimedOut',
    'ActiveCount',
]

web_stats = [
    'errorCount',
    'requestCount',
]

# keys from /core-service/platform-mbean/type/memory?include-runtime=true&recursive=true
memory_stats = [
    'init',
    'used',
    'committed',
    'max',
]

def build_opener(uri, user, password, realm):
    pm = ul.HTTPPasswordMgrWithDefaultRealm()
    pm.add_password(realm=realm, uri=uri, user=user, passwd=password)
    ah = ul.HTTPDigestAuthHandler(pm)
    return ul.build_opener(ah)

def manager_data(url):
    ul.install_opener(build_opener(url, user, password, realm))
    try:
        if current_version >= required_version:
            data = json.loads(ul.urlopen(url, timeout=url_timeout).read())
        else:
            import socket
            socket.setdefaulttimeout(url_timeout)
            data = json.loads(ul.urlopen(url).read())
    except:
        data = 0
    return data

def string_fix(s):
    return re.sub(r"[^a-zA-Z0-9_]", "_", s)

def failed_config():
    print "failed.label FAILED %s" % sys.argv[0]
    print "failed.vlabel ffffailed!!!1"
    print "failed.info Fake field indicating we were unable to connect."
    print "failed.critical 1:"

def failed_fetch():
    print "failed.value 0"

def configure(op_type):
    data = manager_data(manager + options[op_type])
    if op_type == "ds":
        print "multigraph jboss7_ds"
        print "graph_category jboss"
        print "graph_args --base 1000 -l 0"
        print "graph_scale no"
        print "graph_title JBoss7 datasource root graph. Sum of all."
        if data:
            for stat in ds_stats:
                print "%s.label %s" % (stat, stat)
                print "%s.info Connection count" % (stat)
                print "%s.min 0" % (stat)
                print "%s.type DERIVE" % (stat)
            print
            for ds in data[datasource]:
                print "multigraph jboss7_ds.%s" % (string_fix(ds))
                print "graph_title %s" % string_fix(ds)
                print "graph_args --base 1000 -l 0"
                print "graph_scale no"
                print "graph_category jboss"
                for stat in ds_stats:
                    print "%s.label %s" % (stat, stat)
                    print "%s.info Connection count" % (stat)
                    print "%s.min 0" % (stat)
                    print "%s.type DERIVE" % (stat)
                print
        else:
            failed_config()

    if op_type == "web":
        print "multigraph jboss7_web"
        print "graph_category jboss"
        print "graph_args --base 1000 -l 0"
        print "graph_scale no"
        print "graph_title JBoss7 web root graph. Sum of all."
        if data:
            for stat in web_stats:
                print "%s.label %s" % (stat, stat)
                print "%s.info Connection count" % stat
                print "%s.min 0" % stat
                print "%s.type DERIVE" % stat
            print
            for wc in data['connector']:
                print "multigraph jboss7_web.%s" % wc
                print "graph_title %s" % wc
                print "graph_args --base 1000 -l 0"
                print "graph_scale no"
                print "graph_category jboss"
                for stat in web_stats:
                    print "%s.label %s" % (stat, stat)
                    print "%s.info Connection count" % stat
                    print "%s.min 0" % stat
                    print "%s.type DERIVE" % stat
                print
        else:
            failed_config()

    if op_type == "mem":
        print "multigraph jboss7_mem"
        print "graph_category jboss"
        print "graph_args --base 1024 -l 0"
        print "graph_vlabel memory"
        print "graph_title JBoss7 memory root graph. Sum of all."
        if data:
            for stat in memory_stats:
                print "%s.label %s" % (stat, stat)
                print "%s.info Memory status" % stat
                print "%s.min 0" % stat
            print
            # element in data with 'memory' word match
            for el in [x for x in data if re.search('memory', x)]:
                print "multigraph jboss7_mem.%s" % el.replace('-', '_')
                print "graph_title %s" % el
                print "graph_args --base 1024 -l 0"
                print "graph_vlabel memory"
                print "graph_category jboss"
                for stat in memory_stats:
                    print "%s.label %s" % (stat, stat)
                    print "%s.info Memory status" % stat
                    print "%s.min 0" % stat
                print
        else:
            failed_config()

def get_data(op_type):
    data = manager_data(manager + options[op_type])
    if op_type == "ds":
        print "multigraph jboss7_ds"
        if data:
            for stat in ds_stats:
                print "%s.value %s" % (stat, sum([int(data[datasource][ds]['statistics']['pool'][stat]) for ds in data[datasource]]))
            print

            for ds in data[datasource]:
                print "multigraph jboss7_ds.%s" % string_fix(ds)
                for stat in ds_stats:
                    print "%s.value %s" % (stat, data[datasource][ds]['statistics']['pool'][stat])
                print
        else:
            failed_fetch()

    if op_type == "web":
        print "multigraph jboss7_web"
        if data:
            for stat in web_stats:
                print "%s.value %s" % (stat, sum([int(data['connector'][wc][stat]) for wc in data['connector']]))
            print

            for wc in data['connector']:
                print "multigraph jboss7_web.%s" % wc
                for stat in web_stats:
                    print "%s.value %s" % (stat, data['connector'][wc][stat])
                print
        else:
            failed_fetch()

    if op_type == "mem":
        print "multigraph jboss7_mem"
        if data:
            for stat in memory_stats:
                print "%s.value %s" % (stat, sum([data[key][stat] for key in [el for el in data if re.search('memory', el)]]))
            print

            for el in [x for x in data if re.search('memory', x)]:
                print "multigraph jboss7_mem.%s" % el.replace('-', '_')
                for k, v in data[el].items():
                    print "%s.value %s" % (k, v)
                print
        else:
            failed_fetch()

if __name__ == "__main__":
    op_type = sys.argv[0].replace('.py', '').split('_')[-1]

    # silently exit if unknown option
    if op_type not in options:
        sys.exit(0)

    if len(sys.argv) >= 2 and sys.argv[1] == "config":
        configure(op_type)
    else:
        get_data(op_type)