From 69c66072af3b7f233e939eccd3bd2e7f60fadcc3 Mon Sep 17 00:00:00 2001 From: Daniel Goldberg Date: Tue, 12 Nov 2019 19:33:51 +0200 Subject: [PATCH] Turn get_instances into class method. This leads to package_name and package_file both also being class methods. Now, each plugin family can load itself. Reimplemented Fingerprinters and PBAs to use this interface. No more need for fingerprinter collector --- monkey/infection_monkey/monkey.py | 4 +- monkey/infection_monkey/network/HostFinger.py | 14 ++++- .../network/fingerprinters_collector.py | 17 ------ monkey/infection_monkey/post_breach/pba.py | 10 ++++ .../infection_monkey/utils/plugins/plugin.py | 53 +++++++++++++++++++ 5 files changed, 77 insertions(+), 21 deletions(-) delete mode 100644 monkey/infection_monkey/network/fingerprinters_collector.py diff --git a/monkey/infection_monkey/monkey.py b/monkey/infection_monkey/monkey.py index 5b472e0a7..3d0b9d3d7 100644 --- a/monkey/infection_monkey/monkey.py +++ b/monkey/infection_monkey/monkey.py @@ -6,13 +6,13 @@ import sys import time import infection_monkey.tunnel as tunnel +from infection_monkey.network.HostFinger import HostFinger from infection_monkey.utils.monkey_dir import create_monkey_dir, get_monkey_dir_path, remove_monkey_dir from infection_monkey.utils.monkey_log_path import get_monkey_log_path from infection_monkey.utils.environment import is_windows_os from infection_monkey.config import WormConfiguration from infection_monkey.control import ControlClient from infection_monkey.model import DELAY_DELETE_CMD -from infection_monkey.network.fingerprinters_collector import get_fingerprint_instances from infection_monkey.network.firewall import app as firewall from infection_monkey.network.network_scanner import NetworkScanner from infection_monkey.system_info import SystemInfoCollector @@ -146,7 +146,7 @@ class InfectionMonkey(object): self._exploiters = WormConfiguration.exploiter_classes - self._fingerprint = get_fingerprint_instances() + self._fingerprint = HostFinger.get_instances() if not self._keep_running or not WormConfiguration.alive: break diff --git a/monkey/infection_monkey/network/HostFinger.py b/monkey/infection_monkey/network/HostFinger.py index 22e5b0698..692071160 100644 --- a/monkey/infection_monkey/network/HostFinger.py +++ b/monkey/infection_monkey/network/HostFinger.py @@ -1,10 +1,20 @@ -from abc import ABCMeta, abstractproperty, abstractmethod +from abc import abstractmethod from infection_monkey.config import WormConfiguration from infection_monkey.utils.plugins.plugin import Plugin -class HostFinger(Plugin, metaclass=ABCMeta): +class HostFinger(Plugin): + @staticmethod + def base_package_file(): + import infection_monkey.network # avoid circular imports + return infection_monkey.network.__file__ + + @staticmethod + def base_package_name(): + import infection_monkey.network # avoid circular imports + return infection_monkey.network.__package__ + @property @abstractmethod def _SCANNED_SERVICE(self): diff --git a/monkey/infection_monkey/network/fingerprinters_collector.py b/monkey/infection_monkey/network/fingerprinters_collector.py deleted file mode 100644 index a68e00edf..000000000 --- a/monkey/infection_monkey/network/fingerprinters_collector.py +++ /dev/null @@ -1,17 +0,0 @@ -import logging -from typing import Sequence -from infection_monkey.utils.plugins.load_plugins import get_instances -from infection_monkey.network.HostFinger import HostFinger - -LOG = logging.getLogger(__name__) - - -def get_fingerprint_instances() -> Sequence[HostFinger]: - """ - Returns the fingerprint objects according to configuration as a list - :return: A list of HostFinger objects. - """ - # note this currently assumes we're in the same package as the fingerprinters - # if this changes, this file should be updated - # like when they move into a network plugins folder - return get_instances(__package__, __file__, HostFinger) diff --git a/monkey/infection_monkey/post_breach/pba.py b/monkey/infection_monkey/post_breach/pba.py index 1bb73b604..6b689ec23 100644 --- a/monkey/infection_monkey/post_breach/pba.py +++ b/monkey/infection_monkey/post_breach/pba.py @@ -20,6 +20,16 @@ class PBA(Plugin): Post breach action object. Can be extended to support more than command execution on target machine. """ + @staticmethod + def base_package_name(): + import infection_monkey.post_breach.actions # avoid circular imports + return infection_monkey.post_breach.actions.__package__ + + @staticmethod + def base_package_file(): + import infection_monkey.post_breach.actions + return infection_monkey.post_breach.actions.__file__ + def __init__(self, name="unknown", linux_cmd="", windows_cmd=""): """ :param name: Name of post breach action. diff --git a/monkey/infection_monkey/utils/plugins/plugin.py b/monkey/infection_monkey/utils/plugins/plugin.py index ca122f70d..474afbba8 100644 --- a/monkey/infection_monkey/utils/plugins/plugin.py +++ b/monkey/infection_monkey/utils/plugins/plugin.py @@ -1,4 +1,17 @@ +import importlib +import inspect +import logging from abc import ABCMeta, abstractmethod +from os.path import dirname, basename, isfile, join +import glob +from typing import Sequence + +LOG = logging.getLogger(__name__) + + +def _get_candidate_files(base_package_file): + files = glob.glob(join(dirname(base_package_file), "*.py")) + return [basename(f)[:-3] for f in files if isfile(f) and not f.endswith('__init__.py')] class Plugin(metaclass=ABCMeta): @@ -7,3 +20,43 @@ class Plugin(metaclass=ABCMeta): @abstractmethod def should_run(class_name: str) -> bool: raise NotImplementedError() + + @classmethod + def get_instances(cls) -> Sequence[object]: + """ + Returns the type objects from base_package_spec. + base_package name and file must refer to the same package otherwise bad results + :return: A list of parent_class objects. + """ + objects = [] + candidate_files = _get_candidate_files(cls.base_package_file()) + LOG.info("looking for classes of type {} in {}".format(cls.__name__, cls.base_package_name())) + # Go through all of files + for file in candidate_files: + # Import module from that file + module = importlib.import_module('.' + file, cls.base_package_name()) + # Get all classes in a module + # m[1] because return object is (name,class) + classes = [m[1] for m in inspect.getmembers(module, inspect.isclass) + if ((m[1].__module__ == module.__name__) and issubclass(m[1], cls))] + # Get object from class + for class_object in classes: + LOG.debug("Checking if should run object {}".format(class_object.__name__)) + try: + if class_object.should_run(class_object.__name__): + instance = class_object() + objects.append(instance) + LOG.debug("Added {} to list".format(class_object.__name__)) + except Exception as e: + LOG.warning("Exception {} when checking if {} should run".format(str(e), class_object.__name__)) + return objects + + @staticmethod + @abstractmethod + def base_package_file(): + pass + + @staticmethod + @abstractmethod + def base_package_name(): + pass