root/tags/v1.0/provisioner.py

"""
Provisioners contact remote devices and apply changes to them.

Should be implemented as a Protocol/Factory style thing to fit
with the async mechanism. A Provisioner is an interface applied
to a Factory that creates a, for example, ssh session with a
remote site.

Need to define some sort of Change protocol, to handle a
generic send/expect style interface.
"""
import re

from zope.interface import Interface, implements

from twisted.internet import defer, reactor
from twisted.internet import protocol
from twisted.python import log as tlog
from twisted.internet.error import ProcessDone, ProcessTerminated, AlreadyCancelled, AlreadyCalled

from change import ChangeFailed, ChangeConditionFailure, CHANGE_STATE
import util

import logging
import traceback

import debug
log = logging.getLogger('modipy')

from twisted.internet.base import DelayedCall
DelayedCall.debug = True

class IProvisioner(Interface):
    """
    Defines the Provisioner interface
    """

    def __init__(self, name='', namespace={}):
        """
        """

    def perform_change(self, ignored, change):
        """
        Applies a change to the devices the change is defined to affect.
        """

    def apply_change(self, device, change):
        """
        Apply a change to a specific device.
        """

    def backout_change(self, device, change):
        """
        Back out a change that was applied to a specific device.
        """

class UserBailout(Exception):
    """
    User denied a change command from executing in Authoritarian mode.
    This causes the program to bail out complete at exactly that point.
    """

