Merge pull request #257 from pyiron/two_factor

Add support for submitting to systems behind two factor authentication

pyproject.toml

@@ -46,6 +46,7 @@ executor = [
     "pympipool==0.7.9",
     "cloudpickle==3.0.0",
 ]
+twofactor = ["pyauthenticator==0.2.0"]
 
 [project.scripts]
 pysqa = "pysqa.cmd:command_line"

pysqa/ext/remote.py

@@ -23,7 +23,26 @@ def __init__(self, config, directory="~/.queues", execute_command=execute_comman
         self._ssh_known_hosts = os.path.abspath(
             os.path.expanduser(config["known_hosts"])
         )
-        self._ssh_key = os.path.abspath(os.path.expanduser(config["ssh_key"]))
+        if "ssh_key" in config.keys():
+            self._ssh_key = os.path.abspath(os.path.expanduser(config["ssh_key"]))
+        else:
+            self._ssh_key = None
+        if "ssh_password" in config.keys():
+            self._ssh_password = config["ssh_password"]
+        else:
+            self._ssh_password = None
+        if "ssh_key_passphrase" in config.keys():
+            self._ssh_key_passphrase = config["ssh_key_passphrase"]
+        else:
+            self._ssh_key_passphrase = None
+        if "ssh_authenticator_service" in config.keys():
+            self._ssh_authenticator_service = config["ssh_authenticator_service"]
+        else:
+            self._ssh_authenticator_service = None
+        if "ssh_proxy_host" in config.keys():
+            self._ssh_proxy_host = config["ssh_proxy_host"]
+        else:
+            self._ssh_proxy_host = None
         self._ssh_remote_config_dir = config["ssh_remote_config_dir"]
         self._ssh_remote_path = config["ssh_remote_path"]
         self._ssh_local_path = os.path.abspath(
@@ -42,6 +61,7 @@ def __init__(self, config, directory="~/.queues", execute_command=execute_comman
         else:
             self._ssh_continous_connection = False
         self._ssh_connection = None
+        self._ssh_proxy_connection = None
         self._remote_flag = True
 
     def convert_path_to_remote(self, path):
@@ -158,6 +178,8 @@ def transfer_file(self, file, transfer_back=False):
     def __del__(self):
         if self._ssh_connection is not None:
             self._ssh_connection.close()
+        if self._ssh_proxy_connection is not None:
+            self._ssh_proxy_connection.close()
 
     def _check_ssh_connection(self):
         if self._ssh_connection is None:
@@ -191,13 +213,74 @@ def _transfer_files(self, file_dict, sftp=None, transfer_back=False):
     def _open_ssh_connection(self):
         ssh = paramiko.SSHClient()
         ssh.load_host_keys(self._ssh_known_hosts)
-        ssh.connect(
-            hostname=self._ssh_host,
-            port=self._ssh_port,
-            username=self._ssh_username,
-            key_filename=self._ssh_key,
-        )
-        return ssh
+        if self._ssh_key is not None and self._ssh_key_passphrase is not None:
+            ssh.connect(
+                hostname=self._ssh_host,
+                port=self._ssh_port,
+                username=self._ssh_username,
+                key_filename=self._ssh_key,
+                passphrase=self._ssh_key_passphrase,
+            )
+        elif self._ssh_key is not None:
+            ssh.connect(
+                hostname=self._ssh_host,
+                port=self._ssh_port,
+                username=self._ssh_username,
+                key_filename=self._ssh_key,
+            )
+        elif self._ssh_password is not None and self._ssh_authenticator_service is None:
+            ssh.connect(
+                hostname=self._ssh_host,
+                port=self._ssh_port,
+                username=self._ssh_username,
+                password=self._ssh_password,
+            )
+        elif (
+            self._ssh_password is not None
+            and self._ssh_authenticator_service is not None
+        ):
+
+            def authentication(title, instructions, prompt_list):
+                from pyauthenticator import get_two_factor_code
+
+                if len(prompt_list) > 0:
+                    return [
+                        get_two_factor_code(service=self._ssh_authenticator_service)
+                    ]
+                else:
+                    return []
+
+            ssh.connect(
+                hostname=self._ssh_host,
+                port=self._ssh_port,
+                username=self._ssh_username,
+                password=self._ssh_password,
+            )
+
+            ssh._transport.auth_interactive(
+                username=self._ssh_username, handler=authentication, submethods=""
+            )
+        else:
+            raise ValueError("Un-supported authentication method.")
+
+        if self._ssh_proxy_host is not None:
+            client_new = paramiko.SSHClient()
+            client_new.set_missing_host_key_policy(paramiko.AutoAddPolicy())
+            vmtransport = ssh.get_transport()
+            vmchannel = vmtransport.open_channel(
+                kind="direct-tcpip",
+                dest_addr=(self._ssh_proxy_host, self._ssh_port),
+                src_addr=(self._ssh_host, self._ssh_port),
+            )
+            client_new.connect(
+                hostname=self._ssh_proxy_host,
+                username=self._ssh_username,
+                sock=vmchannel,
+            )
+            self._ssh_proxy_connection = ssh
+            return client_new
+        else:
+            return ssh
 
     def _remote_command(self):
         return "python -m pysqa --config_directory " + self._ssh_remote_config_dir + " "