[client] Improve client to support playbook looping control (opencti/#6930)

OpenCTI-Platform/opencti#6930

Author: richard-julien

pycti/api/opencti_api_client.py

@@ -207,6 +207,12 @@ def __init__(
     def set_applicant_id_header(self, applicant_id):
         self.request_headers["opencti-applicant-id"] = applicant_id
 
+    def set_playbook_id_header(self, playbook_id):
+        self.request_headers["opencti-playbook-id"] = playbook_id
+
+    def set_event_id(self, event_id):
+        self.request_headers["opencti-event-id"] = event_id
+
     def set_synchronized_upsert_header(self, synchronized):
         self.request_headers["synchronized-upsert"] = (
             "true" if synchronized is True else "false"

pycti/api/opencti_api_playbook.py

@@ -9,14 +9,15 @@ def playbook_step_execution(self, playbook: dict, bundle: str):
             "Executing playbook step", {"playbook_id": playbook["playbook_id"]}
         )
         query = """
-            mutation PlaybookStepExecution($execution_id: ID!, $execution_start: DateTime!, $data_instance_id: ID!, $playbook_id: ID!, $previous_step_id: ID!, $step_id: ID!, $previous_bundle: String!, $bundle: String!) {
-                playbookStepExecution(execution_id: $execution_id, execution_start: $execution_start, data_instance_id: $data_instance_id, playbook_id: $playbook_id, previous_step_id: $previous_step_id, step_id: $step_id, previous_bundle: $previous_bundle, bundle: $bundle)
+            mutation PlaybookStepExecution($execution_id: ID!, $event_id: ID!, $execution_start: DateTime!, $data_instance_id: ID!, $playbook_id: ID!, $previous_step_id: ID!, $step_id: ID!, $previous_bundle: String!, $bundle: String!) {
+                playbookStepExecution(execution_id: $execution_id, event_id: $event_id, execution_start: $execution_start, data_instance_id: $data_instance_id, playbook_id: $playbook_id, previous_step_id: $previous_step_id, step_id: $step_id, previous_bundle: $previous_bundle, bundle: $bundle)
             }
            """
         self.api.query(
             query,
             {
                 "execution_id": playbook["execution_id"],
+                "event_id": playbook["event_id"],
                 "execution_start": playbook["execution_start"],
                 "playbook_id": playbook["playbook_id"],
                 "data_instance_id": playbook["data_instance_id"],

pycti/connector/opencti_connector_helper.py

@@ -282,18 +282,20 @@ def _data_handler(self, json_data) -> None:
                 is_playbook = "playbook" in json_data["internal"]
                 # If playbook, compute object on data bundle
                 if is_playbook:
-                    execution_id = json_data["internal"]["playbook"]["execution_id"]
                     execution_start = self.helper.date_now()
-                    playbook_id = json_data["internal"]["playbook"]["playbook_id"]
-                    data_instance_id = json_data["internal"]["playbook"][
+                    event_id = json_data["internal"]["playbook"].get("event_id")
+                    execution_id = json_data["internal"]["playbook"].get("execution_id")
+                    playbook_id = json_data["internal"]["playbook"].get("playbook_id")
+                    data_instance_id = json_data["internal"]["playbook"].get(
                         "data_instance_id"
-                    ]
+                    )
                     previous_bundle = json.dumps((json_data["event"]["bundle"]))
                     step_id = json_data["internal"]["playbook"]["step_id"]
                     previous_step_id = json_data["internal"]["playbook"][
                         "previous_step_id"
                     ]
                     playbook_data = {
+                        "event_id": event_id,
                         "execution_id": execution_id,
                         "execution_start": execution_start,
                         "playbook_id": playbook_id,