Added sftp server interface stub by https://github.com/rspivak/sftpserver. Added stub SFTP handler to fake server. Added SFTP functionality tests for SSHClient. Cleaned up fake server startup, doc changes

pkittenis · pkittenis · commit 193ae31d3f1b · 2014-04-04T15:33:27.000+01:00
diff --git a/fake_server/fake_server.py b/fake_server/fake_server.py
@@ -19,6 +19,7 @@
 import logging
 import paramiko
 import time
+from stub_sftp import StubSFTPServer
 
 logger = logging.getLogger("fake_server")
 paramiko_logger = logging.getLogger('paramiko.transport')
@@ -81,70 +82,56 @@ def listen(cmd_req_response, sock, fail_auth = False):
     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 failed: %s' % (str(e),))
         traceback.print_exc()
         return
-    # accept_thread = gevent.spawn(handle_ssh_connection,
-    #                              cmd_req_response, client, addr,
-    #                              fail_auth=fail_auth)
     handle_ssh_connection(cmd_req_response, sock, fail_auth=fail_auth)
-    # accept_thread.start()
-    # return accept_thread
 
-def _handle_ssh_connection(cmd_req_response, t, fail_auth = False):
+def _handle_ssh_connection(cmd_req_response, transport, fail_auth = False):
     try:
-        t.load_server_moduli()
+        transport.load_server_moduli()
     except:
         return
-    t.add_server_key(host_key)
+    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)
     try:
-        t.start_server(server=server)
+        transport.start_server(server=server)
     except paramiko.SSHException, e:
         logger.exception('SSH negotiation failed')
         return
     except Exception:
         logger.exception("Error occured starting server")
         return
-    return _accept_ssh_data(t, server)
-
-def _accept_ssh_data(t, server):
-    chan = t.accept(20)
-    if not chan:
+    channel = transport.accept(20)
+    if not channel:
         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():
+    while transport.is_active():
+        time.sleep(1)
+    while not channel.send_ready():
         time.sleep(.5)
-    chan.close()
+    channel.close()
     
 def handle_ssh_connection(cmd_req_response, sock, fail_auth = False):
-    client, addr = sock.accept()
+    conn, addr = sock.accept()
     logger.info('Got connection..')
     try:
-        t = paramiko.Transport(client)
-        _handle_ssh_connection(cmd_req_response, t, fail_auth=fail_auth)
+        transport = paramiko.Transport(conn)
+        _handle_ssh_connection(cmd_req_response, transport, fail_auth=fail_auth)
     except Exception, e:
         logger.error('*** Caught exception: %s: %s' % (str(e.__class__), str(e),))
         traceback.print_exc()
         try:
-            t.close()
+            transport.close()
         except:
             pass
         return
