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

import os,sys
import optparse
import xmlrpclib
import snake.client, snake.log

try:
    import snake.ksdb
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...)")
    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="URL of SNAKE server")
        parser.add_option("-v", "--verbose",
                    action="store_true", dest="verbose", default=False,
                    help="Output extra information")
    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):
            '''the server is not responding'''
            log.error("The snake server '%s' is not responding. \
 Please check your server URL and try again." % opt.server)
            sys.exit(1)
        server = snake.client.connect(opt.server)
        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 as data source" % opt.server)
    else:
        if not sys.modules.has_key("snake.ksdb"):
            log.error("Unable to use local treedb, is snake-server installed?")
            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
            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'''
                    if not copt.name:
                        raise Exception, "Must provide a --name for kickstarts added through stdin"
                    fd = sys.stdin
                else:
                    fd = open(o,'r')
                snake_ks_add(fd,name,copt.desc)
                fd.close()
                # Need some sort of snake_ks_info(name) cmd
                (ksname,ksdesc) = snake_ks_list({"name":name})[0]
            except Exception, e:
                log.error("Failed to add %s: %s " % (o,e))
                continue
            else:
                print "Added %s: %s" % (ksname,ksdesc)

    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)

        # prepare keyword arguments
        arg = dict()
        if copt.version: arg["version"] = copt.version

        for o in cargs:
            try:
                print snake_ks_generate(o,arg)
            except IOError:
                log.error("Kickstart '%s' not found" % o)
                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

