1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
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
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
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
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
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
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
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
208
210 self.channel.sendall('T%d 0 %d 0\n' % (mtime, atime))
211 self._recv_confirm()
212
214
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
234
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
241 self.channel.sendall('\x00')
242 msg = self.channel.recv(1024)
243 if not msg:
244 break
245 code = msg[0]
246 try:
247 command[code](msg[1:])
248 except KeyError:
249 raise exception.SCPException(str(msg).strip())
250
251 self._set_dirtimes()
252
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
262 self._utime = (atime, mtime)
263
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
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
313 finally:
314 file_hdl.close()
315
316
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
342 self._recv_dir = os.path.split(self._recv_dir)[0]
343
345 try:
346 for d in self._dirtimes:
347 os.utime(d, self._dirtimes[d])
348 finally:
349 self._dirtimes = {}
350