Create ssh_key shared module under client/shared

The SSH creation is a common method, so should share the get_public_key and
setup_ssh_key functions for both server and client.

* Create new ssh_key.py under the client/shared
* Move get_public_key from server/base_utils.py into ssh_key.py
* Move setup_ssh_key from server/hosts/ssh_host.py into ssh_key.py
* Sync ssh_host module with setup_ssh_key functions change

Signed-off-by: Wayne Sun <[email protected]>
Signed-off-by: Alex Jia <[email protected]>

client/shared/ssh_key.py

@@ -0,0 +1,72 @@
+import os
+from autotest.client.shared import pxssh
+from autotest.client.shared import base_utils as utils
+
+def get_public_key():
+    """
+    Return a valid string ssh public key for the user executing autoserv or
+    autotest. If there's no DSA or RSA public key, create a DSA keypair with
+    ssh-keygen and return it.
+    """
+
+    ssh_conf_path = os.path.expanduser('~/.ssh')
+
+    dsa_public_key_path = os.path.join(ssh_conf_path, 'id_dsa.pub')
+    dsa_private_key_path = os.path.join(ssh_conf_path, 'id_dsa')
+
+    rsa_public_key_path = os.path.join(ssh_conf_path, 'id_rsa.pub')
+    rsa_private_key_path = os.path.join(ssh_conf_path, 'id_rsa')
+
+    has_dsa_keypair = os.path.isfile(dsa_public_key_path) and \
+        os.path.isfile(dsa_private_key_path)
+    has_rsa_keypair = os.path.isfile(rsa_public_key_path) and \
+        os.path.isfile(rsa_private_key_path)
+
+    if has_dsa_keypair:
+        print 'DSA keypair found, using it'
+        public_key_path = dsa_public_key_path
+
+    elif has_rsa_keypair:
+        print 'RSA keypair found, using it'
+        public_key_path = rsa_public_key_path
+
+    else:
+        print 'Neither RSA nor DSA keypair found, creating DSA ssh key pair'
+        utils.system('ssh-keygen -t dsa -q -N "" -f %s' % dsa_private_key_path)
+        public_key_path = dsa_public_key_path
+
+    public_key = open(public_key_path, 'r')
+    public_key_str = public_key.read()
+    public_key.close()
+
+    return public_key_str
+
+
+def setup_ssh_key(hostname, user, password, port):
+    logging.debug('Performing SSH key setup on %s:%d as %s.' %
+                  (self.hostname, self.port, self.user))
+
+    try:
+        host = pxssh.pxssh()
+        host.login(hostname, user, password, port)
+        public_key = get_public_key()
+
+        host.sendline('mkdir -p ~/.ssh')
+        host.prompt()
+        host.sendline('chmod 700 ~/.ssh')
+        host.prompt()
+        host.sendline("echo '%s' >> ~/.ssh/authorized_keys; " %
+                        public_key)
+        host.prompt()
+        host.sendline('chmod 600 ~/.ssh/authorized_keys')
+        host.prompt()
+        host.logout()
+
+        logging.debug('SSH key setup complete.')
+
+    except:
+        logging.debug('SSH key setup has failed.')
+        try:
+            host.logout()
+        except:
+            pass

server/base_utils.py

@@ -273,46 +273,6 @@ def parse_machine(machine, user='root', password='', port=22, profile=''):
     return machine, user, password, port, profile
 
 
-def get_public_key():
-    """
-    Return a valid string ssh public key for the user executing autoserv or
-    autotest. If there's no DSA or RSA public key, create a DSA keypair with
-    ssh-keygen and return it.
-    """
-
-    ssh_conf_path = os.path.expanduser('~/.ssh')
-
-    dsa_public_key_path = os.path.join(ssh_conf_path, 'id_dsa.pub')
-    dsa_private_key_path = os.path.join(ssh_conf_path, 'id_dsa')
-
-    rsa_public_key_path = os.path.join(ssh_conf_path, 'id_rsa.pub')
-    rsa_private_key_path = os.path.join(ssh_conf_path, 'id_rsa')
-
-    has_dsa_keypair = os.path.isfile(dsa_public_key_path) and \
-        os.path.isfile(dsa_private_key_path)
-    has_rsa_keypair = os.path.isfile(rsa_public_key_path) and \
-        os.path.isfile(rsa_private_key_path)
-
-    if has_dsa_keypair:
-        print 'DSA keypair found, using it'
-        public_key_path = dsa_public_key_path
-
-    elif has_rsa_keypair:
-        print 'RSA keypair found, using it'
-        public_key_path = rsa_public_key_path
-
-    else:
-        print 'Neither RSA nor DSA keypair found, creating DSA ssh key pair'
-        utils.system('ssh-keygen -t dsa -q -N "" -f %s' % dsa_private_key_path)
-        public_key_path = dsa_public_key_path
-
-    public_key = open(public_key_path, 'r')
-    public_key_str = public_key.read()
-    public_key.close()
-
-    return public_key_str
-
-
 def get_sync_control_file(control, host_name, host_num,
                           instance, num_jobs, port_base=63100):
     """

server/hosts/ssh_host.py

@@ -11,7 +11,7 @@
 """
 
 import sys, re, traceback, logging, subprocess, os
-from autotest.client.shared import error, pxssh
+from autotest.client.shared import error, ssh_key
 from autotest.server import utils
 from autotest.server.hosts import abstract_ssh
 
@@ -205,43 +205,13 @@ def run_grep(self, command, timeout=30, ignore_status=False,
             raise error.AutoservRunError("command execution error", result)
 
 
-    def setup_ssh_key(self):
-        logging.debug('Performing SSH key setup on %s:%d as %s.' %
-                      (self.hostname, self.port, self.user))
-
-        try:
-            host = pxssh.pxssh()
-            host.login(self.hostname, self.user, self.password,
-                        port=self.port)
-            public_key = utils.get_public_key()
-
-            host.sendline('mkdir -p ~/.ssh')
-            host.prompt()
-            host.sendline('chmod 700 ~/.ssh')
-            host.prompt()
-            host.sendline("echo '%s' >> ~/.ssh/authorized_keys; " %
-                            public_key)
-            host.prompt()
-            host.sendline('chmod 600 ~/.ssh/authorized_keys')
-            host.prompt()
-            host.logout()
-
-            logging.debug('SSH key setup complete.')
-
-        except:
-            logging.debug('SSH key setup has failed.')
-            try:
-                host.logout()
-            except:
-                pass
-
-
     def setup_ssh(self):
         if self.password:
             try:
                 self.ssh_ping()
             except error.AutoservSshPingHostError:
-                self.setup_ssh_key()
+                ssh_key.setup_ssh_key(self.hostname, self.user, self.password,
+                                   port=self.port)
 
 
 class AsyncSSHMixin(object):