diff --git a/fake_server/stub_sftp.py b/fake_server/stub_sftp.py
@@ -0,0 +1,206 @@
+# Copyright (C) 2003-2009  Robey Pointer <robeypointer@gmail.com>
+#
+# This file is part of paramiko.
+#
+# Paramiko 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; either version 2.1 of the License, or (at your option)
+# any later version.
+#
+# Paramiko is distrubuted 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 Paramiko; if not, write to the Free Software Foundation, Inc.,
+# 59 Temple Place, Suite 330, Boston, MA  02111-1307  USA.
+
+"""
+A stub SFTP server for loopback SFTP testing.
+"""
+
+import os
+from paramiko import ServerInterface, SFTPServerInterface, SFTPServer, SFTPAttributes, \
+    SFTPHandle, SFTP_OK, AUTH_SUCCESSFUL, OPEN_SUCCEEDED
+
+
+class StubServer (ServerInterface):
+    def check_auth_password(self, username, password):
+        # all are allowed
+        return AUTH_SUCCESSFUL
+
+    def check_channel_request(self, kind, chanid):
+        return OPEN_SUCCEEDED
+
+
+class StubSFTPHandle (SFTPHandle):
+    def stat(self):
+        try:
+            return SFTPAttributes.from_stat(os.fstat(self.readfile.fileno()))
+        except OSError as e:
+            return SFTPServer.convert_errno(e.errno)
+
+    def chattr(self, attr):
+        # python doesn't have equivalents to fchown or fchmod, so we have to
+        # use the stored filename
+        try:
+            SFTPServer.set_file_attr(self.filename, attr)
+            return SFTP_OK
+        except OSError as e:
+            return SFTPServer.convert_errno(e.errno)
+
+
+class StubSFTPServer (SFTPServerInterface):
+    # assume current folder is a fine root
+    # (the tests always create and eventualy delete a subfolder, so there shouldn't be any mess)
+    ROOT = os.getcwd()
+        
+    def _realpath(self, path):
+        return self.ROOT + self.canonicalize(path)
+
+    def list_folder(self, path):
+        path = self._realpath(path)
+        try:
+            out = [ ]
+            flist = os.listdir(path)
+            for fname in flist:
+                attr = SFTPAttributes.from_stat(os.stat(os.path.join(path, fname)))
+                attr.filename = fname
+                out.append(attr)
+            return out
+        except OSError as e:
+            return SFTPServer.convert_errno(e.errno)
+
+    def stat(self, path):
+        path = self._realpath(path)
+        try:
+            return SFTPAttributes.from_stat(os.stat(path))
+        except OSError as e:
+            return SFTPServer.convert_errno(e.errno)
+
+    def lstat(self, path):
+        path = self._realpath(path)
+        try:
+            return SFTPAttributes.from_stat(os.lstat(path))
+        except OSError as e:
+            return SFTPServer.convert_errno(e.errno)
+
+    def open(self, path, flags, attr):
+        path = self._realpath(path)
+        try:
+            binary_flag = getattr(os, 'O_BINARY',  0)
+            flags |= binary_flag
+            mode = getattr(attr, 'st_mode', None)
+            if mode is not None:
+                fd = os.open(path, flags, mode)
+            else:
+                # os.open() defaults to 0777 which is
+                # an odd default mode for files
+                fd = os.open(path, flags, 0o666)
+        except OSError as e:
+            return SFTPServer.convert_errno(e.errno)
+        if (flags & os.O_CREAT) and (attr is not None):
+            attr._flags &= ~attr.FLAG_PERMISSIONS
+            SFTPServer.set_file_attr(path, attr)
+        if flags & os.O_WRONLY:
+            if flags & os.O_APPEND:
+                fstr = 'ab'
+            else:
+                fstr = 'wb'
+        elif flags & os.O_RDWR:
+            if flags & os.O_APPEND:
+                fstr = 'a+b'
+            else:
+                fstr = 'r+b'
+        else:
+            # O_RDONLY (== 0)
+            fstr = 'rb'
+        try:
+            f = os.fdopen(fd, fstr)
+        except OSError as e:
+            return SFTPServer.convert_errno(e.errno)
+        fobj = StubSFTPHandle(flags)
+        fobj.filename = path
+        fobj.readfile = f
+        fobj.writefile = f
+        return fobj
+
+    def remove(self, path):
+        path = self._realpath(path)
+        try:
+            os.remove(path)
+        except OSError as e:
+            return SFTPServer.convert_errno(e.errno)
+        return SFTP_OK
+
+    def rename(self, oldpath, newpath):
+        oldpath = self._realpath(oldpath)
+        newpath = self._realpath(newpath)
+        try:
+            os.rename(oldpath, newpath)
+        except OSError as e:
+            return SFTPServer.convert_errno(e.errno)
+        return SFTP_OK
+
+    def mkdir(self, path, attr):
+        path = self._realpath(path)
+        try:
+            os.mkdir(path)
+            if attr is not None:
+                SFTPServer.set_file_attr(path, attr)
+        except OSError as e:
+            return SFTPServer.convert_errno(e.errno)
+        return SFTP_OK
+
+    def rmdir(self, path):
+        path = self._realpath(path)
+        try:
+            os.rmdir(path)
+        except OSError as e:
+            return SFTPServer.convert_errno(e.errno)
+        return SFTP_OK
+
+    def chattr(self, path, attr):
+        path = self._realpath(path)
+        try:
+            SFTPServer.set_file_attr(path, attr)
+        except OSError as e:
+            return SFTPServer.convert_errno(e.errno)
+        return SFTP_OK
+
+    def symlink(self, target_path, path):
+        path = self._realpath(path)
+        if (len(target_path) > 0) and (target_path[0] == '/'):
+            # absolute symlink
+            target_path = os.path.join(self.ROOT, target_path[1:])
+            if target_path[:2] == '//':
+                # bug in os.path.join
+                target_path = target_path[1:]
+        else:
+            # compute relative to path
+            abspath = os.path.join(os.path.dirname(path), target_path)
+            if abspath[:len(self.ROOT)] != self.ROOT:
+                # this symlink isn't going to work anyway -- just break it immediately
+                target_path = '<error>'
+        try:
+            os.symlink(target_path, path)
+        except OSError as e:
+            return SFTPServer.convert_errno(e.errno)
+        return SFTP_OK
+
+    def readlink(self, path):
+        path = self._realpath(path)
+        try:
+            symlink = os.readlink(path)
+        except OSError as e:
+            return SFTPServer.convert_errno(e.errno)
+        # if it's absolute, remove the root
+        if os.path.isabs(symlink):
+            if symlink[:len(self.ROOT)] == self.ROOT:
+                symlink = symlink[len(self.ROOT):]
+                if (len(symlink) == 0) or (symlink[0] != '/'):
+                    symlink = '/' + symlink
+            else:
+                symlink = '<error>'
+        return symlink
diff --git a/pssh.py b/pssh.py
@@ -226,8 +226,6 @@ def __init__(self, hosts,
         """
         :param hosts: Hosts to connect to
         :type hosts: list(str)
-        :param pool_size: Pool size - how many commands to run in parallel
-        :type pool_size: int
         :param user: (Optional) User to login as. Defaults to logged in user or\
         user from ~/.ssh/config or /etc/ssh/ssh_config if set
         :type user: str
@@ -239,8 +237,7 @@ def __init__(self, hosts,
         :type port: int
         :param pkey: (Optional) Client's private key to be used to connect with
         :type pkey: :mod:`paramiko.PKey`
-        :param pool_size: (Optional) Greenlet pool size. Controls on how many\
-        hosts to execute tasks in parallel. Defaults to 10
+        :param pool_size: Pool size - how many commands to run in parallel
         :type pool_size: int
         :raises: :mod:`pssh.AuthenticationException` on authentication error
         :raises: :mod:`pssh.UnknownHostException` on DNS resolution error
diff --git a/tests/test_ssh_client.py b/tests/test_ssh_client.py
@@ -3,34 +3,50 @@
 """Unittests for parallel-ssh"""
 
 import unittest
-from pssh import SSHClient, ParallelSSHClient, UnknownHostException
-from paramiko import AuthenticationException
+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
+import os
 
-class SSHClientTest(unittest.TestCase):
-
-    def test_ssh_client(self):
-        try:
-            client = SSHClient('testy')
-        except UnknownHostException, e:
-            print e
-            return
-        stdin, stdout, stderr = client.exec_command('ls -ltrh')
-        for line in stdout:
-            print line.strip()
-        client.copy_file("fake file", "fake file")
+# _setup_logger(server_logger)
+# _setup_logger(logger)
+# _setup_logger(paramiko_logger)
 
-class ParallelSSHClientTest(unittest.TestCase):
+class SSHClientTest(unittest.TestCase):
+    
+    def setUp(self):
+        self.fake_cmd = 'fake cmd'
+        self.fake_resp = 'fake response'
+        self.listen_socket = make_socket('127.0.0.1')
+        self.listen_port = self.listen_socket.getsockname()[1]
 
-    def test_parallel_ssh_client(self):
-        client = ParallelSSHClient(['testy'])
-        cmds = client.exec_command('ls -ltrh')
-        try:
-            print [client.get_stdout(cmd) for cmd in cmds]
-        except UnknownHostException, e:
-            print e
-            return
-        cmds = client.copy_file('fake file', 'fake file')
-        client.pool.join()
+    def test_ssh_client_sftp(self):
+        """Test SFTP features of SSHClient. Copy local filename to server,
+        check that data in both files is the same, make new directory on
+        server, remove files and directory."""
+        test_file_data = 'test'
+        local_filename = 'test_file'
+        remote_filename = 'test_file_copy'
+        remote_dir = 'remote_dir'
+        test_file = open(local_filename, 'w')
+        test_file.writelines([test_file_data + os.linesep])
+        test_file.close()
+        server = start_server({ self.fake_cmd : self.fake_resp }, self.listen_socket)
+        client = SSHClient('127.0.0.1', port=self.listen_port)
+        client.copy_file(local_filename, remote_filename)
+        self.assertTrue(os.path.isfile(remote_filename), msg="SFTP copy failed")
+        copied_file = open(remote_filename, 'r')
+        copied_file_data = copied_file.readlines()[0].strip()
+        copied_file.close()
+        self.assertEqual(test_file_data, copied_file_data,
+                         msg="Data in destination file %s does not match source %s" % (copied_file_data, test_file_data))
+        os.unlink(local_filename)
+        os.unlink(remote_filename)
+        client.mkdir(client._make_sftp(), remote_dir)
+        self.assertTrue(os.path.isdir(remote_dir))
+        os.rmdir(remote_dir)
+        del client
+        server.join()
 
 if __name__ == '__main__':
     unittest.main()
