fixed linting #2

kube_hunter/core/events/handler.py

@@ -6,7 +6,7 @@ from threading import Thread
 
 from kube_hunter.conf import get_config
 from kube_hunter.core.types import ActiveHunter, HunterBase
-from kube_hunter.core.events.types import Vulnerability, EventFilterBase
+from kube_hunter.core.events.types import Vulnerability, EventFilterBase, MultipleEventsContainer
 
 logger = logging.getLogger(__name__)
 
@@ -19,11 +19,33 @@ class EventQueue(Queue):
         self.active_hunters = dict()
         self.all_hunters = dict()
 
-        self.hooks = defaultdict(list)
-        self.filters = defaultdict(list)
         self.running = True
         self.workers = list()
 
+        # -- Regular Subscription --
+        # Structure: key: Event Class, value: tuple(Registered Hunter, Predicate Function)
+        self.hooks = defaultdict(list)
+        self.filters = defaultdict(list)
+        # --------------------------
+
+        # -- Multiple Subscription --
+        # Structure: key: Event Class, value: tuple(Registered Hunter, Predicate Function)
+        self.multi_hooks = defaultdict(list)
+
+        # When subscribing to multiple events, this gets populated with required event classes
+        # Structure: key: Hunter Class, value: set(RequiredEventClass1, RequiredEventClass2)
+        self.hook_dependencies = defaultdict(set)
+
+        # To keep track of fulfilled dependencies. we need to have a structure which saves historical instanciated
+        # events mapped to a registered hunter.
+        # We used a 2 dimensional dictionary in order to fulfill two demands:
+        #   * correctly count published required events
+        #   * save historical events fired, easily sorted by their type
+        #
+        # Structure: hook_fulfilled_deps[hunter_class] -> fulfilled_events_for_hunter[event_class] -> [EventObject, EventObject2]
+        self.hook_fulfilled_deps = defaultdict(lambda: defaultdict(list))
+        # ---------------------------
+
         for _ in range(num_worker):
             t = Thread(target=self.worker)
             t.daemon = True
@@ -34,16 +56,66 @@ class EventQueue(Queue):
         t.daemon = True
         t.start()
 
-    # decorator wrapping for easy subscription
+    """
+    ######################################################
+    + ----------------- Public Methods ----------------- +
+    ######################################################
+    """
+
     def subscribe(self, event, hook=None, predicate=None):
+        """
+        The Subscribe Decorator - For Regular Registration
+        Use this to register for one event only. Your hunter will execute each time this event is published
+
+        @param event - Event class to subscribe to
+        @param predicate - Optional: Function that will be called with the published event as a parameter before trigger.
+                            If it's return value is False, the Hunter will not run (default=None).
+        @param hook - Hunter class to register for (ignore when using as a decorator)
+        """
+
         def wrapper(hook):
             self.subscribe_event(event, hook=hook, predicate=predicate)
             return hook
 
         return wrapper
 
-    # wrapper takes care of the subscribe once mechanism
+    def subscribe_many(self, events, hook=None, predicates=None):
+        """
+        The Subscribe Many Decorator - For Multiple Registration,
+        When your attack needs several prerequisites to exist in the cluster, You need to register for multiple events.
+        Your hunter will execute once for every new combination of required events.
+        For example:
+            1. event A was published 3 times
+            2. event B was published once.
+            3. event B was published again
+        Your hunter will execute 2 times:
+            * (on step 2) with the newest version of A
+            * (on step 3) with the newest version of A and newest version of B
+
+        @param events - List of event classes to subscribe to
+        @param predicates - Optional: List of function that will be called with the published event as a parameter before trigger.
+                            If it's return value is False, the Hunter will not run (default=None).
+        @param hook - Hunter class to register for (ignore when using as a decorator)
+        """
+
+        def wrapper(hook):
+            self.subscribe_events(events, hook=hook, predicates=predicates)
+            return hook
+
+        return wrapper
+
     def subscribe_once(self, event, hook=None, predicate=None):
