[client] Add helper for enrichment connectors

Author: SamuelHassine

pycti/connector/opencti_connector_helper.py

@@ -1,5 +1,6 @@
 import asyncio
 import base64
+import copy
 import datetime
 import json
 import logging
@@ -1310,3 +1311,20 @@ def get_attribute_in_mitre_extension(key, object) -> any:
                 "extension-definition--322b8f77-262a-4cb8-a915-1e441e00329b"
             ][key]
         return None
+
+    def get_data_from_enrichment(self, data, opencti_entity):
+        stix_id = data["entity_id"]
+        bundle = data.get("bundle", None)
+        # Extract IPv4, IPv6 and Domain from entity data
+        if bundle is None:
+            # Generate bundle
+            stix_objects = self.api.stix2.prepare_export(
+                self.api.stix2.generate_export(copy.copy(opencti_entity))
+            )
+        else:
+            stix_objects = bundle["objects"]
+        stix_entity = [e for e in stix_objects if e["id"] == stix_id][0]
+        return {
+            "stix_entity": stix_entity,
+            "stix_objects": stix_objects,
+        }

pycti/utils/opencti_stix2.py

@@ -2412,6 +2412,10 @@ def import_bundle(
     def put_attribute_in_extension(
         object, extension_id, key, value, multiple=False
     ) -> any:
+        if ("x_opencti_" + key) in object:
+            del object["x_opencti_" + key]
+        if ("x_mitre_" + key) in object:
+            del object["x_mitre_" + key]
         if "extensions" not in object:
             object["extensions"] = {}
         if extension_id not in object["extensions"]: