Fix private key handling logic (#690)

This PR removes the feature allowing users to upload a private key in the process of SSH setup due to security concerns.

Author: edan-bainglass

aiidalab_widgets_base/computational_resources.py

@@ -12,7 +12,6 @@
 import ipywidgets as ipw
 import jinja2
 import pexpect
-import shortuuid
 import traitlets as tl
 from aiida import common, orm, plugins
 from aiida.orm.utils.builders.computer import ComputerBuilder
@@ -337,17 +336,10 @@ def __init__(self, ssh_folder: Path | None = None, **kwargs):
             style=STYLE,
         )
 
-        self._inp_private_key = ipw.FileUpload(
-            accept="",
-            layout=LAYOUT,
-            description="Private key",
-            multiple=False,
-        )
         self._verification_mode = ipw.Dropdown(
             options=[
                 ("Provide password to remote machine", "password"),
                 ("Manually deploy a public key", "public_key"),
-                ("Upload a custom private key", "private_key"),
                 ("Use multiple factor authentication", "mfa"),
             ],
             layout=LAYOUT,
@@ -379,7 +371,6 @@ def __init__(self, ssh_folder: Path | None = None, **kwargs):
                     <ul>
                         <li><b>Provide password to remote machine</b>: We use your password to deploy a generated public key to the remote machine (<b>recommended</b>)</li>
                         <li><b>Manually deploy a public key</b>: Manually copy a generated public key over to the remote machine</li>
-                        <li><b>Upload a custom private key</b>: Upload a custom private key file</li>
                         <li><b>Use multiple factor authentication</b>: To be used in tandem with the MFA app (beta)</li>
                     </ul>
                 """,
@@ -454,7 +445,7 @@ def _is_in_config(self):
             return False
         return True
 
-    def _write_ssh_config(self, private_key_abs_fname=None):
+    def _write_ssh_config(self):
         """Put host information into the config file."""
         config_path = self._ssh_folder / "config"
 
@@ -474,8 +465,6 @@ def _write_ssh_config(self, private_key_abs_fname=None):
                 file.write(
                     f"  ProxyCommand {self.proxy_command.value.format(username=self.username.value)}\n"
                 )
-            if private_key_abs_fname:
-                file.write(f"  IdentityFile {private_key_abs_fname}\n")
             file.write("  ServerAliveInterval 5\n")
 
     def key_pair_prepare(self):
@@ -511,39 +500,18 @@ def _on_setup_ssh_button_pressed(self, _=None):
 
             self.thread_ssh_copy_id()
 
-        # For not password ssh auth (such as using private_key or 2FA), key pair is not needed (2FA)
+        # For not password ssh auth (e.g., 2FA), key pair is not needed (2FA)
         # or the key pair is ready.
         # There are other mechanism to set up the ssh connection.
         # But we still need to write the ssh config to the ssh config file for such as
         # proxy jump.
 
-        private_key_fname = None
-        if self._verification_mode.value == "private_key":
-            # Write private key in ~/.ssh/ and use the name of upload file,
-            # if exist, generate random string and append to filename then override current name.
-
-            # unwrap private key file and setting temporary private_key content
-            private_key_fname, private_key_content = self._private_key
-            if private_key_fname is None:  # check private key file
-                message = "Please upload your private key file."
-                self.message = wrap_message(message, MessageLevel.ERROR)
-                return
-
-            filename = Path(private_key_fname).name
-
-            # if the private key filename is exist, generate random string and append to filename subfix
-            # then override current name.
-            if filename in [str(p.name) for p in Path(self._ssh_folder).iterdir()]:
-                private_key_fpath = self._ssh_folder / f"{filename}-{shortuuid.uuid()}"
-
-            self._add_private_key(private_key_fpath, private_key_content)
-
         # TODO(danielhollas): I am not sure this is correct. What if the user wants
-        # to overwrite the private key? Or any other config? The configuration would never be written.
+        # to overwrite a config entry? The configuration would never be written.
         # And the user is not notified that we did not write anything.
         # https://github.com/aiidalab/aiidalab-widgets-base/issues/516
         if not self._is_in_config():
-            self._write_ssh_config(private_key_abs_fname=private_key_fname)
+            self._write_ssh_config()
 
     def _ssh_copy_id(self):
         """Run the ssh-copy-id command and follow it until it is completed."""
@@ -647,20 +615,6 @@ def _on_verification_mode_change(self, change):
             clear_output()
             if self._verification_mode.value == "password":
                 self.password_box.layout.display = "block"
-            elif self._verification_mode.value == "private_key":
-                display(self._inp_private_key)
-                self.password_box.layout.display = "none"
-                display(
-                    ipw.HTML(
-                        value="""
-                        <div class="alert alert-warning">
-                            <strong>Warning!</strong> The use of private keys is generally <b>not recommended</b>.
-                            However, some HPC centers do use this method.
-                            Please make sure you know what you are doing.
-                        </div>
-                    """
-                    )
-                )
             elif self._verification_mode.value == "public_key":
                 self.password_box.layout.display = "none"
                 public_key = self._ssh_folder / "id_rsa.pub"
@@ -674,23 +628,6 @@ def _on_verification_mode_change(self, change):
             else:
                 self.password_box.layout.display = "none"
 
-    @property
-    def _private_key(self) -> tuple[str | None, bytes | None]:
-        """Unwrap private key file and setting filename and file content."""
-        if self._inp_private_key.value:
-            (fname, _value), *_ = self._inp_private_key.value.items()
-            content = copy.copy(_value["content"])
-            self._inp_private_key.value.clear()
-            self._inp_private_key._counter = 0  # pylint: disable=protected-access
-            return fname, content
-        return None, None
-
-    @staticmethod
-    def _add_private_key(private_key_fpath: Path, private_key_content: bytes):
-        """Write private key to the private key file in the ssh folder."""
-        private_key_fpath.write_bytes(private_key_content)
-        private_key_fpath.chmod(0o600)
-
     def _reset(self):
         self.hostname.value = ""
         self.port.value = 22

tests/test_computational_resources.py

@@ -53,29 +53,8 @@ def test_ssh_computer_setup_widget(monkeypatch, tmp_path):
     # Check that ssh-keygen operation is successful.
     widget._ssh_keygen()
 
-    # Create non-default private key file.
-    private_key_path = tmp_path / ".ssh" / "my_key_name"
-    widget._add_private_key(private_key_path, b"my_key_content")
-    assert private_key_path.exists()
-    with open(private_key_path) as f:
-        assert f.read() == "my_key_content"
-
-    # set private key with same name to trigger the rename operation
-    widget._verification_mode.value = "private_key"
-    # mock _private_key to mimic the upload of the private key
-    monkeypatch.setattr(
-        "aiidalab_widgets_base.computational_resources.SshComputerSetup._private_key",
-        property(lambda _: ("my_key_name", b"my_key_content_new")),
-    )
-    # check the private key is renamed, monkeypatch the shortuuid to make the test deterministic
-    monkeypatch.setattr("shortuuid.uuid", lambda: "00001111")
-
     widget._on_setup_ssh_button_pressed()
 
-    assert "my_key_name-00001111" in [
-        str(p.name) for p in Path(tmp_path / ".ssh").iterdir()
-    ]
-
     # Setting the ssh_config to an empty dictionary should reset the widget.
     widget.ssh_config = {}
     assert widget.hostname.value == ""