Merging feature branches into master

Author: pkittenis

fake_server/fake_agent.py

@@ -0,0 +1,28 @@
+"""
+Fake SSH Agent for testing ssh agent forwarding and agent based key
+authentication
+"""
+
+import paramiko.agent
+
+class FakeAgent(paramiko.agent.AgentSSH):
+
+    def __init__(self):
+        self._conn = None
+        self.keys = []
+    
+    def add_key(self, key):
+        """Add key to agent.
+        :param key: Key to add
+        :type key: :mod:`paramiko.pkey.PKey`
+        """
+        self.keys.append(key)
+
+    def _connect(self, conn):
+        pass
+
+    def _close(self):
+        self._keys = []
+
+    def get_keys(self):
+        return tuple(self.keys)

fake_server/fake_server.py

@@ -53,6 +53,10 @@ def check_channel_pty_request(self, channel, term, width, height, pixelwidth,
                                   pixelheight, modes):
         return True
 
+    def check_channel_forward_agent_request(self, channel):
+        logger.debug("Forward agent key request for channel %s" % (channel,))
+        return True
+
     def check_channel_exec_request(self, channel, cmd):
         logger.debug("Got exec request on channel %s for cmd %s" % (channel, cmd,))
         # Remove any 'bash -c' and/or quotes from command

pssh.py

@@ -72,7 +72,8 @@ class SSHClient(object):
 
     def __init__(self, host,
                  user=None, password=None, port=None,
-                 pkey=None):
+                 pkey=None, forward_ssh_agent=True,
+                 _agent=None):
         """Connect to host honouring any user set configuration in ~/.ssh/config \
         or /etc/ssh/ssh_config
          
@@ -89,6 +90,15 @@ def __init__(self, host,
         :type port: int
         :param pkey: (Optional) Client's private key to be used to connect with
         :type pkey: :mod:`paramiko.PKey`
+        :param forward_ssh_agent: (Optional) Turn on SSH agent forwarding - \
+        equivalent to `ssh -A` from the `ssh` command line utility.
+        Defaults to True if not set.
+        :type forward_ssh_agent: bool
+        :param _agent: (Optional) Override SSH agent object with the provided. \
+        This allows for overriding of the default paramiko behaviour of \
+        connecting to local SSH agent to lookup keys with our own SSH agent.
+        Only really useful for testing, hence the internal variable prefix.
+        :type _agent: :mod:`paramiko.agent.Agent`
         :raises: :mod:`pssh.AuthenticationException` on authentication error
         :raises: :mod:`pssh.UnknownHostException` on DNS resolution error
         :raises: :mod:`pssh.ConnectionErrorException` on error connecting
@@ -113,6 +123,7 @@ def __init__(self, host,
             user = _user
         client = paramiko.SSHClient()
         client.set_missing_host_key_policy(paramiko.MissingHostKeyPolicy())
+        self.forward_ssh_agent = forward_ssh_agent
         self.client = client
         self.channel = None
         self.user = user
@@ -125,6 +136,8 @@ def __init__(self, host,
         if self.proxy_command:
             logger.debug("Proxy configured for destination host %s - ProxyCommand: '%s'" % (
                 self.host, " ".join(self.proxy_command.cmd),))
+        if _agent:
+            self.client._agent = _agent
         self._connect()
 
     def _connect(self, retries=1):
@@ -177,6 +190,8 @@ def exec_command(self, command, sudo=False, **kwargs):
         stderr are buffers containing command output.
         """
         channel = self.client.get_transport().open_session()
+        if self.forward_ssh_agent:
+            agent_handler = paramiko.agent.AgentRequestHandler(channel)
         channel.get_pty()
         # stdin (unused), stdout, stderr
         (_, stdout, stderr) = (channel.makefile('wb'), channel.makefile('rb'),

tests/test_ssh_client.py

@@ -22,13 +22,18 @@
 from pssh import SSHClient, ParallelSSHClient, UnknownHostException, AuthenticationException, _setup_logger, logger
 from fake_server.fake_server import start_server, make_socket, logger as server_logger, \
     paramiko_logger
+from fake_server.fake_agent import FakeAgent
+import paramiko
 import os
 from test_pssh_client import USER_KEY
 
 # _setup_logger(server_logger)
 # _setup_logger(logger)
 # _setup_logger(paramiko_logger)
 
+USER_KEY = paramiko.RSAKey.from_private_key_file(
+    os.path.sep.join([os.path.dirname(__file__), 'test_client_private_key']))
+
 class SSHClientTest(unittest.TestCase):
 
     def setUp(self):
@@ -70,5 +75,27 @@ def test_ssh_client_sftp(self):
         del client
         server.join()
 
+    def test_ssh_agent_authentication(self):
+        """Test authentication via SSH agent.
+        Do not provide public key to use when creating SSHClient,
+        instead override the client's agent with our own fake SSH agent,
+        add our to key to agent and try to login to server.
+        Key should be automatically picked up from the overriden agent"""
+        agent = FakeAgent()
+        agent.add_key(USER_KEY)
+        server = start_server({ self.fake_cmd : self.fake_resp },
+                                self.listen_socket)
+        client = SSHClient('127.0.0.1', port=self.listen_port,
+                           _agent=agent)
+        channel, host, _stdout, _stderr = client.exec_command(self.fake_cmd)
+        output = (line.strip() for line in _stdout)
+        channel.close()
+        output = list(output)
+        expected = [self.fake_resp]
+        self.assertEqual(expected, output,
+                         msg = "Got unexpected command output - %s" % (output,))
+        del client
+        server.join()
+
 if __name__ == '__main__':
     unittest.main()