Implement #689

grycap · Oct 19, 2018 · d063fd0 · d063fd0
1 parent bd3ba34
commit d063fd0
Showing 3 changed files with 151 additions and 106 deletions.
diff --git a/IM/SSH.py b/IM/SSH.py
@@ -234,9 +234,11 @@ def sftp_get(self, src, dest):
     def _sftp_put(sftp, src, dest):
         f_flags = LIBSSH2_FXF_CREAT | LIBSSH2_FXF_WRITE
         fileinfo = os.stat(src)
-        with open(src, 'rb') as local_fh, sftp.open(dest, f_flags, fileinfo.st_mode) as remote_fh:
+        remote_fh = sftp.open(dest, f_flags, fileinfo.st_mode)
+        with open(src, 'rb') as local_fh:
             for data in local_fh:
                 remote_fh.write(data)
+        remote_fh.close()
 
     def sftp_put(self, src, dest):
         """ Puts a file to the remote server
@@ -396,10 +398,10 @@ def sftp_list(self, directory):
         client = self.connect()
         sftp = client.sftp_init()
         res = []
-        with sftp.opendir(directory) as fh:
-            for _, name, attrs in fh.readdir():
-                print(attrs.permissions)
-                res.append(name)
+        fh = sftp.opendir(directory)
+        for _, name, _ in fh.readdir():
+            res.append(name)
+        fh.close()
         return res
 
     def sftp_list_attr(self, directory):
@@ -415,9 +417,10 @@ def sftp_list_attr(self, directory):
         client = self.connect()
         sftp = client.sftp_init()
         res = []
-        with sftp.opendir(directory) as fh:
-            for _, _, attrs in fh.readdir():
-                res.append(attrs)
+        fh = sftp.opendir(directory)
+        for _, _, attrs in fh.readdir():
+            res.append(attrs)
+        fh.close()
         return res
 
     def getcwd(self):

diff --git a/IM/SSHRetry.py b/IM/SSHRetry.py
@@ -17,7 +17,7 @@
 
 from IM.retry import retry
 from IM.SSH import SSH, AuthenticationException
-import paramiko
+from ssh2.exceptions import AuthenticationError
 
 
 class SSHRetry(SSH):
@@ -26,72 +26,72 @@ class SSHRetry(SSH):
     DELAY = 3
     BACKOFF = 2
 
-    @retry(Exception, (AuthenticationException, paramiko.AuthenticationException),
+    @retry(Exception, (AuthenticationException, AuthenticationError),
            tries=TRIES, delay=DELAY, backoff=BACKOFF)
     def execute(self, command, timeout=None):
         return SSH.execute(self, command, timeout)
 
-    @retry(Exception, (AuthenticationException, paramiko.AuthenticationException),
+    @retry(Exception, (AuthenticationException, AuthenticationError),
            tries=TRIES, delay=DELAY, backoff=BACKOFF)
     def sftp_get(self, src, dest):
         return SSH.sftp_get(self, src, dest)
 
-    @retry(Exception, (AuthenticationException, paramiko.AuthenticationException),
+    @retry(Exception, (AuthenticationException, AuthenticationError),
            tries=TRIES, delay=DELAY, backoff=BACKOFF)
     def sftp_get_files(self, src, dest):
         return SSH.sftp_get_files(self, src, dest)
 
-    @retry(Exception, (AuthenticationException, paramiko.AuthenticationException),
+    @retry(Exception, (AuthenticationException, AuthenticationError),
            tries=TRIES, delay=DELAY, backoff=BACKOFF)
     def sftp_put_files(self, files):
         return SSH.sftp_put_files(self, files)
 
-    @retry(Exception, (AuthenticationException, paramiko.AuthenticationException),
+    @retry(Exception, (AuthenticationException, AuthenticationError),
            tries=TRIES, delay=DELAY, backoff=BACKOFF)
     def sftp_put(self, src, dest):
         return SSH.sftp_put(self, src, dest)
 
-    @retry(Exception, (AuthenticationException, paramiko.AuthenticationException),
+    @retry(Exception, (AuthenticationException, AuthenticationError),
            tries=TRIES, delay=DELAY, backoff=BACKOFF)
     def sftp_put_dir(self, src, dest):
         return SSH.sftp_put_dir(self, src, dest)
 
-    @retry(Exception, (AuthenticationException, paramiko.AuthenticationException),
+    @retry(Exception, (AuthenticationException, AuthenticationError),
            tries=TRIES, delay=DELAY, backoff=BACKOFF)
     def sftp_get_dir(self, src, dest):
         return SSH.sftp_get_dir(self, src, dest)
 
-    @retry(Exception, (AuthenticationException, paramiko.AuthenticationException),
+    @retry(Exception, (AuthenticationException, AuthenticationError),
            tries=TRIES, delay=DELAY, backoff=BACKOFF)
     def sftp_put_content(self, content, dest):
         return SSH.sftp_put_content(self, content, dest)
 
-    @retry(Exception, (AuthenticationException, paramiko.AuthenticationException),
+    @retry(Exception, (AuthenticationException, AuthenticationError),
            tries=TRIES, delay=DELAY, backoff=BACKOFF)
     def sftp_mkdir(self, directory):
         return SSH.sftp_mkdir(self, directory)
 
-    @retry(Exception, (AuthenticationException, paramiko.AuthenticationException),
+    @retry(Exception, (AuthenticationException, AuthenticationError),
            tries=TRIES, delay=DELAY, backoff=BACKOFF)
     def sftp_list(self, directory):
         return SSH.sftp_list(self, directory)
 
-    @retry(Exception, (AuthenticationException, paramiko.AuthenticationException),
+    @retry(Exception, (AuthenticationException, AuthenticationError),
            tries=TRIES, delay=DELAY, backoff=BACKOFF)
     def sftp_list_attr(self, directory):
         return SSH.sftp_list_attr(self, directory)
 
-    @retry(Exception, (AuthenticationException, paramiko.AuthenticationException),
+    @retry(Exception, (AuthenticationException, AuthenticationError),
            tries=TRIES, delay=DELAY, backoff=BACKOFF)
     def getcwd(self):
         return SSH.getcwd(self)
 
-    @retry(Exception, (AuthenticationException, paramiko.AuthenticationException),
+    @retry(Exception, (AuthenticationException, AuthenticationError),
            tries=TRIES, delay=DELAY, backoff=BACKOFF)
     def sftp_remove(self, path):
         return SSH.sftp_remove(self, path)
 
-    @retry(Exception, (AuthenticationException, paramiko.AuthenticationException),
+    @retry(Exception, (AuthenticationException, AuthenticationError),
            tries=TRIES, delay=DELAY, backoff=BACKOFF)
     def sftp_chmod(self, path, mode):
         return SSH.sftp_chmod(self, path, mode)
diff --git a/test/unit/SSH.py b/test/unit/SSH.py
@@ -18,6 +18,7 @@
 
 import unittest
 import os
+from io import BytesIO
 
 from IM.SSHRetry import SSHRetry
 from mock import patch, MagicMock
@@ -34,161 +35,202 @@ class TestSSH(unittest.TestCase):
     Class to test the SSH class
     """
 
