[client] Add support of vocabulary (#300)

Co-authored-by: Julien Richard <[email protected]>

Author: Kedae

pycti/api/opencti_api_client.py

@@ -52,6 +52,7 @@
 from pycti.entities.opencti_stix_sighting_relationship import StixSightingRelationship
 from pycti.entities.opencti_threat_actor import ThreatActor
 from pycti.entities.opencti_tool import Tool
+from pycti.entities.opencti_vocabulary import Vocabulary
 from pycti.entities.opencti_vulnerability import Vulnerability
 from pycti.utils.opencti_stix2 import OpenCTIStix2
 from pycti.utils.opencti_stix2_utils import OpenCTIStix2Utils
@@ -150,6 +151,7 @@ def __init__(
         self.stix2 = OpenCTIStix2(self)
 
         # Define the entities
+        self.vocabulary = Vocabulary(self)
         self.label = Label(self)
         self.marking_definition = MarkingDefinition(self)
         self.external_reference = ExternalReference(self, File)

pycti/entities/opencti_vocabulary.py

@@ -0,0 +1,132 @@
+import json
+
+
+class Vocabulary:
+    def __init__(self, opencti):
+        self.opencti = opencti
+        self.properties = """
+            id
+            name
+            category {
+                key
+                fields {
+                    key
+                }
+            }
+        """
+
+    def list(self, **kwargs):
+        filters = kwargs.get("filters", None)
+        self.opencti.log(
+            "info", "Listing Vocabularies with filters " + json.dumps(filters) + "."
+        )
+        query = (
+            """
+                    query Vocabularies($filters: [VocabularyFiltering!]) {
+                        vocabularies(filters: $filters) {
+                            edges {
+                                node {
+                                    """
+            + self.properties
+            + """
+                        }
+                    }
+                }
+            }
+        """
+        )
+        result = self.opencti.query(
+            query,
+            {
+                "filters": filters,
+            },
+        )
+        return self.opencti.process_multiple(result["data"]["vocabularies"])
+
+    def read(self, **kwargs):
+        id = kwargs.get("id", None)
+        filters = kwargs.get("filters", None)
+        if id is not None:
+            self.opencti.log("info", "Reading vocabulary {" + id + "}.")
+            query = (
+                """
+                        query Vocabulary($id: String!) {
+                            vocabulary(id: $id) {
+                                """
+                + self.properties
+                + """
+                    }
+                }
+            """
+            )
+            result = self.opencti.query(query, {"id": id})
+            return self.opencti.process_multiple_fields(result["data"]["vocabulary"])
+        elif filters is not None:
+            result = self.list(filters=filters)
+            if len(result) > 0:
+                return result[0]
+            else:
+                return None
+        else:
+            self.opencti.log(
+                "error", "[opencti_vocabulary] Missing parameters: id or filters"
+            )
+            return None
+
+    def handle_vocab(self, vocab, cache, field):
+        if "vocab_" + vocab in cache:
+            vocab_data = cache["vocab_" + vocab]
+        else:
+            vocab_data = self.read_or_create_unchecked(
+                name=vocab,
+                required=field["required"],
+                category=cache["category_" + field["key"]],
+            )
+        if vocab_data is not None:
+            cache["vocab_" + vocab] = vocab_data
+        return vocab_data
+
+    def create(self, **kwargs):
+        name = kwargs.get("name", None)
+        category = kwargs.get("category", None)
+
+        if name is not None and category is not None:
+            self.opencti.log(
+                "info", "Creating or Getting aliased Vocabulary {" + name + "}."
+            )
+            query = (
+                """
+                        mutation VocabularyAdd($input: VocabularyAddInput!) {
+                            vocabularyAdd(input: $input) {
+                                """
+                + self.properties
+                + """
+                    }
+                }
+            """
+            )
+            result = self.opencti.query(
+                query,
+                {
+                    "input": {
+                        "name": name,
+                        "category": category,
+                    }
+                },
+            )
+            return result["data"]["vocabularyAdd"]
+        else:
+            self.opencti.log(
+                "error",
+                "[opencti_vocabulary] Missing parameters: name or category",
+            )
+
+    def read_or_create_unchecked(self, **kwargs):
+        value = kwargs.get("name", None)
+        vocab = self.read(filters=[{"key": "name", "values": [value]}])
+        if vocab is None:
+            try:
+                return self.create(**kwargs)
+            except ValueError:
+                return None
+        return vocab

pycti/utils/opencti_stix2.py

@@ -295,6 +295,51 @@ def extract_embedded_relationships(
             if "object_marking_refs" in stix_object
             else []
         )
+
+        # Open vocabularies
+        object_open_vocabularies = {}
+        if self.mapping_cache.get("vocabularies_definition_fields") is None:
+            self.mapping_cache["vocabularies_definition_fields"] = []
+            query = """
+                    query getVocabCategories {
+                      vocabularyCategories {
+                        key
+                        fields{
+                          key
+                          required
+                        }
+                      }
+                    }
+                """
+            result = self.opencti.query(query)
+            for category in result["data"]["vocabularyCategories"]:
+                for field in category["fields"]:
+                    self.mapping_cache["vocabularies_definition_fields"].append(field)
+                    self.mapping_cache["category_" + field["key"]] = category["key"]
+        if any(
+            field["key"] in stix_object
+            for field in self.mapping_cache["vocabularies_definition_fields"]
+        ):
+            for f in self.mapping_cache["vocabularies_definition_fields"]:
+                if stix_object.get(f["key"]) is None:
+                    continue
+                if isinstance(stix_object.get(f["key"]), list):
+                    object_open_vocabularies[f["key"]] = []
+                    for vocab in stix_object[f["key"]]:
+                        object_open_vocabularies[f["key"]].append(
+                            self.opencti.vocabulary.handle_vocab(
+                                vocab, self.mapping_cache, field=f
+                            )["name"]
+                        )
+                else:
+                    object_open_vocabularies[
+                        f["key"]
+                    ] = self.opencti.vocabulary.handle_vocab(
+                        stix_object[f["key"]], self.mapping_cache, field=f
+                    )[
+                        "name"
+                    ]
+
         # Object Labels
         object_label_ids = []
         if (
@@ -569,6 +614,7 @@ def extract_embedded_relationships(
             "created_by": created_by_id,
             "object_marking": object_marking_ids,
             "object_label": object_label_ids,
+            "open_vocabs": object_open_vocabularies,
             "kill_chain_phases": kill_chain_phases_ids,
             "object_refs": object_refs_ids,
             "granted_refs": granted_refs_ids,
@@ -604,6 +650,7 @@ def import_object(
         created_by_id = embedded_relationships["created_by"]
         object_marking_ids = embedded_relationships["object_marking"]
         object_label_ids = embedded_relationships["object_label"]
+        open_vocabs = embedded_relationships["open_vocabs"]
         kill_chain_phases_ids = embedded_relationships["kill_chain_phases"]
         object_refs_ids = embedded_relationships["object_refs"]
         external_references_ids = embedded_relationships["external_references"]
@@ -614,6 +661,7 @@ def import_object(
             "created_by_id": created_by_id,
             "object_marking_ids": object_marking_ids,
             "object_label_ids": object_label_ids,
+            "open_vocabs": open_vocabs,
             "kill_chain_phases_ids": kill_chain_phases_ids,
             "object_ids": object_refs_ids,
             "external_references_ids": external_references_ids,
@@ -717,6 +765,7 @@ def import_observable(
         created_by_id = embedded_relationships["created_by"]
         object_marking_ids = embedded_relationships["object_marking"]
         object_label_ids = embedded_relationships["object_label"]
+        open_vocabs = embedded_relationships["open_vocabs"]
         granted_refs_ids = embedded_relationships["granted_refs"]
         kill_chain_phases_ids = embedded_relationships["kill_chain_phases"]
         object_refs_ids = embedded_relationships["object_refs"]
@@ -728,6 +777,7 @@ def import_observable(
             "created_by_id": created_by_id,
             "object_marking_ids": object_marking_ids,
             "object_label_ids": object_label_ids,
+            "open_vocabs": open_vocabs,
             "granted_refs_ids": granted_refs_ids,
             "kill_chain_phases_ids": kill_chain_phases_ids,
             "object_ids": object_refs_ids,
@@ -854,6 +904,7 @@ def import_relationship(
         created_by_id = embedded_relationships["created_by"]
         object_marking_ids = embedded_relationships["object_marking"]
         object_label_ids = embedded_relationships["object_label"]
+        open_vocabs = embedded_relationships["open_vocabs"]
         granted_refs_ids = embedded_relationships["granted_refs"]
         kill_chain_phases_ids = embedded_relationships["kill_chain_phases"]
         object_refs_ids = embedded_relationships["object_refs"]
@@ -865,6 +916,7 @@ def import_relationship(
             "created_by_id": created_by_id,
             "object_marking_ids": object_marking_ids,
             "object_label_ids": object_label_ids,
+            "open_vocabs": open_vocabs,
             "granted_refs_ids": granted_refs_ids,
             "kill_chain_phases_ids": kill_chain_phases_ids,
             "object_ids": object_refs_ids,
@@ -947,6 +999,7 @@ def import_sighting(
         created_by_id = embedded_relationships["created_by"]
         object_marking_ids = embedded_relationships["object_marking"]
         object_label_ids = embedded_relationships["object_label"]
+        open_vocabs = embedded_relationships["open_vocabs"]
         granted_refs_ids = embedded_relationships["granted_refs"]
         kill_chain_phases_ids = embedded_relationships["kill_chain_phases"]
         object_refs_ids = embedded_relationships["object_refs"]
@@ -958,6 +1011,7 @@ def import_sighting(
             "created_by_id": created_by_id,
             "object_marking_ids": object_marking_ids,
             "object_label_ids": object_label_ids,
+            "open_vocabs": open_vocabs,
             "granted_refs_ids": granted_refs_ids,
             "kill_chain_phases_ids": kill_chain_phases_ids,
             "object_ids": object_refs_ids,

tests/cases/entities.py

@@ -596,7 +596,7 @@ def data(self) -> Dict:
             "name": "The Black Vine Cyberespionage Group",
             "description": "A simple report with an indicator and campaign",
             "published": "2016-01-20T17:00:00.000Z",
-            "report_types": ["campaign"],
+            "report_types": ["threat-report"],
             # "lang": "en",
             # "object_refs": [self.ipv4["id"], self.domain["id"]],
         }
@@ -870,7 +870,6 @@ def data(self) -> Dict:
         return {
             "type": "Tool",
             "description": "The Evil Org threat actor group",
-            "tool_types": ["remote-access"],
             "name": "VNC",
         }
 

tests/conftest.py

@@ -20,9 +20,9 @@ def api_client(pytestconfig):
         )
     else:
         return OpenCTIApiClient(
-            "https://demo.opencti.io",
-            "7e663f91-d048-4a8b-bdfa-cdb55597942b",
-            ssl_verify=True,
+            "http://localhost:4000",
+            "d434ce02-e58e-4cac-8b4c-42bf16748e84",
+            ssl_verify=False,
         )