From a148729bf87b0ab62af9392137cd45b77e78af80 Mon Sep 17 00:00:00 2001
From: whitequark <whitequark@whitequark.org>
Date: Sat, 5 May 2018 02:50:01 +0000
Subject: [PATCH] artiq_flash, artiq_devtool: add support for ssh -J option.

---
 artiq/frontend/artiq_devtool.py |  9 +++++++--
 artiq/frontend/artiq_flash.py   |  5 ++++-
 artiq/remoting.py               | 14 ++++++++++++--
 3 files changed, 23 insertions(+), 5 deletions(-)

diff --git a/artiq/frontend/artiq_devtool.py b/artiq/frontend/artiq_devtool.py
index a17c29b26..3ab31e776 100755
--- a/artiq/frontend/artiq_devtool.py
+++ b/artiq/frontend/artiq_devtool.py
@@ -46,6 +46,9 @@ def get_argparser():
     parser.add_argument("-H", "--host",
                         type=str, default="lab.m-labs.hk",
                         help="SSH host where the development board is located")
+    parser.add_argument("-J", "--jump",
+                        type=str, default=None,
+                        help="SSH host to jump through")
     parser.add_argument("-b", "--board",
                         type=str, default="{board_type}-1",
                         help="board to connect to on the development SSH host")
@@ -57,7 +60,7 @@ def get_argparser():
                         type=str, default="/dev/ttyUSB_{board}",
                         help="TTY device corresponding to the development board")
     parser.add_argument("-d", "--device",
-                        type=str, default="{board}.{host}",
+                        type=str, default="{board}",
                         help="address or domain corresponding to the development board")
     parser.add_argument("-w", "--wait", action="store_true",
                         help="wait for the board to unlock instead of aborting the actions")
@@ -96,7 +99,7 @@ def main():
     device     = args.device.format(board=board, host=args.host)
     serial     = args.serial.format(board=board)
 
-    client = SSHClient(args.host)
+    client = SSHClient(args.host, args.jump)
 
     flock_acquired = False
     flock_file = None # GC root
@@ -158,6 +161,8 @@ def main():
         for _ in range(args.verbose):
             flash_args.append("-v")
         flash_args += ["-H", args.host]
+        if args.jump:
+            flash_args += ["-J", args.jump]
         flash_args += ["-t", board_type]
         flash_args += ["-V", variant]
         flash_args += ["-I", "source {}".format(board_file)]
diff --git a/artiq/frontend/artiq_flash.py b/artiq/frontend/artiq_flash.py
index 8023dac56..0ed2a5fde 100755
--- a/artiq/frontend/artiq_flash.py
+++ b/artiq/frontend/artiq_flash.py
@@ -48,6 +48,9 @@ Prerequisites:
     parser.add_argument("-H", "--host", metavar="HOSTNAME",
                         type=str, default=None,
                         help="SSH host where the development board is located")
+    parser.add_argument("-J", "--jump",
+                        type=str, default=None,
+                        help="SSH host to jump through")
     parser.add_argument("-t", "--target", default="kc705",
                         help="target board, default: %(default)s, one of: "
                              "kc705 kasli sayma")
@@ -299,7 +302,7 @@ def main():
     if args.host is None:
         client = LocalClient()
     else:
-        client = SSHClient(args.host)
+        client = SSHClient(args.host, args.jump)
 
     programmer = config["programmer"](client, preinit_script=args.preinit_command)
 
diff --git a/artiq/remoting.py b/artiq/remoting.py
index 8400b35e8..2c5f8af31 100644
--- a/artiq/remoting.py
+++ b/artiq/remoting.py
@@ -57,8 +57,9 @@ class LocalClient(Client):
 
 
 class SSHClient(Client):
-    def __init__(self, host):
+    def __init__(self, host, jump_host=None):
         self.host = host
+        self.jump_host = jump_host
         self.ssh = None
         self.sftp = None
         self._tmpr = "/tmp/artiq"
@@ -70,10 +71,19 @@ class SSHClient(Client):
         if self.ssh is None:
             import paramiko
             logging.getLogger("paramiko").setLevel(logging.WARNING)
+
+            if self.jump_host:
+                proxy_cmd = "ssh -W {}:22 {}".format(self.host, self.jump_host)
+                logger.debug("Using proxy command '{}'".format(proxy_cmd))
+                proxy = paramiko.proxy.ProxyCommand(proxy_cmd)
+            else:
+                proxy = None
+
             self.ssh = paramiko.SSHClient()
             self.ssh.load_system_host_keys()
             self.ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy())
-            self.ssh.connect(self.host)
+            self.ssh.connect(self.host, sock=proxy)
+
             logger.debug("Connecting to {}".format(self.host))
         return self.ssh