Handle 'smart' scp_if_ssh option for fetch (#18125)

lib/ansible/plugins/connection/ssh.py

@@ -599,6 +599,55 @@ class Connection(ConnectionBase):
 
         return (returncode, stdout, stderr)
 
+    def _file_transport_command(self, in_path, out_path, sftp_action):
+        # scp and sftp require square brackets for IPv6 addresses, but
+        # accept them for hostnames and IPv4 addresses too.
+        host = '[%s]' % self.host
+
+        # since this can be a non-bool now, we need to handle it correctly
+        scp_if_ssh = C.DEFAULT_SCP_IF_SSH
+        if not isinstance(scp_if_ssh, bool):
+            scp_if_ssh = scp_if_ssh.lower()
+            if scp_if_ssh in BOOLEANS:
+                scp_if_ssh = boolean(scp_if_ssh)
+            elif scp_if_ssh != 'smart':
+                raise AnsibleOptionsError('scp_if_ssh needs to be one of [smart|True|False]')
+
+        # create a list of commands to use based on config options
+        methods = ['sftp']
+        if scp_if_ssh == 'smart':
+            methods.append('scp')
+        elif scp_if_ssh:
+            methods = ['scp']
+
+        success = False
+        res = None
+        for method in methods:
+            if method == 'sftp':
+                cmd = self._build_command('sftp', to_bytes(host))
+                in_data = u"{0} {1} {2}\n".format(sftp_action, pipes.quote(in_path), pipes.quote(out_path))
+            elif method == 'scp':
+                cmd = self._build_command('scp', in_path, u'{0}:{1}'.format(host, pipes.quote(out_path)))
+                in_data = None
+
+            in_data = to_bytes(in_data, nonstring='passthru')
+            (returncode, stdout, stderr) = self._run(cmd, in_data, checkrc=False)
+            # Check the return code and rollover to next method if failed
+            if returncode == 0:
+                success = True
+                break
+            else:
+                # If not in smart mode, the data will be printed by the raise below
+                if scp_if_ssh == 'smart':
+                    display.warning(msg='%s transfer mechanism failed on %s. Use ANSIBLE_DEBUG=1 to see detailed information' % (method, host))
+                    display.debug(msg='%s' % to_native(stdout))
+                    display.debug(msg='%s' % to_native(stderr))
+                res = (returncode, stdout, stderr)
+
+        if not success:
+            raise AnsibleError("failed to transfer file to {0}:\n{1}\n{2}"\
+                    .format(to_native(out_path), to_native(res[1]), to_native(res[2])))
+
     #
     # Main public methods
     #
@@ -655,53 +704,7 @@ class Connection(ConnectionBase):
         if not os.path.exists(to_bytes(in_path, errors='surrogate_or_strict')):
             raise AnsibleFileNotFound("file or module does not exist: {0}".format(to_native(in_path)))
 
-        # scp and sftp require square brackets for IPv6 addresses, but
+        self._file_transport_command(in_path, out_path, 'put')
-        # accept them for hostnames and IPv4 addresses too.
-        host = '[%s]' % self.host
-
-        # since this can be a non-bool now, we need to handle it correctly
-        scp_if_ssh = C.DEFAULT_SCP_IF_SSH
-        if not isinstance(scp_if_ssh, bool):
-            scp_if_ssh = scp_if_ssh.lower()
-            if scp_if_ssh in BOOLEANS:
-                scp_if_ssh = boolean(scp_if_ssh)
-            elif scp_if_ssh != 'smart':
-                raise AnsibleOptionsError('scp_if_ssh needs to be one of [smart|True|False]')
-
-        # create a list of commands to use based on config options
-        methods = ['sftp']
-        if scp_if_ssh == 'smart':
-            methods.append('scp')
-        elif scp_if_ssh:
-            methods = ['scp']
-
-        success = False
-        res = None
-        for method in methods:
-            if method == 'sftp':
-                cmd = self._build_command('sftp', to_bytes(host))
-                in_data = u"put {0} {1}\n".format(pipes.quote(in_path), pipes.quote(out_path))
-            elif method == 'scp':
-                cmd = self._build_command('scp', in_path, u'{0}:{1}'.format(host, pipes.quote(out_path)))
-                in_data = None
-
-            in_data = to_bytes(in_data, nonstring='passthru')
-            (returncode, stdout, stderr) = self._run(cmd, in_data, checkrc=False)
-            # Check the return code and rollover to next method if failed
-            if returncode == 0:
-                success = True
-                break
-            else:
-                # If not in smart mode, the data will be printed by the raise below 
-                if scp_if_ssh == 'smart':
-                    display.warning(msg='%s transfer mechanism failed on %s. Use ANSIBLE_DEBUG=1 to see detailed information' % (method, host))
-                    display.debug(msg='%s' % to_native(stdout))
-                    display.debug(msg='%s' % to_native(stderr))
-                res = (returncode, stdout, stderr)
-
-        if not success:
-            raise AnsibleError("failed to transfer file to {0}:\n{1}\n{2}"\
-                    .format(to_native(out_path), to_native(res[1]), to_native(res[2])))
 
     def fetch_file(self, in_path, out_path):
         ''' fetch a file from remote to local '''
@@ -709,23 +712,7 @@ class Connection(ConnectionBase):
         super(Connection, self).fetch_file(in_path, out_path)
 
         display.vvv(u"FETCH {0} TO {1}".format(in_path, out_path), host=self.host)
-
+        self._file_transport_command(in_path, out_path, 'get')
-        # scp and sftp require square brackets for IPv6 addresses, but
-        # accept them for hostnames and IPv4 addresses too.
-        host = '[%s]' % self.host
-
-        if C.DEFAULT_SCP_IF_SSH:
-            cmd = self._build_command('scp', u'{0}:{1}'.format(host, pipes.quote(in_path)), out_path)
-            in_data = None
-        else:
-            cmd = self._build_command('sftp', host)
-            in_data = u"get {0} {1}\n".format(pipes.quote(in_path), pipes.quote(out_path))
-
-        in_data = to_bytes(in_data, nonstring='passthru')
-        (returncode, stdout, stderr) = self._run(cmd, in_data)
-
-        if returncode != 0:
-            raise AnsibleError("failed to transfer file from {0}:\n{1}\n{2}".format(in_path, stdout, stderr))
 
     def close(self):
         # If we have a persistent ssh connection (ControlPersist), we can ask it

test/units/plugins/connections/test_connection_ssh.py

@@ -308,6 +308,19 @@ class TestConnectionBaseClass(unittest.TestCase):
         conn._run.return_value = (0, '', '')
         conn.host = "some_host"
 
+        # Test with C.DEFAULT_SCP_IF_SSH set to smart
+        # Test when SFTP works
+        C.DEFAULT_SCP_IF_SSH = 'smart'
+        expected_in_data = b' '.join((b'put', to_bytes(pipes.quote('/path/to/in/file')), to_bytes(pipes.quote('/path/to/dest/file')))) + b'\n'
+        conn.put_file('/path/to/in/file', '/path/to/dest/file')
+        conn._run.assert_called_with('some command to run', expected_in_data, checkrc=False)
+
+        # Test when SFTP doesn't work but SCP does
+        conn._run.side_effect = [(1, 'stdout', 'some errors'), (0, '', '')]
+        conn.put_file('/path/to/in/file', '/path/to/dest/file')
+        conn._run.assert_called_with('some command to run', None, checkrc=False)
+        conn._run.side_effect = None
+
         # test with C.DEFAULT_SCP_IF_SSH enabled
         C.DEFAULT_SCP_IF_SSH = True
         conn.put_file('/path/to/in/file', '/path/to/dest/file')
@@ -328,6 +341,7 @@ class TestConnectionBaseClass(unittest.TestCase):
         conn.put_file(u'/path/to/in/file/with/unicode-fö〩', u'/path/to/dest/file/with/unicode-fö〩')
         conn._run.assert_called_with('some command to run', expected_in_data, checkrc=False)
 
+
         # test that a non-zero rc raises an error
         conn._run.return_value = (1, 'stdout', 'some errors')
         self.assertRaises(AnsibleError, conn.put_file, '/path/to/bad/file', '/remote/path/to/file')
@@ -348,25 +362,38 @@ class TestConnectionBaseClass(unittest.TestCase):
         conn._run.return_value = (0, '', '')
         conn.host = "some_host"
 
+        # Test with C.DEFAULT_SCP_IF_SSH set to smart
+        # Test when SFTP works
+        C.DEFAULT_SCP_IF_SSH = 'smart'
+        expected_in_data = b' '.join((b'get', to_bytes(pipes.quote('/path/to/in/file')), to_bytes(pipes.quote('/path/to/dest/file')))) + b'\n'
+        conn.fetch_file('/path/to/in/file', '/path/to/dest/file')
+        conn._run.assert_called_with('some command to run', expected_in_data, checkrc=False)
+
+        # Test when SFTP doesn't work but SCP does
+        conn._run.side_effect = [(1, 'stdout', 'some errors'), (0, '', '')]
+        conn.fetch_file('/path/to/in/file', '/path/to/dest/file')
+        conn._run.assert_called_with('some command to run', None, checkrc=False)
+        conn._run.side_effect = None
+
         # test with C.DEFAULT_SCP_IF_SSH enabled
         C.DEFAULT_SCP_IF_SSH = True
         conn.fetch_file('/path/to/in/file', '/path/to/dest/file')
-        conn._run.assert_called_with('some command to run', None)
+        conn._run.assert_called_with('some command to run', None, checkrc=False)
 
         conn.fetch_file(u'/path/to/in/file/with/unicode-fö〩', u'/path/to/dest/file/with/unicode-fö〩')
-        conn._run.assert_called_with('some command to run', None)
+        conn._run.assert_called_with('some command to run', None, checkrc=False)
 
         # test with C.DEFAULT_SCP_IF_SSH disabled
         C.DEFAULT_SCP_IF_SSH = False
         expected_in_data = b' '.join((b'get', to_bytes(pipes.quote('/path/to/in/file')), to_bytes(pipes.quote('/path/to/dest/file')))) + b'\n'
         conn.fetch_file('/path/to/in/file', '/path/to/dest/file')
-        conn._run.assert_called_with('some command to run', expected_in_data)
+        conn._run.assert_called_with('some command to run', expected_in_data, checkrc=False)
 
         expected_in_data = b' '.join((b'get',
             to_bytes(pipes.quote('/path/to/in/file/with/unicode-fö〩')),
             to_bytes(pipes.quote('/path/to/dest/file/with/unicode-fö〩')))) + b'\n'
         conn.fetch_file(u'/path/to/in/file/with/unicode-fö〩', u'/path/to/dest/file/with/unicode-fö〩')
-        conn._run.assert_called_with('some command to run', expected_in_data)
+        conn._run.assert_called_with('some command to run', expected_in_data, checkrc=False)
 
         # test that a non-zero rc raises an error
         conn._run.return_value = (1, 'stdout', 'some errors')