update encoding/decoding to match envelope encryption format, fix ekm endpoint calls

Author: debussyman

.gitignore

@@ -105,6 +105,7 @@ celerybeat.pid
 # Environments
 .env
 .venv
+.venv312
 env/
 venv/
 ENV/

tests/test_vault.py

@@ -59,13 +59,19 @@ def mock_object_versions(self):
 
     @pytest.fixture
     def mock_data_key(self):
-        return MockDataKey(
-            "key_01234567890abcdef", "MDEyMzQ1Njc4OWFiY2RlZjAxMjM0NTY3ODlhYmNkZWY="
-        ).dict()
+        return {
+            "id": "key_01234567890abcdef",
+            "data_key": "MDEyMzQ1Njc4OWFiY2RlZjAxMjM0NTY3ODlhYmNkZWY=",
+        }
 
     @pytest.fixture
     def mock_data_key_pair(self):
-        return MockDataKeyPair().dict()
+        return {
+            "context": {"key": "test-key"},
+            "id": "key_01234567890abcdef",
+            "data_key": "MDEyMzQ1Njc4OWFiY2RlZjAxMjM0NTY3ODlhYmNkZWY=",
+            "encrypted_keys": "ZW5jcnlwdGVkX2tleXNfZGF0YQ==",
+        }
 
     def test_read_object_success(
         self, mock_vault_object, capture_and_mock_http_client_request
@@ -186,7 +192,7 @@ def test_create_object_success(
         assert request_kwargs["url"].endswith("/vault/v1/kv")
         assert request_kwargs["json"]["name"] == "test-secret"
         assert request_kwargs["json"]["value"] == "secret-value"
-        assert request_kwargs["json"]["key_context"] == {"key": "test-key"}
+        assert request_kwargs["json"]["context"] == KeyContext({"key": "test-key"})
         assert vault_object.id == "vault_01234567890abcdef"
         assert vault_object.name == "test-secret"
         assert vault_object.value == "secret-value"
@@ -310,7 +316,7 @@ def test_create_data_key_success(
 
         assert request_kwargs["method"] == "post"
         assert request_kwargs["url"].endswith("/vault/v1/keys/data-key")
-        assert request_kwargs["json"]["key_context"] == {"key": "test-key"}
+        assert request_kwargs["json"]["context"] == KeyContext({"key": "test-key"})
         assert data_key_pair.data_key.id == "key_01234567890abcdef"
         assert data_key_pair.encrypted_keys == "ZW5jcnlwdGVkX2tleXNfZGF0YQ=="
 
@@ -345,7 +351,7 @@ def test_encrypt_success(
         # Verify create_data_key was called
         assert request_kwargs["method"] == "post"
         assert request_kwargs["url"].endswith("/vault/v1/keys/data-key")
-        assert request_kwargs["json"]["key_context"] == {"key": "test-key"}
+        assert request_kwargs["json"]["context"] == KeyContext({"key": "test-key"})
 
         # Verify we got encrypted data back
         assert isinstance(encrypted_data, str)
@@ -371,7 +377,12 @@ def test_encrypt_with_associated_data(
 
     def test_decrypt_success(self, mock_data_key, capture_and_mock_http_client_request):
         # First encrypt some data to get a valid encrypted payload
-        mock_data_key_pair = MockDataKeyPair().dict()
+        mock_data_key_pair = {
+            "context": {"key": "test-key"},
+            "id": "key_01234567890abcdef",
+            "data_key": "MDEyMzQ1Njc4OWFiY2RlZjAxMjM0NTY3ODlhYmNkZWY=",
+            "encrypted_keys": "ZW5jcnlwdGVkX2tleXNfZGF0YQ==",
+        }
 
         # Mock create_data_key for encryption
         capture_and_mock_http_client_request(self.http_client, mock_data_key_pair, 200)
@@ -393,7 +404,12 @@ def test_decrypt_with_associated_data(
         self, mock_data_key, capture_and_mock_http_client_request
     ):
         # First encrypt some data with associated data
-        mock_data_key_pair = MockDataKeyPair().dict()
+        mock_data_key_pair = {
+            "context": {"key": "test-key"},
+            "id": "key_01234567890abcdef",
+            "data_key": "MDEyMzQ1Njc4OWFiY2RlZjAxMjM0NTY3ODlhYmNkZWY=",
+            "encrypted_keys": "ZW5jcnlwdGVkX2tleXNfZGF0YQ==",
+        }
 
         # Mock create_data_key for encryption
         capture_and_mock_http_client_request(self.http_client, mock_data_key_pair, 200)

workos/async_client.py

@@ -14,6 +14,7 @@
 from workos.utils.http_client import AsyncHTTPClient
 from workos.webhooks import WebhooksModule
 from workos.widgets import WidgetsModule
+from workos.vault import VaultModule
 
 
 class AsyncClient(BaseClient):
@@ -112,3 +113,9 @@ def widgets(self) -> WidgetsModule:
         raise NotImplementedError(
             "Widgets APIs are not yet supported in the async client."
         )
+
+    @property
+    def vault(self) -> VaultModule:
+        raise NotImplementedError(
+            "Vault APIs are not yet supported in the async client."
+        )

workos/client.py

@@ -14,6 +14,7 @@
 from workos.user_management import UserManagement
 from workos.utils.http_client import SyncHTTPClient
 from workos.widgets import Widgets
+from workos.vault import Vault
 
 
 class SyncClient(BaseClient):
@@ -116,3 +117,9 @@ def widgets(self) -> Widgets:
         if not getattr(self, "_widgets", None):
             self._widgets = Widgets(http_client=self._http_client)
         return self._widgets
+
+    @property
+    def vault(self) -> Vault:
+        if not getattr(self, "_vault", None):
+            self._vault = Vault(http_client=self._http_client)
+        return self._vault

workos/types/vault/key.py

@@ -1,5 +1,5 @@
 from typing import Dict
-from pydantic import RootModel
+from pydantic import BaseModel, RootModel
 from workos.types.workos_model import WorkOSModel
 
 
@@ -16,3 +16,10 @@ class DataKeyPair(WorkOSModel):
     context: KeyContext
     data_key: DataKey
     encrypted_keys: str
+
+
+class DecodedKeys(BaseModel):
+    iv: bytes
+    tag: bytes
+    keys: str  # Base64-encoded string
+    ciphertext: bytes

workos/vault.py

@@ -1,8 +1,8 @@
 import base64
 import struct
-from typing import Dict, Optional, Protocol, Sequence
+from typing import Dict, Optional, Protocol, Sequence, Tuple
 from workos.types.vault import VaultObject, ObjectVersion
-from workos.types.vault.key import DataKey, DataKeyPair, KeyContext
+from workos.types.vault.key import DataKey, DataKeyPair, KeyContext, DecodedKeys
 from workos.types.list_resource import (
     ListArgs,
     ListMetadata,
@@ -285,7 +285,7 @@ def create_object(
         request_data = {
             "name": name,
             "value": value,
-            "key_context": key_context.root,
+            "context": key_context,
         }
 
         response = self._http_client.request(
@@ -341,7 +341,7 @@ def delete_object(
 
     def create_data_key(self, *, key_context: KeyContext) -> DataKeyPair:
         request_data = {
-            "key_context": key_context.root,
+            "context": key_context,
         }
 
         response = self._http_client.request(
@@ -350,7 +350,13 @@ def create_data_key(self, *, key_context: KeyContext) -> DataKeyPair:
             json=request_data,
         )
 
-        return DataKeyPair.model_validate(response)
+        return DataKeyPair.model_validate(
+            {
+                "context": response["context"],
+                "data_key": {"id": response["id"], "key": response["data_key"]},
+                "encrypted_keys": response["encrypted_keys"],
+            }
+        )
 
     def decrypt_data_key(
         self,
@@ -367,7 +373,9 @@ def decrypt_data_key(
             json=request_data,
         )
 
-        return DataKey.model_validate(response)
+        return DataKey.model_validate(
+            {"id": response["id"], "key": response["data_key"]}
+        )
 
     def encrypt(
         self, *, data: str, context: KeyContext, associated_data: Optional[str] = None
@@ -376,7 +384,7 @@ def encrypt(
 
         key = self._base64_to_bytes(key_pair.data_key.key)
         key_blob = self._base64_to_bytes(key_pair.encrypted_keys)
-        prefix_len_buffer = self._encode_uint32(len(key_blob))
+        prefix_len_buffer = self._encode_u32(len(key_blob))
         aad_buffer = associated_data.encode("utf-8") if associated_data else None
         iv = self._crypto_provider.random_bytes(12)
 
@@ -398,16 +406,16 @@ def decrypt(
         self, *, encrypted_data: str, associated_data: Optional[str] = None
     ) -> str:
         decoded = self._decode(encrypted_data)
-        data_key = self.decrypt_data_key(keys=self._bytes_to_base64(decoded["keys"]))
+        data_key = self.decrypt_data_key(keys=decoded.keys)
 
         key = self._base64_to_bytes(data_key.key)
         aad_buffer = associated_data.encode("utf-8") if associated_data else None
 
         decrypted_bytes = self._crypto_provider.decrypt(
-            ciphertext=decoded["ciphertext"],
+            ciphertext=decoded.ciphertext,
             key=key,
-            iv=decoded["iv"],
-            tag=decoded["tag"],
+            iv=decoded.iv,
+            tag=decoded.tag,
             aad=aad_buffer,
         )
 
@@ -419,30 +427,77 @@ def _base64_to_bytes(self, data: str) -> bytes:
     def _bytes_to_base64(self, data: bytes) -> str:
         return base64.b64encode(data).decode("utf-8")
 
-    def _encode_uint32(self, value: int) -> bytes:
-        return struct.pack(">I", value)  # Big-endian unsigned int (4 bytes)
+    def _encode_u32(self, value: int) -> bytes:
+        """
+        Encode a 32-bit unsigned integer as LEB128.
 
-    def _decode(self, encrypted_data_b64: str) -> Dict[str, bytes]:
+        Returns:
+            bytes: LEB128-encoded representation of the input value.
+        """
+        if value < 0 or value > 0xFFFFFFFF:
+            raise ValueError("Value must be a 32-bit unsigned integer")
+
+        encoded = bytearray()
+        while True:
+            byte = value & 0x7F
+            value >>= 7
+            if value != 0:
+                byte |= 0x80  # Set continuation bit
+            encoded.append(byte)
+            if value == 0:
+                break
+
+        return bytes(encoded)
+
+    def _decode(self, encrypted_data_b64: str) -> DecodedKeys:
         """
         This function extracts IV, tag, keyBlobLength, keyBlob, and ciphertext
-        from a base64-encoded payload. You must define this according to your encoding format.
-        Assumes format: [IV][TAG][4B Length][keyBlob][ciphertext]
+        from a base64-encoded payload.
+        Encoding format: [IV][TAG][4B Length][keyBlob][ciphertext]
         """
-        raw = base64.b64decode(encrypted_data_b64)
-        offset = 0
+        try:
+            payload = base64.b64decode(encrypted_data_b64)
+        except Exception as e:
+            raise ValueError("Base64 decoding failed") from e
+
+        iv = payload[0:12]
+        tag = payload[12:28]
+
+        try:
+            key_len, leb_len = self._decode_u32(payload[28:])
+        except Exception as e:
+            raise ValueError("Failed to decode key length") from e
+
+        keys_index = 28 + leb_len
+        keys_end = keys_index + key_len
+        keys_slice = payload[keys_index:keys_end]
+        keys = base64.b64encode(keys_slice).decode("utf-8")
+        ciphertext = payload[keys_end:]
 
-        iv = raw[offset : offset + 12]
-        offset += 12
+        return DecodedKeys(iv=iv, tag=tag, keys=keys, ciphertext=ciphertext)
+
+    def _decode_u32(self, buf: bytes) -> Tuple[int, int]:
+        """
+        Decode an unsigned LEB128-encoded 32-bit integer from bytes.
+
+        Returns:
+            (value, length_consumed)
+
+        Raises:
+            ValueError if decoding fails or overflows.
+        """
+        res = 0
+        bit = 0
 
-        tag = raw[offset : offset + 16]
-        offset += 16
+        for i, b in enumerate(buf):
+            if i > 4:
+                raise ValueError("LEB128 integer overflow (was more than 4 bytes)")
 
-        key_len = int.from_bytes(raw[offset : offset + 4], byteorder="big")
-        offset += 4
+            res |= (b & 0x7F) << (7 * bit)
 
-        key_blob = raw[offset : offset + key_len]
-        offset += key_len
+            if (b & 0x80) == 0:
+                return res, i + 1
 
-        ciphertext = raw[offset:]
+            bit += 1
 
-        return {"iv": iv, "tag": tag, "keys": key_blob, "ciphertext": ciphertext}
+        raise ValueError("LEB128 integer not found")