diff --git a/pycti/api/opencti_api_work.py b/pycti/api/opencti_api_work.py index 305811721..4ef5dde86 100644 --- a/pycti/api/opencti_api_work.py +++ b/pycti/api/opencti_api_work.py @@ -82,6 +82,24 @@ def add_expectations(self, work_id: str, expectations: int): except: self.api.app_logger.error("Cannot report expectation") + def add_draft_context(self, work_id: str, draft_context: str): + if self.api.bundle_send_to_queue: + self.api.app_logger.info( + "Update draft context", + {"work_id": work_id, "draft_context": draft_context}, + ) + query = """ + mutation addDraftContext($id: ID!, $draftContext: String) { + workEdit(id: $id) { + addDraftContext(draftContext: $draftContext) + } + } + """ + try: + self.api.query(query, {"id": work_id, "draftContext": draft_context}) + except: + self.api.app_logger.error("Cannot report draft context") + def initiate_work(self, connector_id: str, friendly_name: str) -> str: if self.api.bundle_send_to_queue: self.api.app_logger.info("Initiate work", {"connector_id": connector_id}) diff --git a/pycti/connector/opencti_connector_helper.py b/pycti/connector/opencti_connector_helper.py index 1c125df43..2db740349 100644 --- a/pycti/connector/opencti_connector_helper.py +++ b/pycti/connector/opencti_connector_helper.py @@ -1732,6 +1732,8 @@ 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":