operations.server.reboot: survive dead SSH session during askpass cleanup (#1665)

Author: wowi42

src/pyinfra/connectors/util.py

@@ -233,14 +233,24 @@ def write_stdin(stdin, buffer):
 
 
 def remove_any_sudo_askpass_file(host) -> None:
+    # Best-effort cleanup: this is called from host.disconnect(), and the
+    # connection may already be broken (e.g. after `server.reboot`). Swallow
+    # any errors from the remote ``rm`` and still clear the local state so a
+    # reconnect will regenerate a fresh askpass file.
     sudo_askpass_path = host.connector_data.get("sudo_askpass_path")
     if sudo_askpass_path:
-        host.run_shell_command(StringCommand("rm", "-f", QuoteString(sudo_askpass_path)))
+        try:
+            host.run_shell_command(StringCommand("rm", "-f", QuoteString(sudo_askpass_path)))
+        except Exception as e:
+            logger.debug("Could not remove sudo askpass file %s: %s", sudo_askpass_path, e)
         host.connector_data["sudo_askpass_path"] = None
 
     su_askpass_path = host.connector_data.get("su_askpass_path")
     if su_askpass_path:
-        host.run_shell_command(StringCommand("rm", "-f", QuoteString(su_askpass_path)))
+        try:
+            host.run_shell_command(StringCommand("rm", "-f", QuoteString(su_askpass_path)))
+        except Exception as e:
+            logger.debug("Could not remove su askpass file %s: %s", su_askpass_path, e)
         host.connector_data["su_askpass_path"] = None
 
 

src/pyinfra/operations/server.py

@@ -91,6 +91,13 @@ def wait_and_reconnect(state, host):  # pragma: no cover
         sleep(delay)
         max_retries = round(reboot_timeout / interval)
 
+        # The remote askpass files (if any) live on a host that has just
+        # rebooted — the SSH session is dead and there is nothing to clean up.
+        # Clear the stored paths before disconnecting so the disconnect path
+        # does not attempt an ``rm -f`` over the broken connection.
+        host.connector_data["sudo_askpass_path"] = None
+        host.connector_data["su_askpass_path"] = None
+
         host.disconnect()  # make sure we are properly disconnected
         retries = 0
 

tests/test_connectors/test_util.py

@@ -224,3 +224,59 @@ def test_command_exists_su_config_only(self):
         host = state.inventory.get_host("somehost")
         command = make_unix_command_for_host(state, host, "echo Šablony")
         assert command.get_raw_value() == "sh -c 'echo Šablony'"
+
+
+class TestRemoveAnySudoAskpassFile(TestCase):
+    def test_clears_state_and_runs_rm(self):
+        commands = []
+        host = MagicMock()
+        host.connector_data = {
+            "sudo_askpass_path": "/tmp/sudo-askpass",
+            "su_askpass_path": "/tmp/su-askpass",
+        }
+        host.run_shell_command = lambda cmd: commands.append(cmd.get_raw_value())
+
+        remove_any_sudo_askpass_file(host)
+
+        assert commands == [
+            "rm -f /tmp/sudo-askpass",
+            "rm -f /tmp/su-askpass",
+        ]
+        assert host.connector_data["sudo_askpass_path"] is None
+        assert host.connector_data["su_askpass_path"] is None
+
+    def test_swallows_errors_and_clears_state(self):
+        """
+        Regression test for #1645: `host.disconnect()` → `remove_any_sudo_askpass_file`
+        is called after ``server.reboot`` when the SSH session is already dead.
+        The remote ``rm`` must fail gracefully and the stored path must still
+        be cleared so a reconnect will regenerate a fresh askpass file.
+        """
+        host = MagicMock()
+        host.connector_data = {
+            "sudo_askpass_path": "/tmp/sudo-askpass",
+            "su_askpass_path": "/tmp/su-askpass",
+        }
+
+        def failing_run(cmd):
+            raise Exception("SSH session not active")
+
+        host.run_shell_command = failing_run
+
+        # Must not raise — this is a best-effort cleanup.
+        remove_any_sudo_askpass_file(host)
+
+        assert host.connector_data["sudo_askpass_path"] is None
+        assert host.connector_data["su_askpass_path"] is None
+
+    def test_noop_when_no_state(self):
+        host = MagicMock()
+        host.connector_data = {
+            "sudo_askpass_path": None,
+            "su_askpass_path": None,
+        }
+        host.run_shell_command = MagicMock()
+
+        remove_any_sudo_askpass_file(host)
+
+        host.run_shell_command.assert_not_called()