[client] Enhance speed of searching entities

Samuel Hassine · Samuel Hassine · commit 2fa8a8ba56ad · 2020-05-05T13:23:21.000+02:00
diff --git a/examples/export_intrusion_set_stix2.py b/examples/export_intrusion_set_stix2.py
@@ -17,7 +17,7 @@
 
 # Create the bundle
 bundle = opencti_api_client.stix2.export_entity(
-    "intrusion-set", intrusion_set["id"], "full"
+    "indicator", "356fea34-f7f5-4110-937c-47c9a5abb8fa", "full"
 )
 json_bundle = json.dumps(bundle, indent=4)
 
diff --git a/examples/import_stix2_file.py b/examples/import_stix2_file.py
@@ -3,8 +3,8 @@
 from pycti import OpenCTIApiClient
 
 # Variables
-api_url = "http://localhost:4000"
-api_token = "0b23f787-d013-41a8-8078-97bee84cc99d"
+api_url = "https://demo.opencti.io"
+api_token = "2b4f29e3-5ea8-4890-8cf5-a76f61f1e2b2"
 
 # OpenCTI initialization
 opencti_api_client = OpenCTIApiClient(api_url, api_token)
diff --git a/pycti/api/opencti_api_client.py b/pycti/api/opencti_api_client.py
@@ -319,8 +319,10 @@ def process_multiple_fields(self, data):
             if "relation" in data["createdByRef"]:
                 row["remote_relation_id"] = data["createdByRef"]["relation"]["id"]
             data["createdByRef"] = row
+            data["createdByRefId"] = row["id"]
         else:
             data["createdByRef"] = None
+            data["createdByRefId"] = None
         if "markingDefinitions" in data:
             data["markingDefinitions"] = self.process_multiple(
                 data["markingDefinitions"]
diff --git a/pycti/entities/opencti_attack_pattern.py b/pycti/entities/opencti_attack_pattern.py
@@ -360,7 +360,7 @@ def create(self, **kwargs):
                 filters=[{"key": "external_id", "values": [external_id]}]
             )
         if object_result is not None:
-            if update or object_result["createdByRef"] == created_by_ref:
+            if update or object_result["createdByRefId"] == created_by_ref:
                 # name
                 if object_result["name"] != name:
                     self.opencti.stix_domain_entity.update_field(
@@ -411,13 +411,15 @@ def create(self, **kwargs):
                     and object_result["external_id"] != external_id
                 ):
                     self.opencti.stix_domain_entity.update_field(
-                        id=object_result["id"], key="external_id", value=external_id
+                        id=object_result["id"],
+                        key="external_id",
+                        value=str(external_id),
                     )
                     object_result["external_id"] = external_id
                 # confidence
                 if confidence is not None and object_result["confidence"] != confidence:
                     self.opencti.stix_domain_entity.update_field(
-                        id=object_result["id"], key="confidence", value=confidence
+                        id=object_result["id"], key="confidence", value=str(confidence)
                     )
                     object_result["confidence"] = confidence
             return object_result
diff --git a/pycti/entities/opencti_campaign.py b/pycti/entities/opencti_campaign.py
@@ -308,7 +308,7 @@ def create(self, **kwargs):
             customAttributes=custom_attributes,
         )
         if object_result is not None:
-            if update or object_result["createdByRef"] == created_by_ref:
+            if update or object_result["createdByRefId"] == created_by_ref:
                 # name
                 if object_result["name"] != name:
                     self.opencti.stix_domain_entity.update_field(
diff --git a/pycti/entities/opencti_course_of_action.py b/pycti/entities/opencti_course_of_action.py
@@ -296,7 +296,7 @@ def create(self, **kwargs):
             customAttributes=custom_attributes,
         )
         if object_result is not None:
-            if update or object_result["createdByRef"] == created_by_ref:
+            if update or object_result["createdByRefId"] == created_by_ref:
                 # name
                 if object_result["name"] != name:
                     self.opencti.stix_domain_entity.update_field(
diff --git a/pycti/entities/opencti_identity.py b/pycti/entities/opencti_identity.py
@@ -295,7 +295,7 @@ def create(self, **kwargs):
             customAttributes=custom_attributes,
         )
         if object_result is not None:
