testutils/ptest/Ssh.py

# Licensed to the Apache Software Foundation (ASF) under one or more
# contributor license agreements.  See the NOTICE file distributed with
# this work for additional information regarding copyright ownership.
# The ASF licenses this file to You under the Apache License, Version 2.0
# (the "License"); you may not use this file except in compliance with
# the License.  You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from threading import Thread
from Queue import Queue

import Process

class SSHConnection():
    def __init__(self, host, num = None):
        self.host = host
        if num is None:
            self.hostname = host
        else:
            self.hostname = host + '-' + str(num);
        self.pwd = '/'
        self.env = {}
        self.path = []

    def cd(self, path):
        self.pwd = path.format(host = self.hostname)

    def export(self, env, value):
        self.env[env] = value.format(host = self.hostname)

    def add_path(self, path):
        self.path.append(path.format(host = self.hostname))

    def prefix(self, cmd):
        pre = []
        pre.append('cd "{0}"'.format(self.pwd))
        for (e, v) in self.env.iteritems():
            pre.append('export {0}=\\"{1}\\"'.format(e, v))
        for p in self.path:
            pre.append('export PATH="{0}:${{PATH}}"'.format(p))
        pre.append(cmd)
        return ' && '.join(pre)

    def run(self, cmd, warn_only = False, quiet = False, vewy_quiet = False,
            abandon_output = True):
        # Don't use single quotes in `cmd`, this will break and end badly.
        cmd = cmd.format(host = self.hostname)
        cmd = self.prefix(cmd)
        print(self.hostname + ' =>')
        if vewy_quiet:
            # Be vewy, vewy quiet, I'm hunting wabbits.
            print('[command hidden]\n')
            quiet = True
        else:
            print(cmd + '\n')
        cmd = '''ssh '{0}' "bash -c '{1}'"'''.format(self.host, cmd)
        try:
            return Process.run(cmd, quiet, abandon_output)
        except Exception as e:
            if warn_only:
                print(str(e) + '---------- This was only a warning, ' +
                        'it won\'t stop the execution --\n')
                return None
            else:
                raise e

class SSHSet():
    def __init__(self, conn = []):
        self.conn = conn

    def __len__(self):
        return len(self.conn)

    def add(self, conn):
        if isinstance(conn, list):
            self.conn.extend(conn)
        else:
            self.conn.append(conn)

    def cd(self, path):
        for conn in self.conn:
            conn.cd(path)

    def export(self, env, value):
        for conn in self.conn:
            conn.export(env, value)

    def add_path(self, path):
        for conn in self.conn:
            conn.add_path(path)

    def run(self, cmd, parallel = True, quiet = False, vewy_quiet = False,
            abandon_output = True, warn_only = False):
        if not parallel:
            for conn in self.conn:
                conn.run(cmd, quiet = quiet, vewy_quiet = vewy_quiet,
                        abandon_output = abandon_output, warn_only = warn_only)
        else:
            threads = []
            queue = Queue()
            def wrapper(conn, cmd, queue):
                try:
                    conn.run(cmd, quiet = quiet, vewy_quiet = vewy_quiet,
                            abandon_output = abandon_output,
                            warn_only = warn_only)
                except Exception as e:
                    queue.put(Exception(conn.hostname + ' => ' + str(e)))
            for conn in self.conn:
                thread = Thread(target = wrapper, args = (conn, cmd, queue, ))
                thread.start()
                threads.append(thread)
            for thread in threads:
                thread.join()
            if not queue.empty():
                l = []
                while not queue.empty():
                    e = queue.get()
                    l.append(str(e));
                raise Exception('\n'.join(l))