def test_strip_successful_records(self):
"""StreamAlertFirehose - Strip Successful Records"""
batch = [{
'test': 'success'
}, {
'test': 'data'
}, {
'other': 'failure'
}, {
'other': 'info'
}]
response = {
'FailedPutCount':
1,
'RequestResponses': [{
'RecordId': 'rec_id_00'
}, {
'RecordId': 'rec_id_01'
}, {
'ErrorCode': 10,
'ErrorMessage': 'foo'
}, {
'RecordId': 'rec_id_03'
}]
}
expected_batch = [{'other': 'failure'}]
StreamAlertFirehose._strip_successful_records(batch, response)
assert_equal(batch, expected_batch)
def test_limit_record_size(self, mock_logging):
"""StreamAlertFirehose - Record Size Check"""
test_events = [
# unit_test_simple_log
{
'unit_key_01': 1,
'unit_key_02': 'test' * 250001 # is 4 bytes higher than max
},
{
'unit_key_01': 2,
'unit_key_02': 'test'
},
# test_log_type_json_nested
{
'date': 'January 01, 3005',
'unixtime': '32661446400',
'host': 'my-host.name.website.com',
'data': {
'super': 'secret'
}
}
]
StreamAlertFirehose._limit_record_size(test_events)
assert_true(len(test_events), 2)
assert_true(mock_logging.error.called)
def test_record_delivery(self, mock_logging):
"""StreamAlertFirehose - Record Delivery"""
self.__sa_firehose = StreamAlertFirehose(region='us-east-1',
firehose_config={},
log_sources={})
# Add sample categorized payloads
for payload_type, logs in self._sample_categorized_payloads(
).iteritems():
self.__sa_firehose.categorized_payloads[payload_type].extend(logs)
# Setup mocked Delivery Streams
self._mock_delivery_streams([
'streamalert_data_test_log_type_json_nested',
'streamalert_data_unit_test_simple_log'
])
# Send the records
with patch.object(self.__sa_firehose._firehose_client,
'put_record_batch') as firehose_mock:
firehose_mock.return_value = {'FailedPutCount': 0}
self.__sa_firehose.send()
firehose_mock.assert_called()
assert_true(mock_logging.info.called)
def test_record_delivery_failure(self, mock_logging):
"""StreamAlertFirehose - Record Delivery - Failed PutRecord"""
self.__sa_firehose = StreamAlertFirehose(region='us-east-1',
firehose_config={},
log_sources={})
# Add sample categorized payloads
for payload_type, logs in self._sample_categorized_payloads(
).iteritems():
self.__sa_firehose.categorized_payloads[payload_type].extend(logs)
# Setup mocked Delivery Streams
self._mock_delivery_streams([
'streamalert_data_test_log_type_json_nested',
'streamalert_data_unit_test_simple_log'
])
# Send the records
with patch.object(self.__sa_firehose._firehose_client,
'put_record_batch') as firehose_mock:
firehose_mock.return_value = {
'FailedPutCount':
3,
'RequestResponses': [
{
'RecordId': '12345',
'ErrorCode': '300',
'ErrorMessage': 'Bad message!!!'
},
]
}
self.__sa_firehose.send()
firehose_mock.assert_called()
assert_true(mock_logging.error.called)
def test_load_enabled_sources(self):
"""StreamAlertFirehose - Load Enabled Sources"""
config = load_config('tests/unit/conf')
firehose_config = {
'enabled_logs':
['json:regex_key_with_envelope', 'test_cloudtrail', 'cloudwatch']
} # expands to 2 logs
sa_firehose = StreamAlertFirehose(region='us-east-1',
firehose_config=firehose_config,
log_sources=config['logs'])
assert_equal(len(sa_firehose._enabled_logs), 4)
# Make sure the subtitution works properly
assert_true(all([':' not in log for log in sa_firehose.enabled_logs]))
assert_false(sa_firehose.enabled_log_source('test_inspec'))
def test_sanitize_keys(self):
"""StreamAlertFirehose - Sanitize Keys"""
# test_log_type_json_nested
test_event = {
'date': 'January 01, 3005',
'unixtime': '32661446400',
'host': 'my-host.name.website.com',
'data': {
'super-duper': 'secret',
'sanitize_me': 1,
'example-key': 1,
'moar**data': 2,
'even.more': 3
}
}
expected_sanitized_event = {
'date': 'January 01, 3005',
'unixtime': '32661446400',
'host': 'my-host.name.website.com',
'data': {
'super_duper': 'secret',
'sanitize_me': 1,
'example_key': 1,
'moar__data': 2,
'even_more': 3
}
}
sanitized_event = StreamAlertFirehose.sanitize_keys(test_event)
assert_equal(sanitized_event, expected_sanitized_event)
def test_firehose_reset(self, mock_logging):
"""StreamAlertFirehose - Test Reset Firehose Client"""
def test_func():
pass
sa_firehose = StreamAlertFirehose(region='us-east-1',
firehose_config={},
log_sources={})
id_1 = id(sa_firehose._firehose_client)
sa_firehose._backoff_handler_firehose_reset({
'target': test_func,
'wait': '0.134315135',
'tries': 3
})
id_2 = id(sa_firehose._firehose_client)
assert_true(mock_logging.info.called)
assert_not_equal(id_1, id_2)
def test_record_delivery_client_error(self, mock_logging):
"""StreamAlertFirehose - Record Delivery - Client Error"""
sa_firehose = StreamAlertFirehose(region='us-east-1',
firehose_config={},
log_sources={})
test_events = [
# unit_test_simple_log
{
'unit_key_01': 2,
'unit_key_02': 'testtest'
} for _ in range(10)
]
sa_firehose._firehose_request_helper('invalid_stream', test_events)
missing_stream_message = 'Client Error ... An error occurred ' \
'(ResourceNotFoundException) when calling the PutRecordBatch ' \
'operation: Stream invalid_stream under account 123456789012 not found.'
assert_true(mock_logging.error.called_with(missing_stream_message))
def test_load_enabled_sources_invalid_log(self, mock_logging):
"""StreamAlertFirehose - Load Enabled Sources - Invalid Log"""
config = load_config('tests/unit/conf')
firehose_config = {'enabled_logs': ['log-that-doesnt-exist']}
sa_firehose = StreamAlertFirehose(region='us-east-1',
firehose_config=firehose_config,
log_sources=config['logs'])
assert_equal(len(sa_firehose._enabled_logs), 0)
assert_true(mock_logging.error.called)
def test_segment_records_by_size(self):
"""StreamAlertFirehose - Segment Large Records"""
sa_firehose = StreamAlertFirehose(region='us-east-1',
firehose_config={},
log_sources={})
record_batch = [
# unit_test_simple_log
{
'unit_key_01': 2,
'unit_key_02': 'testtest' * 10000
} for _ in range(100)
]
sized_batches = []
for sized_batch in sa_firehose._segment_records_by_size(record_batch):
sized_batches.append(sized_batch)
assert_true(len(str(sized_batches[0])) < 4000000)
assert_equal(len(sized_batches), 4)
assert_true(isinstance(sized_batches[3][0], dict))
def generate_firehose(config, main_dict, logging_bucket):
"""Generate the Firehose Terraform modules
Args:
config (CLIConfig): The loaded StreamAlert Config
main_dict (infinitedict): The Dict to marshal to a file
logging_bucket (str): The name of the global logging bucket
"""
if not config['global']['infrastructure'].get('firehose',
{}).get('enabled'):
return
sa_firehose = StreamAlertFirehose(
config['global']['account']['region'],
config['global']['infrastructure']['firehose'], config['logs'])
firehose_config = config['global']['infrastructure']['firehose']
firehose_s3_bucket_suffix = firehose_config.get('s3_bucket_suffix',
'streamalert.data')
firehose_s3_bucket_name = '{}.{}'.format(
config['global']['account']['prefix'], firehose_s3_bucket_suffix)
# Firehose Setup module
main_dict['module']['kinesis_firehose_setup'] = {
'source': 'modules/tf_stream_alert_kinesis_firehose_setup',
'account_id': config['global']['account']['aws_account_id'],
'prefix': config['global']['account']['prefix'],
'region': config['global']['account']['region'],
's3_logging_bucket': logging_bucket,
's3_bucket_name': firehose_s3_bucket_name,
'kms_key_id': '${aws_kms_key.server_side_encryption.key_id}'
}
# Add the Delivery Streams individually
for enabled_log in sa_firehose.enabled_logs:
main_dict['module']['kinesis_firehose_{}'.format(enabled_log)] = {
'source': 'modules/tf_stream_alert_kinesis_firehose_delivery_stream',
'buffer_size': config['global']['infrastructure']
['firehose'].get('buffer_size', 64),
'buffer_interval': config['global']['infrastructure']
['firehose'].get('buffer_interval', 300),\
'compression_format': config['global']['infrastructure']
['firehose'].get('compression_format', 'GZIP'),
'log_name': enabled_log,
'role_arn': '${module.kinesis_firehose_setup.firehose_role_arn}',
's3_bucket_name': firehose_s3_bucket_name,
'kms_key_arn': '${aws_kms_key.server_side_encryption.arn}'
}
def setup_mock_firehose_delivery_streams(config):
"""Mock Kinesis Firehose Streams for rule testing
Args:
config (CLIConfig): The StreamAlert config
"""
firehose_config = config['global']['infrastructure'].get('firehose')
if not firehose_config:
return
region = config['global']['account']['region']
sa_firehose = StreamAlertFirehose(region, firehose_config, config['logs'])
for log_type in sa_firehose.enabled_logs:
stream_name = 'streamalert_data_{}'.format(log_type)
prefix = '{}/'.format(log_type)
create_delivery_stream(region, stream_name, prefix)
def test_record_delivery_failed_put_count(self, mock_logging):
"""StreamAlertFirehose - Record Delivery - Failed Put Count"""
self.__sa_firehose = StreamAlertFirehose(region='us-east-1',
firehose_config={},
log_sources={})
# Add sample categorized payloads
for payload_type, logs in self._sample_categorized_payloads(
).iteritems():
self.__sa_firehose.categorized_payloads[payload_type].extend(logs)
# Setup mocked Delivery Streams
self._mock_delivery_streams([
'streamalert_data_test_log_type_json_nested',
'streamalert_data_unit_test_simple_log'
])
with patch.object(self.__sa_firehose._firehose_client,
'put_record_batch') as firehose_mock:
firehose_mock.side_effect = [{
'FailedPutCount':
3,
'RequestResponses': [{
"ErrorCode": "ServiceUnavailableException",
"ErrorMessage": "Slow down."
}, {
"ErrorCode": "ServiceUnavailableException",
"ErrorMessage": "Slow down."
}, {
"ErrorCode": "ServiceUnavailableException",
"ErrorMessage": "Slow down."
}]
}, {
'FailedPutCount':
3,
'RequestResponses': [{
"ErrorCode": "ServiceUnavailableException",
"ErrorMessage": "Slow down."
}, {
"ErrorCode": "ServiceUnavailableException",
"ErrorMessage": "Slow down."
}, {
"ErrorCode": "ServiceUnavailableException",
"ErrorMessage": "Slow down."
}]
}, {
'FailedPutCount':
0,
'RequestResponses': [{
"RecordId": "12345678910",
"ErrorCode": "None",
"ErrorMessage": "None"
}, {
"RecordId": "12345678910",
"ErrorCode": "None",
"ErrorMessage": "None"
}, {
"RecordId": "12345678910",
"ErrorCode": "None",
"ErrorMessage": "None"
}]
}]
self.__sa_firehose.send()
firehose_mock.assert_called()
assert_true(mock_logging.info.called)
def rebuild_partitions(table, bucket, config):
"""Rebuild an Athena table's partitions
Steps:
- Get the list of current partitions
- Destroy existing table
- Re-create tables
- Re-create partitions
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
Types of 'data' and 'alert' are accepted, but only 'data' is implemented
config (CLIConfig): Loaded StreamAlert CLI
"""
athena_client = get_athena_client(config)
sa_firehose = StreamAlertFirehose(
config['global']['account']['region'],
config['global']['infrastructure']['firehose'], config['logs'])
sanitized_table_name = sa_firehose.firehose_log_name(table)
# Get the current set of partitions
partitions = athena_client.get_table_partitions(sanitized_table_name)
if not partitions:
LOGGER_CLI.info('No partitions to rebuild for %s, nothing to do',
sanitized_table_name)
return
# Drop the table
LOGGER_CLI.info('Dropping table %s', sanitized_table_name)
success = athena_client.drop_table(sanitized_table_name)
if not success:
return
LOGGER_CLI.info('Creating table %s', sanitized_table_name)
# Re-create the table with previous partitions
create_table(table, bucket, config)
new_partitions_statement = helpers.add_partition_statement(
partitions, bucket, sanitized_table_name)
# Make sure our new alter table statement is within the query API limits
if len(new_partitions_statement) > MAX_QUERY_LENGTH:
LOGGER_CLI.error(
'Partition statement too large, writing to local file')
with open('partitions_{}.txt'.format(sanitized_table_name),
'w') as partition_file:
partition_file.write(new_partitions_statement)
return
LOGGER_CLI.info('Creating %d new partitions for %s', len(partitions),
sanitized_table_name)
success = athena_client.run_query(query=new_partitions_statement)
if not success:
LOGGER_CLI.error('Error re-creating new partitions for %s',
sanitized_table_name)
return
LOGGER_CLI.info('Successfully rebuilt partitions for %s',
sanitized_table_name)
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 CLI
schema_override (set): An optional set of key=value pairs to be used for
overriding the configured column_name=value_type.
"""
athena_client = StreamAlertAthenaClient(
config, results_key_prefix='stream_alert_cli')
sa_firehose = StreamAlertFirehose(
config['global']['account']['region'],
config['global']['infrastructure']['firehose'], config['logs'])
# Convert special characters in schema name to underscores
sanitized_table_name = sa_firehose.firehose_log_name(table)
# Check that the log type is enabled via Firehose
if sanitized_table_name != 'alerts' and sanitized_table_name not in sa_firehose.enabled_logs:
LOGGER_CLI.error(
'Table name %s missing from configuration or '
'is not enabled.', sanitized_table_name)
return
# Check if the table exists
if athena_client.check_table_exists(sanitized_table_name, True):
LOGGER_CLI.info('The \'%s\' table already exists.',
sanitized_table_name)
return
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 = handler_helpers.to_athena_schema(schema)
query = _construct_create_table_statement(schema=athena_schema,
table_name=table,
bucket=bucket)
else: # all other tables are log types
log_info = config['logs'][table.replace('_', ':', 1)]
schema = dict(log_info['schema'])
sanitized_schema = StreamAlertFirehose.sanitize_keys(schema)
athena_schema = handler_helpers.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 = StreamAlertFirehose.sanitize_keys(
envelope_keys)
# Note: this key is wrapped in backticks to be Hive compliant
athena_schema[
'`streamalert:envelope_keys`'] = handler_helpers.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('=')
if not all([column_name, column_type]):
LOGGER_CLI.error(
'Invalid schema override [%s], use column_name=type format',
override)
# 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_CLI.info('Applied schema override: %s:%s',
column_name, column_type)
else:
LOGGER_CLI.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)
create_table_success, _ = athena_client.run_athena_query(
query=query, database=athena_client.sa_database)
if not create_table_success:
LOGGER_CLI.error('The %s table could not be created',
sanitized_table_name)
return
# Update the CLI config
if (table != 'alerts' and bucket not in config['lambda']
['athena_partition_refresh_config']['buckets']):
config['lambda']['athena_partition_refresh_config']['buckets'][
bucket] = 'data'
config.write()
LOGGER_CLI.info('The %s table was successfully created!',
sanitized_table_name)
class TestStreamAlertFirehose(object):
"""Test class for StreamAlertFirehose"""
def __init__(self):
self.__sa_firehose = None
def teardown(self):
"""Setup before each method"""
self.__sa_firehose = None
@staticmethod
def _sample_categorized_payloads():
return {
'unit_test_simple_log': [{
'unit_key_01': 1,
'unit_key_02': 'test'
}, {
'unit_key_01': 2,
'unit_key_02': 'test'
}],
'test_log_type_json_nested': [{
'date': 'January 01, 3005',
'unixtime': '32661446400',
'host': 'my-host.name.website.com',
'data': {
'super': 'secret'
}
}]
}
@mock_kinesis
def _mock_delivery_streams(self, delivery_stream_names):
"""Mock Kinesis Delivery Streams for tests"""
for delivery_stream in delivery_stream_names:
self.__sa_firehose._firehose_client.create_delivery_stream(
DeliveryStreamName=delivery_stream,
S3DestinationConfiguration={
'RoleARN':
'arn:aws:iam::123456789012:role/firehose_delivery_role',
'BucketARN': 'arn:aws:s3:::kinesis-test',
'Prefix': '{}/'.format(delivery_stream),
'BufferingHints': {
'SizeInMBs': 123,
'IntervalInSeconds': 124
},
'CompressionFormat': 'Snappy',
})
@patch('stream_alert.rule_processor.firehose.LOGGER')
@mock_kinesis
def test_record_delivery_failed_put_count(self, mock_logging):
"""StreamAlertFirehose - Record Delivery - Failed Put Count"""
self.__sa_firehose = StreamAlertFirehose(region='us-east-1',
firehose_config={},
log_sources={})
# Add sample categorized payloads
for payload_type, logs in self._sample_categorized_payloads(
).iteritems():
self.__sa_firehose.categorized_payloads[payload_type].extend(logs)
# Setup mocked Delivery Streams
self._mock_delivery_streams([
'streamalert_data_test_log_type_json_nested',
'streamalert_data_unit_test_simple_log'
])
with patch.object(self.__sa_firehose._firehose_client,
'put_record_batch') as firehose_mock:
firehose_mock.side_effect = [{
'FailedPutCount':
3,
'RequestResponses': [{
"ErrorCode": "ServiceUnavailableException",
"ErrorMessage": "Slow down."
}, {
"ErrorCode": "ServiceUnavailableException",
"ErrorMessage": "Slow down."
}, {
"ErrorCode": "ServiceUnavailableException",
"ErrorMessage": "Slow down."
}]
}, {
'FailedPutCount':
3,
'RequestResponses': [{
"ErrorCode": "ServiceUnavailableException",
"ErrorMessage": "Slow down."
}, {
"ErrorCode": "ServiceUnavailableException",
"ErrorMessage": "Slow down."
}, {
"ErrorCode": "ServiceUnavailableException",
"ErrorMessage": "Slow down."
}]
}, {
'FailedPutCount':
0,
'RequestResponses': [{
"RecordId": "12345678910",
"ErrorCode": "None",
"ErrorMessage": "None"
}, {
"RecordId": "12345678910",
"ErrorCode": "None",
"ErrorMessage": "None"
}, {
"RecordId": "12345678910",
"ErrorCode": "None",
"ErrorMessage": "None"
}]
}]
self.__sa_firehose.send()
firehose_mock.assert_called()
assert_true(mock_logging.info.called)
@patch('stream_alert.rule_processor.firehose.LOGGER')
@mock_kinesis
def test_record_delivery(self, mock_logging):
"""StreamAlertFirehose - Record Delivery"""
self.__sa_firehose = StreamAlertFirehose(region='us-east-1',
firehose_config={},
log_sources={})
# Add sample categorized payloads
for payload_type, logs in self._sample_categorized_payloads(
).iteritems():
self.__sa_firehose.categorized_payloads[payload_type].extend(logs)
# Setup mocked Delivery Streams
self._mock_delivery_streams([
'streamalert_data_test_log_type_json_nested',
'streamalert_data_unit_test_simple_log'
])
# Send the records
with patch.object(self.__sa_firehose._firehose_client,
'put_record_batch') as firehose_mock:
firehose_mock.return_value = {'FailedPutCount': 0}
self.__sa_firehose.send()
firehose_mock.assert_called()
assert_true(mock_logging.info.called)
@patch('stream_alert.rule_processor.firehose.LOGGER')
@mock_kinesis
def test_record_delivery_failure(self, mock_logging):
"""StreamAlertFirehose - Record Delivery - Failed PutRecord"""
self.__sa_firehose = StreamAlertFirehose(region='us-east-1',
firehose_config={},
log_sources={})
# Add sample categorized payloads
for payload_type, logs in self._sample_categorized_payloads(
).iteritems():
self.__sa_firehose.categorized_payloads[payload_type].extend(logs)
# Setup mocked Delivery Streams
self._mock_delivery_streams([
'streamalert_data_test_log_type_json_nested',
'streamalert_data_unit_test_simple_log'
])
# Send the records
with patch.object(self.__sa_firehose._firehose_client,
'put_record_batch') as firehose_mock:
firehose_mock.return_value = {
'FailedPutCount':
3,
'RequestResponses': [
{
'RecordId': '12345',
'ErrorCode': '300',
'ErrorMessage': 'Bad message!!!'
},
]
}
self.__sa_firehose.send()
firehose_mock.assert_called()
assert_true(mock_logging.error.called)
@patch('stream_alert.rule_processor.firehose.LOGGER')
@mock_kinesis
def test_record_delivery_client_error(self, mock_logging):
"""StreamAlertFirehose - Record Delivery - Client Error"""
sa_firehose = StreamAlertFirehose(region='us-east-1',
firehose_config={},
log_sources={})
test_events = [
# unit_test_simple_log
{
'unit_key_01': 2,
'unit_key_02': 'testtest'
} for _ in range(10)
]
sa_firehose._firehose_request_helper('invalid_stream', test_events)
missing_stream_message = 'Client Error ... An error occurred ' \
'(ResourceNotFoundException) when calling the PutRecordBatch ' \
'operation: Stream invalid_stream under account 123456789012 not found.'
assert_true(mock_logging.error.called_with(missing_stream_message))
@mock_kinesis
def test_load_enabled_sources(self):
"""StreamAlertFirehose - Load Enabled Sources"""
config = load_config('tests/unit/conf')
firehose_config = {
'enabled_logs':
['json:regex_key_with_envelope', 'test_cloudtrail', 'cloudwatch']
} # expands to 2 logs
sa_firehose = StreamAlertFirehose(region='us-east-1',
firehose_config=firehose_config,
log_sources=config['logs'])
assert_equal(len(sa_firehose._enabled_logs), 4)
# Make sure the subtitution works properly
assert_true(all([':' not in log for log in sa_firehose.enabled_logs]))
assert_false(sa_firehose.enabled_log_source('test_inspec'))
@patch('stream_alert.rule_processor.firehose.LOGGER')
@mock_kinesis
def test_load_enabled_sources_invalid_log(self, mock_logging):
"""StreamAlertFirehose - Load Enabled Sources - Invalid Log"""
config = load_config('tests/unit/conf')
firehose_config = {'enabled_logs': ['log-that-doesnt-exist']}
sa_firehose = StreamAlertFirehose(region='us-east-1',
firehose_config=firehose_config,
log_sources=config['logs'])
assert_equal(len(sa_firehose._enabled_logs), 0)
assert_true(mock_logging.error.called)
def test_segment_records_by_size(self):
"""StreamAlertFirehose - Segment Large Records"""
sa_firehose = StreamAlertFirehose(region='us-east-1',
firehose_config={},
log_sources={})
record_batch = [
# unit_test_simple_log
{
'unit_key_01': 2,
'unit_key_02': 'testtest' * 10000
} for _ in range(100)
]
sized_batches = []
for sized_batch in sa_firehose._segment_records_by_size(record_batch):
sized_batches.append(sized_batch)
assert_true(len(str(sized_batches[0])) < 4000000)
assert_equal(len(sized_batches), 4)
assert_true(isinstance(sized_batches[3][0], dict))
def test_sanitize_keys(self):
"""StreamAlertFirehose - Sanitize Keys"""
# test_log_type_json_nested
test_event = {
'date': 'January 01, 3005',
'unixtime': '32661446400',
'host': 'my-host.name.website.com',
'data': {
'super-duper': 'secret',
'sanitize_me': 1,
'example-key': 1,
'moar**data': 2,
'even.more': 3
}
}
expected_sanitized_event = {
'date': 'January 01, 3005',
'unixtime': '32661446400',
'host': 'my-host.name.website.com',
'data': {
'super_duper': 'secret',
'sanitize_me': 1,
'example_key': 1,
'moar__data': 2,
'even_more': 3
}
}
sanitized_event = StreamAlertFirehose.sanitize_keys(test_event)
assert_equal(sanitized_event, expected_sanitized_event)
@patch('stream_alert.rule_processor.firehose.LOGGER')
def test_limit_record_size(self, mock_logging):
"""StreamAlertFirehose - Record Size Check"""
test_events = [
# unit_test_simple_log
{
'unit_key_01': 1,
'unit_key_02': 'test' * 250001 # is 4 bytes higher than max
},
{
'unit_key_01': 2,
'unit_key_02': 'test'
},
# test_log_type_json_nested
{
'date': 'January 01, 3005',
'unixtime': '32661446400',
'host': 'my-host.name.website.com',
'data': {
'super': 'secret'
}
}
]
StreamAlertFirehose._limit_record_size(test_events)
assert_true(len(test_events), 2)
assert_true(mock_logging.error.called)
def rebuild_partitions(table, bucket, config):
"""Rebuild an Athena table's partitions
Steps:
- Get the list of current partitions
- Destroy existing table
- Re-create tables
- Re-create partitions
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
Types of 'data' and 'alert' are accepted, but only 'data' is implemented
config (CLIConfig): Loaded StreamAlert CLI
"""
athena_client = StreamAlertAthenaClient(
config, results_key_prefix='stream_alert_cli')
sa_firehose = StreamAlertFirehose(
config['global']['account']['region'],
config['global']['infrastructure']['firehose'], config['logs'])
sanitized_table_name = sa_firehose.firehose_log_name(table)
# Get the current set of partitions
partition_success, partitions = athena_client.run_athena_query(
query='SHOW PARTITIONS {}'.format(sanitized_table_name),
database=athena_client.sa_database)
if not partition_success:
LOGGER_CLI.error('An error occurred when loading partitions for %s',
sanitized_table_name)
return
unique_partitions = athena_helpers.unique_values_from_query(partitions)
if not unique_partitions:
LOGGER_CLI.info('No partitions to rebuild for %s, nothing to do',
sanitized_table_name)
return
# Drop the table
LOGGER_CLI.info('Dropping table %s', sanitized_table_name)
drop_success, _ = athena_client.run_athena_query(
query='DROP TABLE {}'.format(sanitized_table_name),
database=athena_client.sa_database)
if not drop_success:
LOGGER_CLI.error('An error occurred when dropping the %s table',
sanitized_table_name)
return
LOGGER_CLI.info('Dropped table %s', sanitized_table_name)
LOGGER_CLI.info('Creating table %s', sanitized_table_name)
# Re-create the table with previous partitions
create_table(table, bucket, config)
new_partitions_statement = athena_helpers.partition_statement(
unique_partitions, bucket, sanitized_table_name)
# Make sure our new alter table statement is within the query API limits
if len(new_partitions_statement) > MAX_QUERY_LENGTH:
LOGGER_CLI.error(
'Partition statement too large, writing to local file')
with open('partitions_{}.txt'.format(sanitized_table_name),
'w') as partition_file:
partition_file.write(new_partitions_statement)
return
LOGGER_CLI.info('Creating %d new partitions for %s',
len(unique_partitions), sanitized_table_name)
new_part_success, _ = athena_client.run_athena_query(
query=new_partitions_statement, database=athena_client.sa_database)
if not new_part_success:
LOGGER_CLI.error('Error re-creating new partitions for %s',
sanitized_table_name)
return
LOGGER_CLI.info('Successfully rebuilt partitions for %s',
sanitized_table_name)
def create_table(athena_client, options, config):
"""Create a 'streamalert' Athena table
Args:
athena_client (boto3.client): Instantiated CLI AthenaClient
options (namedtuple): The parsed args passed from the CLI
config (CLIConfig): Loaded StreamAlert CLI
"""
sa_firehose = StreamAlertFirehose(
config['global']['account']['region'],
config['global']['infrastructure']['firehose'], config['logs'])
if not options.bucket:
LOGGER_CLI.error('Missing command line argument --bucket')
return
if not options.refresh_type:
LOGGER_CLI.error('Missing command line argument --refresh_type')
return
if options.type == 'data':
if not options.table_name:
LOGGER_CLI.error('Missing command line argument --table_name')
return
# Convert special characters in schema name to underscores
sanitized_table_name = sa_firehose.firehose_log_name(
options.table_name)
# Check that the log type is enabled via Firehose
if sanitized_table_name not in sa_firehose.enabled_logs:
LOGGER_CLI.error(
'Table name %s missing from configuration or '
'is not enabled.', sanitized_table_name)
return
# Check if the table exists
if athena_client.check_table_exists(sanitized_table_name):
LOGGER_CLI.info('The \'%s\' table already exists.',
sanitized_table_name)
return
log_info = config['logs'][options.table_name.replace('_', ':', 1)]
schema = dict(log_info['schema'])
sanitized_schema = StreamAlertFirehose.sanitize_keys(schema)
athena_schema = handler_helpers.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 = StreamAlertFirehose.sanitize_keys(
envelope_keys)
# Note: this key is wrapped in backticks to be Hive compliant
athena_schema[
'`streamalert:envelope_keys`'] = handler_helpers.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 options.schema_override:
for override in options.schema_override:
if '=' not in override:
LOGGER_CLI.error(
'Invalid schema override [%s], use column_name=type format',
override)
return
column_name, column_type = override.split('=')
if not all([column_name, column_type]):
LOGGER_CLI.error(
'Invalid schema override [%s], use column_name=type format',
override)
# 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_CLI.info('Applied schema override: %s:%s',
column_name, column_type)
else:
LOGGER_CLI.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=options.bucket)
elif options.type == 'alerts':
if athena_client.check_table_exists(options.type):
LOGGER_CLI.info('The \'alerts\' table already exists.')
return
query = ALERTS_TABLE_STATEMENT.format(bucket=options.bucket)
if query:
create_table_success, _ = athena_client.run_athena_query(
query=query, database='streamalert')
if create_table_success:
# Update the CLI config
config['lambda']['athena_partition_refresh_config'] \
['refresh_type'][options.refresh_type][options.bucket] = options.type
config.write()
table_name = options.type if options.type == 'alerts' else sanitized_table_name
LOGGER_CLI.info('The %s table was successfully created!',
table_name)
class StreamAlert(object):
"""Wrapper class for handling StreamAlert classification and processing"""
config = {}
def __init__(self, context):
"""Initializer
Args:
context (dict): An AWS context object which provides metadata on the currently
executing lambda function.
"""
# Load the config. Validation occurs during load, which will
# raise exceptions on any ConfigErrors
StreamAlert.config = StreamAlert.config or load_config()
# Load the environment from the context arn
self.env = load_env(context)
# Instantiate the send_alerts here to handle sending the triggered alerts to the
# alert processor
self.alert_forwarder = AlertForwarder()
# Instantiate a classifier that is used for this run
self.classifier = StreamClassifier(config=self.config)
self._failed_record_count = 0
self._processed_record_count = 0
self._processed_size = 0
self._alerts = []
rule_import_paths = [
item for location in {'rule_locations', 'matcher_locations'}
for item in self.config['global']['general'][location]
]
# Create an instance of the StreamRules class that gets cached in the
# StreamAlert class as an instance property
self._rules_engine = RulesEngine(self.config, *rule_import_paths)
# Firehose client attribute
self._firehose_client = None
def run(self, event):
"""StreamAlert Lambda function handler.
Loads the configuration for the StreamAlert function which contains
available data sources, log schemas, normalized types, and outputs.
Classifies logs sent into a parsed type.
Matches records against rules.
Args:
event (dict): An AWS event mapped to a specific source/entity
containing data read by Lambda.
Returns:
bool: True if all logs being parsed match a schema
"""
records = event.get('Records', [])
LOGGER.debug('Number of incoming records: %d', len(records))
if not records:
return False
firehose_config = self.config['global'].get('infrastructure',
{}).get('firehose', {})
if firehose_config.get('enabled'):
self._firehose_client = StreamAlertFirehose(
self.env['lambda_region'], firehose_config,
self.config['logs'])
payload_with_normalized_records = []
for raw_record in records:
# Get the service and entity from the payload. If the service/entity
# is not in our config, log and error and go onto the next record
service, entity = self.classifier.extract_service_and_entity(
raw_record)
if not service:
LOGGER.error(
'No valid service found in payload\'s raw record. Skipping '
'record: %s', raw_record)
continue
if not entity:
LOGGER.error(
'Unable to extract entity from payload\'s raw record for service %s. '
'Skipping record: %s', service, raw_record)
continue
# Cache the log sources for this service and entity on the classifier
if not self.classifier.load_sources(service, entity):
continue
# Create the StreamPayload to use for encapsulating parsed info
payload = load_stream_payload(service, entity, raw_record)
if not payload:
continue
payload_with_normalized_records.extend(
self._process_alerts(payload))
# Log normalized records metric
MetricLogger.log_metric(FUNCTION_NAME, MetricLogger.NORMALIZED_RECORDS,
len(payload_with_normalized_records))
# Apply Threat Intel to normalized records in the end of Rule Processor invocation
record_alerts = self._rules_engine.threat_intel_match(
payload_with_normalized_records)
self._alerts.extend(record_alerts)
if record_alerts:
self.alert_forwarder.send_alerts(record_alerts)
MetricLogger.log_metric(FUNCTION_NAME, MetricLogger.TOTAL_RECORDS,
self._processed_record_count)
MetricLogger.log_metric(FUNCTION_NAME,
MetricLogger.TOTAL_PROCESSED_SIZE,
self._processed_size)
LOGGER.debug('Invalid record count: %d', self._failed_record_count)
MetricLogger.log_metric(FUNCTION_NAME, MetricLogger.FAILED_PARSES,
self._failed_record_count)
LOGGER.debug('%s alerts triggered', len(self._alerts))
MetricLogger.log_metric(FUNCTION_NAME, MetricLogger.TRIGGERED_ALERTS,
len(self._alerts))
# Check if debugging logging is on before json dumping alerts since
# this can be time consuming if there are a lot of alerts
if self._alerts and LOGGER.isEnabledFor(LOG_LEVEL_DEBUG):
LOGGER.debug(
'Alerts:\n%s',
json.dumps([alert.output_dict() for alert in self._alerts],
indent=2,
sort_keys=True))
if self._firehose_client:
self._firehose_client.send()
# Only log rule info here if this is not running tests
# During testing, this gets logged at the end and printing here could be confusing
# since stress testing calls this method multiple times
if self.env['lambda_alias'] != 'development':
stats.print_rule_stats(True)
return self._failed_record_count == 0
@property
def alerts(self):
"""Returns list of Alert instances (useful for testing)."""
return self._alerts
def _process_alerts(self, payload):
"""Run the record through the rules, saving any alerts and forwarding them to Dynamo.
Args:
payload (StreamPayload): StreamAlert payload object being processed
"""
payload_with_normalized_records = []
for record in payload.pre_parse():
# Increment the processed size using the length of this record
self._processed_size += len(record.pre_parsed_record)
self.classifier.classify_record(record)
if not record.valid:
if self.env['lambda_alias'] != 'development':
LOGGER.error(
'Record does not match any defined schemas: %s\n%s',
record, record.pre_parsed_record)
self._failed_record_count += 1
continue
# Increment the total processed records to get an accurate assessment of throughput
self._processed_record_count += len(record.records)
LOGGER.debug(
'Classified and Parsed Payload: <Valid: %s, Log Source: %s, Entity: %s>',
record.valid, record.log_source, record.entity)
record_alerts, normalized_records = self._rules_engine.run(record)
payload_with_normalized_records.extend(normalized_records)
LOGGER.debug(
'Processed %d valid record(s) that resulted in %d alert(s).',
len(payload.records), len(record_alerts))
# Add all parsed records to the categorized payload dict only if Firehose is enabled
if self._firehose_client:
# Only send payloads with enabled log sources
if self._firehose_client.enabled_log_source(
payload.log_source):
self._firehose_client.categorized_payloads[
payload.log_source].extend(payload.records)
if not record_alerts:
continue
# Extend the list of alerts with any new ones so they can be returned
self._alerts.extend(record_alerts)
self.alert_forwarder.send_alerts(record_alerts)
return payload_with_normalized_records
def run(self, event):
"""StreamAlert Lambda function handler.
Loads the configuration for the StreamAlert function which contains
available data sources, log schemas, normalized types, and outputs.
Classifies logs sent into a parsed type.
Matches records against rules.
Args:
event (dict): An AWS event mapped to a specific source/entity
containing data read by Lambda.
Returns:
bool: True if all logs being parsed match a schema
"""
records = event.get('Records', [])
LOGGER.debug('Number of incoming records: %d', len(records))
if not records:
return False
firehose_config = self.config['global'].get('infrastructure',
{}).get('firehose', {})
if firehose_config.get('enabled'):
self._firehose_client = StreamAlertFirehose(
self.env['lambda_region'], firehose_config,
self.config['logs'])
payload_with_normalized_records = []
for raw_record in records:
# Get the service and entity from the payload. If the service/entity
# is not in our config, log and error and go onto the next record
service, entity = self.classifier.extract_service_and_entity(
raw_record)
if not service:
LOGGER.error(
'No valid service found in payload\'s raw record. Skipping '
'record: %s', raw_record)
continue
if not entity:
LOGGER.error(
'Unable to extract entity from payload\'s raw record for service %s. '
'Skipping record: %s', service, raw_record)
continue
# Cache the log sources for this service and entity on the classifier
if not self.classifier.load_sources(service, entity):
continue
# Create the StreamPayload to use for encapsulating parsed info
payload = load_stream_payload(service, entity, raw_record)
if not payload:
continue
payload_with_normalized_records.extend(
self._process_alerts(payload))
# Log normalized records metric
MetricLogger.log_metric(FUNCTION_NAME, MetricLogger.NORMALIZED_RECORDS,
len(payload_with_normalized_records))
# Apply Threat Intel to normalized records in the end of Rule Processor invocation
record_alerts = self._rules_engine.threat_intel_match(
payload_with_normalized_records)
self._alerts.extend(record_alerts)
if record_alerts:
self.alert_forwarder.send_alerts(record_alerts)
MetricLogger.log_metric(FUNCTION_NAME, MetricLogger.TOTAL_RECORDS,
self._processed_record_count)
MetricLogger.log_metric(FUNCTION_NAME,
MetricLogger.TOTAL_PROCESSED_SIZE,
self._processed_size)
LOGGER.debug('Invalid record count: %d', self._failed_record_count)
MetricLogger.log_metric(FUNCTION_NAME, MetricLogger.FAILED_PARSES,
self._failed_record_count)
LOGGER.debug('%s alerts triggered', len(self._alerts))
MetricLogger.log_metric(FUNCTION_NAME, MetricLogger.TRIGGERED_ALERTS,
len(self._alerts))
# Check if debugging logging is on before json dumping alerts since
# this can be time consuming if there are a lot of alerts
if self._alerts and LOGGER.isEnabledFor(LOG_LEVEL_DEBUG):
LOGGER.debug(
'Alerts:\n%s',
json.dumps([alert.output_dict() for alert in self._alerts],
indent=2,
sort_keys=True))
if self._firehose_client:
self._firehose_client.send()
# Only log rule info here if this is not running tests
# During testing, this gets logged at the end and printing here could be confusing
# since stress testing calls this method multiple times
if self.env['lambda_alias'] != 'development':
stats.print_rule_stats(True)
return self._failed_record_count == 0
def rebuild_partitions(athena_client, options, config):
"""Rebuild an Athena table's partitions
Steps:
- Get the list of current partitions
- Destroy existing table
- Re-create tables
- Re-create partitions
Args:
athena_client (boto3.client): Instantiated CLI AthenaClient
options (namedtuple): The parsed args passed from the CLI
config (CLIConfig): Loaded StreamAlert CLI
"""
if not options.table_name:
LOGGER_CLI.error('Missing command line argument --table_name')
return
if not options.bucket:
LOGGER_CLI.error('Missing command line argument --bucket')
return
sa_firehose = StreamAlertFirehose(
config['global']['account']['region'],
config['global']['infrastructure']['firehose'], config['logs'])
sanitized_table_name = sa_firehose.firehose_log_name(options.table_name)
if options.type == 'data':
# Get the current set of partitions
partition_success, partitions = athena_client.run_athena_query(
query='SHOW PARTITIONS {}'.format(sanitized_table_name),
database='streamalert')
if not partition_success:
LOGGER_CLI.error('An error occured when loading partitions for %s',
sanitized_table_name)
return
unique_partitions = athena_helpers.unique_values_from_query(partitions)
# Drop the table
LOGGER_CLI.info('Dropping table %s', sanitized_table_name)
drop_success, _ = athena_client.run_athena_query(
query='DROP TABLE {}'.format(sanitized_table_name),
database='streamalert')
if not drop_success:
LOGGER_CLI.error('An error occured when dropping the %s table',
sanitized_table_name)
return
LOGGER_CLI.info('Dropped table %s', sanitized_table_name)
new_partitions_statement = athena_helpers.partition_statement(
unique_partitions, options.bucket, sanitized_table_name)
# Make sure our new alter table statement is within the query API limits
if len(new_partitions_statement) > MAX_QUERY_LENGTH:
LOGGER_CLI.error(
'Partition statement too large, writing to local file')
with open('partitions_{}.txt'.format(sanitized_table_name),
'w') as partition_file:
partition_file.write(new_partitions_statement)
return
# Re-create the table with previous partitions
options.refresh_type = 'add_hive_partition'
create_table(athena_client, options, config)
LOGGER_CLI.info('Creating %d new partitions for %s',
len(unique_partitions), sanitized_table_name)
new_part_success, _ = athena_client.run_athena_query(
query=new_partitions_statement, database='streamalert')
if not new_part_success:
LOGGER_CLI.error('Error re-creating new partitions for %s',
sanitized_table_name)
return
LOGGER_CLI.info('Successfully rebuilt partitions for %s',
sanitized_table_name)
else:
LOGGER_CLI.info('Refreshing alerts tables unsupported')
def create_table(athena_client, options, config):
"""Create a 'streamalert' Athena table
Args:
athena_client (boto3.client): Instantiated CLI AthenaClient
options (namedtuple): The parsed args passed from the CLI
config (CLIConfig): Loaded StreamAlert CLI
"""
sa_firehose = StreamAlertFirehose(
config['global']['account']['region'],
config['global']['infrastructure']['firehose'], config['logs'])
if not options.bucket:
LOGGER_CLI.error('Missing command line argument --bucket')
return
if not options.refresh_type:
LOGGER_CLI.error('Missing command line argument --refresh_type')
return
if options.type == 'data':
if not options.table_name:
LOGGER_CLI.error('Missing command line argument --table_name')
return
sanitized_table_name = sa_firehose.firehose_log_name(
options.table_name)
if sanitized_table_name not in sa_firehose.enabled_logs:
LOGGER_CLI.error(
'Table name %s missing from configuration or '
'is not enabled.', sanitized_table_name)
return
if athena_client.check_table_exists(sanitized_table_name):
LOGGER_CLI.info('The \'%s\' table already exists.',
sanitized_table_name)
return
log_info = config['logs'][options.table_name.replace('_', ':', 1)]
schema = dict(log_info['schema'])
schema_statement = ''
sanitized_schema = StreamAlertFirehose.sanitize_keys(schema)
athena_schema = {}
_add_to_athena_schema(sanitized_schema, athena_schema)
# Support envelope keys
configuration_options = log_info.get('configuration')
if configuration_options:
envelope_keys = configuration_options.get('envelope_keys')
if envelope_keys:
sanitized_envelope_key_schema = StreamAlertFirehose.sanitize_keys(
envelope_keys)
# Note: this key is wrapped in backticks to be Hive compliant
_add_to_athena_schema(sanitized_envelope_key_schema,
athena_schema,
'`streamalert:envelope_keys`')
for key_name, key_type in athena_schema.iteritems():
# Account for nested structs
if isinstance(key_type, dict):
struct_schema = ''.join([
'{0}:{1},'.format(sub_key, sub_type)
for sub_key, sub_type in key_type.iteritems()
])
nested_schema_statement = '{0} struct<{1}>, '.format(
key_name,
# Use the minus index to remove the last comma
struct_schema[:-1])
schema_statement += nested_schema_statement
else:
schema_statement += '{0} {1},'.format(key_name, key_type)
query = (
'CREATE EXTERNAL TABLE {table_name} ({schema}) '
'PARTITIONED BY (dt string) '
'ROW FORMAT SERDE \'org.openx.data.jsonserde.JsonSerDe\' '
'WITH SERDEPROPERTIES ( \'ignore.malformed.json\' = \'true\') '
'LOCATION \'s3://{bucket}/{table_name}/\''.format(
table_name=sanitized_table_name,
# Use the minus index to remove the last comma
schema=schema_statement[:-1],
bucket=options.bucket))
elif options.type == 'alerts':
if athena_client.check_table_exists(options.type):
LOGGER_CLI.info('The \'alerts\' table already exists.')
return
query = ('CREATE EXTERNAL TABLE alerts ('
'log_source string,'
'log_type string,'
'outputs array<string>,'
'record string,'
'rule_description string,'
'rule_name string,'
'source_entity string,'
'source_service string)'
'PARTITIONED BY (dt string)'
'ROW FORMAT SERDE \'org.openx.data.jsonserde.JsonSerDe\''
'LOCATION \'s3://{bucket}/alerts/\''.format(
bucket=options.bucket))
if query:
create_table_success, _ = athena_client.run_athena_query(
query=query, database='streamalert')
if create_table_success:
# Update the CLI config
config['lambda']['athena_partition_refresh_config'] \
['refresh_type'][options.refresh_type][options.bucket] = options.type
config.write()
table_name = options.type if options.type == 'alerts' else sanitized_table_name
LOGGER_CLI.info('The %s table was successfully created!',
table_name)