| # 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 __future__ import with_statement |
| import os, sys, subprocess |
| |
| sys.path.append(sys.path[0] + '/lib') |
| import cmd, signal, atexit, threading |
| import pxssh, pexpect |
| |
| #disable deprecationwarnings |
| import warnings |
| warnings.simplefilter('ignore', DeprecationWarning) |
| |
| |
| progname = os.path.split(sys.argv[0])[-1] |
| |
| class VersionError(Exception): |
| pass |
| |
| if sys.version_info < (2, 5, 0): |
| raise VersionError( |
| '''Error: %s is supported on Python versions 2.5 or greater |
| Please upgrade python installed on this machine.''' % progname) |
| |
| |
| ################ |
| def cli_help(): |
| help_path = os.path.join(sys.path[0], '..', 'docs', 'cli_help', progname + '_help'); |
| f = None |
| try: |
| try: |
| f = open(help_path); |
| return f.read(-1) |
| except: |
| return '' |
| finally: |
| if f: f.close() |
| |
| |
| ################ |
| def parse_port(val): |
| try: return int(val) |
| except ValueError: return 0 |
| |
| |
| ############# |
| def ssh_prefix(host): |
| ssh = 'ssh -o "BatchMode yes" -o "StrictHostKeyChecking no" ' + host |
| return ssh |
| |
| |
| ############# |
| def scp_prefix(): |
| scp = 'scp -o "BatchMode yes" -o "StrictHostKeyChecking no"' |
| return scp |
| |
| ############# |
| class Error(Exception): |
| '''Base class for exceptions in this module.''' |
| pass |
| |
| |
| ############# |
| class HostNameError(Error): |
| def __init__(self, msg, lineno = 0): |
| if lineno: self.msg = ('%s at line %d' % (value, lineno)) |
| else: self.msg = msg |
| def __str__(self): |
| return self.msg |
| |
| |
| ############# |
| class SSHError(Error): |
| def __init__(self, msg): |
| self.msg = msg |
| def __str__(self): |
| return self.msg |
| |
| ############# |
| def ssh(host, cmd): |
| s = 'ssh -o "BatchMode yes" -o "StrictHostKeyChecking no" %s %s 2>& 1' % (host, cmd) |
| p = os.popen(s) |
| if not p: |
| raise SSHError('unable to ssh to %s' % host) |
| yield p |
| out = p.read(-1) |
| rc = p.close() or 0 |
| yield (rc, out) |
| |
| |
| ############# |
| class HostList(): |
| |
| def __init__(self): |
| self.list = [] |
| |
| def get(self): |
| return self.list |
| |
| def add(self, host): |
| '''Add a host to the hostlist.''' |
| if host.find(':') >= 0 or host.find('@') >= 0: |
| raise HostNameError(host) |
| self.list.append(host) |
| return self.list |
| |
| def parseFile(self, path): |
| '''Add lines in a file to the hostlist.''' |
| with open(path) as fp: |
| lineno = 0 |
| for line in fp: |
| lineno = lineno + 1 |
| line = line.strip() |
| if not line: continue; |
| if line[0] == '#': continue |
| if not self.add(line): |
| raise HostNameError(line, lineno) |
| return self.list |
| |
| def batch(self, size=32): |
| '''Return list of hostlists, each is of 'size' elements''' |
| cnt = int((len(self.list) - 1) / size) + 1 |
| b = [] |
| for i in xrange(cnt): |
| start = i * size |
| end = start + size |
| if end > len(self.list): |
| b.append(self.list[start:]) |
| else: |
| b.append(self.list[start:end]) |
| return b |
| |
| def checkSSH(self): |
| '''Check that ssh to hostlist is okay.''' |
| |
| #TODO: Eliminate this function - it's redundant to the Session class |
| |
| for hlist in self.batch(): |
| slist = [] |
| for h in hlist: |
| s = ssh(h, 'echo') |
| s.next() |
| slist.append((h, s)) |
| for (h, s) in slist: |
| (rc, out) = s.next() |
| if out and out.find('Permission denied') >= 0: |
| raise SSHError("Permission denied while ssh to '%s'" % h) |
| if rc: |
| raise SSHError("Unable to ssh to '%s'" % h) |
| return True |
| |
| def filterMultiHomedHosts(self): |
| '''For multiple host that is of the same node, keep only one in the hostlist.''' |
| unique = {} |
| #TODO: use a Session object to do this work so that it runs in parallel |
| |
| for hlist in self.batch(): |
| slist = [] |
| for h in hlist: |
| s = ssh(h, 'hostname') |
| s.next() |
| slist.append((h, s)) |
| for (h, s) in slist: |
| (rc, hname) = s.next() |
| if rc: |
| raise SSHError('Cannot execute hostname on host %s' % h) |
| hname = hname.strip() |
| if not unique.get(hname): |
| unique[hname] = h |
| elif hname == h: |
| unique[hname] = h |
| |
| self.list = unique.values() |
| return self.list |
| |
| ############# |
| # Session is a command session, derived from a base class cmd.Cmd |
| class Session(cmd.Cmd): |
| '''Implements a list of open ssh sessions ready to execute commands''' |
| verbose=False |
| hostList=[] |
| userName=None |
| class SessionError(StandardError): pass |
| class SessionCmdExit(StandardError): pass |
| |
| def __init__(self, hostList=None, userName=None): |
| cmd.Cmd.__init__(self) |
| self.pxssh = [] |
| self.prompt = '=> ' |
| self.peerStringFormatRaw = None |
| if hostList: |
| for host in hostList: |
| self.hostList.append(host) |
| if userName: self.userName=userName |
| |
| def peerStringFormat(self): |
| if self.peerStringFormatRaw: return self.peerStringFormatRaw |
| cnt = 0 |
| for p in self.pxssh: |
| if cnt < len(p.x_peer): cnt = len(p.x_peer) |
| self.peerStringFormatRaw = "[%%%ds]" % cnt |
| return self.peerStringFormatRaw |
| |
| def login(self, hostList=None, userName=None): |
| '''This is the normal entry point used to add host names to the object and log in to each of them''' |
| if self.verbose: print '\n[Reset ...]' |
| if not (self.hostList or hostList): |
| raise self.SessionError('No host list available to Login method') |
| if not (self.userName or userName): |
| raise self.SessionError('No user name available to Login method') |
| |
| #Cleanup |
| self.clean() |
| |
| if hostList: #We have a new hostlist to use, initialize it |
| self.hostList=[] |
| for host in hostList: |
| self.hostList.append(host) |
| if userName: self.userName=userName #We have a new userName to use |
| |
| # MPP-6583. Save off term type and set to nothing before creating ssh process |
| origTERM = os.getenv('TERM', None) |
| os.putenv('TERM', '') |
| |
| for host in hostList: |
| self.hostList.append(host) |
| p = pxssh.pxssh() |
| p.loginAsync(host, self.userName) |
| p.x_peer = host |
| p.x_pid = p.pid |
| self.pxssh.append(p) |
| |
| # Restore terminal type |
| if origTERM: |
| os.putenv('TERM', origTERM) |
| |
| for p in self.pxssh: |
| if self.verbose: print '[INFO] login %s' % p.x_peer |
| if not p.loginWait(set_term_dumb=True): |
| print '[ERROR] unable to login to %s' % p.x_peer |
| print |
| print 'hint: use gpssh-exkeys to setup public-key authentication between hosts' |
| print |
| raise self.SessionError('Error logging into host %s' % p.x_peer) |
| |
| def close(self): |
| self.clean() |
| |
| def reset(self): |
| '''reads from all the ssh connections to make sure we dont have any pending cruft''' |
| for s in self.pxssh: |
| s.readlines() |
| |
| def clean(self): |
| self.closePxsshList(self.pxssh) |
| self.pxssh = [] |
| |
| def emptyline(self): |
| pass |
| |
| def escapeLine(self,line): |
| '''Escape occurrences of \ and $ as needed and package the line as an "eval" shell command''' |
| line = line.strip() |
| if line == 'EOF' or line == 'exit' or line == 'quit': |
| raise self.SessionCmdExit() |
| line = line.split('\\') |
| line = '\\\\'.join(line) |
| line = line.split('"') |
| line = '\\"'.join(line) |
| line = line.split('$') |
| line = '\\$'.join(line) |
| line = 'eval "' + line + '" < /dev/null' |
| |
| return line |
| |
| def executeCommand(self,command): |
| commandoutput=[] |
| |
| #Execute the command in all of the ssh sessions |
| for s in self.pxssh: |
| s.sendline(command) |
| s.flush() |
| |
| #Wait for each command and retrieve the output |
| for s in self.pxssh: |
| #Wait for each command to finish |
| #!! TODO verify that this is a tight wait loop and find another way to do this |
| while not s.prompt(120) and s.isalive() and not s.eof(): pass |
| |
| for s in self.pxssh: |
| #Split the output into an array of lines so that we can add text to the beginning of |
| # each line |
| output = s.before.split('\n') |
| output = output[1:-1] |
| |
| commandoutput.append(output) |
| |
| return commandoutput.__iter__() |
| |
| # Interactive command line handler |
| # Override of base class, handles commands that aren't recognized as part of a predefined set |
| # The "command" argument is a command line to be executed on all available command sessions |
| # The output of the command execution is printed to the standard output, prepended with |
| # the hostname of each machine from which the output came |
| def default(self, command): |
| |
| line = self.escapeLine(command) |
| |
| if self.verbose: print command |
| |
| #Execute the command on our ssh sessions |
| commandoutput=self.executeCommand(command) |
| self.writeCommandOutput(commandoutput) |
| |
| def writeCommandOutput(self,commandoutput): |
| '''Takes a list of output lists as an iterator and writes them to standard output, |
| formatted with the hostname from which each output array was obtained''' |
| for s in self.pxssh: |
| output = commandoutput.next() |
| #Write the output |
| if len(output) == 0: |
| print (self.peerStringFormat() % s.x_peer) |
| else: |
| for line in output: |
| print (self.peerStringFormat() % s.x_peer), line |
| |
| def closePxsshList(self,list): |
| def closePxsshOne(p): p.close(force = True) |
| th = [] |
| for p in list: |
| t = threading.Thread(target=closePxsshOne, args=(p,)) |
| t.start() |
| th.append(t) |
| for t in th: |
| t.join() |
