[client] Add support for x_opencti_organization_class (#91)

Author: maertv

pycti/entities/opencti_identity.py

@@ -1,7 +1,7 @@
 # coding: utf-8
 
 import json
-from pycti.utils.constants import CustomProperties
+from pycti.utils.constants import CustomProperties, IdentityTypes
 from pycti.utils.opencti_stix2 import SPEC_VERSION
 
 
@@ -223,39 +223,53 @@ def create_raw(self, **kwargs):
         created_by_ref = kwargs.get("createdByRef", None)
         marking_definitions = kwargs.get("markingDefinitions", None)
         tags = kwargs.get("tags", None)
+        organization_class = kwargs.get("organization_class", None)
 
         if name is not None and description is not None:
             self.opencti.log("info", "Creating Identity {" + name + "}.")
-            query = (
+
+            input_variables = {
+                "name": name,
+                "description": description,
+                "alias": alias,
+                "internal_id_key": id,
+                "stix_id_key": stix_id_key,
+                "created": created,
+                "modified": modified,
+                "createdByRef": created_by_ref,
+                "markingDefinitions": marking_definitions,
+                "tags": tags,
+            }
+
+            if type == IdentityTypes.ORGANIZATION.value:
+                query = f"""
+                    mutation OrganizationAdd($input: OrganizationAddInput) {{
+                        organizationAdd(input: $input) {{
+                            {self.properties}
+                        }}
+                    }}
                 """
-                mutation IdentityAdd($input: IdentityAddInput) {
-                    identityAdd(input: $input) {
-                         """
-                + self.properties
-                + """
-                    }
-                }
-            """
-            )
-            result = self.opencti.query(
-                query,
-                {
-                    "input": {
-                        "name": name,
-                        "description": description,
-                        "alias": alias,
-                        "type": type,
-                        "internal_id_key": id,
-                        "stix_id_key": stix_id_key,
-                        "created": created,
-                        "modified": modified,
-                        "createdByRef": created_by_ref,
-                        "markingDefinitions": marking_definitions,
-                        "tags": tags,
-                    }
-                },
+
+                input_variables["organization_class"] = organization_class
+
+                result_data_field = "organizationAdd"
+            else:
+                query = f"""
+                    mutation IdentityAdd($input: IdentityAddInput) {{
+                        identityAdd(input: $input) {{
+                            {self.properties}
+                        }}
+                    }}
+                """
+
+                input_variables["type"] = type
+
+                result_data_field = "identityAdd"
+
+            result = self.opencti.query(query, {"input": input_variables,},)
+            return self.opencti.process_multiple_fields(
+                result["data"][result_data_field]
             )
-            return self.opencti.process_multiple_fields(result["data"]["identityAdd"])
         else:
             self.opencti.log("error", "Missing parameters: name and description")
 
@@ -278,13 +292,17 @@ def create(self, **kwargs):
         created_by_ref = kwargs.get("createdByRef", None)
         marking_definitions = kwargs.get("markingDefinitions", None)
         tags = kwargs.get("tags", None)
+        organization_class = kwargs.get("organization_class", None)
         update = kwargs.get("update", False)
         custom_attributes = """
             id
             entity_type
             name
             description 
             alias
+            ... on Organization {
+                organization_class
+            }
             createdByRef {
                 node {
                     id
@@ -326,6 +344,18 @@ def create(self, **kwargs):
                         id=object_result["id"], key="alias", value=new_aliases
                     )
                     object_result["alias"] = new_aliases
+                # organization_class
+                if (
+                    organization_class is not None
+                    and "organization_class" in object_result
+                    and object_result["organization_class"] != organization_class
+                ):
+                    self.opencti.stix_domain_entity.update_field(
+                        id=object_result["id"],
+                        key="organization_class",
+                        value=organization_class,
+                    )
+                    object_result["organization_class"] = organization_class
             return object_result
         else:
             return self.create_raw(
@@ -340,6 +370,7 @@ def create(self, **kwargs):
                 createdByRef=created_by_ref,
                 markingDefinitions=marking_definitions,
                 tags=tags,
+                organization_class=organization_class,
             )
 
     """

pycti/entities/opencti_stix_domain_entity.py

@@ -392,7 +392,7 @@ def read(self, **kwargs):
             )
         elif filters is not None:
             result = self.list(
-                types=types, filters=filters, customAttribute=custom_attributes
+                types=types, filters=filters, customAttributes=custom_attributes
             )
             if len(result) > 0:
                 return result[0]

pycti/utils/opencti_stix2.py

@@ -1730,6 +1730,9 @@ def create_identity(self, stix_object, extras, update=False):
             if "marking_definitions_ids" in extras
             else [],
             tags=extras["tags_ids"] if "tags_ids" in extras else [],
+            organization_class=stix_object[CustomProperties.ORG_CLASS]
+            if CustomProperties.ORG_CLASS in stix_object
+            else None,
             update=update,
         )