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 behavior), 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 transferred. 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 transferred 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 retrieve 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 transferred 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 self._recv_confirm()
177
178 - def _send_recursive(self, files):
179 single_files = [] 180 for base in files: 181 lastdir = base 182 if not os.path.isdir(base): 183 single_files.append(base) 184 continue 185 for root, dirs, fls in os.walk(base): 186 # pop back out to the next dir in the walk 187 while lastdir != os.path.commonprefix([lastdir, root]): 188 self._send_popd() 189 lastdir = os.path.split(lastdir)[0] 190 self._send_pushd(root) 191 lastdir = root 192 self._send_files([os.path.join(root, f) for f in fls]) 193 if single_files: 194 self._send_popd() 195 self._send_files(single_files)
196
197 - def _send_pushd(self, directory):
198 (mode, size, mtime, atime) = self._read_stats(directory) 199 basename = os.path.basename(directory) 200 if self.preserve_times: 201 self._send_time(mtime, atime) 202 self.channel.sendall('D%s 0 %s\n' % (mode, basename)) 203 self._recv_confirm()
204
205 - def _send_popd(self):
206 self.channel.sendall('E\n') 207 self._recv_confirm()
208
209 - def _send_time(self, mtime, atime):
210 self.channel.sendall('T%d 0 %d 0\n' % (mtime, atime)) 211 self._recv_confirm()
212
213 - def _recv_confirm(self):
214 # read scp response 215 msg = '' 216 try: 217 msg = self.channel.recv(512) 218 except SocketTimeout: 219 raise exception.SCPException('Timeout waiting for scp response') 220 if msg and msg[0] == '\x00': 221 return 222 elif msg and msg[0] == '\x01': 223 raise exception.SCPException(msg[1:]) 224 elif self.channel.recv_stderr_ready(): 225 msg = self.channel.recv_stderr(512) 226 raise exception.SCPException(msg) 227 elif not msg: 228 raise exception.SCPException('No response from server') 229 else: 230 raise exception.SCPException( 231 'Invalid response from server: ' + msg)
232
233 - def _recv_all(self):
234 # loop over scp commands, and recive as necessary 235 command = {'C': self._recv_file, 236 'T': self._set_time, 237 'D': self._recv_pushd, 238 'E': self._recv_popd} 239 while not self.channel.closed: 240 # wait for command as long as we're open 241 self.channel.sendall('\x00') 242 msg = self.channel.recv(1024) 243 if not msg: # chan closed while recving 244 break 245 code = msg[0] 246 try: 247 command[code](msg[1:]) 248 except KeyError: 249 raise exception.SCPException(str(msg).strip()) 250 # directory times can't be set until we're done writing files 251 self._set_dirtimes()
252
253 - def _set_time(self, cmd):
254 try: 255 times = cmd.split() 256 mtime = int(times[0]) 257 atime = int(times[2]) or mtime 258 except: 259 self.channel.send('\x01') 260 raise exception.SCPException('Bad time format') 261 # save for later 262 self._utime = (atime, mtime)
263
264 - def _recv_file(self, cmd):
265 chan = self.channel 266 parts = cmd.split() 267 try: 268 mode = int(parts[0], 8) 269 size = int(parts[1]) 270 path = os.path.join(self._recv_dir, parts[2]) 271 if self._rename_file: 272 path = self._recv_dir 273 except: 274 chan.send('\x01') 275 chan.close() 276 raise exception.SCPException('Bad file format') 277 278 try: 279 file_hdl = file(path, 'wb') 280 except IOError, e: 281 chan.send('\x01' + str(e)) 282 chan.close() 283 raise 284 285 buff_size = self.buff_size 286 pos = 0 287 chan.send('\x00') 288 try: 289 while pos < size: 290 # we have to make sure we don't read the final byte 291 if size - pos <= buff_size: 292 buff_size = size - pos 293 file_hdl.write(chan.recv(buff_size)) 294 pos = file_hdl.tell() 295 if self._progress: 296 self._progress(path, size, pos) 297 if size == 0 and self._progress: 298 self._progress(path, 1, 1) 299 300 msg = chan.recv(512) 301 if msg and msg[0] != '\x00': 302 raise exception.SCPException(msg[1:]) 303 except SocketTimeout: 304 chan.close() 305 raise exception.SCPException('Error receiving, socket.timeout') 306 307 file_hdl.truncate() 308 try: 309 os.utime(path, self._utime) 310 self._utime = None 311 os.chmod(path, mode) 312 # should we notify the other end? 313 finally: 314 file_hdl.close()
315 # '\x00' confirmation sent in _recv_all 316
317 - def _recv_pushd(self, cmd):
318 parts = cmd.split() 319 try: 320 mode = int(parts[0], 8) 321 path = os.path.join(self._recv_dir, parts[2]) 322 if self._rename_dir: 323 path = self._recv_dir 324 except: 325 self.channel.send('\x01') 326 raise exception.SCPException('Bad directory format') 327 try: 328 if not os.path.exists(path): 329 os.mkdir(path, mode) 330 elif os.path.isdir(path): 331 os.chmod(path, mode) 332 else: 333 raise exception.SCPException('%s: Not a directory' % path) 334 self._dirtimes[path] = (self._utime) 335 self._utime = None 336 self._recv_dir = path 337 except (OSError, exception.SCPException), e: 338 self.channel.send('\x01' + e[0]) 339 raise
340
341 - def _recv_popd(self, *cmd):
342 self._recv_dir = os.path.split(self._recv_dir)[0]
343
344 - def _set_dirtimes(self):
345 try: 346 for d in self._dirtimes: 347 os.utime(d, self._dirtimes[d]) 348 finally: 349 self._dirtimes = {}
350