add look_for_keys option (#299)

Fix #296.
deepmodeling · Jan 5, 2023 · e4ac4f5 · e4ac4f5
1 parent a812090
commit e4ac4f5
Showing 1 changed file with 6 additions and 2 deletions.
diff --git a/dpdispatcher/ssh_context.py b/dpdispatcher/ssh_context.py
@@ -32,7 +32,8 @@ def __init__(self,
                 passphrase=None,
                 timeout=10,
                 totp_secret=None,
-                tar_compress=True
+                tar_compress=True,
+                look_for_keys=True,
                 ):
 
         self.hostname = hostname
@@ -45,6 +46,7 @@ def __init__(self,
         self.totp_secret = totp_secret
         self.ssh = None
         self.tar_compress = tar_compress
+        self.look_for_keys = look_for_keys
         self._setup_ssh()
 
     # @classmethod
@@ -160,7 +162,7 @@ def _setup_ssh(self):
                         pass
                     if key is not None:
                         break
-        else:
+        elif self.look_for_keys:
             for keytype, name in [
                 (paramiko.RSAKey, "rsa"),
                 (paramiko.DSSKey, "dsa"),
@@ -298,6 +300,7 @@ def arginfo():
         doc_totp_secret = 'Time-based one time password secret. It should be a base32-encoded string' \
                           ' extracted from the 2D code.'
         doc_tar_compress = 'The archive will be compressed in upload and download if it is True. If not, compression will be skipped.'
+        doc_look_for_keys = "enable searching for discoverable private key files in ~/.ssh/"
 
         ssh_remote_profile_args = [
             Argument("hostname", str, optional=False, doc=doc_hostname),
@@ -309,6 +312,7 @@ def arginfo():
             Argument("timeout", int, optional=True, default=10, doc=doc_timeout),
             Argument("totp_secret", str, optional=True, default=None, doc=doc_totp_secret),
             Argument("tar_compress", bool, optional=True, default=True, doc = doc_tar_compress),
+            Argument("look_for_keys", bool, optional=True, default=True, doc=doc_look_for_keys),
         ]
         ssh_remote_profile_format = Argument("ssh_session", dict, ssh_remote_profile_args)
         return ssh_remote_profile_format