ai4eutils/sas_blob_utils.py

"""
Copyright (c) Microsoft Corporation. All rights reserved.
Licensed under the MIT License.

This module contains helper functions for dealing with Shared Access Signatures
(SAS) tokens for Azure Blob Storage.

The default Azure Storage SAS URI format is:
    https://<account>.blob.core.windows.net/<container>/<blob>?<sas_token>

This module assumes azure-storage-blob version 12.5.

Documentation for Azure Blob Storage:
docs.microsoft.com/en-us/azure/developer/python/sdk/storage/storage-blob-readme

Documentation for SAS:
docs.microsoft.com/en-us/azure/storage/common/storage-sas-overview
"""
from datetime import datetime, timedelta
import io
import re
from typing import (
    Any, AnyStr, Dict, IO, Iterable, List, Optional, Set, Tuple, Union)
from urllib import parse
import uuid

from tqdm import tqdm

from azure.storage.blob import (
    BlobClient,
    BlobProperties,
    ContainerClient,
    ContainerSasPermissions,
    generate_container_sas,
    upload_blob_to_url)


def build_azure_storage_uri(
        account: str,
        container: Optional[str] = None,
        blob: Optional[str] = None,
        sas_token: Optional[str] = None,
        account_url_template: str = 'https://{account}.blob.core.windows.net'
        ) -> str:
    """
    Args:
        account: str, name of Azure Storage account
        container: optional str, name of Azure Blob Storage container
        blob: optional str, name of blob, not URL-escaped
            if blob is given, must also specify container
        sas_token: optional str, Shared Access Signature (SAS)
            does not start with '?'
        account_url_template: str, Python 3 string formatting template,
            contains '{account}' placeholder, defaults to default Azure
            Storage URL format. Set this value if using Azurite Azure Storage
            emulator.

    Returns: str, Azure storage URI
    """
    uri = account_url_template.format(account=account)
    if container is not None:
        uri = f'{uri}/{container}'
    if blob is not None:
        assert container is not None
        blob = parse.quote(blob)
        uri = f'{uri}/{blob}'
    if sas_token is not None:
        assert sas_token[0] != '?'
        uri = f'{uri}?{sas_token}'
    return uri


def _get_resource_reference(prefix: str) -> str:
    return '{}{}'.format(prefix, str(uuid.uuid4()).replace('-', ''))


def get_client_from_uri(container_uri: str) -> ContainerClient:
    """Gets a ContainerClient for the given container URI."""
    return ContainerClient.from_container_url(container_uri)


def get_account_from_uri(sas_uri: str) -> str:
    """
    Assumes that sas_uri points to Azure Blob Storage account hosted at
    a default Azure URI. Does not work for locally-emulated Azure Storage
    or Azure Storage hosted at custom endpoints.
    """
    url_parts = parse.urlsplit(sas_uri)
    loc = url_parts.netloc  # "<account>.blob.core.windows.net"
    return loc.split('.')[0]


def get_container_from_uri(sas_uri: str, unquote: bool = True) -> str:
    """Gets the container name from a Azure Blob Storage URI.

    Assumes that sas_uri points to Azure Blob Storage account hosted at
    a default Azure URI. Does not work for locally-emulated Azure Storage
    or Azure Storage hosted at custom endpoints.

    Args:
        sas_uri: str, Azure blob storage URI, may include SAS token
        unquote: bool, whether to replace any %xx escapes by their
            single-character equivalent, default True

    Returns: str, container name

    Raises: ValueError, if sas_uri does not include a container
    """
    url_parts = parse.urlsplit(sas_uri)
    raw_path = url_parts.path.lstrip('/')  # remove leading "/" from path
    container = raw_path.split('/')[0]
    if container == '':
        raise ValueError('Given sas_uri does not include a container.')
    if unquote:
        container = parse.unquote(container)
    return container


def get_blob_from_uri(sas_uri: str, unquote: bool = True) -> str:
    """Return the path to the blob from the root container if this sas_uri
    is for an individual blob; otherwise returns None.

    Args:
        sas_uri: str, Azure blob storage URI, may include SAS token
        unquote: bool, whether to replace any %xx escapes by their
            single-character equivalent, default True

    Returns: str, blob name (path to the blob from the root container)

    Raises: ValueError, if sas_uri does not include a blob name
    """
    # Get the entire path with all slashes after the container
    url_parts = parse.urlsplit(sas_uri)
    raw_path = url_parts.path.lstrip('/')  # remove leading "/" from path
    parts = raw_path.split('/', maxsplit=1)
    if len(parts) < 2 or parts[1] == '':
        raise ValueError('Given sas_uri does not include a blob name')

    blob = parts[1]  # first item is an empty string
    if unquote:
        blob = parse.unquote(blob)
    return blob


