[client] Improve relationship standard id generation (#659)

Author: richard-julien

pycti/entities/opencti_stix_core_relationship.py

@@ -336,6 +336,7 @@ def generate_id(
             start_time = start_time.isoformat()
         if isinstance(stop_time, datetime.datetime):
             stop_time = stop_time.isoformat()
+
         if start_time is not None and stop_time is not None:
             data = {
                 "relationship_type": relationship_type,
@@ -376,6 +377,16 @@ def generate_id(
         :return List of stix_core_relationship objects
     """
 
+    @staticmethod
+    def generate_id_from_data(data):
+        return StixCoreRelationship.generate_id(
+            data["relationship_type"],
+            data["source_ref"],
+            data["target_ref"],
+            data.get("start_time"),
+            data.get("stop_time"),
+        )
+
     def list(self, **kwargs):
         from_or_to_id = kwargs.get("fromOrToId", None)
         element_with_target_types = kwargs.get("elementWithTargetTypes", None)

pycti/entities/opencti_stix_sighting_relationship.py

@@ -262,28 +262,53 @@ def __init__(self, opencti):
         """
 
     @staticmethod
-    def generate_id(source_ref, target_ref, first_seen=None, last_seen=None):
+    def generate_id(
+        relationship_type,
+        sighting_of_ref,
+        where_sighted_refs,
+        first_seen=None,
+        last_seen=None,
+    ):
         if isinstance(first_seen, datetime.datetime):
             first_seen = first_seen.isoformat()
         if isinstance(last_seen, datetime.datetime):
             last_seen = last_seen.isoformat()
 
         if first_seen is not None and last_seen is not None:
             data = {
-                "source_ref": source_ref,
-                "target_ref": target_ref,
+                "type": relationship_type,
+                "sighting_of_ref": sighting_of_ref,
+                "where_sighted_refs": where_sighted_refs,
                 "first_seen": first_seen,
                 "last_seen": last_seen,
             }
+        elif first_seen is not None:
+            data = {
+                "type": relationship_type,
+                "sighting_of_ref": sighting_of_ref,
+                "where_sighted_refs": where_sighted_refs,
+                "first_seen": first_seen,
+            }
         else:
             data = {
-                "source_ref": source_ref,
-                "target_ref": target_ref,
+                "type": relationship_type,
+                "sighting_of_ref": sighting_of_ref,
+                "where_sighted_refs": where_sighted_refs,
             }
         data = canonicalize(data, utf8=False)
         id = str(uuid.uuid5(uuid.UUID("00abedb4-aa42-466c-9c01-fed23315a9b7"), data))
         return "sighting--" + id
 
+    @staticmethod
+    def generate_id_from_data(data):
+        return StixSightingRelationship.generate_id(
+            data["type"],
+            data["sighting_of_ref"],
+            data["where_sighted_refs"],
+            data.get("first_seen"),
+            data.get("last_seen"),
+        )
+
     """
         List stix_sightings objects
 

pycti/utils/opencti_stix2.py

@@ -855,6 +855,7 @@ def get_reader(self, entity_type: str):
     def get_stix_helper(self):
         # Import
         return {
+            # entities
             "attack-pattern": self.opencti.attack_pattern,
             "campaign": self.opencti.campaign,
             "note": self.opencti.note,
@@ -894,6 +895,9 @@ def get_stix_helper(self):
             "task": self.opencti.task,
             "x-opencti-task": self.opencti.task,
             "vocabulary": self.opencti.vocabulary,
+            # relationships
+            "relationship": self.opencti.stix_core_relationship,
+            "sighting": self.opencti.stix_sighting_relationship,
         }
 
     def generate_standard_id_from_stix(self, data):
@@ -2399,23 +2403,21 @@ def prepare_bundle_ids(self, bundle, use_json=True, keep_original_id=False):
         # First iteration to cache all entity ids
         stix_helpers = self.get_stix_helper()
         for item in bundle_data["objects"]:
-            if item["type"] != "relationship" and item["type"] != "sighting":
-                helper = stix_helpers.get(item["type"])
-                if hasattr(helper, "generate_id_from_data"):
-                    standard_id = helper.generate_id_from_data(item)
-                    cache_ids[item["id"]] = standard_id
+            helper = stix_helpers.get(item["type"])
+            if hasattr(helper, "generate_id_from_data"):
+                standard_id = helper.generate_id_from_data(item)
+                cache_ids[item["id"]] = standard_id
         # Second iteration to replace and remap
         for item in bundle_data["objects"]:
             # For entities, try to replace the main id
             # Keep the current one if needed
-            if item["type"] != "relationship" and item["type"] != "sighting":
-                if cache_ids.get(item["id"]):
-                    original_id = item["id"]
-                    item["id"] = cache_ids[original_id]
-                    if keep_original_id:
-                        item["x_opencti_stix_ids"] = item.get(
-                            "x_opencti_stix_ids", []
-                        ) + [original_id]
+            if cache_ids.get(item["id"]):
+                original_id = item["id"]
+                item["id"] = cache_ids[original_id]
+                if keep_original_id:
+                    item["x_opencti_stix_ids"] = item.get("x_opencti_stix_ids", []) + [
+                        original_id
+                    ]
             # For all elements, replace all refs (source_ref, object_refs, ...)
             ref_keys = list(
                 filter(lambda i: i.endswith("_ref") or i.endswith("_refs"), item.keys())

tests/01-unit/stix/test_bundle_ids_rewrite.py

@@ -17,6 +17,7 @@ def load_test_file():
     return bundle_data
 
 
+# !! WARNING !!, this need to be changed along with 01-unit/domain/identifier-test.js
 # fmt: off
 def test_ids_generation():
     gen_id = get_cti_helper().generate_standard_id_from_stix
@@ -95,6 +96,16 @@ def test_ids_generation():
     assert gen_id({"type": "threat-actor", "name": "CARD04", "x_opencti_type": "Threat-Actor-Individual"}) == "threat-actor--af15b6ae-a3dd-54d3-8fa0-3adfe0391d01"
     # vocabulary
     assert gen_id({"type": "vocabulary", "name": "facebook", "category": "account_type_ov"}) == "vocabulary--85ae7185-ff6f-509b-a011-3069921614aa"
+    # relationship
+    base_relationship = {"type": "relationship", "relationship_type": "based-on", "source_ref": "from_id", "target_ref": "to_id"}
+    assert gen_id(base_relationship) == "relationship--0b11fa67-da01-5d34-9864-67d4d71c3740"
+    assert gen_id({**base_relationship, "start_time": "2022-11-25T19:00:05.000Z"}) == "relationship--c5e1e2ce-14d6-535b-911d-267e92119e01"
+    assert gen_id({**base_relationship, "start_time": "2022-11-25T19:00:05.000Z", "stop_time": "2022-11-26T19:00:05.000Z"}) == "relationship--a7778a7d-a743-5193-9912-89f88f9ed0b4"
+    # sighting
+    base_sighting = {"type": "sighting", "sighting_of_ref": "from_id", "where_sighted_refs": ["to_id"]}
+    assert gen_id(base_sighting) == 'sighting--161901df-21bb-527a-b96b-354119279fe2'
+    assert gen_id({**base_sighting, "first_seen": "2022-11-25T19:00:05.000Z"}) == "sighting--3c59ceea-8e41-5adb-a257-d070d19e6d2b"
+    assert gen_id({**base_sighting, "first_seen": "2022-11-25T19:00:05.000Z", "last_seen": "2022-11-26T19:00:05.000Z"}) == "sighting--b4d307b6-d22c-5f22-b530-876c298493da"
 # fmt: on
 
 

tests/data/bundle_ids_sample.json

@@ -59,7 +59,7 @@
       ]
     },
     {
-      "id": "relationship--ba52fced-422a-4bee-816a-85aa21c9eaca",
+      "id": "relationship--ba52fced-422a-4bee-816a-85aa21c9eacc",
       "type": "relationship",
       "spec_version": "2.1",
       "relationship_type": "related-to",
@@ -71,6 +71,19 @@
       "stop_time": "1900-01-01T00:00:00.000Z",
       "created_by_ref": "identity--7b82b010-b1c0-4dae-981f-7756374a17da",
       "object_marking_refs": ["marking-definition--78ca4366-f5b8-4764-83f7-34ce38198e27"]
+    },
+    {
+      "type": "sighting",
+      "spec_version": "2.1",
+      "id": "sighting--ee20065d-2555-424f-ad9e-0f8428623c75",
+      "created": "2016-08-06T20:08:31.000Z",
+      "modified": "2016-09-06T20:08:31.000Z",
+      "sighting_of_ref": "malware--d650c5b9-4b43-5781-8576-ea52bd6c7ce5",
+      "where_sighted_refs": ["identity--7b82b010-b1c0-4dae-981f-7756374a17da"],
+      "first_seen": "2016-08-06T20:08:31.000Z",
+      "last_seen": "2016-08-07T20:08:31.000Z",
+      "count": 12,
+      "x_opencti_negative": true
     }
   ]
 }