fc9c5f0072e19d9f56304ef95a47250bfae7b142-PR-1230

Build # fc9c5f0072e19d9f56304ef95a47250bfae7b142-PR-1230

Build Type

github-actions

Committed by GitHub

Commit Message

<a href="https://github.com/airbnb/streamalert/commit/<a class=hub.com/airbnb/streamalert/commit/<a class="double-link" href="https://git"><a class=hub.com/airbnb/streamalert/commit/<a class="double-link" href="https://git"><a class=hub.com/airbnb/streamalert/commit/<a class="double-link" href="https://git"><a class=hub.com/airbnb/streamalert/commit/fc9c5f0072e19d9f56304ef95a47250bfae7b142">fc9c5f007">&lt;a href=&quot;https://github.com/airbnb/streamalert/commit/</a><a class="double-link" href="https://github.com/airbnb/streamalert/commit/&lt;a class=&quot;double-link&quot; href=&quot;https://git">&lt;a class=</a>hub.com/airbnb/streamalert/commit/&lt;a class=&quot;double-link&quot; href=&quot;https://git">&lt;a class=</a>hub.com/airbnb/streamalert/commit/&lt;a class=&quot;double-link&quot; href=&quot;https://git">&lt;a class=</a>hub.com/airbnb/streamalert/commit/&lt;a class=&quot;double-link&quot; href=&quot;https://git">&lt;a class=</a>hub.com/airbnb/streamalert/commit/fc9c5f0072e19d9f56304ef95a47250bfae7b142">fc9c5f007</a><a href="https://github.com/airbnb/streamalert/commit/fc9c5f0072e19d9f56304ef95a47250bfae7b142">&quot;&gt;&amp;lt;a href=&amp;quot;https://github.com/airbnb/streamalert/commit/&lt;/a&gt;&lt;a class=&quot;double-link&quot; href=&quot;https://github.com/airbnb/streamalert/commit/&amp;lt;a class=&amp;quot;double-link&amp;quot; href=&amp;quot;https://git&quot;&gt;&amp;lt;a class=&lt;/a&gt;hub.com/airbnb/streamalert/commit/&amp;lt;a class=&amp;quot;double-link&amp;quot; href=&amp;quot;https://git&quot;&gt;&amp;lt;a class=&lt;/a&gt;hub.com/airbnb/streamalert/commit/&amp;lt;a class=&amp;quot;double-link&amp;quot; href=&amp;quot;https://git&quot;&gt;&amp;lt;a class=&lt;/a&gt;hub.com/airbnb/streamalert/commit/&amp;lt;a class=&amp;quot;double-link&amp;quot; href=&amp;quot;https://git&quot;&gt;&amp;lt;a class=&lt;/a&gt;hub.com/airbnb/streamalert/commit/fc9c5f0072e19d9f56304ef95a47250bfae7b142&quot;&gt;fc9c5f007&lt;/a&gt;&lt;a href=&quot;https://github.com/airbnb/streamalert/commit/fc9c5f0072e19d9f56304ef95a47250bfae7b142&quot;&gt;&amp;quot;&amp;gt;&amp;amp;lt;a href=&amp;amp;quot;https://github.com/airbnb/streamalert/commit/&amp;lt;/a&amp;gt;&amp;lt;a class=&amp;quot;double-link&amp;quot; href=&amp;quot;https://github.com/airbnb/streamalert/comm... (continued)

Pull Request Pull Request #1230: Normalization reboot - Artifacts inventory

Run Details

91 of 91 new or added lines in 5 files covered. (100.0%)

6407 of 6718 relevant lines covered (95.37%)

0.95 hits per line

Source File
Press 'n' to go to next uncovered line, 'b' for previous

97.42

/streamalert/shared/firehose.py

"""
Copyright 2017-present Airbnb, Inc.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

   http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
"""
from collections import defaultdict
import json
import hashlib
import re

import backoff
import boto3
from botocore.exceptions import ClientError, HTTPClientError
from botocore.exceptions import ConnectionError as BotocoreConnectionError

from streamalert.shared import CLASSIFIER_FUNCTION_NAME as FUNCTION_NAME
import streamalert.shared.helpers.boto as boto_helpers
from streamalert.shared.logger import get_logger
from streamalert.shared.metrics import MetricLogger
from streamalert.shared.backoff_handlers import (
    backoff_handler,
    giveup_handler,
    success_handler
)

LOGGER = get_logger(__name__)


