Changed fake server to threading to make it compatible with paramiko's ServerInterface. Changed tests to use new fake server

Author: pkittenis

fake_server/fake_server.py

@@ -1,17 +1,22 @@
 #!/usr/bin/env python
 
-"""Fake SSH server to test our SSH clients.
-Supports execution of commands via exec_command. Does _not_ support interactive shells,
-our clients do not use them.
-Server private key is hardcoded, server listen code inspired by demo_server.py in paramiko repository"""
+"""
+Fake SSH server to test our SSH clients.
+Supports execution of commands via exec_command. Does _not_ support interactive \
+shells, our clients do not use them.
+Server private key is hardcoded, server listen code inspired by demo_server.py in \
+paramiko repository
+"""
 
 import os
-from gevent import socket
-import gevent.event
+import socket
+import threading
+from threading import Event
 import sys
 import traceback
 import logging
 import paramiko
+import time
 
 logger = logging.getLogger(__name__)
 paramiko_logger = logging.getLogger('paramiko.transport')
@@ -20,7 +25,7 @@
 
 class Server (paramiko.ServerInterface):
     def __init__(self, cmd_req_response = {}, fail_auth = False):
-        self.event = gevent.event.Event()
+        self.event = Event()
         self.cmd_req_response = cmd_req_response
         self.fail_auth = fail_auth
 
@@ -56,60 +61,77 @@ def check_channel_exec_request(self, channel, cmd):
         self.event.set()
         return True
 
-def _make_socket(listen_ip, listen_port):
+def make_socket(listen_ip):
+    """Make socket on given address and available port chosen by OS"""
     try:
         sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
         sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
-        sock.bind((listen_ip, listen_port))
+        sock.bind((listen_ip, 0))
     except Exception, e:
         logger.error('Failed to bind to address - %s' % (str(e),))
         traceback.print_exc()
         return
     return sock
 
-def listen(cmd_req_response, listen_ip = '127.0.0.1', listen_port = 2200, fail_auth = False):
-    """Run a fake ssh server and given a cmd_to_run, send given response"""
-    sock = _make_socket(listen_ip, listen_port)
+def listen(cmd_req_response, sock, fail_auth = False):
+    """Run a fake ssh server and given a cmd_to_run, send given \
+    response to client connection. Returns (server, socket) tuple \
+    where server is a joinable server thread and socket is listening \
+    socket of server."""
+    # sock = _make_socket(listen_ip)
+    listen_ip, listen_port = sock.getsockname()
     if not sock:
         logger.error("Could not establish listening connection on %s:%s", listen_ip, listen_port)
         return 
     try:
         sock.listen(100)
         logger.info('Listening for connection on %s:%s..', listen_ip, listen_port)
-        client, addr = sock.accept()
     except Exception, e:
-        logger.error('*** Listen/accept failed: %s' % (str(e),))
+        logger.error('*** Listen failed: %s' % (str(e),))
         traceback.print_exc()
         return
+    accept_thread = threading.Thread(target=handle_ssh_connection,
+                                     args=(cmd_req_response, sock,),
+                                     kwargs={'fail_auth' : fail_auth},)
+    accept_thread.start()
+    return accept_thread
+
+def _handle_ssh_connection(cmd_req_response, t, client, addr, fail_auth = False):
+    try:
+        t.load_server_moduli()
+    except:
+        return
+    t.add_server_key(host_key)
+    server = Server(cmd_req_response = cmd_req_response, fail_auth = fail_auth)
+    try:
+        t.start_server(server=server)
+    except paramiko.SSHException, _:
+        logger.error('SSH negotiation failed.')
+        return
+    return _accept_ssh_data(t, server)
+
+def _accept_ssh_data(t, server):
+    chan = t.accept(20)
+    if not chan:
+        logger.error("Could not establish channel")
+        return
+    logger.info("Authenticated..")
+    chan.send_ready()
+    server.event.wait(10)
+    if not server.event.isSet():
+        logger.error('Client never sent command')
+        chan.close()
+        return
+    while not chan.send_ready():
+        time.sleep(.5)
+    chan.close()
+    
+def handle_ssh_connection(cmd_req_response, sock, fail_auth = False):
+    client, addr = sock.accept()
     logger.info('Got connection..')
     try:
         t = paramiko.Transport(client)
