#!/usr/bin/env python

'''Bcfg2 Client'''
__revision__ = '$Revision$'

from getopt import getopt, GetoptError
from os import popen, chmod, unlink, _exit
from signal import signal, SIGINT
from sys import argv
from tempfile import mktemp
from ConfigParser import ConfigParser, NoSectionError, NoOptionError
from xmlrpclib import ServerProxy, Fault
from lxml.etree import Element, XML, tostring, XMLSyntaxError
from time import sleep, time
from sys import exc_info
from traceback import extract_tb
import socket

def cb_sigint_handler(signum, frame):
    '''Exit upon CTRL-C'''
    _exit(1)

def if_then(cond, value_if, value_else):
    ''' Replacement for ternary operator '''
    if cond == True:
        return value_if
    else:
        return value_else

class SafeProxy:
    '''Wrapper for proxy'''
    def __init__(self, setup, client): 
        self.retryCount = 0
        self.client = client
        self.setup = setup
        try:
            self.proxy = ServerProxy(self.setup["server"])
        except IOError, io_error:
            self.client.fatal_error("Invalid server URL %s: %s" %
                                    (self.setup["server"], io_error))
        except:
            self.client.critical_error("initialising XML-RPC")

    def run_method(self, operation_desc, method_name, method_args):
        ''' Perform an XMLRPC invocation against the server'''
        method = getattr(self.proxy, method_name)
        instance_retries = 0
        for i in xrange(int(self.setup["retries"])):
            try:
                self.client.cond_print("debug", "Attempting %s (%d of %d)" %
                                      (operation_desc,(i+1),
                                       int(self.setup["retries"])))
                ret = apply(method, (self.setup['user'],
                                     self.setup['password']) + method_args)
                if instance_retries > 0:
                    self.client.warning_error(
                        "during %s:\nRequired %d attempts to contact server (%s)"
                        % (operation_desc, instance_retries,
                           self.setup["server"]))
                self.client.cond_print("debug", "%s completed successfully" %
                                      (operation_desc))
                return ret
            except Fault, fault:
                self.client.fatal_error("%s encountered a server error:\n%s" %
                                        (operation_desc, fault))
            except socket.error:
                instance_retries += 1
                self.retryCount += 1
                sleep(1.0)                
            except:
                self.client.critical_error(operation_desc)

        self.client.fatal_error("%s failed:\nCould not connect to server (%s)" %
                                (operation_desc, self.setup["server"]))
        

