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
"""

#%% Imports

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)


#%% URI management

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).  Leading ? 
            is removed if present.
        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:
        if sas_token[0] == '?':
            sas_token = sas_token[1:]
        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 is_container_uri(sas_uri: str) -> bool:
    """Returns True if the signed resource field in the URI "sr" is a container "c"
    or a directory "d"
    """
    data = get_all_query_parts(sas_uri)
    if 'sr' not in data:
        return False

    if 'c' in data['sr'] or 'd' in data['sr']:
        return True
    else:
        return False

def is_blob_uri(sas_uri: str) -> bool:
    """Returns True if the signed resource field in the URI "sr" is a blob "b".
    """
    data = get_all_query_parts(sas_uri)
    if 'sr' not in data:
        return False

    if 'b' in data['sr']:
        return True
    else:
        return False


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: str, query part of the SAS token (without leading '?'),
        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.
    """
    data = get_all_query_parts(sas_uri)
    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)


#%% Blob 

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

    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. In these cases, create a
    BlobClient using the default constructor, instead of from_blob_url(),
    and use the BlobClient.exists() method directly.

    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 upload_blob(container_uri: str, blob_name: str,
                data: Union[Iterable[AnyStr], IO[AnyStr]],
                overwrite: bool = False) -> 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
        overwrite: bool, whether to overwrite existing blob (if any)

    Returns: str, URL to blob, includes SAS token if container_uri has SAS token
    """
    account_url, container, sas_token = split_container_uri(container_uri)
    with BlobClient(account_url=account_url, container_name=container,
                    blob_name=blob_name, credential=sas_token) as blob_client:
        blob_client.upload_blob(data, overwrite=overwrite)
        return blob_client.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_url, container, sas_token = split_container_uri(container_uri)

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


#%% Container

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,
        verbose: Optional[bool] = True
) -> 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 regex. Can also be a list of regexes, in which case
            blobs matching *any* of the regex's will be returned.
        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.
    """

    if verbose:
        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,disable=(not verbose))]
            i = len(list_blobs)
        else:
            i = 0
            for blob in tqdm(generator,disable=(not verbose)):
                i += 1
                suffix_ok = (blob_suffix is None
                             or blob.name.lower().endswith(blob_suffix))
                regex_ok = False
                if rsearch is None:
                    regex_ok = True
                else:
                    if not isinstance(rsearch, list):
                        rsearch = [rsearch]
                    # Check whether this blob name matches *any* of our regex's
                    for expr in rsearch:
                        if re.search(expr, blob.name) is not None:
                            regex_ok = True
                            break
                if suffix_ok and regex_ok:
                    list_blobs.append(blob.name)
                    if limit is not None and len(list_blobs) == limit:
                        break

    if verbose:
        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
    """
    if account_url is None:
        account_url = build_azure_storage_uri(account=account_name)
    with ContainerClient(account_url=account_url,
                         container_name=container_name,
                         credential=account_key) as container_client:
        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 split_container_uri(container_uri: str) -> Tuple[str, str, Optional[str]]:
    """
    Args:
        container_uri: str, URI to blob storage container
            <account_url>/<container>?<sas_token>

    Returns: account_url, container_name, sas_token
    """
    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)
    return account_url, container_name, sas_token