[client] Set or modify the reliability of an organization (#97)

OpenCTI-Platform · Jun 23, 2020 · 9c5c1f4 · 9c5c1f4
1 parent bc15565
commit 9c5c1f4
Show file tree

Hide file tree

Showing 3 changed files with 47 additions and 14 deletions.
diff --git a/pycti/entities/opencti_identity.py b/pycti/entities/opencti_identity.py
@@ -24,6 +24,7 @@ def __init__(self, opencti):
             updated_at
             ... on Organization {
                 organization_class
+                reliability
             }
             createdByRef {
                 node {
@@ -38,6 +39,7 @@ def __init__(self, opencti):
                     modified
                     ... on Organization {
                         organization_class
+                        reliability
                     }
                 }
                 relation {
@@ -226,6 +228,7 @@ def create_raw(self, **kwargs):
         marking_definitions = kwargs.get("markingDefinitions", None)
         tags = kwargs.get("tags", None)
         organization_class = kwargs.get("organization_class", None)
+        reliability = kwargs.get("reliability", None)
 
         if name is not None and description is not None:
             self.opencti.log("info", "Creating Identity {" + name + "}.")
@@ -254,6 +257,7 @@ def create_raw(self, **kwargs):
                     }
                 """
                 input_variables["organization_class"] = organization_class
+                input_variables["reliability"] = reliability
                 result_data_field = "organizationAdd"
             else:
                 query = """
@@ -296,9 +300,11 @@ def create(self, **kwargs):
         marking_definitions = kwargs.get("markingDefinitions", None)
         tags = kwargs.get("tags", None)
         organization_class = kwargs.get("organization_class", None)
+        reliability = kwargs.get("reliability", None)
         update = kwargs.get("update", False)
         custom_attributes = """
             id
+            stix_id_key
             entity_type
             name
             description 
@@ -308,6 +314,7 @@ def create(self, **kwargs):
             }
             ... on Organization {
                 organization_class
+                reliability
             }
             createdByRef {
                 node {
@@ -373,6 +380,16 @@ def create(self, **kwargs):
                         value=organization_class,
                     )
                     object_result["organization_class"] = organization_class
+                # reliability
+                if (
+                    self.opencti.not_empty(reliability)
+                    and "reliability" in object_result
+                    and object_result["reliability"] != reliability
+                ):
+                    self.opencti.stix_domain_entity.update_field(
+                        id=object_result["id"], key="reliability", value=reliability,
+                    )
+                    object_result["reliability"] = reliability
             return object_result
         else:
             return self.create_raw(
@@ -389,6 +406,7 @@ def create(self, **kwargs):
                 markingDefinitions=marking_definitions,
                 tags=tags,
                 organization_class=organization_class,
+                reliability=reliability,
             )
 
     """
@@ -446,6 +464,9 @@ def import_from_stix2(self, **kwargs):
                 organization_class=stix_object[CustomProperties.ORG_CLASS]
                 if CustomProperties.ORG_CLASS in stix_object
                 else None,
+                reliability=stix_object[CustomProperties.RELIABILITY]
+                if CustomProperties.RELIABILITY in stix_object
+                else None,
                 update=update,
             )
         else:
@@ -494,12 +515,15 @@ def to_stix2(self, **kwargs):
             identity["modified"] = self.opencti.stix2.format_date(entity["modified"])
             if self.opencti.not_empty(entity["alias"]):
                 identity["aliases"] = entity["alias"]
-            if (
-                entity["entity_type"] == "organization"
-                and "organization_class" in entity
-                and self.opencti.not_empty(entity["organization_class"])
-            ):
-                identity[CustomProperties.ORG_CLASS] = entity["organization_class"]
+            if entity["entity_type"] == "organization":
+                if "organization_class" in entity and self.opencti.not_empty(
+                    entity["organization_class"]
+                ):
+                    identity[CustomProperties.ORG_CLASS] = entity["organization_class"]
+                if "reliability" in entity and self.opencti.not_empty(
+                    entity["reliability"]
+                ):
+                    identity[CustomProperties.RELIABILITY] = entity["reliability"]
             identity[CustomProperties.IDENTITY_TYPE] = entity["entity_type"]
             identity[CustomProperties.ID] = entity["id"]
             return self.opencti.stix2.prepare_export(

diff --git a/pycti/utils/constants.py b/pycti/utils/constants.py
@@ -143,6 +143,7 @@ class CustomProperties:
 
     # applies to STIX Identity
     ORG_CLASS = "x_opencti_organization_class"
+    RELIABILITY = "x_opencti_reliability"
     IDENTITY_TYPE = (
         "x_opencti_identity_type"  # this overrides the stix 'identity_class' property!
     )

diff --git a/pycti/utils/opencti_stix2.py b/pycti/utils/opencti_stix2.py
@@ -1409,14 +1409,22 @@ def prepare_export(
             created_by_ref["modified"] = self.format_date(
                 entity_created_by_ref["modified"]
             )
-            if (
-                entity_created_by_ref["entity_type"] == "organization"
-                and "organization_class" in entity_created_by_ref
-                and self.opencti.not_empty(entity_created_by_ref["organization_class"])
-            ):
-                created_by_ref[CustomProperties.ORG_CLASS] = entity_created_by_ref[
-                    "organization_class"
-                ]
+            if entity_created_by_ref["entity_type"] == "organization":
+                if (
+                    "organization_class" in entity_created_by_ref
+                    and self.opencti.not_empty(
+                        entity_created_by_ref["organization_class"]
+                    )
+                ):
+                    created_by_ref[CustomProperties.ORG_CLASS] = entity_created_by_ref[
+                        "organization_class"
+                    ]
+                if "reliability" in entity_created_by_ref and self.opencti.not_empty(
+                    entity_created_by_ref["reliability"]
+                ):
+                    created_by_ref[
+                        CustomProperties.RELIABILITY
+                    ] = entity_created_by_ref["reliability"]
             if self.opencti.not_empty(entity_created_by_ref["alias"]):
                 created_by_ref[CustomProperties.ALIASES] = entity_created_by_ref[
                     "alias"