-    @patch('paramiko.SSHClient')
-    def test_test_connectivity(self, ssh_client):
+    @patch('socket.socket')
+    @patch('IM.SSH.Session')
+    def test_test_connectivity(self, session, socket):
+        sess = MagicMock()
+        session.return_value = sess
         ssh = SSHRetry("host", "user", "passwd", read_file_as_string("../files/privatekey.pem"))
         success = ssh.test_connectivity(5)
         self.assertTrue(success)
+        self.assertEquals(sess.userauth_password.call_args_list[0][0], ("user", "passwd"))
+
+    @patch('socket.socket')
+    @patch('IM.SSH.Session')
+    def test_execute(self, session, socket):
+        sess = MagicMock()
+        session.return_value = sess
 
-    @patch('paramiko.SSHClient')
-    def test_execute(self, ssh_client):
         ssh = SSHRetry("host", "user", "passwd", read_file_as_string("../files/privatekey.pem"))
 
-        client = MagicMock()
-        ssh_client.return_value = client
-        transport = MagicMock()
         channel = MagicMock()
-        client.get_transport.return_value = transport
-        transport.open_session.return_value = channel
-        channel.makefile.return_value = ["out"]
-        channel.makefile_stderr.return_value = ["err"]
-        channel.recv_exit_status.return_value = 0
+        sess.open_session.return_value = channel
+
+        channel.read.side_effect = [(3, "out"), (0, "")]
+        channel.read_stderr.side_effect = [(3, "err"), (0, "")]
+        channel.get_exit_status.return_value = 0
 
         (res_stdout, res_stderr, exit_status) = ssh.execute("ls")
         self.assertEqual(res_stdout, "out")
         self.assertEqual(res_stderr, "err")
         self.assertEqual(exit_status, 0)
 
