"""Real implementations of Azure clients (Vault and Storage) and their factories.""" from datetime import datetime, timedelta, timezone from mimetypes import add_type, guess_type from typing import IO, Self import structlog from azure.core.exceptions import ResourceNotFoundError from azure.identity.aio import ClientSecretCredential, DefaultAzureCredential from azure.keyvault.secrets.aio import SecretClient from azure.storage.blob import BlobSasPermissions, ContentSettings, StandardBlobTier, generate_blob_sas from azure.storage.blob.aio import BlobServiceClient from skyvern.config import settings from skyvern.forge.sdk.api.azure import ( AsyncAzureStorageClient, AsyncAzureVaultClient, AzureClientFactory, AzureUri, ) from skyvern.forge.sdk.schemas.organizations import AzureClientSecretCredential # Register custom mime types for mimetypes guessing add_type("application/json", ".har") add_type("text/plain", ".log") add_type("application/zstd", ".zst") LOG = structlog.get_logger() class RealAsyncAzureVaultClient(AsyncAzureVaultClient): """Real implementation of Azure Vault client using Azure SDK.""" def __init__(self, credential: ClientSecretCredential | DefaultAzureCredential) -> None: self.credential = credential async def __aenter__(self) -> Self: return self async def __aexit__( self, exc_type: type[BaseException] | None, exc_val: BaseException | None, exc_tb: object ) -> None: await self.credential.close() async def get_secret(self, secret_name: str, vault_name: str) -> str | None: secret_client = await self._get_secret_client(vault_name) try: secret = await secret_client.get_secret(secret_name) return secret.value except Exception as e: LOG.exception("Failed to get secret from Azure Key Vault.", secret_name=secret_name, error=e) return None finally: await secret_client.close() async def create_or_update_secret(self, secret_name: str, secret_value: str, vault_name: str) -> str: secret_client = await self._get_secret_client(vault_name) try: secret = await secret_client.set_secret(secret_name, secret_value) return secret.name except Exception as e: LOG.exception("Failed to create secret from Azure Key Vault.", secret_name=secret_name, error=e) raise e finally: await secret_client.close() async def delete_secret(self, secret_name: str, vault_name: str) -> str: secret_client = await self._get_secret_client(vault_name) try: secret = await secret_client.delete_secret(secret_name) return secret.name except Exception as e: LOG.exception("Failed to delete secret from Azure Key Vault.", secret_name=secret_name, error=e) raise e finally: await secret_client.close() async def _get_secret_client(self, vault_name: str) -> SecretClient: # Azure Key Vault URL format: https://.vault.azure.net # Assuming the secret_name is actually the Key Vault URL and the secret name # This needs to be clarified or passed as separate parameters # For now, let's assume secret_name is the actual secret name and Key Vault URL is in settings. key_vault_url = f"https://{vault_name}.vault.azure.net" # Placeholder, adjust as needed return SecretClient(vault_url=key_vault_url, credential=self.credential) async def close(self) -> None: await self.credential.close() class RealAsyncAzureStorageClient(AsyncAzureStorageClient): """Async client for Azure Blob Storage operations. Implements AsyncAzureStorageClient protocol.""" def __init__( self, account_name: str | None = None, account_key: str | None = None, ) -> None: self.account_name = account_name or settings.AZURE_STORAGE_ACCOUNT_NAME self.account_key = account_key or settings.AZURE_STORAGE_ACCOUNT_KEY if not self.account_name or not self.account_key: raise ValueError("Azure Storage account name and key are required") self._blob_service_client: BlobServiceClient | None = None self._verified_containers: set[str] = set() def _get_blob_service_client(self) -> BlobServiceClient: if self._blob_service_client is None: self._blob_service_client = BlobServiceClient( account_url=f"https://{self.account_name}.blob.core.windows.net", credential=self.account_key, ) return self._blob_service_client async def _ensure_container_exists(self, container: str) -> None: if container in self._verified_containers: return client = self._get_blob_service_client() container_client = client.get_container_client(container) try: if not await container_client.exists(): await container_client.create_container() LOG.info("Created Azure container", container=container) except Exception: LOG.debug("Container may already exist", container=container) self._verified_containers.add(container) async def upload_file( self, uri: str, data: bytes, tier: StandardBlobTier = StandardBlobTier.HOT, tags: dict[str, str] | None = None, ) -> None: parsed = AzureUri(uri) await self._ensure_container_exists(parsed.container) client = self._get_blob_service_client() container_client = client.get_container_client(parsed.container) await container_client.upload_blob( name=parsed.blob_path, data=data, overwrite=True, standard_blob_tier=tier, tags=tags, ) async def upload_file_from_path( self, uri: str, file_path: str, tier: StandardBlobTier = StandardBlobTier.HOT, tags: dict[str, str] | None = None, metadata: dict[str, str] | None = None, ) -> None: parsed = AzureUri(uri) await self._ensure_container_exists(parsed.container) client = self._get_blob_service_client() container_client = client.get_container_client(parsed.container) content_type, _ = guess_type(file_path) content_settings = ContentSettings(content_type=content_type) if content_type else None with open(file_path, "rb") as f: await container_client.upload_blob( name=parsed.blob_path, data=f, overwrite=True, standard_blob_tier=tier, tags=tags, metadata=metadata, content_settings=content_settings, ) async def upload_file_stream( self, uri: str, file_obj: IO[bytes], tier: StandardBlobTier = StandardBlobTier.HOT, tags: dict[str, str] | None = None, metadata: dict[str, str] | None = None, ) -> str: parsed = AzureUri(uri) await self._ensure_container_exists(parsed.container) client = self._get_blob_service_client() container_client = client.get_container_client(parsed.container) await container_client.upload_blob( name=parsed.blob_path, data=file_obj, overwrite=True, standard_blob_tier=tier, tags=tags, metadata=metadata, ) return uri async def download_file(self, uri: str, log_exception: bool = True) -> bytes | None: parsed = AzureUri(uri) try: client = self._get_blob_service_client() container_client = client.get_container_client(parsed.container) blob_client = container_client.get_blob_client(parsed.blob_path) download = await blob_client.download_blob() return await download.readall() except ResourceNotFoundError: if log_exception: LOG.warning("Azure blob not found", uri=uri) return None except Exception: if log_exception: LOG.exception("Failed to download from Azure", uri=uri) return None async def get_blob_properties(self, uri: str) -> dict | None: parsed = AzureUri(uri) try: client = self._get_blob_service_client() container_client = client.get_container_client(parsed.container) blob_client = container_client.get_blob_client(parsed.blob_path) props = await blob_client.get_blob_properties() return { "size": props.size, "content_type": props.content_settings.content_type if props.content_settings else None, "last_modified": props.last_modified, "etag": props.etag, "metadata": props.metadata, } except ResourceNotFoundError: return None except Exception: LOG.exception("Failed to get blob properties", uri=uri) return None async def blob_exists(self, uri: str) -> bool: parsed = AzureUri(uri) try: client = self._get_blob_service_client() container_client = client.get_container_client(parsed.container) blob_client = container_client.get_blob_client(parsed.blob_path) return await blob_client.exists() except Exception: return False async def delete_blob(self, uri: str) -> None: parsed = AzureUri(uri) try: client = self._get_blob_service_client() container_client = client.get_container_client(parsed.container) blob_client = container_client.get_blob_client(parsed.blob_path) await blob_client.delete_blob() except ResourceNotFoundError: LOG.debug("Azure blob not found for deletion", uri=uri) except Exception: LOG.exception("Failed to delete Azure blob", uri=uri) raise async def list_blobs(self, container: str, prefix: str | None = None) -> list[str]: try: client = self._get_blob_service_client() container_client = client.get_container_client(container) blobs = [] async for blob in container_client.list_blobs(name_starts_with=prefix): blobs.append(blob.name) return blobs except ResourceNotFoundError: return [] except Exception: LOG.exception("Failed to list Azure blobs", container=container, prefix=prefix) return [] def create_sas_url(self, uri: str, expiry_hours: int = 24) -> str | None: parsed = AzureUri(uri) try: sas_token = generate_blob_sas( account_name=self.account_name, container_name=parsed.container, blob_name=parsed.blob_path, account_key=self.account_key, permission=BlobSasPermissions(read=True), expiry=datetime.now(timezone.utc) + timedelta(hours=expiry_hours), ) return ( f"https://{self.account_name}.blob.core.windows.net/{parsed.container}/{parsed.blob_path}?{sas_token}" ) except Exception: LOG.exception("Failed to create SAS URL", uri=uri) return None async def create_sas_urls(self, uris: list[str], expiry_hours: int = 24) -> list[str] | None: try: sas_urls: list[str] = [] for uri in uris: url = self.create_sas_url(uri, expiry_hours) if url is None: LOG.warning("SAS URL generation failed, aborting batch", failed_uri=uri, uris=uris) return None sas_urls.append(url) return sas_urls except Exception: LOG.exception("Failed to create SAS URLs") return None async def close(self) -> None: if self._blob_service_client: await self._blob_service_client.close() self._blob_service_client = None async def list_files(self, uri: str) -> list[str]: """List files under a URI prefix. Returns blob names relative to container.""" parsed = AzureUri(uri) return await self.list_blobs(parsed.container, parsed.blob_path) async def get_object_info(self, uri: str) -> dict | None: """Get object info including metadata. Returns dict with Metadata and LastModified keys.""" props = await self.get_blob_properties(uri) if props is None: return None return { "Metadata": props.get("metadata", {}), "LastModified": props.get("last_modified"), } async def delete_file(self, uri: str) -> None: """Delete a file at the given URI.""" await self.delete_blob(uri) async def get_file_metadata(self, uri: str, log_exception: bool = True) -> dict[str, str] | None: """Get only the metadata for a file.""" parsed = AzureUri(uri) try: client = self._get_blob_service_client() container_client = client.get_container_client(parsed.container) blob_client = container_client.get_blob_client(parsed.blob_path) props = await blob_client.get_blob_properties() return props.metadata or {} except ResourceNotFoundError: if log_exception: LOG.warning("Azure blob not found for metadata", uri=uri) return None except Exception: if log_exception: LOG.exception("Failed to get blob metadata", uri=uri) return None class RealAzureClientFactory(AzureClientFactory): """Factory for creating real Azure Vault and Storage clients.""" def create_default(self) -> AsyncAzureVaultClient: """Create an Azure Vault client using DefaultAzureCredential.""" return RealAsyncAzureVaultClient(DefaultAzureCredential()) def create_from_client_secret(self, credential: AzureClientSecretCredential) -> AsyncAzureVaultClient: """Create an Azure Vault client using client secret credentials.""" cred = ClientSecretCredential( tenant_id=credential.tenant_id, client_id=credential.client_id, client_secret=credential.client_secret, ) return RealAsyncAzureVaultClient(cred) def create_storage_client(self, storage_account_name: str, storage_account_key: str) -> AsyncAzureStorageClient: """Create an Azure Storage client with the provided credentials.""" return RealAsyncAzureStorageClient(account_name=storage_account_name, account_key=storage_account_key)