+        """
+        The Subscribe Once Decorator - For Single Trigger Registration,
+        Use this when you want your hunter to execute only in your entire program run
+        wraps subscribe_event method
+
+        @param events - List of event classes to subscribe to
+        @param predicates - Optional: List of function that will be called with the published event as a parameter before trigger.
+                            If it's return value is False, the Hunter will not run (default=None).
+        @param hook - Hunter class to register for (ignore when using as a decorator)
+        """
+
         def wrapper(hook):
             # installing a __new__ magic method on the hunter
             # which will remove the hunter from the list upon creation
@@ -58,29 +130,160 @@ class EventQueue(Queue):
 
         return wrapper
 
-    # getting uninstantiated event object
-    def subscribe_event(self, event, hook=None, predicate=None):
+    def publish_event(self, event, caller=None):
+        """
+        The Publish Event Method - For Publishing Events To Kube-Hunter's Queue
+        """
+        # Document that the hunter published a vulnerability (if it's indeed a vulnerability)
+        # For statistics options
+        self._increase_vuln_count(event, caller)
+
+        # sets the event's parent to be it's publisher hunter.
+        self._set_event_chain(event, caller)
+
+        # applying filters on the event, before publishing it to subscribers.
+        # if filter returned None, not proceeding to publish
+        event = self.apply_filters(event)
+        if event:
+            # If event was rewritten, make sure it's linked again
+            self._set_event_chain(event, caller)
+
+            # Regular Hunter registrations - publish logic
+            # Here we iterate over all the registered-to events:
+            for hooked_event in self.hooks.keys():
+                # We check if the event we want to publish is an inherited class of the current registered-to iterated event
+                # Meaning - if this is a relevant event:
+                if hooked_event in event.__class__.__mro__:
+                    # If so, we want to publish to all registerd hunters.
+                    for hook, predicate in self.hooks[hooked_event]:
+                        if predicate and not predicate(event):
+                            continue
+
+                        self.put(hook(event))
+                        logger.debug(f"Event {event.__class__} got published to hunter - {hook} with {event}")
+
+            # Multiple Hunter registrations - publish logic
+            # Here we iterate over all the registered-to events:
+            for hooked_event in self.multi_hooks.keys():
+                # We check if the event we want to publish is an inherited class of the current registered-to iterated event
+                # Meaning - if this is a relevant event:
+                if hooked_event in event.__class__.__mro__:
+                    # now we iterate over the corresponding registered hunters.
+                    for hook, predicate in self.multi_hooks[hooked_event]:
+                        if predicate and not predicate(event):
+                            continue
+
+                        self._update_multi_hooks(hook, event)
+
+                        if self._is_all_fulfilled_for_hunter(hook):
+                            events_container = MultipleEventsContainer(self._get_latest_events_from_multi_hooks(hook))
+                            self.put(hook(events_container))
+                            logger.debug(
+                                f"Multiple subscription requirements were met for hunter {hook}. events container was \
+                                published with {self.hook_fulfilled_deps[hook].keys()}"
+                            )
+
+    """
+    ######################################################
+    + ---------------- Private Methods ----------------- +
+    + ---------------- (Backend Logic) ----------------- +
+    ######################################################
+    """
+
+    def _get_latest_events_from_multi_hooks(self, hook):
+        """
+        Iterates over fulfilled deps for the hunter, and fetching the latest appended events from history
+        """
+        latest_events = list()
+        for event_class in self.hook_fulfilled_deps[hook].keys():
+            latest_events.append(self.hook_fulfilled_deps[hook][event_class][-1])
+        return latest_events
+
+    def _update_multi_hooks(self, hook, event):
+        """
+        Updates published events in the multi hooks fulfilled store.
+        """
+        self.hook_fulfilled_deps[hook][event.__class__].append(event)
+
+    def _is_all_fulfilled_for_hunter(self, hook):
+        """
+        Returns true for multi hook fulfilled, else oterwise
+        """
+        # Check if the first dimension already contains all necessary event classes
+        return len(self.hook_fulfilled_deps[hook].keys()) == len(self.hook_dependencies[hook])
+
+    def _set_event_chain(self, event, caller):
+        """
+        Sets' events attribute chain.
+        In here we link the event with it's publisher (Hunter),
+        so in the next hunter that catches this event, we could access the previous one's attributes.
+
+        @param event: the event object to be chained
+        @param caller: the Hunter object that published this event.
+        """
+        if caller:
+            event.previous = caller.event
+            event.hunter = caller.__class__
+
+    def _register_hunters(self, hook=None):
+        """
+        This method is called when a Hunter registers itself to the handler.
+        this is done in order to track and correctly configure the current run of the program.
+
+        passive_hunters, active_hunters, all_hunters
+        """
         config = get_config()
         if ActiveHunter in hook.__mro__:
             if not config.active:
