#!/usr/bin/python
# snake-ks - manage the kickstart database

import os,sys
import optparse
import xmlrpclib
import snake.client, snake.log
from snake.machineinfo import get_fingerprints

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

ks_commands = ["add", "list", "generate", "remove", "rename", "describe"]

def print_ks_table(kslist, action=None):
    '''Display a formatted table of kickstart templates.
       Accepts arguments:
         - a list of (name,description) pairs
         - optional string representing an action
    '''

    # sort the list for display
    # FIXME: should this be done by the server instead?
    kslist.sort()

    print " %-25s %-50s" % ("Name", "Description")
    print "=%-25s=%-50s" % ("="*25, "="*40)
    if action: print action
    for (name,desc) in kslist:
        print " %-25s %-50s" % (name,desc)

def setup_option_parser(cmd=None):
    parser = None
    if cmd == "add":
        parser=optparse.OptionParser(usage="snake-ks [options] %s [%s-options] filename" % (cmd,cmd))
        parser.add_option("-n", "--name",
                    action="store", dest="name", default="",
                    help="Name of the kickstart (defaults to filename)")
        parser.add_option("-d", "--desc",
                    action="store", dest="desc", default="",
                    help="A brief description of the kickstart")
    elif cmd == "generate":
        parser=optparse.OptionParser(usage="snake-ks [options] %s [%s-options] kickstart" % (cmd,cmd))
        parser.add_option("-v", "--version",
                    action="store", dest="version", default="",
                    help="Kickstart version to generate (e.g. F8, F7, RHEL5, etc...)")
        parser.add_option("-H", "--fingerprint",
                    action="append", dest="fingerprints", default=None,
                    help="Machine fingerprints used during kickstart generation (defaults to system mac addresses)")
        parser.add_option("--ksmeta",
                    action="store", dest="ksmeta", type="string", default="",
                    help="Define kickstart template variables (optional)")
    elif cmd == "rename":
        parser=optparse.OptionParser(usage="snake-ks [options] %s old-name new-name" % (cmd,))
    elif cmd == "describe":
        parser=optparse.OptionParser(usage="snake-ks [options] %s name description" % (cmd,))
    elif cmd is None:
        parser=optparse.OptionParser(usage="snake-ks [options] < %s >" % (", ".join(ks_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 ksdb, offer a conffile cmdline parameter
        if sys.modules.has_key("snake.ksdb"):
            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-ks [options] %s [%s-options]" % (cmd,cmd))

    return parser

def __xmlrpc_prep_and_add__(fd,*args):
    data = fd.read()
    server.kickstart.add(data,*args)

def main(argv):

    global log, server, snake_tree_list, snake_uri_add, snake_uri_remove, snake_tree_remove

    parser = setup_option_parser()

    # search for requested command
    rIndex = len(argv)
    for cmd in argv:
        if cmd in ks_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)
        snake_ks_list     = server.kickstart.list
        snake_ks_add      = __xmlrpc_prep_and_add__  # some work is needed before hand
        snake_ks_remove   = server.kickstart.remove
        snake_ks_generate = server.kickstart.generate
        snake_ks_rename   = server.kickstart.rename
        snake_ks_describe = server.kickstart.describe
        log.debug("Using ksdb on %s:%s as data source" % (opt.server, opt.port))
    else:
        if not sys.modules.has_key("snake.ksdb"):
            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_ks_list     = snake.ksdb.listkickstarts
        snake_ks_add      = snake.ksdb.add
        snake_ks_remove   = snake.ksdb.remove
        snake_ks_generate = snake.ksdb.generate
        snake_ks_rename   = snake.ksdb.rename
        snake_ks_describe = snake.ksdb.describe
        log.debug("Using local ksdb 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 == 'add':
        cmd_parser = setup_option_parser(cmd)
        (copt,cargs) = cmd_parser.parse_args(cmdopts)

        # no command given
        if len(cargs) == 0:
            if not sys.stdin.isatty():
                '''perhaps data waiting for us on stdin'''
                cargs="-"
            else:
                cmd_parser.error("You need to give one or more kickstart files to add.")
                sys.exit(1)

        # process kickstart arguments
        for o in cargs:

            # what should we call the kickstart?
            if copt.name:
                #FIXME - need to make this work for multiple kickstart files
                name = copt.name
            elif not copt.name and o == "-":
                raise Exception, "Must provide a --name when adding kickstarts via stdin"
            else:
                filename = os.path.basename(o)
                (name, ext) = os.path.splitext(filename)

            log.debug("Adding '%s' as '%s' with description '%s'" % (o,name,copt.desc))

            # time to add the kickstart
            try:
                if o == "-":
                    '''read kickstart from stdin'''
                    fd = sys.stdin
                else:
                    fd = open(o,'r')
                snake_ks_add(fd,name,copt.desc)
                fd.close()
            except Exception, e:
                log.error("Failed to add %s: %s " % (o,e))
                sys.exit(1)
            else:
                print "Added %s" % (o,)

    elif cmd == 'list':
        cmd_parser = setup_option_parser(cmd)
        (copt,cargs) = cmd_parser.parse_args(cmdopts)

        filter = dict()
        for o in cargs:
            if '=' in o:
                k,v = o.split('=')
                filter[k]=v

        kslist = snake_ks_list(filter)

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

        print_ks_table(kslist)

    elif cmd == 'generate':
        cmd_parser = setup_option_parser(cmd)
        (copt,cargs) = cmd_parser.parse_args(cmdopts)

        if not cargs:
            log.error("You need to supply the name of a kickstart to generate.")
            sys.exit(1)

        ksname = cargs[0]

        # prepare keyword arguments
        arg = dict()
        if copt.version: arg["version"] = copt.version
        if len(cargs) > 1: arg["commands"] = cargs[1:]
        if copt.ksmeta:
            arg["ksmeta"] = dict()
            for kv in copt.ksmeta.split(" "):
                (k,v) = kv.split("=",1)
                arg["ksmeta"][k] = v

        # provide machine fingerprints
        arg["fingerprints"] = copt.fingerprints or get_fingerprints()

        try:
            print snake_ks_generate(ksname, arg)
        except IOError:
            log.error("Kickstart '%s' not found" % ksname)
            sys.exit(1)
        except Exception, e:
            log.exception("Unhandled exception occured generating '%s' kickstart" % ksname)
            sys.exit(1)

    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 one or more kickstarts.")
            sys.exit(1)

        # build a dict of name: desc templates to remove
        remove = dict()
        for (name,desc) in snake_ks_list():
            if name in cargs:
                remove[name] = desc

        # 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 kickstart templates found for removal?
        if not remove:
            log.error("No templates marked for removal")
        else:
            print_ks_table(remove.items(), "Removing:")
            print ""

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

    elif cmd == "rename":
        # process command-line arguments
        cmd_parser = setup_option_parser(cmd)
        (copt,cargs) = cmd_parser.parse_args(cmdopts)

        if len(cargs) != 2:
            cmd_parser.error("Incorrect arguments supplied")
            sys.exit(1)

        (oldname,newname) = cargs
        filter = {"name": oldname}
        kslist = snake_ks_list(filter)

        if len(kslist) == 0:
            log.error("No matching kickstart list")
            sys.exit(1)

        # attempt rename
        # FIXME - security?  what if newname = /etc/passwd?
        try:
            snake_ks_rename(oldname,newname)
        except Exception,e:
            log.error("Failed to rename %s: %s " % (oldname,e))
        else:
            print "Renamed %s to %s" % (oldname,newname)

    elif cmd == "describe":
        # process command-line arguments
        cmd_parser = setup_option_parser(cmd)
        (copt,cargs) = cmd_parser.parse_args(cmdopts)

        if len(cargs) != 2:
            cmd_parser.error("Incorrect arguments supplied")
            sys.exit(1)

        (name,desc) = cargs
        filter = {"name": name}
        kslist = snake_ks_list(filter)

        if len(kslist) == 0:
            log.error("No matching kickstart found")
            sys.exit(1)

        # attempt update
        try:
            snake_ks_describe(name,desc)
        except Exception,e:
            log.error("Failed to describe %s: %s " % (name,e))
        else:
            print "Updated description of %s: %s" % (name,desc)

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

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

