Merge pull request #78 from CybercentreCanada/cli_improvements

Add flag in CLI to force installation of packages in environment (both virtual and non-)

maco/cli.py

@@ -92,14 +92,17 @@ def process_filesystem(
     force: bool,
     include_base64: bool,
     create_venv: bool = False,
+    skip_install: bool = False,
 ) -> Tuple[int, int, int]:
     """Process filesystem with extractors and print results of extraction.
 
     Returns total number of analysed files, yara hits and successful maco extractions.
     """
     if force:
         logger.warning("force execute will cause errors if an extractor requires a yara rule hit during execution")
-    collected = collector.Collector(path_extractors, include=include, exclude=exclude, create_venv=create_venv)
+    collected = collector.Collector(
+        path_extractors, include=include, exclude=exclude, create_venv=create_venv, skip_install=skip_install
+    )
 
     logger.info(f"extractors loaded: {[x for x in collected.extractors.keys()]}\n")
     for _, extractor in collected.extractors.items():
@@ -191,6 +194,11 @@ def main():
         "This runs much slower than the alternative but may be necessary "
         "when there are many extractors with conflicting dependencies.",
     )
+    parser.add_argument(
+        "--force_install",
+        action="store_true",
+        help="Force installation of Python dependencies for extractors (in both host and virtual environments).",
+    )
     args = parser.parse_args()
     inc = args.include.split(",") if args.include else []
     exc = args.exclude.split(",") if args.exclude else []
@@ -236,6 +244,7 @@ def main():
         force=args.force,
         include_base64=args.base64,
         create_venv=args.create_venv,
+        skip_install=not args.force_install,
     )
 
 

maco/collector.py

@@ -67,6 +67,7 @@ def __init__(
         include: List[str] = None,
         exclude: List[str] = None,
         create_venv: bool = False,
+        skip_install: bool = False,
     ):
         """Discover and load extractors from file system."""
         # maco requires the extractor to be imported directly, so ensure they are available on the path
@@ -135,6 +136,7 @@ def extractor_module_callback(module: ModuleType, venv: str):
                     root_directory=path_extractors,
                     scanner=yara.compile(source=utils.MACO_YARA_RULE),
                     create_venv=create_venv and os.path.isdir(path_extractors),
+                    skip_install=skip_install,
                 ),
             )
             p.start()

maco/utils.py

@@ -450,9 +450,19 @@ def run_extractor(
         key = f"{module_name}_{extractor_class}"
         if key not in _loaded_extractors:
             # dynamic import of extractor
-            mod = importlib.import_module(module_name)
-            extractor_cls = mod.__getattribute__(extractor_class)
-            extractor = extractor_cls()
+            try:
+                # Add the correct directory to the PATH before attempting to load the extractor
+                import_path = module_path[: -4 - len(module_name)]
+                sys.path.insert(1, import_path)
+                mod = importlib.import_module(module_name)
+                extractor_cls = mod.__getattribute__(extractor_class)
+                extractor = extractor_cls()
+
+                # Add to cache
+                _loaded_extractors[key] = extractor
+            finally:
+                sys.path.pop(1)
+
         else:
             # retrieve cached extractor
             extractor = _loaded_extractors[key]