From db6ecfe81c9594baaa6ea94939e3d172cfba110c Mon Sep 17 00:00:00 2001 From: javier_marcos Date: Tue, 12 Dec 2017 15:29:17 -0800 Subject: [PATCH 1/2] Context can be modified in rule --- stream_alert/rule_processor/rules_engine.py | 5 ++- .../test_rules_engine.py | 38 +++++++++++++++++++ 2 files changed, 42 insertions(+), 1 deletion(-) diff --git a/stream_alert/rule_processor/rules_engine.py b/stream_alert/rule_processor/rules_engine.py index 0c4483601..e777b8dce 100644 --- a/stream_alert/rule_processor/rules_engine.py +++ b/stream_alert/rule_processor/rules_engine.py @@ -284,7 +284,10 @@ def process_rule(record, rule): (bool): The return function of the rule """ try: - rule_result = rule.rule_function(record) + if rule.context: + rule_result = rule.rule_function(record, rule.context) + else: + rule_result = rule.rule_function(record) except Exception: # pylint: disable=broad-except rule_result = False LOGGER.exception( diff --git a/tests/unit/stream_alert_rule_processor/test_rules_engine.py b/tests/unit/stream_alert_rule_processor/test_rules_engine.py index 3210618ba..421a3779e 100644 --- a/tests/unit/stream_alert_rule_processor/test_rules_engine.py +++ b/tests/unit/stream_alert_rule_processor/test_rules_engine.py @@ -852,3 +852,41 @@ def match_ipaddress(_): # pylint: disable=unused-variable payload = load_and_classify_payload(toggled_config, service, entity, raw_record) assert_equal(len(new_rules_engine.process(payload)), 1) + + def test_rule_modify_context(self): + """Rules Engine - Testing Context Modification""" + @rule(logs=['test_log_type_json_nested_with_data'], + outputs=['s3:sample_bucket'], + context={'assigned_user': 'not_set', 'assigned_policy': 'not_set2'}) + def modify_context_test(rec, context): # pylint: disable=unused-variable + """Modify context rule""" + context['assigned_user'] = 'valid_user' + context['assigned_policy'] = 'valid_policy' + return rec['application'] == 'web-app' + + kinesis_data = json.dumps({ + 'date': 'Dec 01 2016', + 'unixtime': '1483139547', + 'host': 'host1.web.prod.net', + 'application': 'web-app', + 'environment': 'prod', + 'data': { + 'category': 'web-server', + 'type': '1', + 'source': 'eu' + } + }) + + # prepare the payloads + service, entity = 'kinesis', 'test_kinesis_stream' + raw_record = make_kinesis_raw_record(entity, kinesis_data) + payload = load_and_classify_payload(self.config, service, entity, raw_record) + + # process payloads + alerts = self.rules_engine.process(payload) + + print alerts + + # alert tests + assert_equal(alerts[0]['context']['assigned_user'], 'valid_user') + assert_equal(alerts[0]['context']['assigned_policy'], 'valid_policy') From 308bac491807d2819228ae37640e8cdda481a9b5 Mon Sep 17 00:00:00 2001 From: javier_marcos Date: Tue, 12 Dec 2017 15:47:19 -0800 Subject: [PATCH 2/2] Removed forgotten debug print --- tests/unit/stream_alert_rule_processor/test_rules_engine.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/tests/unit/stream_alert_rule_processor/test_rules_engine.py b/tests/unit/stream_alert_rule_processor/test_rules_engine.py index 421a3779e..78f429c57 100644 --- a/tests/unit/stream_alert_rule_processor/test_rules_engine.py +++ b/tests/unit/stream_alert_rule_processor/test_rules_engine.py @@ -885,8 +885,6 @@ def modify_context_test(rec, context): # pylint: disable=unused-variable # process payloads alerts = self.rules_engine.process(payload) - print alerts - # alert tests assert_equal(alerts[0]['context']['assigned_user'], 'valid_user') assert_equal(alerts[0]['context']['assigned_policy'], 'valid_policy')