workos 1.5.1__py3-none-any.whl → 5.38.0__py3-none-any.whl

workos/vault.py

@@ -0,0 +1,544 @@
+import base64
+from typing import Optional, Protocol, Sequence, Tuple
+from workos.types.vault import VaultObject, ObjectVersion, ObjectDigest, ObjectMetadata
+from workos.types.vault.key import DataKey, DataKeyPair, KeyContext, DecodedKeys
+from workos.types.list_resource import (
+    ListArgs,
+    ListMetadata,
+    ListPage,
+    WorkOSListResource,
+)
+from workos.utils.http_client import SyncHTTPClient
+from workos.utils.pagination_order import PaginationOrder
+from workos.utils.request_helper import (
+    DEFAULT_LIST_RESPONSE_LIMIT,
+    REQUEST_METHOD_DELETE,
+    REQUEST_METHOD_GET,
+    REQUEST_METHOD_POST,
+    REQUEST_METHOD_PUT,
+    RequestHelper,
+)
+from workos.utils.crypto_provider import CryptoProvider
+
+DEFAULT_RESPONSE_LIMIT = DEFAULT_LIST_RESPONSE_LIMIT
+
+VaultObjectList = WorkOSListResource[ObjectDigest, ListArgs, ListMetadata]
+
+
+class VaultModule(Protocol):
+    def read_object(self, *, object_id: str) -> VaultObject:
+        """
+        Get a Vault object with the value decrypted.
+
+        Kwargs:
+            object_id (str): The unique identifier for the object.
+        Returns:
+            VaultObject: A vault object with metadata, name and decrypted value.
+        """
+        ...
+
+    def read_object_by_name(self, *, name: str) -> VaultObject:
+        """
+        Get a Vault object by name with the value decrypted.
+
+        Kwargs:
+            name (str): The unique name of the object.
+        Returns:
+            VaultObject: A vault object with metadata, name and decrypted value.
+        """
+        ...
+
+    def list_objects(
+        self,
+        *,
+        limit: int = DEFAULT_RESPONSE_LIMIT,
+        before: Optional[str] = None,
+        after: Optional[str] = None,
+    ) -> VaultObjectList:
+        """
+        Gets a list of encrypted Vault objects.
+
+        Kwargs:
+            limit (int): The maximum number of objects to return. (Optional)
+            before (str): A cursor to return resources before. (Optional)
+            after (str): A cursor to return resources after. (Optional)
+
+        Returns:
+            VaultObjectList: A list of vault objects with built-in pagination iterator.
+        """
+        ...
+
+    def list_object_versions(
+        self,
+        *,
+        object_id: str,
+    ) -> Sequence[ObjectVersion]:
+        """
+        Gets a list of versions for a specific Vault object.
+
+        Kwargs:
+            object_id (str): The unique identifier for the object.
+
+        Returns:
+            Sequence[ObjectVersion]: A list of object versions.
+        """
+        ...
+
+    def create_object(
+        self,
+        *,
+        name: str,
+        value: str,
+        key_context: KeyContext,
+    ) -> ObjectMetadata:
+        """
+        Create a new Vault encrypted object.
+
+        Kwargs:
+            name (str): The name of the object.
+            value (str): The value to encrypt and store.
+            key_context (KeyContext): A set of key-value dictionary pairs that determines which root keys to use when encrypting data.
+
+        Returns:
+            VaultObject: The created vault object.
+        """
+        ...
+
+    def update_object(
+        self,
+        *,
+        object_id: str,
+        value: str,
+        version_check: Optional[str] = None,
+    ) -> VaultObject:
+        """
+        Update an existing Vault object.
+
+        Kwargs:
+            object_id (str): The unique identifier for the object.
+            value (str): The new value to encrypt and store.
+            version_check (str): A version of the object to prevent clobbering of data during concurrent updates. (Optional)
+
+        Returns:
+            VaultObject: The updated vault object.
+        """
+        ...
+
+    def delete_object(
+        self,
+        *,
+        object_id: str,
+    ) -> None:
+        """
+        Permanently delete a Vault encrypted object. Warning: this cannont be undone.
+
+        Kwargs:
+            object_id (str): The unique identifier for the object.
+        """
+        ...
+
+    def create_data_key(self, *, key_context: KeyContext) -> DataKeyPair:
+        """
+        Generate a data key for local encryption based on the provided key context.
+        The encrypted data key MUST be stored by the application, as it cannot be retrieved after generation.
+
+        Kwargs:
+            key_context (KeyContext): A set of key-value dictionary pairs that determines which root keys to use when encrypting data.
+        """
+        ...
+
+    def decrypt_data_key(
+        self,
+        *,
+        keys: str,
+    ) -> DataKey:
+        """
+        Decrypt encrypted data keys that were previously generated by create_data_key.
+
+        This method takes the encrypted data key blob and uses the WorkOS Vault service
+        to decrypt it, returning the plaintext data key that can be used for local
+        encryption/decryption operations.
+
+        Kwargs:
+            keys (str): The base64-encoded encrypted data key blob returned by create_data_key.
+
+        Returns:
+            DataKey: The decrypted data key containing the key ID and the plaintext key material.
+        """
+        ...
+
+    def encrypt(
+        self,
+        *,
+        data: str,
+        key_context: KeyContext,
+        associated_data: Optional[str] = None,
+    ) -> str:
+        """
+        Encrypt data locally using AES-GCM with a data key derived from the provided context.
+
+        This method generates a new data key for each encryption operation, ensuring that
+        the same plaintext will produce different ciphertext each time it's encrypted.
+        The encrypted data key is embedded in the result so it can be decrypted later.
+
+        Kwargs:
+            data (str): The plaintext data to encrypt.
+            key_context (KeyContext): A set of key-value dictionary pairs that determines which root keys to use when encrypting data.
+            associated_data (str): Additional authenticated data (AAD) that will be authenticated but not encrypted. (Optional)
+
+        Returns:
+            str: Base64-encoded encrypted data containing the IV, authentication tag, encrypted data key, and ciphertext.
+        """
+        ...
+
+    def decrypt(
+        self, *, encrypted_data: str, associated_data: Optional[str] = None
+    ) -> str:
+        """
+        Decrypt data that was previously encrypted using the encrypt method.
+
+        This method extracts the encrypted data key from the encrypted payload,
+        decrypts it using the WorkOS Vault service, and then uses the resulting
+        data key to decrypt the actual data using AES-GCM.
+
+        Kwargs:
+            encrypted_data (str): The base64-encoded encrypted data returned by the encrypt method.
+            associated_data (str): The same additional authenticated data (AAD) that was used during encryption, if any. (Optional)
+
+        Returns:
+            str: The original plaintext data.
+
+        Raises:
+            ValueError: If the encrypted_data format is invalid or if associated_data doesn't match what was used during encryption.
+            cryptography.exceptions.InvalidTag: If the authentication tag verification fails (data has been tampered with).
+        """
+        ...
+
+
+class Vault(VaultModule):
+    _http_client: SyncHTTPClient
+    _crypto_provider: CryptoProvider
+
+    def __init__(self, http_client: SyncHTTPClient):
+        self._http_client = http_client
+        self._crypto_provider = CryptoProvider()
+
+    def read_object(
+        self,
+        *,
+        object_id: str,
+    ) -> VaultObject:
+        if not object_id:
+            raise ValueError("Incomplete arguments: 'object_id' is a required argument")
+
+        response = self._http_client.request(
+            RequestHelper.build_parameterized_url(
+                "vault/v1/kv/{object_id}",
+                object_id=object_id,
+            ),
+            method=REQUEST_METHOD_GET,
+        )
+
+        return VaultObject.model_validate(response)
+
+    def read_object_by_name(
+        self,
+        *,
+        name: str,
+    ) -> VaultObject:
+        if not name:
+            raise ValueError("Incomplete arguments: 'name' is a required argument")
+
+        response = self._http_client.request(
+            RequestHelper.build_parameterized_url(
+                "vault/v1/kv/name/{name}",
+                name=name,
+            ),
+            method=REQUEST_METHOD_GET,
+        )
+
+        return VaultObject.model_validate(response)
+
+    def list_objects(
+        self,
+        *,
+        limit: int = DEFAULT_RESPONSE_LIMIT,
+        before: Optional[str] = None,
+        after: Optional[str] = None,
+    ) -> VaultObjectList:
+        list_params: ListArgs = {
+            "limit": limit,
+            "before": before,
+            "after": after,
+        }
+
+        response = self._http_client.request(
+            "vault/v1/kv",
+            method=REQUEST_METHOD_GET,
+            params=list_params,
+        )
+
+        # Ensure object field is present
+        response_dict = dict(response)
+        if "object" not in response_dict:
+            response_dict["object"] = "list"
+
+        return VaultObjectList(
+            list_method=self.list_objects,
+            list_args=list_params,
+            **ListPage[ObjectDigest](**response_dict).model_dump(),
+        )
+
+    def list_object_versions(
+        self,
+        *,
+        object_id: str,
+    ) -> Sequence[ObjectVersion]:
+        response = self._http_client.request(
+            RequestHelper.build_parameterized_url(
+                "vault/v1/kv/{object_id}/versions",
+                object_id=object_id,
+            ),
+            method=REQUEST_METHOD_GET,
+        )
+
+        return [
+            ObjectVersion.model_validate(version)
+            for version in response.get("data", [])
+        ]
+
+    def create_object(
+        self,
+        *,
+        name: str,
+        value: str,
+        key_context: KeyContext,
+    ) -> ObjectMetadata:
+        if not name or not value:
+            raise ValueError(
+                "Incomplete arguments: 'name' and 'value' are required arguments"
+            )
+
+        request_data = {
+            "name": name,
+            "value": value,
+            "key_context": key_context,
+        }
+
+        response = self._http_client.request(
+            "vault/v1/kv",
+            method=REQUEST_METHOD_POST,
+            json=request_data,
+        )
+
+        return ObjectMetadata.model_validate(response)
+
+    def update_object(
+        self,
+        *,
+        object_id: str,
+        value: str,
+        version_check: Optional[str] = None,
+    ) -> VaultObject:
+        if not object_id:
+            raise ValueError("Incomplete arguments: 'object_id' is a required argument")
+
+        request_data = {
+            "value": value,
+        }
+        if version_check is not None:
+            request_data["version_check"] = version_check
+
+        response = self._http_client.request(
+            RequestHelper.build_parameterized_url(
+                "vault/v1/kv/{object_id}",
+                object_id=object_id,
+            ),
+            method=REQUEST_METHOD_PUT,
+            json=request_data,
+        )
+
+        return VaultObject.model_validate(response)
+
+    def delete_object(
+        self,
+        *,
+        object_id: str,
+    ) -> None:
+        if not object_id:
+            raise ValueError("Incomplete arguments: 'object_id' is a required argument")
+
+        self._http_client.request(
+            RequestHelper.build_parameterized_url(
+                "vault/v1/kv/{object_id}",
+                object_id=object_id,
+            ),
+            method=REQUEST_METHOD_DELETE,
+        )
+
+    def create_data_key(self, *, key_context: KeyContext) -> DataKeyPair:
+        request_data = {
+            "context": key_context,
+        }
+
+        response = self._http_client.request(
+            "vault/v1/keys/data-key",
+            method=REQUEST_METHOD_POST,
+            json=request_data,
+        )
+
+        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,
+        *,
+        keys: str,
+    ) -> DataKey:
+        request_data = {
+            "keys": keys,
+        }
+
+        response = self._http_client.request(
+            "vault/v1/keys/decrypt",
+            method=REQUEST_METHOD_POST,
+            json=request_data,
+        )
+
+        return DataKey.model_validate(
+            {"id": response["id"], "key": response["data_key"]}
+        )
+
+    def encrypt(
+        self,
+        *,
+        data: str,
+        key_context: KeyContext,
+        associated_data: Optional[str] = None,
+    ) -> str:
+        key_pair = self.create_data_key(key_context=key_context)
+
+        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_u32(len(key_blob))
+        aad_buffer = associated_data.encode("utf-8") if associated_data else None
+        iv = self._crypto_provider.random_bytes(12)
+
+        result = self._crypto_provider.encrypt(
+            data.encode("utf-8"), key, iv, aad_buffer
+        )
+
+        combined = (
+            result["iv"]
+            + result["tag"]
+            + prefix_len_buffer
+            + key_blob
+            + result["ciphertext"]
+        )
+
+        return self._bytes_to_base64(combined)
+
+    def decrypt(
+        self, *, encrypted_data: str, associated_data: Optional[str] = None
+    ) -> str:
+        decoded = self._decode(encrypted_data)
+        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,
+            key=key,
+            iv=decoded.iv,
+            tag=decoded.tag,
+            aad=aad_buffer,
+        )
+
+        return decrypted_bytes.decode("utf-8")
+
+    def _base64_to_bytes(self, data: str) -> bytes:
+        return base64.b64decode(data)
+
+    def _bytes_to_base64(self, data: bytes) -> str:
+        return base64.b64encode(data).decode("utf-8")
+
+    def _encode_u32(self, value: int) -> bytes:
+        """
+        Encode a 32-bit unsigned integer as LEB128.
+
+        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.
+        Encoding format: [IV][TAG][4B Length][keyBlob][ciphertext]
+        """
+        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:]
+
+        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
+
+        for i, b in enumerate(buf):
+            if i > 4:
+                raise ValueError("LEB128 integer overflow (was more than 4 bytes)")
+
+            res |= (b & 0x7F) << (7 * bit)
+
+            if (b & 0x80) == 0:
+                return res, i + 1
+
+            bit += 1
+
+        raise ValueError("LEB128 integer not found")