1 files changed, 108 insertions, 0 deletions
diff --git a/benchmarks/harness/ssh_connection.py b/benchmarks/harness/ssh_connection.py
new file mode 100644
index 000000000..e0bf258f1
--- /dev/null
+++ b/benchmarks/harness/ssh_connection.py
@@ -0,0 +1,108 @@
+# python3
+# Copyright 2019 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# 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.
+"""SSHConnection handles the details of SSH connections."""
+
+import os
+import warnings
+
+import paramiko
+
+from benchmarks import harness
+
+# Get rid of paramiko Cryptography Warnings.
+warnings.filterwarnings(action="ignore", module=".*paramiko.*")
+
+
+def send_one_file(client: paramiko.SSHClient, path: str, remote_dir: str):
+  """Sends a single file via an SSH client.
+
+  Args:
+    client: The existing SSH client.
+    path: The local path.
+    remote_dir: The remote directory.
+  """
+  filename = path.split("/").pop()
+  client.exec_command("mkdir -p " + remote_dir)
+  with client.open_sftp() as ftp_client:
+    ftp_client.put(path, os.path.join(remote_dir, filename))
+
+
+class SSHConnection:
+  """SSH connection to a remote machine."""
+
+  def __init__(self, name: str, hostname: str, key_path: str, username: str,
+               **kwargs):
+    """Sets up a paramiko ssh connection to the given hostname."""
+    self._name = name  # Unused.
+    self._hostname = hostname
+    self._username = username
+    self._key_path = key_path  # RSA Key path
+    self._kwargs = kwargs
+    # SSHConnection wraps paramiko. paramiko supports RSA, ECDSA, and Ed25519
+    # keys, and we've chosen to only suport and require RSA keys. paramiko
+    # supports RSA keys that begin with '----BEGIN RSAKEY----'.
+    # https://stackoverflow.com/questions/53600581/ssh-key-generated-by-ssh-keygen-is-not-recognized-by-paramiko
+    self.rsa_key = self._rsa()
+    self.run("true")  # Validate.
+
+  def _client(self) -> paramiko.SSHClient:
+    """Returns a connected SSH client."""
+    client = paramiko.SSHClient()
+    client.set_missing_host_key_policy(paramiko.AutoAddPolicy())
+    client.connect(
+        hostname=self._hostname,
+        port=22,
+        username=self._username,
+        pkey=self.rsa_key,
+        allow_agent=False,
+        look_for_keys=False)
+    return client
+
+  def _rsa(self):
+    if "key_password" in self._kwargs:
+      password = self._kwargs["key_password"]
+    else:
+      password = None
+    rsa = paramiko.RSAKey.from_private_key_file(self._key_path, password)
+    return rsa
+
+  def run(self, cmd: str) -> (str, str):
+    """Runs a command via ssh.
+
+    Args:
+      cmd: The shell command to run.
+
+    Returns:
+      The contents of stdout and stderr.
+    """
+    with self._client() as client:
+      _, stdout, stderr = client.exec_command(command=cmd)
+      stdout.channel.recv_exit_status()
+      stdout = stdout.read().decode("utf-8")
+      stderr = stderr.read().decode("utf-8")
+    return stdout, stderr
+
+  def send_workload(self, name: str) -> str:
+    """Sends a workload tarball to the remote machine.
+
+    Args:
+      name: The workload name.
+
+    Returns:
+      The remote path.
+    """
+    with self._client() as client:
+      send_one_file(client, harness.LOCAL_WORKLOADS_PATH.format(name),
+                    harness.REMOTE_WORKLOADS_PATH.format(name))
+    return harness.REMOTE_WORKLOADS_PATH.format(name)