class FirehoseClient:
    """Handles preparing and sending data from the classifier function to Kinesis Firehose"""
    # Used to detect special characters in payload keys.
    # This is necessary for sanitization of data prior to searching in Athena.
    SPECIAL_CHAR_REGEX = re.compile(r'\W')
    SPECIAL_CHAR_SUB = '_'
    # For PutRecordBatch backoff
    MAX_BACKOFF_ATTEMPTS = 10
    # Adds a max of 20 seconds more to the Lambda function
    MAX_BACKOFF_FIBO_VALUE = 8
    # Set Firehose Limits: http://bit.ly/2fw5UY2
    MAX_BATCH_COUNT = 500
    MAX_BATCH_SIZE = 4000 * 1000
    # The subtraction of 2 accounts for the newline at the end
    MAX_RECORD_SIZE = 1000 * 1000 - 2

    # Default firehose name format, should be formatted with deployment prefix
    DEFAULT_FIREHOSE_FMT = '{}streamalert_{}'

    # Exception for which backoff operations should be performed
    EXCEPTIONS_TO_BACKOFF = (ClientError, BotocoreConnectionError, HTTPClientError)

    # Set of enabled log types for firehose, loaded from configs
    _ENABLED_LOGS = dict()

    # The max length of the firehose stream name is 64. For streamalert data firehose,
    # we reserve 12 chars to have `streamalert_` as part of prefix. Please refer to
    # terraform/modules/tf_kinesis_firehose_delivery_stream/main.tf
    AWS_FIREHOSE_NAME_MAX_LEN = 64

    FIREHOSE_NAME_MIN_HASH_LEN = 8

    def __init__(self, prefix, firehose_config=None, log_sources=None):
        self._prefix = prefix if firehose_config and firehose_config.get('use_prefix', True) else ''
        self._client = boto3.client('firehose', config=boto_helpers.default_config())
        self.load_enabled_log_sources(firehose_config, log_sources, force_load=True)

    @classmethod
    def _records_to_json_list(cls, records):
        """Write the dictionary records to json lines and return the list of lines

        Args:
            records (list<dict>): Records to be written to line-delimited json

        Returns:
            list: JSON serialized records
        """
        # Write the json lines to the object in minimal form
        return [
            json.dumps(record, separators=(',', ':')) + '\n' for record in records
        ]

    @classmethod
    def _record_batches(cls, records):
        """Segment the records into batches that conform to Firehose restrictions

        This will log any single record that is too large to send, and skip it.

        Args:
            records (list): The original records list to be segmented

        Yields:
            list: Batches of JSON serialized records that conform to Firehose restrictions
        """
        records_json = cls._records_to_json_list(records)

        current_batch_size = 0
        current_batch = []
        for record in records_json:
            line_len = len(record)
            # Check if the max size of the batch has been reached or if the current
            # record will exceed the max batch size and start a new batch
            if ((len(current_batch) == cls.MAX_BATCH_COUNT) or
                    (current_batch_size + line_len > cls.MAX_BATCH_SIZE)) and current_batch:
                yield current_batch[:]
                current_batch_size = 0
                del current_batch[:]

            if line_len > cls.MAX_RECORD_SIZE:
                LOGGER.error('Record too large (%d) to send to Firehose:\n%s', line_len, record)
                cls._log_failed(1)
                continue

            # Add the record to the batch
            current_batch_size += line_len
            current_batch.append(record)

        # yield the result of the last batch (no need to copy via slicing)
        if current_batch:
            yield current_batch

    @classmethod
    def sanitized_value(cls, key):
        """Sanitize a key by replacing non-word characters with '_'

        Args:
            key (str): a string needs to be sanitized

        Returns:
            str: sanitized string
        """
        return re.sub(cls.SPECIAL_CHAR_REGEX, cls.SPECIAL_CHAR_SUB, key)

    @classmethod
    def sanitize_keys(cls, record):
        """Remove special characters from parsed record keys

        This is required when searching in Athena.  Keys can only have
        a period or underscore

        Args:
            record (dict): Original parsed record

        Returns:
            dict: A sanitized record
        """
        new_record = {}
        for key, value in record.items():
            sanitized_key = cls.sanitized_value(key)

            # Handle nested objects
            if isinstance(value, dict):
                new_record[sanitized_key] = cls.sanitize_keys(record[key])
            else:
                new_record[sanitized_key] = record[key]

        return new_record

    @staticmethod
    def _strip_successful_records(batch, response):
        """Inspect the response and remove any records records that have successfully to sent

        For each record, the index of the response element is the same as the index
        used in the request array.

        Args:
            batch (list): List of dicts with JSON dumped records that are being
                sent to Firehose. Format is:
                [{'Data': <json-dumped-rec>}, {'Data': <json-dumped-rec>}]
            response (dict): Response object from the boto3.client.put_record_batch call
                that contains metadata on the success status of the call
        """
        success_indices = [idx for idx, rec in enumerate(response['RequestResponses'])
                           if rec.get('RecordId')]

        for idx in sorted(success_indices, reverse=True):
            del batch[idx]

    def _categorize_records(self, payloads):
        """Add the records to the proper list of cached records, based on log type

        Args:
            payloads (list): List of PayloadRecord items that include parsed records
        """
        # Create a dictionary to hold parsed payloads by log type.
        # Firehose needs this information to send to its corresponding
        # delivery stream.
        categorized_records = defaultdict(list)

        for payload in payloads:
            # Only send payloads with enabled log sources
            if not self.enabled_log_source(payload.log_schema_type):
                continue

            # Add the records to the dictionary of categorized records
            categorized_records[payload.log_schema_type].extend(payload.parsed_records)

        return categorized_records

    @classmethod
    def _finalize(cls, response, stream_name, size):
        """Perform any final operations for this response, such as metric logging, etc

        Args:
            response (dict): boto3 client response object
            stream_name (str): The name of the Delivery Stream to send to
            size (int): The original size of the batch being sent
        """
        if not response:
            return  # Could happen in the case of backoff failing entirely

        # Check for failures that occurred in PutRecordBatch after several backoff attempts
        if response.get('FailedPutCount'):
            failed_records = [
                failed for failed in response['RequestResponses']
                if failed.get('ErrorCode')
            ]
            cls._log_failed(response['FailedPutCount'])

            # Only print the first 100 failed records to Cloudwatch logs
            LOGGER.error(
                'Failed to put the following records to firehose %s: %s',
                stream_name,
                json.dumps(failed_records[:1024], indent=2)
            )
            return

        MetricLogger.log_metric(FUNCTION_NAME, MetricLogger.FIREHOSE_RECORDS_SENT, size)
        LOGGER.info(
            'Successfully sent %d message(s) to firehose %s with RequestId \'%s\'',
            size,
            stream_name,
            response.get('ResponseMetadata', {}).get('RequestId', '')
        )

    @classmethod
    def _log_failed(cls, count):
        """Helper to log the failed Firehose records metric

        Args:
            count (int): Number of failed records
        """
        MetricLogger.log_metric(FUNCTION_NAME, MetricLogger.FIREHOSE_FAILED_RECORDS, count)

    def _send_batch(self, stream_name, record_batch):
        """Send record batches to Firehose

        Args:
            stream_name (str): The name of the Delivery Stream to send to
            record_batch (list): The records to send
        """
        @backoff.on_predicate(backoff.fibo,
                              lambda resp: resp['FailedPutCount'] > 0,
                              max_tries=self.MAX_BACKOFF_ATTEMPTS,
                              max_value=self.MAX_BACKOFF_FIBO_VALUE,
                              on_backoff=backoff_handler(debug_only=False),
                              on_success=success_handler(),
                              on_giveup=giveup_handler())
        @backoff.on_exception(backoff.fibo,
                              self.EXCEPTIONS_TO_BACKOFF,
                              max_tries=self.MAX_BACKOFF_ATTEMPTS,
                              on_backoff=backoff_handler(debug_only=False),
                              on_success=success_handler(),
                              on_giveup=giveup_handler())
        def _firehose_request_helper(data):
            """Firehose request wrapper to use with backoff"""
            # Use the current length of data here so we can track failed records that are retried
            LOGGER.debug('Sending %d records to firehose %s', len(data), stream_name)

            response = self._client.put_record_batch(DeliveryStreamName=stream_name, Records=data)

            # Log this as an error for now so it can be picked up in logs
            if response['FailedPutCount'] > 0:
                LOGGER.warning('Received non-zero FailedPutCount: %d', response['FailedPutCount'])
                # Strip out the successful records so only the failed ones are retried. This happens
                # to the list of dictionary objects, so the called function sees the updated list
                self._strip_successful_records(data, response)

            return response

        # The record here already contains a newline, so do not append one
        records_data = [
            {'Data': record}
            for record in record_batch
        ]

        # The try/except here is to catch the raised error at the end of the backoff
        try:
            return _firehose_request_helper(records_data)
        except self.EXCEPTIONS_TO_BACKOFF:
            LOGGER.exception('Firehose request failed')
            # Use the current length of the records_data in case some records were
            # successful but others were not
            self._log_failed(len(records_data))

    @classmethod
    def generate_firehose_name(cls, prefix, log_stream_name):
        """Generate suffix of stream name complaint to firehose naming restriction, no
        longer than 64 characters

        Args:
            prefix (str): The prefix defined in conf/global.json to firehose stream name
            log_stream_name (str): The name of the log from conf/logs.json or conf/schemas/*.json

        Returns:
            str: suffix of stream name
        """
        if prefix:
            prefix += '_'

        # This same substitution method is used when naming the Delivery Streams
        stream_name = cls.sanitized_value(cls.DEFAULT_FIREHOSE_FMT.format(prefix, log_stream_name))
        if len(stream_name) <= cls.AWS_FIREHOSE_NAME_MAX_LEN:
            return stream_name

        base_name = stream_name[:cls.AWS_FIREHOSE_NAME_MAX_LEN - cls.FIREHOSE_NAME_MIN_HASH_LEN]
        if not base_name.endswith('_'):
            # make sure this ends in an underscore, but not 2
            base_name = '{}_'.format(
                base_name[:-1]
            ) if base_name[-2] != '_' else '{}_'.format(base_name[:-2])

        # combine the base_name and first 8 chars of hash result together as new
        # stream name.
        return '{}{}'.format(
            base_name, hashlib.md5(stream_name.encode()).hexdigest()  # nosec
        )[:cls.AWS_FIREHOSE_NAME_MAX_LEN]

    @classmethod
    def enabled_log_source(cls, log_source_name):
        """Check that the incoming record is an enabled log source for Firehose

        Args:
            log_source_name (str): The log source of the record

        Returns:
            bool: Whether or not the log source is enabled to send to Firehose
        """
        if not cls._ENABLED_LOGS:
            LOGGER.error('Enabled logs not loaded')
            return False

        return cls.sanitized_value(log_source_name) in cls._ENABLED_LOGS

    @classmethod
    def load_enabled_log_sources(cls, firehose_config, log_sources, force_load=False):
        """Load and expand all declared and enabled Firehose log sources

        Args:
            firehose_config (dict): Loaded Firehose config from global.json
            log_sources (dict): Loaded logs.json file
            force_load (bool=False): Set to True if the log sources should be reloaded
                even if there is cached values

        Returns:
            dict: Enabled logs, key: sanitized table name, value: log type value
        """
        # Do not reload the logs if they are already cached
        if cls._ENABLED_LOGS and not force_load:
            return cls._ENABLED_LOGS

        # Nothing to load if no configs passed
        if not (firehose_config and log_sources):
            return cls._ENABLED_LOGS

        # Expand enabled logs into specific subtypes
        for enabled_log in firehose_config.get('enabled_logs', {}):
            enabled_log_parts = enabled_log.split(':')

            # Expand to all subtypes
            if len(enabled_log_parts) == 1:
                expanded_logs = {
                    cls.sanitized_value(log_name): log_name
                    for log_name in log_sources
                    if log_name.split(':')[0] == enabled_log_parts[0]
                }

                if not expanded_logs:
                    LOGGER.error('Enabled Firehose log %s not declared in logs.json', enabled_log)
                    continue

                cls._ENABLED_LOGS.update(expanded_logs)

            elif len(enabled_log_parts) == 2:
                if enabled_log not in log_sources:
                    LOGGER.error('Enabled Firehose log %s not declared in logs.json', enabled_log)
                    continue

                cls._ENABLED_LOGS[cls.sanitized_value(enabled_log)] = enabled_log

        return cls._ENABLED_LOGS

    @classmethod
    def load_from_config(cls, prefix, firehose_config, log_sources):
        """Get a Firehose client for sending logs

        Args:
            firehose_config (dict): Loaded Firehose config from global.json
            log_sources (dict): Loaded logs.json file

        Returns:
            FirehoseClient or None: If disabled, this returns None, otherwise it returns an
                instanec of FirehoseClient
        """
        if not firehose_config.get('enabled'):
            return
        return cls(prefix=prefix, firehose_config=firehose_config, log_sources=log_sources)

    @classmethod
    def get_client(cls, prefix, artifact_extractor_config):
        """Get a Firehose client for sending artifacts

        Args:
            prefix (str): Account prefix from global.json
            artifact_extractor_config (dict): Loaded Artifact Extractor config from lambda.json

        Returns:
            FirehoseClient or None: If disabled, this returns None, otherwise it returns an
                instanec of FirehoseClient
        """
        if not artifact_extractor_config.get('enabled'):
            return
        return cls(prefix=prefix)

    def send(self, payloads):
        """Send all classified records to a respective Firehose Delivery Stream

        Args:
            payloads (list): List of PayloadRecord items that include parsed records
        """
        records = self._categorize_records(payloads)

        # Iterate through each set of categorized payloads.
        # Each batch will be processed to their specific Firehose, which lands the data
        # in a specific prefix in S3.
        for log_type, records in records.items():
            # firehose stream name has the length limit, no longer than 64 characters
            formatted_stream_name = self.generate_firehose_name(self._prefix, log_type)

            # Process each record batch in the categorized payload set
            for record_batch in self._record_batches(records):
                batch_size = len(record_batch)
                response = self._send_batch(formatted_stream_name, record_batch)
                self._finalize(response, formatted_stream_name, batch_size)

    def send_artifacts(self, artifacts, stream_name):
        """Send artifacts to artifacts Firehose delievery stream
        Args:
            artifacts (list(dict)): A list of artifacts extracted from normalized records.
            stream_name (str): Stream name of destination Firehose.
        """
        for artifact_batch in self._record_batches(artifacts):
            batch_size = len(artifact_batch)
            response = self._send_batch(stream_name, artifact_batch)
            self._finalize(response, stream_name, batch_size)

1	"""
2	Copyright 2017-present Airbnb, Inc.
3
4	Licensed under the Apache License, Version 2.0 (the "License");
5	you may not use this file except in compliance with the License.
6	You may obtain a copy of the License at
7
8	http://www.apache.org/licenses/LICENSE-2.0
9
10	Unless required by applicable law or agreed to in writing, software
11	distributed under the License is distributed on an "AS IS" BASIS,
12	WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13	See the License for the specific language governing permissions and
14	limitations under the License.
15	"""
16	from collections import defaultdict	1✔
17	import json	1✔
18	import hashlib	1✔
19	import re	1✔
20
21	import backoff	1✔
22	import boto3	1✔
23	from botocore.exceptions import ClientError, HTTPClientError	1✔
24	from botocore.exceptions import ConnectionError as BotocoreConnectionError	1✔
25
26	from streamalert.shared import CLASSIFIER_FUNCTION_NAME as FUNCTION_NAME	1✔
27	import streamalert.shared.helpers.boto as boto_helpers	1✔
28	from streamalert.shared.logger import get_logger	1✔
29	from streamalert.shared.metrics import MetricLogger	1✔
30	from streamalert.shared.backoff_handlers import (	1✔
31	backoff_handler,
32	giveup_handler,
33	success_handler
34	)
35
36	LOGGER = get_logger(__name__)	1✔
37
38
39	class FirehoseClient:	1✔
40	"""Handles preparing and sending data from the classifier function to Kinesis Firehose"""
41	# Used to detect special characters in payload keys.
42	# This is necessary for sanitization of data prior to searching in Athena.
43	SPECIAL_CHAR_REGEX = re.compile(r'\W')	1✔
44	SPECIAL_CHAR_SUB = '_'	1✔
45	# For PutRecordBatch backoff
46	MAX_BACKOFF_ATTEMPTS = 10	1✔
47	# Adds a max of 20 seconds more to the Lambda function
48	MAX_BACKOFF_FIBO_VALUE = 8	1✔
49	# Set Firehose Limits: http://bit.ly/2fw5UY2
50	MAX_BATCH_COUNT = 500	1✔
51	MAX_BATCH_SIZE = 4000 * 1000	1✔
52	# The subtraction of 2 accounts for the newline at the end
53	MAX_RECORD_SIZE = 1000 * 1000 - 2	1✔
54
55	# Default firehose name format, should be formatted with deployment prefix
56	DEFAULT_FIREHOSE_FMT = '{}streamalert_{}'	1✔
57
58	# Exception for which backoff operations should be performed
59	EXCEPTIONS_TO_BACKOFF = (ClientError, BotocoreConnectionError, HTTPClientError)	1✔
60
61	# Set of enabled log types for firehose, loaded from configs
62	_ENABLED_LOGS = dict()	1✔
63
64	# The max length of the firehose stream name is 64. For streamalert data firehose,
65	# we reserve 12 chars to have `streamalert_` as part of prefix. Please refer to
66	# terraform/modules/tf_kinesis_firehose_delivery_stream/main.tf
67	AWS_FIREHOSE_NAME_MAX_LEN = 64	1✔
68
69	FIREHOSE_NAME_MIN_HASH_LEN = 8	1✔
70
71	def __init__(self, prefix, firehose_config=None, log_sources=None):	1✔
72	self._prefix = prefix if firehose_config and firehose_config.get('use_prefix', True) else ''	1✔
73	self._client = boto3.client('firehose', config=boto_helpers.default_config())	1✔
74	self.load_enabled_log_sources(firehose_config, log_sources, force_load=True)	1✔
75
76	@classmethod	1✔
77	def _records_to_json_list(cls, records):
78	"""Write the dictionary records to json lines and return the list of lines
79
80	Args:
81	records (list<dict>): Records to be written to line-delimited json
82
83	Returns:
84	list: JSON serialized records
85	"""
86	# Write the json lines to the object in minimal form
87	return [	1✔
88	json.dumps(record, separators=(',', ':')) + '\n' for record in records
89	]
90
91	@classmethod	1✔
92	def _record_batches(cls, records):
93	"""Segment the records into batches that conform to Firehose restrictions
94
95	This will log any single record that is too large to send, and skip it.
96
97	Args:
98	records (list): The original records list to be segmented
99
100	Yields:
101	list: Batches of JSON serialized records that conform to Firehose restrictions
102	"""
103	records_json = cls._records_to_json_list(records)	1✔
104
105	current_batch_size = 0	1✔
106	current_batch = []	1✔
107	for record in records_json:	1✔
108	line_len = len(record)	1✔
109	# Check if the max size of the batch has been reached or if the current
110	# record will exceed the max batch size and start a new batch
111	if ((len(current_batch) == cls.MAX_BATCH_COUNT) or	1✔
112	(current_batch_size + line_len > cls.MAX_BATCH_SIZE)) and current_batch:
113	yield current_batch[:]	1✔
114	current_batch_size = 0	1✔
115	del current_batch[:]	1✔
116
117	if line_len > cls.MAX_RECORD_SIZE:	1✔
118	LOGGER.error('Record too large (%d) to send to Firehose:\n%s', line_len, record)	1✔
119	cls._log_failed(1)	1✔
120	continue	1✔
121
122	# Add the record to the batch
123	current_batch_size += line_len	1✔
124	current_batch.append(record)	1✔
125
126	# yield the result of the last batch (no need to copy via slicing)
127	if current_batch:	1✔
128	yield current_batch	1✔
129
130	@classmethod	1✔
131	def sanitized_value(cls, key):
132	"""Sanitize a key by replacing non-word characters with '_'
133
134	Args:
135	key (str): a string needs to be sanitized
136
137	Returns:
138	str: sanitized string
139	"""
140	return re.sub(cls.SPECIAL_CHAR_REGEX, cls.SPECIAL_CHAR_SUB, key)	1✔
141
142	@classmethod	1✔
143	def sanitize_keys(cls, record):
144	"""Remove special characters from parsed record keys
145
146	This is required when searching in Athena. Keys can only have
147	a period or underscore
148
149	Args:
150	record (dict): Original parsed record
151
152	Returns:
153	dict: A sanitized record
154	"""
155	new_record = {}	1✔
156	for key, value in record.items():	1✔
157	sanitized_key = cls.sanitized_value(key)	1✔
158
159	# Handle nested objects
160	if isinstance(value, dict):	1✔
161	new_record[sanitized_key] = cls.sanitize_keys(record[key])	1✔
162	else:
163	new_record[sanitized_key] = record[key]	1✔
164
165	return new_record	1✔
166
167	@staticmethod	1✔
168	def _strip_successful_records(batch, response):
169	"""Inspect the response and remove any records records that have successfully to sent
170
171	For each record, the index of the response element is the same as the index
172	used in the request array.
173
174	Args:
175	batch (list): List of dicts with JSON dumped records that are being
176	sent to Firehose. Format is:
177	[{'Data': <json-dumped-rec>}, {'Data': <json-dumped-rec>}]
178	response (dict): Response object from the boto3.client.put_record_batch call
179	that contains metadata on the success status of the call
180	"""
181	success_indices = [idx for idx, rec in enumerate(response['RequestResponses'])	1✔
182	if rec.get('RecordId')]
183
184	for idx in sorted(success_indices, reverse=True):	1✔
185	del batch[idx]	1✔
186
187	def _categorize_records(self, payloads):	1✔
188	"""Add the records to the proper list of cached records, based on log type
189
190	Args:
191	payloads (list): List of PayloadRecord items that include parsed records
192	"""
193	# Create a dictionary to hold parsed payloads by log type.
194	# Firehose needs this information to send to its corresponding
195	# delivery stream.
196	categorized_records = defaultdict(list)	1✔
197
198	for payload in payloads:	1✔
199	# Only send payloads with enabled log sources
200	if not self.enabled_log_source(payload.log_schema_type):	1✔
201	continue	1✔
202
203	# Add the records to the dictionary of categorized records
204	categorized_records[payload.log_schema_type].extend(payload.parsed_records)	1✔
205
206	return categorized_records	1✔
207
208	@classmethod	1✔
209	def _finalize(cls, response, stream_name, size):
210	"""Perform any final operations for this response, such as metric logging, etc
211
212	Args:
213	response (dict): boto3 client response object
214	stream_name (str): The name of the Delivery Stream to send to
215	size (int): The original size of the batch being sent
216	"""
217	if not response:	1✔
218	return # Could happen in the case of backoff failing entirely	×
219
220	# Check for failures that occurred in PutRecordBatch after several backoff attempts
221	if response.get('FailedPutCount'):	1✔
222	failed_records = [	1✔
223	failed for failed in response['RequestResponses']
224	if failed.get('ErrorCode')
225	]
226	cls._log_failed(response['FailedPutCount'])	1✔
227
228	# Only print the first 100 failed records to Cloudwatch logs
229	LOGGER.error(	1✔
230	'Failed to put the following records to firehose %s: %s',
231	stream_name,
232	json.dumps(failed_records[:1024], indent=2)
233	)
234	return	1✔
235
236	MetricLogger.log_metric(FUNCTION_NAME, MetricLogger.FIREHOSE_RECORDS_SENT, size)	1✔
237	LOGGER.info(	1✔
238	'Successfully sent %d message(s) to firehose %s with RequestId \'%s\'',
239	size,
240	stream_name,
241	response.get('ResponseMetadata', {}).get('RequestId', '')
242	)
243
244	@classmethod	1✔
245	def _log_failed(cls, count):
246	"""Helper to log the failed Firehose records metric
247
248	Args:
249	count (int): Number of failed records
250	"""
251	MetricLogger.log_metric(FUNCTION_NAME, MetricLogger.FIREHOSE_FAILED_RECORDS, count)	1✔
252
253	def _send_batch(self, stream_name, record_batch):	1✔
254	"""Send record batches to Firehose
255
256	Args:
257	stream_name (str): The name of the Delivery Stream to send to
258	record_batch (list): The records to send
259	"""
260	@backoff.on_predicate(backoff.fibo,	1✔
261	lambda resp: resp['FailedPutCount'] > 0,
262	max_tries=self.MAX_BACKOFF_ATTEMPTS,
263	max_value=self.MAX_BACKOFF_FIBO_VALUE,
264	on_backoff=backoff_handler(debug_only=False),
265	on_success=success_handler(),
266	on_giveup=giveup_handler())
267	@backoff.on_exception(backoff.fibo,	1✔
268	self.EXCEPTIONS_TO_BACKOFF,
269	max_tries=self.MAX_BACKOFF_ATTEMPTS,
270	on_backoff=backoff_handler(debug_only=False),
271	on_success=success_handler(),
272	on_giveup=giveup_handler())
273	def _firehose_request_helper(data):
274	"""Firehose request wrapper to use with backoff"""
275	# Use the current length of data here so we can track failed records that are retried
276	LOGGER.debug('Sending %d records to firehose %s', len(data), stream_name)	1✔
277
278	response = self._client.put_record_batch(DeliveryStreamName=stream_name, Records=data)	1✔
279
280	# Log this as an error for now so it can be picked up in logs
281	if response['FailedPutCount'] > 0:	1✔
282	LOGGER.warning('Received non-zero FailedPutCount: %d', response['FailedPutCount'])	1✔
283	# Strip out the successful records so only the failed ones are retried. This happens
284	# to the list of dictionary objects, so the called function sees the updated list
285	self._strip_successful_records(data, response)	1✔
286
287	return response	1✔
288
289	# The record here already contains a newline, so do not append one
290	records_data = [	1✔
291	{'Data': record}
292	for record in record_batch
293	]
294
295	# The try/except here is to catch the raised error at the end of the backoff
296	try:	1✔
297	return _firehose_request_helper(records_data)	1✔
298	except self.EXCEPTIONS_TO_BACKOFF:	1✔
299	LOGGER.exception('Firehose request failed')	1✔
300	# Use the current length of the records_data in case some records were
301	# successful but others were not
302	self._log_failed(len(records_data))	1✔
303
304	@classmethod	1✔
305	def generate_firehose_name(cls, prefix, log_stream_name):
306	"""Generate suffix of stream name complaint to firehose naming restriction, no
307	longer than 64 characters
308
309	Args:
310	prefix (str): The prefix defined in conf/global.json to firehose stream name
311	log_stream_name (str): The name of the log from conf/logs.json or conf/schemas/*.json
312
313	Returns:
314	str: suffix of stream name
315	"""
316	if prefix:	1✔
317	prefix += '_'	1✔
318
319	# This same substitution method is used when naming the Delivery Streams
320	stream_name = cls.sanitized_value(cls.DEFAULT_FIREHOSE_FMT.format(prefix, log_stream_name))	1✔
321	if len(stream_name) <= cls.AWS_FIREHOSE_NAME_MAX_LEN:	1✔
322	return stream_name	1✔
323
324	base_name = stream_name[:cls.AWS_FIREHOSE_NAME_MAX_LEN - cls.FIREHOSE_NAME_MIN_HASH_LEN]	1✔
325	if not base_name.endswith('_'):	1✔
326	# make sure this ends in an underscore, but not 2
327	base_name = '{}_'.format(	1✔
328	base_name[:-1]
329	) if base_name[-2] != '_' else '{}_'.format(base_name[:-2])
330
331	# combine the base_name and first 8 chars of hash result together as new
332	# stream name.
333	return '{}{}'.format(	1✔
334	base_name, hashlib.md5(stream_name.encode()).hexdigest() # nosec
335	)[:cls.AWS_FIREHOSE_NAME_MAX_LEN]
336
337	@classmethod	1✔
338	def enabled_log_source(cls, log_source_name):
339	"""Check that the incoming record is an enabled log source for Firehose
340
341	Args:
342	log_source_name (str): The log source of the record
343
344	Returns:
345	bool: Whether or not the log source is enabled to send to Firehose
346	"""
347	if not cls._ENABLED_LOGS:	1✔
348	LOGGER.error('Enabled logs not loaded')	1✔
349	return False	1✔
350
351	return cls.sanitized_value(log_source_name) in cls._ENABLED_LOGS	1✔
352
353	@classmethod	1✔
354	def load_enabled_log_sources(cls, firehose_config, log_sources, force_load=False):	1✔
355	"""Load and expand all declared and enabled Firehose log sources
356
357	Args:
358	firehose_config (dict): Loaded Firehose config from global.json
359	log_sources (dict): Loaded logs.json file
360	force_load (bool=False): Set to True if the log sources should be reloaded
361	even if there is cached values
362
363	Returns:
364	dict: Enabled logs, key: sanitized table name, value: log type value
365	"""
366	# Do not reload the logs if they are already cached
367	if cls._ENABLED_LOGS and not force_load:	1✔
368	return cls._ENABLED_LOGS	1✔
369
370	# Nothing to load if no configs passed
371	if not (firehose_config and log_sources):	1✔
372	return cls._ENABLED_LOGS	1✔
373
374	# Expand enabled logs into specific subtypes
375	for enabled_log in firehose_config.get('enabled_logs', {}):	1✔
376	enabled_log_parts = enabled_log.split(':')	1✔
377
378	# Expand to all subtypes
379	if len(enabled_log_parts) == 1:	1✔
380	expanded_logs = {	1✔
381	cls.sanitized_value(log_name): log_name
382	for log_name in log_sources
383	if log_name.split(':')[0] == enabled_log_parts[0]
384	}
385
386	if not expanded_logs:	1✔
387	LOGGER.error('Enabled Firehose log %s not declared in logs.json', enabled_log)	1✔
388	continue	1✔
389
390	cls._ENABLED_LOGS.update(expanded_logs)	1✔
391
392	elif len(enabled_log_parts) == 2:	1✔
393	if enabled_log not in log_sources:	1✔
394	LOGGER.error('Enabled Firehose log %s not declared in logs.json', enabled_log)	1✔
395	continue	1✔
396
397	cls._ENABLED_LOGS[cls.sanitized_value(enabled_log)] = enabled_log	1✔
398
399	return cls._ENABLED_LOGS	1✔
400
401	@classmethod	1✔
402	def load_from_config(cls, prefix, firehose_config, log_sources):
403	"""Get a Firehose client for sending logs
404
405	Args:
406	firehose_config (dict): Loaded Firehose config from global.json
407	log_sources (dict): Loaded logs.json file
408
409	Returns:
410	FirehoseClient or None: If disabled, this returns None, otherwise it returns an
411	instanec of FirehoseClient
412	"""
413	if not firehose_config.get('enabled'):	1✔
414	return	1✔
415	return cls(prefix=prefix, firehose_config=firehose_config, log_sources=log_sources)	1✔
416
417	@classmethod	1✔
418	def get_client(cls, prefix, artifact_extractor_config):
419	"""Get a Firehose client for sending artifacts
420
421	Args:
422	prefix (str): Account prefix from global.json
423	artifact_extractor_config (dict): Loaded Artifact Extractor config from lambda.json
424
425	Returns:
426	FirehoseClient or None: If disabled, this returns None, otherwise it returns an
427	instanec of FirehoseClient
428	"""
429	if not artifact_extractor_config.get('enabled'):	×
430	return	×
431	return cls(prefix=prefix)	×
432
433	def send(self, payloads):	1✔
434	"""Send all classified records to a respective Firehose Delivery Stream
435
436	Args:
437	payloads (list): List of PayloadRecord items that include parsed records
438	"""
439	records = self._categorize_records(payloads)	1✔
440
441	# Iterate through each set of categorized payloads.
442	# Each batch will be processed to their specific Firehose, which lands the data
443	# in a specific prefix in S3.
444	for log_type, records in records.items():	1✔
445	# firehose stream name has the length limit, no longer than 64 characters
446	formatted_stream_name = self.generate_firehose_name(self._prefix, log_type)	1✔
447
448	# Process each record batch in the categorized payload set
449	for record_batch in self._record_batches(records):	1✔
450	batch_size = len(record_batch)	1✔
451	response = self._send_batch(formatted_stream_name, record_batch)	1✔
452	self._finalize(response, formatted_stream_name, batch_size)	1✔
453
454	def send_artifacts(self, artifacts, stream_name):	1✔
455	"""Send artifacts to artifacts Firehose delievery stream
456	Args:
457	artifacts (list(dict)): A list of artifacts extracted from normalized records.
458	stream_name (str): Stream name of destination Firehose.
459	"""
460	for artifact_batch in self._record_batches(artifacts):	1✔
461	batch_size = len(artifact_batch)	1✔
462	response = self._send_batch(stream_name, artifact_batch)	1✔
463	self._finalize(response, stream_name, batch_size)	1✔

airbnb / streamalert / fc9c5f0072e19d9f56304ef95a47250bfae7b142-PR-1230

Source File Press 'n' to go to next uncovered line, 'b' for previous

Source File
Press 'n' to go to next uncovered line, 'b' for previous