sshctrl.py

#
#    Copyright 2008 Zdravko Bozakov and Michael Bredel
#
#    This file is part of SSHLauncher.
#
#    SSHLauncher is free software: you can redistribute it and/or modify
#    it under the terms of the GNU General Public License as published by
#    the Free Software Foundation, either version 2 of the License, or
#    (at your option) any later version.
#
#    SSHLauncher is distributed in the hope that it will be useful,
#    but WITHOUT ANY WARRANTY; without even the implied warranty of
#    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#    GNU General Public License for more details.
#
#    You should have received a copy of the GNU General Public License
#    along with Foobar.  If not, see <http://www.gnu.org/licenses/>.
#
try:
    import pexpect
    from pexpect import pxssh
except ImportError:
    print("Could not find pexpect module. You can install it using:")
    print("    pip install pexpect")
    exit(1)
import select
import threading
import time
import os
from packaging import version

if version.parse(pexpect.__version__) < version.parse("3.0"):
    print('The installed version of pexpect is not supported.')
    print('Please install pexpect>=3.0.')
    exit(3)

COLOR_TERM = False
if os.environ.get("TERM"):
    COLOR_TERM = True


LOOP_DELAY = .1
WAIT_NOTIFICATION_DELAY = 60  # seconds

DEBUG = False
VERBOSE = True
# LSD-style debug output
DEBUG_LABEL = '\033[43;1;37mDEBUG\033[m'


def session_tag(id_str):
    """Print section [id] in session_tag"""
    h = hash(id_str)
    ansi_fg = 30 + (h & 7)
    # ansi_bg = 40 + (-h & 7)
    ansi_bg = ansi_fg
    return ansi_color('['+id_str+']', bold=True,
                      ansi_bg=ansi_bg,
                      ansi_fg=ansi_fg)


def ansi_color(msg, bold=True, ansi_fg=31, ansi_bg=47):
    """Print section [id] in session_tag"""
    # TODO check if the terminal supports it
    # see also: http://pypi.python.org/pypi/colorama
    if bold:
        ansi_bold = 1
    else:
        ansi_bold = 22
    if COLOR_TERM:
        return "\033[%d;%d;%dm%s\033[0m" % (ansi_bg, ansi_bold, ansi_fg, msg)
    else:
        return msg


def ansi_bold(msg):
    if COLOR_TERM:
        return "\033[1m%s\033[0m" % msg
    else:
        return msg


class ConfigError(Exception):
    def __init__(self):
        pass


class SSHControl (threading.Thread):
    """Starts a thread which establishes an SSH connection and executes a command.

    """
    ID_STR_LEN = 0
    # pxssh instance
    s = None
    # strore all ssh threads
    ssh_threads = []
    # kill all threads if this is set to True
    terminate_threads = False

    thread_lock = None
    ESCAPE = False
    SIMULATE = False

    # start a login (l) or interactive(i) bash shell when issuing remote command
    # TODO this should be an option
    SHELL = '/bin/bash -ls'
    # commands which should always executed before the command string
    SHELL_SETUP = 'unset HISTFILE; '
    # how long should we wait (in seconds) before assuming the ssh connection
    # cannot be established
    SSH_LOGIN_TIMEOUT = 15
    # how long should we wait (in seconds) before attempting to reestablish an
    # ssh connection
    SSH_LOGIN_REPEAT_TIMEOUT = 5
    # how long should we wait (in seconds) for an expect (after) string.
    # Consider increasing for very long experiments.
    PEXPECT_TIMEOUT = 86400     # 24 hours
    # how long should we wait (in seconds) for the prompt
    PROMPT_TIMEOUT = 86400      # 24 hours

    def __init__(self, id, hostname, port, username, passwd, command, after={},
                 sync=[], delay=0):
        """initialize class variables"""
        threading.Thread.__init__(self)

        self.id = id
        SSHControl.ID_STR_LEN = max(len(session_tag(id)), SSHControl.ID_STR_LEN)

        self.connected = False
        self.registered_after = False
        self.registered_sync = False
        self.syncNotified = False
        self.config_ok = False

        self.hostname = hostname
        self.username = username
        self.passwd = passwd
        self.delay = delay
        self.command = command
        self.after = after
        self.sync = sync
        self.port = port
        self.s = None

        # store after strings from other sessions
        self.after_list = {}
        self.syncList = []

        self.lock = threading.Lock()
        self.lockSync = threading.Lock()

        self.daemon = True

    def run(self):
        """thread method"""
        self.name = self.id
        time.sleep(self.delay)

        # wait until all threads have been started
        while not all([t.is_alive() for t in SSHControl.ssh_threads]):
            time.sleep(LOOP_DELAY)

        # register own expect string on thread id specified by after
        if self.after:
            for t in SSHControl.ssh_threads:
                for key in list(self.after.keys()):
                    if t.name == str(key):
                        t.register_after(self.id, self.after[key])
        self.registered_after = True

        # wait until all threads have registered
        while any([not t.registered_after for t in SSHControl.ssh_threads]):
            time.sleep(LOOP_DELAY)

        # register sync expect string on thread id specified by sync
        if self.sync:
            for t in SSHControl.ssh_threads:
                for key in self.sync:
                    if t.name == str(key):
                        t.register_sync(self.id)
        self.registered_sync = True

        # wait until all threads are synchonized
        while any([not t.registered_sync for t in SSHControl.ssh_threads]):
            time.sleep(LOOP_DELAY)

        # now check the configuration
        try:
            self._check_config()
        except ConfigError:
            self.ssh_abort()
            return

        # wait to see if all configurations are ok
        while any([not t.config_ok for t in SSHControl.ssh_threads]):
            time.sleep(LOOP_DELAY)

        # wait until all expect strings registered by current session have been
        # matched
        timer = time.time()
        while (self.after and not self.syncNotified):
            time.sleep(LOOP_DELAY)
            # periodically tell the user what's happening
            if time.time()-timer > WAIT_NOTIFICATION_DELAY:
                timer = time.time()
                self.info('...waiting for %s ' %
                          ' and '.join(['"%s" from ' %v + session_tag(k)
                                        for k, v in self.after.items()]))

        # then connect
        if not self.ssh_connect(self.hostname, self.port,
                                self.username, self.passwd):
            self.ssh_abort()
            return

        # notify all other threads in sync-group
        for tid in self.syncList:
            for t in SSHControl.ssh_threads:
                if t.name == tid:
                    t.notifySync(self.id)

        # wait until all sync-group threads have been connected
        while self.sync:
            time.sleep(LOOP_DELAY)
        # keep execution order
        while (self.after):
            time.sleep(LOOP_DELAY)

        # if simulate only
        if SSHControl.SIMULATE:
            self.command = self.simCommand()

        # execute command string
        self.info("executing: \033[1;34m%s\033[m " % (self.command))
        # escape single quotes in command
        self.command = self.command.replace("'", r"'\''")
        # spawn a new shell session on remote host and pipe command
        # TODO currently we assume that bash is installed
        if SSHControl.ESCAPE:
            self.command = 'echo -e \'' + self.SHELL_SETUP + self.command + '\' | ' + self.SHELL
        else:
            self.command = 'echo \''    + self.SHELL_SETUP + self.command + '\' | ' + self.SHELL

        try:
            # send comand
            self.s.sendline(self.command)

            # start loop to match expect strings and notify other threads
            self.__expectWait()

            # wait for section command to finish (i.e. reach the prompt) and
            # disconnect
            if self.s.prompt(self.PROMPT_TIMEOUT):
                self.ssh_disconnect()
            else:
                self.error("did not reach the prompt! something went wrong")
                print(self.s.before)
                raise SystemError
                # if DEBUG : self.s.interact(local=locals())
        except (select.error, IOError, OSError):
            self.ssh_abort()
            return

    def register_after(self, id, afterCommand):
        """Register an expect string on this block."""
        # FIXME I don't think we really need locks here. appending to a list is
        # an atomic operation
        self.lock.acquire()
        try:
            self.after_list[afterCommand].append(id)
        except KeyError:
            self.after_list[afterCommand] = [id]
        finally:
            self.lock.release()
        self.debug('"%s" registered by %s' % (afterCommand, session_tag(id)))

    def _check_config(self):
        """Very rudimentary check for valid section configuration. Returns True if
        checks pass.

        """
        if self.terminate_threads:
            return False

        if self.after:
            # check for circular refences
            for remote_after in [t for t in SSHControl.ssh_threads if (t.name in list(self.after.keys()))]:
                if (remote_after.after and (self.id in list(remote_after.after.keys()))):
                    self.error("cirular reference %s <-> %s?" %
                               (session_tag(self.id),
                                session_tag(remote_after.id)))
                    raise ConfigError

            # check for invalid section ids in AFTER
            for a in self.after:
                if a not in [t.id for t in SSHControl.ssh_threads]:
                    self.error("'%s' is not a valid section id in AFTER" % a)
                    raise ConfigError

            # check for invalid section ids in SYNC
            if not (self.sync is None):
                for a in self.sync:
                    if a not in [t.id for t in SSHControl.ssh_threads]:
                        self.error("'%s' is not a valid section id in SYNC" % a)
                        self.sync.remove(a)

        self.info("configuration seems ok")
        self.config_ok = True
        return True

    def register_sync(self, id):
        """Register an expect string on this block."""
        self.lockSync.acquire()
        try:
            self.syncList.append(id)
        except KeyError:
            self.syncList = [id]
        finally:
            self.lockSync.release()
        self.debug("synchronized with %s" % session_tag(id))

    def __expectWait(self):
        """Wait until all expect strings matched and remove corresponding entries from
        the threads after: queues.

        """
        try:
            while list(self.after_list.keys()):
                res = self.s.expect(list(self.after_list.keys()), self.PEXPECT_TIMEOUT)
                after = list(self.after_list.keys())[res]
                self.debug('"%s" matched...' % after)

                for tid in self.after_list[after]:
                    for t in SSHControl.ssh_threads:
                        if t.name == tid:
                            t.notifyAfter(self.id, after)

                self.lock.acquire()
                try:
                    del self.after_list[after]
                finally:
                    self.lock.release()

        except pexpect.EOF:
            self.info("EOF!")
        except pexpect.TIMEOUT as e:
            self.info("timed-out waiting for: %s" % (list(self.after_list.keys())))
            self.debug(str(e))

    def __str__(self):
        return session_tag(self.id)

    def info(self, msg, stdout=True):
        with SSHControl.thread_lock:
            id_label = session_tag(self.id).ljust(SSHControl.ID_STR_LEN)
            msg = "%s   %s" % (id_label, msg)
            if stdout:
                print(msg)
            return msg

    def debug(self, msg, stdout=True):
        if not DEBUG:
            return

        with SSHControl.thread_lock:
            id_label = session_tag(self.id).ljust(SSHControl.ID_STR_LEN)
            msg = "%s   %s   %s" % (id_label, msg, DEBUG_LABEL)
            if stdout:
                print(msg)

    def error(self, msg, stdout=True):
        with SSHControl.thread_lock:
            id_label = session_tag(self.id).ljust(SSHControl.ID_STR_LEN)
            msg = "%s   %s" % (id_label, '\033[1;31m'+msg+'\033[0m')
            if stdout:
                print(msg)
            return msg

    def notifyAfter(self, id, after):
        """Notify a thread that the expect string has been matched"""
        self.lock.acquire()  # TODO remove unnecessary locks
        try:
            del self.after[id]
        finally:
            self.lock.release()
        # tell user which match strings are still open
        aremaining = None
        if VERBOSE:
            aremaining = ' and '.join(["\"%s\" from %s" %
                                       (a, session_tag(f)) for
                                       f, a in self.after.items()])
        if aremaining:
            aremaining = " (still waiting for %s)" % aremaining
        self.info("matched \"%s\" from %s%s" % (after, session_tag(id),
                                                aremaining))

    def notifySync(self, id):
        """Notify a thread to start in sync mode"""
        self.info("SYNC: notified from %s" % (session_tag(id)))
        self.lockSync.acquire()
        try:
            self.sync.remove(id)
            self.syncNotified = True
        finally:
            self.lockSync.release()

    def ssh_connect(self, hostname, port, username, passwd):
        """Connect and return True if successfull. Retry if connection fails.

        """
        sshport = int(port)
        self.info("connecting to %s:%s ... " % (ansi_bold(self.hostname),
                                                ansi_bold(self.port)))
        if self.s is None:
            self.s = pxssh.pxssh()
        try:
            # p.hollands suggestion:
            # original_prompt=r"][#$]|~[#$]|bash.*?[#$]|[#$] |.*@.*:.*>" BROKEN
            # original_prompt=r"[#$]|$"
            self.s.login(hostname, username, passwd,
                         login_timeout=self.SSH_LOGIN_TIMEOUT,
                         port=sshport, auto_prompt_reset=True)
            self.s.set_unique_prompt

            # TODO: disabling echo currently only works reliably under bash, so
            # also trying this way
            self.s.setecho(False)
            self.s.sendline('stty -echo;')
            # match the prompt within X seconds
            if not self.s.prompt(self.PROMPT_TIMEOUT):
                print("could not match the prompt!")
                return False

            if not self.s.isalive():
                return False

            self.connected = True
            self.info("connected")

            if DEBUG:
                    self.debug("writing log file to %s" %
                               ansi_bold(self.id+'.log'))
                    fout = open('%s.log' % (self.id), 'wb')
                    self.s.logfile = fout

        except pxssh.ExceptionPxssh as e:
            if e.message == 'password refused':
                self.error('invalid password for %s@%s! exiting...' %
                           (self.username, self.hostname))
                self.ssh_abort(e)
                return False

            self.debug(str(e))

            self.info("SSH login to %s FAILED... retrying in %ss" %
                      (self.hostname, self.SSH_LOGIN_REPEAT_TIMEOUT))

            time.sleep(self.SSH_LOGIN_REPEAT_TIMEOUT)
            self.s.close()
            self.s = None
            self.ssh_connect(self.hostname, self.port,
                             self.username, self.passwd)
        except (select.error, IOError, OSError) as e:
            self.ssh_abort(e)
            return False
        except (pexpect.EOF) as e:
            self.error("could not connect to %s. check hostname" %
                       ansi_bold(self.hostname))
            self.ssh_abort(e)
            return False
        return True

    def ssh_abort(self, e=None):
        SSHControl.terminate_threads = True

    def remove(self):
        SSHControl.ssh_threads.remove(self)

    def ssh_disconnect(self):
        if self.s:
            try:
                self.s.close()
            except OSError:
                pass
            self.s = None
            self.info("disconnected")
        self.connected = False
        self.remove()

    def simCommand(self):
        """TODO @CK"""
        result = ""
        for t in SSHControl.ssh_threads:
            if not (t.after is None):
                if self.id in list(t.after.keys()):
                    result = result + "echo \"" + t.after.get(self.id) + "\"; "
        return result