Key Vault clients use azure-identity credentials

sdk/keyvault/azure-security-keyvault/azure/security/keyvault/__init__.py

@@ -3,9 +3,11 @@
 # Licensed under the MIT License. See LICENSE.txt in the project root for
 # license information.
 # --------------------------------------------------------------------------
+from .keys._client import KeyClient
+from .secrets._client import SecretClient
 from .vault_client import VaultClient
 from .version import VERSION
 
-__all__ = ["VaultClient"]
+__all__ = ["VaultClient", "KeyClient", "SecretClient"]
 
 __version__ = VERSION

sdk/keyvault/azure-security-keyvault/azure/security/keyvault/_internal.py

@@ -3,19 +3,34 @@
 # Licensed under the MIT License. See LICENSE.txt in the project root for
 # license information.
 # --------------------------------------------------------------------------
-from azure.core.pipeline.policies import HTTPPolicy
+from collections import namedtuple
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    # pylint:disable=unused-import
+    from typing import Any, Mapping, Optional
+    from azure.core.credentials import TokenCredential
+    from azure.core.pipeline.transport import HttpTransport
 
 try:
     import urllib.parse as parse
 except ImportError:
     import urlparse as parse  # pylint: disable=import-error
 
-from collections import namedtuple
+from azure.core import Configuration
+from azure.core.pipeline import Pipeline
+from azure.core.pipeline.policies import BearerTokenCredentialPolicy, HTTPPolicy
+from azure.core.pipeline.transport import RequestsTransport
+
+from ._generated import KeyVaultClient
 
 
 _VaultId = namedtuple("VaultId", ["vault_url", "collection", "name", "version"])
 
 
+KEY_VAULT_SCOPES = ("https://vault.azure.net/.default",)
+
+
 def _parse_vault_id(url):
     try:
         parsed_uri = parse.urlparse(url)
@@ -37,13 +52,67 @@ def _parse_vault_id(url):
     )
 
 
-# TODO: integrate with azure.core
-class _BearerTokenCredentialPolicy(HTTPPolicy):
-    def __init__(self, credentials):
-        self._credentials = credentials
+class _KeyVaultClientBase(object):
+    """
+    :param credential:  A credential or credential provider which can be used to authenticate to the vault,
+        a ValueError will be raised if the entity is not provided
+    :type credential: azure.core.credentials.TokenCredential
+    :param str vault_url: The url of the vault to which the client will connect,
+        a ValueError will be raised if the entity is not provided
+    :param ~azure.core.configuration.Configuration config:  The configuration for the KeyClient
+    """
+
+    @staticmethod
+    def create_config(credential, api_version=None, **kwargs):
+        # type: (TokenCredential, Optional[str], Mapping[str, Any]) -> Configuration
+        if api_version is None:
+            api_version = KeyVaultClient.DEFAULT_API_VERSION
+        config = KeyVaultClient.get_configuration_class(api_version, aio=False)(credential, **kwargs)
+        config.authentication_policy = BearerTokenCredentialPolicy(credential, scopes=KEY_VAULT_SCOPES)
+        return config
+
+    def __init__(self, vault_url, credential, config=None, transport=None, api_version=None, **kwargs):
+        # type: (str, TokenCredential, Configuration, Optional[HttpTransport], Optional[str], **Any) -> None
+        if not credential:
+            raise ValueError(
+                "credential should be an object supporting the TokenCredential protocol, such as a credential from azure-identity"
+            )
+        if not vault_url:
+            raise ValueError("vault_url must be the URL of an Azure Key Vault")
+
+        self._vault_url = vault_url.strip(" /")
+
+        client = kwargs.pop("generated_client", None)
+        if client:
+            # caller provided a configured client -> nothing left to initialize
+            self._client = client
+            return
+
+        if api_version is None:
+            api_version = KeyVaultClient.DEFAULT_API_VERSION
+
+        config = config or self.create_config(credential, api_version=api_version, **kwargs)
+        pipeline = kwargs.pop("pipeline", None) or self._build_pipeline(config, transport)
+        self._client = KeyVaultClient(credential, api_version=api_version, pipeline=pipeline, aio=False, **kwargs)
+
+    def _build_pipeline(self, config, transport):
+        # type: (Configuration, HttpTransport) -> Pipeline
+        policies = [
+            config.headers_policy,
+            config.user_agent_policy,
+            config.proxy_policy,
+            config.redirect_policy,
+            config.retry_policy,
+            config.authentication_policy,
+            config.logging_policy,
+        ]
+
+        if transport is None:
+            transport = RequestsTransport(config)
 
-    def send(self, request, **kwargs):
-        auth_header = "Bearer " + self._credentials.token["access_token"]
-        request.http_request.headers["Authorization"] = auth_header
+        return Pipeline(transport, policies=policies)
 