def get_sas_token_from_uri(sas_uri: str) -> Optional[str]:
    """Get the query part of the SAS token that contains permissions, access
    times and signature.

    Args:
        sas_uri: str, Azure blob storage SAS token

    Returns: Query part of the SAS token, or None if URI has no token.
    """
    url_parts = parse.urlsplit(sas_uri)
    sas_token = url_parts.query or None  # None if query is empty string
    return sas_token


def get_resource_type_from_uri(sas_uri: str) -> Optional[str]:
    """Get the resource type pointed to by this SAS token.

    Args:
        sas_uri: str, Azure blob storage URI with SAS token

    Returns: A string (either 'blob' or 'container') or None.
    """
    url_parts = parse.urlsplit(sas_uri)
    data = parse.parse_qs(url_parts.query)
    if 'sr' in data:
        types = data['sr']
        if 'b' in types:
            return 'blob'
        elif 'c' in types:
            return 'container'
    return None


def get_endpoint_suffix(sas_uri):
    """Gets the endpoint at which the blob storage account is served.
    Args:
        sas_uri: str, Azure blob storage URI with SAS token

    Returns: A string, usually 'core.windows.net' or 'core.chinacloudapi.cn', to
        use for the `endpoint` argument in various blob storage SDK functions.
    """
    url_parts = parse.urlsplit(sas_uri)
    suffix = url_parts.netloc.split('.blob.')[1].split('/')[0]
    return suffix


def get_permissions_from_uri(sas_uri: str) -> Set[str]:
    """Get the permissions given by this SAS token.

    Args:
        sas_uri: str, Azure blob storage URI with SAS token

    Returns: A set containing some of 'read', 'write', 'delete' and 'list'.
        Empty set returned if no permission specified in sas_uri.
    """
    url_parts = parse.urlsplit(sas_uri)
    data = parse.parse_qs(url_parts.query)
    permissions_set = set()
    if 'sp' in data:
        permissions = data['sp'][0]
        if 'r' in permissions:
            permissions_set.add('read')
        if 'w' in permissions:
            permissions_set.add('write')
        if 'd' in permissions:
            permissions_set.add('delete')
        if 'l' in permissions:
            permissions_set.add('list')
    return permissions_set


def get_all_query_parts(sas_uri: str) -> Dict[str, Any]:
    """Gets the SAS token parameters."""
    url_parts = parse.urlsplit(sas_uri)
    return parse.parse_qs(url_parts.query)


def check_blob_exists(sas_uri: str, blob_name: Optional[str] = None) -> bool:
    """Checks whether a given URI points to an actual blob.

    Args:
        sas_uri: str, URI to a container or a blob
            if blob_name is given, sas_uri is treated as a container URI
            otherwise, sas_uri is treated as a blob URI
        blob_name: optional str, name of blob, not URL-escaped
            must be given if sas_uri is a URI to a container

    Returns: bool, whether the sas_uri given points to an existing blob
    """
    if blob_name is not None:
        sas_uri = build_blob_uri(
            container_uri=sas_uri, blob_name=blob_name)

    with BlobClient.from_blob_url(sas_uri) as blob_client:
        return blob_client.exists()


def list_blobs_in_container(
        container_uri: str,
        blob_prefix: Optional[str] = None,
        blob_suffix: Optional[Union[str, Tuple[str]]] = None,
        rsearch: Optional[str] = None,
        limit: Optional[int] = None
        ) -> List[str]:
    """Get a sorted list of blob names in this container.

    Args:
        container_uri: str, URI to a container, may include SAS token
        blob_prefix: optional str, returned results will only contain blob names
            to with this prefix
        blob_suffix: optional str or tuple of str, returned results will only
            contain blob names with this/these suffix(es). The blob names will
            be lowercased first before comparing with the suffix(es).
        rsearch: optional str, returned results will only contain blob names
            that match this Python regex pattern at any point in the blob name.
            Use '^' character to only match from the beginning of the blob name.
        limit: int, maximum # of blob names to list
            if None, then returns all blob names

    Returns:
        sorted list of blob names, of length limit or shorter.
    """
    print('listing blobs...')
    if (get_sas_token_from_uri(container_uri) is not None
            and get_resource_type_from_uri(container_uri) != 'container'):
        raise ValueError('The SAS token provided is not for a container.')

    if blob_prefix is not None and not isinstance(blob_prefix, str):
        raise ValueError('blob_prefix must be a str.')

    if (blob_suffix is not None
            and not isinstance(blob_suffix, str)
            and not isinstance(blob_suffix, tuple)):
        raise ValueError('blob_suffix must be a str or a tuple of strings')

    list_blobs = []
    with get_client_from_uri(container_uri) as container_client:
        generator = container_client.list_blobs(
            name_starts_with=blob_prefix)

        if blob_suffix is None and rsearch is None:
            list_blobs = [blob.name for blob in tqdm(generator)]
            i = len(list_blobs)
        else:
            i = 0
            for blob in tqdm(generator):
                i += 1
                suffix_ok = (blob_suffix is None
                             or blob.name.lower().endswith(blob_suffix))
                regex_ok = (rsearch is None
                            or re.search(rsearch, blob.name) is not None)
                if suffix_ok and regex_ok:
                    list_blobs.append(blob.name)
                    if limit is not None and len(list_blobs) == limit:
                        break

    print(f'Enumerated {len(list_blobs)} matching blobs out of {i} total')
    return sorted(list_blobs)  # sort for determinism


