plugin add fix

LDO-CERT · Mar 19, 2024 · ab8ab72 · ab8ab72
1 parent 297640d
commit ab8ab72
Show file tree

Hide file tree

Showing 3 changed files with 62 additions and 57 deletions.
diff --git a/orochi/utils/plugin_install.py b/orochi/utils/plugin_install.py
@@ -1,4 +1,6 @@
+import logging
 import os
+import shutil
 import uuid
 import zipfile
 from pathlib import Path
@@ -11,28 +13,6 @@
 
 
 def plugin_install(plugin_path):
-    bash_script = None
-    reqs_script = False
-    py_name = None
-
-    plugin_folder = Path(settings.VOLATILITY_PLUGIN_PATH)
-    tmp_folder = plugin_folder / str(uuid.uuid4())
-    os.mkdir(tmp_folder)
-    with zipfile.ZipFile(plugin_path, "r") as f:
-        for name in f.namelist():
-            # Read apt script, no need to persist
-            if name.endswith(".sh"):
-                bash_script = f.read(name).decode("utf-8")
-            # Move requirements in temp folder
-            elif name.lower() == "requirements.txt":
-                reqs_script = True
-                with open(tmp_folder / "requirements.txt", "wb") as reqs:
-                    reqs.write(f.read(name))
-            # keep script in custom folder
-            elif name.endswith(".py"):
-                with open(plugin_folder / name, "wb") as reqs:
-                    reqs.write(f.read(name))
-                py_name = Path(name).stem
 
     def install_process(bash_script, reqs_script, tmp_folder):
         if bash_script:
@@ -41,18 +21,43 @@ def install_process(bash_script, reqs_script, tmp_folder):
         if reqs_script:
             os.system(f"pip install -r {tmp_folder}/requirements.txt")
 
-    # Install all on dask and workers
-    install_process(bash_script, reqs_script, tmp_folder)
-    dask_client = get_client(address="tcp://scheduler:8786")
-    dask_client.run(install_process, bash_script, reqs_script, tmp_folder)
-
-    _ = contexts.Context()
-    _ = framework.import_files(volatility3.plugins, True)
-    available_plugins = framework.list_plugins()
-
-    # after install recover name from installed plugin
-    for available_plugin in available_plugins:
-        if available_plugin.startswith(f"custom.{py_name}"):
-            plugin_name = available_plugin
-            return plugin_name
-    return None
+    try:
+        bash_script = None
+        reqs_script = False
+        py_name = None
+        plugin_folder = Path(settings.VOLATILITY_PLUGIN_PATH)
+        tmp_folder = plugin_folder / str(uuid.uuid4())
+        os.mkdir(tmp_folder)
+        with zipfile.ZipFile(plugin_path, "r") as f:
+            for name in f.namelist():
+                # Read apt script, no need to persist
+                if name.endswith(".sh"):
+                    bash_script = f.read(name).decode("utf-8")
+                # Move requirements in temp folder
+                elif name.lower() == "requirements.txt":
+                    reqs_script = True
+                    with open(tmp_folder / "requirements.txt", "wb") as reqs:
+                        reqs.write(f.read(name))
+                # keep script in custom folder
+                elif name.endswith(".py"):
+                    with open(plugin_folder / name, "wb") as reqs:
+                        reqs.write(f.read(name))
+                    py_name = Path(name).stem
+
+        # Install all on dask and workers
+        install_process(bash_script, reqs_script, tmp_folder)
+        dask_client = get_client(address="tcp://scheduler:8786")
+        dask_client.run(install_process, bash_script, reqs_script, tmp_folder)
+        shutil.rmtree(tmp_folder)
+
+        # after install recover name from installed plugin
+        _ = contexts.Context()
+        _ = framework.import_files(volatility3.plugins, True)
+        plugin_names = [
+            x for x in framework.list_plugins() if x.startswith(f"custom.{py_name}")
+        ]
+        logging.debug(f"Installed plugins: {','.join(plugin_names)}")
+        return plugin_names
+    except Exception as e:
+        logging.error(f"Error installing plugin: {e}")
+        return []
diff --git a/orochi/website/management/commands/plugins_sync.py b/orochi/website/management/commands/plugins_sync.py
@@ -21,17 +21,19 @@ def handle(self, *args, **kwargs):
         installed_plugins = [x.name for x in plugins]
         if len(plugins) > 0:
             self.stdout.write(
-                self.style.SUCCESS(
-                    "Plugins in db: {}".format(", ".join(installed_plugins))
-                )
+                self.style.SUCCESS(f'Plugins in db: {", ".join(installed_plugins)}')
             )
         else:
             self.stdout.write(self.style.SUCCESS("No plugins in db"))
 
         _ = contexts.Context()
         _ = framework.import_files(volatility3.plugins, True)
