/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 breezy/transport/ssh.py

  • Committer: Breezy landing bot
  • Author(s): Jelmer Vernooij
  • Date: 2018-11-16 18:59:44 UTC
  • mfrom: (7143.15.15 more-cleanups)
  • Revision ID: breezy.the.bot@gmail.com-20181116185944-biefv1sub37qfybm
Sprinkle some PEP8iness.

Merged from https://code.launchpad.net/~jelmer/brz/more-cleanups/+merge/358611

Show diffs side-by-side

added added

removed removed

Lines of Context:
breezy/transport/ssh.py
1
 
# Copyright (C) 2006-2010 Robey Pointer <robey@lag.net>
 
1
# Copyright (C) 2006-2011 Robey Pointer <robey@lag.net>
2
2
# Copyright (C) 2005, 2006, 2007 Canonical Ltd
3
3
#
4
4
# This program is free software; you can redistribute it and/or modify
17
17
 
18
18
"""Foundation SSH support for SFTP and smart server."""
19
19
 
 
20
from __future__ import absolute_import
 
21
 
20
22
import errno
21
23
import getpass
22
24
import logging
24
26
import socket
25
27
import subprocess
26
28
import sys
 
29
from binascii import hexlify
27
30
 
28
 
from bzrlib import (
 
31
from .. import (
29
32
    config,
30
33
    errors,
31
34
    osutils,
35
38
 
36
39
try:
37
40
    import paramiko
38
 
except ImportError, e:
 
41
except ImportError as e:
39
42
    # If we have an ssh subprocess, we don't strictly need paramiko for all ssh
40
43
    # access
41
44
    paramiko = None
43
46
    from paramiko.sftp_client import SFTPClient
44
47
 
45
48
 
 
49
class StrangeHostname(errors.BzrError):
 
50
    _fmt = "Refusing to connect to strange SSH hostname %(hostname)s"
 
51
 
 
52
 
46
53
SYSTEM_HOSTKEYS = {}
47
 
BZR_HOSTKEYS = {}
48
 
 
49
 
 
50
 
_paramiko_version = getattr(paramiko, '__version_info__', (0, 0, 0))
51
 
 
52
 
# Paramiko 1.5 tries to open a socket.AF_UNIX in order to connect
53
 
# to ssh-agent. That attribute doesn't exist on win32 (it does in cygwin)
54
 
# so we get an AttributeError exception. So we will not try to
55
 
# connect to an agent if we are on win32 and using Paramiko older than 1.6
56
 
_use_ssh_agent = (sys.platform != 'win32' or _paramiko_version >= (1, 6, 0))
 
54
BRZ_HOSTKEYS = {}
57
55
 
58
56
 
59
57
class SSHVendorManager(object):
60
58
    """Manager for manage SSH vendors."""
61
59
 
62
60
    # Note, although at first sign the class interface seems similar to
63
 
    # bzrlib.registry.Registry it is not possible/convenient to directly use
 
61
    # breezy.registry.Registry it is not possible/convenient to directly use
64
62
    # the Registry because the class just has "get()" interface instead of the
65
63
    # Registry's "get(key)".
66
64
 
82
80
        self._cached_ssh_vendor = None
83
81
 
84
82
    def _get_vendor_by_environment(self, environment=None):
85
 
        """Return the vendor or None based on BZR_SSH environment variable.
 
83
        """Return the vendor or None based on BRZ_SSH environment variable.
86
84
 
87
 
        :raises UnknownSSH: if the BZR_SSH environment variable contains
 
85
        :raises UnknownSSH: if the BRZ_SSH environment variable contains
88
86
                            unknown vendor name
89
87
        """
90
88
        if environment is None:
91
89
            environment = os.environ
92
 
        if 'BZR_SSH' in environment:
93
 
            vendor_name = environment['BZR_SSH']
 
90
        if 'BRZ_SSH' in environment:
 
91
            vendor_name = environment['BRZ_SSH']
94
92
            try:
95
93
                vendor = self._ssh_vendors[vendor_name]
96
94
            except KeyError:
107
105
            p = subprocess.Popen(args,
108
106
                                 stdout=subprocess.PIPE,
109
107
                                 stderr=subprocess.PIPE,
 
108
                                 bufsize=0,
110
109
                                 **os_specific_subprocess_params())
111
110
            stdout, stderr = p.communicate()
112
111
        except OSError:
113
 
            stdout = stderr = ''
114
 
        return stdout + stderr
 
112
            stdout = stderr = b''
 
113
        return (stdout + stderr).decode(osutils.get_terminal_encoding())
115
114
 
116
115
    def _get_vendor_by_version_string(self, version, progname):
117
116
        """Return the vendor or None based on output from the subprocess.
126
125
        elif 'SSH Secure Shell' in version:
127
126
            trace.mutter('ssh implementation is SSH Corp.')
128
127
            vendor = SSHCorpSubprocessVendor()
 
128
        elif 'lsh' in version:
 
129
            trace.mutter('ssh implementation is GNU lsh.')
 
130
            vendor = LSHSubprocessVendor()
129
131
        # As plink user prompts are not handled currently, don't auto-detect
130
132
        # it by inspection below, but keep this vendor detection for if a path
131
 
        # is given in BZR_SSH. See https://bugs.launchpad.net/bugs/414743
 
133
        # is given in BRZ_SSH. See https://bugs.launchpad.net/bugs/414743
132
134
        elif 'plink' in version and progname == 'plink':
133
135
            # Checking if "plink" was the executed argument as Windows
134
 
            # sometimes reports 'ssh -V' incorrectly with 'plink' in it's
 
136
            # sometimes reports 'ssh -V' incorrectly with 'plink' in its
135
137
            # version.  See https://bugs.launchpad.net/bzr/+bug/107155
136
138
            trace.mutter("ssh implementation is Putty's plink.")
137
139
            vendor = PLinkSubprocessVendor()
145
147
    def _get_vendor_from_path(self, path):
146
148
        """Return the vendor or None using the program at the given path"""
147
149
        version = self._get_ssh_version_string([path, '-V'])
148
 
        return self._get_vendor_by_version_string(version, 
149
 
            os.path.splitext(os.path.basename(path))[0])
 
150
        return self._get_vendor_by_version_string(version,
 
151
                                                  os.path.splitext(os.path.basename(path))[0])
150
152
 
151
153
    def get_vendor(self, environment=None):
152
154
        """Find out what version of SSH is on the system.
153
155
 
154
156
        :raises SSHVendorNotFound: if no any SSH vendor is found
155
 
        :raises UnknownSSH: if the BZR_SSH environment variable contains
 
157
        :raises UnknownSSH: if the BRZ_SSH environment variable contains
156
158
                            unknown vendor name
157
159
        """
158
160
        if self._cached_ssh_vendor is None:
167
169
            self._cached_ssh_vendor = vendor
168
170
        return self._cached_ssh_vendor
169
171
 
 
172
 
170
173
_ssh_vendor_manager = SSHVendorManager()
171
174
_get_ssh_vendor = _ssh_vendor_manager.get_vendor
172
175
register_default_ssh_vendor = _ssh_vendor_manager.register_default_vendor
199
202
    def recv(self, n):
200
203
        try:
201
204
            return self.__socket.recv(n)
202
 
        except socket.error, e:
 
205
        except socket.error as e:
203
206
            if e.args[0] in (errno.EPIPE, errno.ECONNRESET, errno.ECONNABORTED,
204
207
                             errno.EBADF):
205
208
                # Connection has closed.  Paramiko expects an empty string in
239
242
    def connect_ssh(self, username, password, host, port, command):
240
243
        """Make an SSH connection.
241
244
 
242
 
        :returns: something with a `close` method, and a `get_filelike_channels`
243
 
            method that returns a pair of (read, write) filelike objects.
 
245
        :returns: an SSHConnection.
244
246
        """
245
247
        raise NotImplementedError(self.connect_ssh)
246
248
 
262
264
        sock = socket.socket()
263
265
        try:
264
266
            sock.connect((host, port))
265
 
        except socket.error, e:
 
267
        except socket.error as e:
266
268
            self._raise_connection_error(host, port=port, orig_error=e)
267
269
        return SFTPClient(SocketAsChannelAdapter(sock))
268
270
 
 
271
 
269
272
register_ssh_vendor('loopback', LoopbackVendor())
270
273
 
271
274
 
272
 
class _ParamikoSSHConnection(object):
273
 
    def __init__(self, channel):
274
 
        self.channel = channel
275
 
 
276
 
    def get_filelike_channels(self):
277
 
        return self.channel.makefile('rb'), self.channel.makefile('wb')
278
 
 
279
 
    def close(self):
280
 
        return self.channel.close()
281
 
 
282
 
 
283
275
class ParamikoVendor(SSHVendor):
284
276
    """Vendor that uses paramiko."""
285
277
 
 
278
    def _hexify(self, s):
 
279
        return hexlify(s).upper()
 
280
 
286
281
    def _connect(self, username, password, host, port):
287
 
        global SYSTEM_HOSTKEYS, BZR_HOSTKEYS
 
282
        global SYSTEM_HOSTKEYS, BRZ_HOSTKEYS
288
283
 
289
284
        load_host_keys()
290
285
 
292
287
            t = paramiko.Transport((host, port or 22))
293
288
            t.set_log_channel('bzr.paramiko')
294
289
            t.start_client()
295
 
        except (paramiko.SSHException, socket.error), e:
 
290
        except (paramiko.SSHException, socket.error) as e:
296
291
            self._raise_connection_error(host, port=port, orig_error=e)
297
292
 
298
293
        server_key = t.get_remote_server_key()
299
 
        server_key_hex = paramiko.util.hexify(server_key.get_fingerprint())
 
294
        server_key_hex = self._hexify(server_key.get_fingerprint())
300
295
        keytype = server_key.get_name()
301
296
        if host in SYSTEM_HOSTKEYS and keytype in SYSTEM_HOSTKEYS[host]:
302
297
            our_server_key = SYSTEM_HOSTKEYS[host][keytype]
303
 
            our_server_key_hex = paramiko.util.hexify(
304
 
                our_server_key.get_fingerprint())
305
 
        elif host in BZR_HOSTKEYS and keytype in BZR_HOSTKEYS[host]:
306
 
            our_server_key = BZR_HOSTKEYS[host][keytype]
307
 
            our_server_key_hex = paramiko.util.hexify(
308
 
                our_server_key.get_fingerprint())
 
298
            our_server_key_hex = self._hexify(our_server_key.get_fingerprint())
 
299
        elif host in BRZ_HOSTKEYS and keytype in BRZ_HOSTKEYS[host]:
 
300
            our_server_key = BRZ_HOSTKEYS[host][keytype]
 
301
            our_server_key_hex = self._hexify(our_server_key.get_fingerprint())
309
302
        else:
310
303
            trace.warning('Adding %s host key for %s: %s'
311
304
                          % (keytype, host, server_key_hex))
312
 
            add = getattr(BZR_HOSTKEYS, 'add', None)
313
 
            if add is not None: # paramiko >= 1.X.X
314
 
                BZR_HOSTKEYS.add(host, keytype, server_key)
 
305
            add = getattr(BRZ_HOSTKEYS, 'add', None)
 
306
            if add is not None:  # paramiko >= 1.X.X
 
307
                BRZ_HOSTKEYS.add(host, keytype, server_key)
315
308
            else:
316
 
                BZR_HOSTKEYS.setdefault(host, {})[keytype] = server_key
 
309
                BRZ_HOSTKEYS.setdefault(host, {})[keytype] = server_key
317
310
            our_server_key = server_key
318
 
            our_server_key_hex = paramiko.util.hexify(
319
 
                our_server_key.get_fingerprint())
 
311
            our_server_key_hex = self._hexify(our_server_key.get_fingerprint())
320
312
            save_host_keys()
321
313
        if server_key != our_server_key:
322
314
            filename1 = os.path.expanduser('~/.ssh/known_hosts')
333
325
        t = self._connect(username, password, host, port)
334
326
        try:
335
327
            return t.open_sftp_client()
336
 
        except paramiko.SSHException, e:
 
328
        except paramiko.SSHException as e:
337
329
            self._raise_connection_error(host, port=port, orig_error=e,
338
330
                                         msg='Unable to start sftp client')
339
331
 
344
336
            cmdline = ' '.join(command)
345
337
            channel.exec_command(cmdline)
346
338
            return _ParamikoSSHConnection(channel)
347
 
        except paramiko.SSHException, e:
 
339
        except paramiko.SSHException as e:
348
340
            self._raise_connection_error(host, port=port, orig_error=e,
349
341
                                         msg='Unable to invoke remote bzr')
350
342
 
 
343
 
 
344
_ssh_connection_errors = (EOFError, OSError, IOError, socket.error)
351
345
if paramiko is not None:
352
346
    vendor = ParamikoVendor()
353
347
    register_ssh_vendor('paramiko', vendor)
354
348
    register_ssh_vendor('none', vendor)
355
349
    register_default_ssh_vendor(vendor)
356
 
    _sftp_connection_errors = (EOFError, paramiko.SSHException)
 
350
    _ssh_connection_errors += (paramiko.SSHException,)
357
351
    del vendor
358
 
else:
359
 
    _sftp_connection_errors = (EOFError,)
360
352
 
361
353
 
362
354
class SubprocessVendor(SSHVendor):
363
355
    """Abstract base class for vendors that use pipes to a subprocess."""
364
356
 
 
357
    # In general stderr should be inherited from the parent process so prompts
 
358
    # are visible on the terminal. This can be overriden to another file for
 
359
    # tests, but beware of using PIPE which may hang due to not being read.
 
360
    _stderr_target = None
 
361
 
 
362
    @staticmethod
 
363
    def _check_hostname(arg):
 
364
        if arg.startswith('-'):
 
365
            raise StrangeHostname(hostname=arg)
 
366
 
365
367
    def _connect(self, argv):
366
 
        proc = subprocess.Popen(argv,
367
 
                                stdin=subprocess.PIPE,
368
 
                                stdout=subprocess.PIPE,
 
368
        # Attempt to make a socketpair to use as stdin/stdout for the SSH
 
369
        # subprocess.  We prefer sockets to pipes because they support
 
370
        # non-blocking short reads, allowing us to optimistically read 64k (or
 
371
        # whatever) chunks.
 
372
        try:
 
373
            my_sock, subproc_sock = socket.socketpair()
 
374
            osutils.set_fd_cloexec(my_sock)
 
375
        except (AttributeError, socket.error):
 
376
            # This platform doesn't support socketpair(), so just use ordinary
 
377
            # pipes instead.
 
378
            stdin = stdout = subprocess.PIPE
 
379
            my_sock, subproc_sock = None, None
 
380
        else:
 
381
            stdin = stdout = subproc_sock
 
382
        proc = subprocess.Popen(argv, stdin=stdin, stdout=stdout,
 
383
                                stderr=self._stderr_target,
 
384
                                bufsize=0,
369
385
                                **os_specific_subprocess_params())
370
 
        return SSHSubprocess(proc)
 
386
        if subproc_sock is not None:
 
387
            subproc_sock.close()
 
388
        return SSHSubprocessConnection(proc, sock=my_sock)
371
389
 
372
390
    def connect_sftp(self, username, password, host, port):
373
391
        try:
375
393
                                                  subsystem='sftp')
376
394
            sock = self._connect(argv)
377
395
            return SFTPClient(SocketAsChannelAdapter(sock))
378
 
        except _sftp_connection_errors, e:
379
 
            self._raise_connection_error(host, port=port, orig_error=e)
380
 
        except (OSError, IOError), e:
381
 
            # If the machine is fast enough, ssh can actually exit
382
 
            # before we try and send it the sftp request, which
383
 
            # raises a Broken Pipe
384
 
            if e.errno not in (errno.EPIPE,):
385
 
                raise
 
396
        except _ssh_connection_errors as e:
386
397
            self._raise_connection_error(host, port=port, orig_error=e)
387
398
 
388
399
    def connect_ssh(self, username, password, host, port, command):
390
401
            argv = self._get_vendor_specific_argv(username, host, port,
391
402
                                                  command=command)
392
403
            return self._connect(argv)
393
 
        except (EOFError), e:
394
 
            self._raise_connection_error(host, port=port, orig_error=e)
395
 
        except (OSError, IOError), e:
396
 
            # If the machine is fast enough, ssh can actually exit
397
 
            # before we try and send it the sftp request, which
398
 
            # raises a Broken Pipe
399
 
            if e.errno not in (errno.EPIPE,):
400
 
                raise
 
404
        except _ssh_connection_errors as e:
401
405
            self._raise_connection_error(host, port=port, orig_error=e)
402
406
 
403
407
    def _get_vendor_specific_argv(self, username, host, port, subsystem=None,
418
422
                                  command=None):
419
423
        args = [self.executable_path,
420
424
                '-oForwardX11=no', '-oForwardAgent=no',
421
 
                '-oClearAllForwardings=yes', '-oProtocol=2',
 
425
                '-oClearAllForwardings=yes',
422
426
                '-oNoHostAuthenticationForLocalhost=yes']
423
427
        if port is not None:
424
428
            args.extend(['-p', str(port)])
425
429
        if username is not None:
426
430
            args.extend(['-l', username])
427
431
        if subsystem is not None:
428
 
            args.extend(['-s', host, subsystem])
 
432
            args.extend(['-s', '--', host, subsystem])
429
433
        else:
430
 
            args.extend([host] + command)
 
434
            args.extend(['--', host] + command)
431
435
        return args
432
436
 
 
437
 
433
438
register_ssh_vendor('openssh', OpenSSHSubprocessVendor())
434
439
 
435
440
 
440
445
 
441
446
    def _get_vendor_specific_argv(self, username, host, port, subsystem=None,
442
447
                                  command=None):
 
448
        self._check_hostname(host)
443
449
        args = [self.executable_path, '-x']
444
450
        if port is not None:
445
451
            args.extend(['-p', str(port)])
451
457
            args.extend([host] + command)
452
458
        return args
453
459
 
 
460
 
454
461
register_ssh_vendor('sshcorp', SSHCorpSubprocessVendor())
455
462
 
456
463
 
 
464
class LSHSubprocessVendor(SubprocessVendor):
 
465
    """SSH vendor that uses the 'lsh' executable from GNU"""
 
466
 
 
467
    executable_path = 'lsh'
 
468
 
 
469
    def _get_vendor_specific_argv(self, username, host, port, subsystem=None,
 
470
                                  command=None):
 
471
        self._check_hostname(host)
 
472
        args = [self.executable_path]
 
473
        if port is not None:
 
474
            args.extend(['-p', str(port)])
 
475
        if username is not None:
 
476
            args.extend(['-l', username])
 
477
        if subsystem is not None:
 
478
            args.extend(['--subsystem', subsystem, host])
 
479
        else:
 
480
            args.extend([host] + command)
 
481
        return args
 
482
 
 
483
 
 
484
register_ssh_vendor('lsh', LSHSubprocessVendor())
 
485
 
 
486
 
457
487
class PLinkSubprocessVendor(SubprocessVendor):
458
488
    """SSH vendor that uses the 'plink' executable from Putty."""
459
489
 
461
491
 
462
492
    def _get_vendor_specific_argv(self, username, host, port, subsystem=None,
463
493
                                  command=None):
 
494
        self._check_hostname(host)
464
495
        args = [self.executable_path, '-x', '-a', '-ssh', '-2', '-batch']
465
496
        if port is not None:
466
497
            args.extend(['-P', str(port)])
472
503
            args.extend([host] + command)
473
504
        return args
474
505
 
 
506
 
475
507
register_ssh_vendor('plink', PLinkSubprocessVendor())
476
508
 
477
509
 
482
514
    if username is None:
483
515
        username = auth.get_user('ssh', host, port=port,
484
516
                                 default=getpass.getuser())
485
 
    if _use_ssh_agent:
486
 
        agent = paramiko.Agent()
487
 
        for key in agent.get_keys():
488
 
            trace.mutter('Trying SSH agent key %s'
489
 
                         % paramiko.util.hexify(key.get_fingerprint()))
490
 
            try:
491
 
                paramiko_transport.auth_publickey(username, key)
492
 
                return
493
 
            except paramiko.SSHException, e:
494
 
                pass
 
517
    agent = paramiko.Agent()
 
518
    for key in agent.get_keys():
 
519
        trace.mutter('Trying SSH agent key %s'
 
520
                     % self._hexify(key.get_fingerprint()))
 
521
        try:
 
522
            paramiko_transport.auth_publickey(username, key)
 
523
            return
 
524
        except paramiko.SSHException as e:
 
525
            pass
495
526
 
496
527
    # okay, try finding id_rsa or id_dss?  (posix only)
497
528
    if _try_pkey_auth(paramiko_transport, paramiko.RSAKey, username, 'id_rsa'):
512
543
            paramiko_transport.auth_none(username)
513
544
        finally:
514
545
            paramiko_transport.logger.setLevel(old_level)
515
 
    except paramiko.BadAuthenticationType, e:
 
546
    except paramiko.BadAuthenticationType as e:
516
547
        # Supported methods are in the exception
517
548
        supported_auth_types = e.allowed_types
518
 
    except paramiko.SSHException, e:
 
549
    except paramiko.SSHException as e:
519
550
        # Don't know what happened, but just ignore it
520
551
        pass
521
552
    # We treat 'keyboard-interactive' and 'password' auth methods identically,
527
558
    # requires something other than a single password, but we currently don't
528
559
    # support that.
529
560
    if ('password' not in supported_auth_types and
530
 
        'keyboard-interactive' not in supported_auth_types):
 
561
            'keyboard-interactive' not in supported_auth_types):
531
562
        raise errors.ConnectionError('Unable to authenticate to SSH host as'
532
 
            '\n  %s@%s\nsupported auth types: %s'
533
 
            % (username, host, supported_auth_types))
 
563
                                     '\n  %s@%s\nsupported auth types: %s'
 
564
                                     % (username, host, supported_auth_types))
