/brz/remove-bazaar

To get this branch, use:
bzr branch http://gegoxaren.bato24.eu/bzr/brz/remove-bazaar

« back to all changes in this revision

Viewing changes to bzrlib/transport/sftp.py

  • Committer: Robey Pointer
  • Date: 2005-11-22 01:41:46 UTC
  • mfrom: (1185.40.1)
  • mto: (1185.33.37 bzr.dev)
  • mto: This revision was merged to the branch mainline in revision 1512.
  • Revision ID: robey@lag.net-20051122014146-5186f5e310a15202
make sftp put faster when using paramiko 1.5.1

Show diffs side-by-side

added added

removed removed

Lines of Context:
 
1
# Copyright (C) 2005 Robey Pointer <robey@lag.net>, Canonical Ltd
 
2
 
 
3
# This program is free software; you can redistribute it and/or modify
 
4
# it under the terms of the GNU General Public License as published by
 
5
# the Free Software Foundation; either version 2 of the License, or
 
6
# (at your option) any later version.
 
7
 
 
8
# This program is distributed in the hope that it will be useful,
 
9
# but WITHOUT ANY WARRANTY; without even the implied warranty of
 
10
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 
11
# GNU General Public License for more details.
 
12
 
 
13
# You should have received a copy of the GNU General Public License
 
14
# along with this program; if not, write to the Free Software
 
15
# Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA  02111-1307  USA
 
16
 
 
17
"""Implementation of Transport over SFTP, using paramiko."""
 
18
 
 
19
import errno
 
20
import getpass
 
21
import os
 
22
import re
 
23
import stat
 
24
import sys
 
25
import urllib
 
26
 
 
27
from bzrlib.errors import (FileExists, 
 
28
                           TransportNotPossible, NoSuchFile, NonRelativePath,
 
29
                           TransportError)
 
30
from bzrlib.config import config_dir
 
31
from bzrlib.trace import mutter, warning, error
 
32
from bzrlib.transport import Transport, register_transport
 
33
 
 
34
try:
 
35
    import paramiko
 
36
except ImportError:
 
37
    error('The SFTP transport requires paramiko.')
 
38
    raise
 
39
 
 
40
 
 
41
SYSTEM_HOSTKEYS = {}
 
42
BZR_HOSTKEYS = {}
 
43
 
 
44
def load_host_keys():
 
45
    """
 
46
    Load system host keys (probably doesn't work on windows) and any
 
47
    "discovered" keys from previous sessions.
 
48
    """
 
49
    global SYSTEM_HOSTKEYS, BZR_HOSTKEYS
 
50
    try:
 
51
        SYSTEM_HOSTKEYS = paramiko.util.load_host_keys(os.path.expanduser('~/.ssh/known_hosts'))
 
52
    except Exception, e:
 
53
        mutter('failed to load system host keys: ' + str(e))
 
54
    bzr_hostkey_path = os.path.join(config_dir(), 'ssh_host_keys')
 
55
    try:
 
56
        BZR_HOSTKEYS = paramiko.util.load_host_keys(bzr_hostkey_path)
 
57
    except Exception, e:
 
58
        mutter('failed to load bzr host keys: ' + str(e))
 
59
        save_host_keys()
 
60
 
 
61
def save_host_keys():
 
62
    """
 
63
    Save "discovered" host keys in $(config)/ssh_host_keys/.
 
64
    """
 
65
    global SYSTEM_HOSTKEYS, BZR_HOSTKEYS
 
66
    bzr_hostkey_path = os.path.join(config_dir(), 'ssh_host_keys')
 
67
    if not os.path.isdir(config_dir()):
 
68
        os.mkdir(config_dir())
 
69
    try:
 
70
        f = open(bzr_hostkey_path, 'w')
 
71
        f.write('# SSH host keys collected by bzr\n')
 
72
        for hostname, keys in BZR_HOSTKEYS.iteritems():
 
73
            for keytype, key in keys.iteritems():
 
74
                f.write('%s %s %s\n' % (hostname, keytype, key.get_base64()))
 
75
        f.close()
 
76
    except IOError, e:
 
77
        mutter('failed to save bzr host keys: ' + str(e))
 
78
 
 
79
 
 
80
 
 
81
class SFTPTransportError (TransportError):
 
82
    pass
 
83
 
 
84
 
 
85
class SFTPTransport (Transport):
 
