[client] Add a custom STIX field for file upload from bundle (#51)

Author: Samuel Hassine

examples/upload_file_to_intrusion_set.py

@@ -23,7 +23,7 @@
 print(intrusion_set)
 
 # Upload the file
-file = opencti_api_client.stix_domain_entity.upload_file(
+file = opencti_api_client.stix_domain_entity.add_file(
     id=intrusion_set["id"], file_name="./2005_002_001_14428.pdf",
 )
 print(file)

pycti/api/opencti_api_client.py

@@ -333,6 +333,9 @@ def process_multiple_fields(self, data):
         if "indicators" in data:
             data["indicators"] = self.process_multiple(data["indicators"])
             data["indicatorsIds"] = self.process_multiple_ids(data["indicators"])
+        if "importFiles" in data:
+            data["importFiles"] = self.process_multiple(data["importFiles"])
+            data["importFilesIds"] = self.process_multiple_ids(data["importFiles"])
         return data
 
     def upload_file(self, **kwargs):

pycti/entities/opencti_stix_domain_entity.py

@@ -89,6 +89,15 @@ def __init__(self, opencti, file):
                     }
                 }
             }
+            importFiles {
+                edges {
+                    node {
+                        id
+                        name
+                        size
+                    }
+                }
+            }
             ... on AttackPattern {
                 platform
                 required_permission
@@ -486,32 +495,45 @@ def delete(self, **kwargs):
         :return void
     """
 
-    def upload_file(self, **kwargs):
+    def add_file(self, **kwargs):
         id = kwargs.get("id", None)
         file_name = kwargs.get("file_name", None)
         data = kwargs.get("data", None)
         mime_type = kwargs.get("mime_type", "text/plain")
         if id is not None and file_name is not None:
-            self.opencti.log(
-                "info", "Uploading a file in Stix-Domain-Entity {" + id + "}."
-            )
-            query = """
-                mutation StixDomainEntityEdit($id: ID!, $file: Upload!) {
-                    stixDomainEntityEdit(id: $id) {
-                        importPush(file: $file) {
-                            id
-                            name
+            stix_domain_entity = self.read(id=id)
+            if stix_domain_entity is None:
+                self.opencti.log(
+                    "error", "Cannot add File, entity not found"
+                )
+                return False
+            final_file_name = os.path.basename(file_name)
+            current_files = {}
+            for file in stix_domain_entity['importFiles']:
+                current_files[file['name']] = file
+            if final_file_name in current_files:
+                return current_files[final_file_name]
+            else:
+                self.opencti.log(
+                    "info", "Uploading a file in Stix-Domain-Entity {" + id + "}."
+                )
+                query = """
+                    mutation StixDomainEntityEdit($id: ID!, $file: Upload!) {
+                        stixDomainEntityEdit(id: $id) {
+                            importPush(file: $file) {
+                                id
+                                name
+                            }
                         }
                     }
-                }
-             """
-            if data is None:
-                data = open(file_name, "rb")
-                mime_type = magic.from_file(file_name, mime=True)
+                 """
+                if data is None:
+                    data = open(file_name, "rb")
+                    mime_type = magic.from_file(file_name, mime=True)
 
-            return self.opencti.query(
-                query, {"id": id, "file": (self.file(file_name, data, mime_type))}
-            )
+                return self.opencti.query(
+                    query, {"id": id, "file": (self.file(final_file_name, data, mime_type))}
+                )
         else:
             self.opencti.log(
                 "error",

pycti/utils/constants.py

@@ -62,6 +62,9 @@ class CustomProperties:
     # internal id used by OpenCTI - this will be auto generated
     ID = "x_opencti_id"
 
+    # List of files
+    FILES = "x_opencti_files"
+
     # This should be set on all reports to one of the following values:
     #  "external"
     #  "internal"

pycti/utils/opencti_stix2.py

@@ -4,6 +4,7 @@
 import os
 import json
 import uuid
+import base64
 import datetime
 from typing import List
 
@@ -175,6 +176,12 @@ def extract_embedded_relationships(self, stix_object, types=None):
                     tag_result = self.opencti.tag.read(id=tag["id"])
                     if tag_result is not None:
                         self.mapping_cache[tag["id"]] = {"id": tag_result["id"]}
+                    else:
+                        tag_result = self.opencti.tag.create(
+                            tag_type=tag["tag_type"],
+                            value=tag["value"],
+                            color=tag["color"]
+                        )
                 if tag_result is not None:
                     tags_ids.append(tag_result["id"])
 
@@ -463,7 +470,7 @@ def import_object(self, stix_object, update=False, types=None):
             # Add tags
             for tag_id in tags_ids:
                 self.opencti.stix_entity.add_tag(
-                    id=stix_object_result["id"], marking_definition_id=tag_id,
+                    id=stix_object_result["id"], tag_id=tag_id,
                 )
             # Add external references
             for external_reference_id in external_references_ids:
@@ -501,6 +508,14 @@ def import_object(self, stix_object, update=False, types=None):
                             id=stix_object_result["id"],
                             stix_observable_id=observable_ref["id"],
                         )
+            # Add files
+            if CustomProperties.FILES in stix_object:
+                for file in stix_object[CustomProperties.FILES]:
+                    self.opencti.stix_domain_entity.add_file(
+                        file_name=file['name'],
+                        data=base64.b64decode(file['data']),
+                        mime_type=file['mime_type']
+                    )
 
         return stix_object_results