Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

[Key Vault] Add local-only mode to CryptographyClient #16565

Merged
merged 18 commits into from
Mar 10, 2021
Merged

[Key Vault] Add local-only mode to CryptographyClient #16565

Show file tree
Hide file tree
Changes from 3 commits
Commits
Show all changes
18 commits
Select commit Hold shift + click to select a range
13dd7a2
First prototype of from_jwk
mccoyp Feb 5, 2021
41ef114
Add tests
mccoyp Feb 25, 2021
fad4bc0
Incorporate feedback
mccoyp Mar 2, 2021
9ecdf21
Log before raising
mccoyp Mar 4, 2021
b801f12
Accept JsonWebKey object
mccoyp Mar 4, 2021
ae2967b
Allow None key_id, try JWK's .kid
mccoyp Mar 4, 2021
0d2519f
Use JsonWebKey in tests for .kid
mccoyp Mar 5, 2021
218a312
Add .from_jwk example
mccoyp Mar 5, 2021
8a602c1
Address pylint
mccoyp Mar 5, 2021
8345eb6
Pass JWK to crypto client
mccoyp Mar 5, 2021
ec7afaf
Add JWK support for local providers
mccoyp Mar 5, 2021
a4dd8f8
Local crypto only needs JWK
mccoyp Mar 8, 2021
e6ae519
JsonWebKey imports only for type checking
mccoyp Mar 8, 2021
a2d8c87
Appease pylint
mccoyp Mar 8, 2021
fb36a25
Thanks, Charles!
mccoyp Mar 9, 2021
5ca931d
Address feedback
mccoyp Mar 9, 2021
08f5935
Satisfy frozen six requirement
mccoyp Mar 9, 2021
8afbfe2
Remove unused import
mccoyp Mar 10, 2021
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
49 changes: 48 additions & 1 deletion sdk/keyvault/azure-keyvault-keys/azure/keyvault/keys/crypto/_client.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -98,6 +98,7 @@ class CryptographyClient(KeyVaultClientBase):

def __init__(self, key, credential, **kwargs):
# type: (Union[KeyVaultKey, str], TokenCredential, **Any) -> None
self._local_only = kwargs.pop("_local_only", False)

if isinstance(key, KeyVaultKey):
self._key = key
Expand All @@ -109,7 +110,7 @@ def __init__(self, key, credential, **kwargs):
else:
raise ValueError("'key' must be a KeyVaultKey instance or a key ID string including a version")

if not self._key_id.version:
if not (self._key_id.version or self._local_only):
raise ValueError("'key' must include a version")

self._local_provider = NoLocalCryptography()
Expand All @@ -126,6 +127,16 @@ def key_id(self):
"""
return self._key_id.source_id

@classmethod
def from_jwk(cls, jwk):
# type: (dict) -> CryptographyClient
mccoyp marked this conversation as resolved.
Show resolved Hide resolved
"""Creates a client that can only perform cryptographic operations locally.