86
    """
 
87
    Transport implementation for SFTP access.
 
88
    """
 
89
 
 
90
    _url_matcher = re.compile(r'^sftp://([^:@]*(:[^@]*)?@)?(.*?)(:\d+)?(/.*)?$')
 
91
    
 
92
    def __init__(self, base, clone_from=None):
 
93
        assert base.startswith('sftp://')
 
94
        super(SFTPTransport, self).__init__(base)
 
95
        self._parse_url(base)
 
96
        if clone_from is None:
 
97
            self._sftp_connect()
 
98
        else:
 
99
            # use the same ssh connection, etc
 
100
            self._sftp = clone_from._sftp
 
101
        # super saves 'self.base'
 
102
    
 
103
    def should_cache(self):
 
104
        """
 
105
        Return True if the data pulled across should be cached locally.
 
106
        """
 
107
        return True
 
108
 
 
109
    def clone(self, offset=None):
 
110
        """
 
111
        Return a new SFTPTransport with root at self.base + offset.
 
112
        We share the same SFTP session between such transports, because it's
 
113
        fairly expensive to set them up.
 
114
        """
 
115
        if offset is None:
 
116
            return SFTPTransport(self.base, self)
 
117
        else:
 
118
            return SFTPTransport(self.abspath(offset), self)
 
119
 
 
120
    def abspath(self, relpath):
 
121
        """
 
122
        Return the full url to the given relative path.
 
123
        
 
124
        @param relpath: the relative path or path components
 
125
        @type relpath: str or list
 
126
        """
 
127
        return self._unparse_url(self._abspath(relpath))
 
128
    
 
129
    def _abspath(self, relpath):
 
130
        """Return the absolute path segment without the SFTP URL."""
 
131
        # FIXME: share the common code across transports
 
132
        assert isinstance(relpath, basestring)
 
133
        relpath = [urllib.unquote(relpath)]
 
134
        basepath = self._path.split('/')
 
135
        if len(basepath) > 0 and basepath[-1] == '':
 
136
            basepath = basepath[:-1]
 
137
 
 
138
        for p in relpath:
 
139
            if p == '..':
 
140
                if len(basepath) == 0:
 
141
                    # In most filesystems, a request for the parent
 
142
                    # of root, just returns root.
 
143
                    continue
 
144
                basepath.pop()
 
145
            elif p == '.':
 
146
                continue # No-op
 
147
            else:
 
148
                basepath.append(p)
 
149
 
 
150
        path = '/'.join(basepath)
 
151
        if len(path) and path[0] != '/':
 
152
            path = '/' + path
 
153
        return path
 
154
 
 
155
    def relpath(self, abspath):
 
156
        # FIXME: this is identical to HttpTransport -- share it
 
157
        if not abspath.startswith(self.base):
 
158
            raise NonRelativePath('path %r is not under base URL %r'
 
159
                           % (abspath, self.base))
 
160
        pl = len(self.base)
 
161
        return abspath[pl:].lstrip('/')
 
162
 
 
163
    def has(self, relpath):
 
164
        """
 
165
        Does the target location exist?
 
166
        """
 
167
        try:
 
168
            self._sftp.stat(self._abspath(relpath))
 
169
            return True
 
170
        except IOError:
 
171
            return False
 
172
 
 
173
    def get(self, relpath, decode=False):
 
174
        """
 
175
        Get the file at the given relative path.
 
176
 
 
177
        :param relpath: The relative path to the file
 
178
        """
 
179
        try:
 
180
            path = self._abspath(relpath)
 
181
            f = self._sftp.file(path)
 
182
            try:
 
183
                f.prefetch()
 
184
            except AttributeError:
 
185
                # only works on paramiko 1.5.1 or greater
 
186
                pass
 
187
            return f
 
188
        except (IOError, paramiko.SSHException), x:
 
189
            raise NoSuchFile('Error retrieving %s: %s' % (path, str(x)), x)
 
190
 
 
191
    def get_partial(self, relpath, start, length=None):
 
192
        """
 
193
        Get just part of a file.
 
194
 
 
195
        :param relpath: Path to the file, relative to base
 
196
        :param start: The starting position to read from
 
197
        :param length: The length to read. A length of None indicates
 
198
                       read to the end of the file.
 
199
        :return: A file-like object containing at least the specified bytes.
 
200
                 Some implementations may return objects which can be read
 
201
                 past this length, but this is not guaranteed.
 
202
        """
 
