ssh_support: Close ssh clients properly

The proxy client was not closed after the command has executed, and in
case of an error condition the regular client was also not closed.

Fixes #25

1 files changed, 24 insertions, 11 deletions

diff --git a/rumba/ssh_support.py b/rumba/ssh_support.py
index 9aa312e..5f69439 100644
--- a/rumba/ssh_support.py
+++ b/rumba/ssh_support.py
@@ -64,6 +64,7 @@ def ssh_connect(hostname, port, username, password, time_out, proxy_server):
         time.sleep(retry * 5)
 
         try:
+            proxy_client = None
             if proxy_server is not None:
                 proxy_client = get_ssh_client()
                 # Assume port 22 for the proxy server for now
@@ -80,7 +81,7 @@ def ssh_connect(hostname, port, username, password, time_out, proxy_server):
 
             ssh_client.connect(hostname, port, username, password,
                                look_for_keys=True, timeout=time_out, sock=proxy)
-            return ssh_client
+            return ssh_client, proxy_client
         except paramiko.ssh_exception.BadHostKeyException:
             retry += 1
             logger.error(hostname + ' has a mismatching entry in ' +
@@ -154,9 +155,9 @@ def execute_commands(testbed, ssh_config, commands, time_out=3):
     be used when no timeout is needed
     """
 
-    ssh_client = ssh_connect(ssh_config.hostname, ssh_config.port,
-                             testbed.username, testbed.password, time_out,
-                             ssh_config.proxy_server)
+    ssh_client, proxy_client = ssh_connect(ssh_config.hostname, ssh_config.port,
+                                           testbed.username, testbed.password,
+                                           time_out, ssh_config.proxy_server)
 
     o = ""
     for command in commands:
@@ -174,9 +175,14 @@ def execute_commands(testbed, ssh_config, commands, time_out=3):
             raise SSHException('Failed to execute command')
         o = _print_stream(stdout)
         if (chan.recv_exit_status() != 0):
+            ssh_client.close()
+            if proxy_client is not None:
+                proxy_client.close()
             raise SSHException('A remote command returned an error.\n' + o)
 
     ssh_client.close()
+    if proxy_client is not None:
+        proxy_client.close()
     return o
 
 
@@ -210,9 +216,9 @@ def write_text_to_file(testbed, ssh_config, text, file_name):
     @param file_name: file name (including full path) on the host
     """
 
-    ssh_client = ssh_connect(ssh_config.hostname, ssh_config.port,
-                             testbed.username, testbed.password, None,
-                             ssh_config.proxy_server)
+    ssh_client, proxy_client = ssh_connect(ssh_config.hostname, ssh_config.port,
+                                           testbed.username, testbed.password,
+                                           None, ssh_config.proxy_server)
 
     cmd = "touch " + file_name + "; chmod a+rwx " + file_name
 
@@ -231,7 +237,10 @@ def write_text_to_file(testbed, ssh_config, text, file_name):
 
     except SSHException as e:
         raise SSHException('Failed to write text to remote file')
-
+    finally:
+        ssh_client.close()
+        if proxy_client is not None:
+            proxy_client.close()
 
 def copy_files_to_testbed(testbed, ssh_config, paths, destination):
     """
@@ -245,9 +254,9 @@ def copy_files_to_testbed(testbed, ssh_config, paths, destination):
     if destination is not '' and not destination.endswith('/'):
         destination = destination + '/'
 
-    ssh_client = ssh_connect(ssh_config.hostname, ssh_config.port,
-                             testbed.username, testbed.password, None,
-                             ssh_config.proxy_server)
+    ssh_client, proxy_client = ssh_connect(ssh_config.hostname, ssh_config.port,
+                                           testbed.username, testbed.password,
+                                           None, ssh_config.proxy_server)
 
     try:
         sftp_client = ssh_client.open_sftp()
@@ -265,6 +274,10 @@ def copy_files_to_testbed(testbed, ssh_config, paths, destination):
 
     except Exception as e:
         raise SSHException('Failed to copy files to testbed')
+    finally:
+        ssh_client.close()
+        if proxy_client is not None:
+            proxy_client.close()
 
 
 def copy_file_to_testbed(testbed, ssh_config, path, destination):