[client] Improve helper to better detect entity type to work on (#646)

Author: richard-julien

pycti/connector/opencti_connector_helper.py

@@ -262,17 +262,10 @@ def _data_handler(self, json_data) -> None:
                     raise ValueError(
                         "Internal enrichment must be based on a specific id"
                     )
-                default_reader_type = "Stix-Core-Object"
-                readers = self.helper.api.stix2.get_readers()
-                reader_type = (
-                    entity_type if entity_type is not None else default_reader_type
+                do_read = self.helper.api.stix2.get_reader(
+                    entity_type if entity_type is not None else "Stix-Core-Object"
                 )
-                selected_reader = (
-                    readers[reader_type]
-                    if reader_type in readers
-                    else readers[default_reader_type]
-                )
-                opencti_entity = selected_reader(id=entity_id, withFiles=True)
+                opencti_entity = do_read(id=entity_id, withFiles=True)
                 if opencti_entity is None:
                     raise ValueError(
                         "Unable to read/access to the entity, please check that the connector permission"

pycti/utils/opencti_stix2.py

@@ -754,48 +754,7 @@ def extract_embedded_relationships(
             "reports": reports,
         }
 
-    def get_listers(self):
-        return {
-            "Stix-Core-Object": self.opencti.stix_core_object.list,
-            "Stix-Domain-Object": self.opencti.stix_domain_object.list,
-            "Administrative-Area": self.opencti.location.list,
-            "Attack-Pattern": self.opencti.attack_pattern.list,
-            "Campaign": self.opencti.campaign.list,
-            "Channel": self.opencti.channel.list,
-            "Event": self.opencti.event.list,
-            "Note": self.opencti.note.list,
-            "Observed-Data": self.opencti.observed_data.list,
-            "Opinion": self.opencti.opinion.list,
-            "Report": self.opencti.report.list,
-            "Grouping": self.opencti.grouping.list,
-            "Case-Incident": self.opencti.case_incident.list,
-            "Feedback": self.opencti.feedback.list,
-            "Case-Rfi": self.opencti.case_rfi.list,
-            "Case-Rft": self.opencti.case_rft.list,
-            "Task": self.opencti.task.list,
-            "Course-Of-Action": self.opencti.course_of_action.list,
-            "Data-Component": self.opencti.data_component.list,
-            "Data-Source": self.opencti.data_source.list,
-            "Identity": self.opencti.identity.list,
-            "Indicator": self.opencti.indicator.list,
-            "Infrastructure": self.opencti.infrastructure.list,
-            "Intrusion-Set": self.opencti.intrusion_set.list,
-            "Location": self.opencti.location.list,
-            "Language": self.opencti.language.list,
-            "Malware": self.opencti.malware.list,
-            "Malware-Analysis": self.opencti.malware_analysis.list,
-            "Threat-Actor": self.opencti.threat_actor_group.list,
-            "Threat-Actor-Group": self.opencti.threat_actor_group.list,
-            "Threat-Actor-Individual": self.opencti.threat_actor_individual.list,
-            "Tool": self.opencti.tool.list,
-            "Narrative": self.opencti.narrative.list,
-            "Vulnerability": self.opencti.vulnerability.list,
-            "Incident": self.opencti.incident.list,
-            "Stix-Cyber-Observable": self.opencti.stix_cyber_observable.list,
-            "stix-sighting-relationship": self.opencti.stix_sighting_relationship.list,
-            "stix-core-relationship": self.opencti.stix_core_relationship.list,
-        }
-
+    # Please use get_reader instead of this definition
     def get_readers(self):
         return {
             "Attack-Pattern": self.opencti.attack_pattern.read,
@@ -851,8 +810,11 @@ def get_reader(self, entity_type: str):
             entity_type = "Identity"
         if LocationTypes.has_value(entity_type):
             entity_type = "Location"
+        if entity_type == "Container":
+            entity_type = "Stix-Domain-Object"
         if StixCyberObservableTypes.has_value(entity_type):
             entity_type = "Stix-Cyber-Observable"
+
         readers = self.get_readers()
         return readers.get(
             entity_type, lambda **kwargs: self.unknown_type({"type": entity_type})
@@ -1872,7 +1834,7 @@ def prepare_export(
                 filters=relationships_from_filter
             )
             if len(x) > 0:
-                entity["sighting_of_ref"] = entity["from"]["id"]
+                entity["sighting_of_ref"] = entity["from"]["standard_id"]
                 # handle from and to separately like Stix Core Relationship and call 2 requests
                 objects_to_get.append(
                     entity["from"]
@@ -1886,7 +1848,7 @@ def prepare_export(
                 filters=relationships_to_filter
             )
             if len(y) > 0:
-                entity["where_sighted_refs"] = [entity["to"]["id"]]
+                entity["where_sighted_refs"] = [entity["to"]["standard_id"]]
                 objects_to_get.append(entity["to"])
 
             del entity["from"]
@@ -1903,7 +1865,7 @@ def prepare_export(
                 filters=relationships_from_filter
             )
             if len(x) > 0:
-                entity["source_ref"] = entity["from"]["id"]
+                entity["source_ref"] = entity["from"]["standard_id"]
                 # handle from and to separately like Stix Core Relationship and call 2 requests
                 objects_to_get.append(
                     entity["from"]
@@ -1918,7 +1880,7 @@ def prepare_export(
                 filters=relationships_to_filter
             )
             if len(y) > 0:
-                entity["target_ref"] = entity["to"]["id"]
+                entity["target_ref"] = entity["to"]["standard_id"]
                 objects_to_get.append(entity["to"])
             del entity["to"]
         # Stix Domain Object
@@ -2095,32 +2057,14 @@ def prepare_export(
 
             if no_custom_attributes:
                 del entity["x_opencti_id"]
-            # Export
-            reader = self.get_readers()
             # Get extra objects
             for entity_object in objects_to_get:
-                # Map types
-                if entity_object["entity_type"] == "StixFile":
-                    entity_object["entity_type"] = "File"
-
-                if IdentityTypes.has_value(entity_object["entity_type"]):
-                    entity_object["entity_type"] = "Identity"
-                elif LocationTypes.has_value(entity_object["entity_type"]):
-                    entity_object["entity_type"] = "Location"
-                elif StixCyberObservableTypes.has_value(entity_object["entity_type"]):
-                    entity_object["entity_type"] = "Stix-Cyber-Observable"
-                elif "stix-core-relationship" in entity_object["parent_types"]:
-                    entity_object["entity_type"] = "stix-core-relationship"
-                elif "stix-ref-relationship" in entity_object["parent_types"]:
-                    entity_object["entity_type"] = "stix-ref-relationship"
-
-                do_read = reader.get(
-                    entity_object["entity_type"],
-                    lambda **kwargs: self.unknown_type(
-                        {"type": entity_object["entity_type"]}
-                    ),
-                )
-
+                resolve_type = entity_object["entity_type"]
+                if "stix-core-relationship" in entity_object["parent_types"]:
+                    resolve_type = "stix-core-relationship"
+                if "stix-ref-relationship" in entity_object["parent_types"]:
+                    resolve_type = "stix-ref-relationship"
+                do_read = self.get_reader(resolve_type)
                 query_filters = self.prepare_id_filters_export(
                     entity_object["id"], access_filter
                 )
@@ -2225,20 +2169,7 @@ def get_stix_bundle_or_object_from_entity_id(
             "id": "bundle--" + str(uuid.uuid4()),
             "objects": [],
         }
-
-        if entity_type == "StixFile":
-            entity_type = "File"
-
-        # Map types
-        if IdentityTypes.has_value(entity_type):
-            entity_type = "Identity"
-        if LocationTypes.has_value(entity_type):
-            entity_type = "Location"
-
-        readers = self.get_readers()
-        do_read = readers.get(
-            entity_type, lambda **kwargs: self.unknown_type({"type": entity_type})
-        )
+        do_read = self.get_reader(entity_type)
         entity = do_read(id=entity_id)
         if entity is None:
             self.opencti.app_logger.error(