starcluster.scp

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

# Source Code for Module starcluster.scp

      1  # scp.py
      2  # Copyright (C) 2008 James Bardin <jbardin@bu.edu>
      3  #
      4  # This library is free software; you can redistribute it and/or
      5  # modify it under the terms of the GNU Lesser General Public
      6  # License as published by the Free Software Foundation; either
      7  # version 2.1 of the License, or (at your option) any later version.
      8  #
      9  # This library is distributed in the hope that it will be useful,
     10  # but WITHOUT ANY WARRANTY; without even the implied warranty of
     11  # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
     12  # Lesser General Public License for more details.
     13  #
     14  # You should have received a copy of the GNU Lesser General Public License
     15  # along with this library; if not, write to the Free Software Foundation, Inc.,
     16  # 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301  USA
     17
     18
     19  """
     20  Utilities for sending files over ssh using the scp1 protocol.
     21  """
     22
     23  import os
     24  from socket import timeout as SocketTimeout
     25
     26  from starcluster import exception
     27
     28
     29 -class SCPClient(object):
     30      """
     31      An scp1 implementation, compatible with openssh scp.
     32      Raises SCPException for all transport related errors. Local filesystem
     33      and OS errors pass through.
     34
     35      Main public methods are .put and .get
     36      The get method is controlled by the remote scp instance, and behaves
     37      accordingly. This means that symlinks are resolved, and the transfer is
     38      halted after too many levels of symlinks are detected.
     39      The put method uses os.walk for recursion, and sends files accordingly.
     40      Since scp doesn't support symlinks, we send file symlinks as the file
     41      (matching scp behaviour), but we make no attempt at symlinked directories.
     42      """
     43 -    def __init__(self, transport, buff_size=16384, socket_timeout=5.0,
     44                   progress=None):
     45          """
     46          Create an scp1 client.
     47
     48          @param transport: an existing paramiko L{Transport}
     49          @type transport: L{Transport}
     50          @param buff_size: size of the scp send buffer.
     51          @type buff_size: int
     52          @param socket_timeout: channel socket timeout in seconds
     53          @type socket_timeout: float
     54          @param progress: callback - called with (filename, size, sent) during
     55          transfers
     56          @type progress: function(string, int, int)
     57          """
     58          self.transport = transport
     59          self.buff_size = buff_size
     60          self.socket_timeout = socket_timeout
     61          self.channel = None
     62          self.preserve_times = False
     63          self._progress = progress
     64          self._recv_dir = ''
     65          self._rename_file = False
     66          self._rename_dir = False
     67          self._utime = None
     68          self._dirtimes = {}
     69
     70 -    def put(self, files, remote_path='.',
     71              recursive=False, preserve_times=False):
     72          """
     73          Transfer files to remote host.
     74
     75          @param files: A single path, or a list of paths to be transfered.
     76              recursive must be True to transfer directories.
     77          @type files: string OR list of strings
     78          @param remote_path: path in which to receive the files on the remote
     79              host. defaults to '.'
     80          @type remote_path: str
     81          @param recursive: transfer files and directories recursively
     82          @type recursive: bool
     83          @param preserve_times: preserve mtime and atime of transfered files
     84              and directories.
     85          @type preserve_times: bool
     86          """
     87          self.preserve_times = preserve_times
     88          self.channel = self.transport.open_session()
     89          self.channel.settimeout(self.socket_timeout)
     90          scp_command = ('scp -t %s', 'scp -r -t %s')[recursive]
     91          self.channel.exec_command(scp_command % remote_path)
     92          self._recv_confirm()
     93
     94          if not isinstance(files, (list, tuple)):
     95              files = [files]
     96
     97          if recursive:
     98              self._send_recursive(files)
     99          else:
    100              self._send_files(files)
    101
    102          if self.channel:
    103              self.channel.close()
    104
    105 -    def get(self, remote_path, local_path='',
    106              recursive=False, preserve_times=False):
    107          """
    108          Transfer files from remote host to localhost
    109
    110          @param remote_path: path to retreive from remote host. since this is
    111              evaluated by scp on the remote host, shell wildcards and
    112              environment variables may be used.
    113          @type remote_path: str
    114          @param local_path: path in which to receive files locally
    115          @type local_path: str
    116          @param recursive: transfer files and directories recursively
    117          @type recursive: bool
    118          @param preserve_times: preserve mtime and atime of transfered files
    119              and directories.
    120          @type preserve_times: bool
    121          """
    122          self._recv_dir = local_path or os.getcwd()
    123          if len(remote_path) == 1 and local_path != os.curdir:
    124              if not os.path.isdir(local_path):
    125                  self._rename_dir = recursive
    126                  self._rename_file = not recursive
    127          if len(remote_path) > 1:
    128              if not os.path.exists(self._recv_dir):
    129                  msg = "Local path '%s' does not exist" % self._recv_dir
    130                  raise exception.SCPException(msg)
    131              elif not os.path.isdir(self._recv_dir):
    132                  msg = "Local path '%s' is not a directory" % self._recv_dir
    133                  raise exception.SCPException(msg)
    134          rcsv = ('', ' -r')[recursive]
    135          prsv = ('', ' -p')[preserve_times]
    136          self.channel = self.transport.open_session()
    137          self.channel.settimeout(self.socket_timeout)
    138          if isinstance(remote_path, (list, tuple)):
    139              remote_path = ' '.join(remote_path)
    140          self.channel.exec_command('scp%s%s -f %s' % (rcsv, prsv, remote_path))
    141          self._recv_all()
    142
    143          if self.channel:
    144              self.channel.close()
    145
    146 -    def _read_stats(self, name):
    147          """return just the file stats needed for scp"""
    148          stats = os.stat(name)
    149          mode = oct(stats.st_mode)[-4:]
    150          size = stats.st_size
    151          atime = int(stats.st_atime)
    152          mtime = int(stats.st_mtime)
    153          return (mode, size, mtime, atime)
    154
    155 -    def _send_files(self, files):
    156          for name in files:
    157              basename = os.path.basename(name)
    158              (mode, size, mtime, atime) = self._read_stats(name)
    159              if self.preserve_times:
    160                  self._send_time(mtime, atime)
    161              file_hdl = file(name, 'rb')
    162              self.channel.sendall('C%s %d %s\n' % (mode, size, basename))
    163              self._recv_confirm()
    164              file_pos = 0
    165              buff_size = self.buff_size
    166              chan = self.channel
    167              while file_pos < size:
    168                  chan.sendall(file_hdl.read(buff_size))
    169                  file_pos = file_hdl.tell()
    170                  if self._progress:
    171                      self._progress(basename, size, file_pos)
    172              if size == 0 and self._progress:
    173                  self._progress(basename, 1, 1)
    174              chan.sendall('\x00')
    175              file_hdl.close()
    176
    177 -    def _send_recursive(self, files):
    178          single_files = []
    179          for base in files:
    180              lastdir = base
    181              if not os.path.isdir(base):
    182                  single_files.append(base)
    183                  continue
    184              for root, dirs, fls in os.walk(base):
    185                  # pop back out to the next dir in the walk
    186                  while lastdir != os.path.commonprefix([lastdir, root]):
    187                      self._send_popd()
    188                      lastdir = os.path.split(lastdir)[0]
    189                  self._send_pushd(root)
    190                  lastdir = root
    191                  self._send_files([os.path.join(root, f) for f in fls])
    192          if single_files:
    193              self._send_popd()
    194              self._send_files(single_files)
    195
    196 -    def _send_pushd(self, directory):
    197          (mode, size, mtime, atime) = self._read_stats(directory)
    198          basename = os.path.basename(directory)
    199          if self.preserve_times:
    200              self._send_time(mtime, atime)
    201          self.channel.sendall('D%s 0 %s\n' % (mode, basename))
    202          self._recv_confirm()
    203
    204 -    def _send_popd(self):
    205          self.channel.sendall('E\n')
    206          self._recv_confirm()
    207
    208 -    def _send_time(self, mtime, atime):
    209          self.channel.sendall('T%d 0 %d 0\n' % (mtime, atime))
    210          self._recv_confirm()
    211
    212 -    def _recv_confirm(self):
    213          # read scp response
    214          msg = ''
    215          try:
    216              msg = self.channel.recv(512)
    217          except SocketTimeout:
    218              raise exception.SCPException('Timout waiting for scp response')
    219          if msg and msg[0] == '\x00':
    220              return
    221          elif msg and msg[0] == '\x01':
    222              raise exception.SCPException(msg[1:])
    223          elif self.channel.recv_stderr_ready():
    224              msg = self.channel.recv_stderr(512)
    225              raise exception.SCPException(msg)
    226          elif not msg:
    227              raise exception.SCPException('No response from server')
    228          else:
    229              raise exception.SCPException(
    230                  'Invalid response from server: ' + msg)
    231
    232 -    def _recv_all(self):
    233          # loop over scp commands, and recive as necessary
    234          command = {'C': self._recv_file,
    235                     'T': self._set_time,
    236                     'D': self._recv_pushd,
    237                     'E': self._recv_popd}
    238          while not self.channel.closed:
    239              # wait for command as long as we're open
    240              self.channel.sendall('\x00')
    241              msg = self.channel.recv(1024)
    242              if not msg:  # chan closed while recving
    243                  break
    244              code = msg[0]
    245              try:
    246                  command[code](msg[1:])
    247              except KeyError:
    248                  raise exception.SCPException(str(msg).strip())
    249          # directory times can't be set until we're done writing files
    250          self._set_dirtimes()
    251
    252 -    def _set_time(self, cmd):
    253          try:
    254              times = cmd.split()
    255              mtime = int(times[0])
    256              atime = int(times[2]) or mtime
    257          except:
    258              self.channel.send('\x01')
    259              raise exception.SCPException('Bad time format')
    260          # save for later
    261          self._utime = (atime, mtime)
    262
    263 -    def _recv_file(self, cmd):
    264          chan = self.channel
    265          parts = cmd.split()
    266          try:
    267              mode = int(parts[0], 8)
    268              size = int(parts[1])
    269              path = os.path.join(self._recv_dir, parts[2])
    270              if self._rename_file:
    271                  path = self._recv_dir
    272          except:
    273              chan.send('\x01')
    274              chan.close()
    275              raise exception.SCPException('Bad file format')
    276
    277          try:
    278              file_hdl = file(path, 'wb')
    279          except IOError, e:
    280              chan.send('\x01' + str(e))
    281              chan.close()
    282              raise
    283
    284          buff_size = self.buff_size
    285          pos = 0
    286          chan.send('\x00')
    287          try:
    288              while pos < size:
    289                  # we have to make sure we don't read the final byte
    290                  if size - pos <= buff_size:
    291                      buff_size = size - pos
    292                  file_hdl.write(chan.recv(buff_size))
    293                  pos = file_hdl.tell()
    294                  if self._progress:
    295                      self._progress(path, size, pos)
    296              if size == 0 and self._progress:
    297                  self._progress(path, 1, 1)
    298
    299              msg = chan.recv(512)
    300              if msg and msg[0] != '\x00':
    301                  raise exception.SCPException(msg[1:])
    302          except SocketTimeout:
    303              chan.close()
    304              raise exception.SCPException('Error receiving, socket.timeout')
    305
    306          file_hdl.truncate()
    307          try:
    308              os.utime(path, self._utime)
    309              self._utime = None
    310              os.chmod(path, mode)
    311              # should we notify the other end?
    312          finally:
    313              file_hdl.close()
    314          # '\x00' confirmation sent in _recv_all
    315
    316 -    def _recv_pushd(self, cmd):
    317          parts = cmd.split()
    318          try:
    319              mode = int(parts[0], 8)
    320              path = os.path.join(self._recv_dir, parts[2])
    321              if self._rename_dir:
    322                  path = self._recv_dir
    323          except:
    324              self.channel.send('\x01')
    325              raise exception.SCPException('Bad directory format')
    326          try:
    327              if not os.path.exists(path):
    328                  os.mkdir(path, mode)
    329              elif os.path.isdir(path):
    330                  os.chmod(path, mode)
    331              else:
    332                  raise exception.SCPException('%s: Not a directory' % path)
    333              self._dirtimes[path] = (self._utime)
    334              self._utime = None
    335              self._recv_dir = path
    336          except (OSError, exception.SCPException), e:
    337              self.channel.send('\x01' + e[0])
    338              raise
    339
    340 -    def _recv_popd(self, *cmd):
    341          self._recv_dir = os.path.split(self._recv_dir)[0]
    342
    343 -    def _set_dirtimes(self):
    344          try:
    345              for d in self._dirtimes:
    346                  os.utime(d, self._dirtimes[d])
    347          finally:
    348              self._dirtimes = {}
    349