-        try:
-            t.load_server_moduli()
-        except:
-            raise
-        t.add_server_key(host_key)
-        server = Server(cmd_req_response = cmd_req_response, fail_auth = fail_auth)
-        try:
-            t.start_server(server=server)
-        except paramiko.SSHException, _:
-            logger.error('SSH negotiation failed.')
-            return
-        chan = t.accept(20)
-        if not chan:
-            logger.error("Could not establish channel")
-            return
-        logger.info("Authenticated..")
-        chan.send_ready()
-        server.event.wait(10)
-        if not server.event.isSet():
-            logger.error('Client never sent command')
-            chan.close()
-            return
-        while not chan.send_ready():
-            gevent.sleep(.5)
-        chan.close()
-
+        _handle_ssh_connection(cmd_req_response, t, client, addr, fail_auth=fail_auth)
     except Exception, e:
         logger.error('*** Caught exception: %s: %s' % (str(e.__class__), str(e),))
         traceback.print_exc()
@@ -122,4 +144,6 @@ def listen(cmd_req_response, listen_ip = '127.0.0.1', listen_port = 2200, fail_a
 if __name__ == "__main__":
     logging.basicConfig()
     logger.setLevel(logging.DEBUG)
-    listen({'fake' : 'fake response' + os.linesep})
+    sock = make_socket('127.0.0.1')
+    server = listen({'fake' : 'fake response' + os.linesep}, sock)
+    server.join()

tests/test_pssh_client.py

@@ -3,10 +3,9 @@
 """Unittests for parallel-ssh"""
 
 import unittest
-import gevent
 from pssh import ParallelSSHClient, UnknownHostException, \
     AuthenticationException, ConnectionErrorException, _setup_logger
-from fake_server.fake_server import listen, logger as server_logger
+from fake_server.fake_server import listen, make_socket, logger as server_logger
 import random
 import logging
 
@@ -17,35 +16,33 @@ class ParallelSSHClientTest(unittest.TestCase):
     def setUp(self):
         self.fake_cmd = 'fake cmd'
         self.fake_resp = 'fake response'
+        self.listener = make_socket('127.0.0.1')
+        self.listen_port = self.listener.getsockname()[1]
+
+    def cleanUp(self):
+        del self.listener
 
     def test_pssh_client_exec_command(self):
-        listen_port = random.randint(1026, 65534)
-        server = gevent.spawn(listen, { self.fake_cmd : self.fake_resp }, listen_port = listen_port)
-        client = ParallelSSHClient(['localhost'], port=listen_port)
-        gevent.sleep(0)
+        server = listen({ self.fake_cmd : self.fake_resp }, self.listener)
+        client = ParallelSSHClient(['localhost'], port=self.listen_port)
         cmd = client.exec_command(self.fake_cmd)[0]
         output = client.get_stdout(cmd)
         expected = {'localhost' : {'exit_code' : 0}}
         self.assertEqual(expected, output,
                          msg = "Got unexpected command output - %s" % (output,))
-        server.kill()
         del client
+        server.join()
 
     def test_pssh_client_auth_failure(self):
-        listen_port = random.randint(2048, 65534)
-        server = gevent.spawn(listen, { self.fake_cmd : self.fake_resp },
-                              listen_port=listen_port, fail_auth=True, )
-        client = ParallelSSHClient(['localhost'], port=listen_port)
-        gevent.sleep(0)
-        server.join(1)
+        server = listen({ self.fake_cmd : self.fake_resp },
+                        self.listener, fail_auth=True)
+        client = ParallelSSHClient(['localhost'], port=self.listen_port)
         cmd = client.exec_command(self.fake_cmd)[0]
-        server.join(1)
         # Handle exception
         try:
             cmd.get()
             raise Exception("Expected AuthenticationException, got none")
         except AuthenticationException:
             pass
-        server.kill()
         del client
-        
+        server.join()