[client] add draft context to work when sending bundle in draft (opencti #6577)

Author: JeremyCloarec

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})

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":