From b7982552498b3f7fbcc25edf3ae94471f0fe41c1 Mon Sep 17 00:00:00 2001 From: Ilija Lazoroski Date: Thu, 16 Dec 2021 16:26:40 +0100 Subject: [PATCH] Agent: Add plugin_name attribute to puppet's load_plugin --- monkey/infection_monkey/i_puppet.py | 5 +++-- monkey/infection_monkey/puppet/plugin_registry.py | 6 +++--- monkey/infection_monkey/puppet/puppet.py | 4 ++-- 3 files changed, 8 insertions(+), 7 deletions(-) diff --git a/monkey/infection_monkey/i_puppet.py b/monkey/infection_monkey/i_puppet.py index 50e050dc6..f3567ee77 100644 --- a/monkey/infection_monkey/i_puppet.py +++ b/monkey/infection_monkey/i_puppet.py @@ -27,9 +27,10 @@ PostBreachData = namedtuple("PostBreachData", ["command", "result"]) class IPuppet(metaclass=abc.ABCMeta): @abc.abstractmethod - def load_plugin(self, plugin: object, plugin_type: PluginType) -> None: + def load_plugin(self, plugin_name: str, plugin: object, plugin_type: PluginType) -> None: """ - Loads a plugin into the puppet. + Loads a plugin into the puppet + :param str plugin_name: The plugin class name :param object plugin: The plugin object to load :param PluginType plugin_type: The type of plugin being loaded """ diff --git a/monkey/infection_monkey/puppet/plugin_registry.py b/monkey/infection_monkey/puppet/plugin_registry.py index d86a7491e..76d0d3714 100644 --- a/monkey/infection_monkey/puppet/plugin_registry.py +++ b/monkey/infection_monkey/puppet/plugin_registry.py @@ -23,11 +23,11 @@ class PluginRegistry: """ self._registry = {} - def load_plugin(self, plugin: object, plugin_type: PluginType) -> None: + def load_plugin(self, plugin_name: str, plugin: object, plugin_type: PluginType) -> None: self._registry.setdefault(plugin_type, {}) - self._registry[plugin_type][plugin.__class__.__name__] = plugin + self._registry[plugin_type][plugin_name] = plugin - logger.debug(f"Plugin '{plugin.__class__.__name__}' loaded") + logger.debug(f"Plugin '{plugin_name}' loaded") def get_plugin(self, plugin_name: str, plugin_type: PluginType) -> Optional[object]: try: diff --git a/monkey/infection_monkey/puppet/puppet.py b/monkey/infection_monkey/puppet/puppet.py index 0c36435e0..0fac17a7e 100644 --- a/monkey/infection_monkey/puppet/puppet.py +++ b/monkey/infection_monkey/puppet/puppet.py @@ -20,8 +20,8 @@ class Puppet(IPuppet): def __init__(self) -> None: self._plugin_registry = PluginRegistry() - def load_plugin(self, plugin: object, plugin_type: PluginType) -> None: - self._plugin_registry.load_plugin(plugin, plugin_type) + def load_plugin(self, plugin_name: str, plugin: object, plugin_type: PluginType) -> None: + self._plugin_registry.load_plugin(plugin, plugin_name, plugin_type) def run_sys_info_collector(self, name: str) -> Dict: pass