[client] Improve logging (#303)

Author: czechnology

pycti/api/__init__.py

@@ -0,0 +1,3 @@
+import logging
+
+LOGGER = logging.getLogger(__name__)

pycti/api/opencti_api_client.py

@@ -11,6 +11,7 @@
 import urllib3
 from pythonjsonlogger import jsonlogger
 
+from pycti.api import LOGGER
 from pycti.api.opencti_api_connector import OpenCTIApiConnector
 from pycti.api.opencti_api_work import OpenCTIApiWork
 from pycti.entities.opencti_attack_pattern import AttackPattern
@@ -124,21 +125,19 @@ def __init__(
             raise ValueError("A TOKEN must be set")
 
         # Configure logger
-        self.log_level = log_level
-        numeric_level = getattr(logging, self.log_level.upper(), None)
-        if not isinstance(numeric_level, int):
-            raise ValueError("Invalid log level: " + self.log_level)
+        log_level = log_level.upper()
+        LOGGER.setLevel(log_level)
 
         if json_logging:
             log_handler = logging.StreamHandler()
-            log_handler.setLevel(self.log_level.upper())
+            log_handler.setLevel(log_level)
             formatter = CustomJsonFormatter(
                 "%(timestamp)s %(level)s %(name)s %(message)s"
             )
             log_handler.setFormatter(formatter)
-            logging.basicConfig(handlers=[log_handler], level=numeric_level, force=True)
+            logging.basicConfig(handlers=[log_handler], level=log_level, force=True)
         else:
-            logging.basicConfig(level=numeric_level)
+            logging.basicConfig(level=log_level)
 
         # Define API
         self.api_token = token
@@ -328,19 +327,19 @@ def query(self, query, variables={}):
                     else main_error["message"]
                 )
                 if "data" in main_error and "reason" in main_error["data"]:
-                    logging.error(main_error["data"]["reason"])
+                    LOGGER.error(main_error["data"]["reason"])
                     raise ValueError(
                         {"name": error_name, "message": main_error["data"]["reason"]}
                     )
                 else:
-                    logging.error(main_error["message"])
+                    LOGGER.error(main_error["message"])
                     raise ValueError(
                         {"name": error_name, "message": main_error["message"]}
                     )
             else:
                 return result
         else:
-            logging.info(r.text)
+            LOGGER.info(r.text)
             raise ValueError(r.text)
 
     def fetch_opencti_file(self, fetch_uri, binary=False, serialize=False):
@@ -363,24 +362,6 @@ def fetch_opencti_file(self, fetch_uri, binary=False, serialize=False):
             return base64.b64encode(r.text).decode("utf-8")
         return r.text
 
-    def log(self, level, message):
-        """log a message with defined log level
-
-        :param level: must be a valid logging log level (debug, info, warning, error)
-        :type level: str
-        :param message: the message to log
-        :type message: str
-        """
-
-        if level == "debug":
-            logging.debug(message)
-        elif level == "info":
-            logging.info(message)
-        elif level == "warning":
-            logging.warn(message)
-        elif level == "error":
-            logging.error(message)
-
     def health_check(self):
         """submit an example request to the OpenCTI API.
 
@@ -402,7 +383,7 @@ def get_logs_worker_config(self):
         rtype: dict
         """
 
-        logging.info("Getting logs worker config...")
+        LOGGER.info("Getting logs worker config...")
         query = """
             query LogsWorkerConfig {
                 logsWorkerConfig {
@@ -603,7 +584,7 @@ def upload_file(self, **kwargs):
         data = kwargs.get("data", None)
         mime_type = kwargs.get("mime_type", "text/plain")
         if file_name is not None:
-            self.log("info", "Uploading a file.")
+            LOGGER.info("Uploading a file.")
             query = """
                 mutation UploadImport($file: Upload!) {
                     uploadImport(file: $file) {
@@ -621,17 +602,14 @@ def upload_file(self, **kwargs):
 
             return self.query(query, {"file": (File(file_name, data, mime_type))})
         else:
-            self.log(
-                "error",
-                "[upload] Missing parameters: file_name or data",
-            )
+            LOGGER.error("[upload] Missing parameter: file_name")
             return None
 
     def upload_pending_file(self, **kwargs):
         """upload a file to OpenCTI API
 
         :param `**kwargs`: arguments for file upload (required: `file_name` and `data`)
-        :return: returns the query respons for the file upload
+        :return: returns the query response for the file upload
         :rtype: dict
         """
 
@@ -641,7 +619,7 @@ def upload_pending_file(self, **kwargs):
         entity_id = kwargs.get("entity_id", None)
 
         if file_name is not None:
-            self.log("info", "Uploading a file.")
+            LOGGER.info("Uploading a file.")
             query = """
                     mutation UploadPending($file: Upload!, $entityId: String) {
                         uploadPending(file: $file, entityId: $entityId) {
@@ -661,10 +639,7 @@ def upload_pending_file(self, **kwargs):
                 {"file": (File(file_name, data, mime_type)), "entityId": entity_id},
             )
         else:
-            self.log(
-                "error",
-                "[upload] Missing parameters: file_name or data",
-            )
+            LOGGER.error("[upload] Missing parameter: file_name")
             return None
 
     def get_stix_content(self, id):
@@ -674,7 +649,7 @@ def get_stix_content(self, id):
         rtype: dict
         """
 
-        logging.info("Entity in JSON " + id)
+        LOGGER.info("Entity in JSON %s", id)
         query = """
             query StixQuery($id: String!) {
                 stix(id: $id)

pycti/api/opencti_api_connector.py

@@ -1,7 +1,7 @@
 import json
-import logging
 from typing import Any, Dict
 
+from pycti.api import LOGGER
 from pycti.connector.opencti_connector import OpenCTIConnector
 
 
@@ -18,7 +18,7 @@ def list(self) -> Dict:
         :rtype: dict
         """
 
-        logging.info("Getting connectors ...")
+        LOGGER.info("Getting connectors ...")
         query = """
             query GetConnectors {
                 connectorsForWorker {

pycti/api/opencti_api_work.py

@@ -1,7 +1,8 @@
-import logging
 import time
 from typing import Dict, List
 
+from pycti.api import LOGGER
+
 
 class OpenCTIApiWork:
     """OpenCTIApiJob"""
@@ -10,7 +11,7 @@ def __init__(self, api):
         self.api = api
 
     def to_received(self, work_id: str, message: str):
-        logging.info("Reporting work update_received " + work_id)
+        LOGGER.info("Reporting work update_received %s", work_id)
         query = """
             mutation workToReceived($id: ID!, $message: String) {
                 workEdit(id: $id) {
@@ -21,7 +22,7 @@ def to_received(self, work_id: str, message: str):
         self.api.query(query, {"id": work_id, "message": message})
 
     def to_processed(self, work_id: str, message: str, in_error: bool = False):
-        logging.info("Reporting work update_received " + work_id)
+        LOGGER.info("Reporting work update_received %s", work_id)
         query = """
             mutation workToProcessed($id: ID!, $message: String, $inError: Boolean) {
                 workEdit(id: $id) {
@@ -32,7 +33,7 @@ def to_processed(self, work_id: str, message: str, in_error: bool = False):
         self.api.query(query, {"id": work_id, "message": message, "inError": in_error})
 
     def ping(self, work_id: str):
-        logging.info("Ping work " + work_id)
+        LOGGER.info("Ping work %s", work_id)
         query = """
             mutation pingWork($id: ID!) {
                 workEdit(id: $id) {
@@ -43,7 +44,7 @@ def ping(self, work_id: str):
         self.api.query(query, {"id": work_id})
 
     def report_expectation(self, work_id: str, error):
-        logging.info("Report expectation for " + work_id)
+        LOGGER.info("Report expectation for %s", work_id)
         query = """
             mutation reportExpectation($id: ID!, $error: WorkErrorInput) {
                 workEdit(id: $id) {
@@ -57,9 +58,7 @@ def report_expectation(self, work_id: str, error):
             self.api.log("error", "Cannot report expectation")
 
     def add_expectations(self, work_id: str, expectations: int):
-        logging.info(
-            "Update action expectations " + work_id + " - " + str(expectations)
-        )
+        LOGGER.info("Update action expectations %s - %s", work_id, expectations)
         query = """
             mutation addExpectations($id: ID!, $expectations: Int) {
                 workEdit(id: $id) {
@@ -73,7 +72,7 @@ def add_expectations(self, work_id: str, expectations: int):
             self.api.log("error", "Cannot report expectation")
 
     def initiate_work(self, connector_id: str, friendly_name: str) -> str:
-        logging.info("Initiate work for " + connector_id)
+        LOGGER.info("Initiate work for %s", connector_id)
         query = """
             mutation workAdd($connectorId: String!, $friendlyName: String) {
                 workAdd(connectorId: $connectorId, friendlyName: $friendlyName) {

pycti/connector/__init__.py

@@ -0,0 +1,3 @@
+import logging
+
+LOGGER = logging.getLogger(__name__)