#!/usr/bin/python
# snake-machine - manage the machine database

import os,sys
import optparse
import xmlrpclib
import snake.client, snake.log, snake.machine, snake.machineinfo
from socket import gethostname
import rpmUtils.arch

try:
    import snake.machinedb
    import snake.config
except ImportError, e:
    pass

client_commands = ["add", "register", "update", "list", "whoami", "info", "remove"]

def print_table(machines, action=None, sort=True):
    '''Display a formatted table of machines.
       Accepts arguments:
         - a list of machine objects or machine dictionaries
         - optional string representing an action
         - optional boolean indicating whether sorting is needed
    '''

    # sort the list for display
    # FIXME: should this be done by the server instead?
    if sort:
        undecorated = [isinstance(t,dict) and t or t.__dict__ for t in machines]
        decorated = [(int(dict_['id']), dict_) for dict_ in undecorated]
        decorated.sort()
        machines = [dict_ for (key, dict_) in decorated]

    print " %-4s %-30s %-10s" % ("ID", "Name", "Arch")
    print "=%-4s=%-30s=%-10s" % ("="*4, "="*30, "="*10)
    if action: print action
    for d in machines:
        if isinstance(d, dict):
            m = snake.machine.Machine(dict=d)
        else:
            m = d
        print " %-4s %-30s %-50s" % (m.id, m.nickname, m.arch[0])

# FIXME - perhaps move into snake/machine.py as `print t.info()` ?
# FIXME - yes, some of this is ugly ... but it works
def print_info(t):
    '''Utility function to display formatted information about a machine'''
    for (k,v) in [["ID",            t.id],
                  ["Nickname",      t.nickname],
                  ["Arch",          " ".join(t.arch)],
                  ["Fingerprints",  " ".join(t.fingerprints)],
                  ["Kernel Args",   t.bootargs],
                  ["Kickstart",     t.ksdata],
                 ]:
        print "%-12s : %s" % (k,v)
    print ""

def parse_machine_args(args):
    '''examines the argument list and returns (nicklist,idlist), where
    nicklist is all the given URIs and idlist is all of the given arguments
    that are known machine IDs. As a special case, 'all' will expand to a list
    of all known IDs. Other arguments are ignored.'''
    nicklist = list()
    idlist = list()
    known_ids = [hasattr(t,'id') and t.id or t.get('id') for t in snake_list(dict())]

    # special case: for "all", return a list of all known machine ids.
    if 'all' in args:
        idlist = known_ids
        # remove it from the arg list
        while 'all' in args:
            args.remove('all')

    # Parse the rest of the args given.
    for a in args:
        if not a.isdigit():
            nicklist.append(a)
        elif a in known_ids and a not in idlist:
            idlist.append(a)

    return (nicklist,idlist)

def setup_option_parser(cmd=None):
    parser = None
    if cmd in ["add", "register", "update"]:
        if cmd == "update":
            parser = optparse.OptionParser(usage="snake-machine [options] %s id [%s-options] " % (cmd,cmd))
        else:
            parser = optparse.OptionParser(usage="snake-machine [options] %s [%s-options]" % (cmd,cmd))

        def check_ksdata (option, opt_str, value, parser):
            if os.path.isfile(value):
                fd = open(value, 'r')
                value = fd.read()
                fd.close()
            setattr(parser.values, option.dest, value)

        parser.add_option("-n", "--nickname",
                    action="store", dest="nickname", default=None,
                    help="Nickname of the machine (defaults to %default)")
        parser.add_option("-a", "--arch",
                    action="append", dest="arch", default=None,
                    help="Default architecture (defaults to current OS)")
        parser.add_option("-H", "--fingerprint",
                    action="append", dest="fingerprints", default=None,
                    help="Unique fingerprints associated with this system (typically mac addresses)")
        parser.add_option("-b", "--bootargs",
                    action="store", dest="bootargs", default=None,
                    help="Kernel boot arguments required to install this system")
        parser.add_option("-k", "--ksdata",
                    action="callback", dest="ksdata", type="string", callback=check_ksdata,
                    help="Kickstart commands required to install this system (string or filename allowed)")

    elif cmd is None:
        parser=optparse.OptionParser(usage="snake-machine [options] < %s >" % (", ".join(client_commands)))
        parser.add_option("-s", "--server",
                    action="store", dest="server", default=os.environ.get("SNAKE_SERVER",""),
                    help="Specify the SNAKE server")
        parser.add_option("-p", "--port",
                    action="store", dest="port", default=os.environ.get("SNAKE_PORT","2903"),
                    help="The SNAKE xmlrpc port")
        parser.add_option("-v", "--verbose",
                    action="store_true", dest="verbose", default=False,
                    help="Output extra information")

        # if we are accessing a local machinedb, offer a conffile cmdline parameter
        if sys.modules.has_key("snake.machinedb"):
            parser.add_option("-c", "--config",
                action="store", dest="conffile", default=snake.config.DEFAULT_CONFIGFILE,
                help="Override default configuration file (only for local system)")

    else:
        parser=optparse.OptionParser(usage="snake-machine [options] %s [%s-options]" % (cmd,cmd))

    return parser

