diff --git a/src/scheduler/unpacking_scheduler.py b/src/scheduler/unpacking_scheduler.py index beb28851d..ca1bb2077 100644 --- a/src/scheduler/unpacking_scheduler.py +++ b/src/scheduler/unpacking_scheduler.py @@ -22,6 +22,7 @@ ) from objects.firmware import Firmware from storage.db_interface_backend import BackendDbInterface +from storage.db_interface_base import DbInterfaceError from unpacker.extraction_container import ExtractionContainer from unpacker.unpack import Unpacker from unpacker.unpack_base import ExtractionError @@ -205,8 +206,12 @@ def work_thread(self, task: FileObject, container: ExtractionContainer): logging.info(f'Unpacking completed: {task.uid} (extracted files: {len(extracted_objects)})') # each worker needs its own interface because connections are not thread-safe db_interface = self.db_interface() - db_interface.add_object(task) # save FO before submitting to analysis scheduler - self.post_unpack(task) + try: + db_interface.add_object(task) # save FO before submitting to analysis scheduler + self.post_unpack(task) + except DbInterfaceError as error: + logging.error(str(error)) + extracted_objects = [] self._update_currently_unpacked(task, extracted_objects, db_interface) self._schedule_extracted_files(extracted_objects) diff --git a/src/storage/db_interface_backend.py b/src/storage/db_interface_backend.py index 147251887..29f7afbdf 100644 --- a/src/storage/db_interface_backend.py +++ b/src/storage/db_interface_backend.py @@ -149,6 +149,14 @@ def add_child_to_parent(self, parent_uid: str, child_uid: str): def update_object(self, fw_object: FileObject): if isinstance(fw_object, Firmware): + if not self.is_firmware(fw_object.uid): + # special case: Trying to upload a file as firmware that is already in the DB as part of another + # firmware. This is currently not possible and will likely cause errors + parent_fw = self.get_parent_fw(fw_object.uid) + raise DbInterfaceError( + 'Cannot upload file as firmware that is part of another firmware. ' + f'The file you are trying to upload is already part of the following firmware images: {parent_fw}' + ) self.update_firmware(fw_object) self.update_file_object(fw_object) diff --git a/src/test/integration/storage/test_db_interface_backend.py b/src/test/integration/storage/test_db_interface_backend.py index 555915f90..ec143e8e7 100644 --- a/src/test/integration/storage/test_db_interface_backend.py +++ b/src/test/integration/storage/test_db_interface_backend.py @@ -2,6 +2,7 @@ import pytest +from storage.db_interface_base import DbInterfaceError from test.common_helper import create_test_file_object, create_test_firmware from .helper import TEST_FO, TEST_FW, create_fw_with_child_fo, create_fw_with_parent_and_child, add_included_file @@ -143,6 +144,17 @@ def test_update_duplicate_same_fw(backend_db, frontend_db): assert db_fo.parents == {fw.uid} +def test_update_duplicate_file_as_fw(backend_db): + # special case: trying to upload a file as FW that is already in the DB as part of another FW -> should cause error + fo, fw = create_fw_with_child_fo() + backend_db.insert_multiple_objects(fw, fo) + fw2 = create_test_firmware() + fw2.uid = fo.uid + + with pytest.raises(DbInterfaceError): + backend_db.add_object(fw2) + + def test_analysis_exists(backend_db): assert backend_db.analysis_exists(TEST_FO.uid, 'file_type') is False backend_db.insert_file_object(TEST_FO)