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

Author: richard-julien

pycti/api/opencti_api_connector.py

@@ -100,3 +100,18 @@ def register(self, connector: OpenCTIConnector) -> Dict:
            """
         result = self.api.query(query, connector.to_input())
         return result["data"]["registerConnector"]
+
+    def unregister(self, _id: str) -> Dict:
+        """unregister a connector with OpenCTI
+
+        :param _id: `OpenCTIConnector` connector id
+        :type _id: string
+        :return: the response registerConnector data dict
+        :rtype: dict
+        """
+        query = """
+            mutation ConnectorDeletionMutation($id: ID!) {
+                deleteConnector(id: $id) 
+            }
+        """
+        return self.api.query(query, {"id": _id})

pycti/connector/opencti_connector_helper.py

@@ -106,6 +106,8 @@ def __init__(self, helper, config: Dict, callback) -> None:
         self.user = config["connection"]["user"]
         self.password = config["connection"]["pass"]
         self.queue_name = config["listen"]
+        self.exit_event = threading.Event()
+        self.thread = None
 
     # noinspection PyUnusedLocal
     def _process_message(self, channel, method, properties, body) -> None:
@@ -122,9 +124,9 @@ def _process_message(self, channel, method, properties, body) -> None:
         """
 
         json_data = json.loads(body)
