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 = sentinel.signer
test_encryptor._encryption_materials = self.mock_encryption_materials
test_encryptor._header = MagicMock()
test_encryptor._derived_data_key = sentinel.derived_data_key
test = test_encryptor._read_bytes_to_framed_body(128)
self.mock_serialize_frame.assert_called_once_with(
algorithm=self.mock_encryption_materials.algorithm,
plaintext=self.plaintext[:128],
message_id=test_encryptor._header.message_id,
data_encryption_key=sentinel.derived_data_key,
frame_length=test_encryptor.config.frame_length,
sequence_number=1,
is_final_frame=False,
signer=sentinel.signer
)
assert not self.mock_serialize_footer.called
assert not test_encryptor.source_stream.closed
assert test == b'1234'
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.signer = sentinel.signer
test_encryptor.content_type = sentinel.content_type
test_encryptor._header = sentinel.header
test_encryptor.output_buffer = b''
test_encryptor._encryption_materials = self.mock_encryption_materials
test_encryptor._derived_data_key = sentinel.derived_data_key
test_encryptor._write_header()
self.mock_serialize_header.assert_called_once_with(
header=test_encryptor._header,
signer=sentinel.signer
)
self.mock_serialize_header_auth.assert_called_once_with(
algorithm=self.mock_encryption_materials.algorithm,
header=b'12345',
data_encryption_key=sentinel.derived_data_key,
signer=sentinel.signer
)
assert test_encryptor.output_buffer == b'1234567890'
def test_prep_non_framed(self, mock_non_framed_iv):
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 = sentinel.signer
test_encryptor._encryption_materials = self.mock_encryption_materials
test_encryptor._header = MagicMock()
test_encryptor._derived_data_key = sentinel.derived_data_key
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=self.mock_encryption_materials.algorithm,
key=sentinel.derived_data_key,
associated_data=sentinel.associated_data,
iv=mock_non_framed_iv.return_value)
self.mock_serialize_non_framed_open.assert_called_once_with(
algorithm=self.mock_encryption_materials.algorithm,
iv=sentinel.iv,
plaintext_length=test_encryptor.stream_length,
signer=sentinel.signer)
assert test_encryptor.output_buffer == b'1234567890'
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 = sentinel.signer
test_encryptor._encryption_materials = self.mock_encryption_materials
test_encryptor._header = MagicMock()
test_encryptor._derived_data_key = sentinel.derived_data_key
test = test_encryptor._read_bytes_to_framed_body(51)
self.mock_serialize_frame.assert_has_calls(
calls=(call(algorithm=self.mock_encryption_materials.algorithm,
plaintext=self.plaintext[:50],
message_id=test_encryptor._header.message_id,
data_encryption_key=sentinel.derived_data_key,
frame_length=test_encryptor.config.frame_length,
sequence_number=1,
is_final_frame=False,
signer=sentinel.signer),
call(algorithm=self.mock_encryption_materials.algorithm,
plaintext=b'',
message_id=test_encryptor._header.message_id,
data_encryption_key=sentinel.derived_data_key,
frame_length=test_encryptor.config.frame_length,
sequence_number=2,
is_final_frame=True,
signer=sentinel.signer)),
any_order=False)
assert test == b'FIRSTFINALFOOTER'
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,
materials_manager=self.mock_materials_manager,
algorithm=aws_encryption_sdk.internal.defaults.ALGORITHM,
frame_length=self.mock_frame_length,
commitment_policy=self.mock_commitment_policy,
)
test_encryptor.signer = sentinel.signer
test_encryptor.content_type = sentinel.content_type
test_encryptor._header = sentinel.header
sentinel.header.version = SerializationVersion.V1
test_encryptor.output_buffer = b""
test_encryptor._encryption_materials = self.mock_encryption_materials
test_encryptor._derived_data_key = sentinel.derived_data_key
test_encryptor._write_header()
self.mock_serialize_header.assert_called_once_with(
header=test_encryptor._header, signer=sentinel.signer)
self.mock_serialize_header_auth.assert_called_once_with(
version=sentinel.header.version,
algorithm=self.mock_encryption_materials.algorithm,
header=b"12345",
data_encryption_key=sentinel.derived_data_key,
signer=sentinel.signer,
)
assert test_encryptor.output_buffer == b"1234567890"
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._encryption_materials = self.mock_encryption_materials
test_encryptor._header = MagicMock()
test_encryptor.content_type = None
with pytest.raises(NotSupportedError) as excinfo:
test_encryptor._read_bytes(5)
excinfo.match("Unsupported content type")
assert not mock_read_non_framed.called
assert not mock_read_framed.called
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._encryption_materials = self.mock_encryption_materials
test_encryptor._header = 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(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 = sentinel.signer
test_encryptor._encryption_materials = self.mock_encryption_materials
test_encryptor._header = MagicMock()
test_encryptor._derived_data_key = sentinel.derived_data_key
test_encryptor._read_bytes_to_framed_body(len(self.plaintext) + 1)
self.mock_serialize_footer.assert_called_once_with(sentinel.signer)
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._header = MagicMock()
test_encryptor.signer = None
test_encryptor._encryption_materials = self.mock_encryption_materials
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,
materials_manager=self.mock_materials_manager,
frame_length=50,
commitment_policy=self.mock_commitment_policy,
signature_policy=self.mock_signature_policy,
)
test_encryptor.signer = sentinel.signer
test_encryptor._encryption_materials = self.mock_encryption_materials
test_encryptor._header = MagicMock()
test_encryptor._derived_data_key = sentinel.derived_data_key
test = test_encryptor._read_bytes_to_framed_body(51)
self.mock_serialize_frame.assert_has_calls(
calls=(
call(
algorithm=self.mock_encryption_materials.algorithm,
plaintext=self.plaintext[:50],
message_id=test_encryptor._header.message_id,
data_encryption_key=sentinel.derived_data_key,
frame_length=test_encryptor.config.frame_length,
sequence_number=1,
is_final_frame=False,
signer=sentinel.signer,
),
call(
algorithm=self.mock_encryption_materials.algorithm,
plaintext=b"",
message_id=test_encryptor._header.message_id,
data_encryption_key=sentinel.derived_data_key,
frame_length=test_encryptor.config.frame_length,
sequence_number=2,
is_final_frame=True,
signer=sentinel.signer,
),
),
any_order=False,
)
assert test == b"FIRSTFINALFOOTER"
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,
materials_manager=self.mock_materials_manager,
frame_length=len(self.plaintext),
algorithm=Algorithm.AES_128_GCM_IV12_TAG16,
commitment_policy=self.mock_commitment_policy,
)
test_encryptor.signer = None
test_encryptor._encryption_materials = self.mock_encryption_materials
test_encryptor._header = MagicMock()
test_encryptor._derived_data_key = sentinel.derived_data_key
test_encryptor._read_bytes_to_framed_body(len(self.plaintext) + 1)
assert not self.mock_serialize_footer.called
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 = sentinel.signer
test_encryptor._encryption_materials = self.mock_encryption_materials
test_encryptor._header = MagicMock()
test_encryptor._derived_data_key = sentinel.derived_data_key
test = test_encryptor._read_bytes_to_framed_body(len(self.plaintext) + 1)
self.mock_serialize_frame.assert_has_calls(
calls=[
call(
algorithm=self.mock_encryption_materials.algorithm,
plaintext=self.plaintext,
message_id=test_encryptor._header.message_id,
data_encryption_key=sentinel.derived_data_key,
frame_length=test_encryptor.config.frame_length,
sequence_number=1,
is_final_frame=False,
signer=sentinel.signer
),
call(
algorithm=self.mock_encryption_materials.algorithm,
plaintext=self.plaintext[frame_length:],
message_id=test_encryptor._header.message_id,
data_encryption_key=sentinel.derived_data_key,
frame_length=test_encryptor.config.frame_length,
sequence_number=2,
is_final_frame=False,
signer=sentinel.signer
),
call(
algorithm=self.mock_encryption_materials.algorithm,
plaintext=self.plaintext[frame_length * 2:],
message_id=test_encryptor._header.message_id,
data_encryption_key=sentinel.derived_data_key,
frame_length=test_encryptor.config.frame_length,
sequence_number=3,
is_final_frame=False,
signer=sentinel.signer
),
call(
algorithm=self.mock_encryption_materials.algorithm,
plaintext=self.plaintext[frame_length * 3:],
message_id=test_encryptor._header.message_id,
data_encryption_key=sentinel.derived_data_key,
frame_length=test_encryptor.config.frame_length,
sequence_number=4,
is_final_frame=False,
signer=sentinel.signer
),
call(
algorithm=self.mock_encryption_materials.algorithm,
plaintext=b'',
message_id=test_encryptor._header.message_id,
data_encryption_key=sentinel.derived_data_key,
frame_length=test_encryptor.config.frame_length,
sequence_number=5,
is_final_frame=True,
signer=sentinel.signer
)
],
any_order=False
)
self.mock_serialize_footer.assert_called_once_with(sentinel.signer)
assert test_encryptor.source_stream.closed
assert test == b'1234567890-=FINAL/*-'
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,
materials_manager=self.mock_materials_manager,
frame_length=frame_length,
commitment_policy=self.mock_commitment_policy,
)
test_encryptor.signer = sentinel.signer
test_encryptor._encryption_materials = self.mock_encryption_materials
test_encryptor._header = MagicMock()
test_encryptor._derived_data_key = sentinel.derived_data_key
test = test_encryptor._read_bytes_to_framed_body(
len(self.plaintext) + 1)
self.mock_serialize_frame.assert_has_calls(
calls=[
call(
algorithm=self.mock_encryption_materials.algorithm,
plaintext=self.plaintext,
message_id=test_encryptor._header.message_id,
data_encryption_key=sentinel.derived_data_key,
frame_length=test_encryptor.config.frame_length,
sequence_number=1,
is_final_frame=False,
signer=sentinel.signer,
),
call(
algorithm=self.mock_encryption_materials.algorithm,
plaintext=self.plaintext[frame_length:],
message_id=test_encryptor._header.message_id,
data_encryption_key=sentinel.derived_data_key,
frame_length=test_encryptor.config.frame_length,
sequence_number=2,
is_final_frame=False,
signer=sentinel.signer,
),
call(
algorithm=self.mock_encryption_materials.algorithm,
plaintext=self.plaintext[frame_length * 2:],
message_id=test_encryptor._header.message_id,
data_encryption_key=sentinel.derived_data_key,
frame_length=test_encryptor.config.frame_length,
sequence_number=3,
is_final_frame=False,
signer=sentinel.signer,
),
call(
algorithm=self.mock_encryption_materials.algorithm,
plaintext=self.plaintext[frame_length * 3:],
message_id=test_encryptor._header.message_id,
data_encryption_key=sentinel.derived_data_key,
frame_length=test_encryptor.config.frame_length,
sequence_number=4,
is_final_frame=False,
signer=sentinel.signer,
),
call(
algorithm=self.mock_encryption_materials.algorithm,
plaintext=b"",
message_id=test_encryptor._header.message_id,
data_encryption_key=sentinel.derived_data_key,
frame_length=test_encryptor.config.frame_length,
sequence_number=5,
is_final_frame=True,
signer=sentinel.signer,
),
],
any_order=False,
)
self.mock_serialize_footer.assert_called_once_with(sentinel.signer)
assert test == b"1234567890-=FINAL/*-"
