diff --git a/pyproject.toml b/pyproject.toml index e663fda9..a0f591ac 100644 --- a/pyproject.toml +++ b/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" diff --git a/pysqa/ext/remote.py b/pysqa/ext/remote.py index d03d0dfe..18d08d28 100644 --- a/pysqa/ext/remote.py +++ b/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 + " "