[client] Handle custom marking definitions

Author: SamuelHassine

pycti/entities/opencti_marking_definition.py

@@ -237,31 +237,41 @@ def import_from_stix2(self, **kwargs):
         stix_object = kwargs.get("stixObject", None)
         update = kwargs.get("update", False)
         if stix_object is not None:
-            definition = None
-            definition_type = stix_object["definition_type"]
-            if stix_object["definition_type"] == "tlp":
-                definition_type = definition_type.upper()
-                if "definition" in stix_object:
-                    definition = (
-                        definition_type + ":" + stix_object["definition"]["tlp"].upper()
-                    )
-                elif "name" in stix_object:
-                    definition = stix_object["name"]
+            if (
+                "x_opencti_definition_type" in stix_object
+                and "x_opencti_definition" in stix_object
+            ):
+                definition_type = stix_object["x_opencti_definition_type"]
+                definition = stix_object["x_opencti_definition"]
             else:
-                if "definition" in stix_object:
-                    if isinstance(stix_object["definition"], str):
-                        definition = stix_object["definition"]
-                    elif (
-                        isinstance(stix_object["definition"], dict)
-                        and stix_object["definition_type"] in stix_object["definition"]
-                    ):
-                        definition = stix_object["definition"][
-                            stix_object["definition_type"]
-                        ]
-                    else:
+                definition_type = stix_object["definition_type"]
+                definition = None
+                if stix_object["definition_type"] == "tlp":
+                    definition_type = definition_type.upper()
+                    if "definition" in stix_object:
+                        definition = (
+                            definition_type
+                            + ":"
+                            + stix_object["definition"]["tlp"].upper()
+                        )
+                    elif "name" in stix_object:
+                        definition = stix_object["name"]
+                else:
+                    if "definition" in stix_object:
+                        if isinstance(stix_object["definition"], str):
+                            definition = stix_object["definition"]
+                        elif (
+                            isinstance(stix_object["definition"], dict)
+                            and stix_object["definition_type"]
+                            in stix_object["definition"]
+                        ):
+                            definition = stix_object["definition"][
+                                stix_object["definition_type"]
+                            ]
+                        else:
+                            definition = stix_object["name"]
+                    elif "name" in stix_object:
                         definition = stix_object["name"]
-                elif "name" in stix_object:
-                    definition = stix_object["name"]
 
             # Replace TLP:WHITE
             if definition == "TLP:WHITE":