class Provisioner:
    """
    A Provisioner performs the actual execution of Changesets. It could, for example,
    connect to a remote system via SSH and then execute the Changeset Actions.
    """

    implements( IProvisioner, )

    def __init__(self, name='', namespace={}, authoritarian=False, **kwargs):
        self.name = name
        self.connectedDevice = None
        self.namespace = namespace
        self.authoritarian = authoritarian

        log.debug("My namespace is: %s", self.namespace)

    def parse_config_node(self, node):
        pass

    def perform_change(self, ignored, change, namespace={}, backout=False):
        """
        Perform a change on one or more (potentially) remote entities.

        @param namespace: The global namespace passed in to the provisioner.
               This gets merged with the provisioner, device and change namespaces.
        """
        self.change_ok = {}
        self.change_failed = {}
        self.backout_ok = {}
        self.backout_failed = {}

        log.debug("perform change with namespace: %s", namespace)
        
        # get the list of devices the change should be applied to
        devices = change.devices

        d = defer.succeed(None)
        
        if change.serial_mode:

            # apply the change to each device one after the other,
            log.debug("change %s will be performed on devices: %s", change.name, devices)
            for device in devices:
                log.debug("performing change '%s' on device '%s'", change.name, device)

                namespace['provisioner.name'] = self.name
                namespace['provisioner.type'] = self.__class__.__name__

                namespace['change.name'] = change.name
                namespace['change.type'] = change.__class__.__name__

                namespace['device.name'] = device.name
                namespace['device.fqdn'] = device.fqdn
                namespace['device.ipaddress'] = device.ipaddress

                namespace.update(self.namespace)
                namespace.update(change.namespace)
                namespace.update(device.namespace)

                if backout:
                    # Just do the backout portion
                    d.addCallback(self.backout_change, device, change, namespace)

                else:
                    log.debug("applying change with namespace: %s", namespace)
                    d.addCallback(self.apply_change, device, change, namespace)
                    
                    d.addCallback(self.change_complete_success, change, namespace)
                    d.addErrback(self.change_failure, change, namespace)
                    pass
                pass
            pass
        return d

    def apply_change(self, ignored, device, change, namespace={}):
        """
        Apply the change to a device
        """
        # connect to the entity that is having the change applied to it
        d = self.connect(device, namespace)
        d.addCallback(self.do_pre_apply_check, device, change, namespace)
        d.addCallbacks(self.pre_apply_success, self.pre_apply_failed, callbackArgs=(device, change, namespace), errbackArgs=(device, change, namespace))
        #d.addCallback(self.pre_apply_success, device, change, namespace)
        #d.addErrback(self.pre_apply_failed, change, namespace)
        return d

    def change_apply_success(self, result, device, change, namespace):
        """
        Change was applied successfully.
        """
        log.debug("change applied to device '%s' successfully." % device)
        log.debug("result: %s", result)
        # mark the device as having had the change applied to it
        self.change_ok[device] = change

        # do the post-change checking
        return change.test_apply_success(self, result, namespace)

    def change_apply_failed(self, failure, device, change, namespace):

        log.error("Change '%s' failed to apply", change.name)
        e = failure.check( ChangeConditionFailure, UserBailout )
        if e:
            log.error("  failure was: %s, %s", failure.type, failure.value )

            # If we want to bail out, bail now
            if failure.type == UserBailout:
                return defer.fail(failure)
        else:
            log.error("Unhandled failure in provisioner.change_apply_failed()")
            tlog.err(failure)
        
        self.change_failed[device] = failure
        change.set_state('total_failure')
        
        d = self.backout_change(None, device, change, namespace)

        if change.on_fail_continue:
            log.info("Attempting to continue, despite failure...")
            change.set_state('partial_failure')
            pass

        elif change.backout_all:
            for device in self.change_ok.keys():
                d.addCallback( self.backout_change, device, change, namespace )

                pass
            pass

        return d

    def change_failure(self, failure, change, namespace):
        log.debug("Change failure for %s: %s", change.name, failure.value)
        if change.state in [ CHANGE_STATE['pending'], CHANGE_STATE['retry'] ]:
            # If the change is marked as 'on_fail: retry', let it retry
            if change.can_retry():
                change.state = CHANGE_STATE['retry']
            else:
                change.state = CHANGE_STATE['total_failure']
        else:
            log.error("change_failure unhandled change state: %s", change.state)
            raise ValueError("Invalid change state: %s" % change.state)

        # Propogate a UserBailout failure
        if failure.type == UserBailout:
            return failure
        
        #tlog.err(failure)

    def change_complete_success(self, result, change, namespace):
        if change.state in [ CHANGE_STATE['pending'], CHANGE_STATE['retry'] ]:
            change.state = CHANGE_STATE['success']
            log.info("Change '%s' was a success!", change.name)
        else:
            log.error("change_complete_success in weird state: %s", change.state)
            raise ValueError("wrong change state: %s" % change.state)
            change.state = CHANGE_STATE['backout_ok']
            log.info("Change '%s' backed out successfully", change.name)

    def do_pre_apply_check(self, result, device, change, namespace):
        # perform pre-implementation checks
        log.info("Doing pre-apply check for change '%s' on device '%s'", change.name, device)
        return change.pre_apply_check(self, namespace)

    def pre_apply_success(self, result, device, change, namespace):
        log.info("Pre-apply check passed for change '%s' on device '%s'", change.name, device)
        d = self.do_apply(result, device, change, namespace)
        d.addCallback(self.change_apply_success, device, change, namespace)
        d.addErrback(self.change_apply_failed, device, change, namespace)
        return d

    def pre_apply_failed(self, failure, device, change, namespace):
        log.error("Pre-apply check failed for change '%s' on device '%s'", change.name, device)
        raise failure

    def do_apply(self, result, device, change, namespace):
        """
        Apply the actual change
        """
        log.info("Applying change '%s' to device '%s'", change.name, device)
        return change.apply(self, namespace)

    def backout_change(self, ignored, device, change, namespace):
        """
        Back out a change that was applied to a device.
        """
        log.info("Backing out change '%s' from device '%s'", change.name, device)
        d = change.backout(self, namespace)
        d.addCallback(change.test_backout_success, self, namespace)
        d.addCallback(self.backout_success, device, change, namespace)
        d.addErrback(self.backout_failure, device, change, namespace)
        return d

    def backout_success(self, ignored, device, change, namespace):
        log.info("Successfully backed out change: '%s' from device '%s'", change.name, device)
        change.state = CHANGE_STATE['backout_ok']
        self.backout_ok[device] = change

    def backout_failure(self, failure, device, change, namespace):
        log.critical("Backout of change '%s' failed for device '%s'!", change.name, device)
        change.state = CHANGE_STATE['backout_failed']
        self.backout_failed[device] = change
        #tlog.err(failure)
        return failure

