From d4730153b2886c631ec823bf39967336fe9a4655 Mon Sep 17 00:00:00 2001 From: Prashant Date: Thu, 25 Jun 2026 11:57:40 +0530 Subject: [PATCH] feat: add crypto (KMS / Key Vault) backend and SQS exclude_env_credentials; release v0.2.6 - New cloudrift.crypto module: get_crypto for AWS KMS and Azure Key Vault keys (encrypt/decrypt + base64 string helpers), with CryptoError hierarchy. - SQS from_iam_role gains exclude_env_credentials to drop the env-var credential provider so ambient process credentials can't shadow the instance/task role. - Includes pending storage and messaging (azure_bus / base) enhancements. - Bump version to 0.2.6, add azure-keyvault-keys extra, re-lock. --- CLAUDE.md | 6 + README.md | 32 +++- cloudrift/__init__.py | 4 +- cloudrift/core/exceptions.py | 13 ++ cloudrift/crypto/__init__.py | 50 ++++++ cloudrift/crypto/aws_kms.py | 166 ++++++++++++++++++++ cloudrift/crypto/azure_keyvault_keys.py | 133 ++++++++++++++++ cloudrift/crypto/base.py | 59 +++++++ cloudrift/messaging/__init__.py | 19 ++- cloudrift/messaging/azure_bus.py | 62 ++++---- cloudrift/messaging/base.py | 74 ++++++++- cloudrift/messaging/sqs.py | 130 ++++++++++++---- cloudrift/storage/__init__.py | 5 + cloudrift/storage/s3.py | 60 +++++++- pyproject.toml | 9 +- tests/test_crypto.py | 127 +++++++++++++++ tests/test_messaging.py | 197 +++++++++++++++++++----- tests/test_messaging_azure.py | 84 +++++++--- tests/test_storage.py | 60 ++++++++ uv.lock | 25 ++- 20 files changed, 1171 insertions(+), 144 deletions(-) create mode 100644 cloudrift/crypto/__init__.py create mode 100644 cloudrift/crypto/aws_kms.py create mode 100644 cloudrift/crypto/azure_keyvault_keys.py create mode 100644 cloudrift/crypto/base.py create mode 100644 tests/test_crypto.py diff --git a/CLAUDE.md b/CLAUDE.md index 747418d..629bbfd 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -42,6 +42,12 @@ The document category deliberately has **no `base.py` ABC and no backend wrapper - **Cache** uses an explicit auth-method argument: `get_cache(provider, auth_method, **kwargs)` where `auth_method` is the literal `from_*` method name (e.g. `get_cache("redis", "from_url", url=...)`). - **All five other categories** infer the constructor from **which credential keys are present** in `**kwargs`: `get_storage(provider, **kwargs)` calls `from_access_key` if `aws_access_key_id` is present, `from_connection_string` if `connection_string` is present, and falls through to the managed-identity/IAM-role default. When adding an auth method here, add both the constructor (a `from_*` classmethod, or a `connect_*` function for document) and a routing branch in the factory — for document, in **both** `get_mongodb` and `get_mongodb_sync`. + AWS `from_assume_role` (STS AssumeRole, cross-account) is wired for **SQS** (`get_queue`) and **S3** (`get_storage` + `get_storage_client`): `role_arn` present → assume-role path, checked **before** `aws_access_key_id`. It does a synchronous `boto3` `sts:AssumeRole` (optional `external_id` → `ExternalId`) and threads the temporary creds into the `aioboto3.Session`. Temp creds are not auto-refreshed — construct a fresh backend when the session expires. + +### Messaging payload contract (byte-primitive) + +The messaging primitive is **raw bytes + a flat `attributes` map** (string → string), not JSON. `MessagingBackend.send(body: bytes, attributes=None, delay=0, *, group_id, dedup_id)` and `send_batch(messages: list[OutgoingMessage], ...)` where `OutgoingMessage` (in `base.py`) is `{body: bytes, attributes: dict[str,str] | None}`. Attributes map to SQS `MessageAttributes` (String type) and Service Bus `application_properties`. Received `Message.body` is `bytes`; `Message.attributes` is the stringified provider attributes. JSON convenience: the module-level `send_json(backend, message: dict, attributes=None, delay=0, ...)` helper (works for any backend — `json.dumps(...).encode()` then `send`) plus `Message.json()` (decodes the body with `json.loads`). Receipt-handle / ack / FIFO semantics are unchanged. + ### Redis cache specifics The three Redis backends (`redis`, `elasticache`, `azure_redis`) share **all** their operation logic through `_RedisMixin` in `cache/base.py` — the per-provider modules contain *only* the `from_*` constructors that build the `aioredis.Redis` client with the right auth (URL, IAM SigV4 auto-refresh, access key, managed identity). A new Redis command is implemented **once** in `_RedisMixin` and added as an `@abstractmethod` on `CacheBackend`; never reimplement it per provider. `CacheBackend.pipeline()` ships a no-atomicity `_SequentialPipeline` fallback that the Redis mixin overrides with a real server-side transactional pipeline. diff --git a/README.md b/README.md index 107094b..43fdb87 100644 --- a/README.md +++ b/README.md @@ -98,6 +98,9 @@ s3 = get_storage("s3", bucket="b", region="us-east-1") # I s3 = get_storage("s3", bucket="b", aws_access_key_id="...", # static keys aws_secret_access_key="...", region="us-east-1") s3 = get_storage("s3", bucket="b", profile_name="dev") # ~/.aws/credentials +s3 = get_storage("s3", bucket="b", # STS AssumeRole + role_arn="arn:aws:iam::123456789012:role/cross", + external_id="my-external-id", region="us-east-1") # Azure Blob blob = get_storage("azure_blob", connection_string="...", container="c") @@ -132,21 +135,38 @@ from cloudrift.messaging import get_queue sqs = get_queue("sqs", queue_url="https://sqs.us-east-1.amazonaws.com/.../q", region="us-east-1") +# AWS SQS via STS AssumeRole (cross-account) +sqs = get_queue("sqs", queue_url="https://sqs.../q", + role_arn="arn:aws:iam::123456789012:role/cross-account", + external_id="my-external-id", region="us-east-1") + # Azure Service Bus bus = get_queue("azure_bus", connection_string="...", queue_name="my-queue") bus = get_queue("azure_bus", fully_qualified_namespace="ns.servicebus.windows.net", queue_name="my-queue") # managed identity ``` -**Operations**: +**Operations** — the payload primitive is **raw bytes** plus an optional flat +`attributes` map (string → string). Attributes map to SQS `MessageAttributes` +(String type) / Service Bus `application_properties`: ```python -msg_id = await queue.send({"action": "process", "id": 42}, delay=0) -ids = await queue.send_batch([{"n": 1}, {"n": 2}]) +from cloudrift.messaging import OutgoingMessage, send_json + +# Raw bytes + attributes +msg_id = await queue.send(b"raw payload", attributes={"content_type": "text/plain"}) +ids = await queue.send_batch([ + OutgoingMessage(body=b"a", attributes={"k": "1"}), + OutgoingMessage(body=b"b"), +]) + +# JSON convenience: send_json() dumps+encodes, m.json() loads the bytes body +await send_json(queue, {"action": "process", "id": 42}, attributes={"v": "1"}) messages = await queue.receive(max_messages=10, wait_time=20) # long-poll for m in messages: - handle_job(m.body) + handle_job(m.json()) # or m.body for the raw bytes + print(m.attributes) # str -> str map await queue.delete(m.receipt_handle) # ack # or: await queue.nack(m.receipt_handle) # return for immediate redelivery @@ -168,13 +188,13 @@ queues, pass `group_id` (ordering key) and `dedup_id` (deduplication key): # SQS FIFO — group_id is required, dedup_id optional if the queue has # content-based deduplication enabled fifo = get_queue("sqs", queue_url="https://sqs.../jobs.fifo", region="us-east-1") -await fifo.send({"task": "extract"}, group_id="owner-123", dedup_id="evt-abc") +await send_json(fifo, {"task": "extract"}, group_id="owner-123", dedup_id="evt-abc") # Azure Service Bus — queue must be created with sessions enabled; # pass session_enabled=True so the backend uses session receivers bus = get_queue("azure_bus", connection_string="...", queue_name="jobs", session_enabled=True) -await bus.send({"task": "extract"}, group_id="owner-123", dedup_id="evt-abc") +await send_json(bus, {"task": "extract"}, group_id="owner-123", dedup_id="evt-abc") messages = await fifo.receive(max_messages=10, wait_time=20, visibility_timeout=300) for m in messages: diff --git a/cloudrift/__init__.py b/cloudrift/__init__.py index 851f639..d20879c 100644 --- a/cloudrift/__init__.py +++ b/cloudrift/__init__.py @@ -4,10 +4,11 @@ from cloudrift.cache import get_cache, cache_broker_url from cloudrift.secrets import get_secrets from cloudrift.sql import get_sql +from cloudrift.crypto import get_crypto from cloudrift.pubsub import get_pubsub from cloudrift.email import get_email -__version__ = "0.2.5" +__version__ = "0.2.6" __all__ = [ "get_storage", "get_queue", @@ -17,6 +18,7 @@ "cache_broker_url", "get_secrets", "get_sql", + "get_crypto", "get_pubsub", "get_email", ] diff --git a/cloudrift/core/exceptions.py b/cloudrift/core/exceptions.py index 27ff91b..6a3b7b8 100644 --- a/cloudrift/core/exceptions.py +++ b/cloudrift/core/exceptions.py @@ -78,6 +78,19 @@ class SQLAuthError(SQLError): Azure AD access token or RDS IAM token could not be obtained).""" +# Crypto (KMS / Key Vault keys) exceptions +class CryptoError(CloudRiftError): + """Base exception for cryptographic (KMS / Key Vault) operations.""" + + +class CryptoKeyNotFoundError(CryptoError): + """Raised when the referenced key does not exist or is unavailable.""" + + +class CryptoPermissionError(CryptoError): + """Raised on key access / crypto-operation permission failures.""" + + # Pub/Sub exceptions class PubSubError(CloudRiftError): """Base exception for pub/sub operations.""" diff --git a/cloudrift/crypto/__init__.py b/cloudrift/crypto/__init__.py new file mode 100644 index 0000000..31fd22d --- /dev/null +++ b/cloudrift/crypto/__init__.py @@ -0,0 +1,50 @@ +from cloudrift.crypto.base import CryptoBackend + + +def get_crypto(provider: str, **kwargs) -> CryptoBackend: + """Factory to instantiate a key-management crypto backend. + + The cross-cloud analog of AWS KMS: encrypt/decrypt small payloads against a + managed key. AWS KMS ↔ Azure Key Vault keys. + + Args: + provider: ``"aws_kms"`` or ``"azure_keyvault"``. + **kwargs: Provider-specific config. The AWS factory routes to the + appropriate ``from_*`` classmethod based on which credential keys are + present; the Azure factory routes by whether ``client_secret`` is set. + + Returns: + A :class:`CryptoBackend` instance. + + Examples: + get_crypto("aws_kms", key_id="arn:aws:kms:...:key/abc", region="us-east-1") # IAM/env + get_crypto("aws_kms", key_id="alias/my-key", aws_access_key_id="AKIA...", + aws_secret_access_key="...", region="us-east-1") + get_crypto("azure_keyvault", + key_id="https://myvault.vault.azure.net/keys/mykey", + tenant_id="...", client_id="...", client_secret="...") + get_crypto("azure_keyvault", + key_id="https://myvault.vault.azure.net/keys/mykey") # managed identity + """ + if provider in ("aws_kms", "kms"): + from cloudrift.crypto.aws_kms import AWSKMSBackend + + if "aws_access_key_id" in kwargs: + return AWSKMSBackend.from_access_key(**kwargs) + if "profile_name" in kwargs: + return AWSKMSBackend.from_profile(**kwargs) + return AWSKMSBackend.from_iam_role(**kwargs) + + if provider in ("azure_keyvault", "azure_keyvault_keys"): + from cloudrift.crypto.azure_keyvault_keys import AzureKeyVaultKeysBackend + + if "client_secret" in kwargs: + return AzureKeyVaultKeysBackend.from_service_principal(**kwargs) + return AzureKeyVaultKeysBackend.from_managed_identity(**kwargs) + + raise ValueError( + f"Unknown crypto provider: {provider!r}. Choose 'aws_kms' or 'azure_keyvault'." + ) + + +__all__ = ["CryptoBackend", "get_crypto"] diff --git a/cloudrift/crypto/aws_kms.py b/cloudrift/crypto/aws_kms.py new file mode 100644 index 0000000..4492e73 --- /dev/null +++ b/cloudrift/crypto/aws_kms.py @@ -0,0 +1,166 @@ +import asyncio + +import aioboto3 +from botocore.config import Config +from botocore.exceptions import ClientError + +from cloudrift.core.exceptions import ( + CryptoError, + CryptoKeyNotFoundError, + CryptoPermissionError, +) +from cloudrift.crypto.base import CryptoBackend + + +class AWSKMSBackend(CryptoBackend): + """AWS KMS crypto backend (native async via ``aioboto3``). + + Encrypts/decrypts directly against a symmetric KMS key — the same + ``Encrypt`` / ``Decrypt`` calls the AWS SDK makes, so ciphertext is + interchangeable with anything encrypted by raw boto3 against the same key. + A single async client is created lazily and reused. + + Construct via: + - ``from_access_key`` — static credentials (+ optional session token) + - ``from_iam_role`` — instance profile / environment / ECS task role + - ``from_profile`` — named profile from ``~/.aws/credentials`` + + ``key_id`` (key id, ARN, or alias) is required to encrypt. It is optional for + decrypt-only backends — KMS resolves the key from the ciphertext. + """ + + def __init__( + self, + session: aioboto3.Session, + key_id: str | None = None, + *, + encryption_context: dict | None = None, + endpoint_url: str | None = None, + max_pool_connections: int = 25, + connect_timeout: float = 10.0, + read_timeout: float = 30.0, + client_kwargs: dict | None = None, + ) -> None: + self._session = session + self._key_id = key_id + self._encryption_context = encryption_context or {} + self._endpoint_url = endpoint_url + self._config = Config( + max_pool_connections=max_pool_connections, + connect_timeout=connect_timeout, + read_timeout=read_timeout, + ) + self._client_kwargs = client_kwargs or {} + self._client_cm = None + self._client = None + self._lock = asyncio.Lock() + + # ------------------------------------------------------------------ + # Factory constructors + # ------------------------------------------------------------------ + + @classmethod + def from_access_key( + cls, + aws_access_key_id: str, + aws_secret_access_key: str, + key_id: str | None = None, + region: str = "us-east-1", + aws_session_token: str | None = None, + **kwargs, + ) -> "AWSKMSBackend": + """Authenticate with explicit access key / secret.""" + session = aioboto3.Session( + aws_access_key_id=aws_access_key_id, + aws_secret_access_key=aws_secret_access_key, + aws_session_token=aws_session_token, + region_name=region, + ) + return cls(session, key_id, **kwargs) + + @classmethod + def from_iam_role( + cls, + key_id: str | None = None, + region: str = "us-east-1", + **kwargs, + ) -> "AWSKMSBackend": + """Authenticate via IAM role / instance profile / environment variables.""" + session = aioboto3.Session(region_name=region) + return cls(session, key_id, **kwargs) + + @classmethod + def from_profile( + cls, + profile_name: str, + key_id: str | None = None, + region: str = "us-east-1", + **kwargs, + ) -> "AWSKMSBackend": + """Authenticate using a named profile from ``~/.aws/credentials``.""" + session = aioboto3.Session(profile_name=profile_name, region_name=region) + return cls(session, key_id, **kwargs) + + # ------------------------------------------------------------------ + # Internal lifecycle + # ------------------------------------------------------------------ + + async def _ensure(self): + if self._client is not None: + return self._client + async with self._lock: + if self._client is None: + self._client_cm = self._session.client( + "kms", + endpoint_url=self._endpoint_url, + config=self._config, + **self._client_kwargs, + ) + try: + self._client = await self._client_cm.__aenter__() + except Exception: + self._client_cm = None + raise + return self._client + + async def close(self) -> None: + client_cm, self._client_cm = self._client_cm, None + self._client = None + if client_cm is not None: + await client_cm.__aexit__(None, None, None) + + # ------------------------------------------------------------------ + # CryptoBackend implementation + # ------------------------------------------------------------------ + + async def encrypt(self, plaintext: bytes) -> bytes: + if not self._key_id: + raise CryptoError("AWS KMS key_id is required to encrypt") + client = await self._ensure() + try: + kwargs: dict = {"KeyId": self._key_id, "Plaintext": plaintext} + if self._encryption_context: + kwargs["EncryptionContext"] = self._encryption_context + response = await client.encrypt(**kwargs) + return response["CiphertextBlob"] + except ClientError as e: + self._raise(e) + + async def decrypt(self, ciphertext: bytes) -> bytes: + client = await self._ensure() + try: + kwargs: dict = {"CiphertextBlob": ciphertext} + if self._encryption_context: + kwargs["EncryptionContext"] = self._encryption_context + response = await client.decrypt(**kwargs) + return response["Plaintext"] + except ClientError as e: + self._raise(e) + + def _raise(self, exc: ClientError): + code = exc.response["Error"]["Code"] + if code in ("NotFoundException", "KeyUnavailableException"): + raise CryptoKeyNotFoundError(str(exc)) from exc + if code in ("AccessDeniedException", "UnauthorizedAccess"): + raise CryptoPermissionError(str(exc)) from exc + raise CryptoError(str(exc)) from exc diff --git a/cloudrift/crypto/azure_keyvault_keys.py b/cloudrift/crypto/azure_keyvault_keys.py new file mode 100644 index 0000000..7d79007 --- /dev/null +++ b/cloudrift/crypto/azure_keyvault_keys.py @@ -0,0 +1,133 @@ +import asyncio + +from azure.core.exceptions import ( + ClientAuthenticationError, + ResourceNotFoundError, +) +from azure.identity.aio import ( + ClientSecretCredential, + DefaultAzureCredential, + ManagedIdentityCredential, +) +from azure.keyvault.keys.crypto import EncryptionAlgorithm +from azure.keyvault.keys.crypto.aio import CryptographyClient + +from cloudrift.core.exceptions import ( + CryptoError, + CryptoKeyNotFoundError, + CryptoPermissionError, +) +from cloudrift.crypto.base import CryptoBackend + + +class AzureKeyVaultKeysBackend(CryptoBackend): + """Azure Key Vault *keys* crypto backend — the analog of AWS KMS. + + Encrypts/decrypts against a Key Vault key via ``CryptographyClient``. + ``key_id`` is the full key identifier URL, e.g. + ``https://myvault.vault.azure.net/keys/mykey`` (or pinned to a version + ``.../keys/mykey/``). + + The default algorithm is ``RSA-OAEP-256`` (RSA keys). RSA encryption has a + small payload ceiling (~190 bytes for RSA-2048); pass ``algorithm=`` for a + different key type, or wrap a data key for larger payloads. + + Construct via: + - ``from_service_principal`` — tenant_id / client_id / client_secret + - ``from_managed_identity`` — managed identity (optionally a client_id), + else DefaultAzureCredential + """ + + def __init__( + self, + key_id: str, + credential, + *, + algorithm: "EncryptionAlgorithm | None" = None, + ) -> None: + self._key_id = key_id + self._credential = credential + self._algorithm = algorithm or EncryptionAlgorithm.rsa_oaep_256 + self._client: CryptographyClient | None = None + self._lock = asyncio.Lock() + + # ------------------------------------------------------------------ + # Factory constructors + # ------------------------------------------------------------------ + + @classmethod + def from_service_principal( + cls, + key_id: str, + tenant_id: str, + client_id: str, + client_secret: str, + **kwargs, + ) -> "AzureKeyVaultKeysBackend": + """Authenticate with an Azure AD service principal.""" + credential = ClientSecretCredential( + tenant_id=tenant_id, + client_id=client_id, + client_secret=client_secret, + ) + return cls(key_id, credential, **kwargs) + + @classmethod + def from_managed_identity( + cls, + key_id: str, + client_id: str | None = None, + **kwargs, + ) -> "AzureKeyVaultKeysBackend": + """Authenticate via managed identity (or DefaultAzureCredential).""" + credential = ( + ManagedIdentityCredential(client_id=client_id) + if client_id + else DefaultAzureCredential() + ) + return cls(key_id, credential, **kwargs) + + # ------------------------------------------------------------------ + # Internal lifecycle + # ------------------------------------------------------------------ + + async def _ensure(self) -> CryptographyClient: + if self._client is None: + async with self._lock: + if self._client is None: + self._client = CryptographyClient(self._key_id, self._credential) + return self._client + + async def close(self) -> None: + if self._client is not None: + await self._client.close() + self._client = None + if self._credential is not None: + await self._credential.close() + + # ------------------------------------------------------------------ + # CryptoBackend implementation + # ------------------------------------------------------------------ + + async def encrypt(self, plaintext: bytes) -> bytes: + client = await self._ensure() + try: + result = await client.encrypt(self._algorithm, plaintext) + return result.ciphertext + except Exception as e: + self._raise(e) + + async def decrypt(self, ciphertext: bytes) -> bytes: + client = await self._ensure() + try: + result = await client.decrypt(self._algorithm, ciphertext) + return result.plaintext + except Exception as e: + self._raise(e) + + def _raise(self, exc: Exception): + if isinstance(exc, ResourceNotFoundError): + raise CryptoKeyNotFoundError(str(exc)) from exc + if isinstance(exc, ClientAuthenticationError): + raise CryptoPermissionError(str(exc)) from exc + raise CryptoError(str(exc)) from exc diff --git a/cloudrift/crypto/base.py b/cloudrift/crypto/base.py new file mode 100644 index 0000000..df5c8cc --- /dev/null +++ b/cloudrift/crypto/base.py @@ -0,0 +1,59 @@ +import base64 +from abc import ABC, abstractmethod + + +class CryptoBackend(ABC): + """Abstract base class for cloud key-management crypto backends. + + Provides provider-agnostic encrypt/decrypt against a *managed key* — AWS KMS + or an Azure Key Vault key. Subclasses implement the raw byte operations + (:meth:`encrypt` / :meth:`decrypt`); this base adds base64 string helpers for + the common "encrypt a token, store the ciphertext as text" case. + + Ciphertext is the provider's native format (it is NOT re-wrapped by + cloudrift), so values encrypted by the equivalent native SDK call remain + decryptable through this backend and vice-versa. + + Payload size limits are provider-specific: + - AWS KMS symmetric keys accept up to 4 KB of plaintext. + - Azure Key Vault RSA keys (RSA-OAEP-256) accept far less — roughly + 190 bytes for RSA-2048, ~446 bytes for RSA-4096. + For larger payloads, use this backend to wrap a random data key and encrypt + the payload yourself (envelope encryption). + + Backends hold long-lived async clients. Use ``await backend.close()`` (or + ``async with backend:``) to release them cleanly. + """ + + @abstractmethod + async def encrypt(self, plaintext: bytes) -> bytes: + """Encrypt raw bytes with the managed key; return raw ciphertext bytes.""" + + @abstractmethod + async def decrypt(self, ciphertext: bytes) -> bytes: + """Decrypt raw ciphertext bytes; return the original plaintext bytes.""" + + async def encrypt_str(self, plaintext: str) -> str: + """Encrypt a UTF-8 string, returning base64-encoded ciphertext. + + An empty/None input returns ``""`` (no crypto call), mirroring the common + "encrypt this optional token" pattern. + """ + if not plaintext: + return "" + return base64.b64encode(await self.encrypt(plaintext.encode("utf-8"))).decode("ascii") + + async def decrypt_str(self, ciphertext_b64: str) -> str: + """Decrypt base64 ciphertext produced by :meth:`encrypt_str` to a string.""" + if not ciphertext_b64: + return "" + return (await self.decrypt(base64.b64decode(ciphertext_b64))).decode("utf-8") + + async def close(self) -> None: + """Close the underlying client and release sockets. Default is a no-op.""" + + async def __aenter__(self) -> "CryptoBackend": + return self + + async def __aexit__(self, exc_type, exc, tb) -> None: + await self.close() diff --git a/cloudrift/messaging/__init__.py b/cloudrift/messaging/__init__.py index 1184d48..2c7df2f 100644 --- a/cloudrift/messaging/__init__.py +++ b/cloudrift/messaging/__init__.py @@ -1,4 +1,9 @@ -from cloudrift.messaging.base import Message, MessagingBackend +from cloudrift.messaging.base import ( + Message, + MessagingBackend, + OutgoingMessage, + send_json, +) def get_queue(provider: str, **kwargs) -> MessagingBackend: @@ -14,7 +19,9 @@ def get_queue(provider: str, **kwargs) -> MessagingBackend: Examples: get_queue("sqs", queue_url="https://sqs...", region="us-east-1") + get_queue("sqs", queue_url="...", region="us-east-1", exclude_env_credentials=True) get_queue("sqs", queue_url="...", aws_access_key_id="...", aws_secret_access_key="...") + get_queue("sqs", queue_url="...", role_arn="arn:aws:iam::...:role/x", external_id="...") get_queue("azure_bus", connection_string="...", queue_name="my-queue") get_queue("azure_bus", fully_qualified_namespace="ns.servicebus.windows.net", queue_name="q", client_id="...", client_secret="...", tenant_id="...") @@ -22,6 +29,8 @@ def get_queue(provider: str, **kwargs) -> MessagingBackend: if provider == "sqs": from cloudrift.messaging.sqs import AWSSQSBackend + if "role_arn" in kwargs: + return AWSSQSBackend.from_assume_role(**kwargs) if "aws_access_key_id" in kwargs: return AWSSQSBackend.from_access_key(**kwargs) if "profile_name" in kwargs: @@ -40,4 +49,10 @@ def get_queue(provider: str, **kwargs) -> MessagingBackend: raise ValueError(f"Unknown messaging provider: {provider!r}. Choose 'sqs' or 'azure_bus'.") -__all__ = ["Message", "MessagingBackend", "get_queue"] +__all__ = [ + "Message", + "MessagingBackend", + "OutgoingMessage", + "get_queue", + "send_json", +] diff --git a/cloudrift/messaging/azure_bus.py b/cloudrift/messaging/azure_bus.py index 4920160..2709574 100644 --- a/cloudrift/messaging/azure_bus.py +++ b/cloudrift/messaging/azure_bus.py @@ -1,5 +1,4 @@ import asyncio -import json from azure.core.exceptions import HttpResponseError, ResourceNotFoundError from azure.servicebus import NEXT_AVAILABLE_SESSION, ServiceBusMessage @@ -12,7 +11,7 @@ MessagingError, QueueNotFoundError, ) -from cloudrift.messaging.base import Message, MessagingBackend +from cloudrift.messaging.base import Message, MessagingBackend, OutgoingMessage class AzureServiceBusBackend(MessagingBackend): @@ -135,13 +134,9 @@ async def _ensure(self) -> ServiceBusClient: async with self._lock: if self._client is None: if self._connection_string: - self._client = ServiceBusClient.from_connection_string( - self._connection_string - ) + self._client = ServiceBusClient.from_connection_string(self._connection_string) else: - self._client = ServiceBusClient( - self._namespace, credential=self._credential - ) + self._client = ServiceBusClient(self._namespace, credential=self._credential) return self._client async def close(self) -> None: @@ -163,14 +158,19 @@ async def close(self) -> None: # ------------------------------------------------------------------ def _build_message( - self, message: dict, group_id: str | None, dedup_id: str | None + self, + body: bytes, + attributes: dict[str, str] | None, + group_id: str | None, + dedup_id: str | None, ) -> ServiceBusMessage: if self.session_enabled and not group_id: raise MessageSendError( - f"group_id is required when sending to session-enabled queue " - f"{self.queue_name!r}" + f"group_id is required when sending to session-enabled queue {self.queue_name!r}" ) - sb_message = ServiceBusMessage(json.dumps(message)) + sb_message = ServiceBusMessage(body) + if attributes: + sb_message.application_properties = dict(attributes) if group_id: sb_message.session_id = group_id if dedup_id: @@ -179,14 +179,15 @@ def _build_message( async def send( self, - message: dict, + body: bytes, + attributes: dict[str, str] | None = None, delay: int = 0, *, group_id: str | None = None, dedup_id: str | None = None, ) -> str: client = await self._ensure() - sb_message = self._build_message(message, group_id, dedup_id) + sb_message = self._build_message(body, attributes, group_id, dedup_id) try: async with client.get_queue_sender(self.queue_name) as sender: if delay: @@ -204,7 +205,7 @@ async def send( async def send_batch( self, - messages: list[dict], + messages: list[OutgoingMessage], *, group_id: str | None = None, dedup_ids: list[str] | None = None, @@ -213,7 +214,7 @@ async def send_batch( if dedup_ids is not None and len(dedup_ids) != len(messages): raise MessageSendError("dedup_ids must be parallel to messages") sb_messages = [ - self._build_message(m, group_id, dedup_ids[i] if dedup_ids else None) + self._build_message(m.body, m.attributes, group_id, dedup_ids[i] if dedup_ids else None) for i, m in enumerate(messages) ] try: @@ -272,15 +273,22 @@ async def receive( token = str(m.lock_token) self._pending[token] = (receiver, m) tokens.add(token) + # Stringify user-defined application_properties into a flat map, + # alongside the broker metadata we always surface. + attrs: dict[str, str] = { + "sequence_number": str(m.sequence_number), + "enqueued_time": str(m.enqueued_time_utc), + } + for key, value in (m.application_properties or {}).items(): + key_str = key.decode() if isinstance(key, bytes) else str(key) + val_str = value.decode() if isinstance(value, bytes) else str(value) + attrs[key_str] = val_str messages.append( Message( id=str(m.message_id or ""), - body=json.loads(str(m)), + body=bytes(m), receipt_handle=token, - attributes={ - "sequence_number": m.sequence_number, - "enqueued_time": str(m.enqueued_time_utc), - }, + attributes=attrs, group_id=m.session_id, dedup_id=str(m.message_id) if m.message_id else None, receive_count=(m.delivery_count or 0) + 1, @@ -355,9 +363,7 @@ async def dead_letter(self, receipt_handle: str, reason: str) -> None: ) receiver, message = entry try: - await receiver.dead_letter_message( - message, reason=reason, error_description=reason - ) + await receiver.dead_letter_message(message, reason=reason, error_description=reason) except ResourceNotFoundError as e: raise QueueNotFoundError(str(e)) from e except HttpResponseError as e: @@ -381,9 +387,7 @@ async def get_queue_depth(self) -> int: from azure.servicebus.aio.management import ServiceBusAdministrationClient if self._connection_string: - admin = ServiceBusAdministrationClient.from_connection_string( - self._connection_string - ) + admin = ServiceBusAdministrationClient.from_connection_string(self._connection_string) else: admin = ServiceBusAdministrationClient(self._namespace, credential=self._credential) try: @@ -408,9 +412,7 @@ async def health_check(self) -> bool: async def _purge_receiver(self, receiver) -> None: async with receiver: while True: - messages = await receiver.receive_messages( - max_message_count=100, max_wait_time=5 - ) + messages = await receiver.receive_messages(max_message_count=100, max_wait_time=5) if not messages: break for msg in messages: diff --git a/cloudrift/messaging/base.py b/cloudrift/messaging/base.py index 3e1f520..60d4435 100644 --- a/cloudrift/messaging/base.py +++ b/cloudrift/messaging/base.py @@ -1,21 +1,48 @@ +import json from abc import ABC, abstractmethod from dataclasses import dataclass, field +@dataclass +class OutgoingMessage: + """A message to send via :meth:`MessagingBackend.send_batch`. + + ``body`` is the raw payload bytes; ``attributes`` is an optional flat map of + string metadata that maps to SQS ``MessageAttributes`` (String type) and + Service Bus ``application_properties``. + """ + + body: bytes + attributes: dict[str, str] | None = None + + @dataclass class Message: id: str - body: dict + body: bytes receipt_handle: str - attributes: dict = field(default_factory=dict) + attributes: dict[str, str] = field(default_factory=dict) group_id: str | None = None dedup_id: str | None = None receive_count: int | None = None + def json(self): + """Decode the raw ``body`` bytes as JSON. + + Convenience for the common case where the payload was sent with + :func:`send_json`. Raises ``json.JSONDecodeError`` if the body is not + valid JSON. + """ + return json.loads(self.body) + class MessagingBackend(ABC): """Abstract base class for cloud messaging/queue backends. + The primitive payload is **raw bytes** plus an optional flat ``attributes`` + map (string → string). JSON users should use the :func:`send_json` helper + and :meth:`Message.json` to (de)serialize without touching the byte layer. + Backends hold long-lived async clients. Use ``await backend.close()`` (or ``async with backend:``) to release sockets cleanly. @@ -28,27 +55,30 @@ class MessagingBackend(ABC): @abstractmethod async def send( self, - message: dict, + body: bytes, + attributes: dict[str, str] | None = None, delay: int = 0, *, group_id: str | None = None, dedup_id: str | None = None, ) -> str: - """Send a message. Returns the message ID. + """Send a raw-bytes message with optional attributes. Returns the message ID. - group_id/dedup_id apply to FIFO (SQS) or session-enabled (Service Bus) - queues. SQS FIFO does not support per-message ``delay``. + ``attributes`` map to SQS ``MessageAttributes`` (String type) / + Service Bus ``application_properties``. group_id/dedup_id apply to FIFO + (SQS) or session-enabled (Service Bus) queues. SQS FIFO does not support + per-message ``delay``. """ @abstractmethod async def send_batch( self, - messages: list[dict], + messages: list[OutgoingMessage], *, group_id: str | None = None, dedup_ids: list[str] | None = None, ) -> list[str]: - """Send multiple messages. Returns list of message IDs. + """Send multiple :class:`OutgoingMessage`. Returns list of message IDs. ``group_id`` applies to every message; ``dedup_ids``, if given, must be parallel to ``messages``. @@ -65,6 +95,10 @@ async def receive( ) -> list[Message]: """Receive messages. wait_time is long-poll duration in seconds. + Each :class:`Message` carries the raw ``body`` bytes and an + ``attributes`` map (string → string) populated from the provider's + message attributes / application properties. + ``group_id`` receives from a specific session (Service Bus only; SQS cannot filter by group). ``visibility_timeout`` overrides the queue's visibility timeout on SQS; ignored on Service Bus (lock duration is @@ -117,3 +151,27 @@ async def __aenter__(self) -> "MessagingBackend": async def __aexit__(self, exc_type, exc, tb) -> None: await self.close() + + +async def send_json( + backend: MessagingBackend, + message: dict, + attributes: dict[str, str] | None = None, + delay: int = 0, + *, + group_id: str | None = None, + dedup_id: str | None = None, +) -> str: + """Serialize ``message`` to JSON bytes and send it via ``backend``. + + Backend-agnostic convenience wrapper around :meth:`MessagingBackend.send` + for the common JSON-payload case. Decode the received body with + :meth:`Message.json`. + """ + return await backend.send( + json.dumps(message).encode(), + attributes, + delay, + group_id=group_id, + dedup_id=dedup_id, + ) diff --git a/cloudrift/messaging/sqs.py b/cloudrift/messaging/sqs.py index 4c62a1c..d7cb764 100644 --- a/cloudrift/messaging/sqs.py +++ b/cloudrift/messaging/sqs.py @@ -11,7 +11,7 @@ MessagingError, QueueNotFoundError, ) -from cloudrift.messaging.base import Message, MessagingBackend +from cloudrift.messaging.base import Message, MessagingBackend, OutgoingMessage class AWSSQSBackend(MessagingBackend): @@ -22,9 +22,10 @@ class AWSSQSBackend(MessagingBackend): the underlying connections. Use one of the class methods to construct: - - ``from_access_key`` — static credentials (+ optional session token for assumed roles) - - ``from_iam_role`` — instance profile / environment / ECS task role - - ``from_profile`` — named profile from ``~/.aws/credentials`` + - ``from_access_key`` — static credentials (+ optional session token for assumed roles) + - ``from_iam_role`` — instance profile / environment / ECS task role + - ``from_profile`` — named profile from ``~/.aws/credentials`` + - ``from_assume_role`` — STS AssumeRole into a target role (cross-account) """ def __init__( @@ -55,9 +56,9 @@ def __init__( self._client_cm = None self._client = None self._lock = asyncio.Lock() - # receipt_handle → raw message body (JSON string), retained between - # receive() and delete()/dead_letter() so emulated dead-lettering can - # re-send the original payload to the DLQ. + # receipt_handle → raw message body (str as returned by SQS), retained + # between receive() and delete()/dead_letter() so emulated dead-lettering + # can re-send the original payload to the DLQ. self._pending: dict[str, str] = {} # ------------------------------------------------------------------ @@ -90,12 +91,32 @@ def from_iam_role( queue_url: str, region: str = "us-east-1", endpoint_url: str | None = None, + exclude_env_credentials: bool = False, **kwargs, ) -> "AWSSQSBackend": - """Authenticate via IAM role / instance profile / environment variables.""" - session = aioboto3.Session(region_name=region) + """Authenticate via IAM role / instance profile / environment variables. + + Set ``exclude_env_credentials=True`` to drop the environment-variable + credential provider from the resolver, so that any ``AWS_ACCESS_KEY_ID`` / + ``AWS_SECRET_ACCESS_KEY`` set elsewhere in the process (e.g. per-request + credentials assumed for another service) cannot shadow the long-lived + instance / ECS task role this client should use. + """ + session = cls._build_iam_session(region, exclude_env_credentials) return cls(queue_url, session, endpoint_url=endpoint_url, **kwargs) + @staticmethod + def _build_iam_session(region: str, exclude_env_credentials: bool) -> aioboto3.Session: + if not exclude_env_credentials: + return aioboto3.Session(region_name=region) + import aiobotocore.session + + botocore_session = aiobotocore.session.AioSession() + # The container/instance-role provider auto-refreshes; dropping "env" + # ensures stray process env credentials can't take precedence over it. + botocore_session.get_component("credential_provider").remove("env") + return aioboto3.Session(botocore_session=botocore_session, region_name=region) + @classmethod def from_profile( cls, @@ -109,6 +130,39 @@ def from_profile( session = aioboto3.Session(profile_name=profile_name, region_name=region) return cls(queue_url, session, endpoint_url=endpoint_url, **kwargs) + @classmethod + def from_assume_role( + cls, + queue_url: str, + role_arn: str, + external_id: str | None = None, + region: str = "us-east-1", + session_name: str = "cloudrift-sqs", + endpoint_url: str | None = None, + **kwargs, + ) -> "AWSSQSBackend": + """Authenticate by assuming an IAM role via STS (cross-account access). + + Calls ``sts:AssumeRole`` (optionally with ``ExternalId``) using the + ambient credential chain, then builds the SQS session from the returned + temporary credentials. Note: the temporary credentials are not + auto-refreshed; construct a new backend if the session expires. + """ + import boto3 + + sts = boto3.client("sts", region_name=region) + params: dict = {"RoleArn": role_arn, "RoleSessionName": session_name} + if external_id: + params["ExternalId"] = external_id + creds = sts.assume_role(**params)["Credentials"] + session = aioboto3.Session( + aws_access_key_id=creds["AccessKeyId"], + aws_secret_access_key=creds["SecretAccessKey"], + aws_session_token=creds["SessionToken"], + region_name=region, + ) + return cls(queue_url, session, endpoint_url=endpoint_url, **kwargs) + # ------------------------------------------------------------------ # Internal lifecycle # ------------------------------------------------------------------ @@ -142,9 +196,7 @@ async def close(self) -> None: # MessagingBackend implementation # ------------------------------------------------------------------ - def _fifo_params( - self, group_id: str | None, dedup_id: str | None, delay: int = 0 - ) -> dict: + def _fifo_params(self, group_id: str | None, dedup_id: str | None, delay: int = 0) -> dict: """Validate FIFO/standard constraints and return per-message kwargs.""" if self._is_fifo: if delay: @@ -153,23 +205,33 @@ def _fifo_params( "use a queue-level delivery delay instead" ) if not group_id: - raise MessageSendError( - "group_id is required when sending to an SQS FIFO queue" - ) + raise MessageSendError("group_id is required when sending to an SQS FIFO queue") params: dict = {"MessageGroupId": group_id} if dedup_id: params["MessageDeduplicationId"] = dedup_id return params if group_id or dedup_id: raise FeatureNotSupportedError( - "group_id/dedup_id are only supported on SQS FIFO queues " - f"(queue: {self.queue_url})" + f"group_id/dedup_id are only supported on SQS FIFO queues (queue: {self.queue_url})" ) return {"DelaySeconds": delay} if delay else {} + @staticmethod + def _message_attributes(attributes: dict[str, str] | None) -> dict: + """Map a flat str→str attribute dict to SQS MessageAttributes (String type).""" + if not attributes: + return {} + return { + "MessageAttributes": { + key: {"DataType": "String", "StringValue": value} + for key, value in attributes.items() + } + } + async def send( self, - message: dict, + body: bytes, + attributes: dict[str, str] | None = None, delay: int = 0, *, group_id: str | None = None, @@ -177,10 +239,11 @@ async def send( ) -> str: client = await self._ensure() params = self._fifo_params(group_id, dedup_id, delay) + params.update(self._message_attributes(attributes)) try: response = await client.send_message( QueueUrl=self.queue_url, - MessageBody=json.dumps(message), + MessageBody=body.decode(), **params, ) return response["MessageId"] @@ -189,7 +252,7 @@ async def send( async def send_batch( self, - messages: list[dict], + messages: list[OutgoingMessage], *, group_id: str | None = None, dedup_ids: list[str] | None = None, @@ -200,7 +263,8 @@ async def send_batch( entries = [] for i, msg in enumerate(messages): params = self._fifo_params(group_id, dedup_ids[i] if dedup_ids else None) - entries.append({"Id": str(i), "MessageBody": json.dumps(msg), **params}) + params.update(self._message_attributes(msg.attributes)) + entries.append({"Id": str(i), "MessageBody": msg.body.decode(), **params}) try: response = await client.send_message_batch(QueueUrl=self.queue_url, Entries=entries) if response.get("Failed"): @@ -219,9 +283,7 @@ async def receive( visibility_timeout: int | None = None, ) -> list[Message]: if group_id is not None: - raise FeatureNotSupportedError( - "SQS cannot receive from a specific message group" - ) + raise FeatureNotSupportedError("SQS cannot receive from a specific message group") client = await self._ensure() kwargs: dict = {} if visibility_timeout is not None: @@ -232,21 +294,27 @@ async def receive( MaxNumberOfMessages=min(max_messages, 10), WaitTimeSeconds=wait_time, AttributeNames=["All"], + MessageAttributeNames=["All"], **kwargs, ) messages = [] for m in response.get("Messages", []): - attrs = m.get("Attributes", {}) - receive_count = attrs.get("ApproximateReceiveCount") + system_attrs = m.get("Attributes", {}) + receive_count = system_attrs.get("ApproximateReceiveCount") + # Surface user-defined MessageAttributes as a flat str→str map. + attrs = { + key: spec.get("StringValue", "") + for key, spec in m.get("MessageAttributes", {}).items() + } self._pending[m["ReceiptHandle"]] = m["Body"] messages.append( Message( id=m["MessageId"], - body=json.loads(m["Body"]), + body=m["Body"].encode(), receipt_handle=m["ReceiptHandle"], attributes=attrs, - group_id=attrs.get("MessageGroupId"), - dedup_id=attrs.get("MessageDeduplicationId"), + group_id=system_attrs.get("MessageGroupId"), + dedup_id=system_attrs.get("MessageDeduplicationId"), receive_count=int(receive_count) if receive_count else None, ) ) @@ -356,9 +424,7 @@ async def purge(self) -> None: async def health_check(self) -> bool: try: client = await self._ensure() - await client.get_queue_attributes( - QueueUrl=self.queue_url, AttributeNames=["QueueArn"] - ) + await client.get_queue_attributes(QueueUrl=self.queue_url, AttributeNames=["QueueArn"]) return True except Exception: return False diff --git a/cloudrift/storage/__init__.py b/cloudrift/storage/__init__.py index 2411ce8..4b51f1b 100644 --- a/cloudrift/storage/__init__.py +++ b/cloudrift/storage/__init__.py @@ -19,6 +19,7 @@ def get_storage(provider: str, **kwargs) -> StorageBackend: Examples: get_storage("s3", bucket="my-bucket", region="us-east-1") # IAM/env get_storage("s3", bucket="b", aws_access_key_id="AKIA...", aws_secret_access_key="...") + get_storage("s3", bucket="b", role_arn="arn:aws:iam::...:role/x", external_id="...") get_storage("s3", bucket="b", profile_name="dev") get_storage("azure_blob", connection_string="...", container="my-container") get_storage("azure_blob", account_url="...", account_key="...", container="c") @@ -27,6 +28,8 @@ def get_storage(provider: str, **kwargs) -> StorageBackend: if provider == "s3": from cloudrift.storage.s3 import AWSS3Backend + if "role_arn" in kwargs: + return AWSS3Backend.from_assume_role(**kwargs) if "aws_access_key_id" in kwargs: return AWSS3Backend.from_access_key(**kwargs) if "profile_name" in kwargs: @@ -72,6 +75,8 @@ def get_storage_client(provider: str, **kwargs): if provider == "s3": from cloudrift.storage.s3 import AWSS3Client + if "role_arn" in kwargs: + return AWSS3Client.from_assume_role(**kwargs) if "aws_access_key_id" in kwargs: return AWSS3Client.from_access_key(**kwargs) if "profile_name" in kwargs: diff --git a/cloudrift/storage/s3.py b/cloudrift/storage/s3.py index f7c2498..ec8bf90 100644 --- a/cloudrift/storage/s3.py +++ b/cloudrift/storage/s3.py @@ -21,9 +21,10 @@ class AWSS3Client: was issued from it. Use one of the class methods to construct: - - ``from_access_key`` — static credentials (+ optional session token for assumed roles) - - ``from_iam_role`` — instance profile / environment / ECS task role - - ``from_profile`` — named profile from ``~/.aws/credentials`` + - ``from_access_key`` — static credentials (+ optional session token for assumed roles) + - ``from_iam_role`` — instance profile / environment / ECS task role + - ``from_profile`` — named profile from ``~/.aws/credentials`` + - ``from_assume_role`` — STS AssumeRole into a target role (cross-account) """ def __init__( @@ -94,6 +95,38 @@ def from_profile( session = aioboto3.Session(profile_name=profile_name, region_name=region) return cls(session, endpoint_url=endpoint_url, **kwargs) + @classmethod + def from_assume_role( + cls, + role_arn: str, + external_id: str | None = None, + region: str = "us-east-1", + session_name: str = "cloudrift-s3", + endpoint_url: str | None = None, + **kwargs, + ) -> "AWSS3Client": + """Authenticate by assuming an IAM role via STS (cross-account access). + + Calls ``sts:AssumeRole`` (optionally with ``ExternalId``) using the + ambient credential chain, then builds the S3 session from the returned + temporary credentials. Note: the temporary credentials are not + auto-refreshed; construct a new client if the session expires. + """ + import boto3 + + sts = boto3.client("sts", region_name=region) + params: dict = {"RoleArn": role_arn, "RoleSessionName": session_name} + if external_id: + params["ExternalId"] = external_id + creds = sts.assume_role(**params)["Credentials"] + session = aioboto3.Session( + aws_access_key_id=creds["AccessKeyId"], + aws_secret_access_key=creds["SecretAccessKey"], + aws_session_token=creds["SessionToken"], + region_name=region, + ) + return cls(session, endpoint_url=endpoint_url, **kwargs) + # ------------------------------------------------------------------ # Bucket view factory # ------------------------------------------------------------------ @@ -217,6 +250,27 @@ def from_profile( ) return cls(bucket, client, owns_client=True) + @classmethod + def from_assume_role( + cls, + bucket: str, + role_arn: str, + external_id: str | None = None, + region: str = "us-east-1", + session_name: str = "cloudrift-s3", + endpoint_url: str | None = None, + **kwargs, + ) -> "AWSS3Backend": + client = AWSS3Client.from_assume_role( + role_arn=role_arn, + external_id=external_id, + region=region, + session_name=session_name, + endpoint_url=endpoint_url, + **kwargs, + ) + return cls(bucket, client, owns_client=True) + # ------------------------------------------------------------------ # Lifecycle # ------------------------------------------------------------------ diff --git a/pyproject.toml b/pyproject.toml index 8d4d6d7..b8996b6 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,8 +4,8 @@ build-backend = "hatchling.build" [project] name = "lyzr-cloudrift" -version = "0.2.5" -description = "Cloud-agnostic abstraction for storage, messaging, document databases, cache, secrets, SQL, pub/sub, and email" +version = "0.2.6" +description = "Cloud-agnostic abstraction for storage, messaging, document databases, cache, secrets, SQL, crypto (KMS), pub/sub, and email" readme = "README.md" requires-python = ">=3.11" license = { file = "LICENSE" } @@ -44,6 +44,7 @@ azure = [ "azure-servicebus>=7.11.0", "azure-identity>=1.15.0", "azure-keyvault-secrets>=4.7.0", + "azure-keyvault-keys>=4.8.0", # Key Vault keys crypto (KMS analog) "azure-eventgrid>=4.9.0", "azure-communication-email>=1.0.0", "motor>=3.3.0", # Cosmos DB via MongoDB API @@ -93,6 +94,7 @@ all = [ "azure-servicebus>=7.11.0", "azure-identity>=1.15.0", "azure-keyvault-secrets>=4.7.0", + "azure-keyvault-keys>=4.8.0", "azure-eventgrid>=4.9.0", "azure-communication-email>=1.0.0", "redis[hiredis]>=5.0.0", @@ -107,7 +109,7 @@ all = [ dev = [ "pytest>=8.0", "pytest-asyncio>=0.23", - "moto[s3,sqs,sns,ses,secretsmanager,server]>=5.0", + "moto[s3,sqs,sns,ses,secretsmanager,kms,server]>=5.0", "aioboto3>=13.0.0", "motor>=3.3.0", "pymongo>=4.6.3", @@ -117,6 +119,7 @@ dev = [ "aiosmtplib>=3.0", "azure-communication-email>=1.0.0", "azure-identity>=1.15.0", + "azure-keyvault-keys>=4.8.0", "azure-servicebus>=7.11.0", ] diff --git a/tests/test_crypto.py b/tests/test_crypto.py new file mode 100644 index 0000000..ca318af --- /dev/null +++ b/tests/test_crypto.py @@ -0,0 +1,127 @@ +import boto3 +import pytest +from moto.server import ThreadedMotoServer + +from cloudrift.core.exceptions import CryptoError +from cloudrift.crypto import get_crypto + +REGION = "us-east-1" + + +@pytest.fixture(scope="module") +def moto_server(): + server = ThreadedMotoServer(port=0) + server.start() + host, port = server._server.server_address + yield f"http://{host}:{port}" + server.stop() + + +@pytest.fixture +def kms_key_id(moto_server): + kms = boto3.client( + "kms", + region_name=REGION, + aws_access_key_id="test", + aws_secret_access_key="test", + endpoint_url=moto_server, + ) + return kms.create_key(Description="cloudrift-test")["KeyMetadata"]["KeyId"] + + +@pytest.fixture +async def crypto_backend(moto_server, kms_key_id): + backend = get_crypto( + "aws_kms", + key_id=kms_key_id, + aws_access_key_id="test", + aws_secret_access_key="test", + region=REGION, + endpoint_url=moto_server, + ) + yield backend + await backend.close() + + +# --------------------------------------------------------------------------- +# round-trip +# --------------------------------------------------------------------------- + +async def test_encrypt_decrypt_bytes(crypto_backend): + ciphertext = await crypto_backend.encrypt(b"top-secret-token") + assert ciphertext != b"top-secret-token" + assert await crypto_backend.decrypt(ciphertext) == b"top-secret-token" + + +async def test_encrypt_decrypt_str_roundtrip(crypto_backend): + token = "ya29.a0Af-OAuth-Access-Token-Value" + blob = await crypto_backend.encrypt_str(token) + assert isinstance(blob, str) + assert blob != token + assert await crypto_backend.decrypt_str(blob) == token + + +async def test_ciphertext_is_nondeterministic(crypto_backend): + a = await crypto_backend.encrypt_str("same-input") + b = await crypto_backend.encrypt_str("same-input") + assert a != b # KMS adds randomness + assert await crypto_backend.decrypt_str(a) == "same-input" + assert await crypto_backend.decrypt_str(b) == "same-input" + + +async def test_empty_string_short_circuits(crypto_backend): + assert await crypto_backend.encrypt_str("") == "" + assert await crypto_backend.decrypt_str("") == "" + + +async def test_unicode_payload(crypto_backend): + payload = "héllo — 世界 🔐" + assert await crypto_backend.decrypt_str(await crypto_backend.encrypt_str(payload)) == payload + + +# --------------------------------------------------------------------------- +# context manager + decrypt without key_id +# --------------------------------------------------------------------------- + +async def test_async_context_manager(moto_server, kms_key_id): + async with get_crypto( + "aws_kms", + key_id=kms_key_id, + aws_access_key_id="test", + aws_secret_access_key="test", + region=REGION, + endpoint_url=moto_server, + ) as backend: + blob = await backend.encrypt_str("data") + assert await backend.decrypt_str(blob) == "data" + + +async def test_encrypt_without_key_id_raises(moto_server): + backend = get_crypto( + "aws_kms", + aws_access_key_id="test", + aws_secret_access_key="test", + region=REGION, + endpoint_url=moto_server, + ) + try: + with pytest.raises(CryptoError): + await backend.encrypt(b"data") + finally: + await backend.close() + + +# --------------------------------------------------------------------------- +# factory routing +# --------------------------------------------------------------------------- + +def test_get_crypto_unknown_provider(): + with pytest.raises(ValueError, match="Unknown crypto provider"): + get_crypto("gcp_kms", key_id="x") + + +def test_get_crypto_returns_aws_backend(moto_server): + from cloudrift.crypto.aws_kms import AWSKMSBackend + + backend = get_crypto("aws_kms", key_id="alias/x", region=REGION) + assert isinstance(backend, AWSKMSBackend) diff --git a/tests/test_messaging.py b/tests/test_messaging.py index 3a5cc40..282aaee 100644 --- a/tests/test_messaging.py +++ b/tests/test_messaging.py @@ -5,7 +5,7 @@ from moto.server import ThreadedMotoServer from cloudrift.core.exceptions import FeatureNotSupportedError, MessageSendError -from cloudrift.messaging import get_queue +from cloudrift.messaging import OutgoingMessage, get_queue, send_json REGION = "us-east-1" QUEUE_NAME = "test-queue" @@ -38,15 +38,13 @@ async def sqs_backend(moto_server, sqs_client): # Create a DLQ and wire the source queue to it via RedrivePolicy so # dead_letter() can resolve the target from the queue itself. dlq_url = sqs_client.create_queue(QueueName=DLQ_NAME)["QueueUrl"] - dlq_arn = sqs_client.get_queue_attributes( - QueueUrl=dlq_url, AttributeNames=["QueueArn"] - )["Attributes"]["QueueArn"] + dlq_arn = sqs_client.get_queue_attributes(QueueUrl=dlq_url, AttributeNames=["QueueArn"])[ + "Attributes" + ]["QueueArn"] queue_url = sqs_client.create_queue( QueueName=QUEUE_NAME, Attributes={ - "RedrivePolicy": json.dumps( - {"deadLetterTargetArn": dlq_arn, "maxReceiveCount": "5"} - ) + "RedrivePolicy": json.dumps({"deadLetterTargetArn": dlq_arn, "maxReceiveCount": "5"}) }, )["QueueUrl"] backend = get_queue( @@ -64,29 +62,59 @@ async def sqs_backend(moto_server, sqs_client): sqs_client.delete_queue(QueueUrl=dlq_url) -async def test_send_and_receive(sqs_backend): - msg_id = await sqs_backend.send({"action": "greet", "name": "cloudrift"}) +async def test_send_and_receive_bytes(sqs_backend): + msg_id = await sqs_backend.send(b"raw payload") assert isinstance(msg_id, str) messages = await sqs_backend.receive(max_messages=1) assert len(messages) == 1 - assert messages[0].body == {"action": "greet", "name": "cloudrift"} + assert messages[0].body == b"raw payload" + assert isinstance(messages[0].body, bytes) + + +async def test_send_with_attributes_round_trip(sqs_backend): + await sqs_backend.send(b"hi", attributes={"content_type": "text/plain", "source": "unit-test"}) + [m] = await sqs_backend.receive(max_messages=1) + assert m.body == b"hi" + assert m.attributes["content_type"] == "text/plain" + assert m.attributes["source"] == "unit-test" + + +async def test_send_json_and_message_json_helper(sqs_backend): + msg_id = await send_json(sqs_backend, {"action": "greet", "name": "cloudrift"}) + assert isinstance(msg_id, str) + [m] = await sqs_backend.receive(max_messages=1) + assert m.json() == {"action": "greet", "name": "cloudrift"} async def test_send_batch(sqs_backend): - ids = await sqs_backend.send_batch([{"n": 1}, {"n": 2}, {"n": 3}]) + ids = await sqs_backend.send_batch( + [OutgoingMessage(body=b"1"), OutgoingMessage(body=b"2"), OutgoingMessage(body=b"3")] + ) assert len(ids) == 3 +async def test_send_batch_with_attributes(sqs_backend): + await sqs_backend.send_batch( + [ + OutgoingMessage(body=b"a", attributes={"k": "1"}), + OutgoingMessage(body=b"b", attributes={"k": "2"}), + ] + ) + messages = await sqs_backend.receive(max_messages=10) + by_body = {m.body: m.attributes.get("k") for m in messages} + assert by_body == {b"a": "1", b"b": "2"} + + async def test_delete_message(sqs_backend): - await sqs_backend.send({"x": 1}) + await sqs_backend.send(b"x") messages = await sqs_backend.receive(max_messages=1) assert messages await sqs_backend.delete(messages[0].receipt_handle) async def test_purge(sqs_backend): - await sqs_backend.send({"a": 1}) - await sqs_backend.send({"b": 2}) + await sqs_backend.send(b"a") + await sqs_backend.send(b"b") await sqs_backend.purge() messages = await sqs_backend.receive(max_messages=10) assert messages == [] @@ -108,20 +136,20 @@ async def test_health_check(sqs_backend): async def test_standard_queue_zero_delay_omits_delay_seconds(sqs_backend): # Regression: DelaySeconds must be omitted when delay == 0 so the same # code path works on FIFO queues (which reject the parameter). - msg_id = await sqs_backend.send({"ping": True}) + msg_id = await sqs_backend.send(b"ping") assert isinstance(msg_id, str) async def test_group_id_on_standard_queue_raises(sqs_backend): with pytest.raises(FeatureNotSupportedError): - await sqs_backend.send({"x": 1}, group_id="g1") + await sqs_backend.send(b"x", group_id="g1") # --- dead_letter / get_queue_depth --- async def test_dead_letter_moves_message_to_dlq(sqs_backend, sqs_client): - await sqs_backend.send({"poison": True, "id": 7}) + await send_json(sqs_backend, {"poison": True, "id": 7}) [m] = await sqs_backend.receive(max_messages=1) await sqs_backend.dead_letter(m.receipt_handle, reason="schema mismatch") @@ -168,7 +196,7 @@ async def test_dead_letter_without_dlq_raises(moto_server, sqs_client): endpoint_url=moto_server, ) try: - await backend.send({"x": 1}) + await backend.send(b"x") [m] = await backend.receive(max_messages=1) with pytest.raises(MessagingError, match="No dead-letter queue configured"): await backend.dead_letter(m.receipt_handle, reason="x") @@ -191,7 +219,7 @@ async def test_dead_letter_with_explicit_dlq_url(moto_server, sqs_client): endpoint_url=moto_server, ) try: - await backend.send({"n": 1}) + await send_json(backend, {"n": 1}) [m] = await backend.receive(max_messages=1) await backend.dead_letter(m.receipt_handle, reason="explicit") resp = sqs_client.receive_message(QueueUrl=dlq_url, MaxNumberOfMessages=1) @@ -203,7 +231,7 @@ async def test_dead_letter_with_explicit_dlq_url(moto_server, sqs_client): async def test_nack_drops_pending_entry(sqs_backend): - await sqs_backend.send({"retry": 1}) + await sqs_backend.send(b"retry") [m] = await sqs_backend.receive(max_messages=1) assert m.receipt_handle in sqs_backend._pending await sqs_backend.nack(m.receipt_handle) @@ -215,8 +243,8 @@ async def test_nack_drops_pending_entry(sqs_backend): async def test_get_queue_depth(sqs_backend): assert await sqs_backend.get_queue_depth() == 0 - await sqs_backend.send({"a": 1}) - await sqs_backend.send({"b": 2}) + await sqs_backend.send(b"a") + await sqs_backend.send(b"b") assert await sqs_backend.get_queue_depth() == 2 await sqs_backend.purge() @@ -252,20 +280,20 @@ async def fifo_backend(moto_server): async def test_fifo_send_requires_group_id(fifo_backend): with pytest.raises(MessageSendError, match="group_id is required"): - await fifo_backend.send({"x": 1}) + await fifo_backend.send(b"x") async def test_fifo_send_with_delay_raises(fifo_backend): with pytest.raises(FeatureNotSupportedError): - await fifo_backend.send({"x": 1}, delay=5, group_id="g1") + await fifo_backend.send(b"x", delay=5, group_id="g1") async def test_fifo_send_and_receive_exposes_fifo_fields(fifo_backend): - await fifo_backend.send({"n": 1}, group_id="owner-1", dedup_id="d-1") + await send_json(fifo_backend, {"n": 1}, group_id="owner-1", dedup_id="d-1") messages = await fifo_backend.receive(max_messages=1) assert len(messages) == 1 m = messages[0] - assert m.body == {"n": 1} + assert m.json() == {"n": 1} assert m.group_id == "owner-1" assert m.dedup_id == "d-1" assert m.receive_count == 1 @@ -273,8 +301,8 @@ async def test_fifo_send_and_receive_exposes_fifo_fields(fifo_backend): async def test_fifo_dedup_id_suppresses_duplicate(fifo_backend): - await fifo_backend.send({"n": 1}, group_id="g-dedup", dedup_id="same-id") - await fifo_backend.send({"n": 2}, group_id="g-dedup", dedup_id="same-id") + await fifo_backend.send(b"1", group_id="g-dedup", dedup_id="same-id") + await fifo_backend.send(b"2", group_id="g-dedup", dedup_id="same-id") messages = await fifo_backend.receive(max_messages=10) assert len(messages) == 1 await fifo_backend.delete(messages[0].receipt_handle) @@ -282,42 +310,46 @@ async def test_fifo_dedup_id_suppresses_duplicate(fifo_backend): async def test_fifo_ordering_within_group(fifo_backend): for i in range(3): - await fifo_backend.send({"seq": i}, group_id="g-order", dedup_id=f"ord-{i}") + await send_json(fifo_backend, {"seq": i}, group_id="g-order", dedup_id=f"ord-{i}") received = [] while len(received) < 3: messages = await fifo_backend.receive(max_messages=10) if not messages: break for m in messages: - received.append(m.body["seq"]) + received.append(m.json()["seq"]) await fifo_backend.delete(m.receipt_handle) assert received == [0, 1, 2] async def test_fifo_send_batch_with_group_and_dedup_ids(fifo_backend): ids = await fifo_backend.send_batch( - [{"n": 1}, {"n": 2}], group_id="g-batch", dedup_ids=["b-1", "b-2"] + [OutgoingMessage(body=b"1"), OutgoingMessage(body=b"2")], + group_id="g-batch", + dedup_ids=["b-1", "b-2"], ) assert len(ids) == 2 messages = await fifo_backend.receive(max_messages=10) - assert {m.body["n"] for m in messages} == {1, 2} + assert {m.body for m in messages} == {b"1", b"2"} for m in messages: await fifo_backend.delete(m.receipt_handle) async def test_fifo_send_batch_mismatched_dedup_ids(fifo_backend): with pytest.raises(MessageSendError, match="parallel"): - await fifo_backend.send_batch([{"n": 1}], group_id="g", dedup_ids=["a", "b"]) + await fifo_backend.send_batch( + [OutgoingMessage(body=b"1")], group_id="g", dedup_ids=["a", "b"] + ) async def test_nack_redelivers_immediately(fifo_backend): - await fifo_backend.send({"task": "retry-me"}, group_id="g-nack", dedup_id="nack-1") + await send_json(fifo_backend, {"task": "retry-me"}, group_id="g-nack", dedup_id="nack-1") first = await fifo_backend.receive(max_messages=1, visibility_timeout=300) assert len(first) == 1 await fifo_backend.nack(first[0].receipt_handle) second = await fifo_backend.receive(max_messages=1) assert len(second) == 1 - assert second[0].body == {"task": "retry-me"} + assert second[0].json() == {"task": "retry-me"} assert second[0].receive_count == 2 await fifo_backend.delete(second[0].receipt_handle) @@ -349,17 +381,108 @@ async def test_dlq_redrive_flow(moto_server): main_q = get_queue("sqs", queue_url=main_url, **creds) dlq = get_queue("sqs", queue_url=dlq_url, **creds) try: - await dlq.send({"owner_id": "u1", "messages": ["hi"]}, group_id="u1", dedup_id="orig-1") + await send_json( + dlq, {"owner_id": "u1", "messages": ["hi"]}, group_id="u1", dedup_id="orig-1" + ) failed = await dlq.receive(max_messages=10) assert len(failed) == 1 m = failed[0] + # Re-send the raw bytes body verbatim to the main queue. await main_q.send(m.body, group_id=m.group_id, dedup_id="redrive-abc-123") await dlq.delete(m.receipt_handle) redriven = await main_q.receive(max_messages=1) - assert redriven[0].body == {"owner_id": "u1", "messages": ["hi"]} + assert redriven[0].json() == {"owner_id": "u1", "messages": ["hi"]} assert redriven[0].group_id == "u1" finally: await main_q.close() await dlq.close() sqs.delete_queue(QueueUrl=main_url) sqs.delete_queue(QueueUrl=dlq_url) + + +# --- Assume-role auth (STS) --- + + +async def test_get_queue_routes_role_arn_to_assume_role(): + """role_arn in kwargs routes get_queue -> from_assume_role (STS path).""" + from unittest.mock import MagicMock, patch + + from cloudrift.messaging.sqs import AWSSQSBackend + + fake_sts = MagicMock() + fake_sts.assume_role.return_value = { + "Credentials": { + "AccessKeyId": "ASIA-TEMP", + "SecretAccessKey": "temp-secret", + "SessionToken": "temp-token", + } + } + with patch("boto3.client", return_value=fake_sts) as boto_client: + backend = get_queue( + "sqs", + queue_url="https://sqs.us-east-1.amazonaws.com/123/q", + role_arn="arn:aws:iam::123456789012:role/cross", + external_id="ext-42", + region=REGION, + ) + assert isinstance(backend, AWSSQSBackend) + boto_client.assert_called_once_with("sts", region_name=REGION) + fake_sts.assume_role.assert_called_once() + call_kwargs = fake_sts.assume_role.call_args.kwargs + assert call_kwargs["RoleArn"] == "arn:aws:iam::123456789012:role/cross" + assert call_kwargs["ExternalId"] == "ext-42" + # Temp credentials threaded into the aioboto3 session. + creds = await backend._session.get_credentials() + frozen = await creds.get_frozen_credentials() + assert frozen.access_key == "ASIA-TEMP" + assert frozen.token == "temp-token" + + +def test_assume_role_omits_external_id_when_absent(): + from unittest.mock import MagicMock, patch + + fake_sts = MagicMock() + fake_sts.assume_role.return_value = { + "Credentials": { + "AccessKeyId": "ASIA-X", + "SecretAccessKey": "s", + "SessionToken": "t", + } + } + with patch("boto3.client", return_value=fake_sts): + get_queue( + "sqs", + queue_url="https://sqs.us-east-1.amazonaws.com/123/q", + role_arn="arn:aws:iam::123456789012:role/cross", + region=REGION, + ) + assert "ExternalId" not in fake_sts.assume_role.call_args.kwargs + + +# --------------------------------------------------------------------------- +# exclude_env_credentials (prevent ambient env creds shadowing the task role) +# --------------------------------------------------------------------------- + +def _credential_methods(backend): + resolver = backend._session._session.get_component("credential_provider") + return [p.METHOD for p in resolver.providers] + + +def test_from_iam_role_keeps_env_provider_by_default(): + backend = get_queue( + "sqs", queue_url="https://sqs.us-east-1.amazonaws.com/123/q", region=REGION + ) + assert "env" in _credential_methods(backend) + + +def test_exclude_env_credentials_drops_env_provider(): + backend = get_queue( + "sqs", + queue_url="https://sqs.us-east-1.amazonaws.com/123/q", + region=REGION, + exclude_env_credentials=True, + ) + methods = _credential_methods(backend) + assert "env" not in methods + # the rest of the chain (roles, shared credentials) is preserved + assert "assume-role" in methods diff --git a/tests/test_messaging_azure.py b/tests/test_messaging_azure.py index fe7e4d3..e457f36 100644 --- a/tests/test_messaging_azure.py +++ b/tests/test_messaging_azure.py @@ -12,6 +12,7 @@ from cloudrift.core.exceptions import FeatureNotSupportedError, MessageSendError, MessagingError from cloudrift.messaging.azure_bus import AzureServiceBusBackend +from cloudrift.messaging.base import OutgoingMessage CONN_STR = "Endpoint=sb://test.servicebus.windows.net/;SharedAccessKeyName=x;SharedAccessKey=y" @@ -32,17 +33,40 @@ def _patch_client(backend, client): backend._client = client -def _make_received_message(lock_token="tok-1", session_id=None, message_id="m-1", - delivery_count=0, body='{"n": 1}'): - m = MagicMock() - m.lock_token = lock_token - m.session_id = session_id - m.message_id = message_id - m.delivery_count = delivery_count - m.sequence_number = 7 - m.enqueued_time_utc = "2026-01-01" - m.__str__ = MagicMock(return_value=body) - return m +class _FakeReceivedMessage: + """Minimal stand-in for ServiceBusReceivedMessage. + + The real type exposes its raw payload via ``bytes(message)``; emulate that + deterministically (MagicMock does not reliably support ``__bytes__``). + """ + + def __init__( + self, lock_token, session_id, message_id, delivery_count, body, application_properties + ): + self.lock_token = lock_token + self.session_id = session_id + self.message_id = message_id + self.delivery_count = delivery_count + self.sequence_number = 7 + self.enqueued_time_utc = "2026-01-01" + self.application_properties = application_properties + self._body = body + + def __bytes__(self): + return self._body + + +def _make_received_message( + lock_token="tok-1", + session_id=None, + message_id="m-1", + delivery_count=0, + body=b'{"n": 1}', + application_properties=None, +): + return _FakeReceivedMessage( + lock_token, session_id, message_id, delivery_count, body, application_properties + ) async def test_send_sets_session_and_message_id(): @@ -52,7 +76,7 @@ async def test_send_sets_session_and_message_id(): client.get_queue_sender.return_value = sender _patch_client(backend, client) - await backend.send({"n": 1}, group_id="owner-1", dedup_id="d-1") + await backend.send(b'{"n": 1}', group_id="owner-1", dedup_id="d-1") sent = sender.send_messages.call_args[0][0] assert sent.session_id == "owner-1" @@ -63,7 +87,7 @@ async def test_sessionless_send_to_session_queue_raises(): backend = _make_backend(session_enabled=True) _patch_client(backend, MagicMock()) with pytest.raises(MessageSendError, match="group_id is required"): - await backend.send({"n": 1}) + await backend.send(b'{"n": 1}') async def test_send_without_session_on_plain_queue_ok(): @@ -73,11 +97,23 @@ async def test_send_without_session_on_plain_queue_ok(): client.get_queue_sender.return_value = sender _patch_client(backend, client) - await backend.send({"n": 1}) + await backend.send(b'{"n": 1}') sent = sender.send_messages.call_args[0][0] assert sent.session_id is None +async def test_send_sets_application_properties_from_attributes(): + backend = _make_backend(session_enabled=False) + client = MagicMock() + sender = _mock_sender() + client.get_queue_sender.return_value = sender + _patch_client(backend, client) + + await backend.send(b"raw", attributes={"content_type": "text/plain"}) + sent = sender.send_messages.call_args[0][0] + assert sent.application_properties == {"content_type": "text/plain"} + + async def test_send_batch_sets_per_message_dedup_ids(): backend = _make_backend(session_enabled=True) client = MagicMock() @@ -88,7 +124,9 @@ async def test_send_batch_sets_per_message_dedup_ids(): _patch_client(backend, client) ids = await backend.send_batch( - [{"n": 1}, {"n": 2}], group_id="g", dedup_ids=["a", "b"] + [OutgoingMessage(body=b'{"n": 1}'), OutgoingMessage(body=b'{"n": 2}')], + group_id="g", + dedup_ids=["a", "b"], ) assert ids == ["a", "b"] added = [c.args[0] for c in batch.add_message.call_args_list] @@ -145,7 +183,12 @@ async def test_receive_populates_fifo_fields(): backend = _make_backend(session_enabled=True) client = MagicMock() receiver = AsyncMock() - raw = _make_received_message(session_id="owner-1", message_id="d-1", delivery_count=1) + raw = _make_received_message( + session_id="owner-1", + message_id="d-1", + delivery_count=1, + application_properties={b"content_type": b"text/plain"}, + ) receiver.receive_messages.return_value = [raw] client.get_queue_receiver.return_value = receiver _patch_client(backend, client) @@ -154,7 +197,10 @@ async def test_receive_populates_fifo_fields(): assert m.group_id == "owner-1" assert m.dedup_id == "d-1" assert m.receive_count == 2 # delivery_count + 1 - assert m.body == {"n": 1} + assert m.body == b'{"n": 1}' + assert m.json() == {"n": 1} + # application_properties (bytes keys/values) are stringified into attributes. + assert m.attributes["content_type"] == "text/plain" async def test_nack_abandons_and_releases_receiver(): @@ -231,9 +277,7 @@ async def test_get_queue_depth_uses_admin_client(): admin = AsyncMock() admin.get_queue_runtime_properties.return_value = props admin.__aenter__.return_value = admin - with patch( - "azure.servicebus.aio.management.ServiceBusAdministrationClient" - ) as admin_cls: + with patch("azure.servicebus.aio.management.ServiceBusAdministrationClient") as admin_cls: admin_cls.from_connection_string.return_value = admin depth = await backend.get_queue_depth() assert depth == 5 diff --git a/tests/test_storage.py b/tests/test_storage.py index 24b1b8c..0d479f9 100644 --- a/tests/test_storage.py +++ b/tests/test_storage.py @@ -278,6 +278,66 @@ async def test_view_close_is_noop_when_shared(s3_client_two_buckets): assert await view_b.exists("b.txt") +# --- Assume-role auth (STS) --- + + +async def test_get_storage_routes_role_arn_to_assume_role(): + """role_arn in kwargs routes get_storage -> from_assume_role (STS path).""" + from unittest.mock import MagicMock, patch + + from cloudrift.storage.s3 import AWSS3Backend + + fake_sts = MagicMock() + fake_sts.assume_role.return_value = { + "Credentials": { + "AccessKeyId": "ASIA-TEMP", + "SecretAccessKey": "temp-secret", + "SessionToken": "temp-token", + } + } + with patch("boto3.client", return_value=fake_sts) as boto_client: + backend = get_storage( + "s3", + bucket="b", + role_arn="arn:aws:iam::123456789012:role/cross", + external_id="ext-42", + region=REGION, + ) + assert isinstance(backend, AWSS3Backend) + boto_client.assert_called_once_with("sts", region_name=REGION) + call_kwargs = fake_sts.assume_role.call_args.kwargs + assert call_kwargs["RoleArn"] == "arn:aws:iam::123456789012:role/cross" + assert call_kwargs["ExternalId"] == "ext-42" + # Temp credentials from STS are threaded into the aioboto3 session. + creds = await backend._client._session.get_credentials() + frozen = await creds.get_frozen_credentials() + assert frozen.access_key == "ASIA-TEMP" + assert frozen.token == "temp-token" + + +def test_get_storage_client_routes_role_arn_to_assume_role(): + from unittest.mock import MagicMock, patch + + from cloudrift.storage.s3 import AWSS3Client + + fake_sts = MagicMock() + fake_sts.assume_role.return_value = { + "Credentials": { + "AccessKeyId": "ASIA-X", + "SecretAccessKey": "s", + "SessionToken": "t", + } + } + with patch("boto3.client", return_value=fake_sts): + client = get_storage_client( + "s3", + role_arn="arn:aws:iam::123456789012:role/cross", + region=REGION, + ) + assert isinstance(client, AWSS3Client) + assert "ExternalId" not in fake_sts.assume_role.call_args.kwargs + + async def test_get_storage_view_owns_client_and_closes_it(moto_server): """A backend obtained from get_storage() owns its client and closes it.""" bucket = "owned-close-test" diff --git a/uv.lock b/uv.lock index 49e5bb2..2359e14 100644 --- a/uv.lock +++ b/uv.lock @@ -346,6 +346,21 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/49/9a/417b3a533e01953a7c618884df2cb05a71e7b68bdbce4fbdb62349d2a2e8/azure_identity-1.25.3-py3-none-any.whl", hash = "sha256:f4d0b956a8146f30333e071374171f3cfa7bdb8073adb8c3814b65567aa7447c", size = 192138, upload-time = "2026-03-13T01:12:22.951Z" }, ] +[[package]] +name = "azure-keyvault-keys" +version = "4.11.1" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "azure-core" }, + { name = "cryptography" }, + { name = "isodate" }, + { name = "typing-extensions" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/f4/03/5ce6db28b545427d4ab572f6a4ef2a727b6b4e7bf6941cedddf98822535b/azure_keyvault_keys-4.11.1.tar.gz", hash = "sha256:90caa3a7b2c8f6b53c247ec115cf1c1dad7f107cc3aa9f35aff4838bbce7e562", size = 260915, upload-time = "2026-05-19T20:01:08.041Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/0c/3d/7bed91ae9268cf48124cf6990d8cd2c3daff7a2bb1e91a439d26ee90d705/azure_keyvault_keys-4.11.1-py3-none-any.whl", hash = "sha256:f46cdf6ee7a9baf27f70e6838327032886c8a087041dd56397773b1639da8fe2", size = 200651, upload-time = "2026-05-19T20:01:09.809Z" }, +] + [[package]] name = "azure-keyvault-secrets" version = "4.10.0" @@ -1193,7 +1208,7 @@ wheels = [ [[package]] name = "lyzr-cloudrift" -version = "0.2.5" +version = "0.2.6" source = { editable = "." } [package.optional-dependencies] @@ -1204,6 +1219,7 @@ all = [ { name = "azure-communication-email" }, { name = "azure-eventgrid" }, { name = "azure-identity" }, + { name = "azure-keyvault-keys" }, { name = "azure-keyvault-secrets" }, { name = "azure-servicebus" }, { name = "azure-storage-blob" }, @@ -1226,6 +1242,7 @@ azure = [ { name = "azure-communication-email" }, { name = "azure-eventgrid" }, { name = "azure-identity" }, + { name = "azure-keyvault-keys" }, { name = "azure-keyvault-secrets" }, { name = "azure-servicebus" }, { name = "azure-storage-blob" }, @@ -1241,6 +1258,7 @@ dev = [ { name = "aiosmtplib" }, { name = "azure-communication-email" }, { name = "azure-identity" }, + { name = "azure-keyvault-keys" }, { name = "azure-servicebus" }, { name = "fakeredis" }, { name = "httpx" }, @@ -1309,6 +1327,9 @@ requires-dist = [ { name = "azure-identity", marker = "extra == 'dev'", specifier = ">=1.15.0" }, { name = "azure-identity", marker = "extra == 'sql'", specifier = ">=1.15.0" }, { name = "azure-identity", marker = "extra == 'sql-mssql'", specifier = ">=1.15.0" }, + { name = "azure-keyvault-keys", marker = "extra == 'all'", specifier = ">=4.8.0" }, + { name = "azure-keyvault-keys", marker = "extra == 'azure'", specifier = ">=4.8.0" }, + { name = "azure-keyvault-keys", marker = "extra == 'dev'", specifier = ">=4.8.0" }, { name = "azure-keyvault-secrets", marker = "extra == 'all'", specifier = ">=4.7.0" }, { name = "azure-keyvault-secrets", marker = "extra == 'azure'", specifier = ">=4.7.0" }, { name = "azure-servicebus", marker = "extra == 'all'", specifier = ">=7.11.0" }, @@ -1324,7 +1345,7 @@ requires-dist = [ { name = "databricks-sql-connector", marker = "extra == 'sql-databricks'", specifier = ">=3.0" }, { name = "fakeredis", marker = "extra == 'dev'", specifier = ">=2.20.0" }, { name = "httpx", marker = "extra == 'dev'", specifier = ">=0.25.0" }, - { name = "moto", extras = ["s3", "sqs", "sns", "ses", "secretsmanager", "server"], marker = "extra == 'dev'", specifier = ">=5.0" }, + { name = "moto", extras = ["s3", "sqs", "sns", "ses", "secretsmanager", "kms", "server"], marker = "extra == 'dev'", specifier = ">=5.0" }, { name = "motor", marker = "extra == 'all'", specifier = ">=3.3.0" }, { name = "motor", marker = "extra == 'aws'", specifier = ">=3.3.0" }, { name = "motor", marker = "extra == 'azure'", specifier = ">=3.3.0" },