-    @patch('paramiko.SSHClient')
-    @patch('paramiko.SFTPClient')
-    def test_sftp_get(self, sftp_client, ssh_client):
+    @patch('socket.socket')
+    @patch('IM.SSH.Session')
+    def test_sftp_put(self, session, socket):
+        sess = MagicMock()
+        session.return_value = sess
+        sftp = MagicMock()
+        sess.sftp_init.return_value = sftp
+        remote_fh = MagicMock()
+        sftp.open.return_value = remote_fh
+
+        src = "/tmp/some_file"
+        with open(src, 'wb+') as fh:
+            fh.write("some_data")
+
         ssh = SSHRetry("host", "user", "passwd", read_file_as_string("../files/privatekey.pem"))
+        ssh.sftp_put(src, src)
 
-        ssh.sftp_get("some_file", "some_file")
+        self.assertEqual(remote_fh.write.call_args_list[0][0][0], "some_data")
+
+    @patch('socket.socket')
+    @patch('IM.SSH.Session')
+    def test_sftp_get_files(self, session, socket):
+        sess = MagicMock()
+        session.return_value = sess
+        sftp = MagicMock()
+        sess.sftp_init.return_value = sftp
 
-    @patch('paramiko.SSHClient')
-    @patch('paramiko.SFTPClient')
-    def test_sftp_get_files(self, sftp_client, ssh_client):
         ssh = SSHRetry("host", "user", "passwd", read_file_as_string("../files/privatekey.pem"))
 
-        ssh.sftp_get(["some_file"], ["some_file"])
+        ssh.sftp_get_files(["some_file"], ["some_file"])
 
-    @patch('paramiko.SSHClient')
-    @patch('paramiko.SFTPClient')
-    def test_sftp_put(self, sftp_client, ssh_client):
+    @patch('socket.socket')
+    @patch('IM.SSH.Session')
+    def test_sftp_get(self, session, socket):
         ssh = SSHRetry("host", "user", "passwd", read_file_as_string("../files/privatekey.pem"))
 
-        ssh.sftp_put("some_file", "some_file")
+        ssh.sftp_get("some_file", "some_file")
 
-    @patch('paramiko.SSHClient')
-    @patch('paramiko.SFTPClient')
-    def test_sftp_put_files(self, sftp_client, ssh_client):
+    @patch('socket.socket')
+    @patch('IM.SSH.Session')
+    def test_sftp_put_files(self, session, socket):
         ssh = SSHRetry("host", "user", "passwd", read_file_as_string("../files/privatekey.pem"))
 
         ssh.sftp_put_files([("some_file", "some_file")])
 
