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