Implement abstract private key interfaces from "cryptography" library

[malthe]

Is your feature request related to a problem? Please describe.

The cryptography library (already a dependency) is commonly used across the Python ecosystem.

It defines a number of abstract interfaces such as cryptography.hazmat.primitives.asymmetric.rsa.RSAPrivateKey which can be used as integration points to enable the use of keys in an Azure Key Vault as a drop-in replacement for a locally available key.

For example, the Python client library for Snowflake (since version 3.1.1) allows the use of an arbitrary RSAPrivateKey implementation when using keypair-based authentication.

Describe the solution you'd like

The Azure Key Vault client library should provide an implementation of this interface to allow easy implementation of authentication based on keys managed by Azure Key Vault.

The following is an example implementation that I put together based on code from @arsatiki, initially published as part of the proposal to support this interface in the Snowflake driver.

from cryptography.hazmat.primitives import hashes
from cryptography.hazmat.primitives.asymmetric.rsa import RSAPrivateKey
from cryptography.hazmat.primitives.asymmetric.rsa import RSAPublicKey
from cryptography.hazmat.primitives.asymmetric.rsa import RSAPublicNumbers
from cryptography.hazmat.primitives.asymmetric.rsa import AsymmetricPadding
from cryptography.hazmat.primitives.asymmetric.padding import PKCS1v15
from cryptography.hazmat.primitives.asymmetric import utils as asym_utils


class AzureManagedKey(RSAPrivateKey):
    _algorithm_mapping = {
        hashes.SHA256: SignatureAlgorithm.rs256,
        hashes.SHA384: SignatureAlgorithm.rs384,
        hashes.SHA512: SignatureAlgorithm.rs512,
    }

    # XXX: This is just an example implementation! Fill in these manually ...
    _key_name = KEY_NAME
    _kvclient = KEY_CLIENT

    def decrypt(self, *args) -> bytes:
        raise NotImplementedError()

    @property
    def key_size(self) -> int:
        raise NotImplementedError()

    def sign(
        self,
        data: bytes,
        padding: AsymmetricPadding,
        algorithm: Union[asym_utils.Prehashed, hashes.HashAlgorithm],
    ) -> bytes:
        if not isinstance(padding, PKCS1v15):
            raise ValueError("Unsupported padding: %s" % padding.name)
        sig_algorithm = self._algorithm_mapping.get(type(algorithm))
        if sig_algorithm is None:
            raise ValueError("Unsupported algorithm: %s" % algorithm.name)
        cc = self._kvclient.get_cryptography_client(self._key_name)
        digest = hashes.Hash(algorithm)
        digest.update(data)
        result = cc.sign(sig_algorithm, digest.finalize())
        return result.signature

    def private_numbers(self) -> "RSAPrivateNumbers":
        raise NotImplementedError()

    def private_bytes(self, *args) -> bytes:
        raise NotImplementedError()

    def public_key(self) -> RSAPublicKey:
        pubkey = self._kvclient.get_key(self._key_name)
        e = int.from_bytes(pubkey.key.e, "big")
        n = int.from_bytes(pubkey.key.n, "big")
        rsakey = RSAPublicNumbers(e, n)
        return rsakey.public_key()

Describe alternatives you've considered
The alternative is to contribute this to the cryptography package, or maintain a separate 3rd party package.

Additional context
There's an increase awareness these days of the risks in storing private keys locally. Making it easier to use managed keys is a win for everyone.

I wrote some additional motivation here: https://www.maltheborch.com/2023/09/switching-to-managed-encryption-keys.html that uses Azure Key Vault and Snowflake as a case study.

[xiangyan99]

Thanks for the feedback, we’ll investigate asap.

[lmazuel]

Hi @malthe

Thanks for this very detailled issue with code included, very helpful :). I'm always a fan of those opportunities to integrate better with the ecosystem, but I need to talk to our KeyVault team first as I'm personally not an expert. I'll put that on our agenda for discussion. Thanks again!