[client] move nb refs computation from splitter to static utils method

Author: JeremyCloarec

pycti/connector/opencti_connector_helper.py

@@ -2096,7 +2096,7 @@ def send_stix2_bundle(self, bundle: str, **kwargs) -> list:
             os.rename(write_file, final_write_file)
 
         stix2_splitter = OpenCTIStix2Splitter()
-        (expectations_number, _, bundles, _) = (
+        (expectations_number, _, bundles) = (
             stix2_splitter.split_bundle_with_expectations(
                 bundle=bundle,
                 use_json=True,

pycti/utils/opencti_stix2.py

@@ -32,6 +32,7 @@
     STIX_CORE_OBJECTS,
     STIX_CYBER_OBSERVABLE_MAPPING,
     STIX_META_OBJECTS,
+    OpenCTIStix2Utils,
 )
 
 datefinder.ValueError = ValueError, OverflowError
@@ -3076,8 +3077,8 @@ def import_bundle(
             else None
         )
 
-        stix2_splitter = OpenCTIStix2Splitter(objects_max_refs)
-        _, incompatible_elements, bundles, too_large_elements_bundles = (
+        stix2_splitter = OpenCTIStix2Splitter()
+        _, incompatible_elements, bundles = (
             stix2_splitter.split_bundle_with_expectations(
                 stix_bundle, False, event_version
             )
@@ -3095,23 +3096,27 @@ def import_bundle(
                         + " is incompatible and couldn't be processed",
                     },
                 )
-            for too_large_elements_bundle in too_large_elements_bundles:
-                self.opencti.work.report_expectation(
-                    work_id,
-                    {
-                        "error": "Too large element in bundle",
-                        "source": "Element "
-                        + too_large_elements_bundle["id"]
-                        + " is too large and couldn't be processed",
-                    },
-                )
 
         # Import every element in a specific order
         imported_elements = []
+        too_large_elements_bundles = []
         for bundle in bundles:
             for item in bundle["objects"]:
-                self.import_item(item, update, types, 0, work_id)
-                imported_elements.append({"id": item["id"], "type": item["type"]})
+                nb_refs = OpenCTIStix2Utils.compute_object_refs_number(item)
+                if 0 < objects_max_refs <= nb_refs:
+                    self.opencti.work.report_expectation(
+                        work_id,
+                        {
+                            "error": "Too large element in bundle",
+                            "source": "Element "
+                            + item["id"]
+                            + " is too large and couldn't be processed",
+                        },
+                    )
+                    too_large_elements_bundles.append(item)
+                else:
+                    self.import_item(item, update, types, 0, work_id)
+                    imported_elements.append({"id": item["id"], "type": item["type"]})
 
         return imported_elements, too_large_elements_bundles
 

pycti/utils/opencti_stix2_splitter.py

@@ -33,19 +33,17 @@ def is_id_supported(key):
     return True
 
 
-class OpenCTIStix2Splitter:  # pylint: disable=too-many-instance-attributes
+class OpenCTIStix2Splitter:
     """STIX2 bundle splitter for OpenCTI
 
     Splits large STIX2 bundles into smaller chunks for processing.
     """
 
-    def __init__(self, objects_max_refs: int = 0):
-        self.objects_max_refs = objects_max_refs
+    def __init__(self):
         self.cache_index = {}
         self.cache_refs = {}
         self.elements = []
         self.incompatible_items = []
-        self.too_large_elements = []
 
     def get_internal_ids_in_extension(self, item):
         ids = []
@@ -63,7 +61,6 @@ def enlist_element(
         self, item_id, raw_data, cleanup_inconsistent_bundle, parent_acc
     ):
         nb_deps = 1
-        raw_nb_refs = 0
         if item_id not in raw_data:
             return 0
 
@@ -80,7 +77,6 @@ def enlist_element(
             if key.endswith("_refs") and item[key] is not None:
                 to_keep = []
                 for element_ref in item[key]:
-                    raw_nb_refs += 1
                     # We need to check if this ref is not already a reference
                     is_missing_ref = raw_data.get(element_ref) is None
                     must_be_cleaned = is_missing_ref and cleanup_inconsistent_bundle
@@ -107,7 +103,6 @@ def enlist_element(
                             to_keep.append(element_ref)
                     item[key] = to_keep
             elif key.endswith("_ref"):
-                raw_nb_refs += 1
                 is_missing_ref = raw_data.get(value) is None
                 must_be_cleaned = is_missing_ref and cleanup_inconsistent_bundle
                 not_dependency_ref = (
@@ -134,7 +129,6 @@ def enlist_element(
                     item[key] = None
             # Case for embedded elements (deduplicating and cleanup)
             elif key == "external_references" and item[key] is not None:
-                raw_nb_refs += 1
                 # specific case of splitting external references
                 # reference_ids = []
                 deduplicated_references = []
@@ -161,7 +155,6 @@ def enlist_element(
                         # nb_deps += self.enlist_element(reference_id, raw_data)
                 item[key] = deduplicated_references
             elif key == "kill_chain_phases" and item[key] is not None:
-                raw_nb_refs += 1
                 # specific case of splitting kill_chain phases
                 # kill_chain_ids = []
                 deduplicated_kill_chain = []
@@ -203,9 +196,8 @@ def enlist_element(
                 )
             else:
                 is_compatible = is_id_supported(item_id)
-            if 0 < self.objects_max_refs <= raw_nb_refs:
-                self.too_large_elements.append(item)
-            elif is_compatible:
+
+            if is_compatible:
                 self.elements.append(item)
             else:
                 self.incompatible_items.append(item)
@@ -221,7 +213,7 @@ def split_bundle_with_expectations(
         use_json=True,
         event_version=None,
         cleanup_inconsistent_bundle=False,
-    ) -> Tuple[int, list, list, list]:
+    ) -> Tuple[int, list, list]:
         """splits a valid stix2 bundle into a list of bundles"""
         if use_json:
             try:
@@ -271,28 +263,15 @@ def by_dep_size(elem):
                 )
             )
 
-        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

pycti/utils/opencti_stix2_utils.py

@@ -233,3 +233,17 @@ def retrieveClassForMethod(
                 if hasattr(attribute, method):
                     return attribute
         return None
+
+    @staticmethod
+    def compute_object_refs_number(entity: Dict):
+        refs_number = 0
+        for key in list(entity.keys()):
+            if key.endswith("_refs") and entity[key] is not None:
+                refs_number += len(entity[key])
+            elif key.endswith("_ref"):
+                refs_number += 1
+            elif key == "external_references" and entity[key] is not None:
+                refs_number += len(entity[key])
+            elif key == "kill_chain_phases" and entity[key] is not None:
+                refs_number += len(entity[key])
+        return refs_number

tests/01-unit/utils/test_opencti_stix2_splitter.py

@@ -10,15 +10,15 @@ def test_split_bundle():
     stix_splitter = OpenCTIStix2Splitter()
     with open("./tests/data/enterprise-attack.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 == 7016
 
 
 def test_split_test_bundle():
     stix_splitter = OpenCTIStix2Splitter()
     with open("./tests/data/DATA-TEST-STIX2_v2.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 == 59
     base_bundles = json.loads(content)["objects"]
     for base in base_bundles:
@@ -40,13 +40,13 @@ def test_split_mono_entity_bundle():
     stix_splitter = OpenCTIStix2Splitter()
     with open("./tests/data/mono-bundle-entity.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 == 1
     json_bundle = json.loads(bundles[0])["objects"][0]
     assert json_bundle["created_by_ref"] == "fa42a846-8d90-4e51-bc29-71d5b4802168"
     # 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 == 1
@@ -58,11 +58,11 @@ def test_split_mono_relationship_bundle():
     stix_splitter = OpenCTIStix2Splitter()
     with open("./tests/data/mono-bundle-relationship.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 == 1
     # 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 == 0
@@ -72,19 +72,19 @@ def test_split_capec_bundle():
     stix_splitter = OpenCTIStix2Splitter()
     with open("./tests/data/mitre_att_capec.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 == 2610
 
 
 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
@@ -101,11 +101,11 @@ def test_split_missing_refs_bundle():
     stix_splitter = OpenCTIStix2Splitter()
     with open("./tests/data/missing_refs.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 == 3
@@ -115,7 +115,7 @@ def test_split_cyclic_bundle():
     stix_splitter = OpenCTIStix2Splitter()
     with open("./tests/data/cyclic-bundle.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 == 6
     for bundle in bundles:
         json_bundle = json.loads(bundle)