xssh

#!/usr/bin/env python3

from pssh.pssh_client import ParallelSSHClient
import pssh.utils

pssh.utils.enable_host_logger()

import os, sys, re, stat, tempfile, subprocess

import logging
import yaml
import socket
from gevent import joinall

log=None

def log_config(verbose):
    FORMAT = '[%(asctime)s] {%(pathname)s:%(lineno)d} %(levelname)s - %(message)s'
    logging.basicConfig(format=FORMAT)
    logging.getLogger("pssh.host_logger").setLevel(logging.CRITICAL)
    logging.getLogger("pssh.ssh_client").setLevel(logging.ERROR)

    global log
    log = logging.getLogger("xssh")

    if verbose==1:
        log.setLevel(logging.INFO)
    elif verbose>=2:
        log.setLevel(logging.DEBUG)


def parse_args():
    import argparse

    isSCP = (re.search("scp",sys.argv[0]))

    if isSCP:
        parser = argparse.ArgumentParser(description='Copy files around the cluster')
    else:
        parser = argparse.ArgumentParser(description='Execute a command on cluster')

    parser.add_argument('-R', '--no-prefix', help="Disable host prefix on output", action="store_true")
    parser.add_argument('-c', '--conf', help="Config file", default="/etc/xssh.conf")
    parser.add_argument('-t', '--to', help="One ore more hosts", action="append", default=[])
    parser.add_argument('-v', '--verbose', help="Verbosity", action="count", default=0)
    parser.add_argument('-S', '--sudo', help="Run with sudo", action="store_true")

    if isSCP:
        log.debug({"hosts":hosts,"config":pssh_config, "file": args.file})
        path = args.file[0]
        path=os.path.abspath(path)

        from pathlib import Path

        if args.sudo:
            try:
                command = ["sudo", "stat", "-c", "%a/%F/%U", path]
                log.debug("copy local {}".format(command))
                res = subprocess.check_output(command, encoding='utf-8')
                res = res.rstrip("\n").split("/")
                mode = int(res[0],8)
                isdir = res[1] == "directory"
                user = res[2]
            except:
                log.error("failed to stat {}".format(command))
                raise
        else:
            pobj = Path(path)
            user = pobj.owner()
            mode = stat.S_IMODE(pobj.stat().st_mode)
            isdir = pobj.is_dir()

        if args.remote_dest:
            remote_dest = args.remote_dest
        else:
            remote_dest = path


        exit = 0

        if args.sudo:
            import getpass

            if isdir:
                tmppath = tempfile.mkdtemp(prefix="xscp.", suffix="." + os.path.basename(remote_dest))
                os.rmdir(tmppath)
            else:
                tmppath = tempfile.NamedTemporaryFile(prefix="xscp.", suffix="." + os.path.basename(remote_dest)).name

            command = ["sudo", "cp", "-r", path, tmppath]
            log.debug("copy local {}".format(command))
            exit = os.spawnv(os.P_WAIT, "/usr/bin/sudo", command)
            if exit != 0:
                log.error("failed to copy {}".format(command))

            command = ["sudo", "chown", "-R", getpass.getuser(), tmppath]
            exit = os.spawnv(os.P_WAIT, "/usr/bin/sudo", command)
            if exit != 0:
                log.error("failed to chown {}".format(command))

            path = tmppath

        if exit == 0:
            print("copying " + path + " to " + remote_dest + " on all hosts", file=sys.stderr)

            copy_path = tempfile.NamedTemporaryFile(prefix="xscp.", suffix="." + os.path.basename(remote_dest)).name

            result = client.copy_file(path, copy_path, recurse=True)
            result = joinall(result)

            for g in result:
                if g.exception:
                    log.error("Failed to copy {}".format(g.exception))
                    exit = 7

        if exit == 0:
            args.command = ["chmod", oct(mode)[2:], copy_path]
            exit=run_command(client, aliases, args)

        if exit == 0 and args.sudo:
            args.command = ["chown", user, copy_path]
            exit=run_command(client, aliases, args)

        if exit == 0:
            args.command = ["mv", "-f", "-b", copy_path, remote_dest]
            exit=run_command(client, aliases, args)

    else:
        log.debug({"hosts":hosts,"config":pssh_config, "command": args.command})
        exit=run_command(client, aliases, args)

    sys.exit(exit)

def run_command(client, aliases, args):
    command = ""
    for arg in args.command:
        command += '"' + arg.replace('"','\\"') + '" '

    log.debug("run_command: {}".format(command))
    output = client.run_command(command, stop_on_errors=False, sudo=args.sudo)

    maxlen = 0
    for host, host_output in output.items():
        show_host = aliases.get(host, host)
        maxlen=max(maxlen,len(show_host))

    exit=0
    for host, host_output in output.items():
        show_host = aliases.get(host, host)
        padlen = maxlen-len(show_host)
        pad = " " * padlen
        if args.no_prefix:
            printline = lambda line, stream: print(line, file=stream)
        else:
            printline = lambda line, stream: print(show_host, pad, line, file=stream)


        if host_output.stdout:
            for line in host_output.stdout:
                printline(line, sys.stdout)
        if host_output.stderr:
            for line in host_output.stderr:
                printline(line, sys.stderr)
        if host_output.exception:
            printline(str(host_output.exception), sys.stderr)

        if host_output.exit_code:
            exit=host_output.exit_code
    return exit
if __name__ == "__main__":
    main()