-    @patch('paramiko.SSHClient')
-    @patch('paramiko.SFTPClient')
-    def test_sftp_put_dir(self, sftp_client, ssh_client):
+    @patch('socket.socket')
+    @patch('IM.SSH.Session')
+    def test_sftp_put_dir(self, session, socket):
         ssh = SSHRetry("host", "user", "passwd", read_file_as_string("../files/privatekey.pem"))
 
         ssh.sftp_put_dir("/tmp", "/tmp")
 
-    @patch('paramiko.SSHClient')
-    @patch('paramiko.SFTPClient')
-    def test_sftp_get_dir(self, sftp_client, ssh_client):
+    @patch('socket.socket')
+    @patch('IM.SSH.Session')
+    def test_sftp_get_dir(self, session, socket):
         ssh = SSHRetry("host", "user", "passwd", read_file_as_string("../files/privatekey.pem"))
 
         ssh.sftp_get_dir("/tmp", "/tmp")
 
-    @patch('paramiko.SSHClient')
-    @patch('paramiko.SFTPClient')
-    def test_sftp_put_content(self, sftp_client, ssh_client):
+    @patch('socket.socket')
+    @patch('IM.SSH.Session')
+    def test_sftp_put_content(self, session, socket):
         ssh = SSHRetry("host", "user", "passwd", read_file_as_string("../files/privatekey.pem"))
 
         ssh.sftp_put_content("some_file", "some_content")
 
-    @patch('paramiko.SSHClient')
-    @patch('paramiko.SFTPClient')
-    def test_sftp_mkdir(self, sftp_client, ssh_client):
+    @patch('socket.socket')
+    @patch('IM.SSH.Session')
+    def test_sftp_mkdir(self, session, socket):
         ssh = SSHRetry("host", "user", "passwd", read_file_as_string("../files/privatekey.pem"))
 
         ssh.sftp_mkdir("/some_dir")
 
-    @patch('paramiko.SSHClient')
-    @patch('paramiko.SFTPClient.from_transport')
-    def test_sftp_list(self, from_transport, ssh_client):
-        ssh = SSHRetry("host", "user", "passwd", read_file_as_string("../files/privatekey.pem"))
+    @patch('socket.socket')
+    @patch('IM.SSH.Session')
+    def test_sftp_list(self, session, socket):
+        sess = MagicMock()
+        session.return_value = sess
+        sftp = MagicMock()
+        sess.sftp_init.return_value = sftp
+        dir = MagicMock()
+        dir.readdir.return_value = [("", "file", "")]
+        sftp.opendir.return_value = dir
 
-        client = MagicMock()
-        from_transport.return_value = client
-        client.listdir.return_value = ["file"]
+        ssh = SSHRetry("host", "user", "passwd", read_file_as_string("../files/privatekey.pem"))
 
         res = ssh.sftp_list("/some_dir")
         self.assertEqual(res, ["file"])
 
-    @patch('paramiko.SSHClient')
-    @patch('paramiko.SFTPClient.from_transport')
-    def test_sftp_list_attr(self, from_transport, ssh_client):
-        ssh = SSHRetry("host", "user", "passwd", read_file_as_string("../files/privatekey.pem"))
+    @patch('socket.socket')
+    @patch('IM.SSH.Session')
+    def test_sftp_list_attr(self, session, socket):
+        sess = MagicMock()
+        session.return_value = sess
+        sftp = MagicMock()
+        sess.sftp_init.return_value = sftp
+        dir = MagicMock()
+        dir.readdir.return_value = [("", "", "attrs")]
+        sftp.opendir.return_value = dir
 
-        client = MagicMock()
-        from_transport.return_value = client
-        client.listdir_attr.return_value = ["file"]
+        ssh = SSHRetry("host", "user", "passwd", read_file_as_string("../files/privatekey.pem"))
 
         res = ssh.sftp_list_attr("/some_dir")
