Merge pull request #632 from OpenCTI-Platform/issue/631

[client-python] upload_file now accepts optional arg file_markings (#631)

Author: labo-flg

examples/upload_file.py

@@ -1,5 +1,7 @@
 # coding: utf-8
 
+from stix2 import TLP_GREEN
+
 from pycti import OpenCTIApiClient
 
 # Variables
@@ -9,8 +11,11 @@
 # OpenCTI initialization
 opencti_api_client = OpenCTIApiClient(api_url, api_token)
 
-# Upload the file
+# Get a marking
+TLP_GREEN_CTI = opencti_api_client.marking_definition.read(id=TLP_GREEN["id"])
+
+# Upload the file (note that markings are optional)
 file = opencti_api_client.upload_file(
-    file_name="./upload_file_example.pdf",
+    file_name="./upload_file_example.pdf", file_markings=[TLP_GREEN["id"]]
 )
 print(file)

pycti/api/opencti_api_client.py

@@ -616,18 +616,19 @@ def upload_file(self, **kwargs):
         """upload a file to OpenCTI API
 
         :param `**kwargs`: arguments for file upload (required: `file_name` and `data`)
-        :return: returns the query respons for the file upload
+        :return: returns the query response for the file upload
         :rtype: dict
         """
 
         file_name = kwargs.get("file_name", None)
+        file_markings = kwargs.get("file_markings", None)
         data = kwargs.get("data", None)
         mime_type = kwargs.get("mime_type", "text/plain")
         if file_name is not None:
             self.app_logger.info("Uploading a file.")
             query = """
-                mutation UploadImport($file: Upload!) {
-                    uploadImport(file: $file) {
+                mutation UploadImport($file: Upload!, $fileMarkings: [String]) {
+                    uploadImport(file: $file, fileMarkings: $fileMarkings) {
                         id
                         name
                     }
@@ -639,8 +640,11 @@ def upload_file(self, **kwargs):
                     mime_type = "application/json"
                 else:
                     mime_type = magic.from_file(file_name, mime=True)
-
-            return self.query(query, {"file": (File(file_name, data, mime_type))})
+            query_vars = {"file": (File(file_name, data, mime_type))}
+            # optional file markings
+            if file_markings is not None:
+                query_vars["fileMarkings"] = file_markings
+            return self.query(query, query_vars)
         else:
             self.app_logger.error("[upload] Missing parameter: file_name")
             return None