Refactor SSH testing code out of individual tasks

Catch SSH errors and report failures as hints. Also some cleanups and
2to3 fixes.

1 files changed, 22 insertions, 0 deletions

diff --git a/kpov_judge/kpov_util.py b/kpov_judge/kpov_util.py
index 74b1737..9909973 100755
--- a/kpov_judge/kpov_util.py
+++ b/kpov_judge/kpov_util.py
@@ -1,6 +1,7 @@
 #!/usr/bin/env python
 # -*- coding: utf-8 -*-
 
+import collections
 import itertools
 import random
 import string
@@ -12,6 +13,27 @@ import subprocess
 import glob
 import os
 
+def ssh_test(host, user, password, commands=()):
+    from pexpect import pxssh
+    results = collections.defaultdict(str)
+    try:
+        s = pxssh.pxssh(encoding='utf-8')
+        s.login(host, user, password,
+            original_prompt='~[#$] ',
+            auto_prompt_reset=False)
+        results['ssh'] = True
+        results['motd'] = s.before
+        s.set_unique_prompt()
+        for test, command in commands:
+            s.sendline(command)
+            s.prompt()
+            if test:
+                results[test] = s.before[len(command+'\r\n'):]
+        s.logout()
+    except Exception as ex:
+        results['ssh'] = str(ex)
+    return results
+
 def alnum_gen(r, l=1):
     s = ""
     for i in range(l):