:param dict jwk: the key's cryptographic material, as a dictionary.
"""
key_id = "https://key-vault.vault.azure.net/keys/local-key"
mccoyp marked this conversation as resolved.
Show resolved Hide resolved
return cls(KeyVaultKey(key_id, jwk), object, _local_only=True)
mccoyp marked this conversation as resolved.
Show resolved Hide resolved

@distributed_trace
def _initialize(self, **kwargs):
# type: (**Any) -> None
Expand Down Expand Up @@ -185,7 +196,13 @@ def encrypt(self, algorithm, plaintext, **kwargs):
try:
return self._local_provider.encrypt(algorithm, plaintext)
except Exception as ex: # pylint:disable=broad-except
if self._local_only:
raise
mccoyp marked this conversation as resolved.
Show resolved Hide resolved
_LOGGER.warning("Local encrypt operation failed: %s", ex, exc_info=_LOGGER.isEnabledFor(logging.DEBUG))
elif self._local_only:
raise NotImplementedError(
'This key does not support the "encrypt" operation with algorithm "{}"'.format(algorithm)
)

operation_result = self._client.encrypt(
vault_base_url=self._key_id.vault_url,
Expand Down Expand Up @@ -239,7 +256,13 @@ def decrypt(self, algorithm, ciphertext, **kwargs):
try:
return self._local_provider.decrypt(algorithm, ciphertext)
except Exception as ex: # pylint:disable=broad-except
if self._local_only:
raise
_LOGGER.warning("Local decrypt operation failed: %s", ex, exc_info=_LOGGER.isEnabledFor(logging.DEBUG))
elif self._local_only:
raise NotImplementedError(
'This key does not support the "decrypt" operation with algorithm "{}"'.format(algorithm)
)

operation_result = self._client.decrypt(
vault_base_url=self._key_id.vault_url,
Expand Down Expand Up @@ -276,7 +299,13 @@ def wrap_key(self, algorithm, key, **kwargs):
try:
return self._local_provider.wrap_key(algorithm, key)
except Exception as ex: # pylint:disable=broad-except
if self._local_only:
raise
_LOGGER.warning("Local wrap operation failed: %s", ex, exc_info=_LOGGER.isEnabledFor(logging.DEBUG))
elif self._local_only:
raise NotImplementedError(
'This key does not support the "wrapKey" operation with algorithm "{}"'.format(algorithm)
)

operation_result = self._client.wrap_key(
vault_base_url=self._key_id.vault_url,
Expand Down Expand Up @@ -310,7 +339,13 @@ def unwrap_key(self, algorithm, encrypted_key, **kwargs):
try:
return self._local_provider.unwrap_key(algorithm, encrypted_key)
except Exception as ex: # pylint:disable=broad-except
if self._local_only:
raise
_LOGGER.warning("Local unwrap operation failed: %s", ex, exc_info=_LOGGER.isEnabledFor(logging.DEBUG))
elif self._local_only:
raise NotImplementedError(
'This key does not support the "unwrapKey" operation with algorithm "{}"'.format(algorithm)
)

operation_result = self._client.unwrap_key(
vault_base_url=self._key_id.vault_url,
Expand Down Expand Up @@ -344,7 +379,13 @@ def sign(self, algorithm, digest, **kwargs):
try:
return self._local_provider.sign(algorithm, digest)
except Exception as ex: # pylint:disable=broad-except
if self._local_only:
raise
_LOGGER.warning("Local sign operation failed: %s", ex, exc_info=_LOGGER.isEnabledFor(logging.DEBUG))
elif self._local_only:
raise NotImplementedError(
'This key does not support the "sign" operation with algorithm "{}"'.format(algorithm)
)

operation_result = self._client.sign(
vault_base_url=self._key_id.vault_url,
Expand Down Expand Up @@ -380,7 +421,13 @@ def verify(self, algorithm, digest, signature, **kwargs):
try:
return self._local_provider.verify(algorithm, digest, signature)
except Exception as ex: # pylint:disable=broad-except
if self._local_only:
raise
_LOGGER.warning("Local verify operation failed: %s", ex, exc_info=_LOGGER.isEnabledFor(logging.DEBUG))
elif self._local_only:
raise NotImplementedError(
'This key does not support the "verify" operation with algorithm "{}"'.format(algorithm)
)

operation_result = self._client.verify(
vault_base_url=self._key_id.vault_url,
Expand Down
50 changes: 49 additions & 1 deletion sdk/keyvault/azure-keyvault-keys/azure/keyvault/keys/crypto/aio/_client.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -53,6 +53,8 @@ class CryptographyClient(AsyncKeyVaultClientBase):
"""

def __init__(self, key: "Union[KeyVaultKey, str]", credential: "AsyncTokenCredential", **kwargs: "Any") -> None:
self._local_only = kwargs.pop("_local_only", False)

