diff --git a/.gitignore b/.gitignore index 343d861..4e7cb1a 100644 --- a/.gitignore +++ b/.gitignore @@ -131,5 +131,6 @@ dmypy.json # Miscellaneous .DS_Store .vscode/ +.idea/ .ctf/ diff --git a/ctfcli/__main__.py b/ctfcli/__main__.py index a6e061e..fecc0f5 100644 --- a/ctfcli/__main__.py +++ b/ctfcli/__main__.py @@ -1,8 +1,6 @@ import configparser -import importlib -import os import subprocess -import sys + from pathlib import Path import click @@ -13,7 +11,7 @@ from ctfcli.cli.plugins import Plugins from ctfcli.cli.templates import Templates from ctfcli.cli.pages import Pages -from ctfcli.utils.plugins import get_plugin_dir +from ctfcli.utils.plugins import load_plugins from ctfcli.utils.git import check_if_dir_is_inside_git_repo @@ -93,15 +91,8 @@ def templates(self): def main(): - # Load plugins - plugin_dir = get_plugin_dir() - sys.path.insert(0, plugin_dir) - for plugin in sorted(os.listdir(plugin_dir)): - plugin_path = os.path.join(plugin_dir, plugin, "__init__.py") - print("Loading", plugin_path, "as", plugin) - loaded = importlib.import_module(plugin) - loaded.load(COMMANDS) - sys.path.remove(plugin_dir) + # load plugins + load_plugins(COMMANDS) # Load CLI fire.Fire(CTFCLI) diff --git a/ctfcli/utils/deploy.py b/ctfcli/utils/deploy.py index a6688d9..e5809a3 100644 --- a/ctfcli/utils/deploy.py +++ b/ctfcli/utils/deploy.py @@ -136,7 +136,7 @@ def cloud(challenge, host, protocol): # Could not find the service. Create it using our pushed image. # Deploy the image by creating service service = s.post( - "/api/v1/services", json={"name": slug, "image": location,} + "/api/v1/services", json={"name": slug, "image": location,}, ).json()["data"] # Get connection details diff --git a/ctfcli/utils/plugins.py b/ctfcli/utils/plugins.py index 6f93c21..4098348 100644 --- a/ctfcli/utils/plugins.py +++ b/ctfcli/utils/plugins.py @@ -1,16 +1,49 @@ -import os - import appdirs +import importlib +import logging +import os +import sys +from pathlib import Path from ctfcli import __name__ as pkg_name +def load_plugins(commands: dict): + plugin_dir = get_plugin_dir() + sys.path.insert(0, plugin_dir) + + for plugin in sorted(os.listdir(plugin_dir)): + plugin_path = Path(plugin_dir) / plugin / "__init__.py" + + logging.debug(f"Loading {plugin_path} as {plugin}") + + loaded = importlib.import_module(plugin) + loaded.load(commands) + + sys.path.remove(str(plugin_dir)) + + def get_plugin_dir(): - plugins_path = os.path.join(get_data_dir(), "plugins") - if not os.path.exists(plugins_path): + if os.getenv("CTFCLI_PLUGIN_DIR"): + plugins_path = get_custom_plugin_dir() + else: + plugins_path = get_data_dir() / "plugins" + + if not plugins_path.exists(): os.makedirs(plugins_path) - return os.path.join(plugins_path) + + return str(plugins_path.absolute()) + + +def get_custom_plugin_dir() -> Path: + custom_plugins_path = Path(os.getenv("CTFCLI_PLUGIN_DIR")) + + if custom_plugins_path.is_absolute(): + return custom_plugins_path + + base_dir = Path().parent.parent + return base_dir / custom_plugins_path def get_data_dir(): - return appdirs.user_data_dir(appname=pkg_name) + return Path(appdirs.user_data_dir(appname=pkg_name))