diff --git a/doorman/extensions.py b/doorman/extensions.py index a437f78..9cbef23 100644 --- a/doorman/extensions.py +++ b/doorman/extensions.py @@ -49,8 +49,8 @@ def handle_result(self, data, **kwargs): class RuleManager(object): def __init__(self, app=None): - self.loaded_rules = False self.network = None + self.last_update = None if app is not None: self.init_app(app) @@ -70,19 +70,31 @@ def load_alerters(self): alerters = self.app.config.get('DOORMAN_ALERTER_PLUGINS', {}) self.alerters = {} - with self.app.app_context(): - for name, (plugin, config) in alerters.items(): - package, classname = plugin.rsplit('.', 1) - module = import_module(package) - klass = getattr(module, classname, None) + for name, (plugin, config) in alerters.items(): + package, classname = plugin.rsplit('.', 1) + module = import_module(package) + klass = getattr(module, classname, None) - if klass is None: - raise ValueError('Could not find a class named "{0}" in package "{1}"'.format(classname, package)) + if klass is None: + raise ValueError('Could not find a class named "{0}" in package "{1}"'.format(classname, package)) + + if not issubclass(klass, AbstractAlerterPlugin): + raise ValueError('{0} is not a subclass of AbstractAlerterPlugin'.format(name)) - if not issubclass(klass, AbstractAlerterPlugin): - raise ValueError('{0} is not a subclass of AbstractAlerterPlugin'.format(name)) + self.alerters[name] = klass(config) - self.alerters[name] = klass(config) + def should_reload_rules(self): + """ Checks if we need to reload the set of rules. """ + from doorman.models import Rule + + if self.last_update is None: + return True + + newest_rule = Rule.query.order_by(Rule.updated_at.desc()).limit(1).first() + if self.last_update < newest_rule.updated_at: + return True + + return False def load_rules(self): """ Load rules from the database. """ @@ -90,16 +102,10 @@ def load_rules(self): from doorman.models import Rule from sqlalchemy.exc import SQLAlchemyError - self.rules = [] - with self.app.app_context(): - try: - all_rules = list(Rule.query.all()) - except SQLAlchemyError: - # Ignore DB errors when testing - if self.app.config['TESTING']: - all_rules = [] - else: - raise + if not self.should_reload_rules(): + return + + all_rules = list(Rule.query.all()) self.network = Network() for rule in all_rules: @@ -111,16 +117,19 @@ def load_rules(self): # Create the rule. self.network.parse_query(rule.conditions, alerters=rule.alerters, rule_id=rule.id) + # Save the last updated date + # Note: we do this here, and not in should_reload_rules, because it's + # possible that we've reloaded a rule in between the two functions, and + # thus we accidentally don't reload when we should. + self.last_update = max(r.updated_at for r in all_rules) + def handle_log_entry(self, entry, node): """ The actual entrypoint for handling input log entries. """ from doorman.models import Rule from doorman.rules import RuleMatch from doorman.utils import extract_results - # Need to lazy-load rules - if not self.loaded_rules: - self.load_rules() - self.loaded_rules = True + self.load_rules() to_trigger = [] for name, action, columns, timestamp in extract_results(entry): diff --git a/doorman/manage/views.py b/doorman/manage/views.py index 7c6c6bb..35b7932 100644 --- a/doorman/manage/views.py +++ b/doorman/manage/views.py @@ -1,5 +1,6 @@ # -*- coding: utf-8 -*- import json +import datetime as dt from flask import Blueprint, current_app, flash, jsonify, redirect, render_template, request, url_for from flask_login import login_required @@ -23,7 +24,6 @@ DistributedQuery, DistributedQueryTask, FilePath, Node, Pack, Query, Tag, Rule, StatusLog ) -from doorman.tasks import reload_rules from doorman.utils import ( create_query_pack_from_upload, flash_errors, get_paginate_options ) @@ -469,9 +469,9 @@ def add_rule(): rule = Rule(name=form.name.data, alerters=form.alerters.data, description=form.description.data, - conditions=form.conditions.data) + conditions=form.conditions.data, + updated_at=dt.datetime.utcnow()) rule.save() - reload_rules.delay() return redirect(url_for('manage.rule', rule_id=rule.id)) @@ -489,8 +489,8 @@ def rule(rule_id): rule = rule.update(name=form.name.data, alerters=form.alerters.data, description=form.description.data, - conditions=form.conditions.data) - reload_rules.delay() + conditions=form.conditions.data, + updated_at=dt.datetime.utcnow()) return redirect(url_for('manage.rule', rule_id=rule.id)) form = UpdateRuleForm(request.form, obj=rule) diff --git a/doorman/models.py b/doorman/models.py index a4a3d4f..a11af40 100644 --- a/doorman/models.py +++ b/doorman/models.py @@ -445,12 +445,14 @@ class Rule(SurrogatePK, Model): alerters = Column(ARRAY(db.String), nullable=False) description = Column(db.String, nullable=True) conditions = Column(JSONB) + updated_at = Column(db.DateTime, nullable=False, default=dt.datetime.utcnow) - def __init__(self, name, alerters, description=None, conditions=None): + def __init__(self, name, alerters, description=None, conditions=None, updated_at=None): self.name = name self.description = description self.alerters = alerters self.conditions = conditions + self.updated_at = updated_at class User(UserMixin, SurrogatePK, Model): diff --git a/doorman/tasks.py b/doorman/tasks.py index 3a75a95..c018aac 100644 --- a/doorman/tasks.py +++ b/doorman/tasks.py @@ -12,12 +12,6 @@ def analyze_result(result, node): return -@celery.task() -def reload_rules(): - current_app.rule_manager.load_rules() - return - - @celery.task() def example_task(one, two): print('Adding {0} and {1}'.format(one, two)) diff --git a/migrations/versions/c17f01adbe31_add_updated_at_column_to_rule.py b/migrations/versions/c17f01adbe31_add_updated_at_column_to_rule.py new file mode 100644 index 0000000..e531d30 --- /dev/null +++ b/migrations/versions/c17f01adbe31_add_updated_at_column_to_rule.py @@ -0,0 +1,31 @@ +"""Add "updated_at" column to Rule + +Revision ID: c17f01adbe31 +Revises: b50c705fea80 +Create Date: 2016-05-27 15:51:58.168840 + +""" + +# revision identifiers, used by Alembic. +revision = 'c17f01adbe31' +down_revision = 'b50c705fea80' + +from alembic import op +import sqlalchemy as sa +import doorman.database + + +def upgrade(): + ### commands auto generated by Alembic - please adjust! ### + op.add_column('rule', + sa.Column('updated_at', sa.DateTime(), nullable=False, server_default=sa.func.now())) + ### end Alembic commands ### + + op.create_index('idx__rule__updated_at', 'rule', ['updated_at']) + + +def downgrade(): + op.drop_index('idx__rule__updated_at') + ### commands auto generated by Alembic - please adjust! ### + op.drop_column('rule', 'updated_at') + ### end Alembic commands ### diff --git a/tests/test_functional.py b/tests/test_functional.py index 399bf9d..58c48c3 100644 --- a/tests/test_functional.py +++ b/tests/test_functional.py @@ -980,20 +980,6 @@ class TestCreateTag: class TestAddRule: - def test_will_reload_rules(self, node, app, testapp): - from doorman.tasks import reload_rules - - with mock.patch.object(reload_rules, 'delay', return_value=None) as mock_delay: - resp = testapp.post(url_for('manage.add_rule'), { - 'name': 'Test Rule', - 'type': 'blacklist', - 'action': 'both', - 'alerters': 'debug', - 'config': '{"field_name": "foo", "blacklist": []}', - }) - - assert mock_delay.called - def test_supports_custom_operators(self, node, app, testapp): # Add a rule to the application rule = """ @@ -1047,60 +1033,83 @@ def test_supports_custom_operators(self, node, app, testapp): class TestUpdateRule: + pass - def test_will_reload_rules(self, db, node, app, testapp): - from doorman.tasks import reload_rules - rule_conds = { - "condition": "AND", - "rules": [ - { - "id": "query_name", - "field": "query_name", - "type": "string", - "input": "text", - "operator": "equal", - "value": "foo", - }, - ], +class TestRuleManager: + def test_will_load_rules_on_each_call(self, app, db): + """ + Verifies that each call to handle_log_entry will result in a call to load_rules. + """ + from doorman.rules import Network + + mgr = app.rule_manager + now = dt.datetime.utcnow() + + with mock.patch.object(mgr, 'load_rules', wraps=lambda: []) as mock_load_rules: + with mock.patch.object(mgr, 'network', wraps=Network()) as mock_network: + for i in range(0, 2): + mgr.handle_log_entry({ + 'data': [ + { + "diffResults": { + "added": [{'op': 'added'}], + "removed": "", + }, + "name": "fake", + "hostIdentifier": "hostname.local", + "calendarTime": "%s %s" % (now.ctime(), "UTC"), + "unixTime": now.strftime('%s') + } + ] + }, {'host_identifier': 'yes'}) + + assert mock_load_rules.call_count == 2 + + def test_will_reload_when_changed(self, app, db): + from doorman.models import Rule + + mgr = app.rule_manager + dummy_rule = { + "id": "query_name", + "field": "query_name", + "type": "string", + "input": "text", + "operator": "equal", + "value": "dummy-query", } - r = Rule( - name='Test-Rule', - alerters=['debug'], - description='A test rule', - conditions=rule_conds + now = dt.datetime.utcnow() + next = now + dt.timedelta(minutes=5) + + # Insert a first rule. + rule = Rule( + name='foo', + alerters=[], + conditions={'condition': 'AND', 'rules': [dummy_rule]}, + updated_at=now ) - db.session.add(r) + db.session.add(rule) + db.session.commit() + + # Verify that we will reload these rules + assert mgr.should_reload_rules() is True + + # Actually load them + mgr.load_rules() + + # Verify that (with no changes made), we should NOT reload. + assert mgr.should_reload_rules() is False + + # Make a change to a rule. + rule.update( + conditions={'condition': 'OR', 'rules': [dummy_rule]}, + updated_at=next) + db.session.add(rule) db.session.commit() - # Manually reload the rules here, and verify that we have the right - # rule in our list - app.rule_manager.load_rules() - assert len(app.rule_manager.network.conditions) == 2 - - condition_classes = [x.__class__.__name__ for x in app.rule_manager.network.conditions.values()] - assert sorted(condition_classes) == ['AndCondition', 'EqualCondition'] - - # Fake wrapper that just calls reload - def real_reload(*args, **kwargs): - app.rule_manager.load_rules() - - # Update the rule - rule_conds['condition'] = 'OR' - with mock.patch.object(reload_rules, 'delay', wraps=real_reload) as mock_delay: - resp = testapp.post(url_for('manage.rule', rule_id=r.id), { - 'name': 'Test-Rule', - 'alerters': 'debug', - 'conditions': json.dumps(rule_conds), - }) - - assert mock_delay.called - - # Trigger a manual reload again, and verify that it's been updated - app.rule_manager.load_rules() - condition_classes = [x.__class__.__name__ for x in app.rule_manager.network.conditions.values()] - assert sorted(condition_classes) == ['EqualCondition', 'OrCondition'] + # Verify that we will now reload + assert mgr.should_reload_rules() is True class TestRuleEndToEnd: