diamond-setup

#!/usr/bin/env python
##########################################################################

from __future__ import print_function
import os
import sys
import optparse
import traceback

from configobj import ConfigObj

try:
    from setproctitle import setproctitle
except ImportError:
    setproctitle = None

for path in [
    os.path.abspath(os.path.join(os.path.dirname(__file__), '..', 'src')),
    os.path.join('opt', 'diamond', 'lib'),
]:
    if os.path.exists(os.path.join(path, 'diamond', '__init__.py')):
        sys.path.append(path)
        break

from diamond.collector import Collector
from diamond.collector import str_to_bool


def getIncludePaths(path):
    for f in os.listdir(path):
        cPath = os.path.abspath(os.path.join(path, f))

        if os.path.isfile(cPath) and len(f) > 3 and f[-3:] == '.py':
            sys.path.append(os.path.dirname(cPath))

    for f in os.listdir(path):
        cPath = os.path.abspath(os.path.join(path, f))
        if os.path.isdir(cPath):
            getIncludePaths(cPath)

collectors = {}


def getCollectors(path):
    for f in os.listdir(path):
        cPath = os.path.abspath(os.path.join(path, f))

        if (os.path.isfile(cPath)
                and len(f) > 3
                and f[-3:] == '.py'
                and f[0:4] != 'test'):
            modname = f[:-3]
            try:
                # Import the module
                module = __import__(modname, globals(), locals(), ['*'])

                # Find the name
                for attr in dir(module):
                    cls = getattr(module, attr)
                    try:
                        if (issubclass(cls, Collector)
                                and cls.__name__ not in collectors):
                            collectors[cls.__name__] = module
                            break
                    except TypeError:
                        continue
                # print("Imported module: %s %s" % (modname, cls.__name__))
            except Exception:
                print("Failed to import module: %s. %s" % (
                    modname, traceback.format_exc()))
                collectors[modname] = False
                continue

    for f in os.listdir(path):
        cPath = os.path.abspath(os.path.join(path, f))
        if os.path.isdir(cPath):
            getCollectors(cPath)


def typeToString(key):
    if isinstance(obj.config[key], str):
        user_val = obj.config[key]
    elif isinstance(obj.config[key], bool):
        user_val = str(obj.config[key])
    elif isinstance(obj.config[key], int):
        user_val = str(obj.config[key])
    elif isinstance(obj.config[key], list):
        user_val = str(obj.config[key])[1:-1]
    else:
        raise NotImplementedError("Unknown type!")

    return user_val


def stringToType(key, val):
    if type(obj.config[key]) is type(val):
        config_file[key] = val
    elif isinstance(obj.config[key], str):
        if val.lower() == 'false':
            config_file[key] = False
        elif val.lower() == 'true':
            config_file[key] = True
        else:
            config_file[key] = val
    elif isinstance(obj.config[key], bool):
        if isinstance(val, str):
            config_file[key] = str_to_bool(val)
        else:
            config_file[key] = bool(val)
    elif isinstance(obj.config[key], int):
        config_file[key] = int(val)
    elif isinstance(obj.config[key], list):
        entry = ConfigObj([key + ' = ' + val])
        config_file[key] = entry[key]
    else:
        raise NotImplementedError("Unknown type!")


def boolCheck(val):
    if isinstance(val, str):
        return str_to_bool(val)
    elif isinstance(val, bool):
        return val
    else:
        raise NotImplementedError("Unknown type!")


def configureKey(key):
    if not config_keys[key]:
        return

    try:
        user_val = typeToString(key)
    except NotImplementedError:
        return

    print("\n")
    if key in default_conf_help:
        print(default_conf_help[key])
    val = input(key + ' [' + user_val + ']: ')

    # Empty user input? Default to current value
    if len(val) == 0:
        val = obj.config[key]

    try:
        stringToType(key, val)
    except NotImplementedError:
        return

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

if __name__ == "__main__":

    if setproctitle:
        setproctitle('diamond-setup')

    # Initialize Options
    parser = optparse.OptionParser()
    parser.add_option("-c", "--configfile",
                      dest="configfile",
                      default="/etc/diamond/diamond.conf",
                      help="Path to the config file")
    parser.add_option("-C", "--collector",
                      dest="collector",
                      default=None,
                      help="Configure a single collector")
    parser.add_option("-p", "--print",
                      action="store_true",
                      dest="dump",
                      default=False,
                      help="Just print the defaults")

    # Parse Command Line Args
    (options, args) = parser.parse_args()

    # Initialize Config
    if os.path.exists(options.configfile):
        config = ConfigObj(os.path.abspath(options.configfile))
    else:
        print("ERROR: Config file: %s does not exist." % (
            options.configfile), file=sys.stderr)
        print("Please run python config.py -c /path/to/diamond.conf",
              file=sys.stderr)
        parser.print_help(sys.stderr)
        sys.exit(1)

    if not options.dump:
        print('')
        print('I will be over writing files in')
        print(config['server']['collectors_config_path'])
        print('Please type yes to continue')

        val = input('Are you sure? ')
        if val != 'yes':
            sys.exit()

    getIncludePaths(config['server']['collectors_path'])
    getCollectors(config['server']['collectors_path'])
    tests = []

    foundcollector = False
    for collector in collectors:
        if options.collector and collector != options.collector:
            continue

        # Skip configuring the basic collector object
        if collector == "Collector":
            continue

        foundcollector = True

        config_keys = {}
        config_file = ConfigObj()
        config_file.filename = (config['server']['collectors_config_path']
                                + "/" + collector + ".conf")

        # Find the class and load it from the collector module
        try:

            # We can for the name above, so we dont have to scan here anymore
            if not hasattr(collectors[collector], collector):
                continue

            cls = getattr(collectors[collector], collector)
            obj = cls(config=config, handlers={})

            if options.dump:
                print(collector + " " + str(obj.config))
                continue

            default_conf = obj.get_default_config()
            default_conf_help = obj.get_default_config_help()

            for key in obj.get_default_config():
                config_keys[key] = True

            # Manage Keys
            config_keys['enabled'] = True
            config_keys['path'] = False
            config_keys['path_prefix'] = False
            config_keys['instance_prefix'] = False
            config_keys['interval'] = False

            print("*" * 60)
            print("\n\t\tNow configuring " + collector)
            print(collectors[collector].__doc__)

            print("(%s)" % collector)
            configureKey('enabled')

            if boolCheck(config_file['enabled']):
                for key in config_keys:
                    if key == 'enabled':
                        continue
                    configureKey(key)

            config_file.write()

        except IOError as err:
            print("I/O error({}): {}".format(err.errno, err.strerror))
        except KeyboardInterrupt:
            print()
            sys.exit()
        except:
            continue
    if not foundcollector:
        print("Collector not found.")