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)
