fix(connectors.util): prompt for sudo password up front with --use-sudo-password

src/pyinfra/api/config.py

@@ -47,6 +47,8 @@ class ConfigDefaults:
     PRESERVE_SUDO_ENV: bool = False
     USE_SUDO_LOGIN: bool = False
     SUDO_PASSWORD: str | None = None
+    # Prompt for the sudo password before the first sudo command (if SUDO_PASSWORD isn't set)
+    USE_SUDO_PASSWORD: bool = False
     # Use doas and optional user
     DOAS: bool = False
     DOAS_USER: str | None = None

src/pyinfra/connectors/util.py

@@ -364,6 +364,13 @@ def make_unix_command_for_host(
         if "_sudo_password" not in command_arguments or not command_arguments["_sudo_password"]:
             command_arguments["_sudo_password"] = host.connector_data.get("prompted_sudo_password")
 
+        # If we know a password is needed, ask for it up front rather than first trying
+        # passwordless sudo and prompting on failure.
+        if not command_arguments["_sudo_password"] and state.config.USE_SUDO_PASSWORD:
+            sudo_password = getpass(f"{host.print_prefix}sudo password: ")
+            host.connector_data["prompted_sudo_password"] = sudo_password
+            command_arguments["_sudo_password"] = sudo_password
+
         if command_arguments.get("_sudo_password"):
             command_arguments["_sudo_askpass_path"] = _ensure_sudo_askpass_set_for_host(
                 host, temp_dir=op_temp_dir

tests/test_connectors/test_util.py

@@ -457,3 +457,52 @@ def fake_run(command, *args, **kwargs):
             )
 
         assert "${TMPDIR:=/op/tmp}" in captured["command"]
+
+
+class TestUseSudoPassword(TestCase):
+    """
+    With ``config.USE_SUDO_PASSWORD`` (``--use-sudo-password``) set, pyinfra should prompt for the
+    sudo password before the first sudo command rather than first trying passwordless sudo
+    (issue #1850).
+    """
+
+    def _make_host(self, config):
+        state = State(make_inventory(hosts=("somehost",)), config)
+        host = state.inventory.get_host("somehost")
+        host.init(state)
+        host.run_shell_command = MagicMock(  # type: ignore[method-assign]
+            return_value=(True, CommandOutput([OutputLine("stdout", "/tmp/askpass-XYZ")])),
+        )
+        return host
+
+    @patch("pyinfra.connectors.util.getpass", return_value="supersecret")
+    def test_prompts_before_first_sudo_command(self, fake_getpass):
+        host = self._make_host(Config(USE_SUDO_PASSWORD=True))
+
+        command = make_unix_command_for_host(host.state, host, "uptime", _sudo=True)
+        assert command.get_raw_value() == (
+            "env SUDO_ASKPASS=/tmp/askpass-XYZ PYINFRA_SUDO_PASSWORD=supersecret "
+            "sudo -H -A -k sh -c uptime"
+        )
+
+        # The prompted password is reused for later commands
+        make_unix_command_for_host(host.state, host, "uptime", _sudo=True)
+        fake_getpass.assert_called_once()
+
+    @patch("pyinfra.connectors.util.getpass")
+    def test_no_prompt_when_password_given(self, fake_getpass):
+        host = self._make_host(Config(USE_SUDO_PASSWORD=True, SUDO_PASSWORD="configured"))
+
+        command = make_unix_command_for_host(
+            host.state, host, "uptime", _sudo=True, _sudo_password="configured"
+        )
+        assert "sudo -H -A -k" in command.get_raw_value()
+        fake_getpass.assert_not_called()
+
+    @patch("pyinfra.connectors.util.getpass")
+    def test_no_prompt_without_use_sudo_password(self, fake_getpass):
+        host = self._make_host(Config())
+
+        command = make_unix_command_for_host(host.state, host, "uptime", _sudo=True)
+        assert command.get_raw_value() == "sudo -H -n sh -c uptime"
+        fake_getpass.assert_not_called()