Update load_extra_modules.py

kevoreilly · Aug 29, 2023 · ab9c739 · ab9c739
1 parent c7764d1
commit ab9c739
Showing 1 changed file with 17 additions and 15 deletions.
diff --git a/lib/cuckoo/common/load_extra_modules.py b/lib/cuckoo/common/load_extra_modules.py
@@ -38,23 +38,25 @@ def cape_load_decoders(CUCKOO_ROOT: str):
 
     cape_modules = {}
     cape_decoders = os.path.join(CUCKOO_ROOT, "modules", "processing", "parsers", "CAPE")
-    CAPE_DECODERS = [os.path.basename(decoder)[:-3] for decoder in glob.glob(f"{cape_decoders}/[!_]*.py")]
+    CAPE_DECODERS = {"cape": [os.path.basename(decoder)[:-3] for decoder in glob.glob(f"{cape_decoders}/[!_]*.py")]}
 
     private_cape_decoders = os.path.join(CUCKOO_ROOT, "private", "parsers")
-    CAPE_DECODERS += [os.path.basename(decoder)[:-3] for decoder in glob.glob(f"{private_cape_decoders}/[!_]*.py")]
-
-    for name in CAPE_DECODERS:
-        try:
-            # The name of the module must match what's given as the cape_type for yara
-            # hits with the " Config", " Payload", or " Loader" ending removed and with
-            # spaces replaced with underscores.
-            # For example, a cape_type of "Emotet Payload" would trigger a config parser
-            # named "Emotet.py".
-            cape_modules[name.replace("_", " ")] = importlib.import_module(f"modules.processing.parsers.CAPE.{name}")
-        except (ImportError, IndexError) as e:
-            if "datadirs" in str(e):
-                print("You are using wrong pype32 library. pip3 uninstall pype32 && pip3 install -U pype32-py3")
-            print(f"CAPE parser: No module named {name} - {e}")
+    CAPE_DECODERS.setdefault("private", []).extend([os.path.basename(decoder)[:-3] for decoder in glob.glob(f"{private_cape_decoders}/[!_]*.py")])
+
+    versions = {
+        "cape": "modules.processing.parsers.CAPE",
+        "private": "private.parsers",
+    }
+
+    for version, names in CAPE_DECODERS.items():
+        for name in names:
+            try:
+                # The name of the module must match what's given as the cape_type for yara
+                # hits with the " Config", " Payload", or " Loader" ending removed and with  spaces replaced with underscores.
+                # For example, a cape_type of "Emotet Payload" would trigger a config parser named "Emotet.py".
+                cape_modules[name.replace("_", " ")] = importlib.import_module(f"{versions[version]}.{name}")
+            except (ImportError, IndexError) as e:
+                print(f"CAPE parser: No module named {name} - {e}")
 
     return cape_modules