def test_read_bytes_to_framed_body_single_frame_read(self): self.mock_serialize_frame.return_value = (b'1234', b'') pt_stream = io.BytesIO(self.plaintext * 2) test_encryptor = StreamEncryptor( source=pt_stream, key_provider=self.mock_key_provider, frame_length=128 ) test_encryptor.signer = MagicMock() test_encryptor._header = MagicMock() test_encryptor.encryption_data_key = MagicMock() test = test_encryptor._read_bytes_to_framed_body(128) self.mock_serialize_frame.assert_called_once_with( algorithm=test_encryptor.config.algorithm, plaintext=self.plaintext[:128], message_id=test_encryptor._header.message_id, encryption_data_key=test_encryptor.encryption_data_key, frame_length=test_encryptor.config.frame_length, sequence_number=1, is_final_frame=False, signer=test_encryptor.signer ) assert not self.mock_serialize_footer.called assert not test_encryptor.source_stream.closed assert test == b'1234'
def test_prep_non_framed(self): self.mock_serialize_non_framed_open.return_value = b'1234567890' test_encryptor = StreamEncryptor( source=self.mock_input_stream, key_provider=self.mock_key_provider ) test_encryptor.signer = MagicMock() test_encryptor._header = MagicMock() test_encryptor.encryption_data_key = MagicMock() test_encryptor._prep_non_framed() self.mock_get_aad_content_string.assert_called_once_with( content_type=test_encryptor.content_type, is_final_frame=True ) self.mock_assemble_content_aad.assert_called_once_with( message_id=test_encryptor._header.message_id, aad_content_string=sentinel.aad_content_string, seq_num=1, length=test_encryptor.stream_length ) self.mock_encryptor.assert_called_once_with( algorithm=test_encryptor.config.algorithm, key=test_encryptor.encryption_data_key.data_key, associated_data=sentinel.associated_data, message_id=test_encryptor._header.message_id ) self.mock_serialize_non_framed_open.assert_called_once_with( algorithm=test_encryptor.config.algorithm, iv=sentinel.iv, plaintext_length=test_encryptor.stream_length, signer=test_encryptor.signer ) assert test_encryptor.output_buffer == b'1234567890'
def test_write_header(self): self.mock_serialize_header.return_value = b'12345' self.mock_serialize_header_auth.return_value = b'67890' pt_stream = io.BytesIO(self.plaintext) test_encryptor = StreamEncryptor( source=pt_stream, key_provider=self.mock_key_provider, algorithm=aws_encryption_sdk.internal.defaults.ALGORITHM, frame_length=self.mock_frame_length ) test_encryptor.content_type = sentinel.content_type test_encryptor._header = MagicMock() test_encryptor._header.message_id = VALUES['message_id'] test_encryptor.signer = sentinel.signer test_encryptor.output_buffer = b'' test_encryptor.encryption_data_key = self.mock_encryption_data_key test_encryptor._write_header() self.mock_serialize_header.assert_called_once_with( header=test_encryptor._header, signer=test_encryptor.signer ) self.mock_serialize_header_auth.assert_called_once_with( algorithm=test_encryptor.config.algorithm, header=b'12345', message_id=test_encryptor._header.message_id, encryption_data_key=test_encryptor.encryption_data_key, signer=test_encryptor.signer ) assert test_encryptor.output_buffer == b'1234567890'
def test_close(self, mock_close): self.mock_encryption_data_key.key_provider = VALUES['key_provider'] self.mock_encryption_data_key.encrypted_data_key = VALUES['encrypted_data_key'] pt_stream = io.BytesIO(self.plaintext) test_encryptor = StreamEncryptor( source=pt_stream, key_provider=self.mock_key_provider ) test_encryptor.encryption_data_key = self.mock_encryption_data_key test_encryptor.close() mock_close.assert_called_once_with()
def test_read_bytes_unsupported_type(self, mock_read_non_framed, mock_read_framed): pt_stream = io.BytesIO(self.plaintext) test_encryptor = StreamEncryptor( source=pt_stream, key_provider=self.mock_key_provider ) test_encryptor.signer = MagicMock() test_encryptor._header = MagicMock() test_encryptor.encryption_data_key = MagicMock() test_encryptor.content_type = None with six.assertRaisesRegex(self, NotSupportedError, 'Unsupported content type'): test_encryptor._read_bytes(5) assert not mock_read_non_framed.called assert not mock_read_framed.called
def test_read_bytes_to_framed_body_close_no_signer(self): self.mock_serialize_frame.return_value = (b'1234', b'') pt_stream = io.BytesIO(self.plaintext) test_encryptor = StreamEncryptor( source=pt_stream, key_provider=self.mock_key_provider, frame_length=len(self.plaintext), algorithm=Algorithm.AES_128_GCM_IV12_TAG16 ) test_encryptor.signer = None test_encryptor._header = MagicMock() test_encryptor.encryption_data_key = MagicMock() test_encryptor._read_bytes_to_framed_body(len(self.plaintext) + 1) assert not self.mock_serialize_footer.called assert test_encryptor.source_stream.closed
def test_read_bytes_to_framed_body_close(self): self.mock_serialize_frame.return_value = (b'1234', b'') self.mock_serialize_footer.return_value = b'5678' pt_stream = io.BytesIO(self.plaintext) test_encryptor = StreamEncryptor( source=pt_stream, key_provider=self.mock_key_provider, frame_length=len(self.plaintext) ) test_encryptor.signer = MagicMock() test_encryptor._header = MagicMock() test_encryptor.encryption_data_key = MagicMock() test_encryptor._read_bytes_to_framed_body(len(self.plaintext) + 1) self.mock_serialize_footer.assert_called_once_with(test_encryptor.signer) assert test_encryptor.source_stream.closed
def test_read_bytes_to_non_framed_body_no_signer(self): pt_stream = io.BytesIO(self.plaintext) test_encryptor = StreamEncryptor( source=pt_stream, key_provider=self.mock_key_provider, algorithm=Algorithm.AES_128_GCM_IV12_TAG16 ) test_encryptor.signer = None test_encryptor._header = MagicMock() test_encryptor.encryption_data_key = MagicMock() test_encryptor.encryptor = MagicMock() test_encryptor.encryptor.update.return_value = b'123' test_encryptor.encryptor.finalize.return_value = b'456' test_encryptor.encryptor.tag = sentinel.tag self.mock_serialize_non_framed_close.return_value = b'789' self.mock_serialize_footer.return_value = b'0-=' test_encryptor._read_bytes_to_non_framed_body(len(self.plaintext) + 1) assert not self.mock_serialize_footer.called
def test_read_bytes_to_framed_body_single_frame_with_final(self): self.mock_serialize_frame.side_effect = ( (b'FIRST', b''), (b'FINAL', b'') ) self.mock_serialize_footer.return_value = b'FOOTER' pt_stream = io.BytesIO(self.plaintext[:50]) test_encryptor = StreamEncryptor( source=pt_stream, key_provider=self.mock_key_provider, frame_length=50 ) test_encryptor.signer = MagicMock() test_encryptor._header = MagicMock() test_encryptor.encryption_data_key = MagicMock() test = test_encryptor._read_bytes_to_framed_body(51) self.mock_serialize_frame.assert_has_calls( calls=( call( algorithm=test_encryptor.config.algorithm, plaintext=self.plaintext[:50], message_id=test_encryptor._header.message_id, encryption_data_key=test_encryptor.encryption_data_key, frame_length=test_encryptor.config.frame_length, sequence_number=1, is_final_frame=False, signer=test_encryptor.signer ), call( algorithm=test_encryptor.config.algorithm, plaintext=b'', message_id=test_encryptor._header.message_id, encryption_data_key=test_encryptor.encryption_data_key, frame_length=test_encryptor.config.frame_length, sequence_number=2, is_final_frame=True, signer=test_encryptor.signer ) ), any_order=False ) assert test == b'FIRSTFINALFOOTER'
def test_read_bytes_to_framed_body_multi_frame_read(self): frame_length = int(len(self.plaintext) / 4) self.mock_serialize_frame.side_effect = ( (b'123', self.plaintext[frame_length:]), (b'456', self.plaintext[frame_length * 2:]), (b'789', self.plaintext[frame_length * 3:]), (b'0-=', b''), (b'FINAL', b'') ) self.mock_serialize_footer.return_value = b'/*-' pt_stream = io.BytesIO(self.plaintext) test_encryptor = StreamEncryptor( source=pt_stream, key_provider=self.mock_key_provider, frame_length=frame_length ) test_encryptor.signer = MagicMock() test_encryptor._header = MagicMock() test_encryptor.encryption_data_key = MagicMock() test = test_encryptor._read_bytes_to_framed_body(len(self.plaintext) + 1) self.mock_serialize_frame.assert_has_calls( calls=[ call( algorithm=test_encryptor.config.algorithm, plaintext=self.plaintext, message_id=test_encryptor._header.message_id, encryption_data_key=test_encryptor.encryption_data_key, frame_length=test_encryptor.config.frame_length, sequence_number=1, is_final_frame=False, signer=test_encryptor.signer ), call( algorithm=test_encryptor.config.algorithm, plaintext=self.plaintext[frame_length:], message_id=test_encryptor._header.message_id, encryption_data_key=test_encryptor.encryption_data_key, frame_length=test_encryptor.config.frame_length, sequence_number=2, is_final_frame=False, signer=test_encryptor.signer ), call( algorithm=test_encryptor.config.algorithm, plaintext=self.plaintext[frame_length * 2:], message_id=test_encryptor._header.message_id, encryption_data_key=test_encryptor.encryption_data_key, frame_length=test_encryptor.config.frame_length, sequence_number=3, is_final_frame=False, signer=test_encryptor.signer ), call( algorithm=test_encryptor.config.algorithm, plaintext=self.plaintext[frame_length * 3:], message_id=test_encryptor._header.message_id, encryption_data_key=test_encryptor.encryption_data_key, frame_length=test_encryptor.config.frame_length, sequence_number=4, is_final_frame=False, signer=test_encryptor.signer ), call( algorithm=test_encryptor.config.algorithm, plaintext=b'', message_id=test_encryptor._header.message_id, encryption_data_key=test_encryptor.encryption_data_key, frame_length=test_encryptor.config.frame_length, sequence_number=5, is_final_frame=True, signer=test_encryptor.signer ) ], any_order=False ) self.mock_serialize_footer.assert_called_once_with(test_encryptor.signer) assert test_encryptor.source_stream.closed assert test == b'1234567890-=FINAL/*-'