1# Copyright (C) 2006-2011 Robey Pointer <robey@lag.net> 2# Copyright (C) 2005, 2006, 2007 Canonical Ltd 3# 4# This program is free software; you can redistribute it and/or modify 5# it under the terms of the GNU General Public License as published by 6# the Free Software Foundation; either version 2 of the License, or 7# (at your option) any later version. 8# 9# This program 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 12# GNU General Public License for more details. 13# 14# You should have received a copy of the GNU General Public License 15# along with this program; if not, write to the Free Software 16# Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA 17 18"""Foundation SSH support for SFTP and smart server.""" 19 20import errno 21import getpass 22import logging 23import os 24import socket 25import subprocess 26import sys 27from binascii import hexlify 28 29from .. import ( 30 config, 31 bedding, 32 errors, 33 osutils, 34 trace, 35 ui, 36 ) 37 38try: 39 import paramiko 40except ImportError as e: 41 # If we have an ssh subprocess, we don't strictly need paramiko for all ssh 42 # access 43 paramiko = None 44else: 45 from paramiko.sftp_client import SFTPClient 46 47 48class StrangeHostname(errors.BzrError): 49 _fmt = "Refusing to connect to strange SSH hostname %(hostname)s" 50 51 52SYSTEM_HOSTKEYS = {} 53BRZ_HOSTKEYS = {} 54 55 56class SSHVendorManager(object): 57 """Manager for manage SSH vendors.""" 58 59 # Note, although at first sign the class interface seems similar to 60 # breezy.registry.Registry it is not possible/convenient to directly use 61 # the Registry because the class just has "get()" interface instead of the 62 # Registry's "get(key)". 63 64 def __init__(self): 65 self._ssh_vendors = {} 66 self._cached_ssh_vendor = None 67 self._default_ssh_vendor = None 68 69 def register_default_vendor(self, vendor): 70 """Register default SSH vendor.""" 71 self._default_ssh_vendor = vendor 72 73 def register_vendor(self, name, vendor): 74 """Register new SSH vendor by name.""" 75 self._ssh_vendors[name] = vendor 76 77 def clear_cache(self): 78 """Clear previously cached lookup result.""" 79 self._cached_ssh_vendor = None 80 81 def _get_vendor_by_config(self): 82 vendor_name = config.GlobalStack().get('ssh') 83 if vendor_name is not None: 84 try: 85 vendor = self._ssh_vendors[vendor_name] 86 except KeyError: 87 vendor = self._get_vendor_from_path(vendor_name) 88 if vendor is None: 89 raise errors.UnknownSSH(vendor_name) 90 vendor.executable_path = vendor_name 91 return vendor 92 return None 93 94 def _get_ssh_version_string(self, args): 95 """Return SSH version string from the subprocess.""" 96 try: 97 p = subprocess.Popen(args, 98 stdout=subprocess.PIPE, 99 stderr=subprocess.PIPE, 100 bufsize=0, 101 **os_specific_subprocess_params()) 102 stdout, stderr = p.communicate() 103 except OSError: 104 stdout = stderr = b'' 105 return (stdout + stderr).decode(osutils.get_terminal_encoding()) 106 107 def _get_vendor_by_version_string(self, version, progname): 108 """Return the vendor or None based on output from the subprocess. 109 110 :param version: The output of 'ssh -V' like command. 111 :param args: Command line that was run. 112 """ 113 vendor = None 114 if 'OpenSSH' in version: 115 trace.mutter('ssh implementation is OpenSSH') 116 vendor = OpenSSHSubprocessVendor() 117 elif 'SSH Secure Shell' in version: 118 trace.mutter('ssh implementation is SSH Corp.') 119 vendor = SSHCorpSubprocessVendor() 120 elif 'lsh' in version: 121 trace.mutter('ssh implementation is GNU lsh.') 122 vendor = LSHSubprocessVendor() 123 # As plink user prompts are not handled currently, don't auto-detect 124 # it by inspection below, but keep this vendor detection for if a path 125 # is given in BRZ_SSH. See https://bugs.launchpad.net/bugs/414743 126 elif 'plink' in version and progname == 'plink': 127 # Checking if "plink" was the executed argument as Windows 128 # sometimes reports 'ssh -V' incorrectly with 'plink' in its 129 # version. See https://bugs.launchpad.net/bzr/+bug/107155 130 trace.mutter("ssh implementation is Putty's plink.") 131 vendor = PLinkSubprocessVendor() 132 return vendor 133 134 def _get_vendor_by_inspection(self): 135 """Return the vendor or None by checking for known SSH implementations.""" 136 version = self._get_ssh_version_string(['ssh', '-V']) 137 return self._get_vendor_by_version_string(version, "ssh") 138 139 def _get_vendor_from_path(self, path): 140 """Return the vendor or None using the program at the given path""" 141 version = self._get_ssh_version_string([path, '-V']) 142 return self._get_vendor_by_version_string(version, 143 os.path.splitext(os.path.basename(path))[0]) 144 145 def get_vendor(self): 146 """Find out what version of SSH is on the system. 147 148 :raises SSHVendorNotFound: if no any SSH vendor is found 149 :raises UnknownSSH: if the BRZ_SSH environment variable contains 150 unknown vendor name 151 """ 152 if self._cached_ssh_vendor is None: 153 vendor = self._get_vendor_by_config() 154 if vendor is None: 155 vendor = self._get_vendor_by_inspection() 156 if vendor is None: 157 trace.mutter('falling back to default implementation') 158 vendor = self._default_ssh_vendor 159 if vendor is None: 160 raise errors.SSHVendorNotFound() 161 self._cached_ssh_vendor = vendor 162 return self._cached_ssh_vendor 163 164 165_ssh_vendor_manager = SSHVendorManager() 166_get_ssh_vendor = _ssh_vendor_manager.get_vendor 167register_default_ssh_vendor = _ssh_vendor_manager.register_default_vendor 168register_ssh_vendor = _ssh_vendor_manager.register_vendor 169 170 171def _ignore_signals(): 172 # TODO: This should possibly ignore SIGHUP as well, but bzr currently 173 # doesn't handle it itself. 174 # <https://launchpad.net/products/bzr/+bug/41433/+index> 175 import signal 176 signal.signal(signal.SIGINT, signal.SIG_IGN) 177 # GZ 2010-02-19: Perhaps make this check if breakin is installed instead 178 if signal.getsignal(signal.SIGQUIT) != signal.SIG_DFL: 179 signal.signal(signal.SIGQUIT, signal.SIG_IGN) 180 181 182class SocketAsChannelAdapter(object): 183 """Simple wrapper for a socket that pretends to be a paramiko Channel.""" 184 185 def __init__(self, sock): 186 self.__socket = sock 187 188 def get_name(self): 189 return "bzr SocketAsChannelAdapter" 190 191 def send(self, data): 192 return self.__socket.send(data) 193 194 def recv(self, n): 195 try: 196 return self.__socket.recv(n) 197 except socket.error as e: 198 if e.args[0] in (errno.EPIPE, errno.ECONNRESET, errno.ECONNABORTED, 199 errno.EBADF): 200 # Connection has closed. Paramiko expects an empty string in 201 # this case, not an exception. 202 return '' 203 raise 204 205 def recv_ready(self): 206 # TODO: jam 20051215 this function is necessary to support the 207 # pipelined() function. In reality, it probably should use 208 # poll() or select() to actually return if there is data 209 # available, otherwise we probably don't get any benefit 210 return True 211 212 def close(self): 213 self.__socket.close() 214 215 216class SSHVendor(object): 217 """Abstract base class for SSH vendor implementations.""" 218 219 def connect_sftp(self, username, password, host, port): 220 """Make an SSH connection, and return an SFTPClient. 221 222 :param username: an ascii string 223 :param password: an ascii string 224 :param host: a host name as an ascii string 225 :param port: a port number 226 :type port: int 227 228 :raises: ConnectionError if it cannot connect. 229 230 :rtype: paramiko.sftp_client.SFTPClient 231 """ 232 raise NotImplementedError(self.connect_sftp) 233 234 def connect_ssh(self, username, password, host, port, command): 235 """Make an SSH connection. 236 237 :returns: an SSHConnection. 238 """ 239 raise NotImplementedError(self.connect_ssh) 240 241 def _raise_connection_error(self, host, port=None, orig_error=None, 242 msg='Unable to connect to SSH host'): 243 """Raise a SocketConnectionError with properly formatted host. 244 245 This just unifies all the locations that try to raise ConnectionError, 246 so that they format things properly. 247 """ 248 raise errors.SocketConnectionError(host=host, port=port, msg=msg, 249 orig_error=orig_error) 250 251 252class LoopbackVendor(SSHVendor): 253 """SSH "vendor" that connects over a plain TCP socket, not SSH.""" 254 255 def connect_sftp(self, username, password, host, port): 256 sock = socket.socket() 257 try: 258 sock.connect((host, port)) 259 except socket.error as e: 260 self._raise_connection_error(host, port=port, orig_error=e) 261 return SFTPClient(SocketAsChannelAdapter(sock)) 262 263 264register_ssh_vendor('loopback', LoopbackVendor()) 265 266 267class ParamikoVendor(SSHVendor): 268 """Vendor that uses paramiko.""" 269 270 def _hexify(self, s): 271 return hexlify(s).upper() 272 273 def _connect(self, username, password, host, port): 274 global SYSTEM_HOSTKEYS, BRZ_HOSTKEYS 275 276 load_host_keys() 277 278 try: 279 t = paramiko.Transport((host, port or 22)) 280 t.set_log_channel('bzr.paramiko') 281 t.start_client() 282 except (paramiko.SSHException, socket.error) as e: 283 self._raise_connection_error(host, port=port, orig_error=e) 284 285 server_key = t.get_remote_server_key() 286 server_key_hex = self._hexify(server_key.get_fingerprint()) 287 keytype = server_key.get_name() 288 if host in SYSTEM_HOSTKEYS and keytype in SYSTEM_HOSTKEYS[host]: 289 our_server_key = SYSTEM_HOSTKEYS[host][keytype] 290 our_server_key_hex = self._hexify(our_server_key.get_fingerprint()) 291 elif host in BRZ_HOSTKEYS and keytype in BRZ_HOSTKEYS[host]: 292 our_server_key = BRZ_HOSTKEYS[host][keytype] 293 our_server_key_hex = self._hexify(our_server_key.get_fingerprint()) 294 else: 295 trace.warning('Adding %s host key for %s: %s' 296 % (keytype, host, server_key_hex)) 297 add = getattr(BRZ_HOSTKEYS, 'add', None) 298 if add is not None: # paramiko >= 1.X.X 299 BRZ_HOSTKEYS.add(host, keytype, server_key) 300 else: 301 BRZ_HOSTKEYS.setdefault(host, {})[keytype] = server_key 302 our_server_key = server_key 303 our_server_key_hex = self._hexify(our_server_key.get_fingerprint()) 304 save_host_keys() 305 if server_key != our_server_key: 306 filename1 = os.path.expanduser('~/.ssh/known_hosts') 307 filename2 = _ssh_host_keys_config_dir() 308 raise errors.TransportError( 309 'Host keys for %s do not match! %s != %s' % 310 (host, our_server_key_hex, server_key_hex), 311 ['Try editing %s or %s' % (filename1, filename2)]) 312 313 _paramiko_auth(username, password, host, port, t) 314 return t 315 316 def connect_sftp(self, username, password, host, port): 317 t = self._connect(username, password, host, port) 318 try: 319 return t.open_sftp_client() 320 except paramiko.SSHException as e: 321 self._raise_connection_error(host, port=port, orig_error=e, 322 msg='Unable to start sftp client') 323 324 def connect_ssh(self, username, password, host, port, command): 325 t = self._connect(username, password, host, port) 326 try: 327 channel = t.open_session() 328 cmdline = ' '.join(command) 329 channel.exec_command(cmdline) 330 return _ParamikoSSHConnection(channel) 331 except paramiko.SSHException as e: 332 self._raise_connection_error(host, port=port, orig_error=e, 333 msg='Unable to invoke remote bzr') 334 335 336_ssh_connection_errors = (EOFError, OSError, IOError, socket.error) 337if paramiko is not None: 338 vendor = ParamikoVendor() 339 register_ssh_vendor('paramiko', vendor) 340 register_ssh_vendor('none', vendor) 341 register_default_ssh_vendor(vendor) 342 _ssh_connection_errors += (paramiko.SSHException,) 343 del vendor 344 345 346class SubprocessVendor(SSHVendor): 347 """Abstract base class for vendors that use pipes to a subprocess.""" 348 349 # In general stderr should be inherited from the parent process so prompts 350 # are visible on the terminal. This can be overriden to another file for 351 # tests, but beware of using PIPE which may hang due to not being read. 352 _stderr_target = None 353 354 @staticmethod 355 def _check_hostname(arg): 356 if arg.startswith('-'): 357 raise StrangeHostname(hostname=arg) 358 359 def _connect(self, argv): 360 # Attempt to make a socketpair to use as stdin/stdout for the SSH 361 # subprocess. We prefer sockets to pipes because they support 362 # non-blocking short reads, allowing us to optimistically read 64k (or 363 # whatever) chunks. 364 try: 365 my_sock, subproc_sock = socket.socketpair() 366 osutils.set_fd_cloexec(my_sock) 367 except (AttributeError, socket.error): 368 # This platform doesn't support socketpair(), so just use ordinary 369 # pipes instead. 370 stdin = stdout = subprocess.PIPE 371 my_sock, subproc_sock = None, None 372 else: 373 stdin = stdout = subproc_sock 374 proc = subprocess.Popen(argv, stdin=stdin, stdout=stdout, 375 stderr=self._stderr_target, 376 bufsize=0, 377 **os_specific_subprocess_params()) 378 if subproc_sock is not None: 379 subproc_sock.close() 380 return SSHSubprocessConnection(proc, sock=my_sock) 381 382 def connect_sftp(self, username, password, host, port): 383 try: 384 argv = self._get_vendor_specific_argv(username, host, port, 385 subsystem='sftp') 386 sock = self._connect(argv) 387 return SFTPClient(SocketAsChannelAdapter(sock)) 388 except _ssh_connection_errors as e: 389 self._raise_connection_error(host, port=port, orig_error=e) 390 391 def connect_ssh(self, username, password, host, port, command): 392 try: 393 argv = self._get_vendor_specific_argv(username, host, port, 394 command=command) 395 return self._connect(argv) 396 except _ssh_connection_errors as e: 397 self._raise_connection_error(host, port=port, orig_error=e) 398 399 def _get_vendor_specific_argv(self, username, host, port, subsystem=None, 400 command=None): 401 """Returns the argument list to run the subprocess with. 402 403 Exactly one of 'subsystem' and 'command' must be specified. 404 """ 405 raise NotImplementedError(self._get_vendor_specific_argv) 406 407 408class OpenSSHSubprocessVendor(SubprocessVendor): 409 """SSH vendor that uses the 'ssh' executable from OpenSSH.""" 410 411 executable_path = 'ssh' 412 413 def _get_vendor_specific_argv(self, username, host, port, subsystem=None, 414 command=None): 415 args = [self.executable_path, 416 '-oForwardX11=no', '-oForwardAgent=no', 417 '-oClearAllForwardings=yes', 418 '-oNoHostAuthenticationForLocalhost=yes'] 419 if port is not None: 420 args.extend(['-p', str(port)]) 421 if username is not None: 422 args.extend(['-l', username]) 423 if subsystem is not None: 424 args.extend(['-s', '--', host, subsystem]) 425 else: 426 args.extend(['--', host] + command) 427 return args 428 429 430register_ssh_vendor('openssh', OpenSSHSubprocessVendor()) 431 432 433class SSHCorpSubprocessVendor(SubprocessVendor): 434 """SSH vendor that uses the 'ssh' executable from SSH Corporation.""" 435 436 executable_path = 'ssh' 437 438 def _get_vendor_specific_argv(self, username, host, port, subsystem=None, 439 command=None): 440 self._check_hostname(host) 441 args = [self.executable_path, '-x'] 442 if port is not None: 443 args.extend(['-p', str(port)]) 444 if username is not None: 445 args.extend(['-l', username]) 446 if subsystem is not None: 447 args.extend(['-s', subsystem, host]) 448 else: 449 args.extend([host] + command) 450 return args 451 452 453register_ssh_vendor('sshcorp', SSHCorpSubprocessVendor()) 454 455 456class LSHSubprocessVendor(SubprocessVendor): 457 """SSH vendor that uses the 'lsh' executable from GNU""" 458 459 executable_path = 'lsh' 460 461 def _get_vendor_specific_argv(self, username, host, port, subsystem=None, 462 command=None): 463 self._check_hostname(host) 464 args = [self.executable_path] 465 if port is not None: 466 args.extend(['-p', str(port)]) 467 if username is not None: 468 args.extend(['-l', username]) 469 if subsystem is not None: 470 args.extend(['--subsystem', subsystem, host]) 471 else: 472 args.extend([host] + command) 473 return args 474 475 476register_ssh_vendor('lsh', LSHSubprocessVendor()) 477 478 479class PLinkSubprocessVendor(SubprocessVendor): 480 """SSH vendor that uses the 'plink' executable from Putty.""" 481 482 executable_path = 'plink' 483 484 def _get_vendor_specific_argv(self, username, host, port, subsystem=None, 485 command=None): 486 self._check_hostname(host) 487 args = [self.executable_path, '-x', '-a', '-ssh', '-2', '-batch'] 488 if port is not None: 489 args.extend(['-P', str(port)]) 490 if username is not None: 491 args.extend(['-l', username]) 492 if subsystem is not None: 493 args.extend(['-s', host, subsystem]) 494 else: 495 args.extend([host] + command) 496 return args 497 498 499register_ssh_vendor('plink', PLinkSubprocessVendor()) 500 501 502def _paramiko_auth(username, password, host, port, paramiko_transport): 503 auth = config.AuthenticationConfig() 504 # paramiko requires a username, but it might be none if nothing was 505 # supplied. If so, use the local username. 506 if username is None: 507 username = auth.get_user('ssh', host, port=port, 508 default=getpass.getuser()) 509 agent = paramiko.Agent() 510 for key in agent.get_keys(): 511 trace.mutter('Trying SSH agent key %s' 512 % hexlify(key.get_fingerprint()).upper()) 513 try: 514 paramiko_transport.auth_publickey(username, key) 515 return 516 except paramiko.SSHException as e: 517 pass 518 519 # okay, try finding id_rsa or id_dss? (posix only) 520 if _try_pkey_auth(paramiko_transport, paramiko.RSAKey, username, 'id_rsa'): 521 return 522 if _try_pkey_auth(paramiko_transport, paramiko.DSSKey, username, 'id_dsa'): 523 return 524 525 # If we have gotten this far, we are about to try for passwords, do an 526 # auth_none check to see if it is even supported. 527 supported_auth_types = [] 528 try: 529 # Note that with paramiko <1.7.5 this logs an INFO message: 530 # Authentication type (none) not permitted. 531 # So we explicitly disable the logging level for this action 532 old_level = paramiko_transport.logger.level 533 paramiko_transport.logger.setLevel(logging.WARNING) 534 try: 535 paramiko_transport.auth_none(username) 536 finally: 537 paramiko_transport.logger.setLevel(old_level) 538 except paramiko.BadAuthenticationType as e: 539 # Supported methods are in the exception 540 supported_auth_types = e.allowed_types 541 except paramiko.SSHException as e: 542 # Don't know what happened, but just ignore it 543 pass 544 # We treat 'keyboard-interactive' and 'password' auth methods identically, 545 # because Paramiko's auth_password method will automatically try 546 # 'keyboard-interactive' auth (using the password as the response) if 547 # 'password' auth is not available. Apparently some Debian and Gentoo 548 # OpenSSH servers require this. 549 # XXX: It's possible for a server to require keyboard-interactive auth that 550 # requires something other than a single password, but we currently don't 551 # support that. 552 if ('password' not in supported_auth_types and 553 'keyboard-interactive' not in supported_auth_types): 554 raise errors.ConnectionError('Unable to authenticate to SSH host as' 555 '\n %s@%s\nsupported auth types: %s' 556 % (username, host, supported_auth_types)) 557 558 if password: 559 try: 560 paramiko_transport.auth_password(username, password) 561 return 562 except paramiko.SSHException as e: 563 pass 564 565 # give up and ask for a password 566 password = auth.get_password('ssh', host, username, port=port) 567 # get_password can still return None, which means we should not prompt 568 if password is not None: 569 try: 570 paramiko_transport.auth_password(username, password) 571 except paramiko.SSHException as e: 572 raise errors.ConnectionError( 573 'Unable to authenticate to SSH host as' 574 '\n %s@%s\n' % (username, host), e) 575 else: 576 raise errors.ConnectionError('Unable to authenticate to SSH host as' 577 ' %s@%s' % (username, host)) 578 579 580def _try_pkey_auth(paramiko_transport, pkey_class, username, filename): 581 filename = os.path.expanduser('~/.ssh/' + filename) 582 try: 583 key = pkey_class.from_private_key_file(filename) 584 paramiko_transport.auth_publickey(username, key) 585 return True 586 except paramiko.PasswordRequiredException: 587 password = ui.ui_factory.get_password( 588 prompt=u'SSH %(filename)s password', 589 filename=filename.decode(osutils._fs_enc)) 590 try: 591 key = pkey_class.from_private_key_file(filename, password) 592 paramiko_transport.auth_publickey(username, key) 593 return True 594 except paramiko.SSHException: 595 trace.mutter('SSH authentication via %s key failed.' 596 % (os.path.basename(filename),)) 597 except paramiko.SSHException: 598 trace.mutter('SSH authentication via %s key failed.' 599 % (os.path.basename(filename),)) 600 except IOError: 601 pass 602 return False 603 604 605def _ssh_host_keys_config_dir(): 606 return osutils.pathjoin(bedding.config_dir(), 'ssh_host_keys') 607 608 609def load_host_keys(): 610 """ 611 Load system host keys (probably doesn't work on windows) and any 612 "discovered" keys from previous sessions. 613 """ 614 global SYSTEM_HOSTKEYS, BRZ_HOSTKEYS 615 try: 616 SYSTEM_HOSTKEYS = paramiko.util.load_host_keys( 617 os.path.expanduser('~/.ssh/known_hosts')) 618 except IOError as e: 619 trace.mutter('failed to load system host keys: ' + str(e)) 620 brz_hostkey_path = _ssh_host_keys_config_dir() 621 try: 622 BRZ_HOSTKEYS = paramiko.util.load_host_keys(brz_hostkey_path) 623 except IOError as e: 624 trace.mutter('failed to load brz host keys: ' + str(e)) 625 save_host_keys() 626 627 628def save_host_keys(): 629 """ 630 Save "discovered" host keys in $(config)/ssh_host_keys/. 631 """ 632 global SYSTEM_HOSTKEYS, BRZ_HOSTKEYS 633 bzr_hostkey_path = _ssh_host_keys_config_dir() 634 bedding.ensure_config_dir_exists() 635 636 try: 637 with open(bzr_hostkey_path, 'w') as f: 638 f.write('# SSH host keys collected by bzr\n') 639 for hostname, keys in BRZ_HOSTKEYS.items(): 640 for keytype, key in keys.items(): 641 f.write('%s %s %s\n' % 642 (hostname, keytype, key.get_base64())) 643 except IOError as e: 644 trace.mutter('failed to save bzr host keys: ' + str(e)) 645 646 647def os_specific_subprocess_params(): 648 """Get O/S specific subprocess parameters.""" 649 if sys.platform == 'win32': 650 # setting the process group and closing fds is not supported on 651 # win32 652 return {} 653 else: 654 # We close fds other than the pipes as the child process does not need 655 # them to be open. 656 # 657 # We also set the child process to ignore SIGINT. Normally the signal 658 # would be sent to every process in the foreground process group, but 659 # this causes it to be seen only by bzr and not by ssh. Python will 660 # generate a KeyboardInterrupt in bzr, and we will then have a chance 661 # to release locks or do other cleanup over ssh before the connection 662 # goes away. 663 # <https://launchpad.net/products/bzr/+bug/5987> 664 # 665 # Running it in a separate process group is not good because then it 666 # can't get non-echoed input of a password or passphrase. 667 # <https://launchpad.net/products/bzr/+bug/40508> 668 return {'preexec_fn': _ignore_signals, 669 'close_fds': True, 670 } 671 672 673import weakref 674_subproc_weakrefs = set() 675 676 677def _close_ssh_proc(proc, sock): 678 """Carefully close stdin/stdout and reap the SSH process. 679 680 If the pipes are already closed and/or the process has already been 681 wait()ed on, that's ok, and no error is raised. The goal is to do our best 682 to clean up (whether or not a clean up was already tried). 683 """ 684 funcs = [] 685 for closeable in (proc.stdin, proc.stdout, sock): 686 # We expect that either proc (a subprocess.Popen) will have stdin and 687 # stdout streams to close, or that we will have been passed a socket to 688 # close, with the option not in use being None. 689 if closeable is not None: 690 funcs.append(closeable.close) 691 funcs.append(proc.wait) 692 for func in funcs: 693 try: 694 func() 695 except OSError: 696 # It's ok for the pipe to already be closed, or the process to 697 # already be finished. 698 continue 699 700 701class SSHConnection(object): 702 """Abstract base class for SSH connections.""" 703 704 def get_sock_or_pipes(self): 705 """Returns a (kind, io_object) pair. 706 707 If kind == 'socket', then io_object is a socket. 708 709 If kind == 'pipes', then io_object is a pair of file-like objects 710 (read_from, write_to). 711 """ 712 raise NotImplementedError(self.get_sock_or_pipes) 713 714 def close(self): 715 raise NotImplementedError(self.close) 716 717 718class SSHSubprocessConnection(SSHConnection): 719 """A connection to an ssh subprocess via pipes or a socket. 720 721 This class is also socket-like enough to be used with 722 SocketAsChannelAdapter (it has 'send' and 'recv' methods). 723 """ 724 725 def __init__(self, proc, sock=None): 726 """Constructor. 727 728 :param proc: a subprocess.Popen 729 :param sock: if proc.stdin/out is a socket from a socketpair, then sock 730 should breezy's half of that socketpair. If not passed, proc's 731 stdin/out is assumed to be ordinary pipes. 732 """ 733 self.proc = proc 734 self._sock = sock 735 # Add a weakref to proc that will attempt to do the same as self.close 736 # to avoid leaving processes lingering indefinitely. 737 738 def terminate(ref): 739 _subproc_weakrefs.remove(ref) 740 _close_ssh_proc(proc, sock) 741 _subproc_weakrefs.add(weakref.ref(self, terminate)) 742 743 def send(self, data): 744 if self._sock is not None: 745 return self._sock.send(data) 746 else: 747 return os.write(self.proc.stdin.fileno(), data) 748 749 def recv(self, count): 750 if self._sock is not None: 751 return self._sock.recv(count) 752 else: 753 return os.read(self.proc.stdout.fileno(), count) 754 755 def close(self): 756 _close_ssh_proc(self.proc, self._sock) 757 758 def get_sock_or_pipes(self): 759 if self._sock is not None: 760 return 'socket', self._sock 761 else: 762 return 'pipes', (self.proc.stdout, self.proc.stdin) 763 764 765class _ParamikoSSHConnection(SSHConnection): 766 """An SSH connection via paramiko.""" 767 768 def __init__(self, channel): 769 self.channel = channel 770 771 def get_sock_or_pipes(self): 772 return ('socket', self.channel) 773 774 def close(self): 775 return self.channel.close() 776