195255325

Build # 195255325

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/<a class="double-link" href="https://git"><a class=hub.com/airbnb/streamalert/commit/bd01ef83e9e2dbb65a5000cdfa5da88de2a53e6f">bd01ef83e">&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/bd01ef83e9e2dbb65a5000cdfa5da88de2a53e6f">bd01ef83e</a><a href="https://github.com/airbnb/streamalert/commit/bd01ef83e9e2dbb65a5000cdfa5da88de2a53e6f">&lt;a href=&quot;https://github.com/airbnb/streamalert/commit/bd01ef83e9e2dbb65a5000cdfa5da88de2a53e6f&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/commit/&amp;amp;lt;a class=&amp;amp;quot;double-link&amp;amp;quot; href=&amp;amp;quot;https://git&amp;quot;&amp;gt;&amp;amp;lt;a class=&amp;lt;/a&amp;gt;hub.com/airbnb/streamalert/commit/bd01ef83e9e2dbb65a5000cdfa5da88de2a53e6f&amp;quot;&amp;gt;bd01ef83e&amp;lt;/a&amp;gt;&amp;lt;a href=&amp;quot;https://github.com/airbnb/streamalert/commit/bd01ef83e9e2dbb65a5000cdfa5da88de2a53e6f&amp;quot;&amp;gt;&amp;amp;lt;a href=&amp;amp;quot;https://github.com/airbnb/streamalert/commit/bd01ef83e9e2dbb65a5000cdfa5da88d&amp;lt;/a&amp;gt;e2a53e</a>6f&quot;&gt;&amp;quot;&amp;gt;Merge &amp;lt;/a&amp;gt;&amp;lt;a class=&amp;quot;double-link&amp;quot; href=&amp;quot;https://github.com/airbnb/streamalert/commit/&amp;lt;a class=&amp;quot;double-link&amp;quot; href=&amp;quot;https://github.com/airbnb/streamalert/commit/&lt;a class=&quot;double-link&quot; href=&quot;https://github.com/airbnb/streamalert/commit/a37bf13d6&... (continued)

Pull Request Pull Request #1284: addressing issue with 0 byte files in s3

Run Details

9 of 9 new or added lines in 1 file covered. (100.0%)

6519 of 6835 relevant lines covered (95.38%)

0.95 hits per line

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

99.1

/streamalert/classifier/payload/s3.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.
"""
import json
import logging
import urllib.request
import urllib.parse
import urllib.error
import gzip
import os
import tempfile
import subprocess
import time

import boto3
from botocore.exceptions import ClientError
import jsonlines

from streamalert.classifier.payload.payload_base import (
    PayloadRecord,
    RegisterInput,
    StreamPayload
)
from streamalert.shared import CLASSIFIER_FUNCTION_NAME as FUNCTION_NAME
from streamalert.shared.logger import get_logger
from streamalert.shared.metrics import MetricLogger


LOGGER = get_logger(__name__)
LOGGER_DEBUG_ENABLED = LOGGER.isEnabledFor(logging.DEBUG)


class S3PayloadError(Exception):
    """Exception for S3Payload errors"""


@RegisterInput
class S3Payload(StreamPayload):
    """S3Payload class"""

    MAX_S3_SIZE = 128 * 1024 * 1024

    @property
    def bucket(self):
        return self.raw_record['s3']['bucket']['name']

    @property
    def key(self):
        return self.raw_record['s3']['object']['key']

    @property
    def size(self):
        return int(self.raw_record['s3']['object']['size'])

    @property
    def region(self):
        return self.raw_record['awsRegion']

    @property
    def display_size(self):
        """Calculate and format a size for printing"""
        size_kb = round(self.size / 1024.0, 2)
        size_mb = round(size_kb / 1024.0, 2)
        return '{}MB'.format(size_mb) if size_mb else '{}KB'.format(size_kb)

    @classmethod
    def service(cls):
        return 's3'

    @classmethod
    def _unquote(cls, data):
        # Use the urllib unquote method to decode any url encoded characters
        # (ie - %26 --> &) from the bucket and key names
        return urllib.parse.unquote(data)

    def _check_size(self):
        """Ensure the S3 file's size is not too large to download into the Lambda environment

        Returns:
            bool: True if the file is smaller than 128 MB, False otherwise
        """
        # Ignore 0 size files
        if self.size == 0:
            LOGGER.warning('S3 file size is 0 bytes, skipping: %s/%s', self.bucket, self.key)
            return False

        # size greater than 128MB
        if self.size > self.MAX_S3_SIZE:
            raise S3PayloadError('S3 object {}/{} is too large and cannot be downloaded '
                                 'from S3: {}'.format(self.bucket, self.key, self.display_size))

        return True

    @staticmethod
    def _cleanup():
        """Cleanup method to remove all objects in the Lambda container's temp directory"""
        # Do nothing if this is not running in AWS Lambda
        if 'LAMBDA_RUNTIME_DIR' not in os.environ:
            return

        LOGGER.debug('Shredding temp directory')

        for root, dirs, files in os.walk(tempfile.gettempdir(), topdown=False):
            for name in files:
                subprocess.check_call([  # nosec
                    'shred', '--force', '--iterations=1',
                    '--remove', os.path.join(root, name)])
            for name in dirs:
                os.rmdir(os.path.join(root, name))  # nosec

    @staticmethod
    def _gz_reader(open_file):
        open_file.seek(0)
        reader = gzip.GzipFile(fileobj=open_file, mode='r')
        try:
            # Test to ensure this is gzip data, then rewind
            reader.read(1)
            reader.rewind()
        except IOError:
            # Fall back on the default reader
            reader = open_file
            reader.seek(0)

        # Return either the gzip reader or the original reader
        return reader

    @staticmethod
    def _jsonlines_reader(open_file):
        open_file.seek(0)
        reader = None
        try:
            json_lines = jsonlines.Reader(open_file)
            json_lines.read()
            reader = json_lines
        except ValueError:
            reader = open_file

        open_file.seek(0)
        return reader

    @classmethod
    def _read_downloaded_object(cls, open_file):
        """Read the contents of the downloaded S3 file using the open file handle

        Args:
            open_file (file-like object): File handle from which to read the data

        Yields:
            tuple: line number, contents of the line being read
        """
        reader = cls._gz_reader(open_file)
        try:
            # Try to just load the reader as regular json first
            yield 1, json.load(reader)
            return
        except ValueError:
            pass

        # Iterate over the lines, returning each
        # This could be dicts from a jsonlines.Reader, or raw strings
        for line_num, line in enumerate(cls._jsonlines_reader(reader), start=1):
            yield line_num, line.strip() if isinstance(line, str) else line

    def _read_file(self):
        """Download and read the contents of the S3 file

        Yields:
            tuple: line number, contents of the line being read
        """
        bucket = self._unquote(self.bucket)
        key = self._unquote(self.key)

        # Use tempfile.TemporaryFile to do the download
        # This will automatically close/get garbage collected upon completion
        with tempfile.TemporaryFile() as download:
            client = boto3.resource('s3', region_name=self.region).Bucket(bucket)
            start_time = time.time()
            LOGGER.info('[S3Payload] Starting download from S3: %s/%s [%s]', bucket, key, self.size)

            try:
                client.download_fileobj(key, download)
            except (IOError, ClientError):
                LOGGER.exception('Failed to download object from S3')
                raise

            total_time = time.time() - start_time
            LOGGER.info('Completed download in %s seconds', round(total_time, 2))

            # Log a metric on how long this object took to download
            MetricLogger.log_metric(FUNCTION_NAME, MetricLogger.S3_DOWNLOAD_TIME, total_time)

            for line_num, line in self._read_downloaded_object(download):
                yield line_num, line

            # Reading was a success, so truncate the file contents and return
            download.seek(0)
            download.truncate()

    def _pre_parse(self):
        """Pre-parsing method for S3 objects

        Downloads the s3 object into the system's temp directory for reading. The
        file is kept open as a tempfile.TemporaryFile to ensure proper cleanup
        when reading finishes.

        Yields:
            Instances of PayloadRecord back to the caller containing the current log data
        """
        if not self._check_size():
            return  # _check_size can raise an exception as well

        line_num = 0
        for line_num, data in self._read_file():
            yield PayloadRecord(data)

        MetricLogger.log_metric(FUNCTION_NAME, MetricLogger.TOTAL_S3_RECORDS, line_num)

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	import json	1✔
17	import logging	1✔
18	import urllib.request	1✔
19	import urllib.parse	1✔
20	import urllib.error	1✔
21	import gzip	1✔
22	import os	1✔
23	import tempfile	1✔
24	import subprocess	1✔
25	import time	1✔
26
27	import boto3	1✔
28	from botocore.exceptions import ClientError	1✔
29	import jsonlines	1✔
30
31	from streamalert.classifier.payload.payload_base import (	1✔
32	PayloadRecord,
33	RegisterInput,
34	StreamPayload
35	)
36	from streamalert.shared import CLASSIFIER_FUNCTION_NAME as FUNCTION_NAME	1✔
37	from streamalert.shared.logger import get_logger	1✔
38	from streamalert.shared.metrics import MetricLogger	1✔
39
40
41	LOGGER = get_logger(__name__)	1✔
42	LOGGER_DEBUG_ENABLED = LOGGER.isEnabledFor(logging.DEBUG)	1✔
43
44
45	class S3PayloadError(Exception):	1✔
46	"""Exception for S3Payload errors"""
47
48
49	@RegisterInput	1✔
50	class S3Payload(StreamPayload):	1✔
51	"""S3Payload class"""
52
53	MAX_S3_SIZE = 128 * 1024 * 1024	1✔
54
55	@property	1✔
56	def bucket(self):
57	return self.raw_record['s3']['bucket']['name']	1✔
58
59	@property	1✔
60	def key(self):
61	return self.raw_record['s3']['object']['key']	1✔
62
63	@property	1✔
64	def size(self):
65	return int(self.raw_record['s3']['object']['size'])	1✔
66
67	@property	1✔
68	def region(self):
69	return self.raw_record['awsRegion']	1✔
70
71	@property	1✔
72	def display_size(self):
73	"""Calculate and format a size for printing"""
74	size_kb = round(self.size / 1024.0, 2)	1✔
75	size_mb = round(size_kb / 1024.0, 2)	1✔
76	return '{}MB'.format(size_mb) if size_mb else '{}KB'.format(size_kb)	1✔
77
78	@classmethod	1✔
79	def service(cls):
80	return 's3'	1✔
81
82	@classmethod	1✔
83	def _unquote(cls, data):
84	# Use the urllib unquote method to decode any url encoded characters
85	# (ie - %26 --> &) from the bucket and key names
86	return urllib.parse.unquote(data)	1✔
87
88	def _check_size(self):	1✔
89	"""Ensure the S3 file's size is not too large to download into the Lambda environment
90
91	Returns:
92	bool: True if the file is smaller than 128 MB, False otherwise
93	"""
94	# Ignore 0 size files
95	if self.size == 0:	1✔
96	LOGGER.warning('S3 file size is 0 bytes, skipping: %s/%s', self.bucket, self.key)	1✔
97	return False	1✔
98
99	# size greater than 128MB
100	if self.size > self.MAX_S3_SIZE:	1✔
101	raise S3PayloadError('S3 object {}/{} is too large and cannot be downloaded '	1✔
102	'from S3: {}'.format(self.bucket, self.key, self.display_size))
103
104	return True	1✔
105
106	@staticmethod	1✔
107	def _cleanup():
108	"""Cleanup method to remove all objects in the Lambda container's temp directory"""
109	# Do nothing if this is not running in AWS Lambda
110	if 'LAMBDA_RUNTIME_DIR' not in os.environ:	1✔
111	return	1✔
112
113	LOGGER.debug('Shredding temp directory')	1✔
114
115	for root, dirs, files in os.walk(tempfile.gettempdir(), topdown=False):	1✔
116	for name in files:	1✔
117	subprocess.check_call([ # nosec	1✔
118	'shred', '--force', '--iterations=1',
119	'--remove', os.path.join(root, name)])
120	for name in dirs:	1✔
121	os.rmdir(os.path.join(root, name)) # nosec	1✔
122
123	@staticmethod	1✔
124	def _gz_reader(open_file):
125	open_file.seek(0)	1✔
126	reader = gzip.GzipFile(fileobj=open_file, mode='r')	1✔
127	try:	1✔
128	# Test to ensure this is gzip data, then rewind
129	reader.read(1)	1✔
130	reader.rewind()	1✔
131	except IOError:	1✔
132	# Fall back on the default reader
133	reader = open_file	1✔
134	reader.seek(0)	1✔
135
136	# Return either the gzip reader or the original reader
137	return reader	1✔
138
139	@staticmethod	1✔
140	def _jsonlines_reader(open_file):
141	open_file.seek(0)	1✔
142	reader = None	1✔
143	try:	1✔
144	json_lines = jsonlines.Reader(open_file)	1✔
145	json_lines.read()	1✔
146	reader = json_lines	1✔
147	except ValueError:	1✔
148	reader = open_file	1✔
149
150	open_file.seek(0)	1✔
151	return reader	1✔
152
153	@classmethod	1✔
154	def _read_downloaded_object(cls, open_file):
155	"""Read the contents of the downloaded S3 file using the open file handle
156
157	Args:
158	open_file (file-like object): File handle from which to read the data
159
160	Yields:
161	tuple: line number, contents of the line being read
162	"""
163	reader = cls._gz_reader(open_file)	1✔
164	try:	1✔
165	# Try to just load the reader as regular json first
166	yield 1, json.load(reader)	1✔
167	return	1✔
168	except ValueError:	1✔
169	pass	1✔
170
171	# Iterate over the lines, returning each
172	# This could be dicts from a jsonlines.Reader, or raw strings
173	for line_num, line in enumerate(cls._jsonlines_reader(reader), start=1):	1✔
174	yield line_num, line.strip() if isinstance(line, str) else line	1✔
175
176	def _read_file(self):	1✔
177	"""Download and read the contents of the S3 file
178
179	Yields:
180	tuple: line number, contents of the line being read
181	"""
182	bucket = self._unquote(self.bucket)	1✔
183	key = self._unquote(self.key)	1✔
184
185	# Use tempfile.TemporaryFile to do the download
186	# This will automatically close/get garbage collected upon completion
187	with tempfile.TemporaryFile() as download:	1✔
188	client = boto3.resource('s3', region_name=self.region).Bucket(bucket)	1✔
189	start_time = time.time()	1✔
190	LOGGER.info('[S3Payload] Starting download from S3: %s/%s [%s]', bucket, key, self.size)	1✔
191
192	try:	1✔
193	client.download_fileobj(key, download)	1✔
194	except (IOError, ClientError):	1✔
195	LOGGER.exception('Failed to download object from S3')	1✔
196	raise	1✔
197
198	total_time = time.time() - start_time	1✔
199	LOGGER.info('Completed download in %s seconds', round(total_time, 2))	1✔
200
201	# Log a metric on how long this object took to download
202	MetricLogger.log_metric(FUNCTION_NAME, MetricLogger.S3_DOWNLOAD_TIME, total_time)	1✔
203
204	for line_num, line in self._read_downloaded_object(download):	1✔
205	yield line_num, line	1✔
206
207	# Reading was a success, so truncate the file contents and return
208	download.seek(0)	1✔
209	download.truncate()	1✔
210
211	def _pre_parse(self):	1✔
212	"""Pre-parsing method for S3 objects
213
214	Downloads the s3 object into the system's temp directory for reading. The
215	file is kept open as a tempfile.TemporaryFile to ensure proper cleanup
216	when reading finishes.
217
218	Yields:
219	Instances of PayloadRecord back to the caller containing the current log data
220	"""
221	if not self._check_size():	1✔
222	return # _check_size can raise an exception as well	×
223
224	line_num = 0	1✔
225	for line_num, data in self._read_file():	1✔
226	yield PayloadRecord(data)	1✔
227
228	MetricLogger.log_metric(FUNCTION_NAME, MetricLogger.TOTAL_S3_RECORDS, line_num)	1✔

airbnb / streamalert / 195255325

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