From ab8ab7236ac70db20bb3eb0afe4f2337bc6070b5 Mon Sep 17 00:00:00 2001 From: Davide Arcuri Date: Tue, 19 Mar 2024 14:26:51 +0100 Subject: [PATCH] plugin add fix --- orochi/utils/plugin_install.py | 79 ++++++++++--------- .../management/commands/plugins_sync.py | 24 +++--- orochi/website/views.py | 16 ++-- 3 files changed, 62 insertions(+), 57 deletions(-) diff --git a/orochi/utils/plugin_install.py b/orochi/utils/plugin_install.py index 8c8931a9..04db9994 100644 --- 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 index 3ecbba84..de683c10 100644 --- 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 index 1106c1bc..2acd4c63 100644 --- 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"})