def test_encryptor_finalize(patch_default_backend, patch_cipher):
    tester = Encryptor(algorithm=MagicMock(), key=sentinel.key, associated_data=sentinel.aad, iv=sentinel.iv)

    test = tester.finalize()

    tester._encryptor.finalize.assert_called_once_with()
    assert test is tester._encryptor.finalize.return_value
예제 #2
0
class StreamEncryptor(_EncryptionStream):  # pylint: disable=too-many-instance-attributes
    """Provides a streaming encryptor for encrypting a stream source.
    Behaves as a standard file-like object.

    .. note::
        Take care when encrypting framed messages with large frame length and large non-framed
        messages.  See :class:`aws_encryption_sdk.stream` for more details.

    .. note::
        If config is provided, all other parameters are ignored.

    :param config: Client configuration object (config or individual parameters required)
    :type config: aws_encryption_sdk.streaming_client.EncryptorConfig
    :param source: Source data to encrypt or decrypt
    :type source: str, bytes, io.IOBase, or file
    :param materials_manager: `CryptoMaterialsManager` from which to obtain cryptographic materials
        (either `materials_manager` or `key_provider` required)
    :type materials_manager: aws_encryption_sdk.materials_manager.base.CryptoMaterialsManager
    :param key_provider: `MasterKeyProvider` from which to obtain data keys for encryption
        (either `materials_manager` or `key_provider` required)
    :type key_provider: aws_encryption_sdk.key_providers.base.MasterKeyProvider
    :param int source_length: Length of source data (optional)

        .. note::
            If source_length is not provided and unframed message is being written or read() is called,
            will attempt to seek() to the end of the stream and tell() to find the length of source data.

        .. note::
            .. versionadded:: 1.3.0

            If `source_length` and `materials_manager` are both provided, the total plaintext bytes
            encrypted will not be allowed to exceed `source_length`. To maintain backwards compatibility,
            this is not enforced if a `key_provider` is provided.

    :param dict encryption_context: Dictionary defining encryption context
    :param algorithm: Algorithm to use for encryption
    :type algorithm: aws_encryption_sdk.identifiers.Algorithm
    :param int frame_length: Frame length in bytes
    """

    _config_class = EncryptorConfig

    def __init__(self, **kwargs):  # pylint: disable=unused-argument,super-init-not-called
        """Prepares necessary initial values."""
        self.sequence_number = 1

        self.content_type = aws_encryption_sdk.internal.utils.content_type(
            self.config.frame_length)
        self._bytes_encrypted = 0

        if (self.config.frame_length == 0
                and (self.config.source_length is not None
                     and self.config.source_length > MAX_NON_FRAMED_SIZE)):
            raise SerializationError('Source too large for non-framed message')

    def ciphertext_length(self):
        """Returns the length of the resulting ciphertext message in bytes.

        :rtype: int
        """
        return aws_encryption_sdk.internal.formatting.ciphertext_length(
            header=self.header, plaintext_length=self.stream_length)

    def _prep_message(self):
        """Performs initial message setup.

        :raises MasterKeyProviderError: if primary master key is not a member of supplied MasterKeyProvider
        :raises MasterKeyProviderError: if no Master Keys are returned from key_provider
        """
        message_id = aws_encryption_sdk.internal.utils.message_id()

        try:
            plaintext_length = self.stream_length
        except NotSupportedError:
            plaintext_length = None
        encryption_materials_request = EncryptionMaterialsRequest(
            algorithm=self.config.algorithm,
            encryption_context=self.config.encryption_context.copy(),
            frame_length=self.config.frame_length,
            plaintext_rostream=aws_encryption_sdk.internal.utils.ROStream(
                self.source_stream),
            plaintext_length=plaintext_length)
        self._encryption_materials = self.config.materials_manager.get_encryption_materials(
            request=encryption_materials_request)

        if self.config.algorithm is not None and self._encryption_materials.algorithm != self.config.algorithm:
            raise ActionNotAllowedError(
                ('Cryptographic materials manager provided algorithm suite'
                 ' differs from algorithm suite in request.\n'
                 'Required: {requested}\n'
                 'Provided: {provided}').format(
                     requested=self.config.algorithm,
                     provided=self._encryption_materials.algorithm))

        if self._encryption_materials.signing_key is None:
            self.signer = None
        else:
            self.signer = Signer.from_key_bytes(
                algorithm=self._encryption_materials.algorithm,
                key_bytes=self._encryption_materials.signing_key)
        aws_encryption_sdk.internal.utils.validate_frame_length(
            frame_length=self.config.frame_length,
            algorithm=self._encryption_materials.algorithm)

        self._derived_data_key = derive_data_encryption_key(
            source_key=self._encryption_materials.data_encryption_key.data_key,
            algorithm=self._encryption_materials.algorithm,
            message_id=message_id)

        self._header = MessageHeader(
            version=VERSION,
            type=TYPE,
            algorithm=self._encryption_materials.algorithm,
            message_id=message_id,
            encryption_context=self._encryption_materials.encryption_context,
            encrypted_data_keys=self._encryption_materials.encrypted_data_keys,
            content_type=self.content_type,
            content_aad_length=0,
            header_iv_length=self._encryption_materials.algorithm.iv_len,
            frame_length=self.config.frame_length)
        self._write_header()
        if self.content_type == ContentType.NO_FRAMING:
            self._prep_non_framed()
        self._message_prepped = True

    def _write_header(self):
        """Builds the message header and writes it to the output stream."""
        self.output_buffer += aws_encryption_sdk.internal.formatting.serialize.serialize_header(
            header=self._header, signer=self.signer)
        self.output_buffer += aws_encryption_sdk.internal.formatting.serialize.serialize_header_auth(
            algorithm=self._encryption_materials.algorithm,
            header=self.output_buffer,
            data_encryption_key=self._derived_data_key,
            signer=self.signer)

    def _prep_non_framed(self):
        """Prepare the opening data for a non-framed message."""
        aad_content_string = aws_encryption_sdk.internal.utils.get_aad_content_string(
            content_type=self.content_type, is_final_frame=True)
        associated_data = aws_encryption_sdk.internal.formatting.encryption_context.assemble_content_aad(
            message_id=self._header.message_id,
            aad_content_string=aad_content_string,
            seq_num=1,
            length=self.stream_length)
        self.encryptor = Encryptor(
            algorithm=self._encryption_materials.algorithm,
            key=self._derived_data_key,
            associated_data=associated_data,
            iv=non_framed_body_iv(self._encryption_materials.algorithm))
        self.output_buffer += aws_encryption_sdk.internal.formatting.serialize.serialize_non_framed_open(
            algorithm=self._encryption_materials.algorithm,
            iv=self.encryptor.iv,
            plaintext_length=self.stream_length,
            signer=self.signer)

    def _read_bytes_to_non_framed_body(self, b):
        """Reads the requested number of bytes from source to a streaming non-framed message body.

        :param int b: Number of bytes to read
        :returns: Encrypted bytes from source stream
        :rtype: bytes
        """
        _LOGGER.debug('Reading %s bytes', b)
        plaintext = self.source_stream.read(b)
        if self.tell() + len(plaintext) > MAX_NON_FRAMED_SIZE:
            raise SerializationError('Source too large for non-framed message')

        ciphertext = self.encryptor.update(plaintext)
        self._bytes_encrypted += len(plaintext)
        if self.signer is not None:
            self.signer.update(ciphertext)

        if len(plaintext) < b:
            _LOGGER.debug(
                'Closing encryptor after receiving only %s bytes of %s bytes requested',
                plaintext, b)
            self.source_stream.close()
            closing = self.encryptor.finalize()
            if self.signer is not None:
                self.signer.update(closing)
            closing += aws_encryption_sdk.internal.formatting.serialize.serialize_non_framed_close(
                tag=self.encryptor.tag, signer=self.signer)
            if self.signer is not None:
                closing += aws_encryption_sdk.internal.formatting.serialize.serialize_footer(
                    self.signer)
            return ciphertext + closing
        return ciphertext

    def _read_bytes_to_framed_body(self, b):
        """Reads the requested number of bytes from source to a streaming framed message body.

        :param int b: Number of bytes to read
        :returns: Bytes read from source stream, encrypted, and serialized
        :rtype: bytes
        """
        _LOGGER.debug('collecting %s bytes', b)
        _b = b
        b = int(
            math.ceil(b / float(self.config.frame_length)) *
            self.config.frame_length)
        _LOGGER.debug(
            '%s bytes requested; reading %s bytes after normalizing to frame length',
            _b, b)
        plaintext = self.source_stream.read(b)
        _LOGGER.debug('%s bytes read from source', len(plaintext))
        finalize = False
        if len(plaintext) < b:
            _LOGGER.debug('Final plaintext read from source')
            finalize = True
        output = b''
        final_frame_written = False

        while (
                # If not finalizing on this pass, exit when plaintext is exhausted
            (not finalize and plaintext)
                # If finalizing on this pass, wait until final frame is written
                or (finalize and not final_frame_written)):
            is_final_frame = finalize and len(
                plaintext) < self.config.frame_length
            bytes_in_frame = min(len(plaintext), self.config.frame_length)
            _LOGGER.debug('Writing %s bytes into%s frame %s', bytes_in_frame,
                          ' final' if is_final_frame else '',
                          self.sequence_number)
            self._bytes_encrypted += bytes_in_frame
            ciphertext, plaintext = aws_encryption_sdk.internal.formatting.serialize.serialize_frame(
                algorithm=self._encryption_materials.algorithm,
                plaintext=plaintext,
                message_id=self._header.message_id,
                data_encryption_key=self._derived_data_key,
                frame_length=self.config.frame_length,
                sequence_number=self.sequence_number,
                is_final_frame=is_final_frame,
                signer=self.signer)
            final_frame_written = is_final_frame
            output += ciphertext
            self.sequence_number += 1

        if finalize:
            _LOGGER.debug('Writing footer')
            if self.signer is not None:
                output += aws_encryption_sdk.internal.formatting.serialize.serialize_footer(
                    self.signer)
            self.source_stream.close()
        return output

    def _read_bytes(self, b):
        """Reads the requested number of bytes from a streaming message body.

        :param int b: Number of bytes to read
        :raises NotSupportedError: if content type is not supported
        """
        _LOGGER.debug('%s bytes requested from stream with content type: %s',
                      b, self.content_type)
        if b <= len(self.output_buffer) or self.source_stream.closed:
            _LOGGER.debug(
                'No need to read from source stream or source stream closed')
            return

        if self.content_type == ContentType.FRAMED_DATA:
            _LOGGER.debug('Reading to framed body')
            self.output_buffer += self._read_bytes_to_framed_body(b)
        elif self.content_type == ContentType.NO_FRAMING:
            _LOGGER.debug('Reading to non-framed body')
            self.output_buffer += self._read_bytes_to_non_framed_body(b)
        else:
            raise NotSupportedError('Unsupported content type')

        # To maintain backwards compatibility, only enforce this if a CMM is provided by the caller.
        if self.config.key_provider is None and self.config.source_length is not None:
            # Enforce that if the caller provided a source length value, the total bytes encrypted
            # must not exceed that value.
            if self._bytes_encrypted > self.config.source_length:
                raise CustomMaximumValueExceeded(
                    'Bytes encrypted has exceeded stated source length estimate:\n{actual} > {estimated}'
                    .format(actual=self._bytes_encrypted,
                            estimated=self.config.source))

    def close(self):
        """Closes out the stream."""
        _LOGGER.debug('Closing stream')
        super(StreamEncryptor, self).close()
예제 #3
0
class StreamEncryptor(_EncryptionStream):  # pylint: disable=too-many-instance-attributes
    """Provides a streaming encryptor for encrypting a stream source.
    Behaves as a standard file-like object.

    .. note::
        Take care when encrypting framed messages with large frame length and large non-framed
        messages.  See :class:`aws_encryption_sdk.stream` for more details.

    .. note::
        If config is provided, all other parameters are ignored.

    :param config: Client configuration object (config or individual parameters required)
    :type config: aws_encryption_sdk.streaming_client.EncryptorConfig
    :param source: Source data to encrypt or decrypt
    :type source: str, bytes, io.IOBase, or file
    :param materials_manager: `CryptoMaterialsManager` from which to obtain cryptographic materials
        (either `materials_manager` or `key_provider` required)
    :type materials_manager: aws_encryption_sdk.materials_manager.base.CryptoMaterialsManager
    :param key_provider: `MasterKeyProvider` from which to obtain data keys for encryption
        (either `materials_manager` or `key_provider` required)
    :type key_provider: aws_encryption_sdk.key_providers.base.MasterKeyProvider
    :param int source_length: Length of source data (optional)

        .. note::
            If source_length is not provided and unframed message is being written or read() is called,
            will attempt to seek() to the end of the stream and tell() to find the length of source data.

        .. note::
            .. versionadded:: 1.3.0

            If `source_length` and `materials_manager` are both provided, the total plaintext bytes
            encrypted will not be allowed to exceed `source_length`. To maintain backwards compatibility,
            this is not enforced if a `key_provider` is provided.

    :param dict encryption_context: Dictionary defining encryption context
    :param algorithm: Algorithm to use for encryption
    :type algorithm: aws_encryption_sdk.identifiers.Algorithm
    :param int frame_length: Frame length in bytes
    """

    _config_class = EncryptorConfig

    def __init__(self, **kwargs):  # pylint: disable=unused-argument,super-init-not-called
        """Prepares necessary initial values."""
        self.sequence_number = 1

        self.content_type = aws_encryption_sdk.internal.utils.content_type(
            self.config.frame_length)
        self._bytes_encrypted = 0

        if self.config.frame_length == 0 and (
                self.config.source_length is not None
                and self.config.source_length > MAX_NON_FRAMED_SIZE):
            raise SerializationError("Source too large for non-framed message")

        self.__unframed_plaintext_cache = io.BytesIO()
        self.__message_complete = False

    def ciphertext_length(self):
        """Returns the length of the resulting ciphertext message in bytes.

        :rtype: int
        """
        return aws_encryption_sdk.internal.formatting.ciphertext_length(
            header=self.header, plaintext_length=self.stream_length)

    def _prep_message(self):
        """Performs initial message setup.

        :raises MasterKeyProviderError: if primary master key is not a member of supplied MasterKeyProvider
        :raises MasterKeyProviderError: if no Master Keys are returned from key_provider
        """
        validate_commitment_policy_on_encrypt(self.config.commitment_policy,
                                              self.config.algorithm)

        try:
            plaintext_length = self.stream_length
        except NotSupportedError:
            plaintext_length = None
        encryption_materials_request = EncryptionMaterialsRequest(
            algorithm=self.config.algorithm,
            encryption_context=self.config.encryption_context.copy(),
            frame_length=self.config.frame_length,
            plaintext_rostream=aws_encryption_sdk.internal.utils.streams.
            ROStream(self.source_stream),
            plaintext_length=plaintext_length,
            commitment_policy=self.config.commitment_policy,
        )
        self._encryption_materials = self.config.materials_manager.get_encryption_materials(
            request=encryption_materials_request)

        if self.config.algorithm is not None and self._encryption_materials.algorithm != self.config.algorithm:
            raise ActionNotAllowedError(
                ("Cryptographic materials manager provided algorithm suite"
                 " differs from algorithm suite in request.\n"
                 "Required: {requested}\n"
                 "Provided: {provided}").format(
                     requested=self.config.algorithm,
                     provided=self._encryption_materials.algorithm))

        if self._encryption_materials.signing_key is None:
            self.signer = None
        else:
            self.signer = Signer.from_key_bytes(
                algorithm=self._encryption_materials.algorithm,
                key_bytes=self._encryption_materials.signing_key)
        aws_encryption_sdk.internal.utils.validate_frame_length(
            frame_length=self.config.frame_length,
            algorithm=self._encryption_materials.algorithm)

        message_id = aws_encryption_sdk.internal.utils.message_id(
            self._encryption_materials.algorithm.message_id_length())

        self._derived_data_key = derive_data_encryption_key(
            source_key=self._encryption_materials.data_encryption_key.data_key,
            algorithm=self._encryption_materials.algorithm,
            message_id=message_id,
        )

        self._header = self.generate_header(message_id)

        self._write_header()
        if self.content_type == ContentType.NO_FRAMING:
            self._prep_non_framed()
        self._message_prepped = True

    def generate_header(self, message_id):
        """Generates the header object.

        :param message_id: The randomly generated id for the message
        :type message_id: bytes
        """
        version = VERSION
        if self._encryption_materials.algorithm.message_format_version == 0x02:
            version = SerializationVersion.V2

        kwargs = dict(
            version=version,
            algorithm=self._encryption_materials.algorithm,
            message_id=message_id,
            encryption_context=self._encryption_materials.encryption_context,
            encrypted_data_keys=self._encryption_materials.encrypted_data_keys,
            content_type=self.content_type,
            frame_length=self.config.frame_length,
        )

        if self._encryption_materials.algorithm.is_committing():
            commitment_key = calculate_commitment_key(
                source_key=self._encryption_materials.data_encryption_key.
                data_key,
                algorithm=self._encryption_materials.algorithm,
                message_id=message_id,
            )
            kwargs["commitment_key"] = commitment_key

        if version == SerializationVersion.V1:
            kwargs["type"] = TYPE
            kwargs["content_aad_length"] = 0
            kwargs[
                "header_iv_length"] = self._encryption_materials.algorithm.iv_len

        return MessageHeader(**kwargs)

    def _write_header(self):
        """Builds the message header and writes it to the output stream."""
        self.output_buffer += serialize_header(header=self._header,
                                               signer=self.signer)
        self.output_buffer += serialize_header_auth(
            version=self._header.version,
            algorithm=self._encryption_materials.algorithm,
            header=self.output_buffer,
            data_encryption_key=self._derived_data_key,
            signer=self.signer,
        )

    def _prep_non_framed(self):
        """Prepare the opening data for a non-framed message."""
        try:
            plaintext_length = self.stream_length
            self.__unframed_plaintext_cache = self.source_stream
        except NotSupportedError:
            # We need to know the plaintext length before we can start processing the data.
            # If we cannot seek on the source then we need to read the entire source into memory.
            self.__unframed_plaintext_cache = io.BytesIO()
            self.__unframed_plaintext_cache.write(self.source_stream.read())
            plaintext_length = self.__unframed_plaintext_cache.tell()
            self.__unframed_plaintext_cache.seek(0)

        aad_content_string = aws_encryption_sdk.internal.utils.get_aad_content_string(
            content_type=self.content_type, is_final_frame=True)
        associated_data = assemble_content_aad(
            message_id=self._header.message_id,
            aad_content_string=aad_content_string,
            seq_num=1,
            length=plaintext_length,
        )
        self.encryptor = Encryptor(
            algorithm=self._encryption_materials.algorithm,
            key=self._derived_data_key,
            associated_data=associated_data,
            iv=non_framed_body_iv(self._encryption_materials.algorithm),
        )
        self.output_buffer += serialize_non_framed_open(
            algorithm=self._encryption_materials.algorithm,
            iv=self.encryptor.iv,
            plaintext_length=plaintext_length,
            signer=self.signer,
        )

    def _read_bytes_to_non_framed_body(self, b):
        """Reads the requested number of bytes from source to a streaming non-framed message body.

        :param int b: Number of bytes to read
        :returns: Encrypted bytes from source stream
        :rtype: bytes
        """
        _LOGGER.debug("Reading %d bytes", b)
        plaintext = self.__unframed_plaintext_cache.read(b)
        plaintext_length = len(plaintext)
        if self.tell() + len(plaintext) > MAX_NON_FRAMED_SIZE:
            raise SerializationError("Source too large for non-framed message")

        ciphertext = self.encryptor.update(plaintext)
        self._bytes_encrypted += plaintext_length
        if self.signer is not None:
            self.signer.update(ciphertext)

        if len(plaintext) < b:
            _LOGGER.debug(
                "Closing encryptor after receiving only %d bytes of %d bytes requested",
                plaintext_length, b)

            closing = self.encryptor.finalize()

            if self.signer is not None:
                self.signer.update(closing)

            closing += serialize_non_framed_close(tag=self.encryptor.tag,
                                                  signer=self.signer)

            if self.signer is not None:
                closing += serialize_footer(self.signer)
            self.__message_complete = True
            return ciphertext + closing

        return ciphertext

    def _read_bytes_to_framed_body(self, b):
        """Reads the requested number of bytes from source to a streaming framed message body.

        :param int b: Number of bytes to read
        :returns: Bytes read from source stream, encrypted, and serialized
        :rtype: bytes
        """
        _LOGGER.debug("collecting %d bytes", b)
        _b = b

        if b > 0:
            _frames_to_read = math.ceil(b / float(self.config.frame_length))
            b = int(_frames_to_read * self.config.frame_length)
        _LOGGER.debug(
            "%d bytes requested; reading %d bytes after normalizing to frame length",
            _b, b)

        plaintext = self.source_stream.read(b)
        plaintext_length = len(plaintext)
        _LOGGER.debug("%d bytes read from source", plaintext_length)

        finalize = False

        if b < 0 or plaintext_length < b:
            _LOGGER.debug("Final plaintext read from source")
            finalize = True

        output = b""
        final_frame_written = False

        while (
                # If not finalizing on this pass, exit when plaintext is exhausted
            (not finalize and plaintext)
                # If finalizing on this pass, wait until final frame is written
                or (finalize and not final_frame_written)):
            current_plaintext_length = len(plaintext)
            is_final_frame = finalize and current_plaintext_length < self.config.frame_length
            bytes_in_frame = min(current_plaintext_length,
                                 self.config.frame_length)
            _LOGGER.debug(
                "Writing %d bytes into%s frame %d",
                bytes_in_frame,
                " final" if is_final_frame else "",
                self.sequence_number,
            )
            self._bytes_encrypted += bytes_in_frame
            ciphertext, plaintext = serialize_frame(
                algorithm=self._encryption_materials.algorithm,
                plaintext=plaintext,
                message_id=self._header.message_id,
                data_encryption_key=self._derived_data_key,
                frame_length=self.config.frame_length,
                sequence_number=self.sequence_number,
                is_final_frame=is_final_frame,
                signer=self.signer,
            )
            final_frame_written = is_final_frame
            output += ciphertext
            self.sequence_number += 1

        if finalize:
            _LOGGER.debug("Writing footer")
            if self.signer is not None:
                output += serialize_footer(self.signer)
            self.__message_complete = True
        return output

    def _read_bytes(self, b):
        """Reads the requested number of bytes from a streaming message body.

        :param int b: Number of bytes to read
        :raises NotSupportedError: if content type is not supported
        """
        _LOGGER.debug("%d bytes requested from stream with content type: %s",
                      b, self.content_type)
        if 0 <= b <= len(self.output_buffer) or self.__message_complete:
            _LOGGER.debug(
                "No need to read from source stream or source stream closed")
            return

        if self.content_type == ContentType.FRAMED_DATA:
            _LOGGER.debug("Reading to framed body")
            self.output_buffer += self._read_bytes_to_framed_body(b)
        elif self.content_type == ContentType.NO_FRAMING:
            _LOGGER.debug("Reading to non-framed body")
            self.output_buffer += self._read_bytes_to_non_framed_body(b)
        else:
            raise NotSupportedError("Unsupported content type")

        # To maintain backwards compatibility, only enforce this if a CMM is provided by the caller.
        if self.config.key_provider is None and self.config.source_length is not None:
            # Enforce that if the caller provided a source length value, the total bytes encrypted
            # must not exceed that value.
            if self._bytes_encrypted > self.config.source_length:
                raise CustomMaximumValueExceeded(
                    "Bytes encrypted has exceeded stated source length estimate:\n{actual:d} > {estimated:d}"
                    .format(actual=self._bytes_encrypted,
                            estimated=self.config.source_length))

    def close(self):
        """Closes out the stream."""
        _LOGGER.debug("Closing stream")
        super(StreamEncryptor, self).close()