-        thread = threading.Thread(target=self._data_handler, args=[json_data])
-        thread.start()
-        while thread.is_alive():  # Loop while the thread is processing
+        self.thread = threading.Thread(target=self._data_handler, args=[json_data])
+        self.thread.start()
+        while self.thread.is_alive():  # Loop while the thread is processing
             assert self.pika_connection is not None
             self.pika_connection.sleep(1.0)
         logging.info(
@@ -159,7 +161,7 @@ def _data_handler(self, json_data) -> None:
                 logging.error("Failing reporting the processing")
 
     def run(self) -> None:
-        while True:
+        while not self.exit_event.is_set():
             try:
                 # Connect the broker
                 self.pika_credentials = pika.PlainCredentials(self.user, self.password)
@@ -186,6 +188,11 @@ def run(self) -> None:
                 self.helper.log_error(str(e))
                 time.sleep(10)
 
+    def stop(self):
+        self.exit_event.set()
+        if self.thread:
+            self.thread.join()
+
 
 class PingAlive(threading.Thread):
     def __init__(self, connector_id, api, get_state, set_state) -> None:
@@ -195,9 +202,10 @@ def __init__(self, connector_id, api, get_state, set_state) -> None:
         self.api = api
         self.get_state = get_state
         self.set_state = set_state
+        self.exit_event = threading.Event()
 
     def ping(self) -> None:
-        while True:
+        while not self.exit_event.is_set():
             try:
                 initial_state = self.get_state()
                 result = self.api.connector.ping(self.connector_id, initial_state)
@@ -221,12 +229,16 @@ def ping(self) -> None:
             except Exception:  # pylint: disable=broad-except
                 self.in_error = True
                 logging.error("Error pinging the API")
-            time.sleep(40)
+            self.exit_event.wait(40)
 
     def run(self) -> None:
         logging.info("Starting ping alive thread")
         self.ping()
 
+    def stop(self) -> None:
+        logging.info("Preparing for clean shutdown")
+        self.exit_event.set()
+
 
 class ListenStream(threading.Thread):
     def __init__(
@@ -239,6 +251,7 @@ def __init__(
         self.token = token
         self.verify_ssl = verify_ssl
         self.start_timestamp = start_timestamp
+        self.exit_event = threading.Event()
 
     def run(self) -> None:  # pylint: disable=too-many-branches
         current_state = self.helper.get_state()
@@ -432,6 +445,17 @@ def __init__(self, config: Dict) -> None:
         )
         self.ping.start()
 
+        # self.listen_stream = None
+        self.listen_queue = None
+
+    def stop(self) -> None:
+        if self.listen_queue:
+            self.listen_queue.stop()
+        # if self.listen_stream:
+        #     self.listen_stream.stop()
+        self.ping.stop()
+        self.api.connector.unregister(self.connector_id)
+
     def get_name(self) -> Optional[Union[bool, int, str]]:
         return self.connect_name
 
@@ -470,8 +494,8 @@ def listen(self, message_callback: Callable[[Dict], str]) -> None:
         :type message_callback: Callable[[Dict], str]
         """
 
-        listen_queue = ListenQueue(self, self.config, message_callback)
-        listen_queue.start()
+        self.listen_queue = ListenQueue(self, self.config, message_callback)
+        self.listen_queue.start()
 
     def listen_stream(
         self,
@@ -486,10 +510,10 @@ def listen_stream(
         :param message_callback: callback function to process messages
         """
 
-        listen_stream = ListenStream(
+        self.listen_stream = ListenStream(
             self, message_callback, url, token, verify_ssl, start_timestamp
         )
-        listen_stream.start()
+        self.listen_stream.start()
 
     def get_opencti_url(self) -> Optional[Union[bool, int, str]]:
         return self.opencti_url

pycti/entities/opencti_attack_pattern.py

@@ -406,3 +406,19 @@ def import_from_stix2(self, **kwargs):
             self.opencti.log(
                 "error", "[opencti_attack_pattern] Missing parameters: stixObject"
             )
+
+    def delete(self, **kwargs):
+        id = kwargs.get("id", None)
+        if id is not None:
+            self.opencti.log("info", "Deleting Attack Pattern {" + id + "}.")
+            query = """
+                 mutation AttackPatternEdit($id: ID!) {
+                     attackPatternEdit(id: $id) {
+                         delete
+                     }
+                 }
+             """
+            self.opencti.query(query, {"id": id})
+        else:
+            self.opencti.log("error", "[attack_pattern] Missing parameters: id")
+            return None

pycti/entities/opencti_identity.py

@@ -248,7 +248,7 @@ def create(self, **kwargs):
         x_opencti_lastname = kwargs.get("x_opencti_lastname", None)
         update = kwargs.get("update", False)
 
-        if name is not None and description is not None:
+        if type is not None and name is not None and description is not None:
             self.opencti.log("info", "Creating Identity {" + name + "}.")
             input_variables = {
                 "stix_id": stix_id,
@@ -321,7 +321,7 @@ def create(self, **kwargs):
                 result["data"][result_data_field]
             )
         else:
-            self.opencti.log("error", "Missing parameters: name and description")
+            self.opencti.log("error", "Missing parameters: type, name and description")
 
     """
         Import an Identity object from a STIX2 object

pycti/entities/opencti_kill_chain_phase.py

@@ -219,3 +219,21 @@ def update_field(self, **kwargs):
                 "[opencti_kill_chain] Missing parameters: id and key and value",
             )
             return None
+
+    def delete(self, **kwargs):
+        id = kwargs.get("id", None)
+        if id is not None:
+            self.opencti.log("info", "Deleting Kill-Chain-Phase {" + id + "}.")
+            query = """
+                 mutation KillChainPhaseEdit($id: ID!) {
+                     killChainPhaseEdit(id: $id) {
+                         delete
+                     }
+                 }
+             """
+            self.opencti.query(query, {"id": id})
+        else:
+            self.opencti.log(
+                "error", "[opencti_kill_chain_phase] Missing parameters: id"
+            )
+            return None

pycti/entities/opencti_label.py

@@ -200,4 +200,20 @@ def update_field(self, **kwargs):
                 "error",
                 "[opencti_label] Missing parameters: id and key and value",
             )
-            return None
+            return None
+
+    def delete(self, **kwargs):
+        id = kwargs.get("id", None)
+        if id is not None:
+            self.opencti.log("info", "Deleting Label {" + id + "}.")
+            query = """
+                 mutation LabelEdit($id: ID!) {
+                     labelEdit(id: $id) {
+                         delete
+                     }
+                 }
+             """
+            self.opencti.query(query, {"id": id})
+        else:
+            self.opencti.log("error", "[opencti_label] Missing parameters: id")
+            return None

pycti/entities/opencti_marking_definition.py

@@ -270,3 +270,21 @@ def import_from_stix2(self, **kwargs):
             self.opencti.log(
                 "error", "[opencti_marking_definition] Missing parameters: stixObject"
             )
+
+    def delete(self, **kwargs):
+        id = kwargs.get("id", None)
+        if id is not None:
+            self.opencti.log("info", "Deleting Marking-Definition {" + id + "}.")
+            query = """
+                 mutation MarkingDefinitionEdit($id: ID!) {
+                     markingDefinitionEdit(id: $id) {
+                         delete
+                     }
+                 }
+             """
+            self.opencti.query(query, {"id": id})
+        else:
+            self.opencti.log(
+                "error", "[opencti_marking_definition] Missing parameters: id"
+            )
+            return None

pycti/entities/opencti_observed_data.py

@@ -387,7 +387,7 @@ def create(self, **kwargs):
         else:
             self.opencti.log(
                 "error",
-                "[opencti_observedData] Missing parameters: first_observed or last_observed",
+                "[opencti_observedData] Missing parameters: first_observed, last_observed or objects",
             )
 
     """

pycti/entities/opencti_stix_cyber_observable.py

@@ -2,6 +2,7 @@
 
 import json
 import os
+
 import magic
 
 
@@ -320,11 +321,11 @@ def list(self, **kwargs):
         )
         query = (
             """
-            query StixCyberObservables($types: [String], $filters: [StixCyberObservablesFiltering], $search: String, $first: Int, $after: ID, $orderBy: StixCyberObservablesOrdering, $orderMode: OrderingMode) {
-                stixCyberObservables(types: $types, filters: $filters, search: $search, first: $first, after: $after, orderBy: $orderBy, orderMode: $orderMode) {
-                    edges {
-                        node {
-                            """
+                query StixCyberObservables($types: [String], $filters: [StixCyberObservablesFiltering], $search: String, $first: Int, $after: ID, $orderBy: StixCyberObservablesOrdering, $orderMode: OrderingMode) {
+                    stixCyberObservables(types: $types, filters: $filters, search: $search, first: $first, after: $after, orderBy: $orderBy, orderMode: $orderMode) {
+                        edges {
+                            node {
+                                """
             + (custom_attributes if custom_attributes is not None else self.properties)
             + """
                         }
@@ -398,9 +399,9 @@ def read(self, **kwargs):
             self.opencti.log("info", "Reading StixCyberObservable {" + id + "}.")
             query = (
                 """
-                query StixCyberObservable($id: String!) {
-                    stixCyberObservable(id: $id) {
-                        """
+                    query StixCyberObservable($id: String!) {
+                        stixCyberObservable(id: $id) {
+                            """
                 + (
                     custom_attributes
                     if custom_attributes is not None
@@ -1736,3 +1737,30 @@ def push_list_export(self, file_name, data, list_filters=""):
                 "listFilters": list_filters,
             },
         )
+
+    def ask_for_enrichment(self, **kwargs):
+        id = kwargs.get("id", None)
+        connector_id = kwargs.get("connector_id", None)
+
+        if id is None or connector_id is None:
+            self.opencti.log("error", "Missing parameters: id and connector_id")
+            return False
+
+        query = """
+            mutation StixCoreObjectEnrichmentLinesMutation($id: ID!, $connectorId: ID!) {
+                stixCoreObjectEdit(id: $id) {
+                    askEnrichment(connectorId: $connectorId) {
+                        id
+                    }
+                }
+            }
+            """
+
+        self.opencti.query(
+            query,
+            {
+                "id": id,
+                "connectorId": connector_id,
+            },
+        )
+        return True

requirements.txt

@@ -11,3 +11,4 @@ black==20.8b1
 python-magic==0.4.22; sys_platform == 'linux' or sys_platform == 'darwin'
 python-magic-bin==0.4.14; sys_platform == 'win32'
 pytest_randomly==3.8.0
+pytest-cases==3.6.3