From e45e720e86d61300596294cd06ebe70bb08fabe4 Mon Sep 17 00:00:00 2001
From: zjgemi <liuxin_zijian@163.com>
Date: Wed, 25 Sep 2024 21:56:46 +0800
Subject: [PATCH] add proxy_command
Signed-off-by: zjgemi <liuxin_zijian@163.com>
---
dpdispatcher/contexts/ssh_context.py | 20 +++++++++++++++++---
1 file changed, 17 insertions(+), 3 deletions(-)
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