def rule_staging_handler(options, config):
"""Handle operations related to the rule table (listing, updating, etc)
Args:
options (argparse.Namespace): Various options needed by subcommand
handlers
config (CLIConfig): Loaded configuration from 'conf/' directory
Returns:
bool: False if errors occurred, True otherwise
"""
if options.subcommand == 'enable':
config.toggle_rule_staging(options.enable)
table_name = '{}_streamalert_rules'.format(
config['global']['account']['prefix'])
if options.subcommand == 'status':
print RuleTable(table_name).__str__(options.verbose)
if options.subcommand in {'stage', 'unstage'}:
stage = (options.subcommand == 'stage')
table = RuleTable(table_name)
for rule in options.rules:
table.toggle_staged_state(rule, stage)
return True
def __init__(self):
config = load_config()
prefix = config['global']['account']['prefix']
# Create the rule table class for getting staging information
self._rule_table = RuleTable('{}_streamalert_rules'.format(prefix))
athena_config = config['lambda']['athena_partition_refresh_config']
# Get the name of the athena database to access
db_name = athena_config.get('database_name', self.STREAMALERT_DATABASE.format(prefix))
# Get the S3 bucket to store Athena query results
results_bucket = athena_config.get(
'results_bucket',
's3://{}.streamalert.athena-results'.format(prefix)
)
self._athena_client = AthenaClient(db_name, results_bucket, self.ATHENA_S3_PREFIX)
self._current_time = datetime.utcnow()
# Store the SNS topic arn to send alert stat information to
self._publisher = StatsPublisher(config, self._athena_client, self._current_time)
self._staging_stats = dict()
def _load_rule_table(cls, config):
"""Load and return a RuleTable class for communicating with the DynamoDB rule table
Args:
config (dict): Loaded configuration from 'conf/' directory
Returns:
rule_table.RuleTable: Loaded frontend for DynamoDB rules table
"""
# Ensure the rules table is enabled
rt_config = config['global']['infrastructure']['rules_table']
if not rt_config.get('enabled', False):
return
now = datetime.utcnow()
refresh_delta = timedelta(
minutes=rt_config.get('cache_refresh_minutes', 10))
# The rule table will need 'refreshed' if the refresh interval has been surpassed
needs_refresh = cls._RULE_TABLE_LAST_REFRESH + refresh_delta < now
if not needs_refresh:
LOGGER.debug(
'Rule table does not need refreshed (last refresh time: %s; '
'current time: %s)', cls._RULE_TABLE_LAST_REFRESH, now)
return
LOGGER.info(
'Refreshing rule table (last refresh time: %s; current time: %s)',
cls._RULE_TABLE_LAST_REFRESH, now)
table_name = '{}_streamalert_rules'.format(
config['global']['account']['prefix'])
cls._RULE_TABLE = RuleTable(table_name)
cls._RULE_TABLE_LAST_REFRESH = now
def rule_table_handler(options, config):
"""Handle operations related to the rule table (listing, updating, etc)
Args:
options (argparser.Namespace): Various options needed by subcommand
handlers
config (CLIConfig): Loaded configuration from 'conf/' directory
"""
table_name = '{}_streamalert_rules'.format(
config['global']['account']['prefix'])
if options.subcommand == 'status':
print RuleTable(table_name).__str__(options.verbose)
if options.subcommand in {'stage', 'unstage'}:
stage = (options.subcommand == 'stage')
table = RuleTable(table_name)
for rule in options.rules:
table.toggle_staged_state(rule, stage)
def test_rule_staged_only(self):
"""Rules Engine - Staged Rule"""
@rule(logs=['cloudwatch:test_match_types'], outputs=['foobar'])
def rule_staged_only(_): # pylint: disable=unused-variable
"""Modify context rule"""
return True
kinesis_data = json.dumps({
'account': 123456,
'region': '123456123456',
'source': '1.1.1.2',
'detail': {
'eventName': 'ConsoleLogin',
'sourceIPAddress': '1.1.1.2',
'recipientAccountId': '654321'
}
})
table = RuleTable('table')
table._remote_rule_info = {'rule_staged_only': {'Staged': True}}
self.config['global']['infrastructure']['rules_table'][
'enabled'] = True
with patch.object(RulesEngine, '_RULE_TABLE', table), \
patch.object(RulesEngine, '_RULE_TABLE_LAST_REFRESH', datetime.utcnow()):
self.rules_engine._load_rule_table(self.config)
# prepare the payloads
service, entity = 'kinesis', 'test_kinesis_stream'
raw_record = make_kinesis_raw_record(entity, kinesis_data)
payload = load_and_classify_payload(self.config, service, entity,
raw_record)
# process payloads
alerts, _ = self.rules_engine.run(payload)
# alert tests
assert_equal(list(alerts[0].outputs)[0], 'aws-firehose:alerts')
class RulePromoter(object):
"""Run queries to generate statistics on alerts."""
ATHENA_S3_PREFIX = 'rule_promoter'
STREAMALERT_DATABASE = '{}_streamalert'
def __init__(self):
self._config = load_config()
prefix = self._config['global']['account']['prefix']
# Create the rule table class for getting staging information
self._rule_table = RuleTable('{}_streamalert_rules'.format(prefix))
athena_config = self._config['lambda'][
'athena_partition_refresh_config']
# Get the name of the athena database to access
db_name = athena_config.get('database_name',
self.STREAMALERT_DATABASE.format(prefix))
# Get the S3 bucket to store Athena query results
results_bucket = athena_config.get(
'results_bucket',
's3://{}.streamalert.athena-results'.format(prefix))
self._athena_client = AthenaClient(db_name, results_bucket,
self.ATHENA_S3_PREFIX)
self._current_time = datetime.utcnow()
self._staging_stats = dict()
def _get_staging_info(self):
"""Query the Rule table for rule staging info needed to count each rule's alerts
Example of rule metadata returned by RuleTable.remote_rule_info():
{
'example_rule_name':
{
'Staged': True
'StagedAt': datetime.datetime object,
'StagedUntil': '2018-04-21T02:23:13.332223Z'
}
}
"""
for rule in sorted(self._rule_table.remote_rule_info):
info = self._rule_table.remote_rule_info[rule]
# If the rule is not staged, do not get stats on it
if not info['Staged']:
continue
self._staging_stats[rule] = StagingStatistic(
info['StagedAt'], info['StagedUntil'], self._current_time,
rule)
return len(self._staging_stats) != 0
def _update_alert_count(self):
"""Transform Athena query results into alert counts for rules_engine
Args:
query (str): Athena query to run and wait for results
Returns:
dict: Representation of alert counts, where key is the rule name
and value is the alert count (int) since this rule was staged
"""
query = StagingStatistic.construct_compound_count_query(
self._staging_stats.values())
LOGGER.debug('Running compound query for alert count: \'%s\'', query)
for page, results in enumerate(
self._athena_client.query_result_paginator(query)):
for i, row in enumerate(results['ResultSet']['Rows']):
if page == 0 and i == 0: # skip header row included in first page only
continue
row_values = [data.values()[0] for data in row['Data']]
rule_name, alert_count = row_values[0], int(row_values[1])
LOGGER.debug('Found %d alerts for rule \'%s\'', alert_count,
rule_name)
self._staging_stats[rule_name].alert_count = alert_count
def run(self, send_digest):
"""Perform statistic analysis of currently staged rules
Args:
send_digest (bool): True if the staging statistics digest should be
published, False otherwise
"""
if not self._get_staging_info():
LOGGER.debug('No staged rules to promote')
return
self._update_alert_count()
self._promote_rules()
if send_digest:
publisher = StatsPublisher(self._config, self._athena_client,
self._current_time)
publisher.publish(self._staging_stats.values())
else:
LOGGER.debug('Staging statistics digest will not be sent')
def _promote_rules(self):
"""Promote any rule that has not resulted in any alerts since being staged"""
for rule in self._rules_to_be_promoted:
LOGGER.info('Promoting rule \'%s\' at %s', rule,
self._current_time)
self._rule_table.toggle_staged_state(rule, False)
@property
def _rules_to_be_promoted(self):
"""Returns a list of rules that are eligible for promotion"""
return [
rule for rule, stat in self._staging_stats.iteritems()
if self._current_time > stat.staged_until and stat.alert_count == 0
]
@property
def _rules_failing_promotion(self):
"""Returns a list of rules that are ineligible for promotion"""
return [
rule for rule, stat in self._staging_stats.iteritems()
if stat.alert_count != 0
]
