diff --git a/nisyscfg/errors.py b/nisyscfg/errors.py index 6f96947..12ecb2f 100644 --- a/nisyscfg/errors.py +++ b/nisyscfg/errors.py @@ -23,6 +23,10 @@ def __init__(self, message): super(Error, self).__init__(message) +class InvalidSystemImageError(Error): + """This error is raised when the system image is invalid.""" + + class LibraryError(Error): def __init__(self, code, description): assert _is_error(code), "Should not raise Error if code is not fatal." diff --git a/nisyscfg/system.py b/nisyscfg/system.py index 59399e2..700da43 100644 --- a/nisyscfg/system.py +++ b/nisyscfg/system.py @@ -1,4 +1,10 @@ +from __future__ import annotations + import ctypes +import pathlib +import tempfile +import zipfile +from contextlib import ExitStack import nisyscfg import nisyscfg._library_singleton @@ -1066,3 +1072,75 @@ def save_changes(self) -> SaveChangesResult: nisyscfg.errors.handle_error(self, error_code) nisyscfg.errors.handle_error(self, error_code_2) return SaveChangesResult(restart_required=restart_required.value != 0, details=details) + + def set_system_image( + self, + source: str | pathlib.Path, + auto_restart: bool = True, + encryption_passphrase: str | None = None, + exclude_paths: List[str] | None = None, + original_system_only: bool = False, + network_settings: nisyscfg.enums.NetworkInterfaceSettings = nisyscfg.enums.NetworkInterfaceSettings.PRESERVE_PRIMARY_PRESERVE_OTHERS, + ) -> None: + """Applies an image to a system. + + The system image is a copy of the contents and software on the primary + hard drive of a specified target system. Applying the image to the + system restores it to the state captured when the image was created. + This image can exist as a folder or a zipped archive. + + source - The path to the system image file or folder. + + auto_restart - Restarts the system into install mode by default before + the operation is performed, and restarts back to a running state after + the operation is complete. If you choose not to restart automatically, + and the system is not in install mode, the resulting image may not be + valid. + + encryption_passphrase - A passphrase used to encrypt a portion of the + image that contains sensitive information. + + exclude_paths - Specifies the list of files and folders to exclude from + the target image. Files on the blacklist will not be copied from the + image to the target and they will not be removed from the target. + + original_system_only - Verifies that the target system has the same MAC + address as the system from which the image was originally created. + Selecting True will allow you to restore an image from the exact same + target from which the image was created only. This option is False by + default. When the option is False, this operation can also apply the + system image to other targets of the same device class. + + network_settings - Resets the primary network adapter and disables + secondary adapters by default. + """ + source = pathlib.Path(source) + if not source.exists(): + raise FileNotFoundError(f"The source {source} does not exist.") + if exclude_paths is not None: + raise NotImplementedError("Excluding paths is not implemented yet.") + + with ExitStack() as stack: + if source.is_dir(): + source_folder = str(source) + else: + source_folder = str(stack.enter_context(tempfile.TemporaryDirectory())) + try: + zip_ref = stack.enter_context(zipfile.ZipFile(source, "r")) + except zipfile.BadZipFile: + raise nisyscfg.errors.InvalidSystemImageError( + f"{source} is not a valid zip file." + ) + zip_ref.extractall(source_folder) + + error_code = self._library.SetSystemImageFromFolder2( + self._session, + nisyscfg.enums.Bool(auto_restart), + c_string_encode(source_folder), + c_string_encode(encryption_passphrase), + len(exclude_paths) if exclude_paths else 0, + exclude_paths, + nisyscfg.enums.Bool(original_system_only), + network_settings, + ) + nisyscfg.errors.handle_error(self, error_code) diff --git a/tests/mock_invalid_system_image.bin b/tests/mock_invalid_system_image.bin new file mode 100644 index 0000000..e466dcb --- /dev/null +++ b/tests/mock_invalid_system_image.bin @@ -0,0 +1 @@ +invalid \ No newline at end of file diff --git a/tests/mock_system_image.zip b/tests/mock_system_image.zip new file mode 100644 index 0000000..03314e5 Binary files /dev/null and b/tests/mock_system_image.zip differ diff --git a/tests/test_session.py b/tests/test_session.py index 9775947..caf7a5b 100644 --- a/tests/test_session.py +++ b/tests/test_session.py @@ -1,4 +1,6 @@ import ctypes +import pathlib + import hightime import nisyscfg as nisyscfg import nisyscfg.enums @@ -112,9 +114,9 @@ def initialize_session_mock( def _get_status_description_mock(session_handle, status, detailed_description): - ctypes.cast( - detailed_description, ctypes.POINTER(ctypes.c_void_p) - ).contents.value = STATUS_DESCRIPTION_VOID_P.value + ctypes.cast(detailed_description, ctypes.POINTER(ctypes.c_void_p)).contents.value = ( + STATUS_DESCRIPTION_VOID_P.value + ) return nisyscfg.errors.Status.OK @@ -1086,3 +1088,68 @@ def test_filter_xnet_has_protocol(lib_mock): with nisyscfg.Session() as session: filter = session.create_filter() assert "protocol" in dir(filter.xnet) + + +def test_session_set_system_image_with_folder(lib_mock, tmp_path): + lib_mock.return_value.NISysCfgSetSystemImageFromFolder2.return_value = nisyscfg.errors.Status.OK + + with nisyscfg.Session() as session: + session.set_system_image(tmp_path) + + expected_calls = [ + mock.call(mock.ANY, mock.ANY), + mock.call().NISysCfgInitializeSession( + mock.ANY, mock.ANY, mock.ANY, mock.ANY, mock.ANY, mock.ANY, mock.ANY, mock.ANY + ), + mock.call().NISysCfgSetSystemImageFromFolder2( + CVoidPMatcher(SESSION_HANDLE), + True, + str(tmp_path).encode("ascii"), + None, + 0, + None, + False, + nisyscfg.enums.NetworkInterfaceSettings.PRESERVE_PRIMARY_PRESERVE_OTHERS, + ), + mock.call().NISysCfgCloseHandle(mock.ANY), + ] + assert lib_mock.mock_calls == expected_calls + + +def test_session_set_system_image_with_file(lib_mock, tmp_path): + lib_mock.return_value.NISysCfgSetSystemImageFromFolder2.return_value = nisyscfg.errors.Status.OK + system_image_path = pathlib.Path(__file__).parent / "mock_system_image.zip" + + with mock.patch("tempfile.TemporaryDirectory") as mock_tempdir: + mock_tempdir.return_value.__enter__.return_value = tmp_path + with nisyscfg.Session() as session: + session.set_system_image(system_image_path) + + expected_calls = [ + mock.call(mock.ANY, mock.ANY), + mock.call().NISysCfgInitializeSession( + mock.ANY, mock.ANY, mock.ANY, mock.ANY, mock.ANY, mock.ANY, mock.ANY, mock.ANY + ), + mock.call().NISysCfgSetSystemImageFromFolder2( + CVoidPMatcher(SESSION_HANDLE), + True, + str(tmp_path).encode("ascii"), + None, + 0, + None, + False, + nisyscfg.enums.NetworkInterfaceSettings.PRESERVE_PRIMARY_PRESERVE_OTHERS, + ), + mock.call().NISysCfgCloseHandle(mock.ANY), + ] + assert lib_mock.mock_calls == expected_calls + + +def test_session_set_system_image_with_invalid_file(lib_mock, tmp_path): + system_image_path = pathlib.Path(__file__).parent / "mock_invalid_system_image.bin" + + with mock.patch("tempfile.TemporaryDirectory") as mock_tempdir: + mock_tempdir.return_value.__enter__.return_value = tmp_path + with nisyscfg.Session() as session: + with pytest.raises(nisyscfg.errors.InvalidSystemImageError): + session.set_system_image(system_image_path)