diff --git a/dpdispatcher/contexts/ssh_context.py b/dpdispatcher/contexts/ssh_context.py index 8537894d..ac69bf39 100644 --- a/dpdispatcher/contexts/ssh_context.py +++ b/dpdispatcher/contexts/ssh_context.py @@ -45,6 +45,7 @@ def __init__( tar_compress=True, look_for_keys=True, execute_command=None, + proxy_command=None, ): self.hostname = hostname self.username = username @@ -59,6 +60,7 @@ def __init__( self.look_for_keys = look_for_keys self.execute_command = execute_command self._keyboard_interactive_auth = False + self.proxy_command = proxy_command self._setup_ssh() # @classmethod @@ -141,9 +143,13 @@ def _setup_ssh(self): # ) # assert(self.ssh.get_transport().is_active()) # transport = self.ssh.get_transport() - sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM) - sock.settimeout(self.timeout) - sock.connect((self.hostname, self.port)) + if self.proxy_command is not None: + sock = paramiko.ProxyCommand(self.proxy_command) + sock.settimeout(self.timeout) + else: + sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM) + sock.settimeout(self.timeout) + sock.connect((self.hostname, self.port)) # Make a Paramiko Transport object using the socket ts = paramiko.Transport(sock) @@ -343,6 +349,7 @@ def arginfo(): "enable searching for discoverable private key files in ~/.ssh/" ) doc_execute_command = "execute command after ssh connection is established." + doc_proxy_command = "The command of proxy." ssh_remote_profile_args = [ Argument("hostname", str, optional=False, doc=doc_hostname), Argument("username", str, optional=False, doc=doc_username), @@ -391,6 +398,13 @@ def arginfo(): default=None, doc=doc_execute_command, ), + Argument( + "proxy_command", + str, + optional=True, + default=None, + doc=doc_proxy_command, + ), ] ssh_remote_profile_format = Argument( "ssh_session", dict, ssh_remote_profile_args