starcluster.ssh

Package starcluster :: Module ssh |
[hide private]
[frames] | no frames]

# Source Code for Module starcluster.ssh

      1  """
      2  ssh.py
      3  Friendly Python SSH2 interface.
      4  From http://commandline.org.uk/code/
      5  License: LGPL
      6  modified by justin riley (justin.t.riley@gmail.com)
      7  """
      8
      9  import os
     10  import re
     11  import sys
     12  import stat
     13  import string
     14  import socket
     15  import paramiko
     16  import posixpath
     17
     18  # windows does not have termios...
     19  try:
     20      import termios
     21      import tty
     22      HAS_TERMIOS = True
     23  except ImportError:
     24      HAS_TERMIOS = False
     25
     26  from starcluster import exception
     27  from starcluster.logger import log
     28
     29
     30 -class SSHClient(object):
     31      """
     32      Establishes an SSH connection to a remote host using either password or
     33      private key authentication. Once established, this object allows executing
     34      commands, copying files to/from the remote host, various file querying
     35      similar to os.path.*, and much more.
     36      """
     37
     38 -    def __init__(self,
     39                   host,
     40                   username=None,
     41                   password=None,
     42                   private_key=None,
     43                   private_key_pass=None,
     44                   port=22,
     45                   timeout=30):
     46          self._host = host
     47          self._port = 22
     48          self._pkey = None
     49          self._username = username or os.environ['LOGNAME']
     50          self._password = password
     51          self._timeout = timeout
     52          self._sftp = None
     53          self._transport = None
     54          if private_key:
     55              self._pkey = self.load_private_key(private_key, private_key_pass)
     56          elif not password:
     57              raise exception.SSHNoCredentialsError()
     58
     59 -    def load_private_key(self, private_key, private_key_pass=None):
     60          # Use Private Key.
     61          log.debug('loading private key %s' % private_key)
     62          if private_key.endswith('rsa') or private_key.count('rsa'):
     63              pkey = self._load_rsa_key(private_key, private_key_pass)
     64          elif private_key.endswith('dsa') or private_key.count('dsa'):
     65              pkey = self._load_dsa_key(private_key, private_key_pass)
     66          else:
     67              log.debug("specified key does not end in either rsa or dsa" + \
     68                        ", trying both")
     69              pkey = self._load_rsa_key(private_key, private_key_pass)
     70              if pkey is None:
     71                  pkey = self._load_dsa_key(private_key, private_key_pass)
     72          return pkey
     73
     74 -    def connect(self, host=None, username=None, password=None,
     75                  private_key=None, private_key_pass=None, port=22, timeout=30):
     76          host = host or self._host
     77          username = username or self._username
     78          pkey = self._pkey
     79          if private_key:
     80              pkey = self.load_private_key(private_key, private_key_pass)
     81          log.debug("connecting to host %s on port %d as user %s" % (host, port,
     82                                                                     username))
     83          try:
     84              sock = self._get_socket(host, port)
     85              transport = paramiko.Transport(sock)
     86              transport.banner_timeout = timeout
     87          except socket.error:
     88              raise exception.SSHConnectionError(host, port)
     89          # Authenticate the transport.
     90          try:
     91              transport.connect(username=username, pkey=pkey, password=password)
     92          except paramiko.AuthenticationException:
     93              raise exception.SSHAuthException(username, host)
     94          except paramiko.SSHException, e:
     95              msg = e.args[0]
     96              raise exception.SSHError(msg)
     97          except socket.error:
     98              raise exception.SSHConnectionError(host, port)
     99          except EOFError:
    100              raise exception.SSHConnectionError(host, port)
    101          except Exception, e:
    102              raise exception.SSHError(str(e))
    103          self.close()
    104          self._transport = transport
    105          return self
    106
    107      @property
    108 -    def transport(self):
    109          """
    110          This property attempts to return an active SSH transport
    111          """
    112          if not self._transport or not self._transport.is_active():
    113              self.connect(self._host, self._username, self._password,
    114                           port=self._port, timeout=self._timeout)
    115          return self._transport
    116
    117 -    def get_server_public_key(self):
    118          return self.transport.get_remote_server_key()
    119
    120 -    def is_active(self):
    121          if self._transport:
    122              return self._transport.is_active()
    123          return False
    124
    125 -    def _get_socket(self, hostname, port):
    126          for (family, socktype, proto, canonname, sockaddr) in \
    127          socket.getaddrinfo(hostname, port, socket.AF_UNSPEC,
    128                             socket.SOCK_STREAM):
    129              if socktype == socket.SOCK_STREAM:
    130                  af = family
    131                  break
    132              else:
    133                  raise exception.SSHError(
    134                      'No suitable address family for %s' % hostname)
    135          sock = socket.socket(af, socket.SOCK_STREAM)
    136          sock.settimeout(self._timeout)
    137          sock.connect((hostname, port))
    138          return sock
    139
    140 -    def _load_rsa_key(self, private_key, private_key_pass=None):
    141          private_key_file = os.path.expanduser(private_key)
    142          try:
    143              rsa_key = paramiko.RSAKey.from_private_key_file(private_key_file,
    144                                                              private_key_pass)
    145              log.debug("Using private key %s (rsa)" % private_key)
    146              return rsa_key
    147          except paramiko.SSHException:
    148              log.error('invalid rsa key or passphrase specified')
    149
    150 -    def _load_dsa_key(self, private_key, private_key_pass=None):
    151          private_key_file = os.path.expanduser(private_key)
    152          try:
    153              dsa_key = paramiko.DSSKey.from_private_key_file(private_key_file,
    154                                                              private_key_pass)
    155              log.info("Using private key %s (dsa)" % private_key)
    156              return dsa_key
    157          except paramiko.SSHException:
    158              log.error('invalid dsa key or passphrase specified')
    159
    160      @property
    161 -    def sftp(self):
    162          """Establish the SFTP connection."""
    163          if not self._sftp or self._sftp.sock.closed:
    164              log.debug("creating sftp connection")
    165              self._sftp = paramiko.SFTPClient.from_transport(self.transport)
    166          return self._sftp
    167
    168 -    def generate_rsa_key(self):
    169          return paramiko.RSAKey.generate(2048)
    170
    171 -    def get_public_key(self, key):
    172          return ' '.join([key.get_name(), key.get_base64()])
    173
    174 -    def load_remote_rsa_key(self, remote_filename):
    175          """
    176          Returns paramiko.RSAKey object for an RSA key located on the remote
    177          machine
    178          """
    179          rfile = self.remote_file(remote_filename, 'r')
    180          key = paramiko.RSAKey(file_obj=rfile)
    181          rfile.close()
    182          return key
    183
    184 -    def makedirs(self, path, mode=0755):
    185          """
    186          Same as os.makedirs - makes a new directory and automatically creates
    187          all parent directories if they do not exist.
    188
    189          mode specifies unix permissions to apply to the new dir
    190          """
    191          head, tail = posixpath.split(path)
    192          if not tail:
    193              head, tail = posixpath.split(head)
    194          if head and tail and not self.path_exists(head):
    195              try:
    196                  self.makedirs(head, mode)
    197              except OSError, e:
    198                  # be happy if someone already created the path
    199                  if e.errno != os.errno.EEXIST:
    200                      raise
    201              # xxx/newdir/. exists if xxx/newdir exists
    202              if tail == posixpath.curdir:
    203                  return
    204          self.mkdir(path, mode)
    205
    206 -    def mkdir(self, path, mode=0755, ignore_failure=False):
    207          """
    208          Make a new directory on the remote machine
    209
    210          If parent is True, create all parent directories that do not exist
    211
    212          mode specifies unix permissions to apply to the new dir
    213          """
    214          try:
    215              return self.sftp.mkdir(path, mode)
    216          except IOError:
    217              if not ignore_failure:
    218                  raise
    219
    220 -    def get_remote_file_lines(self, remote_file, regex=None, matching=True):
    221          """
    222          Returns list of lines in a remote_file
    223
    224          If regex is passed only lines that contain a pattern that matches
    225          regex will be returned
    226
    227          If matching is set to False then only lines *not* containing a pattern
    228          that matches regex will be returned
    229          """
    230          f = self.remote_file(remote_file, 'r')
    231          flines = f.readlines()
    232          f.close()
    233          if regex is None:
    234              return flines
    235          r = re.compile(regex)
    236          lines = []
    237          for line in flines:
    238              match = r.search(line)
    239              if matching and match:
    240                  lines.append(line)
    241              elif not matching and not match:
    242                  lines.append(line)
    243          return lines
    244
    245 -    def remove_lines_from_file(self, remote_file, regex):
    246          """
    247          Removes lines matching regex from remote_file
    248          """
    249          if regex in [None, '']:
    250              log.debug('no regex supplied...returning')
    251              return
    252          lines = self.get_remote_file_lines(remote_file, regex, matching=False)
    253          log.debug("new %s after removing regex (%s) matches:\n%s" % \
    254                    (remote_file, regex, ''.join(lines)))
    255          f = self.remote_file(remote_file)
    256          f.writelines(lines)
    257          f.close()
    258
    259 -    def unlink(self, remote_file):
    260          return self.sftp.unlink(remote_file)
    261
    262 -    def remote_file(self, file, mode='w'):
    263          """
    264          Returns a remote file descriptor
    265          """
    266          rfile = self.sftp.open(file, mode)
    267          rfile.name = file
    268          return rfile
    269
    270 -    def path_exists(self, path):
    271          """
    272          Test whether a remote path exists.
    273          Returns False for broken symbolic links
    274          """
    275          try:
    276              self.stat(path)
    277              return True
    278          except IOError:
    279              return False
    280
    281 -    def chown(self, uid, gid, remote_file):
    282          """
    283          Apply permissions (mode) to remote_file
    284          """
    285          f = self.remote_file(remote_file, 'r')
    286          f.chown(uid, gid, remote_file)
    287          f.close()
    288
    289 -    def chmod(self, mode, remote_file):
    290          """
    291          Apply permissions (mode) to remote_file
    292          """
    293          f = self.remote_file(remote_file, 'r')
    294          f.chmod(mode)
    295          f.close()
    296
    297 -    def ls(self, path):
    298          """
    299          Return a list containing the names of the entries in the remote path.
    300          """
    301          return [os.path.join(path, f) for f in self.sftp.listdir(path)]
    302
    303 -    def isdir(self, path):
    304          """
    305          Return true if the remote path refers to an existing directory.
    306          """
    307          try:
    308              s = self.stat(path)
    309          except IOError:
    310              return False
    311          return stat.S_ISDIR(s.st_mode)
    312
    313 -    def isfile(self, path):
    314          """
    315          Return true if the remote path refers to an existing file.
    316          """
    317          try:
    318              s = self.stat(path)
    319          except IOError:
    320              return False
    321          return stat.S_ISREG(s.st_mode)
    322
    323 -    def stat(self, path):
    324          """
    325          Perform a stat system call on the given remote path.
    326          """
    327          return self.sftp.stat(path)
    328
    329 -    def get(self, remotepath, localpath=None):
    330          """
    331          Copies a file between the remote host and the local host.
    332          """
    333          if not localpath:
    334              localpath = os.path.split(remotepath)[1]
    335          self.sftp_connect()
    336          self.sftp.get(remotepath, localpath)
    337
    338 -    def put(self, localpath, remotepath=None):
    339          """
    340          Copies a file between the local host and the remote host.
    341          """
    342          if not remotepath:
    343              remotepath = os.path.split(localpath)[1]
    344          self.sftp.put(localpath, remotepath)
    345
    346 -    def execute_async(self, command):
    347          """
    348          Executes a remote command without blocking
    349
    350          NOTE: this method will not block, however, if your process does not
    351          complete or background itself before the python process executing this
    352          code exits, it will not persist on the remote machine
    353          """
    354
    355          channel = self.transport.open_session()
    356          channel.exec_command(command)
    357
    358 -    def execute(self, command, silent=True, only_printable=False,
    359                  ignore_exit_status=False, log_output=True):
    360          """
    361          Execute a remote command and return stdout/stderr
    362
    363          NOTE: this function blocks until the process finishes
    364
    365          kwargs:
    366          silent - do not print output
    367          only_printable - filter the command's output to allow only printable
    368                          characters
    369          returns List of output lines
    370          """
    371          channel = self.transport.open_session()
    372          channel.exec_command(command)
    373          #stdin = channel.makefile('wb', -1)
    374          stdout = channel.makefile('rb', -1)
    375          stderr = channel.makefile_stderr('rb', -1)
    376          output = []
    377          line = None
    378          if silent:
    379              output = stdout.readlines() + stderr.readlines()
    380          else:
    381              while line != '':
    382                  line = stdout.readline()
    383                  if only_printable:
    384                      line = ''.join(c for c in line if c in string.printable)
    385                  if line != '':
    386                      output.append(line)
    387                      print line,
    388              for line in stderr.readlines():
    389                  output.append(line)
    390                  print line
    391          if only_printable:
    392              output = map(lambda line: ''.join(c for c in line if c in
    393                                                string.printable), output)
    394          output = map(lambda line: line.strip(), output)
    395          exit_status = channel.recv_exit_status()
    396          if exit_status != 0:
    397              if not ignore_exit_status:
    398                  log.error("command '%s' failed with status %d" % (command,
    399                                                                    exit_status))
    400              else:
    401                  log.debug("command %s failed with status %d" % (command,
    402                                                                  exit_status))
    403          if log_output:
    404              for line in output:
    405                  log.debug(line.strip())
    406          return output
    407
    408 -    def has_required(self, progs):
    409          """
    410          Same as check_required but returns False if not all commands exist
    411          """
    412          try:
    413              return self.check_required(progs)
    414          except exception.RemoteCommandNotFound:
    415              return False
    416
    417 -    def check_required(self, progs):
    418          """
    419          Checks that all commands in the progs list exist on the remote system.
    420          Returns True if all commands exist and raises exception.CommandNotFound
    421          if not.
    422          """
    423          for prog in progs:
    424              if not self.which(prog):
    425                  raise exception.RemoteCommandNotFound(prog)
    426          return True
    427
    428 -    def which(self, prog):
    429          return self.execute('which %s' % prog, ignore_exit_status=True)
    430
    431 -    def get_path(self):
    432          """Returns the PATH environment variable on the remote machine"""
    433          return self.get_env()['PATH']
    434
    435 -    def get_env(self):
    436          """Returns the remote machine's environment as a dictionary"""
    437          env = {}
    438          for line in self.execute('env'):
    439              key, val = line.split('=', 1)
    440              env[key] = val
    441          return env
    442
    443 -    def close(self):
    444          """Closes the connection and cleans up."""
    445          if self._sftp:
    446              self._sftp.close()
    447          if self._transport:
    448              self._transport.close()
    449
    450 -    def interactive_shell(self, user='root'):
    451          if user and self.transport.get_username() != user:
    452              self.connect(username=user)
    453          try:
    454              chan = self.transport.open_session()
    455              chan.get_pty()
    456              chan.invoke_shell()
    457              log.info('Starting interactive shell...')
    458              if HAS_TERMIOS:
    459                  self._posix_shell(chan)
    460              else:
    461                  self._windows_shell(chan)
    462              chan.close()
    463          except Exception, e:
    464              import traceback
    465              print '*** Caught exception: %s: %s' % (e.__class__, e)
    466              traceback.print_exc()
    467
    468 -    def _posix_shell(self, chan):
    469          import select
    470
    471          oldtty = termios.tcgetattr(sys.stdin)
    472          try:
    473              tty.setraw(sys.stdin.fileno())
    474              tty.setcbreak(sys.stdin.fileno())
    475              chan.settimeout(0.0)
    476
    477              # needs to be sent to give vim correct size FIX
    478              chan.send('eval $(resize)\n')
    479
    480              while True:
    481                  r, w, e = select.select([chan, sys.stdin], [], [])
    482                  if chan in r:
    483                      try:
    484                          x = chan.recv(1024)
    485                          if len(x) == 0:
    486                              print '\r\n*** EOF\r\n',
    487                              break
    488                          sys.stdout.write(x)
    489                          sys.stdout.flush()
    490                      except socket.timeout:
    491                          pass
    492                  if sys.stdin in r:
    493                      # fixes up arrow problem
    494                      x = os.read(sys.stdin.fileno(), 1)
    495                      if len(x) == 0:
    496                          break
    497                      chan.send(x)
    498          finally:
    499              termios.tcsetattr(sys.stdin, termios.TCSADRAIN, oldtty)
    500
    501      # thanks to Mike Looijmans for this code
    502 -    def _windows_shell(self, chan):
    503          import threading
    504
    505          sys.stdout.write("Line-buffered terminal emulation. " + \
    506                           "Press F6 or ^Z to send EOF.\r\n\r\n")
    507
    508          def writeall(sock):
    509              while True:
    510                  data = sock.recv(256)
    511                  if not data:
    512                      sys.stdout.write('\r\n*** EOF ***\r\n\r\n')
    513                      sys.stdout.flush()
    514                      break
    515                  sys.stdout.write(data)
    516                  sys.stdout.flush()
    517
    518          writer = threading.Thread(target=writeall, args=(chan,))
    519          writer.start()
    520
    521          # needs to be sent to give vim correct size FIX
    522          chan.send('eval $(resize)\n')
    523
    524          try:
    525              while True:
    526                  d = sys.stdin.read(1)
    527                  if not d:
    528                      break
    529                  chan.send(d)
    530          except EOFError:
    531              # user hit ^Z or F6
    532              pass
    533
    534 -    def __del__(self):
    535          """Attempt to clean up if not explicitly closed."""
    536          log.debug('__del__ called')
    537          self.close()
    538
    539
    540  # for backwards compatibility
    541  Connection = SSHClient
    542
    543
    544 -def main():
    545      """Little test when called directly."""
    546      # Set these to your own details.
    547      myssh = SSHClient('somehost.domain.com')
    548      print myssh.execute('hostname')
    549      #myssh.put('ssh.py')
    550      myssh.close()
    551
    552  if __name__ == "__main__":
    553      main()
    554