diff --git a/copytonodes b/copytonodes new file mode 100755 index 0000000..20e467d --- /dev/null +++ b/copytonodes @@ -0,0 +1,43 @@ +#!/usr/bin/env python + +import argparse +import getpass +from noderange import expand +import nodeconnection +import os +import paramiko +from paramiko import SSHException +import sys + + +def copytonodes(nodespec, src_file=None, dest_file=None, dshbak=False, + verbose=False, user=None, password=None): + if src_file is None or dest_file is None: + print "copytonodes, src_file or dest_file is None" + sys.exit(1) + + fd = os.open(src_file, os.O_RDONLY) + file_txt = os.read(fd, 1024) + cmd = "/bin/cat << EOF > " + dest_file + " " + file_txt + "\nEOF" + nodeconnection.runonnodes(nodespec, cmd, dshbak, verbose, user) + cmd = "/bin/chmod +x " + dest_file + nodeconnection.runonnodes(nodespec, cmd, dshbak, verbose, user) + +def main(): + parser = argparse.ArgumentParser( + description="Run a command on a set of nodes") + parser.add_argument('-s', action="store", dest="src_file") + parser.add_argument('-d', action="store", dest="dest_file") + parser.add_argument('-t', action="store_true", default=False) + parser.add_argument('-v', action="store_true", default=False) + parser.add_argument('-w', action="store", dest="where") + parser.add_argument('-u', action="store", dest="user") + parser.add_argument('-p', action="store", dest="password") + args = parser.parse_args(sys.argv[1:]) + + copytonodes(args.where, args.src_file, args.dest_file, + dshbak=args.t, verbose=args.v, + user=args.user, password=args.password) + +if __name__ == "__main__": + main() diff --git a/nodeconnection.py b/nodeconnection.py new file mode 100755 index 0000000..c76d3c3 --- /dev/null +++ b/nodeconnection.py @@ -0,0 +1,117 @@ +#!/usr/bin/env python + +import argparse +import getpass +from noderange import expand +import io +import os +import paramiko +from paramiko import SSHException +import string +import select +import sys + +class NodeConnection(object): + _password = None + def __init__(self, name, user=None, password=None): + self.name = name + self.ssh = None + self.stdin = None + self.stdout = None + self.stderr = None + if user: + self._user = user + else: + self._user = getpass.getuser() + if not password is None: + NodeConnection._password = password + + def connect(self, verbose=False, port=22): + if verbose: + print "Connection to host '%s'" % (self.name) + self.ssh = paramiko.SSHClient() + self.ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy()) + try: + self.ssh.connect(self.name, 22, self._user, + password=NodeConnection._password) + except SSHException as e: + #Try again with a password + NodeConnection._password = \ + getpass.getpass("Key based logon did not work, please " + "provide password for %s@%s: " % + (self._user, self.name), stream=sys.stderr) + self.ssh.connect(self.name, 22, self._user, + password=NodeConnection._password) + + def exec_command(self, cmd): + if not self.ssh: + raise Exception("exec_command: not connected") + self.stdin, self.stdout, self.stderr = self.ssh.exec_command(cmd) + + def exec_sudo_command(self, cmd): + self.stdout = io.StringIO() + self.stdin = io.StringIO() + self.stderr = io.StringIO() + chan = self.ssh.get_transport().open_session() + chan.get_pty() + chan.exec_command(cmd) + prompt = chan.recv(1024) + if string.find(prompt, "password") != -1: + if NodeConnection._password is None: + NodeConnection._password = getpass.getpass( + "Sudo access requires a password, please" + " provide password for %s@%s: " % + (self._user, self.name), stream=sys.stderr) + plen = chan.send(NodeConnection._password + '\n') + data = chan.recv(1024) + while chan.exit_status_ready() != True: + rl, wl, xl = select.select([chan],[],[],0.0) + data += chan.recv(1024) + self.stdout = io.StringIO(unicode(data)) + else: + self.stdout = io.StringIO(unicode(prompt)) + + def print_output(self, leader='', output=False): + #FIXME: dshbak (leader) mode nott yet implemented correctly due to + # readline() returning a character at a time rather than a line. + outbuf = [] + if not self.stdin or not self.stdout or not self.stderr: + raise Exception("print_output: not connected") + for line in self.stdout.readlines(): + if output is True: + outbuf.append(line) + else: + sys.stdout.write(line) + sys.stdout.flush() + for line in self.stderr.readlines(): + if output is True: + outbuf.append(line) + else: + sys.stderr.write(line) + sys.stderr.flush() + if output: + return outbuf + +def runonnodes(nodespec, cmd, dshbak=False, verbose=False, + user=None, password=None, output=False, + sudo=False): + nodes = expand(nodespec) + + if len(nodes) == 0: + print "Need at least one node to run on" + sys.exit(1) + + for node in nodes: + nc = NodeConnection(node, user, password) + nc.connect(verbose=verbose) + if sudo is False: + nc.exec_command(cmd) + else: + nc.exec_sudo_command(cmd) + if not dshbak: + print "--------------- %s ---------------" % (node) + outbuf = nc.print_output(output=output) + else: + outbuf = nc.print_output(str(node) + ": ", output=output) + if output == True: + return outbuf diff --git a/noderange.py b/noderange.py index c844a95..f6a621e 100644 --- a/noderange.py +++ b/noderange.py @@ -9,6 +9,9 @@ def expand(range): compute[1-3], return a list of nodes """ nodes = [] + if range is None or len(range) < 1: + print "Invalid range, empty" + return [] tidyrange = string.replace(range, " ", "").strip() diff --git a/runonnodes b/runonnodes index a00fb64..40ac0e1 100755 --- a/runonnodes +++ b/runonnodes @@ -1,80 +1,31 @@ #!/usr/bin/env python +import argparse +import getpass from noderange import expand -import sys +import nodeconnection +import os import paramiko from paramiko import SSHException -import getpass -import argparse - -def runonnodes(nodespec, cmd, dshbak=False, verbose=False, user=None): - nodes = expand(nodespec) - - if len(nodes) == 0: - print "Need at least one node to run on" - sys.exit(1) - - for node in nodes: - nc = NodeConnection(node, user) - nc.connect(verbose=verbose) - nc.exec_command(cmd) - if not dshbak: - print "--------------- %s ---------------" % (node) - nc.print_output() - else: - nc.print_output(str(node) + ": ") - - -class NodeConnection(object): - def __init__(self, name, user=None): - self.name = name - self.ssh = None - self.stdin = None - self.stdout = None - self.stderr = None - if user: - self.user = user - else: - self.user = getpass.getuser() - - def connect(self, verbose=False, port=22): - if verbose: - print "Connection to host '%s'" % (self.name) - self.ssh = paramiko.SSHClient() - self.ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy()) - try: - self.ssh.connect(self.name, 22, self.user) - except SSHException as e: - #Try again with a password - self.ssh.connect(self.name, 22, self.user, password=getpass.getpass("Key based logon did not work, please provide password for %s@%s: " % (self.user, self.name), stream=sys.stderr,)) - - def exec_command(self, cmd): - if not self.ssh: - raise Exception("exec_command: not connected") - self.stdin, self.stdout, self.stderr = self.ssh.exec_command(cmd) +import sys - def print_output(self, leader=''): - #FIXME: dshbak (leader) mode nott yet implemented correctly due to - # readline() returning a character at a time rather than a line. - if not self.stdin or not self.stdout or not self.stderr: - raise Exception("print_output: not connected") - for line in self.stdout.readline(): - sys.stdout.write(line) - sys.stdout.flush() - for line in self.stderr.readline(): - sys.stderr.write(line) - sys.stderr.flush() def main(): - parser = argparse.ArgumentParser(description="Run a command on a set of nodes") + parser = argparse.ArgumentParser( + description="Run a command on a set of nodes") parser.add_argument('-t', action="store_true", default=False) parser.add_argument('-v', action="store_true", default=False) + parser.add_argument('-s', action="store_true", default=False) parser.add_argument('-w', action="store", dest="where") parser.add_argument('-u', action="store", dest="user") + parser.add_argument('-p', action="store", dest="password") parser.add_argument('command', nargs='*', action="store") args = parser.parse_args(sys.argv[1:]) - runonnodes(args.where, " ".join(args.command), dshbak=args.t, verbose=args.v, user=args.user) + nodeconnection.runonnodes(args.where, " ".join(args.command), + dshbak=args.t, verbose=args.v, + user=args.user, password=args.password, + sudo=args.s) if __name__ == "__main__": main()