Yardstick Plugin: add support for ssh login using key

JIRA: YARDSTICK-442

Change-Id: I4736e8cb8331d7a74c8c9946e21edd791b0c8ba9
Signed-off-by: JingLu5 <lvjing5@huawei.com>

2 files changed, 52 insertions, 25 deletions

diff --git a/plugin/storperf.yaml b/plugin/storperf.yaml
index d08e26eb6..074a82067 100644
--- a/plugin/storperf.yaml
+++ b/plugin/storperf.yaml
@@ -10,4 +10,5 @@ plugins:
 deployment:
   ip: 192.168.23.2
   user: root
+  # Remove 'password' if log into deployment location using key file
   password: root
diff --git a/yardstick/cmd/commands/plugin.py b/yardstick/cmd/commands/plugin.py
index 10e5cdfbe..cf6612529 100644
--- a/yardstick/cmd/commands/plugin.py
+++ b/yardstick/cmd/commands/plugin.py
@@ -85,24 +85,31 @@ class PluginCommands(object):
 
         deployment_user = deployment.get("user")
         deployment_ssh_port = deployment.get("ssh_port", ssh.DEFAULT_PORT)
-        deployment_ip = deployment.get("ip")
-        deployment_password = deployment.get("password")
+        deployment_ip = deployment.get("ip", None)
+        deployment_password = deployment.get("password", None)
+        deployment_key_filename = deployment.get("key_filename",
+                                                 "/root/.ssh/id_rsa")
 
         if deployment_ip == "local":
             installer_ip = os.environ.get("INSTALLER_IP", None)
 
-            LOG.info("user:%s, host:%s", deployment_user, installer_ip)
-            self.client = ssh.SSH(deployment_user, installer_ip,
-                                  password=deployment_password,
-                                  port=deployment_ssh_port)
-            self.client.wait(timeout=600)
+            if deployment_password is not None:
+                self._login_via_password(deployment_user, installer_ip,
+                                         deployment_password,
+                                         deployment_ssh_port)
+            else:
+                self._login_via_key(self, deployment_user, installer_ip,
+                                    deployment_key_filename,
+                                    deployment_ssh_port)
         else:
-            LOG.info("user:%s, host:%s", deployment_user, deployment_ip)
-            self.client = ssh.SSH(deployment_user, deployment_ip,
-                                  password=deployment_password,
-                                  port=deployment_ssh_port)
-            self.client.wait(timeout=600)
-
+            if deployment_password is not None:
+                self._login_via_password(deployment_user, deployment_ip,
+                                         deployment_password,
+                                         deployment_ssh_port)
+            else:
+                self._login_via_key(self, deployment_user, deployment_ip,
+                                    deployment_key_filename,
+                                    deployment_ssh_port)
         # copy script to host
         cmd = "cat > ~/%s.sh" % plugin_name
 
@@ -117,23 +124,31 @@ class PluginCommands(object):
 
         deployment_user = deployment.get("user")
         deployment_ssh_port = deployment.get("ssh_port", ssh.DEFAULT_PORT)
-        deployment_ip = deployment.get("ip")
-        deployment_password = deployment.get("password")
+        deployment_ip = deployment.get("ip", None)
+        deployment_password = deployment.get("password", None)
+        deployment_key_filename = deployment.get("key_filename",
+                                                 "/root/.ssh/id_rsa")
 
         if deployment_ip == "local":
             installer_ip = os.environ.get("INSTALLER_IP", None)
 
-            LOG.info("user:%s, host:%s", deployment_user, installer_ip)
-            self.client = ssh.SSH(deployment_user, installer_ip,
-                                  password=deployment_password,
-                                  port=deployment_ssh_port)
-            self.client.wait(timeout=600)
+            if deployment_password is not None:
+                self._login_via_password(deployment_user, installer_ip,
+                                         deployment_password,
+                                         deployment_ssh_port)
+            else:
+                self._login_via_key(self, deployment_user, installer_ip,
+                                    deployment_key_filename,
+                                    deployment_ssh_port)
         else:
-            LOG.info("user:%s, host:%s", deployment_user, deployment_ip)
-            self.client = ssh.SSH(deployment_user, deployment_ip,
-                                  password=deployment_password,
-                                  port=deployment_ssh_port)
-            self.client.wait(timeout=600)
+            if deployment_password is not None:
+                self._login_via_password(deployment_user, deployment_ip,
+                                         deployment_password,
+                                         deployment_ssh_port)
+            else:
+                self._login_via_key(self, deployment_user, deployment_ip,
+                                    deployment_key_filename,
+                                    deployment_ssh_port)
 
         # copy script to host
         cmd = "cat > ~/%s.sh" % plugin_name
@@ -141,6 +156,17 @@ class PluginCommands(object):
         LOG.info("copying script to host: %s", cmd)
         self.client.run(cmd, stdin=open(self.script, 'rb'))
 
+    def _login_via_password(self, user, ip, password, ssh_port):
+        LOG.info("Log in via pw, user:%s, host:%s", user, ip)
+        self.client = ssh.SSH(user, ip, password=password, port=ssh_port)
+        self.client.wait(timeout=600)
+
+    def _login_via_key(self, user, ip, key_filename, ssh_port):
+        LOG.info("Log in via key, user:%s, host:%s", user, ip)
+        self.client = ssh.SSH(user, ip, key_filename=key_filename,
+                              port=ssh_port)
+        self.client.wait(timeout=600)
+
     def _run(self, plugin_name):
         '''Run installation script '''
         cmd = "sudo bash %s" % plugin_name + ".sh"