← Back to branch summary

~bzr-pqm/bzr/bzr.dev

« back to all changes in this revision

Viewing changes to bzrlib/transport/sftp.py

[merge] sftp fixes

Show diffs side-by-side

added added

removed removed

Lines of Context:
bzrlib/transport/sftp.py
23
23
import stat
24
24
import sys
25
25
import urllib
 
26
import time
 
27
import random
26
28
 
27
29
from bzrlib.errors import (FileExists, 
28
30
                           TransportNotPossible, NoSuchFile, NonRelativePath,
29
 
                           TransportError)
 
31
                           TransportError,
 
32
                           LockError)
30
33
from bzrlib.config import config_dir
31
34
from bzrlib.trace import mutter, warning, error
32
35
from bzrlib.transport import Transport, register_transport
36
39
except ImportError:
37
40
    error('The SFTP transport requires paramiko.')
38
41
    raise
 
42
else:
 
43
    from paramiko.sftp import (SFTP_FLAG_WRITE, SFTP_FLAG_CREATE,
 
44
                               SFTP_FLAG_EXCL, SFTP_FLAG_TRUNC,
 
45
                               CMD_HANDLE, CMD_OPEN)
 
46
    from paramiko.sftp_attr import SFTPAttributes
 
47
    from paramiko.sftp_file import SFTPFile
39
48
 
40
49
 
41
50
SYSTEM_HOSTKEYS = {}
81
90
class SFTPTransportError (TransportError):
82
91
    pass
83
92
 
 
93
class SFTPLock(object):
 
94
    """This fakes a lock in a remote location."""
 
95
    __slots__ = ['path', 'lock_path', 'lock_file', 'transport']
 
96
    def __init__(self, path, transport):
 
97
        assert isinstance(transport, SFTPTransport)
 
98
 
 
99
        self.lock_file = None
 
100
        self.path = path
 
101
        self.lock_path = path + '.write-lock'
 
102
        self.transport = transport
 
103
        try:
 
104
            self.lock_file = transport._sftp_open_exclusive(self.lock_path)
 
105
        except FileExists:
 
106
            raise LockError('File %r already locked' % (self.path,))
 
107
 
 
108
    def __del__(self):
 
109
        """Should this warn, or actually try to cleanup?"""
 
110
        if self.lock_file:
 
111
            warn("SFTPLock %r not explicitly unlocked" % (self.path,))
 
112
            self.unlock()
 
113
 
 
114
    def unlock(self):
 
115
        if not self.lock_file:
 
116
            return
 
117
        self.lock_file.close()
 
118
        self.lock_file = None
 
119
        try:
 
120
            self.transport.delete(self.lock_path)
 
121
        except (NoSuchFile,):
 
122
            # What specific errors should we catch here?
 
123
            pass
84
124
 
85
125
class SFTPTransport (Transport):
86
126
    """
87
127
    Transport implementation for SFTP access.
88
128
    """
89
129
 
90
 
    _url_matcher = re.compile(r'^sftp://([^:@]*(:[^@]*)?@)?(.*?)(:\d+)?(/.*)?$')
 
130
    _url_matcher = re.compile(r'^sftp://([^:@]*(:[^@]*)?@)?(.*?)(:[^/]+)?(/.*)?$')
91
131
    
92
132
    def __init__(self, base, clone_from=None):
93
133
        assert base.startswith('sftp://')
148
188
                basepath.append(p)
149
189
 
150
190
        path = '/'.join(basepath)
151
 
        if len(path) and path[0] != '/':
152
 
            path = '/' + path
 
191
        # could still be a "relative" path here, but relative on the sftp server
153
192
        return path
154
193
 
155
194
    def relpath(self, abspath):
156
195
        # FIXME: this is identical to HttpTransport -- share it
157
 
        if not abspath.startswith(self.base):
 
196
        m = self._url_matcher.match(abspath)
 
197
        path = m.group(5)
 
198
        if not path.startswith(self._path):
158
199
            raise NonRelativePath('path %r is not under base URL %r'
159
200
                           % (abspath, self.base))
160
201
        pl = len(self.base)
178
219
        """
179
220
        try:
180
221
            path = self._abspath(relpath)
181
 
            return self._sftp.file(path)
 
222
            f = self._sftp.file(path)
 
223
            try:
 
224
                f.prefetch()
 
225
            except AttributeError:
 
226
                # only works on paramiko 1.5.1 or greater
 
227
                pass
 
228
            return f
182
229
        except (IOError, paramiko.SSHException), x:
183
230
            raise NoSuchFile('Error retrieving %s: %s' % (path, str(x)), x)
184
231
 
194
241
                 Some implementations may return objects which can be read
195
242
                 past this length, but this is not guaranteed.
196
243
        """
 
