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/*-'
