[client] Adapt attack pattern and relationship stix ids

Author: richard-julien

pycti/entities/opencti_attack_pattern.py

@@ -232,7 +232,8 @@ def generate_id(name, x_mitre_id=None):
 
     @staticmethod
     def generate_id_from_data(data):
-        return AttackPattern.generate_id(data.get("name"), data.get("x_mitre_id"))
+        external_id = data.get("x_mitre_id") or data.get("x_opencti_external_id")
+        return AttackPattern.generate_id(data.get("name"), external_id)
 
     """
         List Attack-Pattern objects

pycti/entities/opencti_stix_core_relationship.py

@@ -610,6 +610,7 @@ def create(self, **kwargs):
         kill_chain_phases = kwargs.get("killChainPhases", None)
         granted_refs = kwargs.get("objectOrganization", None)
         x_opencti_workflow_id = kwargs.get("x_opencti_workflow_id", None)
+        x_opencti_stix_ids = kwargs.get("x_opencti_stix_ids", None)
         update = kwargs.get("update", False)
 
         self.opencti.app_logger.info(
@@ -653,6 +654,7 @@ def create(self, **kwargs):
                     "externalReferences": external_references,
                     "killChainPhases": kill_chain_phases,
                     "x_opencti_workflow_id": x_opencti_workflow_id,
+                    "x_opencti_stix_ids": x_opencti_stix_ids,
                     "update": update,
                 }
             },
@@ -1143,6 +1145,10 @@ def import_from_stix2(self, **kwargs):
         default_date = kwargs.get("defaultDate", False)
         if stix_relation is not None:
             # Search in extensions
+            if "x_opencti_stix_ids" not in stix_relation:
+                stix_relation["x_opencti_stix_ids"] = (
+                    self.opencti.get_attribute_in_extension("stix_ids", stix_relation)
+                )
             if "x_opencti_granted_refs" not in stix_relation:
                 stix_relation["x_opencti_granted_refs"] = (
                     self.opencti.get_attribute_in_extension(
@@ -1224,6 +1230,11 @@ def import_from_stix2(self, **kwargs):
                     if "x_opencti_workflow_id" in stix_relation
                     else None
                 ),
+                x_opencti_stix_ids=(
+                    stix_relation["x_opencti_stix_ids"]
+                    if "x_opencti_stix_ids" in stix_relation
+                    else None
+                ),
                 update=update,
             )
         else:

pycti/utils/opencti_stix2.py

@@ -2412,18 +2412,31 @@ def prepare_bundle_ids(self, bundle, use_json=True, keep_original_id=False):
             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
+                if item["id"] != standard_id:
+                    cache_ids[item["id"]] = standard_id
+
+        # If no id to rewrite, return the processed bundle
+        if not bool(cache_ids):
+            # As rewrite can have multiple roundtrip
+            # we kept the first original id in a meta attribute
+            # final behavior is to rewrite this in the x_opencti_stix_ids
+            for item in bundle_data["objects"]:
+                if "x_keep_original_id" in item:
+                    item["x_opencti_stix_ids"] = item.get("x_opencti_stix_ids", []) + [
+                        item["x_keep_original_id"]
+                    ]
+                    del item["x_keep_original_id"]
+            return json.dumps(bundle_data) if use_json else bundle_data
+
         # 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 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 keep_original_id and "x_keep_original_id" not in item:
+                    item["x_keep_original_id"] = 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())
@@ -2436,7 +2449,9 @@ def prepare_bundle_ids(self, bundle, use_json=True, keep_original_id=False):
                 else:
                     item[ref_key] = cache_ids.get(item[ref_key], item[ref_key])
 
-        return json.dumps(bundle_data) if use_json else bundle_data
+        return self.prepare_bundle_ids(
+            bundle_data, False, keep_original_id=keep_original_id
+        )
 
     def import_item(
         self,

tests/01-unit/stix/test_bundle_ids_rewrite.py

@@ -23,7 +23,8 @@ def test_ids_generation():
     gen_id = get_cti_helper().generate_standard_id_from_stix
     # attack-pattern
     assert gen_id({"type": "attack-pattern", "name": "attack"}) =='attack-pattern--25f21617-8de8-5d5e-8cd4-b7e88547ba76'
-    assert gen_id({"type": "attack-pattern", "name": "attack", "x_mitre_id": 'MITREID'}) == 'attack-pattern--b74cfee2-7b14-585e-862f-fea45e802da9'
+    assert gen_id({"type": "attack-pattern", "name": "attack", "x_opencti_external_id": 'MITREID'}) == 'attack-pattern--b74cfee2-7b14-585e-862f-fea45e802da9'
+    assert gen_id({"type": "attack-pattern", "name": "Spear phishing messages with malicious links", "x_mitre_id": 'T1368'}) == 'attack-pattern--a01046cc-192f-5d52-8e75-6e447fae3890'
     assert gen_id({"type": "attack-pattern", "x_mitre_id": "MITREID"}) == 'attack-pattern--b74cfee2-7b14-585e-862f-fea45e802da9'
     # campaign
     assert gen_id({"type": "campaign", "name": "attack"}) == 'campaign--25f21617-8de8-5d5e-8cd4-b7e88547ba76'
@@ -101,6 +102,7 @@ def test_ids_generation():
     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"
+    assert gen_id({"type": "relationship", 'relationship_type': 'uses', 'source_ref': 'malware--21c45dbe-54ec-5bb7-b8cd-9f27cc518714', 'start_time': '2020-02-29T22:30:00.000Z', 'stop_time': '2020-02-29T22:30:00.000Z', 'target_ref': 'attack-pattern--fd8179dd-1632-5ec8-8b93-d2ae121e05a4'}) == 'relationship--67f5f01f-6b15-5154-ae31-019a75fedcff'
     # 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'
@@ -112,18 +114,23 @@ def test_ids_generation():
 def test_prepare_bundle_ids_keep_original():
     helper = get_cti_helper()
     bundle_data = load_test_file()
-    malware_source = bundle_data["objects"][0]
-    assert malware_source["id"] == "malware--d650c5b9-4b43-5781-8576-ea52bd6c7ce5"
-    assert malware_source.get("x_opencti_stix_ids") is None
+    # malware_source = bundle_data["objects"][0]
+    # assert malware_source["id"] == "malware--d650c5b9-4b43-5781-8576-ea52bd6c7ce5"
+    # assert malware_source.get("x_opencti_stix_ids") is None
     prepared_bundle = helper.prepare_bundle_ids(
         bundle=bundle_data, use_json=False, keep_original_id=True
     )
-    print(json.dumps(prepared_bundle))
     malware_target = prepared_bundle["objects"][0]
     assert malware_target["id"] == "malware--d650c5b9-4b43-5781-8576-ea52bd6c7ce0"
     assert malware_target.get("x_opencti_stix_ids") == [
         "malware--d650c5b9-4b43-5781-8576-ea52bd6c7ce5"
     ]
+    sighting = prepared_bundle["objects"][5]
+    assert sighting["id"] == "sighting--287d622a-9ffd-5e7f-bb0b-67f1e320f752"
+    assert (
+        sighting["x_opencti_stix_ids"][0]
+        == "sighting--ee20065d-2555-424f-ad9e-0f8428623c75"
+    )
 
 
 def test_prepare_bundle_ids():
@@ -135,7 +142,9 @@ def test_prepare_bundle_ids():
     prepared_bundle = helper.prepare_bundle_ids(
         bundle=bundle_data, use_json=False, keep_original_id=False
     )
-    print(json.dumps(prepared_bundle))
     malware_target = prepared_bundle["objects"][0]
     assert malware_target["id"] == "malware--d650c5b9-4b43-5781-8576-ea52bd6c7ce0"
     assert malware_target.get("x_opencti_stix_ids") is None
+    sighting = prepared_bundle["objects"][5]
+    assert sighting["id"] == "sighting--287d622a-9ffd-5e7f-bb0b-67f1e320f752"
+    assert ("x_opencti_stix_ids" in sighting) == False