Added direct-tcpip tunneling support to fake server. Added test for proxying an SSH connection through another SSH server using direct-tcpip. Resolves #12

Author: Dan

fake_server/fake_server.py

@@ -25,6 +25,8 @@
 paramiko repository
 """
 
+from gevent import monkey
+monkey.patch_all()
 import gevent
 import os
 import socket
@@ -36,8 +38,7 @@
 import paramiko
 import time
 from stub_sftp import StubSFTPServer
-from gevent import monkey
-monkey.patch_all()
+from tunnel import Tunneler
 
 
 logger = logging.getLogger("fake_server")
@@ -46,10 +47,11 @@
 host_key = paramiko.RSAKey(filename = os.path.sep.join([os.path.dirname(__file__), 'rsa.key']))
 
 class Server (paramiko.ServerInterface):
-    def __init__(self, cmd_req_response = {}, fail_auth=False):
+    def __init__(self, transport, cmd_req_response = {}, fail_auth=False):
         self.event = Event()
         self.cmd_req_response = cmd_req_response
         self.fail_auth = fail_auth
+        self.transport = transport
 
     def check_channel_request(self, kind, chanid):
         return paramiko.OPEN_SUCCEEDED
@@ -72,6 +74,20 @@ def check_channel_pty_request(self, channel, term, width, height, pixelwidth,
                                   pixelheight, modes):
         return True
 
+    def check_channel_direct_tcpip_request(self, chanid, origin, destination):
+        logger.debug("Proxy connection %s -> %s requested", origin, destination,)
+        extra = {'username' : self.transport.get_username()}
+        logger.debug("Starting proxy connection %s -> %s",
+                     origin, destination, extra=extra)
+        try:
+            tunnel = Tunneler(destination, self.transport, chanid)
+            tunnel.start()
+        except Exception, ex:
+            logger.error("Error creating proxy connection to %s - %s",
+                         destination, ex,)
+            return paramiko.OPEN_FAILED_CONNECT_FAILED
+        return paramiko.OPEN_SUCCEEDED
+
     def check_channel_forward_agent_request(self, channel):
         logger.debug("Forward agent key request for channel %s" % (channel,))
         return True
@@ -142,7 +158,8 @@ def _handle_ssh_connection(cmd_req_response, transport, fail_auth=False):
         return
     transport.add_server_key(host_key)
     transport.set_subsystem_handler('sftp', paramiko.SFTPServer, StubSFTPServer)
-    server = Server(cmd_req_response = cmd_req_response, fail_auth=fail_auth)
+    server = Server(transport, cmd_req_response=cmd_req_response,
+                    fail_auth=fail_auth)
     try:
         transport.start_server(server=server)
     except paramiko.SSHException, e:
@@ -157,9 +174,9 @@ def _handle_ssh_connection(cmd_req_response, transport, fail_auth=False):
         return
     while transport.is_active():
         logger.debug("Transport active, waiting..")
-        time.sleep(1)
+        gevent.sleep(1)
     while not channel.send_ready():
-        time.sleep(.5)
+        gevent.sleep(.2)
     channel.close()
 
 def handle_ssh_connection(cmd_req_response, sock,

fake_server/tunnel.py

@@ -0,0 +1,70 @@
+# This file is part of parallel-ssh.
+
+# Copyright (C) 2014 Panos Kittenis
+
+# This library is free software; you can redistribute it and/or
+# modify it under the terms of the GNU Lesser General Public
+# License as published by the Free Software Foundation, version 2.1.
+
+# This library is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
+# Lesser General Public License for more details.
+
+# You should have received a copy of the GNU Lesser General Public
+# License along with this library; if not, write to the Free Software
+# Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301  USA
+
+"""
+Generic 'Tunneler' module for tunneling between source <-> destination
+network connections asynchronously with gevent.
+"""
+
+import gevent
+from gevent import socket, select
+import logging
+
+logger = logging.getLogger("fake_server")
+
+class Tunneler(gevent.Greenlet):    
+    def __init__(self, address, transport, chanid):
+        gevent.Greenlet.__init__(self)
+        self.socket = socket.create_connection(address)
+        self.transport = transport
+        self.chanid = chanid
+
+    def close(self):
+        try:
+            self.transport.close()
+        except Exception:
+            pass
+        return
+
+    def tunnel(self, dest_socket, source_chan):
+        try:
+            while True:
+                logger.debug("Tunnel waiting for data..")
+                data = source_chan.recv(1024)
+                dest_socket.sendall(data)
+                response_data = dest_socket.recv(1024)
+                source_chan.sendall(response_data)
+                logger.debug("Tunnel sent data..")
+                gevent.sleep(1)
+        finally:
+            source_chan.close()
+            dest_socket.close()
+
+    def run(self):
+        channel = self.transport.accept(20)
+        if not channel:
+            return
+        if not channel.get_id() == self.chanid:
+            return
+        peer = self.socket.getpeername()
+        logger.debug("Start tunneling with peer %s, user %s", peer,
+                     self.transport.get_username())
+        try:
+            self.tunnel(self.socket, channel)
+        except Exception, ex:
+            logger.exception("Got exception creating tunnel - %s", ex,)
+        logger.debug("Finished tunneling")

tests/test_pssh_client.py

@@ -21,7 +21,7 @@
 
 import unittest
 from pssh import ParallelSSHClient, UnknownHostException, \
-     AuthenticationException, ConnectionErrorException
+     AuthenticationException, ConnectionErrorException, logger as pssh_logger
 from fake_server.fake_server import start_server, make_socket, \
      logger as server_logger, paramiko_logger
 import random
@@ -35,6 +35,10 @@
 USER_KEY = paramiko.RSAKey.from_private_key_file(
     os.path.sep.join([os.path.dirname(__file__), 'test_client_private_key']))
 
+# server_logger.setLevel(logging.DEBUG)
+# pssh_logger.setLevel(logging.DEBUG)
+# logging.basicConfig()
+
 class ParallelSSHClientTest(unittest.TestCase):
 
     def setUp(self):
@@ -67,7 +71,7 @@ def test_pssh_client_exec_command(self):
         output = client.get_stdout(cmd)
         expected = {'127.0.0.1' : {'exit_code' : 0}}
         self.assertEqual(expected, output,
-                         msg = "Got unexpected command output - %s" % (output,))
+                         msg="Got unexpected command output - %s" % (output,))
         del client
         server.join()
 
@@ -85,15 +89,15 @@ def test_pssh_client_exec_command_get_buffers(self):
         stdout = list(output['127.0.0.1']['stdout'])
         stderr = list(output['127.0.0.1']['stderr'])
         self.assertEqual(expected_exit_code, exit_code,
-                         msg = "Got unexpected exit code - %s, expected %s" %
+                         msg="Got unexpected exit code - %s, expected %s" %
                          (exit_code,
                           expected_exit_code,))
         self.assertEqual(expected_stdout, stdout,
-                         msg = "Got unexpected stdout - %s, expected %s" % 
+                         msg="Got unexpected stdout - %s, expected %s" % 
                          (stdout,
                           expected_stdout,))
         self.assertEqual(expected_stderr, stderr,
-                         msg = "Got unexpected stderr - %s, expected %s" % 
+                         msg="Got unexpected stderr - %s, expected %s" % 
                          (stderr,
                           expected_stderr,))
         del client
@@ -112,15 +116,15 @@ def test_pssh_client_run_command_get_output(self):
         stdout = list(output['127.0.0.1']['stdout'])
         stderr = list(output['127.0.0.1']['stderr'])
         self.assertEqual(expected_exit_code, exit_code,
-                         msg = "Got unexpected exit code - %s, expected %s" %
+                         msg="Got unexpected exit code - %s, expected %s" %
                          (exit_code,
                           expected_exit_code,))
         self.assertEqual(expected_stdout, stdout,
-                         msg = "Got unexpected stdout - %s, expected %s" %
+                         msg="Got unexpected stdout - %s, expected %s" %
                          (stdout,
                           expected_stdout,))
         self.assertEqual(expected_stderr, stderr,
-                         msg = "Got unexpected stderr - %s, expected %s" %
+                         msg="Got unexpected stderr - %s, expected %s" %
                          (stderr,
                           expected_stderr,))
         del client
@@ -141,15 +145,15 @@ def test_pssh_client_run_command_get_output_explicit(self):
         stdout = list(output['127.0.0.1']['stdout'])
         stderr = list(output['127.0.0.1']['stderr'])
         self.assertEqual(expected_exit_code, exit_code,
-                         msg = "Got unexpected exit code - %s, expected %s" %
+                         msg="Got unexpected exit code - %s, expected %s" %
                          (exit_code,
                           expected_exit_code,))
         self.assertEqual(expected_stdout, stdout,
-                         msg = "Got unexpected stdout - %s, expected %s" % 
+                         msg="Got unexpected stdout - %s, expected %s" % 
                          (stdout,
                           expected_stdout,))
         self.assertEqual(expected_stderr, stderr,
-                         msg = "Got unexpected stderr - %s, expected %s" % 
+                         msg="Got unexpected stderr - %s, expected %s" % 
                          (stderr,
                           expected_stderr,))
         del client
@@ -298,3 +302,28 @@ def test_pssh_pool_size(self):
         self.assertEqual(expected, actual,
                          msg="Expected pool size to be %s, got %s" % (
                              expected, actual,))
+
+    def test_ssh_proxy(self):
+        """Test connecting to remote destination via SSH proxy
+        client -> proxy -> destination
+        Proxy SSH server accepts no commands and sends no responses, only
+        proxies to destination. Destination accepts a command as usual."""
+        proxy_server_socket = make_socket('127.0.0.1')
+        proxy_server_port = proxy_server_socket.getsockname()[1]
+        proxy_server = start_server({}, proxy_server_socket)
+        server = start_server({ self.fake_cmd : self.fake_resp },
+                              self.listen_socket)
+        client = ParallelSSHClient(['127.0.0.1'], port=self.listen_port,
+                                   pkey=self.user_key,
+                                   proxy_host='127.0.0.1',
+                                   proxy_port=proxy_server_port
+                                   )
+        output = client.run_command(self.fake_cmd)
+        stdout = list(output['127.0.0.1']['stdout'])
+        expected_stdout = [self.fake_resp]
+        self.assertEqual(expected_stdout, stdout,
+                         msg="Got unexpected stdout - %s, expected %s" % 
+                         (stdout,
+                          expected_stdout,))
+        server.kill()
+        proxy_server.kill()