From a148729bf87b0ab62af9392137cd45b77e78af80 Mon Sep 17 00:00:00 2001 From: whitequark 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