-        return self.next.send(request, **kwargs)
+    @property
+    def vault_url(self):
+        # type: () -> str
+        return self._vault_url

sdk/keyvault/azure-security-keyvault/azure/security/keyvault/aio/_internal.py

@@ -3,25 +3,112 @@
 # Licensed under the MIT License. See LICENSE.txt in the project root for
 # license information.
 # --------------------------------------------------------------------------
-from typing import Any, Callable
+from typing import Any, Callable, Mapping, Optional, TYPE_CHECKING
+
+if TYPE_CHECKING:
+    try:
+        from azure.core.credentials import TokenCredential
+    except ImportError:
+        # TokenCredential is a typing_extensions.Protocol; we don't depend on that package
+        pass
 
 from azure.core.async_paging import AsyncPagedMixin
+from azure.core.configuration import Configuration
+from azure.core.pipeline import AsyncPipeline
+from azure.core.pipeline.policies import AsyncBearerTokenCredentialPolicy
+from azure.core.pipeline.transport import AsyncioRequestsTransport, HttpTransport
 from msrest.serialization import Model
 
+from .._generated import KeyVaultClient
+from .._internal import KEY_VAULT_SCOPES
+
 
 class AsyncPagingAdapter:
     """For each item in an AsyncPagedMixin, returns the result of applying fn to that item.
     Python 3.6 added syntax that could replace this (yield within async for)."""
-    def __init__(self, pages, fn):
-        # type: (AsyncPagedMixin, Callable[[Model], Any]) -> None
+
+    def __init__(self, pages: AsyncPagedMixin, fn: Callable[[Model], Any]) -> None:
         self._pages = pages
         self._fn = fn
 
     def __aiter__(self):
         return self
 
-    async def __anext__(self):
+    async def __anext__(self) -> Any:
         item = await self._pages.__anext__()
         if not item:
             raise StopAsyncIteration
         return self._fn(item)
+
+
+class _AsyncKeyVaultClientBase:
+    """
+    :param credential:  A credential or credential provider which can be used to authenticate to the vault,
+        a ValueError will be raised if the entity is not provided
+    :type credential: azure.authentication.Credential or azure.authentication.CredentialProvider
+    :param str vault_url: The url of the vault to which the client will connect,
+        a ValueError will be raised if the entity is not provided
+    :param ~azure.core.configuration.Configuration config:  The configuration for the SecretClient
+    """
+
+    @staticmethod
+    def create_config(
+        credential: "TokenCredential", api_version: str = None, **kwargs: Mapping[str, Any]
+    ) -> Configuration:
+        if api_version is None:
+            api_version = KeyVaultClient.DEFAULT_API_VERSION
+        config = KeyVaultClient.get_configuration_class(api_version, aio=True)(credential, **kwargs)
+        config.authentication_policy = AsyncBearerTokenCredentialPolicy(credential, scopes=KEY_VAULT_SCOPES)
+        return config
+
+    def __init__(
+        self,
+        vault_url: str,
+        credential: "TokenCredential",
+        config: Configuration = None,
+        transport: HttpTransport = None,
+        api_version: str = None,
+        **kwargs: Any
+    ) -> None:
+        if not credential:
+            raise ValueError(
+                "credential should be an object supporting the TokenCredential protocol, such as a credential from azure-identity"
+            )
+        if not vault_url:
+            raise ValueError("vault_url must be the URL of an Azure Key Vault")
+
+        self._vault_url = vault_url.strip(" /")
+
+        client = kwargs.pop("generated_client", None)
+        if client:
+            # caller provided a configured client -> nothing left to initialize
+            self._client = client
+            return
+
+        if api_version is None:
+            api_version = KeyVaultClient.DEFAULT_API_VERSION
+
+        config = config or self.create_config(credential, api_version=api_version, **kwargs)
+        pipeline = kwargs.pop("pipeline", None) or self._build_pipeline(config, transport=transport)
+        self._client = KeyVaultClient(credential, api_version=api_version, pipeline=pipeline, aio=True)
+
+    @staticmethod
+    def _build_pipeline(config: Configuration, transport: HttpTransport) -> AsyncPipeline:
+        policies = [
+            config.headers_policy,
+            config.user_agent_policy,
+            config.proxy_policy,
+            config.redirect_policy,
+            config.retry_policy,
+            config.authentication_policy,
+            config.logging_policy,
+        ]
+
+        if transport is None:
+            transport = AsyncioRequestsTransport(config)
+
+        return AsyncPipeline(transport, policies=policies)
+
+    @property
+    def vault_url(self) -> str:
+        return self._vault_url