203
        f = self.get(relpath)
 
204
        f.seek(start)
 
205
        try:
 
206
            f.prefetch()
 
207
        except AttributeError:
 
208
            # only works on paramiko 1.5.1 or greater
 
209
            pass
 
210
        return f
 
211
 
 
212
    def put(self, relpath, f):
 
213
        """
 
214
        Copy the file-like or string object into the location.
 
215
 
 
216
        :param relpath: Location to put the contents, relative to base.
 
217
        :param f:       File-like or string object.
 
218
        """
 
219
        # FIXME: should do something atomic or locking here, this is unsafe
 
220
        try:
 
221
            path = self._abspath(relpath)
 
222
            fout = self._sftp.file(path, 'wb')
 
223
        except IOError, e:
 
224
            self._translate_io_exception(e, relpath)
 
225
        except (IOError, paramiko.SSHException), x:
 
226
            raise SFTPTransportError('Unable to write file %r' % (path,), x)
 
227
        try:
 
228
            self._pump(f, fout)
 
229
        finally:
 
230
            fout.close()
 
231
 
 
232
    def iter_files_recursive(self):
 
233
        """Walk the relative paths of all files in this transport."""
 
234
        queue = list(self.list_dir('.'))
 
235
        while queue:
 
236
            relpath = urllib.quote(queue.pop(0))
 
237
            st = self.stat(relpath)
 
238
            if stat.S_ISDIR(st.st_mode):
 
239
                for i, basename in enumerate(self.list_dir(relpath)):
 
240
                    queue.insert(i, relpath+'/'+basename)
 
241
            else:
 
242
                yield relpath
 
243
 
 
244
    def mkdir(self, relpath):
 
245
        """Create a directory at the given path."""
 
246
        try:
 
247
            path = self._abspath(relpath)
 
248
            self._sftp.mkdir(path)
 
249
        except IOError, e:
 
250
            self._translate_io_exception(e, relpath)
 
251
        except (IOError, paramiko.SSHException), x:
 
252
            raise SFTPTransportError('Unable to mkdir %r' % (path,), x)
 
253
 
 
254
    def _translate_io_exception(self, e, relpath):
 
255
        # paramiko seems to generate detailless errors.
 
256
        if (e.errno == errno.ENOENT or
 
257
            e.args == ('No such file or directory',) or
 
258
            e.args == ('No such file',)):
 
259
            raise NoSuchFile(relpath)
 
260
        if (e.args == ('mkdir failed',)):
 
261
            raise FileExists(relpath)
 
262
        # strange but true, for the paramiko server.
 
263
        if (e.args == ('Failure',)):
 
264
            raise FileExists(relpath)
 
265
        raise
 
266
 
 
267
    def append(self, relpath, f):
 
268
        """
 
269
        Append the text in the file-like object into the final
 
270
        location.
 
271
        """
 
272
        try:
 
273
            path = self._abspath(relpath)
 
274
            fout = self._sftp.file(path, 'ab')
 
275
            self._pump(f, fout)
 
276
        except (IOError, paramiko.SSHException), x:
 
277
            raise SFTPTransportError('Unable to append file %r' % (path,), x)
 
278
 
 
279
    def copy(self, rel_from, rel_to):
 
280
        """Copy the item at rel_from to the location at rel_to"""
 
281
        path_from = self._abspath(rel_from)
 
282
        path_to = self._abspath(rel_to)
 
283
        try:
 
284
            fin = self._sftp.file(path_from, 'rb')
 
285
            try:
 
286
                fout = self._sftp.file(path_to, 'wb')
 
287
                try:
 
288
                    fout.set_pipelined(True)
 
289
                    self._pump(fin, fout)
 
290
                finally:
 
291
                    fout.close()
 
292
            finally:
 
293
                fin.close()
 
294
        except (IOError, paramiko.SSHException), x:
 
295
            raise SFTPTransportError('Unable to copy %r to %r' % (path_from, path_to), x)
 
296
 
 
297
    def move(self, rel_from, rel_to):
 
298
        """Move the item at rel_from to the location at rel_to"""
 
299
        path_from = self._abspath(rel_from)
 
300
        path_to = self._abspath(rel_to)
 
301
        try:
 
