From e46b16d12ae92615fb8c33b9c21c7dad8e8f6cf9 Mon Sep 17 00:00:00 2001
From: Julien Richard <julien.richard@filigran.io>
Date: Sat, 8 Feb 2025 01:39:37 +0100
Subject: [PATCH] [client] Add markings support for workbench

---
 pycti/api/opencti_api_client.py             | 11 ++++++++---
 pycti/connector/opencti_connector_helper.py |  2 ++
 2 files changed, 10 insertions(+), 3 deletions(-)

diff --git a/pycti/api/opencti_api_client.py b/pycti/api/opencti_api_client.py
index 022649a01..1853c5795 100644
--- a/pycti/api/opencti_api_client.py
+++ b/pycti/api/opencti_api_client.py
@@ -712,12 +712,13 @@ def upload_pending_file(self, **kwargs):
         data = kwargs.get("data", None)
         mime_type = kwargs.get("mime_type", "text/plain")
         entity_id = kwargs.get("entity_id", None)
+        file_markings = kwargs.get("file_markings", [])
 
         if file_name is not None:
             self.app_logger.info("Uploading a file.")
             query = """
-                    mutation UploadPending($file: Upload!, $entityId: String) {
-                        uploadPending(file: $file, entityId: $entityId) {
+                    mutation UploadPending($file: Upload!, $entityId: String, $file_markings: [String!]) {
+                        uploadPending(file: $file, entityId: $entityId, file_markings: $file_markings) {
                             id
                             name
                         }
@@ -731,7 +732,11 @@ def upload_pending_file(self, **kwargs):
                     mime_type = magic.from_file(file_name, mime=True)
             return self.query(
                 query,
-                {"file": (File(file_name, data, mime_type)), "entityId": entity_id},
+                {
+                    "file": (File(file_name, data, mime_type)),
+                    "entityId": entity_id,
+                    "file_markings": file_markings,
+                },
             )
         else:
             self.app_logger.error("[upload] Missing parameter: file_name")
diff --git a/pycti/connector/opencti_connector_helper.py b/pycti/connector/opencti_connector_helper.py
index 2db740349..353451aac 100644
--- a/pycti/connector/opencti_connector_helper.py
+++ b/pycti/connector/opencti_connector_helper.py
@@ -1581,6 +1581,7 @@ def send_stix2_bundle(self, bundle: str, **kwargs) -> list:
         event_version = kwargs.get("event_version", None)
         bypass_validation = kwargs.get("bypass_validation", False)
         entity_id = kwargs.get("entity_id", None)
+        file_markings = kwargs.get("file_markings", None)
         file_name = kwargs.get("file_name", None)
         bundle_send_to_queue = kwargs.get("send_to_queue", self.bundle_send_to_queue)
         cleanup_inconsistent_bundle = kwargs.get("cleanup_inconsistent_bundle", False)
@@ -1648,6 +1649,7 @@ def send_stix2_bundle(self, bundle: str, **kwargs) -> list:
                     data=bundle,
                     mime_type="application/json",
                     entity_id=entity_id,
+                    file_markings=file_markings,
                 )
                 return []
             elif validation_mode == "draft" and not draft_id: