def test_compute_common_top_level(self): """Alert Class - Compute Common - No Nested Dictionaries""" record1 = {'a': 1, 'b': 2, 'c': 3} record2 = {'b': 2, 'c': 3, 'd': 4} record3 = {'c': 3, 'd': 4, 'e': 5} assert_equal({'c': 3}, Alert._compute_common([record1, record2, record3]))
def get_random_alert(key_count, rule_name, omit_rule_desc=False): """This loop generates key/value pairs with a key of length 6 and value of length 148. when formatted, each line should consume 160 characters, account for newline and asterisk for bold. For example: '*000001:* 6D829150B0154BF9BAC733FD25C61FA3D8CD3868AC2A92F19EEE119B 9CE8D6094966AA7592CE371002F1F7D82617673FCC9A9DB2A8F432AA791D74AB80BBCAD9\n' Therefore, 25*160 = 4000 character message size (exactly the 4000 limit) Anything over 4000 characters will result in multi-part slack messages: 55*160 = 8800 & 8800/4000 = ceil(2.2) = 3 messages needed """ # This default value is set in the rules engine rule_description = 'No rule description provided' if omit_rule_desc else 'rule test description' return Alert( rule_name, { '{:06}'.format(key): '{:0148X}'.format(random.randrange(16**128)) for key in range(key_count) }, {'slack:unit_test_channel'}, rule_description=rule_description )
def test_create_from_dynamo_record(self): """Alert Class - Create Alert from Dynamo Record""" alert = self._customized_alert() # Converting to a Dynamo record and back again should result in the exact same alert record = alert.dynamo_record() new_alert = Alert.create_from_dynamo_record(record) assert_equal(alert.dynamo_record(), new_alert.dynamo_record())
def test_merge_groups_too_recent(self): """Alert Merger - Alert Collection - All Alerts Too Recent to Merge""" alerts = [ Alert('', {'key': True}, set(), merge_by_keys=['key'], merge_window=timedelta(minutes=10)) ] assert_equal([], main.AlertMerger._merge_groups(alerts))
def test_add_mergeable(self): """Alert Merger - Merge Group - Add Alert to Group""" alert = Alert('', {'key': True}, set(), merge_by_keys=['key'], merge_window=timedelta(minutes=5)) group = main.AlertMergeGroup(alert) assert_true(group.add(alert)) # An alert can always merge with itself assert_equal([alert, alert], group.alerts)
def test_compute_common_partial_nested(self): """Alert Class - Compute Common - Some Common Features in Nested Dictionary""" # This is the example given in the docstring record1 = {'abc': 123, 'nested': {'A': 1, 'B': 2}} record2 = {'abc': 123, 'def': 456, 'nested': {'A': 1}} assert_equal({ 'abc': 123, 'nested': { 'A': 1 } }, Alert._compute_common([record1, record2]))
def test_merge_groups_single(self): """Alert Merger - Alert Collection - Single Merge Group""" alerts = [ Alert('', {'key': True}, set(), created=datetime(year=2000, month=1, day=1), merge_by_keys=['key'], merge_window=timedelta(minutes=5)), Alert('', { 'key': True, 'other': True }, set(), created=datetime(year=2000, month=1, day=1), merge_by_keys=['key'], merge_window=timedelta(minutes=5)) ] groups = main.AlertMerger._merge_groups(alerts) assert_equal(1, len(groups)) assert_equal(alerts, groups[0].alerts)
def test_dispatch(self, mock_logger): """Alert Merger - Dispatch to Alert Processor Lambda""" self.merger.lambda_client = MagicMock() self.merger.table.add_alerts([ # An alert without any merge criteria Alert('no_merging', {}, {'output'}), # 2 Alerts which will be merged (and will be be too large to send the entire record) Alert('merge_me', {'key': True}, {'output'}, created=datetime(year=2000, month=1, day=1), merge_by_keys=['key'], merge_window=timedelta(minutes=5)), Alert('merge_me', { 'key': True, 'other': 'abc' * 50 }, {'output'}, created=datetime(year=2000, month=1, day=1, minute=1), merge_by_keys=['key'], merge_window=timedelta(minutes=5)), # Alert which has already sent successfully (will be deleted) Alert('already_sent', {}, {'output'}, outputs_sent={'output'}) ]) self.merger.dispatch() # NOTE (Bobby): The following assertion was modified during the py2 -> py3 # conversion to disregard order of calls. mock_logger.assert_has_calls([ call.info('Merged %d alerts into a new alert with ID %s', 2, ANY), call.info('Dispatching %s to %s (attempt %d)', ANY, _ALERT_PROCESSOR, 1), call.info('Dispatching %s to %s (attempt %d)', ANY, _ALERT_PROCESSOR, 1) ], any_order=True)
def test_merge_groups_limit_reached(self): """Alert Merger - Alert Collection - Max Alerts Per Group""" alerts = [ Alert('same_rule_name', {'key': 'A'}, set(), created=datetime(year=2000, month=1, day=1), merge_by_keys=['key'], merge_window=timedelta(minutes=5)), ] * 5 # Since max alerts per group is 2, it should create 3 merged groups. groups = main.AlertMerger._merge_groups(alerts) assert_equal(3, len(groups)) assert_equal(alerts[0:2], groups[0].alerts) assert_equal(alerts[2:4], groups[1].alerts) assert_equal(alerts[4:], groups[2].alerts)
def test_can_merge_merge_keys_absent(self): """Alert Class - Can Merge - True if Merge Keys Do Not Exist in Either Record""" alert1 = Alert('', {}, set(), merge_by_keys=['key'], merge_window=timedelta(minutes=10)) alert2 = Alert('', {}, set(), merge_by_keys=['key'], merge_window=timedelta(minutes=10)) assert_true(alert1.can_merge(alert2)) assert_true(alert2.can_merge(alert1))
def test_can_merge_different_values(self): """Alert Class - Can Merge - False if Merge Key has Different Values""" alert1 = Alert('', {'key': True}, set(), merge_by_keys=['key'], merge_window=timedelta(minutes=10)) alert2 = Alert('', {'key': False}, set(), merge_by_keys=['key'], merge_window=timedelta(minutes=10)) assert_false(alert1.can_merge(alert2)) assert_false(alert2.can_merge(alert1))
def test_can_merge_key_not_common(self): """Alert Class - Can Merge - False if Merge Key Not Present in Both Records""" alert1 = Alert('', {'key': True}, set(), merge_by_keys=['key'], merge_window=timedelta(minutes=10)) alert2 = Alert('', {'other': True}, set(), merge_by_keys=['key'], merge_window=timedelta(minutes=10)) assert_false(alert1.can_merge(alert2)) assert_false(alert2.can_merge(alert1))
def test_can_merge_different_merge_keys(self): """Alert Class - Can Merge - False if Different Merge Keys Defined""" alert1 = Alert('', {'key': True}, set(), merge_by_keys=['key'], merge_window=timedelta(minutes=10)) alert2 = Alert('', {'key': True}, set(), merge_by_keys=['other'], merge_window=timedelta(minutes=10)) assert_false(alert1.can_merge(alert2)) assert_false(alert2.can_merge(alert1))
def test_alert_generator(self, mock_logger): """Alert Merger - Sorted Alerts - Invalid Alerts are Logged""" records = [ Alert('test_rule', {}, {'output'}).dynamo_record(), { 'Nonsense': 'Record' } ] with patch.object(self.merger.table, 'get_alert_records', return_value=records): result = list(self.merger._alert_generator('test_rule')) # Valid record is returned assert_equal(1, len(result)) assert_equal(records[0]['AlertID'], result[0].alert_id) # Invalid record logs an exception mock_logger.exception.assert_called_once_with( 'Invalid alert record %s', records[1])
def test_compute_common_many_nested(self): """Alert Class - Compute Common - Multiple Levels of Nesting""" record1 = {'a': {'b': {'c': 3, 'd': 4}, 'e': {'h': {'i': 9}}, 'j': {}}} record2 = { 'a': { 'b': { 'c': 3, }, 'e': { 'f': { 'g': 8 }, 'h': {} }, 'j': {} } } expected = {'a': {'b': {'c': 3}, 'j': {}}} assert_equal(expected, Alert._compute_common([record1, record2]))
def dispatch(self): """Find and dispatch all pending alerts to the alert processor.""" # To reduce the API calls to Dynamo, batch all additions and deletions until the end. merged_alerts = [] # List of newly created merge alerts alerts_to_delete = [] # List of alerts which can be deleted for rule_name in self.table.rule_names_generator(): merge_enabled_alerts = [] for alert in self._alert_generator(rule_name): if alert.remaining_outputs: # If an alert still has pending outputs, it needs to be sent immediately. # For example, all alerts are sent to the default firehose now even if they will # later be merged when sending to other outputs. self._dispatch_alert(alert) elif alert.merge_enabled: # This alert has finished sending to non-merged outputs; it is now a candidate # for alert merging. merge_enabled_alerts.append(alert) else: # This alert has sent successfully but doesn't need to be merged. # It should have been deleted by the alert processor, but we can do it now. alerts_to_delete.append(alert) for group in self._merge_groups(merge_enabled_alerts): # Create a new merged Alert. new_alert = Alert.merge(group.alerts) LOGGER.info('Merged %d alerts into a new alert with ID %s', len(group.alerts), new_alert.alert_id) merged_alerts.append(new_alert) # Since we already guaranteed that the original alerts have sent to the unmerged # outputs (e.g. default firehose), they can be safely marked for deletion. alerts_to_delete.extend(group.alerts) if merged_alerts: # Add new merged alerts to the alerts table and send them to the alert processor. self.table.add_alerts(merged_alerts) for alert in merged_alerts: self._dispatch_alert(alert) if alerts_to_delete: self.table.delete_alerts([(alert.rule_name, alert.alert_id) for alert in alerts_to_delete])
def _customized_alert(): return Alert('test_rule', {'abc': 123}, { 'aws-firehose:alerts', 'aws-sns:test-output', 'aws-s3:other-output' }, alert_id='abc-123', attempts=1, cluster='', context={'rule': 'context'}, created=datetime.utcnow(), dispatched=datetime.utcnow(), log_source='source', log_type='csv', merge_by_keys=['abc'], merge_window=timedelta(minutes=5), outputs_sent={'aws-sns:test-output'}, rule_description='A Test Rule', source_entity='entity', source_service='s3', staged=True)
def test_can_merge_too_far_apart(self): """Alert Class - Can Merge - False if Outside Merge Window""" alert1 = Alert('', {'key': True}, set(), created=datetime(year=2000, month=1, day=1, minute=0), merge_by_keys=['key'], merge_window=timedelta(minutes=10)) alert2 = Alert('', {'key': True}, set(), created=datetime(year=2000, month=1, day=1, minute=11), merge_by_keys=['key'], merge_window=timedelta(minutes=10)) assert_false(alert1.can_merge(alert2)) assert_false(alert2.can_merge(alert1))
def _alert_generator(self, rule_name): """ Returns a generator that yields Alert instances triggered from the given rule name. To limit memory consumption, the generator yields a maximum number of alerts, defined by self._alert_generator_limit. """ generator = self.table.get_alert_records(rule_name, self.alert_proc_timeout) for idx, record in enumerate(generator, start=1): try: yield Alert.create_from_dynamo_record(record) except AlertCreationError: LOGGER.exception('Invalid alert record %s', record) continue if idx >= self._alert_generator_limit: LOGGER.warning( 'Alert Merger reached alert limit of %d for rule "%s"', self._alert_generator_limit, rule_name) return
def test_can_merge_true(self): """Alert Class - Can Merge - True Result""" alert1 = Alert('', {'key': True}, set(), created=datetime(year=2000, month=1, day=1, minute=0), merge_by_keys=['key'], merge_window=timedelta(minutes=10)) alert2 = Alert('', { 'key': True, 'other': True }, set(), created=datetime(year=2000, month=1, day=1, minute=10), merge_by_keys=['key'], merge_window=timedelta(minutes=10)) assert_true(alert1.can_merge(alert2)) assert_true(alert2.can_merge(alert1))
def get_alert(context=None): """This function generates a sample alert for testing purposes Args: context (dict): Optional alert context """ return Alert( 'cb_binarystore_file_added', { 'compressed_size': '9982', 'timestamp': '1496947381.18', 'node_id': '1', 'cb_server': 'cbserver', 'size': '21504', 'type': 'binarystore.file.added', 'file_path': '/tmp/5DA/AD8/0F9AA55DA3BDE84B35656AD8911A22E1.zip', 'md5': '0F9AA55DA3BDE84B35656AD8911A22E1' }, {'slack:unit_test_channel'}, alert_id='79192344-4a6d-4850-8d06-9c3fef1060a4', context=context, log_source='carbonblack:binarystore.file.added', log_type='json', rule_description='Info about this rule and what actions to take', source_entity='corp-prefix.prod.cb.region', source_service='s3')
def run(self, event): """Run the alert processor! Args: event (dict): Lambda invocation event containing at least the rule name and alert ID. Returns: dict: Maps output (str) to whether it sent successfully (bool). An empty dict is returned if the Alert was improperly formatted. """ # Grab the alert record from Dynamo (if needed). if set(event) == {'AlertID', 'RuleName'}: LOGGER.info('Retrieving %s from alerts table', event) alert_record = self.alerts_table.get_alert_record( event['RuleName'], event['AlertID']) if not alert_record: LOGGER.error('%s does not exist in the alerts table', event) return {} else: alert_record = event # Convert record to an Alert instance. try: alert = Alert.create_from_dynamo_record(alert_record) except AlertCreationError: LOGGER.exception('Invalid alert %s', event) return {} # Remove normalization key from the record. # TODO: Consider including this in at least some outputs, e.g. default Athena firehose if Normalizer.NORMALIZATION_KEY in alert.record: del alert.record[Normalizer.NORMALIZATION_KEY] result = self._send_to_outputs(alert) self._update_table(alert, result) return result
def test_merge(self): """Alert Class - Merge - Create Merged Alert""" # Example based on a CarbonBlack log record1 = { 'alliance_data_virustotal': [], 'alliance_link_virustotal': '', 'alliance_score_virustotal': 0, 'cmdline': 'whoami', 'comms_ip': '1.2.3.4', 'hostname': 'my-computer-name', 'path': '/usr/bin/whoami', 'streamalert:ioc': { 'hello': 'world' }, 'timestamp': 1234.5678, 'username': '******' } alert1 = Alert('RuleName', record1, {'aws-sns:topic'}, created=datetime(year=2000, month=1, day=1), merge_by_keys=['hostname', 'username'], merge_window=timedelta(minutes=5)) # Second alert has slightly different record and different outputs record2 = copy.deepcopy(record1) record2['streamalert:ioc'] = {'goodbye': 'world'} record2['timestamp'] = 9999 alert2 = Alert('RuleName', record2, {'slack:channel'}, created=datetime(year=2000, month=1, day=2), merge_by_keys=['hostname', 'username'], merge_window=timedelta(minutes=5)) merged = Alert.merge([alert1, alert2]) assert_is_instance(merged, Alert) assert_equal({'slack:channel'}, merged.outputs) # Most recent outputs were used expected_record = { 'AlertCount': 2, 'AlertTimeFirst': '2000-01-01T00:00:00.000000Z', 'AlertTimeLast': '2000-01-02T00:00:00.000000Z', 'MergedBy': { 'hostname': 'my-computer-name', 'username': '******' }, 'OtherCommonKeys': { 'alliance_data_virustotal': [], 'alliance_link_virustotal': '', 'alliance_score_virustotal': 0, 'cmdline': 'whoami', 'comms_ip': '1.2.3.4', 'path': '/usr/bin/whoami', }, 'ValueDiffs': { '2000-01-01T00:00:00.000000Z': { 'streamalert:ioc': { 'hello': 'world' }, 'timestamp': 1234.5678 }, '2000-01-02T00:00:00.000000Z': { 'streamalert:ioc': { 'goodbye': 'world' }, 'timestamp': 9999 } } } assert_equal(expected_record, merged.record)
def test_compute_diff_top_level(self): """Alert Class - Compute Diff - Top Level Keys""" common = {'c': 3} record = {'a': 1, 'b': 2, 'c': 3} assert_equal({'a': 1, 'b': 2}, Alert._compute_diff(common, record))
def test_merge_groups_complex(self): """Alert Merger - Alert Collection - Complex Merge Groups""" alerts = [ # Merge group 1 - key 'A' minutes 0-5 Alert('same_rule_name', {'key': 'A'}, set(), created=datetime(year=2000, month=1, day=1), merge_by_keys=['key'], merge_window=timedelta(minutes=5)), Alert('same_rule_name', {'key': 'A'}, set(), created=datetime(year=2000, month=1, day=1, minute=1), merge_by_keys=['key'], merge_window=timedelta(minutes=5)), # Merge group 2 - Key B minutes 0-5 Alert('same_rule_name', {'key': 'B'}, set(), created=datetime(year=2000, month=1, day=1, minute=2), merge_by_keys=['key'], merge_window=timedelta(minutes=5)), Alert('same_rule_name', {'key': 'B'}, set(), created=datetime(year=2000, month=1, day=1, minute=2, second=30), merge_by_keys=['key'], merge_window=timedelta(minutes=5)), Alert('same_rule_name', {'key': 'B'}, set(), created=datetime(year=2000, month=1, day=1, minute=3), merge_by_keys=['key'], merge_window=timedelta(minutes=5)), # Merge group 3 - Different merge keys Alert('same_rule_name', { 'key': 'A', 'other': 'B' }, set(), created=datetime(year=2000, month=1, day=1, minute=4), merge_by_keys=['key', 'other'], merge_window=timedelta(minutes=5)), Alert('same_rule_name', { 'key': 'A', 'other': 'B' }, set(), created=datetime(year=2000, month=1, day=1, minute=5), merge_by_keys=['key', 'other'], merge_window=timedelta(minutes=5)), # Merge group 4 - key A minutes 50-55 Alert('same_rule_name', {'key': 'A'}, set(), created=datetime(year=2000, month=1, day=1, minute=50), merge_by_keys=['key'], merge_window=timedelta(minutes=5)), # This alert (created now) is too recent to fit in any merge group. Alert('same_rule_name', {'key': 'A'}, set(), merge_by_keys=['key'], merge_window=timedelta(minutes=10)) ] groups = main.AlertMerger._merge_groups(alerts) assert_equal(4, len(groups)) assert_equal(alerts[0:2], groups[0].alerts) assert_equal(alerts[2:5], groups[1].alerts) assert_equal(alerts[5:7], groups[2].alerts) assert_equal([alerts[7]], groups[3].alerts)
def test_compute_diff_different_types(self): """Alert Class - Compute Diff - Type Mismatch Short-Circuits Recursion""" common = {'b': 2} record = {'a': 1, 'b': {'nested': 'stuff'}} assert_equal(record, Alert._compute_diff(common, record))
def test_compute_diff_nested(self): """Alert Class - Compute Diff - Difference in Nested Dictionary""" # This is the example given in the docstring common = {'abc': 123, 'nested': {'A': 1}} record = {'abc': 123, 'nested': {'A': 1, 'B': 2}} assert_equal({'nested': {'B': 2}}, Alert._compute_diff(common, record))
class TestAlertProcessor: """Tests for alert_processor/main.py""" # pylint: disable=no-member,no-self-use,protected-access @patch('streamalert.alert_processor.main.load_config', Mock(return_value=load_config('tests/unit/conf/', validate=True))) @patch.dict(os.environ, MOCK_ENV) @patch.object(AlertProcessor, 'BACKOFF_MAX_TRIES', 1) @patch('streamalert.alert_processor.main.AlertTable', MagicMock()) def setup(self): """Alert Processor - Test Setup""" # pylint: disable=attribute-defined-outside-init self.processor = AlertProcessor() self.alert = Alert('hello_world', { 'abc': 123, Normalizer.NORMALIZATION_KEY: {} }, {'slack:unit-test-channel'}) def test_init(self): """Alert Processor - Initialization""" assert_is_instance(self.processor.config, dict) @patch('streamalert.alert_processor.main.LOGGER') def test_create_dispatcher_invalid(self, mock_logger): """Alert Processor - Create Dispatcher - Invalid Output""" assert_is_none(self.processor._create_dispatcher('helloworld')) mock_logger.error.called_once_with(ANY, 'helloworld') @patch('streamalert.alert_processor.main.LOGGER') def test_create_dispatcher_output_doesnt_exist(self, mock_logger): """Alert Processor - Create Dispatcher - Output Does Not Exist""" assert_is_none( self.processor._create_dispatcher('slack:no-such-channel')) mock_logger.error.called_once_with('The output \'%s\' does not exist!', 'slack:no-such-channel') @patch.dict(os.environ, MOCK_ENV) def test_create_dispatcher(self): """Alert Processor - Create Dispatcher - Success""" dispatcher = self.processor._create_dispatcher( 'aws-s3:unit_test_bucket') assert_is_instance(dispatcher, OutputDispatcher) @patch.object(AlertProcessor, '_create_dispatcher') def test_send_alerts_success(self, mock_create_dispatcher): """Alert Processor - Send Alerts Success""" mock_create_dispatcher.return_value.dispatch.return_value = True result = self.processor._send_to_outputs(self.alert) mock_create_dispatcher.assert_called_once() mock_create_dispatcher.return_value.dispatch.assert_called_once() assert_equal({'slack:unit-test-channel': True}, result) assert_equal(self.alert.outputs, self.alert.outputs_sent) @patch.object(AlertProcessor, '_create_dispatcher') def test_send_alerts_failure(self, mock_create_dispatcher): """Alert Processor - Send Alerts Failure""" mock_create_dispatcher.return_value.dispatch.return_value = False result = self.processor._send_to_outputs(self.alert) mock_create_dispatcher.assert_called_once() mock_create_dispatcher.return_value.dispatch.assert_called_once() assert_equal({'slack:unit-test-channel': False}, result) assert_equal(set(), self.alert.outputs_sent) @patch.object(AlertProcessor, '_create_dispatcher', return_value=None) def test_send_alerts_skip_invalid_outputs(self, mock_create_dispatcher): """Alert Processor - Send Alerts With Invalid Outputs""" result = self.processor._send_to_outputs(self.alert) mock_create_dispatcher.assert_called_once() assert_equal({'slack:unit-test-channel': False}, result) def test_update_alerts_table_none(self): """Alert Processor - Update Alerts Table - Empty Results""" self.processor.alerts_table.delete_alert = MagicMock() self.processor.alerts_table.update_retry_outputs = MagicMock() self.processor._update_table(self.alert, {}) self.processor.alerts_table.delete_alert.assert_not_called() self.processor.alerts_table.update_retry_outputs.assert_not_called() def test_update_alerts_table_delete(self): """Alert Processor - Update Alerts Table - Delete Item""" self.processor._update_table(self.alert, {'out1': True, 'out2': True}) self.processor.alerts_table.delete_alerts.assert_called_once_with([ (self.alert.rule_name, self.alert.alert_id) ]) def test_update_alerts_table_update(self): """Alert Processor - Update Alerts Table - Update With Failed Outputs""" self.processor._update_table(self.alert, { 'out1': True, 'out2': False, 'out3': False }) self.processor.alerts_table.update_sent_outputs.assert_called_once_with( self.alert) @patch.object(AlertProcessor, '_send_to_outputs', return_value={'slack:unit-test-channel': True}) @patch.object(AlertProcessor, '_update_table') def test_run_full_event(self, mock_send_alerts, mock_update_table): """Alert Processor - Run With the Full Alert Record""" result = self.processor.run(self.alert.dynamo_record()) assert_equal({'slack:unit-test-channel': True}, result) mock_send_alerts.assert_called_once() mock_update_table.assert_called_once() @patch('streamalert.alert_processor.main.LOGGER') def test_run_invalid_alert(self, mock_logger): """Alert Processor - Run With an Invalid Alert""" result = self.processor.run({'Record': 'Nonsense'}) assert_equal({}, result) mock_logger.exception.called_once_with('Invalid alert %s', {'Record': 'Nonsense'}) @patch.object(AlertProcessor, '_send_to_outputs', return_value={'slack:unit-test-channel': True}) @patch.object(AlertProcessor, '_update_table') def test_run_get_alert_from_dynamo(self, mock_send_alerts, mock_update_table): """Alert Processor - Run With Just the Alert Key""" self.processor.alerts_table.get_alert_record = MagicMock( return_value=self.alert.dynamo_record()) result = self.processor.run(self.alert.dynamo_key) assert_equal({'slack:unit-test-channel': True}, result) self.processor.alerts_table.get_alert_record.assert_called_once_with( self.alert.rule_name, self.alert.alert_id) mock_send_alerts.assert_called_once() mock_update_table.assert_called_once() @patch('streamalert.alert_processor.main.LOGGER') def test_run_alert_does_not_exist(self, mock_logger): """Alert Processor - Run - Alert Does Not Exist""" self.processor.alerts_table.get_alert_record = MagicMock( return_value=None) self.processor.run(self.alert.dynamo_key) mock_logger.error.assert_called_once_with( '%s does not exist in the alerts table', self.alert.dynamo_key) @patch.dict(os.environ, MOCK_ENV) @patch.object(AlertProcessor, 'run', return_value={'output': True}) def test_handler(self, mock_run): """Alert Processor - Lambda Handler""" event = {'AlertID': 'abc', 'RuleName': 'hello_world'} result = handler(event, None) assert_equal({'output': True}, result) mock_run.assert_called_once_with(event)
def create_table(table, bucket, config, schema_override=None): """Create a 'streamalert' Athena table Args: table (str): The name of the table being rebuilt bucket (str): The s3 bucket to be used as the location for Athena data table_type (str): The type of table being refreshed config (CLIConfig): Loaded StreamAlert config schema_override (set): An optional set of key=value pairs to be used for overriding the configured column_name=value_type. Returns: bool: False if errors occurred, True otherwise """ enabled_logs = FirehoseClient.load_enabled_log_sources( config['global']['infrastructure']['firehose'], config['logs']) # Convert special characters in schema name to underscores sanitized_table_name = FirehoseClient.firehose_log_name(table) # Check that the log type is enabled via Firehose if sanitized_table_name != 'alerts' and sanitized_table_name not in enabled_logs: LOGGER.error( 'Table name %s missing from configuration or ' 'is not enabled.', sanitized_table_name) return False athena_client = get_athena_client(config) config_data_bucket = firehose_data_bucket(config) if not config_data_bucket: LOGGER.error('The \'firehose\' module is not enabled in global.json') return False # Check if the table exists if athena_client.check_table_exists(sanitized_table_name): LOGGER.info('The \'%s\' table already exists.', sanitized_table_name) return False if table == 'alerts': # get a fake alert so we can get the keys needed and their types alert = Alert('temp_rule_name', {}, {}) output = alert.output_dict() schema = record_to_schema(output) athena_schema = helpers.logs_schema_to_athena_schema(schema) # Use the bucket if supplied, otherwise use the default alerts bucket bucket = bucket or firehose_alerts_bucket(config) query = _construct_create_table_statement(schema=athena_schema, table_name=table, bucket=bucket) else: # all other tables are log types # Use the bucket if supplied, otherwise use the default data bucket bucket = bucket or config_data_bucket log_info = config['logs'][table.replace('_', ':', 1)] schema = dict(log_info['schema']) sanitized_schema = FirehoseClient.sanitize_keys(schema) athena_schema = helpers.logs_schema_to_athena_schema(sanitized_schema) # Add envelope keys to Athena Schema configuration_options = log_info.get('configuration') if configuration_options: envelope_keys = configuration_options.get('envelope_keys') if envelope_keys: sanitized_envelope_key_schema = FirehoseClient.sanitize_keys( envelope_keys) # Note: this key is wrapped in backticks to be Hive compliant athena_schema[ '`streamalert:envelope_keys`'] = helpers.logs_schema_to_athena_schema( sanitized_envelope_key_schema) # Handle Schema overrides # This is useful when an Athena schema needs to differ from the normal log schema if schema_override: for override in schema_override: column_name, column_type = override.split('=') # Columns are escaped to avoid Hive issues with special characters column_name = '`{}`'.format(column_name) if column_name in athena_schema: athena_schema[column_name] = column_type LOGGER.info('Applied schema override: %s:%s', column_name, column_type) else: LOGGER.error( 'Schema override column %s not found in Athena Schema, skipping', column_name) query = _construct_create_table_statement( schema=athena_schema, table_name=sanitized_table_name, bucket=bucket) success = athena_client.run_query(query=query) if not success: LOGGER.error('The %s table could not be created', sanitized_table_name) return False # Update the CLI config if table != 'alerts' and bucket != config_data_bucket: # Only add buckets to the config if they are not one of the default/configured buckets # Ensure 'buckets' exists in the config (since it is not required) config['lambda']['athena_partition_refresh_config']['buckets'] = ( config['lambda']['athena_partition_refresh_config'].get( 'buckets', {})) if bucket not in config['lambda']['athena_partition_refresh_config'][ 'buckets']: config['lambda']['athena_partition_refresh_config']['buckets'][ bucket] = 'data' config.write() LOGGER.info('The %s table was successfully created!', sanitized_table_name) return True
def test_merge_nested(self): """Alert Class - Merge - Merge with Nested Keys""" record1 = { 'NumMatchedRules': 1, 'FileInfo': { 'Deleted': None, 'Nested': [1, 2, 'three'] }, 'MatchedRules': { 'Rule1': 'MatchedStrings' } } alert1 = Alert('RuleName', record1, {'slack:channel'}, created=datetime(year=2000, month=1, day=1), merge_by_keys=['Nested'], merge_window=timedelta(minutes=5)) record2 = { 'NumMatchedRules': 2, 'FileInfo': { 'Deleted': None, 'Nested': [1, 2, 'three'] }, 'MatchedRules': { 'Rule1': 'MatchedStrings' } } alert2 = Alert('RuleName', record2, {'slack:channel'}, created=datetime(year=2000, month=1, day=2), merge_by_keys=['Nested'], merge_window=timedelta(minutes=5)) record3 = { 'MatchedRules': { 'Rule1': 'MatchedStrings' }, 'Nested': [1, 2, 'three'] # This is in a different place in the record } alert3 = Alert('RuleName', record3, {'slack:channel'}, created=datetime(year=2000, month=1, day=3), merge_by_keys=['Nested'], merge_window=timedelta(minutes=5)) merged = Alert.merge([alert1, alert2, alert3]) expected_record = { 'AlertCount': 3, 'AlertTimeFirst': '2000-01-01T00:00:00.000000Z', 'AlertTimeLast': '2000-01-03T00:00:00.000000Z', 'MergedBy': { 'Nested': [1, 2, 'three'] }, 'OtherCommonKeys': { 'MatchedRules': { 'Rule1': 'MatchedStrings' } }, 'ValueDiffs': { '2000-01-01T00:00:00.000000Z': { 'NumMatchedRules': 1, 'FileInfo': { 'Deleted': None } }, '2000-01-02T00:00:00.000000Z': { 'NumMatchedRules': 2, 'FileInfo': { 'Deleted': None } }, '2000-01-03T00:00:00.000000Z': {} } } assert_equal(expected_record, merged.record)