-        self.assertEqual(res, ["file"])
+        self.assertEqual(res, ["attrs"])
+
+    @patch('socket.socket')
+    @patch('IM.SSH.Session')
+    def test_getcwd(self, session, socket):
+        sess = MagicMock()
+        session.return_value = sess
 
-    @patch('paramiko.SSHClient')
-    @patch('paramiko.SFTPClient.from_transport')
-    def test_getcwd(self, from_transport, ssh_client):
         ssh = SSHRetry("host", "user", "passwd", read_file_as_string("../files/privatekey.pem"))
 
-        client = MagicMock()
-        from_transport.return_value = client
-        client.getcwd.return_value = "some_dir"
+        channel = MagicMock()
+        sess.open_session.return_value = channel
+
+        channel.read.side_effect = [(3, "some_dir"), (0, "")]
+        channel.read_stderr.side_effect = [(3, "err"), (0, "")]
+        channel.get_exit_status.return_value = 0
 
         res = ssh.getcwd()
         self.assertEqual(res, "some_dir")
 
-    @patch('paramiko.SSHClient')
-    def test_execute_timeout(self, ssh_client):
+    @patch('socket.socket')
+    @patch('IM.SSH.Session')
+    def test_execute_timeout(self, session, socket):
+        sess = MagicMock()
+        session.return_value = sess
+
         ssh = SSHRetry("host", "user", "passwd", read_file_as_string("../files/privatekey.pem"))
 
-        client = MagicMock()
-        ssh_client.return_value = client
-        tansport = MagicMock()
-        client.get_transport.return_value = tansport
         channel = MagicMock()
-        tansport.open_session.return_value = channel
-        channel.makefile.return_value = "out"
-        channel.makefile_stderr.return_value = "err"
-        channel.recv_exit_status.return_value = 0
+        sess.open_session.return_value = channel
+
+        channel.read.side_effect = [(3, "out"), (0, "")]
+        channel.read_stderr.side_effect = [(3, "err"), (0, "")]
+        channel.get_exit_status.return_value = 0
 
         res_stdout, res_stderr, code = ssh.execute_timeout("ls", 5)
         self.assertEqual(res_stdout, "out")
         self.assertEqual(res_stderr, "err")
         self.assertEqual(code, 0)
 
-    @patch('paramiko.SSHClient')
-    @patch('paramiko.SFTPClient.from_transport')
-    def test_sftp_remove(self, from_transport, ssh_client):
-        ssh = SSHRetry("host", "user", "passwd", read_file_as_string("../files/privatekey.pem"))
+    @patch('socket.socket')
+    @patch('IM.SSH.Session')
+    def test_sftp_remove(self, session, socket):
+        sess = MagicMock()
+        session.return_value = sess
+        sftp = MagicMock()
+        sess.sftp_init.return_value = sftp
+        sftp.stat.side_effect = Exception()
 
-        client = MagicMock()
-        from_transport.return_value = client
-        client.remove.return_value = True
+        ssh = SSHRetry("host", "user", "passwd", read_file_as_string("../files/privatekey.pem"))
 
         res = ssh.sftp_remove("some_file")
         self.assertTrue(res)
 
-    @patch('paramiko.SSHClient')
-    @patch('paramiko.SFTPClient.from_transport')
-    def test_sftp_chmod(self, from_transport, ssh_client):
-        ssh = SSHRetry("host", "user", "passwd", read_file_as_string("../files/privatekey.pem"))
+    @patch('socket.socket')
+    @patch('IM.SSH.Session')
+    def test_sftp_chmod(self, session, socket):
+        sess = MagicMock()
+        session.return_value = sess
+        sftp = MagicMock()
+        sess.sftp_init.return_value = sftp
 
-        client = MagicMock()
-        from_transport.return_value = client
+        ssh = SSHRetry("host", "user", "passwd", read_file_as_string("../files/privatekey.pem"))
 
         res = ssh.sftp_chmod("some_file", 0o644)
         self.assertTrue(res)