class CommandProvisioner(Provisioner):
    """
    A command provisioner provides methods for implementing command changes on
    remote devices.
    """

    command_re = re.compile(r'^(?P<start>.*)"(?P<quoted>.*)"(?P<end>.*)$')

    def parse_config_node(self, node):
        Provisioner.parse_config_node(self, node)
        pass

    def __init__(self, name='', namespace={}, authoritarian=False, command_timeout=300):
        Provisioner.__init__(self, name, namespace, authoritarian)
        self.command_timeout = int(command_timeout)

    def split_command(self, cmdstring):
        """
        Split a command string, returning everything before
        a quoted section, the quoted section, and then everything
        afterwards.
        """
        cmd_split = []
        m = self.command_re.match(cmdstring)
        if m:
            cmd_split.append( m.group('quoted') )
            cmd_split.extend(m.group('end').split())

            splitpart = self.split_command(m.group('start'))
            splitpart.extend(cmd_split)
            cmd_split = splitpart
            return cmd_split

        return cmdstring.split()

    def run_command(self, ignored, cmdstring):
        """
        Execute a command on the element, returning the results of the command.

        @returns: (return_code, stdout, stderr) where
          return_code: is the exit code of the command, as an integer
          stdout: is a string containing anything printed to STDOUT
          stderr: is a string containing anything printed to STDERR
        """
        log.debug("Executing provisioning command: %s" % cmdstring)
        return defer.succeed( (0, '', '') )

    def connect(self, device):
        pass

class ConnectingProvisioner(CommandProvisioner):
    """
    A ConnectingProvisioner runs a command to connect to the remote device
    and then runs a series of commands by sending them to the remote
    device via STDIN, and receives the output from STDOUT and STDERR.
    """
    def parse_config_node(self, node):
        CommandProvisioner.parse_config_node(self, node)

    def connect(self, device):
        self.connectedDevice = device
        
        self.childConnectingDefer = defer.Deferred()
        self.childConnectingTimeout = reactor.callLater(5, self.childConnectTimeout)

        log.debug("connecting to device '%s'", device)

        cmd = [ '/usr/bin/ssh',
                '-T',
                '-o','BatchMode=yes',
                '-i','/home/daedalus/.ssh/localconnect',
                'localhost' ]
        #cmd = [ '/usr/bin/ssh', 'localhost' ]
        self.ep = ExpectProtocol(self)
        self.p = reactor.spawnProcess(self.ep, cmd[0], cmd )

        #log.debug("timeout thing: %s" % self.childConnectingTimeout)

        return self.childConnectingDefer

    def childConnectTimeout(self):
        """
        This gets called if my connection via the child process doesn't
        complete within a certain timeout value.
        """
        log.error("Connection to child timed out!")
        self.ep.transport.loseConnection()
        self.childConnectingDefer.errback( Exception("Connection timed out") )

    def childConnectSuccess(self):
        """
        This gets called by the child when it finishes connecting.
        Cancel the timeout and notify waiting things that the connection
        completed successfully.
        """
        log.debug("successfully connected to child")
        self.childConnectingTimeout.cancel()
        log.debug("sent timeout cancel")
        self.childConnectingDefer.callback(self)

    def childConnectFailure(self):
        """
        This gets called by the child when it finishes connecting.
        Cancel the timeout and notify waiting things that the connection
        completed successfully.
        """
        log.error("connection by child failed")
        self.childConnectingTimeout.cancel()
        self.childConnectingDefer.errback(ValueError("Child connection failure!"))

    def run_commands(self, ignored, expectset, namespace):
        log.debug('running commands on remote host: %s' % expectset)

        return self.ep.run_remote_commands(expectset)

        # Run the command using the ExpectProtocol
        #return defer.succeed( (0, '', '') )
    
