forked from p15670423/monkey
Agent: Remove disused Plugin abstract class
This commit is contained in:
Mike Salvatore
parent
4316329384
commit
703dc315bc
|
|
@ -1,91 +0,0 @@
|
|||
import glob
|
||||
import importlib
|
||||
import inspect
|
||||
import logging
|
||||
from abc import ABCMeta, abstractmethod
|
||||
from os.path import basename, dirname, isfile, join
|
||||
from typing import Callable, Sequence, Type, TypeVar
|
||||
|
||||
logger = 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")]
|
||||
|
||||
|
||||
PluginType = TypeVar("PluginType", bound="Plugin")
|
||||
|
||||
|
||||
class Plugin(metaclass=ABCMeta):
|
||||
@staticmethod
|
||||
@abstractmethod
|
||||
def should_run(class_name: str) -> bool:
|
||||
raise NotImplementedError()
|
||||
|
||||
@classmethod
|
||||
def get_classes(cls) -> Sequence[Callable]:
|
||||
"""
|
||||
Returns the class 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 classes.
|
||||
"""
|
||||
objects = []
|
||||
candidate_files = _get_candidate_files(cls.base_package_file())
|
||||
logger.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:
|
||||
logger.debug("Checking if should run object {}".format(class_object.__name__))
|
||||
try:
|
||||
if class_object.should_run(class_object.__name__):
|
||||
objects.append(class_object)
|
||||
logger.debug("Added {} to list".format(class_object.__name__))
|
||||
except Exception as e:
|
||||
logger.warning(
|
||||
"Exception {} when checking if {} should run".format(
|
||||
str(e), class_object.__name__
|
||||
)
|
||||
)
|
||||
return objects
|
||||
|
||||
@classmethod
|
||||
def get_instances(cls) -> Sequence[Type[PluginType]]:
|
||||
"""
|
||||
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.
|
||||
"""
|
||||
class_objects = cls.get_classes()
|
||||
instances = []
|
||||
for class_object in class_objects:
|
||||
try:
|
||||
instance = class_object()
|
||||
instances.append(instance)
|
||||
except Exception as e:
|
||||
logger.warning(
|
||||
"Exception {} when initializing {}".format(str(e), class_object.__name__)
|
||||
)
|
||||
return instances
|
||||
|
||||
@staticmethod
|
||||
@abstractmethod
|
||||
def base_package_file():
|
||||
pass
|
||||
|
||||
@staticmethod
|
||||
@abstractmethod
|
||||
def base_package_name():
|
||||
pass
|
||||
|
|
@ -1,7 +0,0 @@
|
|||
from tests.unit_tests.infection_monkey.utils.plugins.pluginTests.PluginTestClass import ( # noqa: F401, E501
|
||||
PluginTester,
|
||||
)
|
||||
|
||||
|
||||
class SomeDummyPlugin:
|
||||
pass
|
||||
|
|
@ -1,6 +0,0 @@
|
|||
from tests.unit_tests.infection_monkey.utils.plugins.pluginTests.PluginTestClass import PluginTester
|
||||
|
||||
|
||||
class BadPluginInit(PluginTester):
|
||||
def __init__(self):
|
||||
raise Exception("TestException")
|
||||
|
|
@ -1,10 +0,0 @@
|
|||
from tests.unit_tests.infection_monkey.utils.plugins.pluginTests.PluginTestClass import PluginTester
|
||||
|
||||
|
||||
class BadInit(PluginTester):
|
||||
def __init__(self):
|
||||
raise Exception("TestException")
|
||||
|
||||
|
||||
class ProperClass(PluginTester):
|
||||
pass
|
||||
|
|
@ -1,23 +0,0 @@
|
|||
import tests.unit_tests.infection_monkey.utils.plugins.pluginTests
|
||||
|
||||
from infection_monkey.utils.plugins.plugin import Plugin
|
||||
|
||||
|
||||
class PluginTester(Plugin):
|
||||
classes_to_load = []
|
||||
|
||||
@staticmethod
|
||||
def should_run(class_name):
|
||||
"""
|
||||
Decides if post breach action is enabled in config
|
||||
:return: True if it needs to be ran, false otherwise
|
||||
"""
|
||||
return class_name in PluginTester.classes_to_load
|
||||
|
||||
@staticmethod
|
||||
def base_package_file():
|
||||
return tests.unit_tests.infection_monkey.utils.plugins.pluginTests.__file__
|
||||
|
||||
@staticmethod
|
||||
def base_package_name():
|
||||
return tests.unit_tests.infection_monkey.utils.plugins.pluginTests.__package__
|
||||
|
|
@ -1,5 +0,0 @@
|
|||
from tests.unit_tests.infection_monkey.utils.plugins.pluginTests.PluginTestClass import PluginTester
|
||||
|
||||
|
||||
class PluginWorking(PluginTester):
|
||||
pass
|
||||
|
|
@ -1,38 +0,0 @@
|
|||
from unittest import TestCase
|
||||
|
||||
from tests.unit_tests.infection_monkey.utils.plugins.pluginTests.BadImport import SomeDummyPlugin
|
||||
from tests.unit_tests.infection_monkey.utils.plugins.pluginTests.BadInit import BadPluginInit
|
||||
from tests.unit_tests.infection_monkey.utils.plugins.pluginTests.ComboFile import (
|
||||
BadInit,
|
||||
ProperClass,
|
||||
)
|
||||
from tests.unit_tests.infection_monkey.utils.plugins.pluginTests.PluginTestClass import PluginTester
|
||||
from tests.unit_tests.infection_monkey.utils.plugins.pluginTests.PluginWorking import PluginWorking
|
||||
|
||||
|
||||
class TestPlugin(TestCase):
|
||||
def test_combo_file(self):
|
||||
PluginTester.classes_to_load = [BadInit.__name__, ProperClass.__name__]
|
||||
to_init = PluginTester.get_classes()
|
||||
self.assertEqual(len(to_init), 2)
|
||||
objects = PluginTester.get_instances()
|
||||
self.assertEqual(len(objects), 1)
|
||||
|
||||
def test_bad_init(self):
|
||||
PluginTester.classes_to_load = [BadPluginInit.__name__]
|
||||
to_init = PluginTester.get_classes()
|
||||
self.assertEqual(len(to_init), 1)
|
||||
objects = PluginTester.get_instances()
|
||||
self.assertEqual(len(objects), 0)
|
||||
|
||||
def test_bad_import(self):
|
||||
PluginTester.classes_to_load = [SomeDummyPlugin.__name__]
|
||||
to_init = PluginTester.get_classes()
|
||||
self.assertEqual(len(to_init), 0)
|
||||
|
||||
def test_flow(self):
|
||||
PluginTester.classes_to_load = [PluginWorking.__name__]
|
||||
to_init = PluginTester.get_classes()
|
||||
self.assertEqual(len(to_init), 1)
|
||||
objects = PluginTester.get_instances()
|
||||
self.assertEqual(len(objects), 1)
|
||||
Loading…
Reference in New Issue