source: fedd/federation/proxy_segment.py @ e91c492

#!/usr/local/bin/python

import sys, os
import re

import tempfile
import subprocess
import logging 
import time
import signal

from service_error import service_error

class proxy_segment:
    """
    Base class for segment starter classes that access their underlying testebd
    remotely using ssh.  It is promarily a code repository for commonly used
    ssh commands for moving code and logging in.
    """
    class ssh_cmd_timeout(RuntimeError): pass

    def __init__(self, log=None, keyfile=None, debug=False):
        """
        log is the logging.log to print messages to, keyfile is the private key
        for ssh interactions and if debug is true, commands are not executed
        using ssh.
        """
        self.log = log or logging.getLogger('fedd.access.proxy_segment')
        self.ssh_privkey_file = keyfile
        self.debug = debug
        self.ssh_exec="/usr/bin/ssh"
        self.scp_exec = "/usr/bin/scp"
        self.ssh_cmd_timeout = proxy_segment.ssh_cmd_timeout

    def scp_file(self, file, user, host, dest=""):
        """
        scp a file to the remote host.  If debug is set the action is only
        logged.
        """

        scp_cmd = [self.scp_exec, '-o', 'IdentitiesOnly yes', 
                '-o', 'StrictHostKeyChecking no', '-o', 'ForwardX11 no','-i', 
                self.ssh_privkey_file, file, 
                "%s@%s:%s" % (user, host, dest)]
        rv = 0

        try:
            dnull = open("/dev/null", "w")
        except EnvironmentError:
            self.log.debug("[ssh_file]: failed to open " + \
                    "/dev/null for redirect")
            dnull = Null

        self.log.debug("[scp_file]: %s" % " ".join(scp_cmd))
        if not self.debug:
            rv = subprocess.call(scp_cmd, stdout=dnull, 
                    stderr=dnull, close_fds=True)

        return rv == 0

    def ssh_cmd(self, user, host, cmd, wname=None, timeout=None):
        """
        Run a remote command on host as user.  If debug is set, the action
        is only logged.  Commands are run without stdin, to avoid stray
        SIGTTINs.
        """
        sh_str = ("%s -n -o 'IdentitiesOnly yes' -o " + \
                "'StrictHostKeyChecking no' -o 'ForwardX11 no' " +
                "-i %s %s@%s %s") % \
                (self.ssh_exec, self.ssh_privkey_file, 
                        user, host, cmd)

        try:
            dnull = open("/dev/null", "w")
        except EnvironmentError:
            self.log.debug("[ssh_cmd]: failed to open /dev/null " + \
                    "for redirect")
            dnull = Null

        self.log.debug("[ssh_cmd]: %s" % sh_str)
        if not self.debug:
            if dnull:
                sub = subprocess.Popen(sh_str, shell=True, stdout=dnull,
                        stderr=dnull, close_fds=True)
            else:
                sub = subprocess.Popen(sh_str, shell=True, close_fds=True)
            if timeout:
                i = 0
                rv = sub.poll()
                while i < timeout:
                    if rv is not None: break
                    else:
                        time.sleep(1)
                        rv = sub.poll()
                        i += 1
                else:
                    self.log.debug("Process exceeded runtime: %s" % sh_str)
                    os.kill(sub.pid, signal.SIGKILL)
                    raise self.ssh_cmd_timeout();
                return rv == 0
            else:
                return sub.wait() == 0
        else:
            if timeout == 0:
                self.log.debug("debug timeout raised on %s " % sh_str)
                raise self.ssh_cmd_timeout()
            else:
                return True