From 002447e749dc239355aa3bb36348308f25ed8e81 Mon Sep 17 00:00:00 2001 From: VakarisZ Date: Thu, 16 May 2019 12:57:15 +0300 Subject: [PATCH] Allows to implement helper classes in pba files --- monkey/infection_monkey/post_breach/post_breach_handler.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/monkey/infection_monkey/post_breach/post_breach_handler.py b/monkey/infection_monkey/post_breach/post_breach_handler.py index 9dafa8f45..65b11276c 100644 --- a/monkey/infection_monkey/post_breach/post_breach_handler.py +++ b/monkey/infection_monkey/post_breach/post_breach_handler.py @@ -40,7 +40,8 @@ class PostBreach(object): # 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__] + pba_classes = [m[1] for m in inspect.getmembers(module, inspect.isclass) + if ((m[1].__module__ == module.__name__) and getattr(m[1], "get_pba", False))] # Get post breach action object from class for pba_class in pba_classes: pba = pba_class.get_pba()