302
            self._sftp.rename(path_from, path_to)
 
303
        except (IOError, paramiko.SSHException), x:
 
304
            raise SFTPTransportError('Unable to move %r to %r' % (path_from, path_to), x)
 
305
 
 
306
    def delete(self, relpath):
 
307
        """Delete the item at relpath"""
 
308
        path = self._abspath(relpath)
 
309
        try:
 
310
            self._sftp.remove(path)
 
311
        except (IOError, paramiko.SSHException), x:
 
312
            raise SFTPTransportError('Unable to delete %r' % (path,), x)
 
313
            
 
314
    def listable(self):
 
315
        """Return True if this store supports listing."""
 
316
        return True
 
317
 
 
318
    def list_dir(self, relpath):
 
319
        """
 
320
        Return a list of all files at the given location.
 
321
        """
 
322
        # does anything actually use this?
 
323
        path = self._abspath(relpath)
 
324
        try:
 
325
            return self._sftp.listdir(path)
 
326
        except (IOError, paramiko.SSHException), x:
 
327
            raise SFTPTransportError('Unable to list folder %r' % (path,), x)
 
328
 
 
329
    def stat(self, relpath):
 
330
        """Return the stat information for a file."""
 
331
        path = self._abspath(relpath)
 
332
        try:
 
333
            return self._sftp.stat(path)
 
334
        except (IOError, paramiko.SSHException), x:
 
335
            raise SFTPTransportError('Unable to stat %r' % (path,), x)
 
336
 
 
337
    def lock_read(self, relpath):
 
338
        """
 
339
        Lock the given file for shared (read) access.
 
340
        :return: A lock object, which should be passed to Transport.unlock()
 
341
        """
 
342
        # FIXME: there should be something clever i can do here...
 
343
        class BogusLock(object):
 
344
            def __init__(self, path):
 
345
                self.path = path
 
346
            def unlock(self):
 
347
                pass
 
348
        return BogusLock(relpath)
 
349
 
 
350
    def lock_write(self, relpath):
 
351
        """
 
352
        Lock the given file for exclusive (write) access.
 
353
        WARNING: many transports do not support this, so trying avoid using it
 
354
 
 
355
        :return: A lock object, which should be passed to Transport.unlock()
 
356
        """
 
357
        # FIXME: there should be something clever i can do here...
 
358
        class BogusLock(object):
 
359
            def __init__(self, path):
 
360
                self.path = path
 
361
            def unlock(self):
 
362
                pass
 
363
        return BogusLock(relpath)
 
364
 
 
365
 
 
366
    def _unparse_url(self, path=None):
 
367
        if path is None:
 
368
            path = self._path
 
369
        if self._port == 22:
 
370
            return 'sftp://%s@%s%s' % (self._username, self._host, path)
 
371
        return 'sftp://%s@%s:%d%s' % (self._username, self._host, self._port, path)
 
372
 
 
373
    def _parse_url(self, url):
 
374
        assert url[:7] == 'sftp://'
 
375
        m = self._url_matcher.match(url)
 
376
        if m is None:
 
377
            raise SFTPTransportError('Unable to parse SFTP URL %r' % (url,))
 
378
        self._username, self._password, self._host, self._port, self._path = m.groups()
 
379
        if self._username is None:
 
380
            self._username = getpass.getuser()
 
381
        else:
 
382
            self._username = self._username[:-1]
 
383
        if self._password:
 
384
            self._password = self._password[1:]
 
385
            self._username = self._username[len(self._password)+1:]
 
386
        if self._port is None:
 
387
            self._port = 22
 
388
        else:
 
389
            self._port = int(self._port[1:])
 
390
        if (self._path is None) or (self._path == ''):
 
391
            self._path = '/'
 
392
 
 
393
    def _sftp_connect(self):
 
394
        global SYSTEM_HOSTKEYS, BZR_HOSTKEYS
 
395
        
 
396
        load_host_keys()
 
397
        
 
398
        try:
 
399
            t = paramiko.Transport((self._host, self._port))
 
400
            t.start_client()
 
401
        except paramiko.SSHException:
 
402
            raise SFTPTransportError('Unable to reach SSH host %s:%d' % (self._host, self._port))
 
403
            
 
404
        server_key = t.get_remote_server_key()
 
405
        server_key_hex = paramiko.util.hexify(server_key.get_fingerprint())
 
