diff --git a/monkey/infection_monkey/network/HostFinger.py b/monkey/infection_monkey/network/HostFinger.py index 1eecabe6f..22e5b0698 100644 --- a/monkey/infection_monkey/network/HostFinger.py +++ b/monkey/infection_monkey/network/HostFinger.py @@ -1,9 +1,10 @@ from abc import ABCMeta, abstractproperty, abstractmethod from infection_monkey.config import WormConfiguration +from infection_monkey.utils.plugins.plugin import Plugin -class HostFinger(metaclass=ABCMeta): +class HostFinger(Plugin, metaclass=ABCMeta): @property @abstractmethod def _SCANNED_SERVICE(self): diff --git a/monkey/infection_monkey/network/fingerprinters_collector.py b/monkey/infection_monkey/network/fingerprinters_collector.py index a823d7bfb..a68e00edf 100644 --- a/monkey/infection_monkey/network/fingerprinters_collector.py +++ b/monkey/infection_monkey/network/fingerprinters_collector.py @@ -1,6 +1,6 @@ import logging from typing import Sequence -from infection_monkey.utils.load_plugins import get_instances +from infection_monkey.utils.plugins.load_plugins import get_instances from infection_monkey.network.HostFinger import HostFinger LOG = logging.getLogger(__name__) diff --git a/monkey/infection_monkey/post_breach/pba.py b/monkey/infection_monkey/post_breach/pba.py index 57bf0aaf7..1bb73b604 100644 --- a/monkey/infection_monkey/post_breach/pba.py +++ b/monkey/infection_monkey/post_breach/pba.py @@ -6,6 +6,7 @@ from infection_monkey.telemetry.post_breach_telem import PostBreachTelem from infection_monkey.utils.environment import is_windows_os from infection_monkey.config import WormConfiguration from infection_monkey.telemetry.attack.t1064_telem import T1064Telem +from infection_monkey.utils.plugins.plugin import Plugin LOG = logging.getLogger(__name__) @@ -14,7 +15,7 @@ __author__ = 'VakarisZ' EXECUTION_WITHOUT_OUTPUT = "(PBA execution produced no output)" -class PBA(object): +class PBA(Plugin): """ Post breach action object. Can be extended to support more than command execution on target machine. """ diff --git a/monkey/infection_monkey/post_breach/post_breach_handler.py b/monkey/infection_monkey/post_breach/post_breach_handler.py index a98439d99..e3c81cd23 100644 --- a/monkey/infection_monkey/post_breach/post_breach_handler.py +++ b/monkey/infection_monkey/post_breach/post_breach_handler.py @@ -1,6 +1,6 @@ import logging from infection_monkey.utils.environment import is_windows_os -from infection_monkey.utils.load_plugins import get_instances +from infection_monkey.utils.plugins.load_plugins import get_instances from infection_monkey.post_breach.pba import PBA import infection_monkey.post_breach.actions diff --git a/monkey/infection_monkey/utils/plugins/__init__.py b/monkey/infection_monkey/utils/plugins/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/monkey/infection_monkey/utils/load_plugins.py b/monkey/infection_monkey/utils/plugins/load_plugins.py similarity index 87% rename from monkey/infection_monkey/utils/load_plugins.py rename to monkey/infection_monkey/utils/plugins/load_plugins.py index 4cf676e5f..9a4faf771 100644 --- a/monkey/infection_monkey/utils/load_plugins.py +++ b/monkey/infection_monkey/utils/plugins/load_plugins.py @@ -4,6 +4,8 @@ import logging from os.path import dirname, basename, isfile, join import glob +from infection_monkey.utils.plugins.plugin import Plugin + LOG = logging.getLogger(__name__) @@ -12,9 +14,11 @@ def _get_candidate_files(base_package_file): return [basename(f)[:-3] for f in files if isfile(f) and not f.endswith('__init__.py')] -def get_instances(base_package_name, base_package_file, parent_class): +def get_instances(base_package_name, base_package_file, parent_class: Plugin): """ - Returns the parent_class type objects from base_package_spec according to configuration + Returns the parent_class type objects from base_package_spec. + parent_class must be a class object that inherits from Plugin + base_package name and file must refer to the same package otherwise bad results :return: A list of parent_class objects. """ objects = [] diff --git a/monkey/infection_monkey/utils/plugins/plugin.py b/monkey/infection_monkey/utils/plugins/plugin.py new file mode 100644 index 000000000..ca122f70d --- /dev/null +++ b/monkey/infection_monkey/utils/plugins/plugin.py @@ -0,0 +1,9 @@ +from abc import ABCMeta, abstractmethod + + +class Plugin(metaclass=ABCMeta): + + @staticmethod + @abstractmethod + def should_run(class_name: str) -> bool: + raise NotImplementedError()