[Blob][SAS] Added support for identity SAS (#7020)

sdk/storage/azure-storage-blob/azure/storage/blob/_shared/models.py

@@ -420,3 +420,36 @@ def __str__(self):
 Services.BLOB = Services(blob=True)
 Services.QUEUE = Services(queue=True)
 Services.FILE = Services(file=True)
+
+
+class UserDelegationKey(object):
+    """
+    Represents a user delegation key, provided to the user by Azure Storage
+    based on their Azure Active Directory access token.
+
+    The fields are saved as simple strings since the user does not have to interact with this object;
+    to generate an identify SAS, the user can simply pass it to the right API.
+
+    :ivar str signed_oid:
+        Object ID of this token.
+    :ivar str signed_tid:
+        Tenant ID of the tenant that issued this token.
+    :ivar str signed_start:
+        The datetime this token becomes valid.
+    :ivar str signed_expiry:
+        The datetime this token expires.
+    :ivar str signed_service:
+        What service this key is valid for.
+    :ivar str signed_version:
+        The version identifier of the REST service that created this token.
+    :ivar str value:
+        The user delegation key.
+    """
+    def __init__(self):
+        self.signed_oid = None
+        self.signed_tid = None
+        self.signed_start = None
+        self.signed_expiry = None
+        self.signed_service = None
+        self.signed_version = None
+        self.value = None

sdk/storage/azure-storage-blob/azure/storage/blob/_shared/response_handlers.py

@@ -19,7 +19,8 @@
     ClientAuthenticationError,
     DecodeError)
 
-from .models import StorageErrorCode
+from .parser import _to_utc_datetime
+from .models import StorageErrorCode, UserDelegationKey
 
 
 if TYPE_CHECKING:
@@ -131,3 +132,15 @@ def process_storage_error(storage_error):
     error.error_code = error_code
     error.additional_info = additional_data
     raise error
+
+
+def parse_to_internal_user_delegation_key(service_user_delegation_key):
+    internal_user_delegation_key = UserDelegationKey()
+    internal_user_delegation_key.signed_oid = service_user_delegation_key.signed_oid
+    internal_user_delegation_key.signed_tid = service_user_delegation_key.signed_tid
+    internal_user_delegation_key.signed_start = _to_utc_datetime(service_user_delegation_key.signed_start)
+    internal_user_delegation_key.signed_expiry = _to_utc_datetime(service_user_delegation_key.signed_expiry)
+    internal_user_delegation_key.signed_service = service_user_delegation_key.signed_service
+    internal_user_delegation_key.signed_version = service_user_delegation_key.signed_version
+    internal_user_delegation_key.value = service_user_delegation_key.value
+    return internal_user_delegation_key

sdk/storage/azure-storage-blob/azure/storage/blob/_shared_access_signature.py

@@ -188,7 +188,7 @@ def generate_container(self, container_name, permission=None, expiry=None,
                                           content_encoding, content_language,
                                           content_type)
         sas.add_resource_signature(self.account_name, self.account_key, container_name,
-                                   user_delegation_key=None)
+                                   user_delegation_key=self.user_delegation_key)
         return sas.get_token()
 
 

sdk/storage/azure-storage-blob/azure/storage/blob/aio/blob_service_client_async.py

@@ -18,8 +18,10 @@
 from .._shared.policies_async import ExponentialRetry
 from .._shared.base_client_async import AsyncStorageAccountHostsMixin
 from .._shared.response_handlers import return_response_headers, process_storage_error
+from .._shared.parser import _to_utc_datetime
+from .._shared.response_handlers import parse_to_internal_user_delegation_key
 from .._generated.aio import AzureBlobStorage
-from .._generated.models import StorageErrorException, StorageServiceProperties
+from .._generated.models import StorageErrorException, StorageServiceProperties, KeyInfo
 from ..blob_service_client import BlobServiceClient as BlobServiceClientBase
 from .container_client_async import ContainerClient
 from .blob_client_async import BlobClient
@@ -110,6 +112,36 @@ def __init__(
         self._client = AzureBlobStorage(url=self.url, pipeline=self._pipeline, loop=loop)
         self._loop = loop
 
+    @distributed_trace_async
+    async def get_user_delegation_key(self, key_start_time,  # type: datetime
+                                      key_expiry_time,  # type: datetime
+                                      timeout=None,  # type: Optional[int]
+                                      **kwargs  # type: Any
+                                      ):
+        # type: (datetime, datetime, Optional[int]) -> UserDelegationKey
+        """
+        Obtain a user delegation key for the purpose of signing SAS tokens.
+        A token credential must be present on the service object for this request to succeed.
+
+        :param datetime key_start_time:
+            A DateTime value. Indicates when the key becomes valid.
+        :param datetime key_expiry_time:
+            A DateTime value. Indicates when the key stops being valid.
+        :param int timeout:
+            The timeout parameter is expressed in seconds.
+        :return: The user delegation key.
+        :rtype: ~azure.storage.blob._shared.models.UserDelegationKey
+        """
+        key_info = KeyInfo(start=_to_utc_datetime(key_start_time), expiry=_to_utc_datetime(key_expiry_time))
+        try:
+            user_delegation_key = await self._client.service.get_user_delegation_key(key_info=key_info,
+                                                                                     timeout=timeout,
+                                                                                     **kwargs)  # type: ignore
+        except StorageErrorException as error:
+            process_storage_error(error)
+
+        return parse_to_internal_user_delegation_key(user_delegation_key)  # type: ignore
+
     @distributed_trace_async
     async def get_account_information(self, **kwargs): # type: ignore
         # type: (Optional[int]) -> Dict[str, str]

sdk/storage/azure-storage-blob/azure/storage/blob/blob_client.py

@@ -29,7 +29,7 @@
     validate_and_format_range_headers)
 from ._shared.response_handlers import return_response_headers, process_storage_error
 from ._generated import AzureBlobStorage
