[client] handle too large items

Author: JeremyCloarec

pycti/utils/opencti_stix2.py

@@ -8,7 +8,7 @@
 import time
 import traceback
 import uuid
-from typing import Any, Dict, List, Optional, Union
+from typing import Any, Dict, List, Optional, Union, Tuple
 
 import datefinder
 import dateutil.parser
@@ -196,7 +196,7 @@ def import_bundle_from_file(
         file_path: str,
         update: bool = False,
         types: List = None,
-    ) -> Optional[List]:
+    ) -> Optional[Tuple[list, list]]:
         """import a stix2 bundle from a file
 
         :param file_path: valid path to the file
@@ -221,7 +221,8 @@ def import_bundle_from_json(
         update: bool = False,
         types: List = None,
         work_id: str = None,
-    ) -> List:
+        objects_max_deps: int = 0,
+    ) -> Tuple[list, list]:
         """import a stix2 bundle from JSON data
 
         :param json_data: JSON data
@@ -231,11 +232,13 @@ def import_bundle_from_json(
         :param types: list of stix2 types, defaults to None
         :type types: list, optional
         :param work_id work_id: str, optional
-        :return: list of imported stix2 objects
-        :rtype: List
+        :param objects_max_deps: max deps amount of objects, reject object import if larger than configured amount
+        :type objects_max_deps: int, optional
+        :return: list of imported stix2 objects and a list of stix2 objects with too many deps
+        :rtype: Tuple[List,List]
         """
         data = json.loads(json_data)
-        return self.import_bundle(data, update, types, work_id)
+        return self.import_bundle(data, update, types, work_id, objects_max_deps)
 
     def resolve_author(self, title: str) -> Optional[Identity]:
         if "fireeye" in title.lower() or "mandiant" in title.lower():
@@ -3054,7 +3057,8 @@ def import_bundle(
         update: bool = False,
         types: List = None,
         work_id: str = None,
-    ) -> List:
+        objects_max_deps: int = 0,
+    ) -> Tuple[list, list]:
         # Check if the bundle is correctly formatted
         if "type" not in stix_bundle or stix_bundle["type"] != "bundle":
             raise ValueError("JSON data type is not a STIX2 bundle")
@@ -3066,8 +3070,8 @@ def import_bundle(
             else None
         )
 
-        stix2_splitter = OpenCTIStix2Splitter()
-        _, incompatible_elements, bundles = (
+        stix2_splitter = OpenCTIStix2Splitter(objects_max_deps)
+        _, incompatible_elements, bundles, too_large_elements_bundles = (
             stix2_splitter.split_bundle_with_expectations(
                 stix_bundle, False, event_version
             )
@@ -3093,7 +3097,7 @@ def import_bundle(
                 self.import_item(item, update, types, 0, work_id)
                 imported_elements.append({"id": item["id"], "type": item["type"]})
 
-        return imported_elements
+        return imported_elements, too_large_elements_bundles
 
     @staticmethod
     def put_attribute_in_extension(

pycti/utils/opencti_stix2_splitter.py

@@ -33,17 +33,19 @@ def is_id_supported(key):
     return True
 
 
-class OpenCTIStix2Splitter:
+class OpenCTIStix2Splitter:  # pylint: disable=too-many-instance-attributes
     """STIX2 bundle splitter for OpenCTI
 
     Splits large STIX2 bundles into smaller chunks for processing.
     """
 
-    def __init__(self):
+    def __init__(self, objects_max_deps: int = 0):
+        self.objects_max_deps = objects_max_deps
         self.cache_index = {}
         self.cache_refs = {}
         self.elements = []
         self.incompatible_items = []
+        self.too_large_elements = []
 
     def get_internal_ids_in_extension(self, item):
         ids = []
@@ -196,7 +198,9 @@ def enlist_element(
                 )
             else:
                 is_compatible = is_id_supported(item_id)
-            if is_compatible:
+            if self.objects_max_deps is not 0 and nb_deps >= self.objects_max_deps:
+                self.too_large_elements.append(item)
+            elif is_compatible:
                 self.elements.append(item)
             else:
                 self.incompatible_items.append(item)
@@ -212,7 +216,7 @@ def split_bundle_with_expectations(
         use_json=True,
         event_version=None,
         cleanup_inconsistent_bundle=False,
-    ) -> Tuple[int, list, list]:
+    ) -> Tuple[int, list, list, list]:
         """splits a valid stix2 bundle into a list of bundles"""
         if use_json:
             try:
@@ -262,11 +266,28 @@ def by_dep_size(elem):
                 )
             )
 
-        return number_expectations, self.incompatible_items, bundles
+        too_large_elements_bundles = []
+        for too_large_element in self.too_large_elements:
+            too_large_elements_bundles.append(
+                self.stix2_create_bundle(
+                    bundle_data["id"],
+                    too_large_element["nb_deps"],
+                    [too_large_element],
+                    use_json,
+                    event_version,
+                )
+            )
+
+        return (
+            number_expectations,
+            self.incompatible_items,
+            bundles,
+            too_large_elements_bundles,
+        )
 
     @deprecated("Use split_bundle_with_expectations instead")
     def split_bundle(self, bundle, use_json=True, event_version=None) -> list:
-        _, _, bundles = self.split_bundle_with_expectations(
+        _, _, bundles, _ = self.split_bundle_with_expectations(
             bundle, use_json, event_version
         )
         return bundles

tests/01-unit/utils/test_opencti_stix2_splitter.py

@@ -80,11 +80,11 @@ def test_split_internal_ids_bundle():
     stix_splitter = OpenCTIStix2Splitter()
     with open("./tests/data/bundle_with_internal_ids.json") as file:
         content = file.read()
-    expectations, _, bundles = stix_splitter.split_bundle_with_expectations(content)
+    expectations, _, bundles, _ = stix_splitter.split_bundle_with_expectations(content)
     assert expectations == 4
     # Split with cleanup_inconsistent_bundle
     stix_splitter = OpenCTIStix2Splitter()
-    expectations, _, bundles = stix_splitter.split_bundle_with_expectations(
+    expectations, _, bundles, _ = stix_splitter.split_bundle_with_expectations(
         bundle=content, cleanup_inconsistent_bundle=True
     )
     assert expectations == 4

tests/02-integration/entities/test_malware.py

@@ -5,7 +5,7 @@ def test_malware_import_with_sample_refs(api_client):
     with open("tests/data/basicMalwareWithSample.json", "r") as content_file:
         content = content_file.read()
 
-    imported_malware_bundle = api_client.stix2.import_bundle_from_json(
+    imported_malware_bundle, _ = api_client.stix2.import_bundle_from_json(
         json_data=content
     )
     assert imported_malware_bundle is not None

tests/02-integration/utils/test_stix_crud.py

@@ -21,7 +21,7 @@ def test_entity_create(entity_class, api_stix, opencti_splitter):
     stix_object = stix_class(**class_data)
     bundle = Bundle(objects=[stix_object]).serialize()
     split_bundle = opencti_splitter.split_bundle(bundle, True, None)[0]
-    bundles_sent = api_stix.import_bundle_from_json(split_bundle, False, None, None)
+    bundles_sent, _ = api_stix.import_bundle_from_json(split_bundle, False, None, None)
 
     assert len(bundles_sent) == 1
     assert bundles_sent[0]["id"] == stix_object["id"]