class Client:
    ''' The main bcfg2 client class '''
    def __init__(self, args):
        self.toolset = None
        self.config = None
        self.options = {
            'verbose': 'v',
            'quick': 'q',
            'debug': 'd',
            'dryrun': 'n',
            'build': 'B',
            'paranoid': 'P',
            'bundle': 'b',
            'file': 'f',
            'cache': 'c',
            'profile': 'p',
            'image': 'i',
            'remove': 'r',
            'help': 'h',
            'setup': 's',
            'server': 'S',
            'user': 'u',
            'password': 'x',
            'retries': 'R'
            }
        self.argOptions = {
            'v': 'verbose',
            'q': 'quick',
            'd': 'debug',
            'n': 'dryrun',
            'B': 'build',
            'P': 'paranoid',
            'b': 'bundle',
            'f': 'file',
            'c': 'cache',
            'p': 'profile',
            'i': 'image',
            'r': 'remove',
            'h': 'help',
            's': 'setup',
            'S': 'server',
            'u': 'user',
            'x': 'password',
            'R': 'retries'
            }
        self.descriptions = {
            'verbose': "enable verbose output",
            'quick': "disable some checksum verification",
            'debug': "enable debugging output",
            'dryrun': "do not actually change the system",
            'build': "disable service control (implies -q)",
            'paranoid': "make automatic backups of config files",
            'bundle': "only configure the given bundle",
            'file': "configure from a file rather than querying the server",
            'cache': "store the configuration in a file",
            'image': "assert the given image for the host",
            'profile': "assert the given profile for the host",
            'remove': "force removal of additional configuration items",
            'help': "print this help message",
            'setup': "use given setup file (default /etc/bcfg2.conf)",
            'server': 'the server hostname to connect to',
            'user': 'the user to provide for authentication',
            'password': 'the password to use',
            'retries': 'the number of times to retry network communication'            
            }
        self.argumentDescriptions = {
            'bundle': "<bundle name>",
            'file': "<cache file>",
            'cache': "<cache file>",
            'profile': "<profile name>",
            'image': "<image name>",
            'remove': "(pkgs | svcs | all)",
            'setup': "<setup file>",
            'server': '<hostname>   ',
            'user': '<user name>  ',
            'password': '<password>  ',
            'retries': '<number of retries>'            
            }

        self.setup = {} 
        self.get_setup(args)

        self.cond_print_setup('debug')

    def cond_print_setup(self, state):
        ''' Display the clients current setup information '''
        for (key, value) in self.setup.iteritems():
            if self.setup[key]:            
                self.cond_print(state, "%s => %s" % (key, value))


    def load_toolset(self, toolset_name):
        '''Import client toolset modules'''
        
        toolset_packages = {
            'debian': "Bcfg2.Client.Debian",
            'rh': "Bcfg2.Client.Redhat",
            'solaris': "Bcfg2.Client.Solaris"
            }

        if toolset_packages.has_key(toolset_name):
            toolset_class = toolset_packages[toolset_name]
        else:
            toolset_class = toolset_name

        try:
            mod = __import__(toolset_class, globals(), locals(), ['*'])
        except:
            self.fatal_error("got unsupported toolset %s from server."
                             % (toolset_name))
                
        try:
            self.toolset = mod.ToolsetImpl(self.config, self.setup)
            
            self.cond_print('debug', "Selected %s toolset..." %
                            (toolset_name))
        except:
            self.critical_error("instantiating toolset %s" %
                                (toolset_name))        

    def run_probe(self, probe):
        '''Execute probe'''
        probe_name = probe.attrib['name']
        ret = Element("probe-data", probe_name, source=probe.attrib['source'])
        try:
            script = open(mktemp(), 'w+')
            try:
                script.write("#!%s\n" %
                             (probe.attrib.get('interpreter', '/bin/sh')))
                script.write(probe.text)
                script.close()
                chmod(script.name, 0755)             
                ret.text = popen(script.name).read()                
            finally:
                unlink(script.name)
        except:
            self.critical_error("executing probe %s" % (probe_name))
        return ret

    def critical_error(self, operation):
        '''Print tracebacks in unexpected cases'''
        print "Traceback information (please include in any bug report):"
        (ttype, value, trace) = exc_info()
        for line in extract_tb(trace):        
            print "File %s, line %i, in %s\n   %s\n" % (line)
            print "%s: %s\n" % (ttype, value)

        self.fatal_error("An unexpected failure occurred in %s" % (operation) )

    def fatal_error(self, message):
        '''Signal a fatal error'''
        print "Fatal error: %s" % (message)
        raise SystemExit, 1

    def warning_error(self, message):
        '''Warn about a problem but continue'''
        print "Warning: %s" % (message)

    def usage_error(self, message):
        '''Die because script was called the wrong way'''
        print "Usage error: %s" % (message)
        self.print_usage()
        raise SystemExit, 2

    def cond_print(self, state, message):
        '''Output debugging information'''
        if self.setup[state]:
            print "bcfg2[%s]: %s" % (state, message)
        
    def print_usage(self):
        ''' Display usage information for bcfg2 '''
        print "bcfg2 usage:"
        for arg in self.options.iteritems():
            if self.argumentDescriptions.has_key(arg[0]):
                print " -%s %s\t%s" % (arg[1],
                                       self.argumentDescriptions[arg[0]],
                                       self.descriptions[arg[0]])
            else:
                print " -%s\t\t\t%s" % (arg[1], self.descriptions[arg[0]])

    def fill_setup_from_file(self, setup_file, ret):
        ''' Read any missing configuration information from a file'''
        default = {
            'server': 'http://localhost:6789/',
            'user': 'root',
            'retries': '6'
            }
        config_locations = {
            'server': ('components', 'bcfg2'),
            'user': ('communication', 'user'),
            'password': ('communication', 'password'),
            'retries': ('communicaton', 'retries')
            }

        self.cond_print_setup('debug')

        config_parser = None        

        for (key, (section, option)) in config_locations.iteritems():
            try:
                if not (ret.has_key(key) and ret[key]):
                    if config_parser == None:
                        self.cond_print('debug', "no %s provided, reading setup info from %s" %
                                       (key, setup_file))
                        config_parser = ConfigParser()
                        config_parser.read(setup_file)
                    try:
                        ret[key] = config_parser.get(section, option)
                    except (NoSectionError, NoOptionError):
                        if default.has_key(key):
                            ret[key] = default[key]
                        else:
                            self.fatal_error(
                                "%s does not contain a value for %s (in %s)"  %
                                (setup_file, option, section))
            except IOError, io_error:
                self.fatal_error("unable to read %s: %s" %
                                 (setup_file, io_error))
            except SystemExit:
                raise
            except:
                self.critical_error("reading config file")            

    def get_setup(self, args):
        '''parse options into a dictionary'''

        for option in self.options.keys():
            self.setup[option] = False
        
        gstr = "".join([self.options[option] +
                        if_then(self.argumentDescriptions.has_key(option),
                                ':', '')
                        for option in self.options.keys()])

        try:
            ginfo = getopt(args, gstr)
        except GetoptError, gerr:
            self.usage_error(gerr)

        for (gopt, garg) in ginfo[0]:
            option = self.argOptions[gopt[1:]]
            if self.argumentDescriptions.has_key(option):
                self.setup[option] = garg
            else:
                self.setup[option] = True

        if (self.setup["file"] != False) and (self.setup["cache"] != False):
            self.usage_error("cannot use -f and -c together")

        if self.setup["help"] == True:
            self.print_usage()
            raise SystemExit, 0

        if self.setup["setup"]:
            setup_file = self.setup["setup"]
        else:
            setup_file = '/etc/bcfg2.conf'

        self.fill_setup_from_file(setup_file, self.setup)
        
    def run(self):
        ''' Perform client execution phase '''
        times = {}

        # begin configuration
        times['start'] = time()

        if self.setup['file']:
            # read config from file
            try:
                self.cond_print('debug', "reading cached configuration from %s" %
                               (self.setup['file']))
                configfile = open(self.setup['file'], 'r')
                rawconfig = configfile.read()
                configfile.close()
            except IOError:
                self.fatal_error("failed to read cached configuration from: %s"
                                 % (self.setup['file']))
        else:
            # retrieve config from server
            proxy = SafeProxy(self.setup, self)
        
            probe_data = proxy.run_method("probe download", "GetProbes", ())

            times['probe_download'] = time()
        
            try:
                probes = XML(probe_data)
            except XMLSyntaxError, syntax_error:
                self.fatal_error(
                    "server returned invalid probe requests: %s" %
                    (syntax_error))
            
            # execute probes
            try:
                probe_info = [self.run_probe(probe)
                              for probe in probes.findall(".//probe")]
            except:
                self.critical_error("executing probes")

            # upload probe responses
            proxy.run_method("probe data upload", "RecvProbeData",
                             (probe_info, ))
        
            times['probe_upload'] = time()

            rawconfig = proxy.run_method("configuration download", "GetConfig",
                                          (self.setup['image'],
                                           self.setup['profile']))

            times['config_download'] = time()

        if self.setup['cache']:
            try:
                open(self.setup['cache'], 'w').write(rawconfig)
            except IOError:
                self.warning_error("failed to write config cache file %s" %
                                   (self.setup['cache']))
            times['caching'] = time()
            
        try:
            self.config = XML(rawconfig)
        except XMLSyntaxError, syntax_error:
            self.fatal_error("the configuration could not be parsed: %s" %
                             (syntax_error))

        times['config_parse'] = time()
    
        if self.config.tag == 'error':
            self.fatal_error("server error: %s" % (self.config.text))

        # Get toolset from server
        try:
            toolset_name = self.config.get('toolset')
        except:
            self.fatal_error("server did not specify a toolset")

        if self.setup['bundle']:
            replacement_xml = Element("Configuration", version='2.0')
            for child in self.config.getroot().getchildren():
                if ((child.tag == 'Bundle') and
                    (child.attrib['name'] == self.setup['bundle'])):
                    replacement_xml.append(child)
            self.config = replacement_xml

        # Create toolset handle
        self.load_toolset(toolset_name)

        times['initialization'] = time()
        
        # verify state
        self.toolset.Inventory()

        times['inventory'] = time()
    
        # summarize current state
        self.toolset.CondDisplayState('verbose', 'initial')

        # install incorrect aspects of configuration
        self.toolset.Install()

        self.toolset.CondDisplayState('verbose', "final")

        times['install'] = time()
        times['finished'] = time()

        if not self.setup['file']:
            # upload statistics
            feedback = Element("upload-statistics")
            timeinfo = Element("OpStamps")
            for (event, timestamp) in times.iteritems():
                timeinfo.set(event, str(timestamp))
            stats = self.toolset.GenerateStats(__revision__)
            stats.append(timeinfo)
            feedback.append(stats)

            proxy.run_method("uploading statistics",
                             "RecvStats", (tostring(feedback),))


if __name__ == '__main__':
    signal(SIGINT, cb_sigint_handler)
    Client(argv[1:]).run()