534
565
 
535
566
    if password:
536
567
        try:
537
568
            paramiko_transport.auth_password(username, password)
538
569
            return
539
 
        except paramiko.SSHException, e:
 
570
        except paramiko.SSHException as e:
540
571
            pass
541
572
 
542
573
    # give up and ask for a password
545
576
    if password is not None:
546
577
        try:
547
578
            paramiko_transport.auth_password(username, password)
548
 
        except paramiko.SSHException, e:
 
579
        except paramiko.SSHException as e:
549
580
            raise errors.ConnectionError(
550
581
                'Unable to authenticate to SSH host as'
551
582
                '\n  %s@%s\n' % (username, host), e)
562
593
        return True
563
594
    except paramiko.PasswordRequiredException:
564
595
        password = ui.ui_factory.get_password(
565
 
            prompt='SSH %(filename)s password', filename=filename)
 
596
            prompt=u'SSH %(filename)s password',
 
597
            filename=filename.decode(osutils._fs_enc))
566
598
        try:
567
599
            key = pkey_class.from_private_key_file(filename, password)
568
600
            paramiko_transport.auth_publickey(username, key)
583
615
    Load system host keys (probably doesn't work on windows) and any
584
616
    "discovered" keys from previous sessions.
585
617
    """
586
 
    global SYSTEM_HOSTKEYS, BZR_HOSTKEYS
 
618
    global SYSTEM_HOSTKEYS, BRZ_HOSTKEYS
587
619
    try:
588
620
        SYSTEM_HOSTKEYS = paramiko.util.load_host_keys(
589
621
            os.path.expanduser('~/.ssh/known_hosts'))
590
 
    except IOError, e:
 
622
    except IOError as e:
591
623
        trace.mutter('failed to load system host keys: ' + str(e))
592
 
    bzr_hostkey_path = osutils.pathjoin(config.config_dir(), 'ssh_host_keys')
 
624
    brz_hostkey_path = osutils.pathjoin(config.config_dir(), 'ssh_host_keys')
593
625
    try:
594
 
        BZR_HOSTKEYS = paramiko.util.load_host_keys(bzr_hostkey_path)
595
 
    except IOError, e:
596
 
        trace.mutter('failed to load bzr host keys: ' + str(e))
 
626
        BRZ_HOSTKEYS = paramiko.util.load_host_keys(brz_hostkey_path)
 
627
    except IOError as e:
 
628
        trace.mutter('failed to load brz host keys: ' + str(e))
597
629
        save_host_keys()
598
630
 
599
631
 
601
633
    """
602
634
    Save "discovered" host keys in $(config)/ssh_host_keys/.
603
635
    """
604
 
    global SYSTEM_HOSTKEYS, BZR_HOSTKEYS
 
636
    global SYSTEM_HOSTKEYS, BRZ_HOSTKEYS
605
637
    bzr_hostkey_path = osutils.pathjoin(config.config_dir(), 'ssh_host_keys')
606
638
    config.ensure_config_dir_exists()
607
639
 
608
640
    try:
609
 
        f = open(bzr_hostkey_path, 'w')
610
 
        f.write('# SSH host keys collected by bzr\n')
611
 
        for hostname, keys in BZR_HOSTKEYS.iteritems():
612
 
            for keytype, key in keys.iteritems():
613
 
                f.write('%s %s %s\n' % (hostname, keytype, key.get_base64()))
614
 
        f.close()
615
 
    except IOError, e:
 
641
        with open(bzr_hostkey_path, 'w') as f:
 
642
            f.write('# SSH host keys collected by bzr\n')
 
643
            for hostname, keys in BRZ_HOSTKEYS.items():
 
644
                for keytype, key in keys.items():
 
645
                    f.write('%s %s %s\n' %
 
646
                            (hostname, keytype, key.get_base64()))
 
647
    except IOError as e:
616
648
        trace.mutter('failed to save bzr host keys: ' + str(e))
617
649
 
618
650
 
641
673
                'close_fds': True,
642
674
                }
643
675
 
 
676
 
644
677
import weakref
645
678
_subproc_weakrefs = set()
646
679
 
647
 
def _close_ssh_proc(proc):
648
 
    for func in [proc.stdin.close, proc.stdout.close, proc.wait]:
 
680
 
 
681
def _close_ssh_proc(proc, sock):
 
682
    """Carefully close stdin/stdout and reap the SSH process.
 
683
 
 
684
    If the pipes are already closed and/or the process has already been
 
685
    wait()ed on, that's ok, and no error is raised.  The goal is to do our best
 
686
    to clean up (whether or not a clean up was already tried).
 
687
    """
 
688
    funcs = []
 
689
    for closeable in (proc.stdin, proc.stdout, sock):
 
690
        # We expect that either proc (a subprocess.Popen) will have stdin and
 
691
        # stdout streams to close, or that we will have been passed a socket to
 
692
        # close, with the option not in use being None.
 
693
        if closeable is not None:
 
694
            funcs.append(closeable.close)
 
695
    funcs.append(proc.wait)
 
696
    for func in funcs:
649
697
        try:
650
698
            func()
651
699
        except OSError:
652
 
            pass
653
 
 
654
 
 
655
 
class SSHSubprocess(object):
656
 
    """A socket-like object that talks to an ssh subprocess via pipes."""
657
 
 
658
 
    def __init__(self, proc):
 
700
            # It's ok for the pipe to already be closed, or the process to
 
701
            # already be finished.
 
702
            continue
 
703
 
 
704
 
 
705
class SSHConnection(object):
 
706
    """Abstract base class for SSH connections."""
 
707
 
 
708
    def get_sock_or_pipes(self):
 
709
        """Returns a (kind, io_object) pair.
 
710
 
 
711
        If kind == 'socket', then io_object is a socket.
 
712
 
 
713
        If kind == 'pipes', then io_object is a pair of file-like objects
 
714
        (read_from, write_to).
 
715
        """
 
716
        raise NotImplementedError(self.get_sock_or_pipes)
 
717
 
 
718
    def close(self):
 
719
        raise NotImplementedError(self.close)
 
720
 
 
721
 
 
722
class SSHSubprocessConnection(SSHConnection):
 
723
    """A connection to an ssh subprocess via pipes or a socket.
 
724
 
 
725
    This class is also socket-like enough to be used with
 
726
    SocketAsChannelAdapter (it has 'send' and 'recv' methods).
 
727
    """
 
728
 
 
729
    def __init__(self, proc, sock=None):
 
730
        """Constructor.
 
731
 
 
732
        :param proc: a subprocess.Popen
 
733
        :param sock: if proc.stdin/out is a socket from a socketpair, then sock
 
734
            should breezy's half of that socketpair.  If not passed, proc's
 
735
            stdin/out is assumed to be ordinary pipes.
 
736
        """
659
737
        self.proc = proc
 
738
        self._sock = sock
660
739
        # Add a weakref to proc that will attempt to do the same as self.close
661
740
        # to avoid leaving processes lingering indefinitely.
 
741
 
662
742
        def terminate(ref):
663
743
            _subproc_weakrefs.remove(ref)
664
 
            _close_ssh_proc(proc)
 
744
            _close_ssh_proc(proc, sock)
665
745
        _subproc_weakrefs.add(weakref.ref(self, terminate))
666
746
 
667
747
    def send(self, data):
668
 
        return os.write(self.proc.stdin.fileno(), data)
 
748
        if self._sock is not None:
 
749
            return self._sock.send(data)
 
750
        else:
 
751
            return os.write(self.proc.stdin.fileno(), data)
669
752
 
670
753
    def recv(self, count):
671
 
        return os.read(self.proc.stdout.fileno(), count)
672
 
 
673
 
    def close(self):
674
 
        _close_ssh_proc(self.proc)
675
 
 
676
 
    def get_filelike_channels(self):
677
 
        return (self.proc.stdout, self.proc.stdin)
678
 
 
 
754
        if self._sock is not None:
 
755
            return self._sock.recv(count)
 
756
        else:
 
757
            return os.read(self.proc.stdout.fileno(), count)
 
758
 
 
759
    def close(self):
 
760
        _close_ssh_proc(self.proc, self._sock)
 
761
 
 
762
    def get_sock_or_pipes(self):
 
763
        if self._sock is not None:
 
764
            return 'socket', self._sock
 
765
        else:
 
766
            return 'pipes', (self.proc.stdout, self.proc.stdin)
 
767
 
 
768
 
 
769
class _ParamikoSSHConnection(SSHConnection):
 
770
    """An SSH connection via paramiko."""
 
771
 
 
772
    def __init__(self, channel):
 
773
        self.channel = channel
 
774
 
 
775
    def get_sock_or_pipes(self):
 
776
        return ('socket', self.channel)
 
777
 
 
778
    def close(self):
 
779
        return self.channel.close()