-        available_plugins = framework.list_plugins()
-        self.stdout.write("Available Plugins: {}".format(", ".join(available_plugins)))
+        available_plugins = {
+            x: y
+            for x, y in framework.list_plugins().items()
+            if not x.startswith("volatility3.cli.")
+        }
+        self.stdout.write(f'Available Plugins: {", ".join(available_plugins)}')
 
         # If plugin doesn't exists anymore disable it
         for plugin in plugins:
@@ -40,9 +42,7 @@ def handle(self, *args, **kwargs):
                 plugin.save()
                 self.stdout.write(
                     self.style.ERROR(
-                        "Plugin {} disabled. It is not available anymore!".format(
-                            plugin
-                        )
+                        f"Plugin {plugin} disabled. It is not available anymore!"
                     )
                 )
 
@@ -59,7 +59,7 @@ def handle(self, *args, **kwargs):
                     plugin = Plugin(name=plugin, operating_system="Other")
                 plugin.comment = plugin_class.__doc__
                 plugin.save()
-                self.stdout.write(self.style.SUCCESS("Plugin {} added!".format(plugin)))
+                self.stdout.write(self.style.SUCCESS(f"Plugin {plugin} added!"))
 
                 # Add new plugin in old dump
                 for dump in Dump.objects.all():
@@ -71,7 +71,7 @@ def handle(self, *args, **kwargs):
                             up.result = RESULT_STATUS_NOT_STARTED
                             up.save()
                 self.stdout.write(
-                    self.style.SUCCESS("Plugin {} added to old dumps!".format(plugin))
+                    self.style.SUCCESS(f"Plugin {plugin} added to old dumps!")
                 )
 
             else:
@@ -85,7 +85,5 @@ def handle(self, *args, **kwargs):
                 up, created = UserPlugin.objects.get_or_create(user=user, plugin=plugin)
                 if created:
                     self.stdout.write(
-                        self.style.SUCCESS(
-                            "Plugin {} added to {}!".format(plugin, user)
-                        )
+                        self.style.SUCCESS(f"Plugin {plugin} added to {user}!")
                     )
diff --git a/orochi/website/views.py b/orochi/website/views.py
@@ -318,14 +318,16 @@ def install_plugin(request):
         f = NamedTemporaryFile(mode="wb", suffix=".zip", delete=False)
         f.write(r.content)
         f.close()
-        if plugin_name := plugin_install(f.name):
-            Plugin(
-                name=plugin_name,
-                operating_system=operating_system,
-                local=True,
-                local_date=datetime.now(),
-            )
+        if plugin_names := plugin_install(f.name):
+            for plugin_name in plugin_names:
+                Plugin(
+                    name=plugin_name,
+                    operating_system=operating_system,
+                    local=True,
+                    local_date=datetime.now(),
+                ).save()
             return JsonResponse({"ok": True})
+        return JsonResponse({"status_code": 404, "error": "Issues installing plugin"})
     return JsonResponse({"status_code": 404, "error": "Issues installing plugin"})