-from ._generated.models import (
+from ._generated.models import ( # pylint: disable=unused-import
     DeleteSnapshotsOptionType,
     BlobHTTPHeaders,
     BlockLookupList,
@@ -38,7 +38,9 @@
     ModifiedAccessConditions,
     SequenceNumberAccessConditions,
     StorageErrorException,
+    UserDelegationKey,
     CpkInfo)
+
 from ._deserialize import deserialize_blob_properties, deserialize_blob_stream
 from ._upload_helpers import (
     upload_block_blob,
@@ -134,6 +136,7 @@ def __init__(
         except AttributeError:
             raise ValueError("Blob URL must be a string.")
         parsed_url = urlparse(blob_url.rstrip('/'))
+
         if not parsed_url.path and not (container and blob):
             raise ValueError("Please specify a container and blob name.")
         if not parsed_url.netloc:
@@ -228,12 +231,14 @@ def generate_shared_access_signature(
             policy_id=None,  # type: Optional[str]
             ip=None,  # type: Optional[str]
             protocol=None,  # type: Optional[str]
+            account_name=None,  # type: Optional[str]
             cache_control=None,  # type: Optional[str]
             content_disposition=None,  # type: Optional[str]
             content_encoding=None,  # type: Optional[str]
             content_language=None,  # type: Optional[str]
-            content_type=None  # type: Optional[str]
-        ):
+            content_type=None,  # type: Optional[str]
+            user_delegation_key=None  # type: Optional[UserDelegationKey]
+            ):
         # type: (...) -> Any
         """
         Generates a shared access signature for the blob.
@@ -275,6 +280,8 @@ def generate_shared_access_signature(
             restricts the request to those IP addresses.
         :param str protocol:
             Specifies the protocol permitted for a request made. The default value is https.
+        :param str account_name:
+            Specifies the account_name when using oauth token as credential. If you use oauth token as credential.
         :param str cache_control:
             Response header value for Cache-Control when resource is accessed
             using this shared access signature.
@@ -290,12 +297,24 @@ def generate_shared_access_signature(
         :param str content_type:
             Response header value for Content-Type when resource is accessed
             using this shared access signature.
+        :param ~azure.storage.blob._shared.models.UserDelegationKey user_delegation_key:
+            Instead of an account key, the user could pass in a user delegation key.
+            A user delegation key can be obtained from the service by authenticating with an AAD identity;
+            this can be accomplished by calling get_user_delegation_key.
+            When present, the SAS is signed with the user delegation key instead.
         :return: A Shared Access Signature (sas) token.
         :rtype: str
         """
-        if not hasattr(self.credential, 'account_key') or not self.credential.account_key:
-            raise ValueError("No account SAS key available.")
-        sas = BlobSharedAccessSignature(self.credential.account_name, self.credential.account_key)
+        if user_delegation_key is not None:
+            if not hasattr(self.credential, 'account_name') and not account_name:
+                raise ValueError("No account_name available. Please provide account_name parameter.")
+
+            account_name = self.credential.account_name if hasattr(self.credential, 'account_name') else account_name
+            sas = BlobSharedAccessSignature(account_name, user_delegation_key=user_delegation_key)
+        else:
+            if not hasattr(self.credential, 'account_key') or not self.credential.account_key:
+                raise ValueError("No account SAS key available.")
+            sas = BlobSharedAccessSignature(self.credential.account_name, self.credential.account_key)
         return sas.generate_blob(
             self.container_name,
             self.blob_name,

sdk/storage/azure-storage-blob/azure/storage/blob/blob_service_client.py

@@ -9,6 +9,7 @@
     Union, Optional, Any, Iterable, Dict, List,
     TYPE_CHECKING
 )
+
 try:
     from urllib.parse import urlparse
 except ImportError:
@@ -18,11 +19,13 @@
 from azure.core.tracing.decorator import distributed_trace
 
 from ._shared.shared_access_signature import SharedAccessSignature
-from ._shared.models import LocationMode, Services
+from ._shared.models import LocationMode, Services, UserDelegationKey
 from ._shared.base_client import StorageAccountHostsMixin, parse_connection_str, parse_query
-from ._shared.response_handlers import return_response_headers, process_storage_error
+from ._shared.parser import _to_utc_datetime
+from ._shared.response_handlers import return_response_headers, process_storage_error, \
+    parse_to_internal_user_delegation_key
 from ._generated import AzureBlobStorage
-from ._generated.models import StorageErrorException, StorageServiceProperties
+from ._generated.models import StorageErrorException, StorageServiceProperties, KeyInfo
 from .container_client import ContainerClient
 from .blob_client import BlobClient
 from .models import ContainerProperties, ContainerPropertiesPaged
@@ -223,6 +226,36 @@ def generate_shared_access_signature(
             protocol=protocol
         ) # type: ignore
 
+    @distributed_trace
+    def get_user_delegation_key(self, key_start_time,  # type: datetime
+                                key_expiry_time,  # type: datetime
+                                timeout=None,  # type: Optional[int]
+                                **kwargs  # type: Any
+                                ):
+        # type: (datetime, datetime, Optional[int]) -> UserDelegationKey
+        """
+        Obtain a user delegation key for the purpose of signing SAS tokens.
+        A token credential must be present on the service object for this request to succeed.
+
+        :param datetime key_start_time:
+            A DateTime value. Indicates when the key becomes valid.
+        :param datetime key_expiry_time:
+            A DateTime value. Indicates when the key stops being valid.
+        :param int timeout:
+            The timeout parameter is expressed in seconds.
+        :return: The user delegation key.
+        :rtype: ~azure.storage.blob._shared.models.UserDelegationKey
+        """
+        key_info = KeyInfo(start=_to_utc_datetime(key_start_time), expiry=_to_utc_datetime(key_expiry_time))
+        try:
+            user_delegation_key = self._client.service.get_user_delegation_key(key_info=key_info,
+                                                                               timeout=timeout,
+                                                                               **kwargs)  # type: ignore
+        except StorageErrorException as error:
+            process_storage_error(error)
+
+        return parse_to_internal_user_delegation_key(user_delegation_key)  # type: ignore
+
     @distributed_trace
     def get_account_information(self, **kwargs): # type: ignore
         # type: (Optional[int]) -> Dict[str, str]

sdk/storage/azure-storage-blob/azure/storage/blob/container_client.py

@@ -119,6 +119,7 @@ def __init__(
         except AttributeError:
             raise ValueError("Container URL must be a string.")
         parsed_url = urlparse(container_url.rstrip('/'))
+
         if not parsed_url.path and not container:
             raise ValueError("Please specify a container name.")
         if not parsed_url.netloc:
@@ -188,11 +189,13 @@ def generate_shared_access_signature(
             policy_id=None,  # type: Optional[str]
             ip=None,  # type: Optional[str]
             protocol=None,  # type: Optional[str]
+            account_name=None,  # type: Optional[str]
             cache_control=None,  # type: Optional[str]
             content_disposition=None,  # type: Optional[str]
             content_encoding=None,  # type: Optional[str]
             content_language=None,  # type: Optional[str]
-            content_type=None  # type: Optional[str]
+            content_type=None,  # type: Optional[str]
+            user_delegation_key=None  # type Optional[]
         ):
         # type: (...) -> Any
         """Generates a shared access signature for the container.
@@ -233,6 +236,8 @@ def generate_shared_access_signature(
             restricts the request to those IP addresses.
         :param str protocol:
             Specifies the protocol permitted for a request made. The default value is https.
+        :param str account_name:
+            Specifies the account_name when using oauth token as credential. If you use oauth token as credential.
         :param str cache_control:
             Response header value for Cache-Control when resource is accessed
             using this shared access signature.
@@ -248,6 +253,11 @@ def generate_shared_access_signature(
         :param str content_type:
             Response header value for Content-Type when resource is accessed
             using this shared access signature.
+        :param ~azure.storage.blob._shared.models.UserDelegationKey user_delegation_key:
+            Instead of an account key, the user could pass in a user delegation key.
+            A user delegation key can be obtained from the service by authenticating with an AAD identity;
+            this can be accomplished by calling get_user_delegation_key.
+            When present, the SAS is signed with the user delegation key instead.
         :return: A Shared Access Signature (sas) token.
         :rtype: str
 
@@ -259,9 +269,16 @@ def generate_shared_access_signature(
                 :dedent: 12
                 :caption: Generating a sas token.
         """
-        if not hasattr(self.credential, 'account_key') and not self.credential.account_key:
-            raise ValueError("No account SAS key available.")
-        sas = BlobSharedAccessSignature(self.credential.account_name, self.credential.account_key)
+        if user_delegation_key is not None:
+            if not hasattr(self.credential, 'account_name') and not account_name:
+                raise ValueError("No account_name available. Please provide account_name parameter.")
+
+            account_name = self.credential.account_name if hasattr(self.credential, 'account_name') else account_name
+            sas = BlobSharedAccessSignature(account_name, user_delegation_key=user_delegation_key)
+        else:
+            if not hasattr(self.credential, 'account_key') and not self.credential.account_key:
+                raise ValueError("No account SAS key available.")
+            sas = BlobSharedAccessSignature(self.credential.account_name, self.credential.account_key)
         return sas.generate_container(
             self.container_name,
             permission=permission,