-            if update or object_result["createdByRef"] == created_by_ref:
+            if update or object_result["createdByRefId"] == created_by_ref:
                 # name
                 if object_result["name"] != name:
                     self.opencti.stix_domain_entity.update_field(
diff --git a/pycti/entities/opencti_incident.py b/pycti/entities/opencti_incident.py
@@ -314,7 +314,7 @@ def create(self, **kwargs):
             customAttributes=custom_attributes,
         )
         if object_result is not None:
-            if update or object_result["createdByRef"] == created_by_ref:
+            if update or object_result["createdByRefId"] == created_by_ref:
                 # name
                 if object_result["name"] != name:
                     self.opencti.stix_domain_entity.update_field(
diff --git a/pycti/entities/opencti_indicator.py b/pycti/entities/opencti_indicator.py
@@ -20,6 +20,8 @@ def __init__(self, opencti):
             graph_data
             indicator_pattern
             pattern_type
+            detection
+            confidence
             valid_from
             valid_until
             score
@@ -243,7 +245,7 @@ def read(self, **kwargs):
             result = self.opencti.query(query, {"id": id})
             return self.opencti.process_multiple_fields(result["data"]["indicator"])
         elif filters is not None:
-            result = self.list(filters=filters)
+            result = self.list(filters=filters, customAttributes=custom_attributes)
             if len(result) > 0:
                 return result[0]
             else:
@@ -390,7 +392,7 @@ def create(self, **kwargs):
         """
         object_result = None
         if stix_id_key is not None:
-            object_result = self.opencti.indicator.read(
+            object_result = self.read(
                 id=stix_id_key, customAttributes=custom_attributes
             )
         if object_result is None:
@@ -405,7 +407,7 @@ def create(self, **kwargs):
                 customAttributes=custom_attributes,
             )
         if object_result is not None:
-            if update or object_result["createdByRef"] == created_by_ref:
+            if update or object_result["createdByRefId"] == created_by_ref:
                 # name
                 if name is not None and object_result["name"] != name:
                     self.opencti.stix_domain_entity.update_field(
@@ -424,19 +426,21 @@ def create(self, **kwargs):
                 # score
                 if score is not None and object_result["score"] != score:
                     self.opencti.stix_domain_entity.update_field(
-                        id=object_result["id"], key="score", value=score
+                        id=object_result["id"], key="score", value=str(score)
                     )
                     object_result["score"] = score
                 # confidence
                 if confidence is not None and object_result["confidence"] != confidence:
                     self.opencti.stix_domain_entity.update_field(
-                        id=object_result["id"], key="confidence", value=confidence
+                        id=object_result["id"], key="confidence", value=str(confidence)
                     )
                     object_result["confidence"] = confidence
                 # detection
                 if detection is not None and object_result["detection"] != detection:
                     self.opencti.stix_domain_entity.update_field(
-                        id=object_result["id"], key="detection", value=detection
+                        id=object_result["id"],
+                        key="detection",
+                        value=str(detection).lower(),
                     )
                     object_result["detection"] = detection
             return object_result
diff --git a/pycti/entities/opencti_intrusion_set.py b/pycti/entities/opencti_intrusion_set.py
@@ -330,7 +330,7 @@ def create(self, **kwargs):
             customAttributes=custom_attributes,
         )
         if object_result is not None:
-            if update or object_result["createdByRef"] == created_by_ref:
+            if update or object_result["createdByRefId"] == created_by_ref:
                 # name
                 if object_result["name"] != name:
                     self.opencti.stix_domain_entity.update_field(
diff --git a/pycti/entities/opencti_malware.py b/pycti/entities/opencti_malware.py
@@ -304,7 +304,7 @@ def create(self, **kwargs):
                 node {
                     id
                 }
-            }            
+            }
         """
         object_result = self.opencti.stix_domain_entity.get_by_stix_id_or_name(
             types=["Malware"],
@@ -313,7 +313,7 @@ def create(self, **kwargs):
             customAttributes=custom_attributes,
         )
         if object_result is not None:
-            if update or object_result["createdByRef"] == created_by_ref:
+            if update or object_result["createdByRefId"] == created_by_ref:
                 # name
                 if object_result["name"] != name:
                     self.opencti.stix_domain_entity.update_field(
@@ -332,7 +332,9 @@ def create(self, **kwargs):
                 # is_family
                 if is_family is not None and object_result["is_family"] != is_family:
                     self.opencti.stix_domain_entity.update_field(
-                        id=object_result["id"], key="is_family", value=is_family
+                        id=object_result["id"],
+                        key="is_family",
+                        value=str(is_family).lower(),
                     )
                     object_result["is_family"] = is_family
                 # alias
diff --git a/pycti/entities/opencti_note.py b/pycti/entities/opencti_note.py
@@ -419,7 +419,7 @@ def create(self, **kwargs):
                 custom_attributes=custom_attributes,
             )
         if object_result is not None:
-            if update or object_result["createdByRef"] == created_by_ref:
+            if update or object_result["createdByRefId"] == created_by_ref:
                 if name is not None and object_result["name"] != name:
                     self.opencti.stix_domain_entity.update_field(
                         id=object_result["id"], key="name", value=name
diff --git a/pycti/entities/opencti_opinion.py b/pycti/entities/opencti_opinion.py
@@ -429,7 +429,7 @@ def create(self, **kwargs):
                 custom_attributes=custom_attributes,
             )
         if object_result is not None:
-            if update or object_result["createdByRef"] == created_by_ref:
+            if update or object_result["createdByRefId"] == created_by_ref:
                 if name is not None and object_result["name"] != name:
                     self.opencti.stix_domain_entity.update_field(
                         id=object_result["id"], key="name", value=name
diff --git a/pycti/entities/opencti_report.py b/pycti/entities/opencti_report.py
@@ -443,7 +443,7 @@ def create(self, **kwargs):
                 custom_attributes=custom_attributes,
             )
         if object_result is not None:
-            if update or object_result["createdByRef"] == created_by_ref:
+            if update or object_result["createdByRefId"] == created_by_ref:
                 if object_result["name"] != name:
                     self.opencti.stix_domain_entity.update_field(
                         id=object_result["id"], key="name", value=name
diff --git a/pycti/entities/opencti_stix_domain_entity.py b/pycti/entities/opencti_stix_domain_entity.py
@@ -144,6 +144,7 @@ def __init__(self, opencti, file):
                 }
             }
             ... on Malware {
+                is_family
                 killChainPhases {
                     edges {
                         node {
@@ -280,6 +281,7 @@ def list(self, **kwargs):
         after = kwargs.get("after", None)
         order_by = kwargs.get("orderBy", None)
         order_mode = kwargs.get("orderMode", None)
+        custom_attributes = kwargs.get("customAttributes", None)
         get_all = kwargs.get("getAll", False)
         if get_all:
             first = 500
@@ -295,7 +297,7 @@ def list(self, **kwargs):
                         edges {
                             node {
                                 """
-            + self.properties
+            + (custom_attributes if custom_attributes is not None else self.properties)
             + """
                         }
                     }
@@ -386,7 +388,9 @@ def read(self, **kwargs):
                 result["data"]["stixDomainEntity"]
             )
         elif filters is not None:
-            result = self.list(types=types, filters=filters)
+            result = self.list(
+                types=types, filters=filters, customAttribute=custom_attributes
+            )
             if len(result) > 0:
                 return result[0]
             else:
diff --git a/pycti/entities/opencti_stix_observable.py b/pycti/entities/opencti_stix_observable.py
@@ -124,6 +124,7 @@ def list(self, **kwargs):
         after = kwargs.get("after", None)
         order_by = kwargs.get("orderBy", None)
         order_mode = kwargs.get("orderMode", None)
+        custom_attributes = kwargs.get("customAttributes", None)
         get_all = kwargs.get("getAll", False)
         if get_all:
             first = 500
@@ -138,7 +139,7 @@ def list(self, **kwargs):
                     edges {
                         node {
                             """
-            + self.properties
+            + (custom_attributes if custom_attributes is not None else self.properties)
             + """
                         }
                     }
@@ -225,7 +226,7 @@ def read(self, **kwargs):
                 result["data"]["stixObservable"]
             )
         elif filters is not None:
-            result = self.list(filters=filters)
+            result = self.list(filters=filters, customAttributes=custom_attributes)
             if len(result) > 0:
                 return result[0]
             else:
@@ -329,7 +330,7 @@ def create(self, **kwargs):
             customAttributes=custom_attributes,
         )
         if object_result is not None:
-            if update or object_result["createdByRef"] == created_by_ref:
+            if update or object_result["createdByRefId"] == created_by_ref:
                 if (
                     description is not None
                     and object_result["description"] != "description"
diff --git a/pycti/entities/opencti_stix_sighting.py b/pycti/entities/opencti_stix_sighting.py
@@ -550,51 +550,28 @@ def to_stix2(self, **kwargs):
         if id is not None and entity is None:
             entity = self.read(id=id)
         if entity is not None:
-            roles = self.opencti.resolve_role(
-                entity["relationship_type"],
-                entity["from"]["entity_type"],
-                entity["to"]["entity_type"],
-            )
-            if roles is not None:
-                final_from_id = entity["from"]["stix_id_key"]
-                final_to_id = entity["to"]["stix_id_key"]
-            else:
-                roles = self.opencti.resolve_role(
-                    entity["relationship_type"],
-                    entity["to"]["entity_type"],
-                    entity["from"]["entity_type"],
-                )
-                if roles is not None:
-                    final_from_id = entity["to"]["stix_id_key"]
-                    final_to_id = entity["from"]["stix_id_key"]
-
             stix_sighting = dict()
             stix_sighting["id"] = entity["stix_id_key"]
-            stix_sighting["type"] = "relationship"
+            stix_sighting["type"] = "sighting"
             stix_sighting["spec_version"] = SPEC_VERSION
-            stix_sighting["relationship_type"] = entity["relationship_type"]
+            stix_sighting["sighting_of_ref"] = entity["from"]["stix_id_key"]
+            stix_sighting["where_sighted_refs"] = entity["to"]["stix_id_key"]
             if self.opencti.not_empty(entity["description"]):
                 stix_sighting["description"] = entity["description"]
-            stix_sighting["source_ref"] = final_from_id
-            stix_sighting["target_ref"] = final_to_id
-            stix_sighting[CustomProperties.SOURCE_REF] = final_from_id
-            stix_sighting[CustomProperties.TARGET_REF] = final_to_id
             stix_sighting["created"] = self.opencti.stix2.format_date(entity["created"])
             stix_sighting["modified"] = self.opencti.stix2.format_date(
                 entity["modified"]
             )
             if self.opencti.not_empty(entity["first_seen"]):
-                stix_sighting[
-                    CustomProperties.FIRST_SEEN
-                ] = self.opencti.stix2.format_date(entity["first_seen"])
+                stix_sighting["first_seen"] = self.opencti.stix2.format_date(
+                    entity["first_seen"]
+                )
             if self.opencti.not_empty(entity["last_seen"]):
-                stix_sighting[
-                    CustomProperties.LAST_SEEN
-                ] = self.opencti.stix2.format_date(entity["last_seen"])
-            if self.opencti.not_empty(entity["weight"]):
-                stix_sighting[CustomProperties.WEIGHT] = entity["weight"]
-            if self.opencti.not_empty(entity["role_played"]):
-                stix_sighting[CustomProperties.ROLE_PLAYED] = entity["role_played"]
+                stix_sighting["last_seen"] = self.opencti.stix2.format_date(
+                    entity["last_seen"]
+                )
+            if self.opencti.not_empty(entity["confidence"]):
+                stix_sighting["confidence"] = entity["confidence"]
             stix_sighting[CustomProperties.ID] = entity["id"]
             return self.opencti.stix2.prepare_export(
                 entity, stix_sighting, mode, max_marking_definition_entity
diff --git a/pycti/entities/opencti_threat_actor.py b/pycti/entities/opencti_threat_actor.py
@@ -333,7 +333,7 @@ def create(self, **kwargs):
             customAttributes=custom_attributes,
         )
         if object_result is not None:
-            if update or object_result["createdByRef"] == created_by_ref:
+            if update or object_result["createdByRefId"] == created_by_ref:
                 # name
                 if object_result["name"] != name:
                     self.opencti.stix_domain_entity.update_field(
diff --git a/pycti/entities/opencti_tool.py b/pycti/entities/opencti_tool.py
@@ -289,7 +289,7 @@ def create(self, **kwargs):
             customAttributes=custom_attributes,
         )
         if object_result is not None:
-            if update or object_result["createdByRef"] == created_by_ref:
+            if update or object_result["createdByRefId"] == created_by_ref:
                 # name
                 if object_result["name"] != name:
                     self.opencti.stix_domain_entity.update_field(
diff --git a/pycti/entities/opencti_vulnerability.py b/pycti/entities/opencti_vulnerability.py
@@ -341,7 +341,7 @@ def create(self, **kwargs):
             customAttributes=custom_attributes,
         )
         if object_result is not None:
-            if update or object_result["createdByRef"] == created_by_ref:
+            if update or object_result["createdByRefId"] == created_by_ref:
                 # name
                 if object_result["name"] != name:
                     self.opencti.stix_domain_entity.update_field(
diff --git a/pycti/utils/opencti_stix2.py b/pycti/utils/opencti_stix2.py
@@ -2168,7 +2168,7 @@ def import_bundle(self, stix_bundle, update=False, types=None) -> List:
                 imported_elements.append({"id": item["id"], "type": item["type"]})
         end_time = time.time()
         self.opencti.log(
-            "info", "Relationships imported in: %ssecs" % round(end_time - start_time)
+            "info", "Sightings imported in: %ssecs" % round(end_time - start_time)
         )
 
         # Reports

Original file line number	Diff line number	Diff line change
`@@ -17,7 +17,7 @@`
`17`	`17`
`18`	`18`	`# Create the bundle`
`19`	`19`	`bundle = opencti_api_client.stix2.export_entity(`
`20`		`- "intrusion-set", intrusion_set["id"], "full"`
	`20`	`+ "indicator", "356fea34-f7f5-4110-937c-47c9a5abb8fa", "full"`
`21`	`21`	`)`
`22`	`22`	`json_bundle = json.dumps(bundle, indent=4)`
`23`	`23`
Original file line number	Diff line number	Diff line change
`@@ -308,7 +308,7 @@ def create(self, **kwargs):`
`308`	`308`	`customAttributes=custom_attributes,`
`309`	`309`	`)`
`310`	`310`	`if object_result is not None:`
`311`		`- if update or object_result["createdByRef"] == created_by_ref:`
	`311`	`+ if update or object_result["createdByRefId"] == created_by_ref:`
`312`	`312`	`# name`
`313`	`313`	`if object_result["name"] != name:`
`314`	`314`	`self.opencti.stix_domain_entity.update_field(`
Original file line number	Diff line number	Diff line change
`@@ -296,7 +296,7 @@ def create(self, **kwargs):`
`296`	`296`	`customAttributes=custom_attributes,`
`297`	`297`	`)`
`298`	`298`	`if object_result is not None:`
`299`		`- if update or object_result["createdByRef"] == created_by_ref:`
	`299`	`+ if update or object_result["createdByRefId"] == created_by_ref:`
`300`	`300`	`# name`
`301`	`301`	`if object_result["name"] != name:`
`302`	`302`	`self.opencti.stix_domain_entity.update_field(`
Original file line number	Diff line number	Diff line change
`@@ -295,7 +295,7 @@ def create(self, **kwargs):`
`295`	`295`	`customAttributes=custom_attributes,`
`296`	`296`	`)`
`297`	`297`	`if object_result is not None:`
`298`		`- if update or object_result["createdByRef"] == created_by_ref:`
	`298`	`+ if update or object_result["createdByRefId"] == created_by_ref:`
`299`	`299`	`# name`
`300`	`300`	`if object_result["name"] != name:`
`301`	`301`	`self.opencti.stix_domain_entity.update_field(`
Original file line number	Diff line number	Diff line change
`@@ -314,7 +314,7 @@ def create(self, **kwargs):`
`314`	`314`	`customAttributes=custom_attributes,`
`315`	`315`	`)`
`316`	`316`	`if object_result is not None:`
`317`		`- if update or object_result["createdByRef"] == created_by_ref:`
	`317`	`+ if update or object_result["createdByRefId"] == created_by_ref:`
`318`	`318`	`# name`
`319`	`319`	`if object_result["name"] != name:`
`320`	`320`	`self.opencti.stix_domain_entity.update_field(`
Original file line number	Diff line number	Diff line change
`@@ -330,7 +330,7 @@ def create(self, **kwargs):`
`330`	`330`	`customAttributes=custom_attributes,`
`331`	`331`	`)`
`332`	`332`	`if object_result is not None:`
`333`		`- if update or object_result["createdByRef"] == created_by_ref:`
	`333`	`+ if update or object_result["createdByRefId"] == created_by_ref:`
`334`	`334`	`# name`
`335`	`335`	`if object_result["name"] != name:`
`336`	`336`	`self.opencti.stix_domain_entity.update_field(`