def main(argv):

    global log, server, snake_list, snake_add, snake_update, snake_remove

    parser = setup_option_parser()

    # search for requested command
    rIndex = len(argv)
    for cmd in argv:
        if cmd in client_commands:
            rIndex = argv.index(cmd) + 1
            break

    (opt, args) = parser.parse_args(argv[1:rIndex])

    # Set up logging
    log = snake.log.setup_logger(opt.verbose and snake.log.DEBUG or snake.log.INFO)

    # Use XMLRPC?
    if opt.server:
        if not snake.client.check_server(opt.server, opt.port):
            '''the server is not responding'''
            log.error("The snake server '%s:%s' is not responding. \
 Please check your server URL and try again." % (opt.server, opt.port))
            sys.exit(1)
        server = snake.client.connect(opt.server, opt.port)
        supported_methods = server.system.listMethods()
        for method in ["machine.list", "machine.add", "machine.update", "machine.remove"]:
            if method not in supported_methods:
                log.error("The snake server '%s:%s' does not support the method '%s'." % (opt.server, opt.port, method))
                sys.exit(1)
        snake_list     = server.machine.list
        snake_add      = server.machine.add
        snake_update   = server.machine.update
        snake_remove   = server.machine.remove
        log.debug("Using machinedb on %s:%s as data source" % (opt.server, opt.port))
    else:
        if not sys.modules.has_key("snake.machinedb"):
            log.error("Unable to use local treedb, is snake-server installed?")
            sys.exit(1)

        # load server configuration file
        try:
            snake.config.getconfig(filename=opt.conffile)
        except Exception, e:
            log.error("Error reading configuration from '%s'" % opt.conffile)
            log.debug(e)
            sys.exit(1)

        snake_list     = snake.machinedb.dictmachines
        snake_add      = snake.machinedb.add
        snake_update   = snake.machinedb.update
        snake_remove   = snake.machinedb.remove
        log.debug("Using local machinedb as data source")

    # no command given
    if len(args) == 0:
        parser.error("You need to give some command")
        sys.exit(1)

    cmd = args[0]
    cmdopts = argv[rIndex:]

    # FIXME: try/except handling!
    if cmd in ["add", "register"]:
        cmd_parser = setup_option_parser(cmd)
        (copt,cargs) = cmd_parser.parse_args(cmdopts)

        # build out hwinfo dictionary
        hwinfo = dict()
        hwinfo["nickname"] = copt.nickname or gethostname()
        hwinfo["arch"] = copt.arch or [rpmUtils.arch.getBaseArch()]
        hwinfo["fingerprints"] = copt.fingerprints or snake.machineinfo.get_fingerprints()
        hwinfo["bootargs"] = copt.bootargs or ""
        hwinfo["ksdata"] = copt.ksdata or ""

        log.debug("Adding '%s'" % (hwinfo["nickname"],))

        # time to add the machine
        try:
            snake_add(hwinfo["nickname"], hwinfo)
        except Exception, e:
            log.error("Failed to add machine %s - %s" % (hwinfo["nickname"],e))
            if opt.verbose:
                log.exception("Unhandled exception")
        else:
            print "%sed machine %s" % (cmd.capitalize(), hwinfo["nickname"],)

    elif cmd == "update":
        cmd_parser = setup_option_parser(cmd)
        (copt,cargs) = cmd_parser.parse_args(cmdopts)

        if len(cargs) == 1:
            id = cargs.pop()
        else:
            parser.error("Required system id missing")
            sys.exit(1)

        # build out hwinfo dictionary
        hwinfo = dict()
        hwinfo["id"] = id
        # only initialize a value if user provided on the cmdline (e.g. no default values)
        for key in ["nickname", "arch", "fingerprints", "bootargs", "ksdata"]:
            if getattr(copt, key) is not None:
                hwinfo[key] = getattr(copt, key)

        log.debug("Updating machine '%s' - %s" % (id,hwinfo))

        # time to add the machine
        try:
            snake_update(id, hwinfo)
        except Exception, e:
            log.error("Failed to update machine id %s - %s" % (id,e))
            if opt.verbose:
                log.exception("Unhandled exception")
        else:
            print "Updated machine %s" % (id,)

    elif cmd in ["list", "whoami"]:
        cmd_parser = setup_option_parser(cmd)
        (copt,cargs) = cmd_parser.parse_args(cmdopts)

        filter = dict()
        if cmd == "whoami":
            filter["fingerprints"] = snake.machineinfo.get_fingerprints()
        else:
            for o in cargs:
                if '=' in o:
                    k,v = o.split('=')
                    if filter.has_key(k):
                        if isinstance(filter[k], str):
                            filter[k] = [filter[k]]
                        filter[k].append(v)
                    else:
                        filter[k]=v

        log.debug("Searching filter: %s" % filter)
        machinelist = snake_list(filter)

        if len(machinelist) == 0:
            log.error("No matching machines to list")
            sys.exit(1)

        if cmd == "whoami":
            if len(machinelist) > 1:
                log.warn("Multiple matching registrations")
            for local_m in machinelist:
                m = snake.machine.Machine(dict=local_m)
                print_info(m)
        else:
            print_table(machinelist)

    elif cmd == "info":
        (nicklist,idlist) = parse_machine_args(cmdopts)

        if len(idlist) + len(nicklist) == 0:
            log.error("No matching Machines to list")
            sys.exit(1)

        for id in idlist:
            local_m = snake_list({'id':id}).pop()
            m = snake.machine.Machine(dict=local_m)
            print_info(m)

        for u in nicklist:
            try:
                local_m = snake_list({'nickname':u}).pop()
                m = snake.machine.Machine(dict=local_m)
                print_info(m)
            except Exception, e:
                log.error("Not recognized as a valid nickname: %s" % u)

    elif cmd == "remove":
        cmd_parser = setup_option_parser(cmd)
        (copt,cargs) = cmd_parser.parse_args(cmdopts)

        if not cargs:
            log.error("You need to supply a machine nickname or id.")
            sys.exit(1)

        # build a dict of {name: __dict__} for future removal
        remove = dict()
        for d in snake_list():
            m = snake.machine.Machine(dict=d)
            if m.nickname in cargs or m.id in cargs:
                remove[m.id] = m
                # Remove nicknames from cargs ... just quiets the check to
                # ensure all cargs were processed
                if m.nickname in cargs: cargs.remove(m.nickname)

        # report if a requested argument was not found
        for o in cargs:
            if not remove.has_key(o):
                log.error("No match for argument: %s" % o)

        # were any machines found for removal?
        if not remove:
            log.error("No machines marked for removal")
        else:
            print_table(remove.values(), "Removing:")
            print ""

            # prompt for user confirmation
            if snake.client.userconfirm():
                for (id, m) in remove.items():
                    try:
                        snake_remove(id)
                    except Exception,e:
                        log.error("Failed to remove %s: %s " % (m.nickname,e))
                    else:
                        print "Removed %s" % (m.nickname,)
            else:
                log.error("Exiting on user command")

    else:
        parser.error("Unrecognized command: %s" % cmd)
        sys.exit(1)

if __name__ == '__main__':
    try:
        main(sys.argv)
    except KeyboardInterrupt:
        pass