244
        # TODO: implement get_partial_multi to help with knit support
197
245
        f = self.get(relpath)
198
246
        f.seek(start)
 
247
        try:
 
248
            f.prefetch()
 
249
        except AttributeError:
 
250
            # only works on paramiko 1.5.1 or greater
 
251
            pass
199
252
        return f
200
253
 
201
254
    def put(self, relpath, f):
205
258
        :param relpath: Location to put the contents, relative to base.
206
259
        :param f:       File-like or string object.
207
260
        """
208
 
        # FIXME: should do something atomic or locking here, this is unsafe
209
 
        try:
210
 
            path = self._abspath(relpath)
211
 
            fout = self._sftp.file(path, 'wb')
212
 
        except IOError, e:
213
 
            self._translate_io_exception(e, relpath)
214
 
        except (IOError, paramiko.SSHException), x:
215
 
            raise SFTPTransportError('Unable to write file %r' % (path,), x)
216
 
        try:
217
 
            self._pump(f, fout)
218
 
        finally:
219
 
            fout.close()
 
261
        final_path = self._abspath(relpath)
 
262
        tmp_relpath = '%s.tmp.%.9f.%d.%d' % (relpath, time.time(),
 
263
                        os.getpid(), random.randint(0,0x7FFFFFFF))
 
264
        tmp_abspath = self._abspath(tmp_relpath)
 
265
        fout = self._sftp_open_exclusive(tmp_relpath)
 
266
 
 
267
        try:
 
268
            try:
 
269
                self._pump(f, fout)
 
270
            except IOError, e:
 
271
                self._translate_io_exception(e, relpath)
 
272
            except paramiko.SSHException, x:
 
273
                raise SFTPTransportError('Unable to write file %r' % (path,), x)
 
274
        except Exception, e:
 
275
            # If we fail, try to clean up the temporary file
 
276
            # before we throw the exception
 
277
            # but don't let another exception mess things up
 
278
            try:
 
279
                fout.close()
 
280
                self._sftp.remove(tmp_abspath)
 
281
            except:
 
282
                pass
 
283
            raise e
 
284
        else:
 
285
            # sftp rename doesn't allow overwriting, so play tricks:
 
286
            tmp_safety = 'bzr.tmp.%.9f.%d.%d' % (time.time(), os.getpid(), random.randint(0, 0x7FFFFFFF))
 
287
            tmp_safety = self._abspath(tmp_safety)
 
288
            try:
 
289
                self._sftp.rename(final_path, tmp_safety)
 
290
                file_existed = True
 
291
            except:
 
292
                file_existed = False
 
293
            try:
 
294
                self._sftp.rename(tmp_abspath, final_path)
 
295
            except IOError, e:
 
296
                self._translate_io_exception(e, relpath)
 
297
            except paramiko.SSHException, x:
 
298
                raise SFTPTransportError('Unable to rename into file %r' 
 
299
                                          % (path,), x)
 
300
            if file_existed:
 
301
                self._sftp.unlink(tmp_safety)
220
302
 
221
303
    def iter_files_recursive(self):
222
304
        """Walk the relative paths of all files in this transport."""
326
408
    def lock_read(self, relpath):
327
409
        """
328
410
        Lock the given file for shared (read) access.
329
 
        :return: A lock object, which should be passed to Transport.unlock()
 
411
        :return: A lock object, which has an unlock() member function
330
412
        """
331
413
        # FIXME: there should be something clever i can do here...
332
414
        class BogusLock(object):
341
423
        Lock the given file for exclusive (write) access.
342
424
        WARNING: many transports do not support this, so trying avoid using it
343
425
 
344
 
        :return: A lock object, which should be passed to Transport.unlock()
 
426
        :return: A lock object, which has an unlock() member function
345
427
        """
346
 
        # FIXME: there should be something clever i can do here...
347
 
        class BogusLock(object):
348
 
            def __init__(self, path):
349
 
                self.path = path
350
 
            def unlock(self):
351
 
                pass
352
 
        return BogusLock(relpath)
 
428
        # This is a little bit bogus, but basically, we create a file
 
429
        # which should not already exist, and if it does, we assume
 
430
        # that there is a lock, and if it doesn't, the we assume
 
431
        # that we have taken the lock.
 
432
        return SFTPLock(relpath, self)
353
433
 
354
434
 
355
435
    def _unparse_url(self, path=None):
356
436
        if path is None:
357
437
            path = self._path
