fix two small bugs related to ssh key auth

Signed-off-by: ashiven <nevisha@pm.me>

Author: ashiven

embark/workers/tasks.py

@@ -724,17 +724,22 @@ def _update_or_create_worker(config: Configuration, ip_address: str):
         )
         worker.save()
         worker.configurations.set([config])
-    finally:
         try:
             init_sudoers_file(config, worker)
             setup_ssh_key(config, worker)
+        except BaseException:
+            logger.error("Failed to setup SSH key or sudoers file for worker %s", worker.name)
+            worker.delete()
+            return
+    finally:
+        try:
             update_system_info(worker)
             update_dependencies_info(worker)
         except BaseException:
             pass
 
 
-def _scan_for_worker(config: Configuration, ip_address: str, port: int = 22, timeout: int = 1) -> str:
+def _scan_for_worker(config: Configuration, ip_address: str, port: int = 22, timeout: int = 1, ssh_auth_check: bool = True) -> str:
     """
     Checks if a host is reachable, has valid SSH creds, and sudo perms.
     Updates the DB accordingly.
@@ -743,6 +748,7 @@ def _scan_for_worker(config: Configuration, ip_address: str, port: int = 22, tim
     :param ip_address: Host's IP address
     :param port: SSH port (default: 22)
     :param timeout: Connection timeout in seconds
+    :param ssh_auth_check: If True, tries to connect to the ip address with the config's ssh credentials and checks whether the user has sudo privileges
     :return: ip_address on success, None otherwise
     """
     try:  # Check TCP socket first before trying SSH
@@ -756,23 +762,26 @@ def _scan_for_worker(config: Configuration, ip_address: str, port: int = 22, tim
 
     client = None
     try:
-        client = new_autoadd_client()
-        client.connect(
-            ip_address, port=port,
-            username=config.ssh_user, password=config.ssh_password,
-            timeout=timeout
-        )
-
-        logger.info("[%s@%s] Connected via SSH. Now testing for sudo privileges.", config.ssh_user, ip_address)
-
-        stdin, stdout, _ = client.exec_command("sudo -v", get_pty=True)  # nosec
-        stdin.write(config.ssh_password + "\n")
-        stdin.flush()
-        if stdout.channel.recv_exit_status():
-            logger.info("[%s@%s] Can't register worker: No sudo permission.", config.ssh_user, ip_address)
-            # Delete worker if SSH configuration changed
-            Worker.objects.filter(ip_address=ip_address).delete()
-            return None
+        # We only want to perform this check for newly created configs, not for scans of existing configs
+        # (ssh pw auth would have been disabled for workers belonging to existing configs)
+        if ssh_auth_check:
+            client = new_autoadd_client()
+            client.connect(
+                ip_address, port=port,
+                username=config.ssh_user, password=config.ssh_password,
+                timeout=timeout
+            )
+
+            logger.info("[%s@%s] Connected via SSH. Now testing for sudo privileges.", config.ssh_user, ip_address)
+
+            stdin, stdout, _ = client.exec_command("sudo -v", get_pty=True)  # nosec
+            stdin.write(config.ssh_password + "\n")
+            stdin.flush()
+            if stdout.channel.recv_exit_status():
+                logger.info("[%s@%s] Can't register worker: No sudo permission.", config.ssh_user, ip_address)
+                # Delete worker if SSH configuration changed
+                Worker.objects.filter(ip_address=ip_address).delete()
+                return None
 
         _update_or_create_worker(config, ip_address)
         logger.info("[%s@%s] Worker is reachable via SSH.", config.ssh_user, ip_address)
@@ -795,13 +804,14 @@ def config_worker_scan_task(configuration_id: int):
     """
     try:
         config = Configuration.objects.get(id=configuration_id)
+        ssh_auth_check = not config.scan_status == Configuration.ScanStatus.FINISHED
         config.scan_status = Configuration.ScanStatus.SCANNING
         config.save()
 
         ip_network = ipaddress.ip_network(config.ip_range, strict=False)
         ip_addresses = [str(ip) for ip in ip_network.hosts()]
         with ThreadPoolExecutor(max_workers=50) as executor:
-            results = executor.map(partial(_scan_for_worker, config), ip_addresses)
+            results = executor.map(partial(_scan_for_worker, config, ssh_auth_check=ssh_auth_check), ip_addresses)
             reachable = set(results) - {None}
 
         unreachable_workers = [worker for worker in config.workers.all() if worker.ip_address not in reachable]

embark/workers/update/update.py

@@ -263,12 +263,33 @@ def undo_sudoers_file(configuration: Configuration, worker: Worker):
             client.close()
 
 
+def _create_ssh_config_dir(configuration: Configuration, worker: Worker):
+    """
+    Connects to the worker and creates the .ssh directory if it does not exist.
+    If this is not done before calling ssh-copy-id, the command may fail.
+    :params configuration: The worker configuration that includes the ssh credentials for the worker
+    :params worker: The worker on which to create the SSH config directory
+    """
+    client = None
+    try:
+        client = worker.ssh_connect(use_password=True)
+        homedir = f"/home/{configuration.ssh_user}" if configuration.ssh_user != "root" else "/root"
+        exec_blocking_ssh(client, f"sudo mkdir -p {homedir}/.ssh && sudo chown {configuration.ssh_user}:{configuration.ssh_user} {homedir}/.ssh")
+        logger.info("_create_ssh_config_dir: SSH config directory created for user %s on worker %s", configuration.ssh_user, worker.ip_address)
+    except paramiko.SSHException as ssh_error:
+        logger.error("_create_ssh_config_dir: Failed to create SSH config directory on worker %s: %s", worker.ip_address, ssh_error)
+    finally:
+        if client:
+            client.close()
+
+
 def setup_ssh_key(configuration: Configuration, worker: Worker):
     """
     Install SSH key on provided worker and disable PW auth
     :params configuration: The SSH keys
     :params worker: The worker to update
     """
+    _create_ssh_config_dir(configuration, worker)
     _, public_key = configuration.ensure_ssh_keys()
 
     logger.info("setup_ssh_key: Starting ssh-copy-id on worker %s", worker.ip_address)