Honeycomb setup and utils
[csit.git] / resources / libraries / python / ssh.py
index 2de6f4a..6914d52 100644 (file)
@@ -10,6 +10,7 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+import socket
 import paramiko
 from paramiko import RSAKey
 import StringIO
@@ -18,7 +19,6 @@ from time import time
 from robot.api import logger
 from interruptingcow import timeout
 from robot.utils.asserts import assert_equal, assert_not_equal
-from socket import timeout as socket_timeout
 
 __all__ = ["exec_cmd", "exec_cmd_no_error"]
 
@@ -56,7 +56,8 @@ class SSH(object):
             self._ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy())
 
             self._ssh.connect(node['host'], username=node['username'],
-                              password=node.get('password'), pkey=pkey)
+                              password=node.get('password'), pkey=pkey,
+                              port=node['port'])
 
             SSH.__existing_connections[node_hash] = self._ssh
 
@@ -67,6 +68,17 @@ class SSH(object):
                 format(self._ssh.get_transport().getpeername()))
         logger.debug('Connections: {0}'.format(str(SSH.__existing_connections)))
 
+    def disconnect(self, node):
+        """Close SSH connection to the node.
+
+        :param node: The node to disconnect from.
+        :type node: dict
+        """
+        node_hash = self._node_hash(node)
+        if node_hash in SSH.__existing_connections:
+            ssh = SSH.__existing_connections.pop(node_hash)
+            ssh.close()
+
     def exec_command(self, cmd, timeout=10):
         """Execute SSH command on a new channel on the connected Node.
 
@@ -84,25 +96,23 @@ class SSH(object):
             self._ssh.get_transport().getpeername(), end-start))
 
         stdout = ""
-        while True:
-            try:
+        try:
+            while True:
                 buf = chan.recv(self.__MAX_RECV_BUF)
                 stdout += buf
                 if not buf:
                     break
-            except socket_timeout:
-                logger.trace('Channels stdout timeout occurred')
-                break
+        except socket.timeout:
+            logger.error('Caught timeout exception, current contents '
+                         'of buffer: {0}'.format(stdout))
+            raise
+
 
         stderr = ""
         while True:
-            try:
-                buf = chan.recv_stderr(self.__MAX_RECV_BUF)
-                stderr += buf
-                if not buf:
-                    break
-            except socket_timeout:
-                logger.trace('Channels stderr timeout occurred')
+            buf = chan.recv_stderr(self.__MAX_RECV_BUF)
+            stderr += buf
+            if not buf:
                 break
 
         return_code = chan.recv_exit_status()