From 9c40b4a022f9d6663fb4ebb883f4011f4e471927 Mon Sep 17 00:00:00 2001 From: Daniel Goldberg Date: Wed, 30 Oct 2019 15:14:27 +0200 Subject: [PATCH] Refactor PBA to use generic plugin --- .../post_breach/actions/__init__.py | 11 -------- .../post_breach/post_breach_handler.py | 27 +++++-------------- 2 files changed, 6 insertions(+), 32 deletions(-) diff --git a/monkey/infection_monkey/post_breach/actions/__init__.py b/monkey/infection_monkey/post_breach/actions/__init__.py index 17007f1e6..e69de29bb 100644 --- a/monkey/infection_monkey/post_breach/actions/__init__.py +++ b/monkey/infection_monkey/post_breach/actions/__init__.py @@ -1,11 +0,0 @@ -from os.path import dirname, basename, isfile, join -import glob - - -def get_pba_files(): - """ - Gets all files under current directory(/actions) - :return: list of all files without .py ending - """ - files = glob.glob(join(dirname(__file__), "*.py")) - return [basename(f)[:-3] for f in files if isfile(f) and not f.endswith('__init__.py')] diff --git a/monkey/infection_monkey/post_breach/post_breach_handler.py b/monkey/infection_monkey/post_breach/post_breach_handler.py index b5dfa93c7..a98439d99 100644 --- a/monkey/infection_monkey/post_breach/post_breach_handler.py +++ b/monkey/infection_monkey/post_breach/post_breach_handler.py @@ -1,9 +1,8 @@ import logging -import inspect -import importlib -from infection_monkey.post_breach.pba import PBA -from infection_monkey.post_breach.actions import get_pba_files from infection_monkey.utils.environment import is_windows_os +from infection_monkey.utils.load_plugins import get_instances +from infection_monkey.post_breach.pba import PBA +import infection_monkey.post_breach.actions LOG = logging.getLogger(__name__) @@ -16,6 +15,7 @@ class PostBreach(object): """ This class handles post breach actions execution """ + def __init__(self): self.os_is_linux = not is_windows_os() self.pba_list = self.config_to_pba_list() @@ -38,20 +38,5 @@ class PostBreach(object): Passes config to each post breach action class and aggregates results into a list. :return: A list of PBA objects. """ - pba_list = [] - pba_files = get_pba_files() - # Go through all of files in ./actions - for pba_file in pba_files: - # Import module from that file - module = importlib.import_module(PATH_TO_ACTIONS + pba_file) - # Get all classes in a module - pba_classes = [m[1] for m in inspect.getmembers(module, inspect.isclass) - if ((m[1].__module__ == module.__name__) and issubclass(m[1], PBA))] - # Get post breach action object from class - for pba_class in pba_classes: - LOG.debug("Checking if should run PBA {}".format(pba_class.__name__)) - if pba_class.should_run(pba_class.__name__): - pba = pba_class() - pba_list.append(pba) - LOG.debug("Added PBA {} to PBA list".format(pba_class.__name__)) - return pba_list + return get_instances(infection_monkey.post_breach.actions.__package__, + infection_monkey.post_breach.actions.__file__, PBA)