-                return
-            self.active_hunters[hook] = hook.__doc__
+                return False
+            else:
+                self.active_hunters[hook] = hook.__doc__
         elif HunterBase in hook.__mro__:
             self.passive_hunters[hook] = hook.__doc__
 
         if HunterBase in hook.__mro__:
             self.all_hunters[hook] = hook.__doc__
 
+        return True
+
+    def _register_filter(self, event, hook=None, predicate=None):
+        if hook not in self.filters[event]:
+            self.filters[event].append((hook, predicate))
+            logging.debug("{} filter subscribed to {}".format(hook, event))
+
+    def _register_hook(self, event, hook=None, predicate=None):
+        if hook not in self.hooks[event]:
+            self.hooks[event].append((hook, predicate))
+            logging.debug("{} subscribed to {}".format(hook, event))
+
+    def subscribe_event(self, event, hook=None, predicate=None):
+        if not self._register_hunters(hook):
+            return
+
         # registering filters
         if EventFilterBase in hook.__mro__:
-            if hook not in self.filters[event]:
-                self.filters[event].append((hook, predicate))
-                logger.debug(f"{hook} filter subscribed to {event}")
-
+            self._register_filter(event, hook, predicate)
         # registering hunters
-        elif hook not in self.hooks[event]:
-            self.hooks[event].append((hook, predicate))
-            logger.debug(f"{hook} subscribed to {event}")
+        else:
+            self._register_hook(event, hook, predicate)
+
+    def subscribe_events(self, events, hook=None, predicates=None):
+        if not self._register_hunters(hook):
+            return False
+
+        if predicates is None:
+            predicates = [None] * len(events)
+
+        # registering filters.
+        if EventFilterBase in hook.__mro__:
+            for event, predicate in zip(events, predicates):
+                self._register_filter(event, hook, predicate)
+        # registering hunters.
+        else:
+            for event, predicate in zip(events, predicates):
+                self.multi_hooks[event].append((hook, predicate))
+
+            self.hook_dependencies[hook] = frozenset(events)
 
     def apply_filters(self, event):
         # if filters are subscribed, apply them on the event
@@ -97,36 +300,11 @@ class EventQueue(Queue):
                         return None
         return event
 
-    # getting instantiated event object
-    def publish_event(self, event, caller=None):
+    def _increase_vuln_count(self, event, caller):
         config = get_config()
-
-        # setting event chain
-        if caller:
-            event.previous = caller.event
-            event.hunter = caller.__class__
-
-        # applying filters on the event, before publishing it to subscribers.
-        # if filter returned None, not proceeding to publish
-        event = self.apply_filters(event)
-        if event:
-            # If event was rewritten, make sure it's linked to its parent ('previous') event
-            if caller:
-                event.previous = caller.event
-                event.hunter = caller.__class__
-
-            for hooked_event in self.hooks.keys():
-                if hooked_event in event.__class__.__mro__:
-                    for hook, predicate in self.hooks[hooked_event]:
-                        if predicate and not predicate(event):
-                            continue
-
-                        if config.statistics and caller:
-                            if Vulnerability in event.__class__.__mro__:
-                                caller.__class__.publishedVulnerabilities += 1
-
-                        logger.debug(f"Event {event.__class__} got published with {event}")
-                        self.put(hook(event))
+        if config.statistics and caller:
+            if Vulnerability in event.__class__.__mro__:
+                caller.__class__.publishedVulnerabilities += 1
 
     # executes callbacks on dedicated thread as a daemon
     def worker(self):

