diff --git a/pycti/connector/opencti_connector_helper.py b/pycti/connector/opencti_connector_helper.py index 0441d88e2..1b9cb7a78 100644 --- a/pycti/connector/opencti_connector_helper.py +++ b/pycti/connector/opencti_connector_helper.py @@ -1669,6 +1669,8 @@ def send_stix2_bundle(self, bundle: str, **kwargs) -> list: if not draft_id: self.connector_logger.error("Draft couldn't be created") return [] + if work_id: + self.api.work.add_draft_context(work_id, draft_id) # If directory setup, write the bundle to the target directory if bundle_send_to_directory and bundle_send_to_directory_path is not None: @@ -1744,8 +1746,6 @@ def send_stix2_bundle(self, bundle: str, **kwargs) -> list: if bundle_send_to_queue: if work_id: self.api.work.add_expectations(work_id, expectations_number) - if draft_id: - self.api.work.add_draft_context(work_id, draft_id) if entities_types is None: entities_types = [] if self.queue_protocol == "amqp": diff --git a/pycti/utils/opencti_stix2.py b/pycti/utils/opencti_stix2.py index 2eeb6fcfa..26ece14bb 100644 --- a/pycti/utils/opencti_stix2.py +++ b/pycti/utils/opencti_stix2.py @@ -2431,7 +2431,11 @@ def apply_patch_files(self, item): ) def apply_patch(self, item): - field_patch = item["opencti_field_patch"] + field_patch = self.opencti.get_attribute_in_extension( + "opencti_field_patch", item + ) + if field_patch is None: + field_patch = item["opencti_field_patch"] field_patch_without_files = [ op for op in field_patch if op["key"] != "x_opencti_files" ] @@ -2458,6 +2462,19 @@ def apply_patch(self, item): ) self.apply_patch_files(item) + def apply_opencti_operation(self, item, operation): + if operation == "delete": + delete_id = item["id"] + self.opencti.stix.delete(id=delete_id) + elif operation == "merge": + target_id = item["merge_target_id"] + source_ids = item["merge_source_ids"] + self.opencti.stix.merge(id=target_id, object_ids=source_ids) + elif operation == "patch": + self.apply_patch(item=item) + else: + raise ValueError("Not supported opencti_operation") + def import_item( self, item, @@ -2469,18 +2486,13 @@ def import_item( worker_logger = self.opencti.logger_class("worker") try: self.opencti.set_retry_number(processing_count) - if "opencti_operation" in item: - if item["opencti_operation"] == "delete": - delete_id = item["id"] - self.opencti.stix.delete(id=delete_id) - elif item["opencti_operation"] == "merge": - target_id = item["merge_target_id"] - source_ids = item["merge_source_ids"] - self.opencti.stix.merge(id=target_id, object_ids=source_ids) - elif item["opencti_operation"] == "patch": - self.apply_patch(item=item) - else: - raise ValueError("Not supported opencti_operation") + opencti_operation = self.opencti.get_attribute_in_extension( + "opencti_operation", item + ) + if opencti_operation is not None: + self.apply_opencti_operation(item, opencti_operation) + elif "opencti_operation" in item: + self.apply_opencti_operation(item, item["opencti_operation"]) elif item["type"] == "relationship": # Import relationship self.import_relationship(item, update, types)