#!/usr/bin/env python

   ##########################################################
  #                                                         #
 # Written 2/22/2012 by Jeff Schroeder released GPLv2 only# #
##########################################################  #
#                                                        #  #
# nhs - Connect to cobbler and get a list of hosts  for  #  #
#       running commands on via onall, etc. The name is  #  #
#       Nifty Hostname Script and it is  authoritative.  #  #
#                                                        # #
##########################################################

import re
import os
import sys
import stat
import time
import signal
import socket
import optparse
import textwrap
import xmlrpclib

# Properly handle SIGPIPE ie: nhs | head ...
signal.signal(signal.SIGPIPE, signal.SIG_DFL)

CACHE_FILE = os.path.join("/", "tmp", ".nhscache")
# Cache time in minutes
CACHE_TIME = 5

def die(msg, parser):
    sys.stderr.write("ERROR: %s\n\n" % msg)
    parser.print_help()
    parser.exit(1)

EPILOG = textwrap.dedent("""If --use-cache and --systems are
specified, the cache file will be written out with only systems matching
the --systems glob. This will cause future lookups with --use-cache
to not display the full results available in cobbler.""")

def main():
    """
    Parse arguments and do the needfuls
    """
    parser = optparse.OptionParser()
    parser.description = "Nifty Hosts Script! - Get hosts from cobbler"
    parser.epilog = EPILOG

    # Default to the "cobbler" hostname
    parser.add_option("-m", "--master", action="store",
        help="Cobbler server (default: %default)",
        default="cobbler",
        dest="cobbler_server",
    )
    parser.add_option("-r", "--regex", action="store",
        help="System records to list in regular expression format (default: %default)",
        default=None,
        dest="regex",
    )
    parser.add_option("-s", "--systems", action="store",
        help="System records to list in glob format (default: '%default')",
        dest="systems",
        default="*",
    )
    parser.add_option("-e", "--exclude", action="store",
        help="Regex of hosts to exclude",
        dest="exclude",
        default=None,
    )
    parser.add_option("-q", "--quiet", action="store_true",
        help="Don't print an error if no systems were found",
        dest="quiet",
        default=False,
    )

    # All of the cache settings are related so group them together
    group = optparse.OptionGroup(parser, "Cache Related Options")

    group.add_option("-U", "--use-cache", action="store_true",
        help="Enable the cache to speed up things (default: %default)",
        dest="usecache",
        default=False,
    )
    group.add_option("-C", "--cache-file", action="store",
        help="File to stored cached systems (default: %default)",
        dest="cachefile",
        default=CACHE_FILE,
    )
    group.add_option("-T", "--cache-time", action="store",
        help="Time the cache is considered 'fresh' in minutes (default: %default)",
        dest="cachetime",
        default=CACHE_TIME,
    )
    parser.add_option_group(group)

    opts,args = parser.parse_args()

    return {
        'quiet'         : opts.quiet,
        'regex'         : opts.regex,
        'exclude'       : opts.exclude,
        'systems'       : opts.systems,
        'usecache'      : opts.usecache,
        'cachefile'     : opts.cachefile,
        'cachetime'     : opts.cachetime,
        'cobbler_server': opts.cobbler_server,
    }

def filter_systems(regex, data, remove=False):
    """
    Positive or negative filtering of a list of hosts using regex

    If remove=True, any hosts *not* matching regex are  returned.
    """
    if isinstance(regex, basestring):
        regex = re.compile(regex)

    if not isinstance(data, (list, tuple)):
        raise ValueError("Data needs to be a list or tuple")

    filtered = filter(regex.match, data)
    if remove:
        ret = list(set(data).difference(filtered))
    else:
        ret = filtered
    return sorted(ret)

def check_cache_age(cachefile, cachetime=5):
    """
    Check the age of a cached file

    cachefile - filename to store cached data
    cachetime - time in minutes to keep cache
    """
    if not os.path.exists(cachefile):
        return False
    mtime = os.stat(cachefile).st_mtime
    age   = int(time.time() - mtime) / 60
    if age > cachetime:
        return False
    return True

def update_cache_file(cachefile, data):
    """
    Update data in the cachefile
    """
    try:
        FH = open(cachefile, "w")
        # Join the list of systems with newlines and
        # make sure to not forget the  trailing  \n.
        FH.writelines('\n'.join(data) + "\n")
        FH.close()
    except (IOError, OSError):
        return False
    return True

def get_hosts_from_cache(cachefile):
    ret = []
    try:
        ret = [i.rstrip() for i in open(cachefile, "r")]
        if not ret:
            raise OSError("No systems")
    except (IOError, OSError):
        pass
    return ret

def get_hosts_from_cobbler(cobbler_server, system_glob):
    server = xmlrpclib.ServerProxy("http://%s/cobbler_api" % cobbler_server)

    try:
        systems = server.find_system({'hostname': system_glob})
    except KeyboardInterrupt:
        raise SystemExit("\nExiting on Ctrl-c")
    except Exception, exc:
        raise SystemExit("ERROR: Problem connecting to cobbler on %s" % cobbler_server)
    return systems


def get_hosts(system_glob, cobbler_server, use_cache=False, cachefile=CACHE_FILE, cachetime=CACHE_TIME):
    """
    Get a list of systems. If the cache is enabled, use it, and populate
    it, otherwise pull all of the system records from cobbler.
    """
    if use_cache:
        if not check_cache_age(cachefile, cachetime):
            systems = get_hosts_from_cobbler(cobbler_server, system_glob)
            update_cache_file(cachefile, systems)
            return systems

        # Read the data from the cachefile
        systems = get_hosts_from_cache(cachefile)
        if systems:
            return systems
    systems = get_hosts_from_cobbler(cobbler_server, system_glob)
    return systems

if __name__ == "__main__":
    data = main()

    # Easier to construct the arguments for get_hosts() this way
    kwargs = {
        "system_glob": data["systems"],
        "cobbler_server": data["cobbler_server"]
    }

    if data.get("usecache"):
        kwargs.update({
            "use_cache": data["usecache"],
            "cachefile": data["cachefile"],
            "cachetime": data["cachetime"],
        })

    # Get all of the systems from either cobbler or the cache  before
    # filtering them out for the various options to trim results down
    systems = get_hosts(**kwargs)

    # Filter hosts based on the arguments
    if data.get("exclude"):
        systems = filter_systems(data["exclude"], systems, remove=True)

    if data.get("regex"):
        try:
            regex = re.compile(data.get("regex"))
        except re.error:
            raise SystemExit("ERROR: '%s' is an invalid regex" % data.get("regex"))

        systems = filter_systems(regex, systems)

    # Exit gracefully if not systems were found in cobbler
    if not systems:
        if data.get("quiet"):
            sys.exit(0)
        raise SystemExit("ERROR: No systems found in %s\n" % data["cobbler_server"])

    print "\n".join(systems)