class ExpectProtocol(protocol.Protocol):
    """
    The ExpectProtocol allows you to supply a series of
    expect/send pairs that will be used to handle the
    processing of received data.

    Each expect/send pair is a tuple of the form:
    (expect_re, send_string), where:

    expect_re: is a regular expression that will be used
    by the re module to parse received data. Each time
    dataReceived() is triggered, the data is added to a
    buffer, which is matched against the re. If a match is
    successful, the corresponding send_dict is processed to
    generate the data to send in response, and the receive
    buffer is emptied.
    If a match is unsuccessful, execution returns to wait
    for more data.

    send_string: is a series of octets to send in response
    to the expect string that was received. % operator replacement
    is possible using %(name)s syntax, with the (name) values
    being populated from the previous expect re.match operation.
    """

    def __init__(self, provisioner):
        
        self.parent = provisioner
        self.saidHello = False
        self.waitingForCommand = None
        self.databuf = ''
        self.data_wait_timeout = None

        self.results = ''

    def connectionMade(self):
        """
        Called when the connection to the remote entity is
        made, whatever it is.
        """
        log.debug("Connected to remote thing! Hurray!")
        self.parent.childConnectSuccess()

    def childConnectionLost(self, fd):
        log.debug("Lost connection to child process fd %s" % fd)

    def childDataReceived(self, fd, data):
        log.debug("received data from child process %s: >%s<" % (fd, data) )

        try:
            self.databuf += data
            log.debug("databuf: >%s<" % self.databuf)
            self.do_expect_processing()

        # force any unexpected errors to call the errback
        except Exception, e:
            self.waitingForCommand.errback(e)

        # ignore stuff from STDERR?
        #if fd == 2:
        #    pass

    def processEnded(self, status_object):
        log.debug("process ended: %s" % status_object)

        if isinstance(status_object.value, ProcessDone):
            log.debug("process exited ok")
        elif isinstance(status_object.value, ProcessTerminated):
            log.debug("process did not exit ok: %s" % status_object.value.exitCode)
            self.parent.childConnectFailure()

        log.debug("handled error")

    def connectionFinished(self, status):
        self.parent.childConnected()

    def run_remote_commands(self, expectset):
        """
        Run a set of commands, using expect processing to do it.
        """
        self.expectset = expectset
        self.waitingForCommand = defer.Deferred()
        self.do_expect_processing()
        return self.waitingForCommand

    def do_expect_processing(self):

        if self.data_wait_timeout:
            self.data_wait_timeout.cancel()

        # get the first expression/cmd pair
        expr, cmdstring = self.expectset[0]
        self.expectset = self.expectset[1:]
        
        # check the expr to see if we should process it
        if expr is None:
            if cmdstring is not None:
                log.debug("No expression to wait for. Running command '%s' immediately." % cmdstring)
                self.issue_command(cmdstring)
            else:
                log.debug("cmdstring is 'None', so nothing to do")

        else:
            log.debug("checking for '%s' in databuf: %s", expr, self.databuf)
            re_expr = re.compile(expr)
            log.debug("compiled expr")

            # Match in multi-line mode, and '.' character matches \n
            match = re_expr.search(self.databuf)
            log.debug("match object is: %s", match)
            if match:
                # check the current recv buffer to see if the expect string is a match
                log.debug("Found expr '%s' in databuf! Yay!" % expr)

                # reset the databuffer whenever we match something
                # FIXME: maybe not do this?
                self.results += self.databuf
                self.databuf = ''
                
                if cmdstring is not None:
                    self.issue_command(cmdstring)
                    pass
                else:
                    # If the command is None, this is the end of processing,
                    # so shut down the connection and trigger command success processing.
                    self.transport.loseConnection()
                    self.waitingForCommand.callback(self.results)
                    
            else:
                log.debug("cannot find '%s' in databuf. Waiting for more data..." % expr)
                self.expectset.insert(0, (expr, cmdstring) )
                self.data_wait_timeout = reactor.callLater(10, self.data_wait_too_long)
                pass
            pass

    def issue_command(self, cmdstring):
        """
        Write the command to the remote device.
        """
        # If we're in authoritarian mode, wait for confirmation
        # that we should execute the command.
        if self.parent.authoritarian:
            log.debug("Authoritarian mode. Waiting for ok to proceed...")
            isok = raw_input("Issue command (y/n)[n]?> ")
            if isok.startswith('y'):
                log.debug("Ok! Let's continue!")
            else:
                log.info(" Bailing out at your command.")
                raise UserBailout("Bailing out at your command")
                #self.waitingForCommand.errback("User requested manual bailout.")
                return
            pass
        log.debug("running command: %s" % cmdstring)                    
        self.transport.writeToChild(0, '%s\n' % cmdstring)
        self.data_wait_timeout = reactor.callLater(10, self.data_wait_too_long)

    def data_wait_too_long(self):
        """
        We have waited for data for too long, so return an error
        """
        self.waitingForCommand.errback( ValueError("Timed out waiting for command output.") )