358
 
        if self._port == 22:
359
 
            return 'sftp://%s@%s%s' % (self._username, self._host, path)
360
 
        return 'sftp://%s@%s:%d%s' % (self._username, self._host, self._port, path)
 
438
        host = self._host
 
439
        username = urllib.quote(self._username)
 
440
        if self._password:
 
441
            username += ':' + urllib.quote(self._password)
 
442
        if self._port != 22:
 
443
            host += ':%d' % self._port
 
444
        return 'sftp://%s@%s/%s' % (username, host, urllib.quote(path))
361
445
 
362
446
    def _parse_url(self, url):
363
447
        assert url[:7] == 'sftp://'
368
452
        if self._username is None:
369
453
            self._username = getpass.getuser()
370
454
        else:
371
 
            self._username = self._username[:-1]
372
 
        if self._password:
373
 
            self._password = self._password[1:]
374
 
            self._username = self._username[len(self._password)+1:]
 
455
            if self._password:
 
456
                # username field is 'user:pass@' in this case, and password is ':pass'
 
457
                username_len = len(self._username) - len(self._password) - 1
 
458
                self._username = urllib.unquote(self._username[:username_len])
 
459
                self._password = urllib.unquote(self._password[1:])
 
460
            else:
 
461
                self._username = urllib.unquote(self._username[:-1])
375
462
        if self._port is None:
376
463
            self._port = 22
377
464
        else:
378
465
            self._port = int(self._port[1:])
379
466
        if (self._path is None) or (self._path == ''):
380
 
            self._path = '/'
 
467
            self._path = ''
 
468
        else:
 
469
            # remove leading '/'
 
470
            self._path = urllib.unquote(self._path[1:])
381
471
 
382
472
    def _sftp_connect(self):
383
473
        global SYSTEM_HOSTKEYS, BZR_HOSTKEYS
446
536
                pass
447
537
 
448
538
        # give up and ask for a password
449
 
        password = getpass.getpass('SSH %s@%s password: ' % (self._username, self._host))
 
539
        # FIXME: shouldn't be implementing UI this deep into bzrlib
 
540
        enc = sys.stdout.encoding
 
541
        password = getpass.getpass('SSH %s@%s password: ' %
 
542
            (self._username.encode(enc, 'replace'), self._host.encode(enc, 'replace')))
450
543
        try:
451
544
            transport.auth_password(self._username, password)
452
545
        except paramiko.SSHException:
460
553
            transport.auth_publickey(self._username, key)
461
554
            return True
462
555
        except paramiko.PasswordRequiredException:
463
 
            password = getpass.getpass('SSH %s password: ' % (os.path.basename(filename),))
 
556
            # FIXME: shouldn't be implementing UI this deep into bzrlib
 
557
            enc = sys.stdout.encoding
 
558
            password = getpass.getpass('SSH %s password: ' % 
 
559
                (os.path.basename(filename).encode(enc, 'replace'),))
464
560
            try:
465
561
                key = pkey_class.from_private_key_file(filename, password)
466
562
                transport.auth_publickey(self._username, key)
472
568
        except IOError:
473
569
            pass
474
570
        return False
 
571
 
 
572
    def _sftp_open_exclusive(self, relpath):
 
573
        """Open a remote path exclusively.
 
574
 
 
575
        SFTP supports O_EXCL (SFTP_FLAG_EXCL), which fails if
 
576
        the file already exists. However it does not expose this
 
577
        at the higher level of SFTPClient.open(), so we have to
 
578
        sneak away with it.
 
579
 
 
580
        WARNING: This breaks the SFTPClient abstraction, so it
 
581
        could easily break against an updated version of paramiko.
 
582
 
 
583
        :param relpath: The relative path, where the file should be opened
 
584
        """
 
585
        path = self._abspath(relpath)
 
586
        attr = SFTPAttributes()
 
587
        mode = (SFTP_FLAG_WRITE | SFTP_FLAG_CREATE 
 
588
                | SFTP_FLAG_TRUNC | SFTP_FLAG_EXCL)
 
589
        try:
 
590
            t, msg = self._sftp._request(CMD_OPEN, path, mode, attr)
 
591
            if t != CMD_HANDLE:
 
592
                raise SFTPTransportError('Expected an SFTP handle')
 
593
            handle = msg.get_string()
 
594
            return SFTPFile(self._sftp, handle, 'w', -1)
 
595
        except IOError, e:
 
596
            self._translate_io_exception(e, relpath)
 
597
        except paramiko.SSHException, x:
 
598
            raise SFTPTransportError('Unable to open file %r' % (path,), x)
 
599