1 files changed, 82 insertions, 0 deletions
diff --git a/benchmarks/harness/tunnel_dispatcher.py b/benchmarks/harness/tunnel_dispatcher.py
new file mode 100644
index 000000000..8dfe2862a
--- /dev/null
+++ b/benchmarks/harness/tunnel_dispatcher.py
@@ -0,0 +1,82 @@
+# 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.
+"""Tunnel handles setting up connections to remote machines."""
+
+import os
+import tempfile
+import time
+
+import docker
+import pexpect
+
+SSH_TUNNEL_COMMAND = """ssh
+ -o GlobalKnownHostsFile=/dev/null
+ -o UserKnownHostsFile=/dev/null
+ -o StrictHostKeyChecking=no
+ -nNT -L {filename}:/var/run/docker.sock
+ -i {key_path}
+ {username}@{hostname}"""
+
+
+class Tunnel:
+  """The tunnel object represents the tunnel via ssh.
+
+  This connects a local unix domain socket with a remote socket.
+  """
+
+  def __init__(self, name, hostname: str, username: str, key_path: str,
+               **kwargs):
+    self._filename = tempfile.NamedTemporaryFile(prefix=name).name
+    self._hostname = hostname
+    self._username = username
+    self._key_path = key_path
+    self._kwargs = kwargs
+    self._process = None
+
+  def connect(self):
+    """Connects the SSH tunnel."""
+    cmd = SSH_TUNNEL_COMMAND.format(
+        filename=self._filename,
+        key_path=self._key_path,
+        username=self._username,
+        hostname=self._hostname)
+    self._process = pexpect.spawn(cmd, timeout=10)
+
+    # If given a password, assume we'll be asked for it.
+    if "key_password" in self._kwargs:
+      self._process.expect(["Enter passphrase for key .*: "])
+      self._process.sendline(self._kwargs["key_password"])
+
+    while True:
+      # Wait for the tunnel to appear.
+      if self._process.exitstatus is not None:
+        raise ConnectionError("Error in setting up ssh tunnel")
+      if os.path.exists(self._filename):
+        return
+      time.sleep(0.1)
+
+  def path(self):
+    """Return the socket file."""
+    return self._filename
+
+  def get_docker_client(self):
+    """Returns a docker client for this Tunne0l."""
+    return docker.DockerClient(base_url="unix:/" + self._filename)
+
+  def __del__(self):
+    """Closes the ssh connection process and deletes the socket file."""
+    if self._process:
+      self._process.close()
+    if os.path.exists(self._filename):
+      os.remove(self._filename)