azure_keyvault_configuration.py

import config
from azure.keyvault.secrets import SecretClient
from azure.identity import ClientSecretCredential
from azure.core.exceptions import ResourceNotFoundError
import time


class AzureKeyVaultConfiguration(config.Configuration):
    """
    Configuration class.

    The Configuration class takes azure keyvault credentials and behaves like a drop-in replacement
    for the regular configuration.
    """

    def __init__(self, az_client_id, az_client_secret, az_tenant_id, az_vault_name,
                 cache_expiration=5*60,
                 lowercase_keys: bool = False):
        """
        Constructor.

        :param config_: a mapping of configuration values. Keys need to be strings.
        :param lowercase_keys: whether to convert every key to lower case.
        """
        credentials = ClientSecretCredential(
            client_id=az_client_id,
            client_secret=az_client_secret,
            tenant_id=az_tenant_id
        )
        vault_url = "https://{az_vault_name}.vault.azure.net/".format(az_vault_name=az_vault_name)
        self.kvclient = SecretClient(vault_url=vault_url, credential=credentials)
        self.cache_expiration = cache_expiration
        self.__cache = {}

    def __eq__(self, other):  # type: ignore
        """Equality operator."""
        return self.as_dict() == other.as_dict()

    def __get_secret(self, key):
        key = key.replace('_', '-')  # Normalize for Azure Keyvault
        now = time.time()
        from_cache = self.__cache.get(key)
        if (not (from_cache is None)) and (from_cache[1] + self.cache_expiration > now):
            return from_cache[0]
        try:
            secret = self.kvclient.get_secret(key)
            self.__cache[key] = (secret.value, now)
            return secret.value
        except ResourceNotFoundError:
            if (key in self.__cache):
                del self.__cache[key]
            return None


    def __getitem__(self, item: str):
        secret = self.__get_secret(item)
        if secret is None:
            raise KeyError(item)
        else:
            return secret

    def __getattr__(self, item: str):
        secret = self.__get_secret(item)
        if secret is None:
            raise KeyError(item)
        else:
            return secret

    def get(self, key: str, default = None):
        """
        Get the configuration values corresponding to :attr:`key`.

        :param key: key to retrieve
        :param default: default value in case the key is missing
        :return: the value found or a default
        """
        secret = self.__get_secret(key)
        if secret is None:
            return default
        else:
            return secret