def test_get_key_id_no_prefix(self):
# Test without the 'alias/' prefix
kms_client = KMSClient(alias=self.kms_alias)
kms_client._client = mock.Mock()
kms_client._client.list_aliases = mock.MagicMock(
return_value=self.kms_key_list)
self.assertEqual(kms_client._get_key_id(), 'id-1')
def test_get_key_alias(self):
kms_client = KMSClient(alias=None)
kms_client._client = mock.Mock()
kms_client._client.list_aliases = mock.MagicMock(
return_value=self.kms_key_list)
self.assertIsNone(kms_client.get_key_alias('unknown-id'))
self.assertEqual(kms_client.get_key_alias('id-2'),
self.kms_key_list['Aliases'][1]['AliasName'])
def test_decrypt(self):
kms_client = KMSClient(alias=self.kms_alias)
kms_client._client = mock.Mock()
kms_client._client.decrypt = mock.MagicMock(
return_value={'Plaintext': 'test'})
result = kms_client.decrypt(encrypted_content='a',
encryption_context={'t': 1})
self.assertEqual(result, 'test')
kms_client._client.decrypt.assert_called_once_with(
CiphertextBlob='a', EncryptionContext={'t': 1})
def test_generate_envelope_key(self):
with self.assertRaises(KMSAliasNotFoundError):
kms_client = KMSClient(alias='unknown-kms-alias')
kms_client._client.list_aliases = mock.MagicMock(
return_value=self.kms_key_list)
kms_client.generate_envelope_key()
kms_generated_envelope_key = {
'Plaintext': 'test_plaintext',
'CiphertextBlob': 'test_ciphertext_blob'
}
kms_client = KMSClient(alias=self.kms_alias)
kms_client._client = mock.Mock()
kms_client._client.generate_data_key = mock.MagicMock(
return_value=kms_generated_envelope_key.copy())
kms_client._client.list_aliases = mock.MagicMock(
return_value=self.kms_key_list)
kms_generated_envelope_key.update(
dict(EncryptionContext={'kms_cmk_id': 'id-1'}))
result = kms_client.generate_envelope_key()
self.assertDictEqual(result, kms_generated_envelope_key)
kms_client._client.generate_data_key.assert_called_once_with(
KeyId='id-1',
EncryptionContext={'kms_cmk_id': 'id-1'},
KeySpec="AES_256")
def test_get_key_id(self):
kms_client = KMSClient(alias=self.prefixed_kms_alias)
kms_client._client = mock.Mock()
kms_client._client.list_aliases = mock.MagicMock(
return_value=self.kms_key_list)
self.assertEqual(kms_client._get_key_id(), 'id-1')
self.kms_key_list['Aliases'] = self.kms_key_list['Aliases'][1:]
self.assertIsNone(kms_client._get_key_id())
# Raise Error if used without alias
kms_client = KMSClient(alias=None)
with self.assertRaises(KMSAliasNotFoundError):
kms_client.generate_envelope_key()
def test_init(self):
with mock.patch('boto3.client') as mock_client:
kms_client = KMSClient(alias=self.kms_alias)
self.assertIsNone(kms_client._region)
self.assertEqual(kms_client._alias, self.prefixed_kms_alias)
mock_client.assert_called_with('kms', region_name=None)
kms_client = KMSClient(alias=self.kms_alias, region='us-test-1')
self.assertEqual(kms_client._region, 'us-test-1')
self.assertEqual(kms_client._alias, self.prefixed_kms_alias)
mock_client.assert_called_with('kms', region_name='us-test-1')
# Test with non-string (None) KMS alias
kms_client = KMSClient(alias=None)
self.assertIsNone(kms_client._alias)
def __init__(self, kms_alias, region=None, display=None, stream=None):
stream = stream or sys.stderr
logger = display or module_logger
self._logger = logger
# If no display is provided, the module_logger is configured and used.
if isinstance(self._logger, logging.Logger):
self._configure_logger(stream=stream)
self._encryption_backend = S3EncryptionClient(
key_provider=KMSClient(alias=kms_alias, region=region)
)
def test_set_alias(self):
kms_client = KMSClient(alias='test')
self.assertEqual(kms_client.alias, 'alias/test')
self.assertFalse(kms_client.set_alias('something-else'))
self.assertEqual(kms_client.alias, 'alias/test')
kms_client = KMSClient(alias=None)
self.assertIsNone(kms_client.alias)
self.assertTrue(kms_client.set_alias('something-else'))
self.assertEqual(kms_client.alias, 'alias/something-else')
def setUp(self):
self.s3_bucket = 'myev-tests-123'
# Sample data encrypted and uploaded by the Ruby SDK
self.s3_metadata = {
'x-amz-cek-alg':
'AES/CBC/PKCS5Padding',
'x-amz-iv':
'wyqRlcTfa8qtvYc2AZM1vg==',
'x-amz-key-v2':
'AQEDAHgwrsK9FliCzxfzK2/qJ7ZEc5ZWM8q9WzJgY/ldZTG'
'jigAAAH4wfAYJKoZIhvcNAQcGoG8wbQIBADBoBgkqhkiG9w0BBwEwHgYJYIZIAW'
'UDBAEuMBEEDJ2oZHkbpznDeSlu3AIBEIA7C9G+Y5BtdSsd2hns8r3dvQKzZ0R/fV'
'd9gbiHZ7/aIoC0+Hwthq+/grWjWlQvOnfDk9F8sK7dO7rWXeM=',
'x-amz-matdesc':
'{"kms_cmk_id":"3d78336b-f5c1-4bde-ab0e-1a8a2d3f1c4a"}',
'x-amz-unencrypted-content-length':
'45',
'x-amz-wrap-alg':
'kms'
}
self.plaintext_data = '{"A": "1", "B": "abcdefghijklmnopqrstuvxyz"}\n'
self.base64_encrypted_data = \
'Vjg8GFTZds+gVhCTP9H+YgfiVifJDSCljtwde46LSLesamFf1Vfz/4v1YGITDlUM'
self.base64_plaintext_envelope_key = \
'tUQZbnbmv1woZhiW7UulhWe3xfNHERnxRoVcjHk5a7w='
self.kms_alias = 'alias/test-kms-alias'
self.kms_key_list = {
'Aliases': [
{
'AliasName': self.kms_alias,
'TargetKeyId': 'id-1'
},
{
'AliasName': 'alias/another-kms-alias',
'TargetKeyId': 'id-2'
},
]
}
self.kms_client = KMSClient(alias=self.kms_alias)
def test_get_object_update_alias_from_metadata(self):
s3_key = 'dummy'
kms_client = KMSClient(alias=None)
encryption_client = S3EncryptionClient(key_provider=kms_client)
original_set_alias = kms_client.set_alias
self.assertIsNone(encryption_client._key_provider.alias)
encryption_client._client = mock.MagicMock()
encryption_client._client.get_object = mock.MagicMock(
return_value={
'Metadata': self.s3_metadata,
'Body': StringIO(base64.b64decode(self.base64_encrypted_data))
})
encryption_client._key_provider = mock.MagicMock()
encryption_client._key_provider.decrypt = \
mock.MagicMock(
return_value=base64.b64decode(
self.base64_plaintext_envelope_key
)
)
mock_get_key_alias = encryption_client._key_provider.get_key_alias = \
mock.MagicMock(return_value='test')
encryption_client._key_provider.alias = None
mock_set_alias = encryption_client._key_provider.set_alias = \
mock.MagicMock(side_effect=original_set_alias)
encryption_client.get_object(bucket=self.s3_bucket, key=s3_key)
mock_get_key_alias.assert_called_once_with(
key_id=json.loads(self.s3_metadata['x-amz-matdesc'])['kms_cmk_id'])
mock_set_alias.assert_called_once_with('test')
self.assertEqual(kms_client._alias, 'alias/test')