[mitre] Refactor to use new helper and icon management

OpenBAS-Platform · Mar 30, 2024 · b84f333 · b84f333
1 parent 9151e90
commit b84f333
Show file tree

Hide file tree

Showing 3 changed files with 9 additions and 3 deletions.
diff --git a/mitre-attack/src/config.yml.sample b/mitre-attack/src/config.yml.sample
@@ -6,5 +6,6 @@ collector:
   id: 'changeme'
   type: 'mitre_attack'
   name: 'Mitre Attack'
+  period: 3600
   log_level: 'info'
 
diff --git a/mitre-attack/src/img/mitre.png b/mitre-attack/src/img/mitre.png
diff --git a/mitre-attack/src/openbas_mitre.py b/mitre-attack/src/openbas_mitre.py
@@ -1,5 +1,5 @@
 import requests
-from _injectors.injector_helper import OpenBASCollectorHelper, OpenBASConfigHelper
+from helpers import OpenBASCollectorHelper, OpenBASConfigHelper
 
 ENTERPRISE_ATTACK_URI = (
     "https://github.com/mitre/cti/raw/master/enterprise-attack/enterprise-attack.json"
@@ -31,9 +31,13 @@ def __init__(self):
                     "env": "COLLECTOR_TYPE",
                     "file_path": ["collector", "type"],
                 },
+                "collector_period": {
+                    "env": "COLLECTOR_PERIOD",
+                    "file_path": ["collector", "period"],
+                },
             },
         )
-        self.helper = OpenBASCollectorHelper(self.config)
+        self.helper = OpenBASCollectorHelper(self.config, open("img/mitre.png", "rb"))
 
     def _kill_chain_phases(self, tactics):
         kill_chain_name = "mitre-attack"
@@ -127,7 +131,8 @@ def _process_message(self) -> None:
 
     # Start the main loop
     def start(self):
-        self.helper.schedule(message_callback=self._process_message, delay=3600)
+        period = self.config.get_conf("collector_period")
+        self.helper.schedule(message_callback=self._process_message, delay=period)
 
 
 if __name__ == "__main__":