class MultiConnectingProvisioner(ConnectingProvisioner):
    """
    A provisioner that runs a command that sets up the connection to
    the remote device for each and every command that needs to be run.
    The command to be run is appended to the connecting command.
    It is designed to be run with something like ssh, eg:
    ssh remotehost <command1>
    ssh remotehost <command2>
    """

    def parse_config_node(self, node):
        """
        Deals with additional configuration that I expect to be provided with.
        """
        ConnectingProvisioner.parse_config_node(self, node)
        log.debug("Parsing config node: %s", node)
        # parse command definition
        if node.tag == 'command':
            self.command = str(node.text)
            log.debug("Set provisioner command to: %s", self.command)

    def connect(self, device, namespace):
        """
        Don't actually connect, but take note of the device.
        """
        self.device = device
        # perform variable substitution on the command

##         self.cmd = [ '/usr/bin/ssh',
##                      '-T',
##                      '-o', 'BatchMode=yes',
##                      '-i','/home/daedalus/.ssh/localconnect',
##                      ]
##        self.cmd.append('%s' % str(device))

        #log.debug("Set up command to run: %s", self.cmd)
        log.debug("Provisioner ready for device: %s", device)
        return defer.succeed(None)
    
    def run_commands(self, ignored, expectset, namespace):
        """
        Run each command, one after the other.
        """
        log.debug("inside provisioner running expectset: %s" % expectset)
        self.cmdoutput = ''
        self.exitcode = None

        # This deferred is called when all commands have finished
        self.all_commands_defer = defer.Deferred()

        # Set up a sequential chain of commands
        d = defer.succeed(None)
        for (expr, cmdstring) in expectset:

            try:
                # Perform variable substitution on the command string
                log.debug("determining commandstring from template: %s", cmdstring)
                cmdstring = str(util.substituteVariables(cmdstring, namespace))
                log.debug("commandstring is: %s", cmdstring)

                # add the cmdstring to the namespace
                namespace['command.send'] = cmdstring

                log.debug("Determining command base from template: %s", self.command)
                command = str(util.substituteVariables(self.command, namespace))
            except KeyError, e:
                log.error("KeyError in commands: %s" % e)
                self.all_commands_defer.errback( e )
                return self.all_commands_defer
            
            #log.debug("checking expr: %s" % expr)
            log.debug("cmdstring is: %s" % command)
            if command is not None:
                log.debug("splitting command: %s", command)
                cmd = self.split_command(command)
                log.debug("command is: %s", cmd)
                d.addCallback(self.spawn_command, cmd, self.command_timeout)
                d.addCallbacks(self.command_completed, self.command_failed)
                pass
            pass

        d.addCallbacks(self.all_commands_done, self.all_commands_failure)
        return self.all_commands_defer

    def spawn_command(self, ignored, cmd, timeout=300):
        """
        Run the command in a sub-process, returning a deferred that
        will fire based on the results.
        """
        log.debug("spawning command...")
        self.waitingForCommand = defer.Deferred()

        # If we're in authoritarian mode, wait for confirmation
        # that we should execute the command.
        log.debug("Checking for authoritarian mode...")
        if self.authoritarian:
            log.debug("Authoritarian mode. Waiting for ok to proceed...")
            isok = raw_input("Issue command: %s\n(y/n)[n]?> " % ' '.join(cmd) )
            if isok.startswith('y'):
                log.debug("Ok! Let's continue!")
            else:
                log.info("Bailing out at your command.")
                self.exitcode = -1
                self.cmdoutput = 'User requested manual bailout'
                log.critical("User requested manual bailout")
                self.waitingForCommand.errback( UserBailout("User requested manual bailout.") )
                return self.waitingForCommand
            pass

        log.debug("running command remotely: %s" % cmd)
        self.ep = SingleCommandProtocol(self, timeout)
        self.p = reactor.spawnProcess(self.ep, cmd[0], cmd)
        log.debug("Spawned process: %s", self.p)

        return self.waitingForCommand

    def command_completed(self, result):
        """
        When the command is completed, add its output to a
        global resultset. This sets the exitcode to the
        exitcode of the last command that completed, and
        """
        log.debug("a command has completed with result: %s", result)
        self.exitcode = result[0]
        self.cmdoutput += result[1]
        log.debug("current results: exit '%d', output: %s", self.exitcode, self.cmdoutput)
        
    def command_failed(self, failure):
        errorstr = "%s: %s" % (self.exitcode, self.cmdoutput)
        log.error("Command failed: %s", errorstr)
        #self.all_commands_defer.errback( Exception(errorstr) )
        return failure
    
    def all_commands_done(self, result):
        """
        Called when all the commands have completed.
        """
        log.debug("All commands have finished.")
        self.all_commands_defer.callback( (self.exitcode, self.cmdoutput) )

    def all_commands_failure(self, failure):
        """
        Some kind of failure in the commands occurred.
        """
        e = failure.check( UserBailout )
        if e:
            #log.error("User bailout detected.")
            self.all_commands_defer.errback( failure )