def generate_writable_container_sas(account_name: str,
                                    account_key: str,
                                    container_name: str,
                                    access_duration_hrs: float,
                                    account_url: Optional[str] = None
                                    ) -> str:
    """Creates a container and returns a SAS URI with read/write/list
    permissions.

    Args:
        account_name: str, name of blob storage account
        account_key: str, account SAS token or account shared access key
        container_name: str, name of container to create, must not match an
            existing container in the given storage account
        access_duration_hrs: float
        account_url: str, optional, defaults to default Azure Storage URL

    Returns: str, URL to newly created container

    Raises: azure.core.exceptions.ResourceExistsError, if container already
        exists

    NOTE: This method currently fails on non-default Azure Storage URLs. The
    initializer for ContainerClient() assumes the default Azure Storage URL
    format, which is a bug that has been reported here:
        https://github.com/Azure/azure-sdk-for-python/issues/12568
    """
    if account_url is None:
        account_url = build_azure_storage_uri(account=account_name)
    container_client = ContainerClient(account_url=account_url,
                                       container_name=container_name,
                                       credential=account_key)
    container_client.create_container()

    permissions = ContainerSasPermissions(read=True, write=True, list=True)
    container_sas_token = generate_container_sas(
        account_name=account_name,
        container_name=container_name,
        account_key=account_key,
        permission=permissions,
        expiry=datetime.utcnow() + timedelta(hours=access_duration_hrs))

    return f'{account_url}/{container_name}?{container_sas_token}'


def upload_blob(container_uri: str, blob_name: str,
                data: Union[Iterable[AnyStr], IO[AnyStr]]) -> str:
    """Creates a new blob of the given name from an IO stream.

    Args:
        container_uri: str, URI to a container, may include SAS token
        blob_name: str, name of blob to upload
        data: str, bytes, or IO stream
            if str, assumes utf-8 encoding

    Returns: str, URL to blob, includes SAS token if container_uri has SAS token
    """
    blob_url = build_blob_uri(container_uri, blob_name)
    upload_blob_to_url(blob_url, data=data)
    return blob_url


def download_blob_to_stream(sas_uri: str) -> Tuple[io.BytesIO, BlobProperties]:
    """Downloads a blob to an IO stream.

    Args:
        sas_uri: str, URI to a blob

    Returns:
        output_stream: io.BytesIO, remember to close it when finished using
        blob_properties: BlobProperties

    Raises: azure.core.exceptions.ResourceNotFoundError, if sas_uri points
        to a non-existant blob
    """
    with BlobClient.from_blob_url(sas_uri) as blob_client:
        output_stream = io.BytesIO()
        blob_client.download_blob().readinto(output_stream)
        output_stream.seek(0)
        blob_properties = blob_client.get_blob_properties()
    return output_stream, blob_properties


def build_blob_uri(container_uri: str, blob_name: str) -> str:
    """
    Args:
        container_uri: str, URI to blob storage container
            <account_url>/<container>?<sas_token>
        blob_name: str, name of blob, not URL-escaped

    Returns: str, blob URI <account_url>/<container>/<blob_name>?<sas_token>,
        <blob_name> is URL-escaped
    """
    account_container = container_uri.split('?', maxsplit=1)[0]
    account_url, container_name = account_container.rsplit('/', maxsplit=1)
    sas_token = get_sas_token_from_uri(container_uri)

    blob_name = parse.quote(blob_name)
    blob_uri = f'{account_url}/{container_name}/{blob_name}'
    if sas_token is not None:
        blob_uri += f'?{sas_token}'
    return blob_uri