Add optional profiling for MVT modules

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:

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",

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)