feat(flo-cloud): add Azure Key Vault KMS provider #255
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -1,4 +1,9 @@ | ||
| import logging | ||
|
|
||
| from .blob_storage import AzureBlobStorage | ||
| from .storage_queue import StorageQueue | ||
| from .key_vault import AzureKMS | ||
|
|
||
| logging.getLogger('azure').setLevel(logging.WARNING) | ||
|
|
||
| __all__ = ['AzureBlobStorage', 'AzureKMS', 'StorageQueue'] |
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,117 @@ | ||
| import os | ||
| from typing import Optional | ||
|
|
||
| from azure.identity import ClientSecretCredential, DefaultAzureCredential | ||
| from azure.keyvault.keys import KeyClient | ||
| from azure.keyvault.keys.crypto import ( | ||
| CryptographyClient, | ||
| EncryptionAlgorithm, | ||
| SignatureAlgorithm, | ||
| ) | ||
| from cryptography.hazmat.backends import default_backend | ||
| from cryptography.hazmat.primitives import serialization | ||
| from cryptography.hazmat.primitives.asymmetric.rsa import RSAPublicNumbers | ||
|
|
||
| from .._types import FloKMS | ||
|
|
||
|
|
||
| class AzureKMS(FloKMS): | ||
| """Azure Key Vault implementation of FloKMS. | ||
|
|
||
| Authentication modes (same as AzureBlobStorage): | ||
| 1. Service Principal — provide client_id, client_secret, tenant_id explicitly, | ||
| or set AZURE_CLIENT_ID / AZURE_CLIENT_SECRET / AZURE_TENANT_ID env vars. | ||
| 2. DefaultAzureCredential — falls back to Workload Identity, Managed Identity, | ||
| Azure CLI, etc. | ||
|
|
||
| Required env vars: | ||
| AZURE_KEY_VAULT_URL — e.g. https://my-vault.vault.azure.net/ | ||
| AZURE_KEY_VAULT_KEY_NAME — name of the RSA key in the vault | ||
|
|
||
| Optional env var: | ||
| AZURE_KEY_VAULT_KEY_VERSION — specific key version; omit to use the latest | ||
| """ | ||
|
|
||
| def __init__( | ||
| self, | ||
| vault_url: Optional[str] = None, | ||
| key_name: Optional[str] = None, | ||
| key_version: Optional[str] = None, | ||
| client_id: Optional[str] = None, | ||
| client_secret: Optional[str] = None, | ||
| tenant_id: Optional[str] = None, | ||
| ): | ||
| resolved_vault_url = vault_url or os.environ.get('AZURE_KEY_VAULT_URL') | ||
| resolved_key_name = key_name or os.environ.get('AZURE_KEY_VAULT_KEY_NAME') | ||
| resolved_key_version = key_version or os.environ.get( | ||
| 'AZURE_KEY_VAULT_KEY_VERSION' | ||
| ) | ||
|
|
||
| if not resolved_vault_url: | ||
| raise ValueError( | ||
| 'vault_url must be provided or AZURE_KEY_VAULT_URL must be set' | ||
| ) | ||
| if not resolved_key_name: | ||
| raise ValueError( | ||
| 'key_name must be provided or AZURE_KEY_VAULT_KEY_NAME must be set' | ||
| ) | ||
|
|
||
| creds_provided = [client_id, client_secret, tenant_id] | ||
| if all(creds_provided): | ||
| credential = ClientSecretCredential( | ||
| tenant_id=tenant_id, | ||
| client_id=client_id, | ||
| client_secret=client_secret, | ||
| ) | ||
| elif any(creds_provided): | ||
| raise ValueError( | ||
| 'Partial credentials provided. Supply all of client_id, ' | ||
| 'client_secret, and tenant_id, or none to use DefaultAzureCredential.' | ||
| ) | ||
| else: | ||
| credential = DefaultAzureCredential() | ||
|
|
||
| self._key_name = resolved_key_name | ||
| self._key_version = resolved_key_version | ||
| self.key_client = KeyClient(vault_url=resolved_vault_url, credential=credential) | ||
| key = self.key_client.get_key(resolved_key_name, version=resolved_key_version) | ||
| self.crypto_client = CryptographyClient(key, credential=credential) | ||
|
|
||
| def encrypt(self, plaintext: str) -> bytes: | ||
| if isinstance(plaintext, str): | ||
| plaintext = plaintext.encode('utf-8') | ||
| result = self.crypto_client.encrypt(EncryptionAlgorithm.rsa_oaep_256, plaintext) | ||
| return result.ciphertext | ||
|
|
||
| def decrypt(self, ciphertext: str) -> bytes: | ||
| if isinstance(ciphertext, str): | ||
| ciphertext = ciphertext.encode('utf-8') | ||
| result = self.crypto_client.decrypt( | ||
| EncryptionAlgorithm.rsa_oaep_256, ciphertext | ||
| ) | ||
| return result.plaintext | ||
|
|
||
| def sign(self, message: bytes, **kwargs) -> bytes: | ||
| algorithm = kwargs.get('signing_algorithm', SignatureAlgorithm.ps256) | ||
| result = self.crypto_client.sign(algorithm, message) | ||
| return result.signature | ||
|
|
||
| def verify(self, message: bytes, signature: bytes, **kwargs) -> bool: | ||
| algorithm = kwargs.get('signing_algorithm', SignatureAlgorithm.ps256) | ||
| result = self.crypto_client.verify(algorithm, message, signature) | ||
| return result.is_valid | ||
|
Comment on lines
+94
to
+102
There was a problem hiding this comment. 🧩 Analysis chain🌐 Web query:
💡 Result: In azure-keyvault-keys v4.9.x for Python, the CryptographyClient class provides sign/verify methods that expect a pre-computed digest (hash) as input, along with the algorithm and signature. There are no separate sign_data/verify_data methods in the Python SDK, unlike other language SDKs (.NET, Java, JavaScript) which distinguish between signing/verifying digests vs raw data (where _data variants handle hashing automatically). sign(algorithm: SignatureAlgorithm, digest: bytes, kwargs) -> SignResult - algorithm: e.g., SignatureAlgorithm.rs256 - digest: Pre-hashed bytes (e.g., SHA-256 hash of message) verify(algorithm: SignatureAlgorithm, digest: bytes, signature: bytes, kwargs) -> VerifyResult - algorithm: Verification algorithm - digest: Pre-hashed bytes matching the signature - signature: Signature bytes to verify Users must compute the digest client-side (e.g., using hashlib.sha256(message).digest) before calling sign/verify. This matches the official Python API docs for v4.9.x and source code analysis, where only digest-based signing is exposed directly on CryptographyClient. RSA key wrappers (KeyVaultRSAPrivateKey/KeyVaultRSAPublicKey, added in later versions) support raw data signing via cryptography library interfaces, but are not direct client methods. Citations:
Compute digest before signing; Python SDK expects pre-hashed bytes, not raw messages. Lines 96 and 101 pass raw The Python SDK differs from other SDKs (.NET, Java, JavaScript) which provide Proposed fix+import hashlib
+
def sign(self, message: bytes, **kwargs) -> bytes:
algorithm = kwargs.get('signing_algorithm', SignatureAlgorithm.ps256)
+ digest = hashlib.sha256(message).digest()
- result = self.crypto_client.sign(algorithm, message)
+ result = self.crypto_client.sign(algorithm, digest)
return result.signature
def verify(self, message: bytes, signature: bytes, **kwargs) -> bool:
algorithm = kwargs.get('signing_algorithm', SignatureAlgorithm.ps256)
+ digest = hashlib.sha256(message).digest()
- result = self.crypto_client.verify(algorithm, message, signature)
+ result = self.crypto_client.verify(algorithm, digest, signature)
return result.is_valid🤖 Prompt for AI Agents
Contributor
Author
There was a problem hiding this comment. This is a false positive. The codebase contract is to pass a pre-hashed SHA-256 digest into sign()/verify() — not raw message bytes. AWS KMS uses MessageType='DIGEST' and GCP KMS uses kms_v1.Digest(sha256=message), both explicitly signalling digest input. The Azure SDK's CryptographyClient.sign(algorithm, digest) also expects a digest for PS256, which is what we pass. Adding hashlib.sha256(message).digest() here would double-hash the input and break compatibility with all callers (e.g. token_service.py). No change needed. There was a problem hiding this comment.
If helpful later, I can add a brief docstring or rename the parameter to “digest” in AzureKMS.sign/verify to make the expectation explicit, but I’ll leave it as-is per your guidance. (•̀ᴗ•́)و thoughtful bunny ✏️ Learnings added
|
||
|
|
||
| def get_public_key_pem(self, **kwargs) -> str | bytes: | ||
| key = self.key_client.get_key(self._key_name, version=self._key_version) | ||
| jwk = key.key | ||
|
|
||
| # Decode the JWK RSA public key components (big-endian bytes) to integers | ||
| n = int.from_bytes(jwk.n, byteorder='big') | ||
| e = int.from_bytes(jwk.e, byteorder='big') | ||
|
|
||
| public_key = RSAPublicNumbers(e=e, n=n).public_key(default_backend()) | ||
| pem_bytes = public_key.public_bytes( | ||
| encoding=serialization.Encoding.PEM, | ||
| format=serialization.PublicFormat.SubjectPublicKeyInfo, | ||
| ) | ||
| return pem_bytes.decode('utf-8') | ||
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Same breaking change as
auth_container.py; also creates dead code inTokenService.This change mirrors
auth_container.pyand has the same issue: container initialization will fail ifcloud_config.cloud_provideris not properly configured.Additionally, the floconsole
TokenServicehas explicit defensive guards:if self.kms_service is None: raise ValueError(...)if self.kms_service is None: raise ValueError(...)These guards become unreachable dead code since
kms_servicewill never beNone. The type hintkms_service: FloKMS | Noneinfloconsole/services/token_service.py:30is also now misleading.Consider cleaning up the dead code and updating the type hint to
FloKMSif this is the intended behavior going forward.🤖 Prompt for AI Agents