Imported Upstream version 1.10.1upstream/1.10.1

1 files changed, 100 insertions, 65 deletions

diff --git a/paramiko/client.py b/paramiko/client.py
index 4a65477..5b71958 100644
--- a/paramiko/client.py
+++ b/paramiko/client.py
@@ -28,16 +28,16 @@ import warnings
 
 from paramiko.agent import Agent
 from paramiko.common import *
+from paramiko.config import SSH_PORT
 from paramiko.dsskey import DSSKey
 from paramiko.hostkeys import HostKeys
 from paramiko.resource import ResourceManager
 from paramiko.rsakey import RSAKey
 from paramiko.ssh_exception import SSHException, BadHostKeyException
 from paramiko.transport import Transport
+from paramiko.util import retry_on_signal
 
 
-SSH_PORT = 22
-
 class MissingHostKeyPolicy (object):
     """
     Interface for defining the policy that L{SSHClient} should use when the
@@ -82,7 +82,7 @@ class RejectPolicy (MissingHostKeyPolicy):
     def missing_host_key(self, client, hostname, key):
         client._log(DEBUG, 'Rejecting %s host key for %s: %s' %
                     (key.get_name(), hostname, hexlify(key.get_fingerprint())))
-        raise SSHException('Unknown server %s' % hostname)
+        raise SSHException('Server %r not found in known_hosts' % hostname)
 
 
 class WarningPolicy (MissingHostKeyPolicy):
@@ -228,7 +228,7 @@ class SSHClient (object):
 
     def connect(self, hostname, port=SSH_PORT, username=None, password=None, pkey=None,
                 key_filename=None, timeout=None, allow_agent=True, look_for_keys=True,
-                compress=False):
+                compress=False, sock=None):
         """
         Connect to an SSH server and authenticate to it.  The server's host key
         is checked against the system host keys (see L{load_system_host_keys})
@@ -271,6 +271,9 @@ class SSHClient (object):
         @type look_for_keys: bool
         @param compress: set to True to turn on compression
         @type compress: bool
+        @param sock: an open socket or socket-like object (such as a
+            L{Channel}) to use for communication to the target host
+        @type sock: socket
 
         @raise BadHostKeyException: if the server's host key could not be
             verified
@@ -279,21 +282,23 @@ class SSHClient (object):
             establishing an SSH session
         @raise socket.error: if a socket error occurred while connecting
         """
-        for (family, socktype, proto, canonname, sockaddr) in socket.getaddrinfo(hostname, port, socket.AF_UNSPEC, socket.SOCK_STREAM):
-            if socktype == socket.SOCK_STREAM:
-                af = family
-                addr = sockaddr
-                break
-        else:
-            # some OS like AIX don't indicate SOCK_STREAM support, so just guess. :(
-            af, _, _, _, addr = socket.getaddrinfo(hostname, port, socket.AF_UNSPEC, socket.SOCK_STREAM)
-        sock = socket.socket(af, socket.SOCK_STREAM)
-        if timeout is not None:
-            try:
-                sock.settimeout(timeout)
-            except:
-                pass
-        sock.connect(addr)
+        if not sock:
+            for (family, socktype, proto, canonname, sockaddr) in socket.getaddrinfo(hostname, port, socket.AF_UNSPEC, socket.SOCK_STREAM):
+                if socktype == socket.SOCK_STREAM:
+                    af = family
+                    addr = sockaddr
+                    break
+            else:
+                # some OS like AIX don't indicate SOCK_STREAM support, so just guess. :(
+                af, _, _, _, addr = socket.getaddrinfo(hostname, port, socket.AF_UNSPEC, socket.SOCK_STREAM)
+            sock = socket.socket(af, socket.SOCK_STREAM)
+            if timeout is not None:
+                try:
+                    sock.settimeout(timeout)
+                except:
+                    pass
+            retry_on_signal(lambda: sock.connect(addr))
+
         t = self._transport = Transport(sock)
         t.use_compression(compress=compress)
         if self._log_channel is not None:
@@ -344,7 +349,7 @@ class SSHClient (object):
             self._agent.close()
             self._agent = None
 
-    def exec_command(self, command, bufsize=-1):
+    def exec_command(self, command, bufsize=-1, timeout=None, get_pty=False):
         """
         Execute a command on the SSH server.  A new L{Channel} is opened and
         the requested command is executed.  The command's input and output
@@ -355,19 +360,25 @@ class SSHClient (object):
         @type command: str
         @param bufsize: interpreted the same way as by the built-in C{file()} function in python
         @type bufsize: int
+        @param timeout: set command's channel timeout. See L{Channel.settimeout}.settimeout
+        @type timeout: int
         @return: the stdin, stdout, and stderr of the executing command
         @rtype: tuple(L{ChannelFile}, L{ChannelFile}, L{ChannelFile})
 
         @raise SSHException: if the server fails to execute the command
         """
         chan = self._transport.open_session()
+        if(get_pty):
+            chan.get_pty()
+        chan.settimeout(timeout)
         chan.exec_command(command)
         stdin = chan.makefile('wb', bufsize)
         stdout = chan.makefile('rb', bufsize)
         stderr = chan.makefile_stderr('rb', bufsize)
         return stdin, stdout, stderr
 
-    def invoke_shell(self, term='vt100', width=80, height=24):
+    def invoke_shell(self, term='vt100', width=80, height=24, width_pixels=0,
+                height_pixels=0):
         """
         Start an interactive shell session on the SSH server.  A new L{Channel}
         is opened and connected to a pseudo-terminal using the requested
@@ -379,13 +390,17 @@ class SSHClient (object):
         @type width: int
         @param height: the height (in characters) of the terminal window
         @type height: int
+        @param width_pixels: the width (in pixels) of the terminal window
+        @type width_pixels: int
+        @param height_pixels: the height (in pixels) of the terminal window
+        @type height_pixels: int
         @return: a new channel connected to the remote shell
         @rtype: L{Channel}
 
         @raise SSHException: if the server fails to invoke a shell
         """
         chan = self._transport.open_session()
-        chan.get_pty(term, width, height)
+        chan.get_pty(term, width, height, width_pixels, height_pixels)
         chan.invoke_shell()
         return chan
 
@@ -418,68 +433,86 @@ class SSHClient (object):
             - Any "id_rsa" or "id_dsa" key discoverable in ~/.ssh/ (if allowed).
             - Plain username/password auth, if a password was given.
 
-        (The password might be needed to unlock a private key.)
+        (The password might be needed to unlock a private key, or for
+        two-factor authentication [for which it is required].)
         """
         saved_exception = None
+        two_factor = False
+        allowed_types = []
 
         if pkey is not None:
             try:
                 self._log(DEBUG, 'Trying SSH key %s' % hexlify(pkey.get_fingerprint()))
-                self._transport.auth_publickey(username, pkey)
-                return
+                allowed_types = self._transport.auth_publickey(username, pkey)
+                two_factor = (allowed_types == ['password'])
+                if not two_factor:
+                    return
             except SSHException, e:
                 saved_exception = e
 
-        for key_filename in key_filenames:
-            for pkey_class in (RSAKey, DSSKey):
-                try:
-                    key = pkey_class.from_private_key_file(key_filename, password)
-                    self._log(DEBUG, 'Trying key %s from %s' % (hexlify(key.get_fingerprint()), key_filename))
-                    self._transport.auth_publickey(username, key)
-                    return
-                except SSHException, e:
-                    saved_exception = e
-
-        if allow_agent:
+        if not two_factor:
+            for key_filename in key_filenames:
+                for pkey_class in (RSAKey, DSSKey):
+                    try:
+                        key = pkey_class.from_private_key_file(key_filename, password)
+                        self._log(DEBUG, 'Trying key %s from %s' % (hexlify(key.get_fingerprint()), key_filename))
+                        self._transport.auth_publickey(username, key)
+                        two_factor = (allowed_types == ['password'])
+                        if not two_factor:
+                            return
+                        break
+                    except SSHException, e:
+                        saved_exception = e
+
+        if not two_factor and allow_agent:
             if self._agent == None:
                 self._agent = Agent()
 
             for key in self._agent.get_keys():
                 try:
                     self._log(DEBUG, 'Trying SSH agent key %s' % hexlify(key.get_fingerprint()))
