634 lines
20 KiB
Python
634 lines
20 KiB
Python
# Copyright 2018 Amazon.com, Inc. or its affiliates. All Rights Reserved.
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License"). You
|
|
# may not use this file except in compliance with the License. A copy of
|
|
# the License is located at
|
|
#
|
|
# http://aws.amazon.com/apache2.0/
|
|
#
|
|
# or in the "license" file accompanying this file. This file 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.
|
|
"""Binary Event Stream Decoding """
|
|
|
|
from binascii import crc32
|
|
from struct import unpack
|
|
|
|
from botocore.exceptions import EventStreamError
|
|
|
|
# byte length of the prelude (total_length + header_length + prelude_crc)
|
|
_PRELUDE_LENGTH = 12
|
|
_MAX_HEADERS_LENGTH = 128 * 1024 # 128 Kb
|
|
_MAX_PAYLOAD_LENGTH = 16 * 1024**2 # 16 Mb
|
|
|
|
|
|
class ParserError(Exception):
|
|
"""Base binary flow encoding parsing exception."""
|
|
|
|
pass
|
|
|
|
|
|
class DuplicateHeader(ParserError):
|
|
"""Duplicate header found in the event."""
|
|
|
|
def __init__(self, header):
|
|
message = 'Duplicate header present: "%s"' % header
|
|
super().__init__(message)
|
|
|
|
|
|
class InvalidHeadersLength(ParserError):
|
|
"""Headers length is longer than the maximum."""
|
|
|
|
def __init__(self, length):
|
|
message = 'Header length of {} exceeded the maximum of {}'.format(
|
|
length,
|
|
_MAX_HEADERS_LENGTH,
|
|
)
|
|
super().__init__(message)
|
|
|
|
|
|
class InvalidPayloadLength(ParserError):
|
|
"""Payload length is longer than the maximum."""
|
|
|
|
def __init__(self, length):
|
|
message = 'Payload length of {} exceeded the maximum of {}'.format(
|
|
length,
|
|
_MAX_PAYLOAD_LENGTH,
|
|
)
|
|
super().__init__(message)
|
|
|
|
|
|
class ChecksumMismatch(ParserError):
|
|
"""Calculated checksum did not match the expected checksum."""
|
|
|
|
def __init__(self, expected, calculated):
|
|
message = (
|
|
'Checksum mismatch: expected 0x{:08x}, calculated 0x{:08x}'.format(
|
|
expected,
|
|
calculated,
|
|
)
|
|
)
|
|
super().__init__(message)
|
|
|
|
|
|
class NoInitialResponseError(ParserError):
|
|
"""An event of type initial-response was not received.
|
|
|
|
This exception is raised when the event stream produced no events or
|
|
the first event in the stream was not of the initial-response type.
|
|
"""
|
|
|
|
def __init__(self):
|
|
message = 'First event was not of the initial-response type'
|
|
super().__init__(message)
|
|
|
|
|
|
class DecodeUtils:
|
|
"""Unpacking utility functions used in the decoder.
|
|
|
|
All methods on this class take raw bytes and return a tuple containing
|
|
the value parsed from the bytes and the number of bytes consumed to parse
|
|
that value.
|
|
"""
|
|
|
|
UINT8_BYTE_FORMAT = '!B'
|
|
UINT16_BYTE_FORMAT = '!H'
|
|
UINT32_BYTE_FORMAT = '!I'
|
|
INT8_BYTE_FORMAT = '!b'
|
|
INT16_BYTE_FORMAT = '!h'
|
|
INT32_BYTE_FORMAT = '!i'
|
|
INT64_BYTE_FORMAT = '!q'
|
|
PRELUDE_BYTE_FORMAT = '!III'
|
|
|
|
# uint byte size to unpack format
|
|
UINT_BYTE_FORMAT = {
|
|
1: UINT8_BYTE_FORMAT,
|
|
2: UINT16_BYTE_FORMAT,
|
|
4: UINT32_BYTE_FORMAT,
|
|
}
|
|
|
|
@staticmethod
|
|
def unpack_true(data):
|
|
"""This method consumes none of the provided bytes and returns True.
|
|
|
|
:type data: bytes
|
|
:param data: The bytes to parse from. This is ignored in this method.
|
|
|
|
:rtype: tuple
|
|
:rtype: (bool, int)
|
|
:returns: The tuple (True, 0)
|
|
"""
|
|
return True, 0
|
|
|
|
@staticmethod
|
|
def unpack_false(data):
|
|
"""This method consumes none of the provided bytes and returns False.
|
|
|
|
:type data: bytes
|
|
:param data: The bytes to parse from. This is ignored in this method.
|
|
|
|
:rtype: tuple
|
|
:rtype: (bool, int)
|
|
:returns: The tuple (False, 0)
|
|
"""
|
|
return False, 0
|
|
|
|
@staticmethod
|
|
def unpack_uint8(data):
|
|
"""Parse an unsigned 8-bit integer from the bytes.
|
|
|
|
:type data: bytes
|
|
:param data: The bytes to parse from.
|
|
|
|
:rtype: (int, int)
|
|
:returns: A tuple containing the (parsed integer value, bytes consumed)
|
|
"""
|
|
value = unpack(DecodeUtils.UINT8_BYTE_FORMAT, data[:1])[0]
|
|
return value, 1
|
|
|
|
@staticmethod
|
|
def unpack_uint32(data):
|
|
"""Parse an unsigned 32-bit integer from the bytes.
|
|
|
|
:type data: bytes
|
|
:param data: The bytes to parse from.
|
|
|
|
:rtype: (int, int)
|
|
:returns: A tuple containing the (parsed integer value, bytes consumed)
|
|
"""
|
|
value = unpack(DecodeUtils.UINT32_BYTE_FORMAT, data[:4])[0]
|
|
return value, 4
|
|
|
|
@staticmethod
|
|
def unpack_int8(data):
|
|
"""Parse a signed 8-bit integer from the bytes.
|
|
|
|
:type data: bytes
|
|
:param data: The bytes to parse from.
|
|
|
|
:rtype: (int, int)
|
|
:returns: A tuple containing the (parsed integer value, bytes consumed)
|
|
"""
|
|
value = unpack(DecodeUtils.INT8_BYTE_FORMAT, data[:1])[0]
|
|
return value, 1
|
|
|
|
@staticmethod
|
|
def unpack_int16(data):
|
|
"""Parse a signed 16-bit integer from the bytes.
|
|
|
|
:type data: bytes
|
|
:param data: The bytes to parse from.
|
|
|
|
:rtype: tuple
|
|
:rtype: (int, int)
|
|
:returns: A tuple containing the (parsed integer value, bytes consumed)
|
|
"""
|
|
value = unpack(DecodeUtils.INT16_BYTE_FORMAT, data[:2])[0]
|
|
return value, 2
|
|
|
|
@staticmethod
|
|
def unpack_int32(data):
|
|
"""Parse a signed 32-bit integer from the bytes.
|
|
|
|
:type data: bytes
|
|
:param data: The bytes to parse from.
|
|
|
|
:rtype: tuple
|
|
:rtype: (int, int)
|
|
:returns: A tuple containing the (parsed integer value, bytes consumed)
|
|
"""
|
|
value = unpack(DecodeUtils.INT32_BYTE_FORMAT, data[:4])[0]
|
|
return value, 4
|
|
|
|
@staticmethod
|
|
def unpack_int64(data):
|
|
"""Parse a signed 64-bit integer from the bytes.
|
|
|
|
:type data: bytes
|
|
:param data: The bytes to parse from.
|
|
|
|
:rtype: tuple
|
|
:rtype: (int, int)
|
|
:returns: A tuple containing the (parsed integer value, bytes consumed)
|
|
"""
|
|
value = unpack(DecodeUtils.INT64_BYTE_FORMAT, data[:8])[0]
|
|
return value, 8
|
|
|
|
@staticmethod
|
|
def unpack_byte_array(data, length_byte_size=2):
|
|
"""Parse a variable length byte array from the bytes.
|
|
|
|
The bytes are expected to be in the following format:
|
|
[ length ][0 ... length bytes]
|
|
where length is an unsigned integer represented in the smallest number
|
|
of bytes to hold the maximum length of the array.
|
|
|
|
:type data: bytes
|
|
:param data: The bytes to parse from.
|
|
|
|
:type length_byte_size: int
|
|
:param length_byte_size: The byte size of the preceeding integer that
|
|
represents the length of the array. Supported values are 1, 2, and 4.
|
|
|
|
:rtype: (bytes, int)
|
|
:returns: A tuple containing the (parsed byte array, bytes consumed).
|
|
"""
|
|
uint_byte_format = DecodeUtils.UINT_BYTE_FORMAT[length_byte_size]
|
|
length = unpack(uint_byte_format, data[:length_byte_size])[0]
|
|
bytes_end = length + length_byte_size
|
|
array_bytes = data[length_byte_size:bytes_end]
|
|
return array_bytes, bytes_end
|
|
|
|
@staticmethod
|
|
def unpack_utf8_string(data, length_byte_size=2):
|
|
"""Parse a variable length utf-8 string from the bytes.
|
|
|
|
The bytes are expected to be in the following format:
|
|
[ length ][0 ... length bytes]
|
|
where length is an unsigned integer represented in the smallest number
|
|
of bytes to hold the maximum length of the array and the following
|
|
bytes are a valid utf-8 string.
|
|
|
|
:type data: bytes
|
|
:param bytes: The bytes to parse from.
|
|
|
|
:type length_byte_size: int
|
|
:param length_byte_size: The byte size of the preceeding integer that
|
|
represents the length of the array. Supported values are 1, 2, and 4.
|
|
|
|
:rtype: (str, int)
|
|
:returns: A tuple containing the (utf-8 string, bytes consumed).
|
|
"""
|
|
array_bytes, consumed = DecodeUtils.unpack_byte_array(
|
|
data, length_byte_size
|
|
)
|
|
return array_bytes.decode('utf-8'), consumed
|
|
|
|
@staticmethod
|
|
def unpack_uuid(data):
|
|
"""Parse a 16-byte uuid from the bytes.
|
|
|
|
:type data: bytes
|
|
:param data: The bytes to parse from.
|
|
|
|
:rtype: (bytes, int)
|
|
:returns: A tuple containing the (uuid bytes, bytes consumed).
|
|
"""
|
|
return data[:16], 16
|
|
|
|
@staticmethod
|
|
def unpack_prelude(data):
|
|
"""Parse the prelude for an event stream message from the bytes.
|
|
|
|
The prelude for an event stream message has the following format:
|
|
[total_length][header_length][prelude_crc]
|
|
where each field is an unsigned 32-bit integer.
|
|
|
|
:rtype: ((int, int, int), int)
|
|
:returns: A tuple of ((total_length, headers_length, prelude_crc),
|
|
consumed)
|
|
"""
|
|
return (unpack(DecodeUtils.PRELUDE_BYTE_FORMAT, data), _PRELUDE_LENGTH)
|
|
|
|
|
|
def _validate_checksum(data, checksum, crc=0):
|
|
# To generate the same numeric value across all Python versions and
|
|
# platforms use crc32(data) & 0xffffffff.
|
|
computed_checksum = crc32(data, crc) & 0xFFFFFFFF
|
|
if checksum != computed_checksum:
|
|
raise ChecksumMismatch(checksum, computed_checksum)
|
|
|
|
|
|
class MessagePrelude:
|
|
"""Represents the prelude of an event stream message."""
|
|
|
|
def __init__(self, total_length, headers_length, crc):
|
|
self.total_length = total_length
|
|
self.headers_length = headers_length
|
|
self.crc = crc
|
|
|
|
@property
|
|
def payload_length(self):
|
|
"""Calculates the total payload length.
|
|
|
|
The extra minus 4 bytes is for the message CRC.
|
|
|
|
:rtype: int
|
|
:returns: The total payload length.
|
|
"""
|
|
return self.total_length - self.headers_length - _PRELUDE_LENGTH - 4
|
|
|
|
@property
|
|
def payload_end(self):
|
|
"""Calculates the byte offset for the end of the message payload.
|
|
|
|
The extra minus 4 bytes is for the message CRC.
|
|
|
|
:rtype: int
|
|
:returns: The byte offset from the beginning of the event stream
|
|
message to the end of the payload.
|
|
"""
|
|
return self.total_length - 4
|
|
|
|
@property
|
|
def headers_end(self):
|
|
"""Calculates the byte offset for the end of the message headers.
|
|
|
|
:rtype: int
|
|
:returns: The byte offset from the beginning of the event stream
|
|
message to the end of the headers.
|
|
"""
|
|
return _PRELUDE_LENGTH + self.headers_length
|
|
|
|
|
|
class EventStreamMessage:
|
|
"""Represents an event stream message."""
|
|
|
|
def __init__(self, prelude, headers, payload, crc):
|
|
self.prelude = prelude
|
|
self.headers = headers
|
|
self.payload = payload
|
|
self.crc = crc
|
|
|
|
def to_response_dict(self, status_code=200):
|
|
message_type = self.headers.get(':message-type')
|
|
if message_type == 'error' or message_type == 'exception':
|
|
status_code = 400
|
|
return {
|
|
'status_code': status_code,
|
|
'headers': self.headers,
|
|
'body': self.payload,
|
|
}
|
|
|
|
|
|
class EventStreamHeaderParser:
|
|
"""Parses the event headers from an event stream message.
|
|
|
|
Expects all of the header data upfront and creates a dictionary of headers
|
|
to return. This object can be reused multiple times to parse the headers
|
|
from multiple event stream messages.
|
|
"""
|
|
|
|
# Maps header type to appropriate unpacking function
|
|
# These unpacking functions return the value and the amount unpacked
|
|
_HEADER_TYPE_MAP = {
|
|
# boolean_true
|
|
0: DecodeUtils.unpack_true,
|
|
# boolean_false
|
|
1: DecodeUtils.unpack_false,
|
|
# byte
|
|
2: DecodeUtils.unpack_int8,
|
|
# short
|
|
3: DecodeUtils.unpack_int16,
|
|
# integer
|
|
4: DecodeUtils.unpack_int32,
|
|
# long
|
|
5: DecodeUtils.unpack_int64,
|
|
# byte_array
|
|
6: DecodeUtils.unpack_byte_array,
|
|
# string
|
|
7: DecodeUtils.unpack_utf8_string,
|
|
# timestamp
|
|
8: DecodeUtils.unpack_int64,
|
|
# uuid
|
|
9: DecodeUtils.unpack_uuid,
|
|
}
|
|
|
|
def __init__(self):
|
|
self._data = None
|
|
|
|
def parse(self, data):
|
|
"""Parses the event stream headers from an event stream message.
|
|
|
|
:type data: bytes
|
|
:param data: The bytes that correspond to the headers section of an
|
|
event stream message.
|
|
|
|
:rtype: dict
|
|
:returns: A dicionary of header key, value pairs.
|
|
"""
|
|
self._data = data
|
|
return self._parse_headers()
|
|
|
|
def _parse_headers(self):
|
|
headers = {}
|
|
while self._data:
|
|
name, value = self._parse_header()
|
|
if name in headers:
|
|
raise DuplicateHeader(name)
|
|
headers[name] = value
|
|
return headers
|
|
|
|
def _parse_header(self):
|
|
name = self._parse_name()
|
|
value = self._parse_value()
|
|
return name, value
|
|
|
|
def _parse_name(self):
|
|
name, consumed = DecodeUtils.unpack_utf8_string(self._data, 1)
|
|
self._advance_data(consumed)
|
|
return name
|
|
|
|
def _parse_type(self):
|
|
type, consumed = DecodeUtils.unpack_uint8(self._data)
|
|
self._advance_data(consumed)
|
|
return type
|
|
|
|
def _parse_value(self):
|
|
header_type = self._parse_type()
|
|
value_unpacker = self._HEADER_TYPE_MAP[header_type]
|
|
value, consumed = value_unpacker(self._data)
|
|
self._advance_data(consumed)
|
|
return value
|
|
|
|
def _advance_data(self, consumed):
|
|
self._data = self._data[consumed:]
|
|
|
|
|
|
class EventStreamBuffer:
|
|
"""Streaming based event stream buffer
|
|
|
|
A buffer class that wraps bytes from an event stream providing parsed
|
|
messages as they become available via an iterable interface.
|
|
"""
|
|
|
|
def __init__(self):
|
|
self._data = b''
|
|
self._prelude = None
|
|
self._header_parser = EventStreamHeaderParser()
|
|
|
|
def add_data(self, data):
|
|
"""Add data to the buffer.
|
|
|
|
:type data: bytes
|
|
:param data: The bytes to add to the buffer to be used when parsing
|
|
"""
|
|
self._data += data
|
|
|
|
def _validate_prelude(self, prelude):
|
|
if prelude.headers_length > _MAX_HEADERS_LENGTH:
|
|
raise InvalidHeadersLength(prelude.headers_length)
|
|
|
|
if prelude.payload_length > _MAX_PAYLOAD_LENGTH:
|
|
raise InvalidPayloadLength(prelude.payload_length)
|
|
|
|
def _parse_prelude(self):
|
|
prelude_bytes = self._data[:_PRELUDE_LENGTH]
|
|
raw_prelude, _ = DecodeUtils.unpack_prelude(prelude_bytes)
|
|
prelude = MessagePrelude(*raw_prelude)
|
|
self._validate_prelude(prelude)
|
|
# The minus 4 removes the prelude crc from the bytes to be checked
|
|
_validate_checksum(prelude_bytes[: _PRELUDE_LENGTH - 4], prelude.crc)
|
|
return prelude
|
|
|
|
def _parse_headers(self):
|
|
header_bytes = self._data[_PRELUDE_LENGTH : self._prelude.headers_end]
|
|
return self._header_parser.parse(header_bytes)
|
|
|
|
def _parse_payload(self):
|
|
prelude = self._prelude
|
|
payload_bytes = self._data[prelude.headers_end : prelude.payload_end]
|
|
return payload_bytes
|
|
|
|
def _parse_message_crc(self):
|
|
prelude = self._prelude
|
|
crc_bytes = self._data[prelude.payload_end : prelude.total_length]
|
|
message_crc, _ = DecodeUtils.unpack_uint32(crc_bytes)
|
|
return message_crc
|
|
|
|
def _parse_message_bytes(self):
|
|
# The minus 4 includes the prelude crc to the bytes to be checked
|
|
message_bytes = self._data[
|
|
_PRELUDE_LENGTH - 4 : self._prelude.payload_end
|
|
]
|
|
return message_bytes
|
|
|
|
def _validate_message_crc(self):
|
|
message_crc = self._parse_message_crc()
|
|
message_bytes = self._parse_message_bytes()
|
|
_validate_checksum(message_bytes, message_crc, crc=self._prelude.crc)
|
|
return message_crc
|
|
|
|
def _parse_message(self):
|
|
crc = self._validate_message_crc()
|
|
headers = self._parse_headers()
|
|
payload = self._parse_payload()
|
|
message = EventStreamMessage(self._prelude, headers, payload, crc)
|
|
self._prepare_for_next_message()
|
|
return message
|
|
|
|
def _prepare_for_next_message(self):
|
|
# Advance the data and reset the current prelude
|
|
self._data = self._data[self._prelude.total_length :]
|
|
self._prelude = None
|
|
|
|
def next(self):
|
|
"""Provides the next available message parsed from the stream
|
|
|
|
:rtype: EventStreamMessage
|
|
:returns: The next event stream message
|
|
"""
|
|
if len(self._data) < _PRELUDE_LENGTH:
|
|
raise StopIteration()
|
|
|
|
if self._prelude is None:
|
|
self._prelude = self._parse_prelude()
|
|
|
|
if len(self._data) < self._prelude.total_length:
|
|
raise StopIteration()
|
|
|
|
return self._parse_message()
|
|
|
|
def __next__(self):
|
|
return self.next()
|
|
|
|
def __iter__(self):
|
|
return self
|
|
|
|
|
|
class EventStream:
|
|
"""Wrapper class for an event stream body.
|
|
|
|
This wraps the underlying streaming body, parsing it for individual events
|
|
and yielding them as they come available through the iterator interface.
|
|
|
|
The following example uses the S3 select API to get structured data out of
|
|
an object stored in S3 using an event stream.
|
|
|
|
**Example:**
|
|
::
|
|
from botocore.session import Session
|
|
|
|
s3 = Session().create_client('s3')
|
|
response = s3.select_object_content(
|
|
Bucket='bucketname',
|
|
Key='keyname',
|
|
ExpressionType='SQL',
|
|
RequestProgress={'Enabled': True},
|
|
Expression="SELECT * FROM S3Object s",
|
|
InputSerialization={'CSV': {}},
|
|
OutputSerialization={'CSV': {}},
|
|
)
|
|
# This is the event stream in the response
|
|
event_stream = response['Payload']
|
|
end_event_received = False
|
|
with open('output', 'wb') as f:
|
|
# Iterate over events in the event stream as they come
|
|
for event in event_stream:
|
|
# If we received a records event, write the data to a file
|
|
if 'Records' in event:
|
|
data = event['Records']['Payload']
|
|
f.write(data)
|
|
# If we received a progress event, print the details
|
|
elif 'Progress' in event:
|
|
print(event['Progress']['Details'])
|
|
# End event indicates that the request finished successfully
|
|
elif 'End' in event:
|
|
print('Result is complete')
|
|
end_event_received = True
|
|
if not end_event_received:
|
|
raise Exception("End event not received, request incomplete.")
|
|
"""
|
|
|
|
def __init__(self, raw_stream, output_shape, parser, operation_name):
|
|
self._raw_stream = raw_stream
|
|
self._output_shape = output_shape
|
|
self._operation_name = operation_name
|
|
self._parser = parser
|
|
self._event_generator = self._create_raw_event_generator()
|
|
|
|
def __iter__(self):
|
|
for event in self._event_generator:
|
|
parsed_event = self._parse_event(event)
|
|
if parsed_event:
|
|
yield parsed_event
|
|
|
|
def _create_raw_event_generator(self):
|
|
event_stream_buffer = EventStreamBuffer()
|
|
for chunk in self._raw_stream.stream():
|
|
event_stream_buffer.add_data(chunk)
|
|
yield from event_stream_buffer
|
|
|
|
def _parse_event(self, event):
|
|
response_dict = event.to_response_dict()
|
|
parsed_response = self._parser.parse(response_dict, self._output_shape)
|
|
if response_dict['status_code'] == 200:
|
|
return parsed_response
|
|
else:
|
|
raise EventStreamError(parsed_response, self._operation_name)
|
|
|
|
def get_initial_response(self):
|
|
try:
|
|
initial_event = next(self._event_generator)
|
|
event_type = initial_event.headers.get(':event-type')
|
|
if event_type == 'initial-response':
|
|
return initial_event
|
|
except StopIteration:
|
|
pass
|
|
raise NoInitialResponseError()
|
|
|
|
def close(self):
|
|
"""Closes the underlying streaming body."""
|
|
self._raw_stream.close()
|