host/deploy/cmdexec.py

"""This file contains classes that facilitate command execution on local and
remote hosts.
"""

import logging
from subprocess import Popen, PIPE, STDOUT

LOG = logging.getLogger()
SSH_ARGS = [
    '-o', 'StrictHostKeyChecking=no',
    '-o', 'UserKnownHostsFile=/dev/null',
    '-q',
]


class _SshExecBackend(object):
    def __init__(self, server):
        self._server = server
        self._shell = Popen(['ssh', server] + SSH_ARGS,
                            stdin=PIPE, stdout=PIPE, stderr=STDOUT, close_fds=True)


    def shell(self):
        """Get shell process.

        Returns:
            Popen class interfacing shell.
        """
        return self._shell


    def copy(self, source, target):
        """Copy source file to target location.

        This command may throw if execution was unsuccessful.

        Args:
            source Source file location (either local or remote)
            target Target file location (either local or remote)

        Returns:
            None
        """
        LOG.info('Transferring file %s -> %s' % (source, target))
        proc = Popen(['scp'] + SSH_ARGS + [source, '%s:%s' % (self._server, target)],
                     stdin=PIPE, stdout=PIPE, close_fds=True)
        result = proc.wait()
        if result != 0:
            raise Exception('Could not complete file transfer: scp returned %d' % result)


class _LocalExecBackend(object):
    def __init__(self):
        self._shell = Popen(['/bin/sh'], stdin=PIPE, stdout=PIPE, stderr=STDOUT, close_fds=True)


    def shell(self):
        """Get shell process.

        Returns:
            Popen class interfacing shell.
        """
        return self._shell


    def copy(self, source, target):
        """Copy source file to target location.

        This command may throw if execution was unsuccessful.

        Args:
            source Source file location (must be local)
            target Target file location (must be local)

        Returns:
            None
        """
        LOG.info('Copying file %s -> %s' % (source, target))
        proc = Popen(['cp', source, target], stdin=PIPE, stdout=PIPE, close_fds=True)
        result = proc.wait()
        if result != 0:
            raise Exception('Could not complete file transfer: scp returned %d' % result)


class Target(object):
    """RemoteServer allows faster command execution on remote server by keeping a
    single channel open at all times.
    """
    def __init__(self, backend):
        self.backend = backend


    @staticmethod
    def for_localhost():
        """Create new command executor for localhost.

        Commands will be invoked via /bin/sh.

        Returns:
            New instance of CmdExec.
        """
        return Target(_LocalExecBackend())


    @staticmethod
    def for_remote_host(server):
        """Create new command executor for remote server.

        Commands will be invoked via SSH.

        Args:
            server SSH server name (either IP address or server configured via ssh_config)

        Returns:
            New instance of CmdExec.
        """
        return Target(_SshExecBackend(server))


    def execute_no_wait(self, command):
        """Execute operation that may never return, such as subshell."""
        LOG.info('Executing: %s', command)
        self.backend.shell().stdin.write(command)


    def execute(self, command):
        """Execute supplied command on remote server.

        Args:
            command Command or expression to be executed on remote server.

        Returns:
            command output.
        """
        result = self.backend.shell().poll()
        if result:
            raise Exception('Server connection died:', result)

        header = '-=-=-=- COMMAND HEADER -=-=-=-'
        footer = '-=-=-=- COMMAND FOOTER -=-=-=-'

        LOG.info('Executing: %s', command)
        self.backend.shell().stdin.write(
            '\necho \'%s\'; ( %s ) || true; echo \'%s\'\n' % (header, command, footer))
        out = []
        found_header = False

        while True:
            line = self.backend.shell().stdout.readline().strip()
            # Search for header
            if not found_header:
                if line == header:
                    found_header = True
                continue

            if line == footer:
                break

            out.append(line)

        return out


    def copy(self, source, target):
        """Copy source file to target location.

        This call may throw if command execution was not successful.

        Args:
            source Source file name
            target Target file name
        """
        self.backend.copy(source, target)