diff --git a/src/core/events/handler.py b/src/core/events/handler.py index 51d501f..ba9ccac 100644 --- a/src/core/events/handler.py +++ b/src/core/events/handler.py @@ -23,6 +23,9 @@ class EventQueue(Queue, object): self.hooks = defaultdict(list) self.filters = defaultdict(list) + self.multi_hooks = defaultdict(list) # for hooks which listens for multiple events. + self.hook_dependencies = defaultdict(set) # dependencies which were defined for a given hook. + self.hook_fulfilled_deps = defaultdict(set) # keep track of already fulfilled dependencies. self.running = True self.workers = list() @@ -44,6 +47,13 @@ class EventQueue(Queue, object): return wrapper + def multiple_subscribe(self, events, predicates=None): + def wrapper(hook): + self.subscribe_events(events, hook=hook, predicates=predicates) + return hook + + return wrapper + # wrapper takes care of the subscribe once mechanism def subscribe_once(self, event, hook=None, predicate=None): def wrapper(hook): @@ -59,10 +69,12 @@ class EventQueue(Queue, object): return wrapper # getting uninstantiated event object - def subscribe_event(self, event, hook=None, predicate=None): + + # register hunter in tables, return True if the hunter should be used otherwise False (active mode vs passive mode) + def _register_hunters(self, hook=None): if ActiveHunter in hook.__mro__: if not config.active: - return + return False else: self.active_hunters[hook] = hook.__doc__ elif HunterBase in hook.__mro__: @@ -71,20 +83,49 @@ class EventQueue(Queue, object): if HunterBase in hook.__mro__: self.all_hunters[hook] = hook.__doc__ - # registering filters - if EventFilterBase in hook.__mro__: - if hook not in self.filters[event]: - self.filters[event].append((hook, predicate)) - logging.debug('{} filter subscribed to {}'.format(hook, event)) + return True - # registering hunters - elif hook not in self.hooks[event]: + 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__: + self._register_filter(event, hook, predicate) + # registering hunters + 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] = set(events) def apply_filters(self, event): - # if filters are subscribed, apply them on the event + # if filters are subscribed, apply them on the event for hooked_event in self.filters.keys(): if hooked_event in event.__class__.__mro__: for filter_hook, predicate in self.filters[hooked_event]: @@ -98,21 +139,39 @@ class EventQueue(Queue, object): return None return event - # getting instantiated event object - def publish_event(self, event, caller=None): - # setting event chain + def _set_event_chain(self, event, caller): if caller: event.previous = caller.event event.hunter = caller.__class__ - # applying filters on the event, before publishing it to subscribers. + def _increase_vuln_count(self, event, caller): + if config.statistics and caller: + if Vulnerability in event.__class__.__mro__: + caller.__class__.publishedVulnerabilities += 1 + + def mark_and_fire_if_possible(self, hooked_event, hook): + # Sanity check. + assert (hooked_event in self.hook_dependencies[hook]) + + self.hook_fulfilled_deps[hook].add(hooked_event) + + if len(self.hook_fulfilled_deps[hook]) == len(self.hook_dependencies[hook]): + # fire it! + self.put(hook(self.hook_fulfilled_deps[hook])) + # reset the state. + self.hook_fulfilled_deps[hook] = set() + + # getting instantiated event object + def publish_event(self, event, caller=None): + # setting event chain + 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 to its parent ('previous') event - if caller: - event.previous = caller.event - event.hunter = caller.__class__ + self._set_event_chain(event, caller) for hooked_event in self.hooks.keys(): if hooked_event in event.__class__.__mro__: @@ -120,13 +179,20 @@ class EventQueue(Queue, object): if predicate and not predicate(event): continue - if config.statistics and caller: - if Vulnerability in event.__class__.__mro__: - caller.__class__.publishedVulnerabilities += 1 + self._increase_vuln_count(event, caller) logging.debug('Event {} got published with {}'.format(event.__class__, event)) self.put(hook(event)) + for hook, predicate in self.multi_hooks[hooked_event]: + if predicate and not predicate(event): + continue + + logging.debug('Event {} got published with {}'.format(event.__class__, event)) + self._increase_vuln_count(event, caller) + self.mark_and_fire_if_possible(hooked_event, hook) + + # executes callbacks on dedicated thread as a daemon def worker(self): while self.running: diff --git a/tests/core/test_subscribe.py b/tests/core/test_subscribe.py index adf42ec..e320b9c 100644 --- a/tests/core/test_subscribe.py +++ b/tests/core/test_subscribe.py @@ -26,25 +26,38 @@ class RegularHunter(Hunter): global counter counter += 1 +@handler.multiple_subscribe([RegularEvent, OnceOnlyEvent]) +class SmartHunter(Hunter): + def __init__(self, events): + global counter + counter += 1 def test_subscribe_mechanism(): global counter - + # first test normal subscribe and publish works handler.publish_event(RegularEvent()) handler.publish_event(RegularEvent()) handler.publish_event(RegularEvent()) - + time.sleep(0.02) assert counter == 3 counter = 0 - + + # testing the multiple subscription mechanism + handler.publish_event(OnceOnlyEvent()) + + time.sleep(0.02) + # SmartHunter and OnceHunter should run only. + assert counter == 2 + counter = 0 + # testing the subscribe_once mechanism handler.publish_event(OnceOnlyEvent()) handler.publish_event(OnceOnlyEvent()) handler.publish_event(OnceOnlyEvent()) time.sleep(0.02) - # should have been triggered once - assert counter == 1 - \ No newline at end of file + # should have been triggered once before, so never here. + assert counter == 0 +