llm/Lib/site-packages/cassio/config/bundle_download.py

"""
Facilities to manage the download of a secure-connect-bundle from an Astra DB
token.
"""
import logging
from typing import Optional

import requests

DEFAULT_GET_BUNDLE_URL_TEMPLATE = (
    "https://api.astra.datastax.com/v2/databases/{database_id}/secureBundleURL"
)


logger = logging.getLogger(__name__)


def get_astra_bundle_url(
    database_id: str, token: str, bundle_url_template: Optional[str] = None
) -> str:
    """
    Given a database ID and a token, provide a (temporarily-valid)
    URL for downloading the secure-connect-bundle.

    Courtesy of @phact (thank you!).
    """
    if not bundle_url_template:
        bundle_url_template = DEFAULT_GET_BUNDLE_URL_TEMPLATE
    url = bundle_url_template.format(database_id=database_id)
    headers = {"Authorization": f"Bearer {token}", "Content-Type": "application/json"}

    logger.debug(f"Obtaining SCB URL from: {url}")
    response = requests.post(url, headers=headers)
    logger.debug(f"SCB URL response: {response.text}")

    response_json = response.json()
    # Print the response
    if response_json is not None:
        if "errors" in response_json:
            # Handle the errors
            errors = response_json["errors"]
            if any(
                "jwt not valid" in (err.get("message") or "").lower() for err in errors
            ):
                raise ValueError("Invalid or insufficient token.")
            elif any(
                "malformed" in (err.get("message") or "").lower() for err in errors
            ):
                raise ValueError("Invalid or insufficient token.")
            else:
                raise ValueError(
                    "Generic error when fetching the URL to the secure-bundle."
                )
        else:
            return str(response_json["downloadURL"])
    else:
        raise ValueError(
            "Cannot get the secure-bundle URL. "
            "Check the provided database_id and token."
        )


def download_astra_bundle_url(
    database_id: str,
    token: str,
    out_file_path: str,
    bundle_url_template: Optional[str] = None,
) -> None:
    """
    Obtain the secure-connect-bundle and save it to a specified file.
    """
    bundle_url = get_astra_bundle_url(
        database_id=database_id, token=token, bundle_url_template=bundle_url_template
    )
    logger.debug(f"Downloading SCB from: {bundle_url}")
    bundle_data = requests.get(bundle_url)
    with open(out_file_path, "wb") as f:
        f.write(bundle_data.content)
    return