Add IOS pre-transfer get free space check (#371)

* Added get_fre_space

* Update 371.added

Co-authored-by: Jeff Kala <48843785+jeffkala@users.noreply.github.com>

---------

Co-authored-by: Jeff Kala <48843785+jeffkala@users.noreply.github.com>

Author: mattmiller87

changes/371.added

@@ -0,0 +1 @@
+Added free space validation for file copy operations on IOS devices.

pyntc/devices/ios_device.py

@@ -141,6 +141,22 @@ def _get_file_system(self):
         log.error("host %s: File system not found with command 'dir'.")
         raise FileSystemNotFoundError(hostname=self.hostname, command="dir")
 
+    def _get_free_space(self, file_system=None):
+        """Return free bytes on ``file_system`` as reported by IOS ``dir`` output."""
+        if file_system is None:
+            file_system = self._get_file_system()
+
+        raw_data = self.show(f"dir {file_system}")
+        # Example: 16777216 bytes total (1592488 bytes free)
+        match = re.search(r"\((\d+)\s+bytes\s+free\)", raw_data)
+        if match is None:
+            log.error("Host %s: could not parse free space from '%s'.", self.host, f"dir {file_system}")
+            raise CommandError(command=f"dir {file_system}", message="Unable to parse free space from dir output.")
+
+        free_bytes = int(match.group(1))
+        log.debug("Host %s: %s bytes free on %s.", self.host, free_bytes, file_system)
+        return free_bytes
+
     # Get the version of the image that is booted into on the device
     def _image_booted(self, image_name, image_pattern=r".*\.(\d+\.\d+\.\w+)\.SPA.+", **vendor_specifics):
         version_data = self.show("show version")
@@ -730,6 +746,7 @@ def file_copy(self, src, dest=None, file_system=None):
         log.debug("Host %s: Local checksum for file %s is %s.", self.host, src, local_checksum)
 
         if not self.verify_file(local_checksum, dest, file_system=file_system):
+            self._check_free_space(os.path.getsize(src), file_system=file_system)
             file_copy = self._file_copy_instance(src, dest, file_system=file_system)
             #        if not self.fc.verify_space_available():
             #            raise FileTransferError('Not enough space available.')
@@ -786,6 +803,7 @@ def remote_file_copy(self, src: FileCopyModel, dest=None, file_system=None, **kw
         if dest is None:
             dest = src.file_name
         if not self.verify_file(src.checksum, dest, hashing_algorithm=src.hashing_algorithm, file_system=file_system):
+            self._pre_transfer_space_check(src, file_system)
             current_prompt = self.native.find_prompt()
 
             # Define prompt mapping for expected prompts during file copy

tests/unit/test_devices/test_ios_device.py

@@ -8,6 +8,7 @@
 from pyntc.devices import IOSDevice
 from pyntc.devices import ios_device as ios_module
 from pyntc.devices.base_device import RollbackError
+from pyntc.errors import NotEnoughFreeSpaceError
 from pyntc.utils.models import FileCopyModel
 
 from .device_mocks.ios import send_command, send_command_expect
@@ -140,9 +141,13 @@ def test_file_copy_remote_exists_not(self, mock_ft):
 
     @mock.patch.object(IOSDevice, "get_local_checksum")
     @mock.patch.object(IOSDevice, "verify_file")
+    @mock.patch.object(IOSDevice, "_check_free_space")
+    @mock.patch("pyntc.devices.ios_device.os.path.getsize", return_value=1024)
     @mock.patch("pyntc.devices.ios_device.FileTransfer", autospec=True)
     @mock.patch.object(IOSDevice, "open")
-    def test_file_copy(self, mock_open, mock_ft, mock_verify_file, mock_get_local_checksum):
+    def test_file_copy(
+        self, mock_open, mock_ft, _mock_getsize, _mock_check_space, mock_verify_file, mock_get_local_checksum
+    ):
         self.device.native.send_command.side_effect = None
         self.device.native.send_command.return_value = "flash: /dev/null"
 
@@ -159,9 +164,13 @@ def test_file_copy(self, mock_open, mock_ft, mock_verify_file, mock_get_local_ch
 
     @mock.patch.object(IOSDevice, "get_local_checksum")
     @mock.patch.object(IOSDevice, "verify_file")
+    @mock.patch.object(IOSDevice, "_check_free_space")
+    @mock.patch("pyntc.devices.ios_device.os.path.getsize", return_value=1024)
     @mock.patch("pyntc.devices.ios_device.FileTransfer", autospec=True)
     @mock.patch.object(IOSDevice, "open")
-    def test_file_copy_different_dest(self, mock_open, mock_ft, mock_verify_file, mock_get_local_checksum):
+    def test_file_copy_different_dest(
+        self, mock_open, mock_ft, _mock_getsize, _mock_check_space, mock_verify_file, mock_get_local_checksum
+    ):
         self.device.native.send_command_timing.side_effect = None
         self.device.native.send_command.return_value = "flash: /dev/null"
         mock_ft_instance = mock_ft.return_value
@@ -177,9 +186,13 @@ def test_file_copy_different_dest(self, mock_open, mock_ft, mock_verify_file, mo
 
     @mock.patch.object(IOSDevice, "get_local_checksum")
     @mock.patch.object(IOSDevice, "verify_file")
+    @mock.patch.object(IOSDevice, "_check_free_space")
+    @mock.patch("pyntc.devices.ios_device.os.path.getsize", return_value=1024)
     @mock.patch("pyntc.devices.ios_device.FileTransfer", autospec=True)
     @mock.patch.object(IOSDevice, "open")
-    def test_file_copy_fail(self, mock_open, mock_ft, mock_verify_file, mock_get_local_checksum):
+    def test_file_copy_fail(
+        self, mock_open, mock_ft, _mock_getsize, _mock_check_space, mock_verify_file, mock_get_local_checksum
+    ):
         self.device.native.send_command_timing.side_effect = None
         self.device.native.send_command.return_value = "flash: /dev/null"
         mock_ft_instance = mock_ft.return_value
@@ -193,9 +206,13 @@ def test_file_copy_fail(self, mock_open, mock_ft, mock_verify_file, mock_get_loc
 
     @mock.patch.object(IOSDevice, "get_local_checksum")
     @mock.patch.object(IOSDevice, "verify_file")
+    @mock.patch.object(IOSDevice, "_check_free_space")
+    @mock.patch("pyntc.devices.ios_device.os.path.getsize", return_value=1024)
     @mock.patch("pyntc.devices.ios_device.FileTransfer", autospec=True)
     @mock.patch.object(IOSDevice, "open")
-    def test_file_copy_socket_closed_good_md5(self, mock_open, mock_ft, mock_verify_file, mock_get_local_checksum):
+    def test_file_copy_socket_closed_good_md5(
+        self, mock_open, mock_ft, _mock_getsize, _mock_check_space, mock_verify_file, mock_get_local_checksum
+    ):
         self.device.native.send_command_timing.side_effect = None
         self.device.native.send_command.return_value = "flash: /dev/null"
         mock_ft_instance = mock_ft.return_value
@@ -214,9 +231,13 @@ def test_file_copy_socket_closed_good_md5(self, mock_open, mock_ft, mock_verify_
 
     @mock.patch.object(IOSDevice, "get_local_checksum")
     @mock.patch.object(IOSDevice, "verify_file")
+    @mock.patch.object(IOSDevice, "_check_free_space")
+    @mock.patch("pyntc.devices.ios_device.os.path.getsize", return_value=1024)
     @mock.patch("pyntc.devices.ios_device.FileTransfer", autospec=True)
     @mock.patch.object(IOSDevice, "open")
-    def test_file_copy_fail_socket_closed_bad_md5(self, mock_open, mock_ft, mock_verify_file, mock_get_local_checksum):
+    def test_file_copy_fail_socket_closed_bad_md5(
+        self, mock_open, mock_ft, _mock_getsize, _mock_check_space, mock_verify_file, mock_get_local_checksum
+    ):
         self.device.native.send_command_timing.side_effect = None
         self.device.native.send_command.return_value = "flash: /dev/null"
         mock_ft_instance = mock_ft.return_value
@@ -479,6 +500,15 @@ def test_get_remote_checksum(self):
             with self.assertRaises(ValueError):
                 self.device.get_remote_checksum("file.txt", hashing_algorithm="invalid_algo", file_system="flash:")
 
+    def test_get_free_space(self):
+        self.device.native.send_command.return_value = "16777216 bytes total (1592488 bytes free)"
+        self.assertEqual(self.device._get_free_space(file_system="flash:"), 1592488)
+
+    def test_get_free_space_raises_when_unparsable(self):
+        self.device.native.send_command.return_value = "Directory of flash:/\nUnable to read totals"
+        with self.assertRaises(ios_module.CommandError):
+            self.device._get_free_space(file_system="flash:")
+
     @mock.patch.object(IOSDevice, "verify_file")
     def test_remote_file_copy_success(self, mock_verify):
         # Setup file model
@@ -594,6 +624,49 @@ def test_remote_file_copy_failure_on_error_output(self, mock_verify):
             with self.assertRaises(FileTransferError):
                 self.device.remote_file_copy(src)
 
+    @mock.patch.object(IOSDevice, "verify_file")
+    def test_remote_file_copy_raises_not_enough_free_space(self, mock_verify):
+        src = FileCopyModel(
+            download_url="http://1.1.1.1/test.bin",
+            checksum="12345",
+            file_name="test.bin",
+            hashing_algorithm="md5",
+            file_size=2,
+            file_size_unit="gigabytes",
+        )
+        mock_verify.return_value = False
+        self.device.native.send_command.return_value = "16777216 bytes total (1592488 bytes free)"
+
+        with self.assertRaises(NotEnoughFreeSpaceError):
+            self.device.remote_file_copy(src, file_system="flash:")
+
+        assert not any(
+            "copy " in str(call.kwargs.get("command_string", call.args[0] if call.args else ""))
+            for call in self.device.native.send_command.call_args_list
+        )
+
+    @mock.patch.object(IOSDevice, "verify_file")
+    @mock.patch.object(IOSDevice, "_check_free_space")
+    def test_remote_file_copy_skips_space_check_when_file_size_omitted(self, mock_check_free_space, mock_verify):
+        src = FileCopyModel(
+            download_url="sftp://1.1.1.1/test.bin",
+            checksum="12345",
+            file_name="test.bin",
+            hashing_algorithm="md5",
+            timeout=300,
+        )
+        mock_verify.side_effect = [False, True]
+        self.device.native.send_command.side_effect = [
+            "Address or name of remote host [1.1.1.1]?",
+            "123456 bytes copied in 10.2 secs. Copy complete.",
+        ]
+        self.device.native.find_prompt.return_value = "Router#"
+
+        self.device.remote_file_copy(src, file_system="flash:")
+
+        mock_check_free_space.assert_not_called()
+        self.device.native.send_command.assert_called()
+
 
 if __name__ == "__main__":
     unittest.main()