From 57f142e68a207b439c187e13ead4c668644aa7e0 Mon Sep 17 00:00:00 2001
From: Axel Fahy <axelfahy@users.noreply.github.com>
Date: Mon, 25 Sep 2023 00:07:34 +0200
Subject: [PATCH] [mitre] add metrics (#1423)

---
 external-import/mitre/docker-compose.yml    | 2 +-
 external-import/mitre/src/config.yml.sample | 3 ++-
 external-import/mitre/src/connector.py      | 9 ++++++++-
 3 files changed, 11 insertions(+), 3 deletions(-)

diff --git a/external-import/mitre/docker-compose.yml b/external-import/mitre/docker-compose.yml
index c4927b157c..208bda12e0 100644
--- a/external-import/mitre/docker-compose.yml
+++ b/external-import/mitre/docker-compose.yml
@@ -14,4 +14,4 @@ services:
       - CONNECTOR_RUN_AND_TERMINATE=false
       - CONNECTOR_LOG_LEVEL=error
       - MITRE_INTERVAL=7 # In days
-    restart: always
\ No newline at end of file
+    restart: always
diff --git a/external-import/mitre/src/config.yml.sample b/external-import/mitre/src/config.yml.sample
index c0b85bf96f..4e4cb70696 100644
--- a/external-import/mitre/src/config.yml.sample
+++ b/external-import/mitre/src/config.yml.sample
@@ -11,6 +11,7 @@ connector:
   update_existing_data: false
   run_and_terminate: false
   log_level: 'info'
+  expose_metrics: False
 
 mitre:
-  interval: 7 # In days
\ No newline at end of file
+  interval: 7 # In days
diff --git a/external-import/mitre/src/connector.py b/external-import/mitre/src/connector.py
index cef9f14c77..763bea8eb3 100644
--- a/external-import/mitre/src/connector.py
+++ b/external-import/mitre/src/connector.py
@@ -56,7 +56,7 @@ class Mitre:
     def __init__(self):
         config_file_path = os.path.dirname(os.path.abspath(__file__)) + "/config.yml"
         config = (
-            yaml.load(open(config_file_path), Loader=yaml.FullLoader)
+            yaml.load(open(config_file_path, encoding="utf8"), Loader=yaml.FullLoader)
             if os.path.isfile(config_file_path)
             else {}
         )
@@ -153,6 +153,7 @@ def retrieve_data(self, url: str) -> Optional[dict]:
             urllib.error.ContentTooShortError,
         ) as urllib_error:
             self.helper.log_error(f"Error retrieving url {url}: {urllib_error}")
+            self.helper.metric.inc("client_error_count")
         return None
 
     def add_confidence_to_bundle_objects(self, stix_bundle: dict):
@@ -189,6 +190,9 @@ def process_data(self):
             return
 
         self.helper.log_info(f"Connector will run now {time_now}.")
+        self.helper.metric.inc("run_count")
+        self.helper.metric.state("running")
+
         friendly_name = f"MITRE run @ {time_now}"
         work_id = self.helper.api.work.initiate_work(
             self.helper.connect_id, friendly_name
@@ -208,6 +212,7 @@ def process_data(self):
                 update=self.update_existing_data,
                 work_id=work_id,
             )
+            self.helper.metric.inc("record_send", len(data["objects"]))
 
         message = f"Connector successfully run, storing last_run as {time_now}"
         self.helper.log_info(message)
@@ -227,10 +232,12 @@ def run(self):
                 self.process_data()
             except (KeyboardInterrupt, SystemExit):
                 self.helper.log_info("Connector stop")
+                self.helper.metric.state("stopped")
                 sys.exit(0)
             except Exception as e:
                 self.helper.log_error(str(e))
             finally:
+                self.helper.metric.state("idle")
                 time.sleep(60)