kube_hunter/core/events/types.py

@@ -62,6 +62,20 @@ class Event:
         return history
 
 
+class MultipleEventsContainer(Event):
+    """
+    This is the class of the object an hunter will get if he was registered to multiple events.
+    """
+
+    def __init__(self, events):
+        self.events = events
+
+    def get_by_class(self, event_class):
+        for event in self.events:
+            if event.__class__ == event_class:
+                return event
+
+
 class Service:
     def __init__(self, name, path="", secure=True):
         self.name = name

tests/core/test_subscribe.py

@@ -6,6 +6,8 @@ from kube_hunter.core.events.types import Event, Service
 from kube_hunter.core.events import handler
 
 counter = 0
+first_run = True
+
 set_config(Config())
 
 
@@ -19,6 +21,16 @@ class RegularEvent(Service, Event):
         Service.__init__(self, "Test Service")
 
 
+class AnotherRegularEvent(Service, Event):
+    def __init__(self):
+        Service.__init__(self, "Test Service (another)")
+
+
+class DifferentRegularEvent(Service, Event):
+    def __init__(self):
+        Service.__init__(self, "Test Service (different)")
+
+
 @handler.subscribe_once(OnceOnlyEvent)
 class OnceHunter(Hunter):
     def __init__(self, event):
@@ -33,8 +45,36 @@ class RegularHunter(Hunter):
         counter += 1
 
 
+@handler.subscribe_many([DifferentRegularEvent, AnotherRegularEvent])
+class SmartHunter(Hunter):
+    def __init__(self, events):
+        global counter, first_run
+        counter += 1
+
+        # we add an attribute on the second scan.
+        # here we test that we get the latest event
+        different_event = events.get_by_class(DifferentRegularEvent)
+        if first_run:
+            first_run = False
+            assert not different_event.new_value
+        else:
+            assert different_event.new_value
+
+
+@handler.subscribe_many([DifferentRegularEvent, AnotherRegularEvent])
+class SmartHunter2(Hunter):
+    def __init__(self, events):
+        global counter
+        counter += 1
+
+        # check if we can access the events
+        assert events.get_by_class(DifferentRegularEvent).__class__ == DifferentRegularEvent
+        assert events.get_by_class(AnotherRegularEvent).__class__ == AnotherRegularEvent
+
+
 def test_subscribe_mechanism():
     global counter
+    counter = 0
 
     # first test normal subscribe and publish works
     handler.publish_event(RegularEvent())
@@ -43,13 +83,47 @@ def test_subscribe_mechanism():
 
     time.sleep(0.02)
     assert counter == 3
+
+
+def test_subscribe_once_mechanism():
+    global counter
     counter = 0
 
-    # testing the subscribe_once mechanism
-    handler.publish_event(OnceOnlyEvent())
-    handler.publish_event(OnceOnlyEvent())
+    # testing the multiple subscription mechanism
     handler.publish_event(OnceOnlyEvent())
 
     time.sleep(0.02)
-    # should have been triggered once
     assert counter == 1
+    counter = 0
+
+    handler.publish_event(OnceOnlyEvent())
+    handler.publish_event(OnceOnlyEvent())
+    handler.publish_event(OnceOnlyEvent())
+    time.sleep(0.02)
+
+    assert counter == 0
+
+
+def test_subscribe_many_mechanism():
+    global counter
+    counter = 0
+
+    # testing the multiple subscription mechanism
+    handler.publish_event(DifferentRegularEvent())
+    handler.publish_event(DifferentRegularEvent())
+    handler.publish_event(DifferentRegularEvent())
+    handler.publish_event(DifferentRegularEvent())
+    handler.publish_event(DifferentRegularEvent())
+    handler.publish_event(AnotherRegularEvent())
+
+    time.sleep(0.02)
+    # We expect SmartHunter and SmartHunter2 to be executed once. hence the counter should be 2
+    assert counter == 2
+    counter = 0
+
+    # Test using most recent event
+    newer_version_event = DifferentRegularEvent()
+    newer_version_event.new_value = True
+    handler.publish_event(newer_version_event)
+
+    assert counter == 2