def run_query_for_results(self, query):
"""Run Athena queries and get results back. This is a blocking operation
Args:
query (str): SQL query to execute
Returns:
dict: Response object with the result of the running query
"""
execution_id = self._execute_and_wait(query)
if not execution_id:
return
query_results = self._client.get_query_results(
QueryExecutionId=execution_id)
# The idea here is to leave the processing logic to the calling functions.
# No data being returned isn't always an indication that something is wrong.
# When handling the query result data, iterate over each element in the Row,
# and parse the Data key.
# Reference: https://bit.ly/2tWOQ2N
if not query_results['ResultSet']['Rows']:
LOGGER.debug('The query %s returned empty rows of data', query)
return query_results
def check_query_status(self, query_execution_id):
"""Check in on the running query, back off if the job is running or queued
Args:
query_execution_id (str): Athena query execution ID
Returns:
bool: True if the query state is SUCCEEDED, False otherwise
Reference https://bit.ly/2uuRtda.
"""
states_to_backoff = {'QUEUED', 'RUNNING'}
@backoff.on_predicate(backoff.fibo,
lambda resp: \
resp['QueryExecution']['Status']['State'] in states_to_backoff,
max_value=10,
jitter=backoff.full_jitter,
on_backoff=backoff_handler(),
on_success=success_handler(True))
def _check_status(query_execution_id):
return self._client.get_query_execution(
QueryExecutionId=query_execution_id)
execution_result = _check_status(query_execution_id)
state = execution_result['QueryExecution']['Status']['State']
if state != 'SUCCEEDED':
reason = execution_result['QueryExecution']['Status'][
'StateChangeReason']
LOGGER.error('Query %s %s with reason %s, exiting!',
query_execution_id, state, reason)
return False
return True
def process(cls, matcher_from_rule, record):
"""Process will ensure this record is valid for this matcher
Args:
matcher_from_rule (str): Name of matcher supplied in the rule decorator
record (dict): Record that the matcher should be applied against
Returns:
bool: True if the matcher applied to this record, False otherwise
"""
func = Matcher._matchers.get(matcher_from_rule)
if not func:
# TODO: previously, if a matcher used in a rule did not exist,
# we would log an error but not return False. Should an invalid
# matcher in a rule throw and error or be ignored? See here:
# https://github.com/airbnb/streamalert/blob/c33709129ee2bd9cd38f50f3e95fc7d01518e539/stream_alert/rule_processor/rules_engine.py#L162-L163
# TODO: is there any reason we shouldn't just import thee matcher into
# the rule file, instead of referring to it by the function's name?
LOGGER.error('The matcher [%s] does not exist!', matcher_from_rule)
return False
try:
return func(record)
except Exception: # pylint: disable=broad-except
LOGGER.exception('Encountered error with matcher: %s',
func.__name__)
return False
def print_rule_stats(reset=False):
"""Print some additional rule stats
Args:
reset (bool): Optional flag to reset the tracking statistics after printing
"""
if not RULE_STATS:
LOGGER.error('No rule statistics to print')
return
max_rule_name_len = max([len(rule) for rule in RULE_STATS.keys()])
stat_lines = []
for rule, stat in sorted(RULE_STATS.iteritems(), key=lambda (k, v):
(v, k)):
stat_lines.append('{rule: <{pad}}{stat}'.format(rule=rule,
pad=max_rule_name_len +
4,
stat=stat))
LOGGER.info('Rule statistics:\n%s', '\n'.join(stat_lines))
# Clear the dictionary that is storing statistics
# This allows for resetting when cumulative stats are not wanted
if reset:
RULE_STATS.clear()
def _wrapped(details):
message = '[Backoff]: Giving up calling \'{}\' after {:f} seconds and {:d} tries'.format(
details['target'].__name__, details['elapsed'], details['tries'])
if not debug_only:
LOGGER.info(message)
else:
LOGGER.debug(message)
def toggle_staged_state(self, rule_name, stage):
"""Mark the specified rule as staged=True or staged=False
Args:
rule_name (string): The name of the rule being staged
stage (bool): True if this rule should be staged and False if
this rule should be promoted out of staging.
"""
LOGGER.debug('Toggling staged state for rule \'%s\' to: %s', rule_name, stage)
update_expressions = ['set Staged = :staged']
expression_attributes = [':staged']
expression_values = [stage]
# If staging, add some additonal staging context to the expression
if stage:
update_expressions.extend(['StagedAt = :staged_at', 'StagedUntil = :staged_until'])
expression_attributes.extend([':staged_at', ':staged_until'])
expression_values.extend(self._staged_window())
args = {
'Key': {
'RuleName': rule_name
},
'UpdateExpression': ','.join(update_expressions),
'ExpressionAttributeValues': dict(zip(expression_attributes, expression_values)),
'ConditionExpression': 'attribute_exists(RuleName)'
}
try:
self._table.update_item(**args)
except ClientError as error:
if error.response['Error']['Code'] != 'ConditionalCheckFailedException':
raise
LOGGER.exception('Could not toggle rule staging status')
def _wrapped(details):
message = '[Backoff]: Calling \'{}\' again in {:f} seconds with {:d} tries so far'.format(
details['target'].__name__, details['wait'], details['tries'])
if not debug_only:
LOGGER.info(message)
else:
LOGGER.debug(message)
def giveup_handler(details):
"""Backoff logging handler for when backoff gives up.
Args:
details (dict): Backoff context containing the number of tries,
target function currently executing, kwargs, args, value,
and wait time.
"""
LOGGER.info('[Backoff]: Exiting after %d tries calling %s',
details['tries'], details['target'].__name__)
def success_handler(details):
"""Backoff logging handler for when backoff succeeds.
Args:
details (dict): Backoff context containing the number of tries,
target function currently executing, kwargs, args, value,
and wait time.
"""
LOGGER.debug('[Backoff]: Completed after %d tries calling %s',
details['tries'], details['target'].__name__)
def _add_new_rules(self, skip_staging=False):
"""Add any new local rules (renamed rules included) to the remote database"""
# If the table is empty, no rules have been added yet
# Add them all as unstaged to avoid demoting rules from production status
# Also, allow the user to bypass staging with the skip_staging flag
skip_staging = skip_staging or (len(self.remote_rule_names) == 0)
with self._table.batch_writer() as batch:
for rule_name in self.local_not_remote:
LOGGER.debug('Adding rule \'%s\' (skip staging=%s)', rule_name, skip_staging)
batch.put_item(self._dynamo_record(rule_name, skip_staging))
def _wrapped(details):
message = '[Backoff]: Successfully called \'{}\' after {:f} seconds and {:d} tries'.format(
details['target'].__name__,
details['elapsed'],
details['tries'],
)
# We will only want to log backoff on_success when tries more than 1.
if not debug_only and int(details['tries']) > 1:
LOGGER.info(message)
else:
LOGGER.debug(message)
def backoff_handler(details):
"""Backoff logging handler for when polling occurs.
Args:
details (dict): Backoff context containing the number of tries,
target function currently executing, kwargs, args, value,
and wait time.
"""
LOGGER.info(
'[Backoff]: Trying again in %f seconds after %d tries calling %s',
details['wait'], details['tries'], details['target'].__name__)
def timed(*args, **kw):
"""Wrapping function"""
time_start = time.time()
result = func(*args, **kw)
time_end = time.time()
message = '(module) {} (method) {} (time): {:>.4f}ms'.format(
func.__module__, func.__name__, (time_end - time_start) * 1000)
LOGGER.debug(message)
return result
def drop_all_tables(self):
"""Drop all table in the database
Returns:
bool: True if all tables were dropped successfully, False otherwise
"""
result = self.run_query_for_results('SHOW TABLES')
if not result:
LOGGER.error('There was an issue getting all tables')
return False
unique_tables = self._unique_values_from_query(result)
return all(self.drop_table(table) for table in unique_tables)
def drop_table(self, table_name):
"""Drop a specific table in the database
Args:
table_name (str): Table name which should be dropped from the database
Returns:
bool: True if the table was successfully dropped, False otherwise
"""
success = self.run_query('DROP TABLE {}'.format(table_name))
if not success:
LOGGER.error('Unable to drop table: %s', table_name)
return False
LOGGER.info('Successfully dropped table: %s', table_name)
return True
def get_table_partitions(self, table_name):
"""Get the list of partitions for a specific table in the database
Args:
table_name (str): Table name for which partitions should be returned
Returns:
set: Unique set of partitions for the given table
"""
partitions = self.run_query_for_results(
'SHOW PARTITIONS {}'.format(table_name))
if not partitions:
LOGGER.error('An error occurred when loading partitions for %s',
table_name)
return
return self._unique_values_from_query(partitions)
def download_s3_objects(self):
"""Download S3 files (json format) from S3 buckets into memory.
Returns:
dict: A dictionary contains information loaded from S3. The file name
will be the key, and value is file content in json format.
"""
_lookup_tables = {}
for bucket, files in self._buckets_info.iteritems():
for json_file in files:
try:
start_time = time.time()
s3_object = self._s3_client.Object(bucket, json_file).get()
size_kb = round(s3_object.get('ContentLength') / 1024.0, 2)
size_mb = round(size_kb / 1024.0, 2)
display_size = '{}MB'.format(size_mb) if size_mb else '{}KB'.format(size_kb)
LOGGER.info('Downloaded S3 file size %s and updated lookup table %s',
display_size, json_file)
data = s3_object.get('Body').read()
except ClientError as err:
LOGGER.error('Encounterred error while downloading %s from %s, %s',
json_file, bucket, err.response['Error']['Message'])
return _lookup_tables
except(Timeout, TimeoutError):
# Catching TimeoutError will catch both `ReadTimeoutError` and
# `ConnectionTimeoutError`.
LOGGER.error('Reading %s from S3 is timed out.', json_file)
return _lookup_tables
# The lookup data can optionally be compressed, so try to decompress
# This will fall back and use the original data if decompression fails
try:
data = zlib.decompress(data, 47)
except zlib.error:
LOGGER.debug('Data in \'%s\' is not compressed', json_file)
table_name = os.path.splitext(json_file)[0]
_lookup_tables[table_name] = json.loads(data)
total_time = time.time() - start_time
LOGGER.info('Downloaded S3 file %s seconds', round(total_time, 2))
return _lookup_tables
def _execute_query(self, query):
"""Execute an Athena query on the current database. This operation is non-blocking
Args:
query (str): SQL query to execute
Returns:
dict: Response object with the status of the running query
"""
LOGGER.debug('Executing query: %s', query)
try:
return self._client.start_query_execution(
QueryString=query,
QueryExecutionContext={'Database': self.database},
ResultConfiguration={'OutputLocation': self._s3_results_path})
except ClientError:
LOGGER.exception('Athena query failed')
def _execute_and_wait(self, query):
"""Execute an Athena query on the current database. This is a blocking operation
Args:
query (str): SQL query to execute
Returns:
str: Athena execution ID for the query that was executed
"""
response = self._execute_query(query)
if not response:
return
exeuction_id = response['QueryExecutionId']
success = self.check_query_status(exeuction_id)
if not success:
LOGGER.error('Athena query failed:\n%s', query)
return
return exeuction_id
def log_metric(cls, lambda_function, metric_name, value):
"""Log a metric using the logger the list of metrics to be sent to CloudWatch
Args:
metric_name (str): Name of metric to publish to. Choices are in `Metrics.Name` above
value (num): Numeric information to post to metric. AWS expects
this to be of type 'float' but will accept any numeric value that
is not super small (negative) or super large.
"""
# Do not log any metrics if they have been disabled by the user
if not ENABLE_METRICS:
return
if lambda_function not in cls._available_metrics:
LOGGER.error('Function \'%s\' not defined in available metrics. Options are: %s',
lambda_function,
', '.join('\'{}\''.format(key) for key in cls._available_metrics
if cls._available_metrics[key]))
return
if metric_name not in cls._available_metrics[lambda_function]:
LOGGER.error('Metric name (\'%s\') not defined for \'%s\' function. Options are: %s',
metric_name,
lambda_function,
', '.join('\'{}\''.format(value)
for value in cls._available_metrics[lambda_function]))
return
# Use a default format for logging this metric that will get picked up by the filters
LOGGER.info('{"metric_name": "%s", "metric_value": %s}', metric_name, value)
def check_query_status(self, execution_id):
"""Check in on the running query, back off if the job is running or queued
Args:
query_execution_id (str): Athena query execution ID
Returns:
bool: True if the query state is SUCCEEDED, False otherwise
Reference https://bit.ly/2uuRtda.
Raises:
AthenaQueryExecutionError: If any failure occurs while checking the status of the
query, this exception will be raised
"""
LOGGER.debug('Checking status of query with execution ID: %s',
execution_id)
states_to_backoff = {'QUEUED', 'RUNNING'}
@backoff.on_predicate(backoff.fibo,
lambda resp: \
resp['QueryExecution']['Status']['State'] in states_to_backoff,
max_value=10,
jitter=backoff.full_jitter,
on_backoff=backoff_handler(),
on_success=success_handler(True))
def _check_status(query_execution_id):
return self._client.get_query_execution(
QueryExecutionId=query_execution_id)
execution_result = _check_status(execution_id)
state = execution_result['QueryExecution']['Status']['State']
if state == 'SUCCEEDED':
return
# When the state is not SUCCEEDED, something bad must have occurred, so raise an exception
reason = execution_result['QueryExecution']['Status'][
'StateChangeReason']
raise AthenaQueryExecutionError(
'Query \'{}\' {} with reason \'{}\', exiting'.format(
execution_id, state, reason))
def process(self, record):
"""Process will call this rule's function on the passed record
Args:
record (dict): Record that this rule should be run against
Returns:
bool: True if this rule triggers for the passed record, False otherwise
"""
try:
# The initial_context object must be copied. This avoids
# bleed over from other runs of the rule using the same
# context object
if self.initial_context is not None:
self.context = deepcopy(self.initial_context)
return self.func(record, self.context)
return self.func(record)
except Exception: # pylint: disable=broad-except
LOGGER.exception('Encountered error with rule: %s', self.rule_name)
return False
def _execute_query(self, query):
"""Execute an Athena query on the current database. This operation is non-blocking
Args:
query (str): SQL query to execute
Returns:
dict: Response object with the status of the running query
Raises:
AthenaQueryExecutionError: If any failure occurs during the execution of the
query, this exception will be raised
"""
LOGGER.debug('Executing query: %s', query)
try:
return self._client.start_query_execution(
QueryString=query,
QueryExecutionContext={'Database': self.database},
ResultConfiguration={'OutputLocation': self._s3_results_path})
except ClientError as err:
raise AthenaQueryExecutionError(
'Athena query failed:\n{}'.format(err))
def checksum(self):
"""Produce an md5 for the contents of this rule.
This logic applies to expressions within the function only. It does not take
into account: the function name, docstring, comments, or decorator arguments
"""
if not self._checksum:
try:
code = inspect.getsource(self.func)
root = ast.parse(code)
md5 = hashlib.md5() # nosec
for expression in root.body[0].body:
# This check is necessary to ensure changes to the docstring
# are allowed without altering the checksum
if not isinstance(expression, ast.Expr):
md5.update(ast.dump(expression))
self._checksum = md5.hexdigest()
except (TypeError, IndentationError, IndexError):
LOGGER.exception('Could not checksum rule function')
self._checksum = self.CHECKSUM_UNKNOWN
return self._checksum
def toggle_staged_state(self, rule_name, stage):
"""Mark the specified rule as staged=True or staged=False
Args:
rule_name (string): The name of the rule being staged
stage (bool): True if this rule should be staged and False if
this rule should be promoted out of staging.
"""
if rule_name not in self.remote_rule_info:
LOGGER.error(
'Staging status for rule \'%s\' cannot be set to %s; rule does not exist',
rule_name, stage)
return
if self.remote_rule_info[rule_name]['Staged'] and stage:
LOGGER.info(
'Rule \'%s\' is already staged and will have its staging window updated',
rule_name)
LOGGER.debug('Toggling staged state for rule \'%s\' to: %s', rule_name,
stage)
update_expressions = ['set Staged = :staged']
expression_attributes = [':staged']
expression_values = [stage]
# If staging, add some additonal staging context to the expression
if stage:
update_expressions.extend(
['StagedAt = :staged_at', 'StagedUntil = :staged_until'])
expression_attributes.extend([':staged_at', ':staged_until'])
expression_values.extend(self._staged_window())
args = {
'UpdateExpression':
','.join(update_expressions),
'ExpressionAttributeValues':
dict(zip(expression_attributes, expression_values))
}
args.update(self._default_dynamo_kwargs(rule_name))
self._table.update_item(**args)
def load_lookup_tables(cls, config):
"""Load arbitrary json files to memory from S3 buckets when lookup table enabled
The lookup tables will also be refreshed based on "cache_refresh_minutes" setting
in the config.
Args:
config (dict): Loaded configuration from 'conf/' directory
Returns:
Return False if lookup table enabled or missing config. Otherwise, it
will return an instance of LookupTables class.
"""
lookup_tables = config['global']['infrastructure'].get('lookup_tables')
if not (lookup_tables and lookup_tables.get('enabled', False)):
return False
buckets_info = lookup_tables.get('buckets')
if not buckets_info:
LOGGER.error('Buckets not defined')
return False
lookup_refresh_interval = lookup_tables.get('cache_refresh_minutes',
10)
now = datetime.utcnow()
refresh_delta = timedelta(minutes=lookup_refresh_interval)
needs_refresh = cls._LOOKUP_TABLES_LAST_REFRESH + refresh_delta < now
if not needs_refresh:
LOGGER.debug(
'lookup tables do not need refresh (last refresh time: %s; '
'current time: %s)', cls._LOOKUP_TABLES_LAST_REFRESH, now)
return False
LOGGER.info(
'Refreshing lookup tables (last refresh time: %s; current time: %s)',
cls._LOOKUP_TABLES_LAST_REFRESH, now)
cls._LOOKUP_TABLES_LAST_REFRESH = now
return cls(buckets_info)
def download_s3_objects(self):
"""Download S3 files (json format) from S3 buckets into memory.
Returns:
dict: A dictionary contains information loaded from S3. The file name
will be the key, and value is file content in json format.
"""
_lookup_tables = {}
for bucket, files in self._buckets_info.iteritems():
for json_file in files:
table_name = os.path.splitext(json_file)[0]
try:
start_time = time.time()
s3_object = self._s3_client.Object(bucket, json_file).get()
size_kb = round(s3_object.get('ContentLength') / 1024.0, 2)
size_mb = round(size_kb / 1024.0, 2)
display_size = '{}MB'.format(
size_mb) if size_mb else '{}KB'.format(size_kb)
LOGGER.info(
'Downloaded S3 file size %s and updated lookup table %s',
display_size, table_name)
_lookup_tables[table_name] = json.loads(
s3_object.get('Body').read())
except ClientError as err:
LOGGER.error(
'Encounterred error while downloading %s from %s, %s',
json_file, bucket, err.response['Error']['Message'])
return _lookup_tables
total_time = time.time() - start_time
LOGGER.info('Downloaded S3 file %s seconds',
round(total_time, 2))
return _lookup_tables
def _del_old_rules(self):
"""Delete any rules that exist in the rule database but not locally"""
with self._table.batch_writer() as batch:
for rule_name in self.remote_not_local:
LOGGER.debug('Deleting rule \'%s\'', rule_name)
batch.delete_item(Key={'RuleName': rule_name})
import os
from stream_alert.shared import (
ALERT_PROCESSOR_NAME,
ATHENA_PARTITION_REFRESH_NAME,
LOGGER,
RULE_PROCESSOR_NAME
)
CLUSTER = os.environ.get('CLUSTER', 'unknown_cluster')
try:
ENABLE_METRICS = bool(int(os.environ.get('ENABLE_METRICS', 0)))
except ValueError as err:
ENABLE_METRICS = False
LOGGER.error('Invalid value for metric toggling, expected 0 or 1: %s',
err.message)
if not ENABLE_METRICS:
LOGGER.debug('Logging of metric data is currently disabled.')
class MetricLogger(object):
"""Class to hold metric logging to be picked up by log metric filters.
This basically acts as an enum, allowing for the use of dot notation for
accessing properties and avoids doing dict lookups a ton.
"""
# Constant metric names used for CloudWatch
FAILED_PARSES = 'FailedParses'
S3_DOWNLOAD_TIME = 'S3DownloadTime'
