pyinfra/pyinfra/storage/storage_info.py

from dataclasses import dataclass

import requests
from azure.storage.blob import BlobServiceClient
from minio import Minio

from pyinfra.config import Config
from pyinfra.exception import UnknownStorageBackend
from pyinfra.storage.storages.azure import AzureStorage
from pyinfra.storage.storages.interface import Storage
from pyinfra.storage.storages.s3 import S3Storage
from pyinfra.utils.cipher import decrypt
from pyinfra.utils.url_parsing import validate_and_parse_s3_endpoint


@dataclass(frozen=True)
class StorageInfo:
    bucket_name: str


@dataclass(frozen=True)
class AzureStorageInfo(StorageInfo):
    connection_string: str

    def __hash__(self):
        return hash(self.connection_string)

    def __eq__(self, other):
        if not isinstance(other, AzureStorageInfo):
            return False
        return self.connection_string == other.connection_string


@dataclass(frozen=True)
class S3StorageInfo(StorageInfo):
    secure: bool
    endpoint: str
    access_key: str
    secret_key: str
    region: str

    def __hash__(self):
        return hash((self.secure, self.endpoint, self.access_key, self.secret_key, self.region))

    def __eq__(self, other):
        if not isinstance(other, S3StorageInfo):
            return False
        return (
            self.secure == other.secure
            and self.endpoint == other.endpoint
            and self.access_key == other.access_key
            and self.secret_key == other.secret_key
            and self.region == other.region
        )


def get_storage_from_storage_info(storage_info: StorageInfo) -> Storage:
    if isinstance(storage_info, AzureStorageInfo):
        return AzureStorage(BlobServiceClient.from_connection_string(conn_str=storage_info.connection_string))
    elif isinstance(storage_info, S3StorageInfo):
        return S3Storage(
            Minio(
                secure=storage_info.secure,
                endpoint=storage_info.endpoint,
                access_key=storage_info.access_key,
                secret_key=storage_info.secret_key,
                region=storage_info.region,
            )
        )
    else:
        raise UnknownStorageBackend()


def get_storage_info_from_endpoint(public_key: str, endpoint: str, x_tenant_id: str) -> StorageInfo:
    resp = requests.get(f"{endpoint}/{x_tenant_id}").json()

    maybe_azure = resp.get("azureStorageConnection")
    maybe_s3 = resp.get("s3StorageConnection")
    assert not (maybe_azure and maybe_s3)

    if maybe_azure:
        connection_string = decrypt(public_key, maybe_azure["connectionString"])
        storage_info = AzureStorageInfo(
            connection_string=connection_string,
            bucket_name=maybe_azure["containerName"],
        )
    elif maybe_s3:
        secure, endpoint = validate_and_parse_s3_endpoint(maybe_s3["endpoint"])
        secret = decrypt(public_key, maybe_s3["secret"])

        storage_info = S3StorageInfo(
            secure=secure,
            endpoint=endpoint,
            access_key=maybe_s3["key"],
            secret_key=secret,
            region=maybe_s3["region"],
            bucket_name=maybe_s3["bucketName"],
        )
    else:
        raise UnknownStorageBackend()

    return storage_info


def get_storage_info_from_config(config: Config) -> StorageInfo:
    if config.storage_backend == "s3":
        storage_info = S3StorageInfo(
            secure=config.storage_secure_connection,
            endpoint=config.storage_endpoint,
            access_key=config.storage_key,
            secret_key=config.storage_secret,
            region=config.storage_region,
            bucket_name=config.storage_bucket,
        )

    elif config.storage_backend == "azure":
        storage_info = AzureStorageInfo(
            connection_string=config.storage_azureconnectionstring,
            bucket_name=config.storage_bucket,
        )

    else:
        raise UnknownStorageBackend(f"Unknown storage backend '{config.storage_backend}'.")

    return storage_info