Move sudo password retrieve into make host unix command util.

Author: Fizzadar

pyinfra/api/connectors/chroot.py

@@ -12,7 +12,7 @@
 from pyinfra.progress import progress_spinner
 
 from .local import run_shell_command as run_local_shell_command
-from .util import get_sudo_password, make_unix_command_for_host
+from .util import make_unix_command_for_host
 
 
 @memoize
@@ -57,16 +57,8 @@ def run_shell_command(
     print_output=False,
     print_input=False,
     return_combined_output=False,
-    use_sudo_password=False,
     **command_kwargs
 ):
-
-    if use_sudo_password:
-        command_kwargs['use_sudo_password'] = get_sudo_password(
-            host,
-            use_sudo_password,
-        )
-
     chroot_directory = host.host_data['chroot_directory']
 
     command = make_unix_command_for_host(state, host, command, **command_kwargs)

pyinfra/api/connectors/local.py

@@ -14,7 +14,6 @@
 
 from .util import (
     execute_command_with_sudo_retry,
-    get_sudo_password,
     make_unix_command_for_host,
     run_local_process,
     split_combined_output,
@@ -43,7 +42,6 @@ def run_shell_command(
     print_output=False,
     print_input=False,
     return_combined_output=False,
-    use_sudo_password=False,
     **command_kwargs
 ):
     '''
@@ -63,11 +61,6 @@ def run_shell_command(
         stdout and stderr are both lists of strings from each buffer.
     '''
 
-    if use_sudo_password:
-        command_kwargs['use_sudo_password'] = get_sudo_password(
-            host, use_sudo_password,
-        )
-
     def execute_command():
         unix_command = make_unix_command_for_host(state, host, command, **command_kwargs)
         actual_command = unix_command.get_raw_value()

pyinfra/api/connectors/ssh.py

@@ -33,7 +33,6 @@
 from .sshuserclient import SSHClient
 from .util import (
     execute_command_with_sudo_retry,
-    get_sudo_password,
     make_unix_command_for_host,
     read_buffers_into_queue,
     run_local_process,
@@ -261,7 +260,6 @@ def run_shell_command(
     print_output=False,
     print_input=False,
     return_combined_output=False,
-    use_sudo_password=False,
     **command_kwargs
 ):
     '''
@@ -282,11 +280,6 @@ def run_shell_command(
         stdout and stderr are both lists of strings from each buffer.
     '''
 
-    if use_sudo_password:
-        command_kwargs['use_sudo_password'] = get_sudo_password(
-            host, use_sudo_password,
-        )
-
     def execute_command():
         unix_command = make_unix_command_for_host(state, host, command, **command_kwargs)
         actual_command = unix_command.get_raw_value()

pyinfra/api/connectors/util.py

@@ -57,10 +57,7 @@ def execute_command_with_sudo_retry(host, command_kwargs, 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
-            )
+            command_kwargs['use_sudo_password'] = True  # ask for the password
             return_code, combined_output = execute_command()
 
     return return_code, combined_output
@@ -166,7 +163,7 @@ def write_stdin(stdin, buffer):
     buffer.close()
 
 
-def get_sudo_password(host, use_sudo_password):
+def _get_sudo_password(host, use_sudo_password):
     sudo_askpass_uploaded = host.connector_data.get('sudo_askpass_uploaded', False)
     if not sudo_askpass_uploaded:
         host.put_file(get_sudo_askpass_exe(), SUDO_ASKPASS_EXE_FILENAME)
@@ -182,7 +179,7 @@ def get_sudo_password(host, use_sudo_password):
     else:
         sudo_password = use_sudo_password
 
-    return (SUDO_ASKPASS_EXE_FILENAME, shlex_quote(sudo_password))
+    return shlex_quote(sudo_password)
 
 
 def remove_any_sudo_askpass_file(host):
@@ -230,6 +227,10 @@ def make_unix_command_for_host(state, host, *command_args, **command_kwargs):
             ('use_su_login', 'preserve_su_env', 'su_shell'),
         )
 
+    use_sudo_password = command_kwargs.pop('use_sudo_password', None)
+    if use_sudo_password:
+        command_kwargs['sudo_password'] = _get_sudo_password(host, use_sudo_password)
+
     return make_unix_command(*command_args, **command_kwargs)
 
 
@@ -247,7 +248,7 @@ def make_unix_command(
     sudo=False,
     sudo_user=None,
     use_sudo_login=False,
-    use_sudo_password=False,
+    sudo_password=False,
     preserve_sudo_env=False,
     # Doas config
     doas=False,
@@ -284,18 +285,17 @@ def make_unix_command(
         if doas_user:
             command_bits.extend(['-u', doas_user])
 
-    if use_sudo_password:
-        askpass_filename, sudo_password = use_sudo_password
+    if sudo_password:
         command_bits.extend([
             'env',
-            'SUDO_ASKPASS={0}'.format(askpass_filename),
+            'SUDO_ASKPASS={0}'.format(SUDO_ASKPASS_EXE_FILENAME),
             MaskString('{0}={1}'.format(SUDO_ASKPASS_ENV_VAR, sudo_password)),
         ])
 
     if sudo:
         command_bits.extend(['sudo', '-H'])
 
-        if use_sudo_password:
+        if sudo_password:
             command_bits.extend(['-A', '-k'])  # use askpass, disable cache
         else:
             command_bits.append('-n')  # disable prompt/interactivity

tests/test_connectors/test_ssh.py

@@ -569,13 +569,13 @@ def test_run_shell_command_sudo_password_automatic_prompt(
     #
 
     @patch('pyinfra.api.connectors.ssh.SSHClient')
-    @patch('pyinfra.api.connectors.util.get_sudo_password')
+    @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_get_sudo_password.return_value = 'PASSWORD'
 
         fake_ssh = MagicMock()
         fake_stdin = MagicMock()
@@ -599,7 +599,7 @@ def test_run_shell_command_retry_for_sudo_password(
         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'",
+            "env SUDO_ASKPASS=pyinfra-sudo-askpass PYINFRA_SUDO_PASSWORD=PASSWORD sh -c 'echo hi'",
             get_pty=False,
         )