Refactor SSH connection to make adding passwords easier. (#602)

* Refactor SSH method.

* Tidy up.

* Update datashuttle/utils/rclone.py

Co-authored-by: Shrey Singh <[email protected]>

* Update datashuttle/utils/ssh.py

Co-authored-by: Shrey Singh <[email protected]>

* Fix docstring.

* Small tidy ups.

---------

Co-authored-by: Shrey Singh <[email protected]>

Author: JoeZiminski

datashuttle/configs/config_class.py

@@ -62,7 +62,6 @@ def __init__(
 
         self.logging_path: Path
         self.hostkeys_path: Path
-        self.ssh_key_path: Path
         self.project_metadata_path: Path
 
     def setup_after_load(self) -> None:
@@ -300,8 +299,6 @@ def init_paths(self) -> None:
             self.project_name
         )
 
-        self.ssh_key_path = datashuttle_path / f"{self.project_name}_ssh_key"
-
         self.hostkeys_path = datashuttle_path / "hostkeys"
 
         self.logging_path = self.make_and_get_logging_path()

datashuttle/datashuttle_class.py

@@ -28,7 +28,6 @@
         TopLevelFolder,
     )
 
-import paramiko
 import yaml
 
 from datashuttle.configs import (
@@ -820,46 +819,26 @@ def setup_ssh_connection(self) -> None:
             "setup-ssh-connection-to-central-server", local_vars=locals()
         )
 
-        verified = ssh.verify_ssh_central_host(
+        verified = ssh.verify_ssh_central_host_api(
             self.cfg["central_host_id"],
             self.cfg.hostkeys_path,
             log=True,
         )
 
         if verified:
-            ssh.setup_ssh_key(self.cfg, log=True)
-            self._setup_rclone_central_ssh_config(log=True)
+            private_key_str = ssh.setup_ssh_key_api(self.cfg, log=True)
+
+            self._setup_rclone_central_ssh_config(private_key_str, log=True)
 
             rclone.check_successful_connection_and_raise_error_on_fail(
                 self.cfg
             )
 
-        ds_logger.close_log_filehandler()
-
-    @requires_ssh_configs
-    @check_is_not_local_project
-    def write_public_key(self, filepath: str) -> None:
-        """Save the public SSH key to a specified filepath.
-
-        By default, only the SSH private key is stored in the
-        datashuttle configs folder. Use this function to save
-        the public key.
-
-        Parameters
-        ----------
-        filepath
-            Full filepath (including filename) to write the
-            public key to.
-
-        """
-        key: paramiko.RSAKey
-        key = paramiko.RSAKey.from_private_key_file(
-            self.cfg.ssh_key_path.as_posix()
-        )
+            utils.log_and_message(
+                "SSH key pair setup successfully. SSH key saved to the RClone config file."
+            )
 
-        with open(filepath, "w") as public:
-            public.write(key.get_base64())
-        public.close()
+        ds_logger.close_log_filehandler()
 
     # -------------------------------------------------------------------------
     # Google Drive
@@ -1588,11 +1567,13 @@ def _make_project_metadata_if_does_not_exist(self) -> None:
         """
         folders.create_folders(self.cfg.project_metadata_path, log=False)
 
-    def _setup_rclone_central_ssh_config(self, log: bool) -> None:
+    def _setup_rclone_central_ssh_config(
+        self, private_key_str: str, log: bool
+    ) -> None:
         rclone.setup_rclone_config_for_ssh(
             self.cfg,
             self.cfg.get_rclone_config_name("ssh"),
-            self.cfg.ssh_key_path,
+            private_key_str,
             log=log,
         )
 

datashuttle/tui/interface.py

@@ -511,10 +511,14 @@ def setup_key_pair_and_rclone_config(
     ) -> InterfaceOutput:
         """Set up SSH key pair and associated rclone configuration."""
         try:
+            rsa_key, private_key_str = ssh.generate_ssh_key_strings()
+
             ssh.add_public_key_to_central_authorized_keys(
-                self.project.cfg, password, log=False
+                self.project.cfg, rsa_key, password, log=False
+            )
+            self.project._setup_rclone_central_ssh_config(
+                private_key_str, log=False
             )
-            self.project._setup_rclone_central_ssh_config(log=False)
 
             rclone.check_successful_connection_and_raise_error_on_fail(
                 self.project.cfg

datashuttle/tui/screens/setup_ssh.py

@@ -152,10 +152,7 @@ def use_password_to_setup_ssh_key_pairs(self) -> None:
         )
 
         if success:
-            message = (
-                f"Connection successful! SSH key "
-                f"saved to {self.interface.get_configs().ssh_key_path}"
-            )
+            message = "Connection successful! SSH key saved to the RClone config file."
             self.query_one("#setup_ssh_ok_button").label = "Finish"
             self.query_one("#setup_ssh_cancel_button").disabled = True
             self.stage += 1

datashuttle/utils/rclone.py

@@ -3,8 +3,6 @@
 from typing import TYPE_CHECKING, Dict, List, Literal, Optional
 
 if TYPE_CHECKING:
-    from pathlib import Path
-
     from datashuttle.configs.config_class import Configs
     from datashuttle.utils.custom_types import TopLevelFolder
 
@@ -176,7 +174,7 @@ def setup_rclone_config_for_local_filesystem(
 def setup_rclone_config_for_ssh(
     cfg: Configs,
     rclone_config_name: str,
-    ssh_key_path: Path,
+    private_key_str: str,
     log: bool = True,
 ) -> None:
     """Set the RClone remote config for ssh.
@@ -194,25 +192,27 @@ def setup_rclone_config_for_ssh(
         canonical config name, generated by
         datashuttle.cfg.get_rclone_config_name()
 
-    ssh_key_path
-        path to the ssh key used for connecting to
-        ssh central filesystem
+    private_key_str
+        PEM encoded ssh private key to pass to RClone.
 
     log
         whether to log, if True logger must already be initialised.
 
     """
-    call_rclone(
+    key_escaped = private_key_str.replace("\n", "\\n")
+
+    command = (
         f"config create "
         f"{rclone_config_name} "
         f"sftp "
         f"host {cfg['central_host_id']} "
         f"user {cfg['central_host_username']} "
         f"port {canonical_configs.get_default_ssh_port()} "
-        f"key_file {ssh_key_path.as_posix()}",
-        pipe_std=True,
+        f'-- key_pem "{key_escaped}"'
     )
 
+    call_rclone(command, pipe_std=True)
+
     if log:
         log_rclone_config_output()
 
@@ -274,7 +274,7 @@ def setup_rclone_config_for_gdrive(
         f"{client_secret_key_value}"
         f"scope drive "
         f"root_folder_id {cfg['gdrive_root_folder_id']} "
-        f"{extra_args}",
+        f"{extra_args}"
     )
 
     return process