class CommandFailure(Exception):
    """
    Used to contain sub-process errors.
    """

class SingleCommandProtocol(protocol.Protocol):
    """
    Used to run a single command in a sub-process.
    """
    
    def __init__(self, provisioner, timeout_after=300):

        self.parent = provisioner
        self.databuf = ''
        self.exitCode = None
        self.timeout_after = timeout_after

    def connectionMade(self):
        #log.debug("Pretend connection made! Hurray!")

        # FIXME: provide a way to set this via the config file
        self.timeout = reactor.callLater(self.timeout_after, self.timedOut)
        pass
    
    def childConnectionLost(self, fd):
        #log.debug("Lost connection to child process fd %s" % fd)
        pass

    def childDataReceived(self, fd, data):
        #log.debug("received data from child process %s: >%s<" % (fd, data) )

        self.databuf += data
        log.debug("databuf: >%s<" % self.databuf)

    def processEnded(self, status_object):
        log.debug("process ended: %s" % status_object)
        try:
            if isinstance(status_object.value, ProcessDone):
                log.debug("process exited ok: %s" % status_object.value.exitCode)
                self.exitCode = 0
                self.timeout.cancel()
                self.parent.waitingForCommand.callback( (self.exitCode, self.databuf) )
            
            elif isinstance(status_object.value, ProcessTerminated):
                log.debug("process did not exit ok: %s: [%s: %s] %s" % (status_object.value.exitCode, status_object.value.signal, status_object.value.status, self.databuf))
                self.exitCode = status_object.value.exitCode
                self.timeout.cancel()
                self.parent.waitingForCommand.callback( (self.exitCode, self.databuf) )
                pass

        except AlreadyCalled, AlreadyCancelled:
            pass
                
        except Exception, e:
            log.error("Exception raised during end of process processing. That's quite bad.")
            raise
##             traceback.print_exc(e)
##             self.timeout.cancel()
##             self.parent.waitingForCommand.errback( e )

    def timedOut(self):
        # FIXME: If the command times out, we need to ignore anything that comes
        # back from the process, or we'll notify of error conditions more than once
        # I'm hoping the AlreadyCalled an AlreadyCancelled catch above will take care of this.
        log.error("Timed out waiting for command to exit")
        self.transport.loseConnection()
        self.parent.waitingForCommand.errback( (1, '%s: command timeout' % self.__class__) )