-                    self._transport.auth_publickey(username, key)
-                    return
+                    # for 2-factor auth a successfully auth'd key will result in ['password']
+                    allowed_types = self._transport.auth_publickey(username, key)
+                    two_factor = (allowed_types == ['password'])
+                    if not two_factor:
+                        return
+                    break
                 except SSHException, e:
                     saved_exception = e
 
-        keyfiles = []
-        rsa_key = os.path.expanduser('~/.ssh/id_rsa')
-        dsa_key = os.path.expanduser('~/.ssh/id_dsa')
-        if os.path.isfile(rsa_key):
-            keyfiles.append((RSAKey, rsa_key))
-        if os.path.isfile(dsa_key):
-            keyfiles.append((DSSKey, dsa_key))
-        # look in ~/ssh/ for windows users:
-        rsa_key = os.path.expanduser('~/ssh/id_rsa')
-        dsa_key = os.path.expanduser('~/ssh/id_dsa')
-        if os.path.isfile(rsa_key):
-            keyfiles.append((RSAKey, rsa_key))
-        if os.path.isfile(dsa_key):
-            keyfiles.append((DSSKey, dsa_key))
-
-        if not look_for_keys:
+        if not two_factor:
             keyfiles = []
-
-        for pkey_class, filename in keyfiles:
-            try:
-                key = pkey_class.from_private_key_file(filename, password)
-                self._log(DEBUG, 'Trying discovered key %s in %s' % (hexlify(key.get_fingerprint()), filename))
-                self._transport.auth_publickey(username, key)
-                return
-            except SSHException, e:
-                saved_exception = e
-            except IOError, e:
-                saved_exception = e
+            rsa_key = os.path.expanduser('~/.ssh/id_rsa')
+            dsa_key = os.path.expanduser('~/.ssh/id_dsa')
+            if os.path.isfile(rsa_key):
+                keyfiles.append((RSAKey, rsa_key))
+            if os.path.isfile(dsa_key):
+                keyfiles.append((DSSKey, dsa_key))
+            # look in ~/ssh/ for windows users:
+            rsa_key = os.path.expanduser('~/ssh/id_rsa')
+            dsa_key = os.path.expanduser('~/ssh/id_dsa')
+            if os.path.isfile(rsa_key):
+                keyfiles.append((RSAKey, rsa_key))
+            if os.path.isfile(dsa_key):
+                keyfiles.append((DSSKey, dsa_key))
+
+            if not look_for_keys:
+                keyfiles = []
+
+            for pkey_class, filename in keyfiles:
+                try:
+                    key = pkey_class.from_private_key_file(filename, password)
+                    self._log(DEBUG, 'Trying discovered key %s in %s' % (hexlify(key.get_fingerprint()), filename))
+                    # for 2-factor auth a successfully auth'd key will result in ['password']
+                    allowed_types = self._transport.auth_publickey(username, key)
+                    two_factor = (allowed_types == ['password'])
+                    if not two_factor:
+                        return
+                    break
+                except SSHException, e:
+                    saved_exception = e
+                except IOError, e:
+                    saved_exception = e
 
         if password is not None:
             try:
@@ -487,6 +520,8 @@ class SSHClient (object):
                 return
             except SSHException, e:
                 saved_exception = e
+        elif two_factor:
+            raise SSHException('Two-factor authentication requires a password')
 
         # if we got an auth-failed exception earlier, re-raise it
         if saved_exception is not None: