From 2b01ed7179b4e6dc43579e1df182c49ca2e3e82d Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Donncha=20=C3=93=20Cearbhaill?= <donncha@donncha.is>
Date: Thu, 29 Jun 2023 13:22:43 +0200
Subject: [PATCH] Add optional profiling for MVT modules

---
 mvt/common/cmd_check_iocs.py | 3 ++-
 mvt/common/module.py         | 6 ++++--
 mvt/common/utils.py          | 9 +++++++++
 3 files changed, 15 insertions(+), 3 deletions(-)

diff --git a/mvt/common/cmd_check_iocs.py b/mvt/common/cmd_check_iocs.py
index 6e01bd3..2277241 100644
--- a/mvt/common/cmd_check_iocs.py
+++ b/mvt/common/cmd_check_iocs.py
@@ -8,6 +8,7 @@ import os
 from typing import Optional
 
 from mvt.common.command import Command
+from mvt.common.utils import exec_or_profile
 
 log = logging.getLogger(__name__)
 
@@ -69,7 +70,7 @@ class CmdCheckIOCS(Command):
                     m.indicators.log = m.log
 
                 try:
-                    m.check_indicators()
+                    exec_or_profile("m.check_indicators()", globals(), locals())
                 except NotImplementedError:
                     continue
                 else:
diff --git a/mvt/common/module.py b/mvt/common/module.py
index 0dc02d1..2b7f21d 100644
--- a/mvt/common/module.py
+++ b/mvt/common/module.py
@@ -11,6 +11,8 @@ from typing import Any, Dict, List, Optional, Union
 
 import simplejson as json
 
+from .utils import exec_or_profile
+
 
 class DatabaseNotFoundError(Exception):
     pass
@@ -162,7 +164,7 @@ def run_module(module: MVTModule) -> None:
     module.log.info("Running module %s...", module.__class__.__name__)
 
     try:
-        module.run()
+        exec_or_profile("module.run()", globals(), locals())
     except NotImplementedError:
         module.log.exception(
             "The run() procedure of module %s was not implemented yet!",
@@ -192,7 +194,7 @@ def run_module(module: MVTModule) -> None:
         )
     else:
         try:
-            module.check_indicators()
+            exec_or_profile("module.check_indicators()", globals(), locals())
         except NotImplementedError:
             module.log.info(
                 "The %s module does not support checking for indicators",
diff --git a/mvt/common/utils.py b/mvt/common/utils.py
index 0e33d47..961b6f0 100644
--- a/mvt/common/utils.py
+++ b/mvt/common/utils.py
@@ -8,6 +8,7 @@ import hashlib
 import logging
 import os
 import re
+import cProfile
 from typing import Any, Iterator, Union
 
 from rich.logging import RichHandler
@@ -225,3 +226,11 @@ def set_verbose_logging(verbose: bool = False):
         handler.setLevel(logging.DEBUG)
     else:
         handler.setLevel(logging.INFO)
+
+
+def exec_or_profile(module, globals, locals):
+    """Hook for profiling MVT modules"""
+    if int(os.environ.get("MVT_PROFILE", False)):
+        cProfile.runctx(module, globals, locals)
+    else:
+        exec(module, globals, locals)