Merge branch 'master' of github.com:OpenCTI-Platform/client-python

Author: Samuel Hassine

.gitignore

@@ -12,4 +12,5 @@ test.py
 examples/*.json
 examples/*.pdf
 docs/_build
-Pipfile*
+Pipfile*
+.mypy_cache/

docs/conf.py

@@ -34,7 +34,12 @@
 # Add any Sphinx extension module names here, as strings. They can be
 # extensions coming with Sphinx (named 'sphinx.ext.*') or your custom
 # ones.
-extensions = ["sphinx.ext.autodoc", "sphinx.ext.inheritance_diagram", "autoapi.sphinx"]
+extensions = [
+    "sphinx.ext.autodoc",
+    "sphinx.ext.inheritance_diagram",
+    "autoapi.sphinx",
+    "sphinx_autodoc_typehints",
+]
 
 # Add any paths that contain templates here, relative to this directory.
 templates_path = ["_templates"]

docs/requirements.txt

@@ -1 +1,3 @@
-autoapi>=2.0.1
+autoapi>=2.0.1
+sphinx_rtd_theme>=0.4.3
+sphinx-autodoc-typehints>=1.10.3

pycti/connector/opencti_connector_helper.py

@@ -8,27 +8,21 @@
 import uuid
 import os
 
-from typing import Callable, Dict, List, Union
+from typing import Callable, Dict, List, Optional, Union
 from pika.exceptions import UnroutableError, NackError
 from pycti.api.opencti_api_client import OpenCTIApiClient
 from pycti.connector.opencti_connector import OpenCTIConnector
 
 
 def get_config_variable(
-    env_var: str, yaml_path: str, config={}, isNumber=False
-) -> Union[str, int]:
+    env_var: str, yaml_path: str, config: Dict = {}, isNumber: Optional[bool] = False
+) -> Union[bool, int, None, str]:
     """[summary]
 
     :param env_var: environnement variable name
-    :type env_var: str
     :param yaml_path: path to yaml config
-    :type yaml_path: str
     :param config: client config dict, defaults to {}
-    :type config: dict, optional
     :param isNumber: specify if the variable is a number, defaults to False
-    :type isNumber: bool, optional
-    :return: either a str or a int variable value
-    :rtype: str or int
     """
 
     if os.getenv(env_var) is not None:

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

@@ -1736,6 +1736,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,
         )
 

setup.py

@@ -62,5 +62,8 @@ def run(self):
         "python-magic-bin==0.4.14;sys.platform=='win32'",
     ],
     cmdclass={"verify": VerifyVersionCommand},
-    extras_require={"dev": ["black", "wheel"]},  # Optional
+    extras_require={
+        "dev": ["black", "wheel",],
+        "doc": ["autoapi", "sphinx_rtd_theme", "sphinx-autodoc-typehints"],
+    },  # Optional
 )