Implement automatic sudo password prompting (#609)

* Implement automatic sudo password prompting for the local connector.

This works by executing the command, detecting the sudo failed password
message output and then propmts the user for the password and re-executes
the command using that.

* Implement automatic sudo password prompting for the SSH connector.
* Fix command debug log calls.
* Add test for automatic sudo prompting in SSH connector.

Author: Fizzadar

pyinfra/api/connectors/local.py

@@ -13,6 +13,7 @@
 from pyinfra.api.util import get_file_io
 
 from .util import (
+    execute_command_with_sudo_retry,
     get_sudo_password,
     make_unix_command,
     run_local_process,
@@ -69,20 +70,25 @@ def run_shell_command(
             put_file=put_file,
         )
 
-    command = make_unix_command(command, state=state, **command_kwargs)
-    actual_command = command.get_raw_value()
+    def execute_command():
+        unix_command = make_unix_command(command, state=state, **command_kwargs)
+        actual_command = unix_command.get_raw_value()
 
-    logger.debug('--> Running command on localhost: {0}'.format(command))
+        logger.debug('--> Running command on localhost: {0}'.format(unix_command))
 
-    if print_input:
-        click.echo('{0}>>> {1}'.format(host.print_prefix, command), err=True)
+        if print_input:
+            click.echo('{0}>>> {1}'.format(host.print_prefix, unix_command), err=True)
 
-    return_code, combined_output = run_local_process(
-        actual_command,
-        stdin=stdin,
-        timeout=timeout,
-        print_output=print_output,
-        print_prefix=host.print_prefix,
+        return run_local_process(
+            actual_command,
+            stdin=stdin,
+            timeout=timeout,
+            print_output=print_output,
+            print_prefix=host.print_prefix,
+        )
+
+    return_code, combined_output = execute_command_with_sudo_retry(
+        host, command_kwargs, execute_command,
     )
 
     if success_exit_codes:

pyinfra/api/connectors/ssh.py

@@ -32,6 +32,7 @@
 
 from .sshuserclient import SSHClient
 from .util import (
+    execute_command_with_sudo_retry,
     get_sudo_password,
     make_unix_command,
     read_buffers_into_queue,
@@ -273,41 +274,48 @@ def run_shell_command(
             put_file=put_file,
         )
 
-    command = make_unix_command(command, state=state, **command_kwargs)
-    actual_command = command.get_raw_value()
+    def execute_command():
+        unix_command = make_unix_command(command, state=state, **command_kwargs)
+        actual_command = unix_command.get_raw_value()
 
-    logger.debug('Running command on {0}: (pty={1}) {2}'.format(
-        host.name, get_pty, command,
-    ))
+        logger.debug('Running command on {0}: (pty={1}) {2}'.format(
+            host.name, get_pty, unix_command,
+        ))
 
-    if print_input:
-        click.echo('{0}>>> {1}'.format(host.print_prefix, command), err=True)
+        if print_input:
+            click.echo('{0}>>> {1}'.format(host.print_prefix, unix_command), err=True)
 
-    # Run it! Get stdout, stderr & the underlying channel
-    stdin_buffer, stdout_buffer, stderr_buffer = host.connection.exec_command(
-        actual_command,
-        get_pty=get_pty,
-    )
+        # Run it! Get stdout, stderr & the underlying channel
+        stdin_buffer, stdout_buffer, stderr_buffer = host.connection.exec_command(
+            actual_command,
+            get_pty=get_pty,
+        )
 
-    if stdin:
-        write_stdin(stdin, stdin_buffer)
+        if stdin:
+            write_stdin(stdin, stdin_buffer)
 
-    combined_output = read_buffers_into_queue(
-        stdout_buffer,
-        stderr_buffer,
-        timeout=timeout,
-        print_output=print_output,
-        print_prefix=host.print_prefix,
-    )
+        combined_output = read_buffers_into_queue(
+            stdout_buffer,
+            stderr_buffer,
+            timeout=timeout,
+            print_output=print_output,
+            print_prefix=host.print_prefix,
+        )
+
+        logger.debug('Waiting for exit status...')
+        exit_status = stdout_buffer.channel.recv_exit_status()
+        logger.debug('Command exit status: {0}'.format(exit_status))
 
-    logger.debug('Waiting for exit status...')
-    exit_status = stdout_buffer.channel.recv_exit_status()
-    logger.debug('Command exit status: {0}'.format(exit_status))
+        return exit_status, combined_output
+
+    return_code, combined_output = execute_command_with_sudo_retry(
+        host, command_kwargs, execute_command,
+    )
 
     if success_exit_codes:
-        status = exit_status in success_exit_codes
+        status = return_code in success_exit_codes
     else:
-        status = exit_status == 0
+        status = return_code == 0
 
     if return_combined_output:
         return status, combined_output

pyinfra/api/connectors/util.py

@@ -51,6 +51,21 @@ def _print(line):
             _print(line)
 
 
+def execute_command_with_sudo_retry(host, command_kwargs, execute_command):
+    return_code, combined_output = execute_command()
+
+    if return_code != 0 and combined_output:
+        last_line = combined_output[-1][1]
+        if last_line == 'sudo: a password is required':
+            command_kwargs['use_sudo_password'] = get_sudo_password(
+                host,
+                use_sudo_password=True,  # ask for the password
+            )
+            return_code, combined_output = execute_command()
+
+    return return_code, combined_output
+
+
 def run_local_process(
     command,
     stdin=None,
@@ -151,11 +166,11 @@ def write_stdin(stdin, buffer):
     buffer.close()
 
 
-def get_sudo_password(state, host, use_sudo_password, run_shell_command, put_file):
+def get_sudo_password(host, use_sudo_password):
     sudo_askpass_uploaded = host.connector_data.get('sudo_askpass_uploaded', False)
     if not sudo_askpass_uploaded:
-        put_file(state, host, get_sudo_askpass_exe(), SUDO_ASKPASS_EXE_FILENAME)
-        run_shell_command(state, host, 'chmod +x {0}'.format(SUDO_ASKPASS_EXE_FILENAME))
+        host.put_file(get_sudo_askpass_exe(), SUDO_ASKPASS_EXE_FILENAME)
+        host.run_shell_command('chmod +x {0}'.format(SUDO_ASKPASS_EXE_FILENAME))
         host.connector_data['sudo_askpass_uploaded'] = True
 
     if use_sudo_password is True:

tests/test_connectors/test_ssh.py

@@ -463,6 +463,41 @@ def test_run_shell_command_error(self, fake_ssh_client):
         assert len(out) == 3
         assert out[0] is False
 
+    @patch('pyinfra.api.connectors.ssh.SSHClient')
+    @patch('pyinfra.api.connectors.util.get_sudo_password')
+    def test_run_shell_command_retry_for_sudo_password(
+        self,
+        fake_get_sudo_password,
+        fake_ssh_client,
+    ):
+        fake_get_sudo_password.return_value = ['FILENAME', 'PASSWORD']
+
+        fake_ssh = MagicMock()
+        fake_stdin = MagicMock()
+        fake_stdout = MagicMock()
+        fake_stderr = ['sudo: a password is required']
+        fake_ssh.exec_command.return_value = fake_stdin, fake_stdout, fake_stderr
+
+        fake_ssh_client.return_value = fake_ssh
+
+        inventory = make_inventory(hosts=('somehost',))
+        state = State(inventory, Config())
+        host = inventory.get_host('somehost')
+        host.connect(state)
+
+        command = 'echo hi'
+        return_values = [1, 0]  # return 0 on the second call
+        fake_stdout.channel.recv_exit_status.side_effect = lambda: return_values.pop(0)
+
+        out = host.run_shell_command(command)
+        assert len(out) == 3
+        assert out[0] is True
+        assert fake_get_sudo_password.called
+        fake_ssh.exec_command.assert_called_with(
+            "env SUDO_ASKPASS=FILENAME PYINFRA_SUDO_PASSWORD=PASSWORD sh -c 'echo hi'",
+            get_pty=False,
+        )
+
     @patch('pyinfra.api.connectors.ssh.SSHClient')
     @patch('pyinfra.api.connectors.ssh.SFTPClient')
     def test_put_file(self, fake_sftp_client, fake_ssh_client):