if isinstance(key, KeyVaultKey):
self._key = key
self._key_id = parse_key_vault_id(key.id)
Expand All @@ -63,7 +65,7 @@ def __init__(self, key: "Union[KeyVaultKey, str]", credential: "AsyncTokenCreden
else:
raise ValueError("'key' must be a KeyVaultKey instance or a key ID string including a version")

if not self._key_id.version:
if not (self._key_id.version or self._local_only):
raise ValueError("'key' must include a version")

self._local_provider = NoLocalCryptography()
Expand All @@ -79,6 +81,16 @@ def key_id(self) -> str:
"""
return self._key_id.source_id

@classmethod
def from_jwk(cls, jwk):
# type: (dict) -> CryptographyClient
"""Creates a client that can only perform cryptographic operations locally.

:param dict jwk: the key's cryptographic material, as a dictionary.
"""
key_id = "https://key-vault.vault.azure.net/keys/local-key"
return cls(KeyVaultKey(key_id, jwk), object, _local_only=True)

@distributed_trace_async
async def _initialize(self, **kwargs):
# type: (**Any) -> None
Expand Down Expand Up @@ -137,7 +149,13 @@ async def encrypt(self, algorithm: "EncryptionAlgorithm", plaintext: bytes, **kw
try:
return self._local_provider.encrypt(algorithm, plaintext)
except Exception as ex: # pylint:disable=broad-except
if self._local_only:
raise
_LOGGER.warning("Local encrypt operation failed: %s", ex, exc_info=_LOGGER.isEnabledFor(logging.DEBUG))
elif self._local_only:
raise NotImplementedError(
'This key does not support the "encrypt" operation with algorithm "{}"'.format(algorithm)
)

operation_result = await self._client.encrypt(
vault_base_url=self._key_id.vault_url,
Expand Down Expand Up @@ -190,7 +208,13 @@ async def decrypt(self, algorithm: "EncryptionAlgorithm", ciphertext: bytes, **k
try:
return self._local_provider.decrypt(algorithm, ciphertext)
except Exception as ex: # pylint:disable=broad-except
if self._local_only:
raise
_LOGGER.warning("Local decrypt operation failed: %s", ex, exc_info=_LOGGER.isEnabledFor(logging.DEBUG))
elif self._local_only:
raise NotImplementedError(
'This key does not support the "decrypt" operation with algorithm "{}"'.format(algorithm)
)

operation_result = await self._client.decrypt(
vault_base_url=self._key_id.vault_url,
Expand Down Expand Up @@ -226,7 +250,13 @@ async def wrap_key(self, algorithm: "KeyWrapAlgorithm", key: bytes, **kwargs: "A
try:
return self._local_provider.wrap_key(algorithm, key)
except Exception as ex: # pylint:disable=broad-except
if self._local_only:
raise
_LOGGER.warning("Local wrap operation failed: %s", ex, exc_info=_LOGGER.isEnabledFor(logging.DEBUG))
elif self._local_only:
raise NotImplementedError(
'This key does not support the "wrapKey" operation with algorithm "{}"'.format(algorithm)
)

operation_result = await self._client.wrap_key(
vault_base_url=self._key_id.vault_url,
Expand Down Expand Up @@ -259,7 +289,13 @@ async def unwrap_key(self, algorithm: "KeyWrapAlgorithm", encrypted_key: bytes,
try:
return self._local_provider.unwrap_key(algorithm, encrypted_key)
except Exception as ex: # pylint:disable=broad-except
if self._local_only:
raise
_LOGGER.warning("Local unwrap operation failed: %s", ex, exc_info=_LOGGER.isEnabledFor(logging.DEBUG))
elif self._local_only:
raise NotImplementedError(
'This key does not support the "unwrapKey" operation with algorithm "{}"'.format(algorithm)
)

operation_result = await self._client.unwrap_key(
vault_base_url=self._key_id.vault_url,
Expand Down Expand Up @@ -293,7 +329,13 @@ async def sign(self, algorithm: "SignatureAlgorithm", digest: bytes, **kwargs: "
try:
return self._local_provider.sign(algorithm, digest)
except Exception as ex: # pylint:disable=broad-except
if self._local_only:
raise
_LOGGER.warning("Local sign operation failed: %s", ex, exc_info=_LOGGER.isEnabledFor(logging.DEBUG))
elif self._local_only:
raise NotImplementedError(
'This key does not support the "sign" operation with algorithm "{}"'.format(algorithm)
)

operation_result = await self._client.sign(
vault_base_url=self._key_id.vault_url,
Expand Down Expand Up @@ -330,7 +372,13 @@ async def verify(
try:
return self._local_provider.verify(algorithm, digest, signature)
except Exception as ex: # pylint:disable=broad-except
if self._local_only:
raise
_LOGGER.warning("Local verify operation failed: %s", ex, exc_info=_LOGGER.isEnabledFor(logging.DEBUG))
elif self._local_only:
raise NotImplementedError(
'This key does not support the "verify" operation with algorithm "{}"'.format(algorithm)
)

operation_result = await self._client.verify(
vault_base_url=self._key_id.vault_url,
Expand Down
80 changes: 78 additions & 2 deletions sdk/keyvault/azure-keyvault-keys/tests/test_crypto_client.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -12,9 +12,9 @@
except ImportError:
import mock

from azure.core.exceptions import HttpResponseError
from azure.core.exceptions import AzureError, HttpResponseError
from azure.core.pipeline.policies import SansIOHTTPPolicy
from azure.keyvault.keys import JsonWebKey, KeyClient, KeyCurveName, KeyType, KeyVaultKey
from azure.keyvault.keys import JsonWebKey, KeyClient, KeyCurveName, KeyOperation, KeyVaultKey
from azure.keyvault.keys.crypto import CryptographyClient, EncryptionAlgorithm, KeyWrapAlgorithm, SignatureAlgorithm
from azure.keyvault.keys.crypto._key_validity import _UTC
from azure.keyvault.keys._shared import HttpChallengeCache
Expand Down Expand Up @@ -491,6 +491,82 @@ def test_calls_service_for_operations_unsupported_locally():
assert supports_nothing.wrap_key.call_count == 0


def test_local_only_mode_no_service_calls():
"""A local-only CryptographyClient shouldn't call the service if an operation can't be performed locally"""

mock_client = mock.Mock()
jwk = mock.Mock(spec=JsonWebKey)
client = CryptographyClient.from_jwk(jwk=jwk)
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

We expect no request, so I would use a mock transport that raises when the client tries to send e.g.

mock.Mock(send=mock.Mock(side_effect=Exception("client shouldn't send a request"))

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

(That is, if it would be a bug for the client to send a request in local-only mode.)

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This is one that I've made a note of to revisit -- I'm not too familiar with Mock and wasn't able to randomly turn enough dials to get a better understanding. I figure that asserting that the underlying client doesn't call get_key is sufficient for the time being

client._client = mock_client

supports_nothing = mock.Mock(supports=mock.Mock(return_value=False))
with mock.patch(CryptographyClient.__module__ + ".get_local_cryptography_provider", lambda *_: supports_nothing):
with pytest.raises(NotImplementedError):
client.decrypt(EncryptionAlgorithm.rsa_oaep, b"...")
assert mock_client.decrypt.call_count == 0

with pytest.raises(NotImplementedError):
client.encrypt(EncryptionAlgorithm.rsa_oaep, b"...")
assert mock_client.encrypt.call_count == 0

with pytest.raises(NotImplementedError):
client.sign(SignatureAlgorithm.rs256, b"...")
assert mock_client.sign.call_count == 0

with pytest.raises(NotImplementedError):
client.verify(SignatureAlgorithm.rs256, b"...", b"...")
assert mock_client.verify.call_count == 0

with pytest.raises(NotImplementedError):
client.unwrap_key(KeyWrapAlgorithm.rsa_oaep, b"...")
assert mock_client.unwrap_key.call_count == 0

with pytest.raises(NotImplementedError):
client.wrap_key(KeyWrapAlgorithm.rsa_oaep, b"...")
assert mock_client.wrap_key.call_count == 0


def test_local_only_mode_raise():
"""A local-only CryptographyClient should raise an exception if an operation can't be performed locally"""

jwk = {"kty":"RSA", "key_ops":["decrypt", "verify", "unwrapKey"], "n":b"10011", "e":b"10001"}
client = CryptographyClient.from_jwk(jwk=jwk)

# Algorithm not supported locally
with pytest.raises(NotImplementedError) as ex:
client.decrypt(EncryptionAlgorithm.a256_gcm, b"...")
assert EncryptionAlgorithm.a256_gcm in str(ex.value)
assert KeyOperation.decrypt in str(ex.value)

# Operation not included in JWK permissions
with pytest.raises(AzureError) as ex:
client.encrypt(EncryptionAlgorithm.rsa_oaep, b"...")
assert KeyOperation.encrypt in str(ex.value)

# Algorithm not supported locally
with pytest.raises(NotImplementedError) as ex:
client.verify(SignatureAlgorithm.es256, b"...", b"...")
assert SignatureAlgorithm.es256 in str(ex.value)
assert KeyOperation.verify in str(ex.value)

# Algorithm not supported locally, and operation not included in JWK permissions
with pytest.raises(NotImplementedError) as ex:
client.sign(SignatureAlgorithm.rs256, b"...")
assert SignatureAlgorithm.rs256 in str(ex.value)
assert KeyOperation.sign in str(ex.value)

# Algorithm not supported locally
with pytest.raises(NotImplementedError) as ex:
client.unwrap_key(KeyWrapAlgorithm.aes_256, b"...")
assert KeyWrapAlgorithm.aes_256 in str(ex.value)
assert KeyOperation.unwrap_key in str(ex.value)

# Operation not included in JWK permissions
with pytest.raises(AzureError) as ex:
client.wrap_key(KeyWrapAlgorithm.rsa_oaep, b"...")
assert KeyOperation.wrap_key in str(ex.value)


def test_prefers_local_provider():
"""The client should complete operations locally whenever possible"""

Expand Down
Loading