406
        keytype = server_key.get_name()
 
407
        if SYSTEM_HOSTKEYS.has_key(self._host) and SYSTEM_HOSTKEYS[self._host].has_key(keytype):
 
408
            our_server_key = SYSTEM_HOSTKEYS[self._host][keytype]
 
409
            our_server_key_hex = paramiko.util.hexify(our_server_key.get_fingerprint())
 
410
        elif BZR_HOSTKEYS.has_key(self._host) and BZR_HOSTKEYS[self._host].has_key(keytype):
 
411
            our_server_key = BZR_HOSTKEYS[self._host][keytype]
 
412
            our_server_key_hex = paramiko.util.hexify(our_server_key.get_fingerprint())
 
413
        else:
 
414
            warning('Adding %s host key for %s: %s' % (keytype, self._host, server_key_hex))
 
415
            if not BZR_HOSTKEYS.has_key(self._host):
 
416
                BZR_HOSTKEYS[self._host] = {}
 
417
            BZR_HOSTKEYS[self._host][keytype] = server_key
 
418
            our_server_key = server_key
 
419
            our_server_key_hex = paramiko.util.hexify(our_server_key.get_fingerprint())
 
420
            save_host_keys()
 
421
        if server_key != our_server_key:
 
422
            filename1 = os.path.expanduser('~/.ssh/known_hosts')
 
423
            filename2 = os.path.join(config_dir(), 'ssh_host_keys')
 
424
            raise SFTPTransportError('Host keys for %s do not match!  %s != %s' % \
 
425
                (self._host, our_server_key_hex, server_key_hex),
 
426
                ['Try editing %s or %s' % (filename1, filename2)])
 
427
 
 
428
        self._sftp_auth(t, self._username, self._host)
 
429
        
 
430
        try:
 
431
            self._sftp = t.open_sftp_client()
 
432
        except paramiko.SSHException:
 
433
            raise BzrError('Unable to find path %s on SFTP server %s' % \
 
434
                (self._path, self._host))
 
435
 
 
436
    def _sftp_auth(self, transport, username, host):
 
437
        agent = paramiko.Agent()
 
438
        for key in agent.get_keys():
 
439
            mutter('Trying SSH agent key %s' % paramiko.util.hexify(key.get_fingerprint()))
 
440
            try:
 
441
                transport.auth_publickey(self._username, key)
 
442
                return
 
443
            except paramiko.SSHException, e:
 
444
                pass
 
445
        
 
446
        # okay, try finding id_rsa or id_dss?  (posix only)
 
447
        if self._try_pkey_auth(transport, paramiko.RSAKey, 'id_rsa'):
 
448
            return
 
449
        if self._try_pkey_auth(transport, paramiko.DSSKey, 'id_dsa'):
 
450
            return
 
451
 
 
452
        if self._password:
 
453
            try:
 
454
                transport.auth_password(self._username, self._password)
 
455
                return
 
456
            except paramiko.SSHException, e:
 
457
                pass
 
458
 
 
459
        # give up and ask for a password
 
460
        password = getpass.getpass('SSH %s@%s password: ' % (self._username, self._host))
 
461
        try:
 
462
            transport.auth_password(self._username, password)
 
463
        except paramiko.SSHException:
 
464
            raise SFTPTransportError('Unable to authenticate to SSH host as %s@%s' % \
 
465
                (self._username, self._host))
 
466
 
 
467
    def _try_pkey_auth(self, transport, pkey_class, filename):
 
468
        filename = os.path.expanduser('~/.ssh/' + filename)
 
469
        try:
 
470
            key = pkey_class.from_private_key_file(filename)
 
471
            transport.auth_publickey(self._username, key)
 
472
            return True
 
473
        except paramiko.PasswordRequiredException:
 
474
            password = getpass.getpass('SSH %s password: ' % (os.path.basename(filename),))
 
475
            try:
 
476
                key = pkey_class.from_private_key_file(filename, password)
 
477
                transport.auth_publickey(self._username, key)
 
478
                return True
 
479
            except paramiko.SSHException:
 
480
                mutter('SSH authentication via %s key failed.' % (os.path.basename(filename),))
 
481
        except paramiko.SSHException:
 
482
            mutter('SSH authentication via %s key failed.' % (os.path.basename(filename),))
 
483
        except IOError:
 
484
            pass
 
485
        return False