diff --git a/orochi/utils/plugin_install.py b/orochi/utils/plugin_install.py index ebced56e..d88ee429 100644 --- a/orochi/utils/plugin_install.py +++ b/orochi/utils/plugin_install.py @@ -5,9 +5,10 @@ import zipfile from pathlib import Path -import volatility3.plugins from distributed import get_client from django.conf import settings + +import volatility3.plugins from volatility3 import framework from volatility3.framework import contexts @@ -24,7 +25,7 @@ def install_process(bash_script, reqs_script, tmp_folder): try: bash_script = None reqs_script = False - py_name = None + py_names = [] plugin_folder = Path(settings.VOLATILITY_PLUGIN_PATH) tmp_folder = plugin_folder / str(uuid.uuid4()) os.mkdir(tmp_folder) @@ -42,7 +43,7 @@ def install_process(bash_script, reqs_script, tmp_folder): elif name.endswith(".py"): with open(plugin_folder / name, "wb") as reqs: reqs.write(f.read(name)) - py_name = Path(name).stem + py_names.append(Path(name).stem) # Install all on dask and workers install_process(bash_script, reqs_script, tmp_folder) @@ -53,11 +54,11 @@ def install_process(bash_script, reqs_script, tmp_folder): # after install recover name from installed plugin _ = contexts.Context() _ = framework.import_files(volatility3.plugins, True) - plugin_names = [ - {x: y} - for x, y in framework.list_plugins().items() - if x.startswith(f"custom.{py_name}") - ] + plugin_names = [] + for x, y in framework.list_plugins().items(): + for py_name in py_names: + if x.startswith(f"custom.{py_name}"): + plugin_names.append({x: y}) logging.debug("Plugins installed successfully.") return plugin_names except Exception as e: