rumba: Add better error handling of SSH connect

If an SSH connect doesn't succeed at first, it is retried up to 5
times with an exponential backoff. It also adds an error message in
case of an offending host key.

3 files changed, 45 insertions, 24 deletions

diff --git a/rumba/prototypes/enroll.py b/rumba/prototypes/enroll.py
index 774ba52..35cc136 100755
--- a/rumba/prototypes/enroll.py
+++ b/rumba/prototypes/enroll.py
@@ -28,6 +28,7 @@ import argparse
 import socket
 import time
 import re
+import sys
 
 def printalo(byt):
     print(repr(byt).replace('\\n', '\n'))
@@ -133,5 +134,6 @@ if connected:
 
 else:
     print('Failed to connect to "%s"' % socket_name)
+    sys.exit(-1)
 
 s.close()
diff --git a/rumba/prototypes/irati.py b/rumba/prototypes/irati.py
index f104054..48b013a 100644
--- a/rumba/prototypes/irati.py
+++ b/rumba/prototypes/irati.py
@@ -79,7 +79,7 @@ class Experiment(mod.Experiment):
                 "export https_proxy=\"https://proxy.atlantis.ugent.be:8080\"; "
                 + self.sudo("apt-get install g++ gcc "
                             "protobuf-compiler libprotobuf-dev git --yes "
-                            "pkg-config "
+                            "pkg-config ethtool socat "
                             "libnl-3-dev libnl-genl-3-dev"),
                 self.sudo("rm -rf ~/irati"),
                 "cd ~; "
diff --git a/rumba/ssh_support.py b/rumba/ssh_support.py
index dfa290a..64d17e6 100644
--- a/rumba/ssh_support.py
+++ b/rumba/ssh_support.py
@@ -25,9 +25,15 @@
 
 import os
 import paramiko
+import time
 
 import rumba.log as log
 
+# Fix Python 2.x.
+try:
+    input = raw_input
+except NameError:
+    pass
 
 logger = log.get_logger(__name__)
 
@@ -51,6 +57,30 @@ def _print_stream(stream):
                 logger.debug(oi)
         return o.rstrip()
 
+def ssh_connect(ssh_client, hostname, port, username, password, time_out,
+                proxy):
+    retry = 0
+    max_retries = 5
+    while retry < max_retries:
+        time.sleep(retry * 5)
+        try:
+            ssh_client.connect(hostname, port, username, password,
+                               look_for_keys=True, timeout=time_out, sock=proxy)
+            return
+        except paramiko.ssh_exception.SSHException:
+            retry += 1
+            logger.error('Failed to connect to host, retrying: ' +
+                         str(retry) + '/' + str(max_retries) + ' retries')
+        except paramiko.ssh_exception.BadHostKeyException:
+            retry += 1
+            logger.error(hostname + ' has a mismatching entry in ' +
+                         '~/.ssh/known_hosts')
+            logger.error('If you are sure this is not a man in the ' +
+                         'middle attack, edit that file to remove the' +
+                         'entry and then hit return to try again.')
+            input('Hit Enter when ready')
+    if retry == max_retries:
+        raise SSHException('Failed to connect to host')
 
 def execute_proxy_commands(testbed, ssh_config, commands, time_out=3):
     """
@@ -109,20 +139,15 @@ def execute_commands(testbed, ssh_config, commands, time_out=3):
     no result received in given number of seconds, the value None can
     be used when no timeout is needed
     """
-    ssh_client = get_ssh_client()
-
     if ssh_config.proxy_command is not None:
         proxy = paramiko.ProxyCommand(ssh_config.proxy_command)
     else:
         proxy = None
 
-    try:
-        ssh_client.connect(ssh_config.hostname, ssh_config.port,
-                           testbed.username, testbed.password,
-                           look_for_keys=True, timeout=time_out,
-                           sock=proxy)
-    except paramiko.ssh_exception.SSHException as e:
-        raise SSHException('Failed to connect to host')
+    ssh_client = get_ssh_client()
+
+    ssh_connect(ssh_client, ssh_config.hostname, ssh_config.port,
+                testbed.username, testbed.password, time_out, proxy)
 
     o = ""
     for command in commands:
@@ -182,16 +207,12 @@ def write_text_to_file(testbed, ssh_config, text, file_name):
     else:
         proxy = None
 
-    try:
-        ssh_client.connect(ssh_config.hostname, ssh_config.port,
-                           testbed.username,
-                           testbed.password,
-                           look_for_keys=True,
-                           sock=proxy)
+    ssh_connect(ssh_client, ssh_config.hostname, ssh_config.port,
+                testbed.username, testbed.password, None, proxy)
 
-        cmd = "touch " + file_name + \
-              "; chmod a+rwx " + file_name
+    cmd = "touch " + file_name + "; chmod a+rwx " + file_name
 
+    try:
         stdin, stdout, stderr = ssh_client.exec_command(cmd)
         del stdin, stdout
         err = str(stderr.read()).strip('b\'\"\\n')
@@ -227,13 +248,11 @@ def copy_files_to_testbed(testbed, ssh_config, paths, destination):
     if destination is not '' and not destination.endswith('/'):
         destination = destination + '/'
 
-    try:
-        ssh_client.connect(ssh_config.hostname, ssh_config.port,
-                           testbed.username,
-                           testbed.password,
-                           look_for_keys=True,
-                           sock=proxy)
 
+    ssh_connect(ssh_client, ssh_config.hostname, ssh_config.port,
+                testbed.username, testbed.password, None, proxy)
+
+    try:
         sftp_client = ssh_client.open_sftp()
 
         for path in paths: