otdf-python 0.4.0__py3-none-any.whl → 0.4.2__py3-none-any.whl

otdf_python/kas_client.py

@@ -1,6 +1,4 @@
-"""
-KASClient: Handles communication with the Key Access Service (KAS).
-"""
+"""KASClient: Handles communication with the Key Access Service (KAS)."""
 
 import base64
 import hashlib
@@ -22,6 +20,8 @@ from .sdk_exceptions import SDKException
 
 @dataclass
 class KeyAccess:
+    """Key access response from KAS."""
+
     url: str
     wrapped_key: str
     ephemeral_public_key: str | None = None
@@ -29,6 +29,8 @@ class KeyAccess:
 
 
 class KASClient:
+    """Client for communicating with the Key Access Service (KAS)."""
+
     def __init__(
         self,
         kas_url=None,
@@ -37,6 +39,7 @@ class KASClient:
         use_plaintext=False,
         verify_ssl=True,
     ):
+        """Initialize KAS client."""
         self.kas_url = kas_url
         self.token_source = token_source
         self.cache = cache or KASKeyCache()
@@ -62,15 +65,33 @@ class KASClient:
             self._dpop_public_key
         )
 
-    def _normalize_kas_url(self, url: str) -> str:
+    def __enter__(self):
+        """Enter context manager."""
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        """Exit context manager and clean up resources."""
+        self.close()
+
+    def close(self):
+        """Close the KAS client and release resources.
+
+        This method should be called when the client is no longer needed
+        to properly clean up resources. It's also called automatically
+        when using the client as a context manager.
         """
-        Normalize KAS URLs based on client security settings.
+        if self.connect_rpc_client:
+            self.connect_rpc_client.close()
+
+    def _normalize_kas_url(self, url: str) -> str:
+        """Normalize KAS URLs based on client security settings.
 
         Args:
             url: The KAS URL to normalize
 
         Returns:
             Normalized URL with appropriate protocol and port
+
         """
         from urllib.parse import urlparse
 
@@ -78,7 +99,7 @@ class KASClient:
             # Parse the URL
             parsed = urlparse(url)
         except Exception as e:
-            raise SDKException(f"error trying to parse URL [{url}]", e)
+            raise SDKException(f"error trying to parse URL [{url}]: {e}") from e
 
         # Check if we have a host or if this is likely a hostname:port combination
         if parsed.hostname is None:
@@ -100,10 +121,10 @@ class KASClient:
                 try:
                     port = int(port_str)
                     return f"{scheme}://{host}:{port}"
-                except ValueError:
+                except ValueError as err:
                     raise SDKException(
                         f"error trying to create URL for host and port [{url}]"
-                    )
+                    ) from err
             else:
                 # Hostname with or without path, add default port
                 if "/" in url:
@@ -116,7 +137,7 @@ class KASClient:
         except Exception as e:
             raise SDKException(
                 f"error trying to create URL for host and port [{url}]", e
-            )
+            ) from e
 
     def _handle_existing_scheme(self, parsed) -> str:
         """Handle URLs with existing scheme by normalizing protocol and port."""
@@ -138,17 +159,17 @@ class KASClient:
             logging.debug(f"normalized url [{parsed.geturl()}] to [{normalized_url}]")
             return normalized_url
         except Exception as e:
-            raise SDKException("error creating KAS address", e)
+            raise SDKException(f"error creating KAS address: {e}") from e
 
     def _get_wrapped_key_base64(self, key_access):
-        """
-        Extract and normalize the wrapped key to base64-encoded string.
+        """Extract and normalize the wrapped key to base64-encoded string.
 
         Args:
             key_access: KeyAccess object
 
         Returns:
             Base64-encoded wrapped key string
+
         """
         wrapped_key = getattr(key_access, "wrappedKey", None) or getattr(
             key_access, "wrapped_key", None
@@ -166,14 +187,14 @@ class KASClient:
         return wrapped_key
 
     def _build_key_access_dict(self, key_access):
-        """
-        Build key access dictionary from KeyAccess object, handling both old and new field names.
+        """Build key access dictionary from KeyAccess object, handling both old and new field names.
 
         Args:
             key_access: KeyAccess object
 
         Returns:
             Dictionary with key access information
+
         """
         wrapped_key = self._get_wrapped_key_base64(key_access)
 
@@ -197,12 +218,12 @@ class KASClient:
         return key_access_dict
 
     def _add_optional_fields(self, key_access_dict, key_access):
-        """
-        Add optional fields to key access dictionary.
+        """Add optional fields to key access dictionary.
 
         Args:
             key_access_dict: Dictionary to add fields to
             key_access: KeyAccess object to extract fields from
+
         """
         # Policy binding
         policy_binding = getattr(key_access, "policyBinding", None) or getattr(
@@ -244,14 +265,14 @@ class KASClient:
             key_access_dict["header"] = base64.b64encode(header).decode("utf-8")
 
     def _get_algorithm_from_session_key_type(self, session_key_type):
-        """
-        Convert session key type to algorithm string for KAS.
+        """Convert session key type to algorithm string for KAS.
 
         Args:
             session_key_type: Session key type (EC_KEY_TYPE or RSA_KEY_TYPE)
 
         Returns:
             Algorithm string or None
+
         """
         if session_key_type == EC_KEY_TYPE:
             return "ec:secp256r1"  # Default EC curve for NanoTDF
@@ -262,8 +283,7 @@ class KASClient:
     def _build_rewrap_request(
         self, policy_json, client_public_key, key_access_dict, algorithm, has_header
     ):
-        """
-        Build the unsigned rewrap request structure.
+        """Build the unsigned rewrap request structure.
 
         Args:
             policy_json: Policy JSON string
@@ -274,6 +294,7 @@ class KASClient:
 
         Returns:
             Dictionary with unsigned rewrap request
+
         """
         import json
 
@@ -316,8 +337,7 @@ class KASClient:
     def _create_signed_request_jwt(
         self, policy_json, client_public_key, key_access, session_key_type=None
     ):
-        """
-        Create a signed JWT for the rewrap request.
+        """Create a signed JWT for the rewrap request.
         The JWT is signed with the DPoP private key.
 
         Args:
@@ -325,6 +345,7 @@ class KASClient:
             client_public_key: Client public key PEM string
             key_access: KeyAccess object
             session_key_type: Optional session key type (RSA_KEY_TYPE or EC_KEY_TYPE)
+
         """
         # Build key access dictionary handling both old and new field names
         key_access_dict = self._build_key_access_dict(key_access)
@@ -354,8 +375,7 @@ class KASClient:
         return jwt.encode(payload, self._dpop_private_key_pem, algorithm="RS256")
 
     def _create_connect_rpc_signed_token(self, key_access, policy_json):
-        """
-        Create a signed token specifically for Connect RPC requests.
+        """Create a signed token specifically for Connect RPC requests.
         For now, this delegates to the existing JWT creation method.
         """
         return self._create_signed_request_jwt(
@@ -363,8 +383,7 @@ class KASClient:
         )
 
     def _create_dpop_proof(self, method, url, access_token=None):
-        """
-        Create a DPoP proof JWT as per RFC 9449.
+        """Create a DPoP proof JWT as per RFC 9449.
 
         Args:
             method: HTTP method (e.g., "POST")
@@ -373,6 +392,7 @@ class KASClient:
 
         Returns:
             DPoP proof JWT string
+
         """
         now = int(time.time())
 
@@ -424,8 +444,7 @@ class KASClient:
         )
 
     def get_public_key(self, kas_info):
-        """
-        Get KAS public key using Connect RPC.
+        """Get KAS public key using Connect RPC.
         Checks cache first if available.
         """
         try:
@@ -448,10 +467,7 @@ class KASClient:
             raise
 
     def _get_public_key_with_connect_rpc(self, kas_info):
-        """
-        Get KAS public key using Connect RPC.
-        """
-
+        """Get KAS public key using Connect RPC."""
         # Get access token for authentication if token source is available
         access_token = None
         if self.token_source:
@@ -483,17 +499,17 @@ class KASClient:
                 f"Connect RPC public key request failed: {type(e).__name__}: {e}"
             )
             logging.error(f"Full traceback: {error_details}")
-            raise SDKException(f"Connect RPC public key request failed: {e}")
+            raise SDKException(f"Connect RPC public key request failed: {e}") from e
 
     def _normalize_session_key_type(self, session_key_type):
-        """
-        Normalize session key type to the appropriate enum value.
+        """Normalize session key type to the appropriate enum value.
 
         Args:
             session_key_type: Type of the session key (KeyType enum or string "RSA"/"EC")
 
         Returns:
             Normalized key type enum
+
         """
         if isinstance(session_key_type, str):
             if session_key_type.upper() == "RSA":
@@ -511,14 +527,14 @@ class KASClient:
         return session_key_type
 
     def _prepare_ec_keypair(self, session_key_type):
-        """
-        Prepare EC key pair for unwrapping.
+        """Prepare EC key pair for unwrapping.
 
         Args:
             session_key_type: EC key type with curve information
 
         Returns:
             ECKeyPair instance and client public key
+
         """
         from .eckeypair import ECKeyPair
 
@@ -528,12 +544,12 @@ class KASClient:
         return ec_key_pair, client_public_key
 
     def _prepare_rsa_keypair(self):
-        """
-        Prepare RSA key pair for unwrapping, reusing if possible.
+        """Prepare RSA key pair for unwrapping, reusing if possible.
         Uses separate ephemeral keys for encryption (not DPoP keys).
 
         Returns:
             Client public key PEM for the ephemeral encryption key
+
         """
         if self.decryptor is None:
             # Generate ephemeral keys for encryption (separate from DPoP keys)
@@ -543,8 +559,7 @@ class KASClient:
         return self.client_public_key
 
     def _unwrap_with_ec(self, wrapped_key, ec_key_pair, response_data):
-        """
-        Unwrap a key using EC cryptography.
+        """Unwrap a key using EC cryptography.
 
         Args:
             wrapped_key: The wrapped key to decrypt
@@ -553,6 +568,7 @@ class KASClient:
 
         Returns:
             Unwrapped key as bytes
+
         """
         if ec_key_pair is None:
             raise SDKException(
@@ -581,9 +597,7 @@ class KASClient:
         return gcm.decrypt(wrapped_key)
 
     def _ensure_client_keypair(self, session_key_type):
-        """
-        Ensure client keypair is generated and stored.
-        """
+        """Ensure client keypair is generated and stored."""
         if session_key_type == RSA_KEY_TYPE:
             if self.decryptor is None:
                 private_key, public_key = CryptoUtils.generate_rsa_keypair()
@@ -600,15 +614,13 @@ class KASClient:
                 self.client_public_key = CryptoUtils.get_rsa_public_key_pem(public_key)
 
     def _parse_and_decrypt_response(self, response):
-        """
-        Parse JSON response and decrypt the wrapped key.
-        """
+        """Parse JSON response and decrypt the wrapped key."""
         try:
             response_data = response.json()
         except Exception as e:
             logging.error(f"Failed to parse JSON response: {e}")
             logging.error(f"Raw response content: {response.content}")
-            raise SDKException(f"Invalid JSON response from KAS: {e}")
+            raise SDKException(f"Invalid JSON response from KAS: {e}") from e
 
         entity_wrapped_key = response_data.get("entityWrappedKey")
         if not entity_wrapped_key:
@@ -621,8 +633,7 @@ class KASClient:
         return self.decryptor.decrypt(encrypted_key)
 
     def unwrap(self, key_access, policy_json, session_key_type=None) -> bytes:
-        """
-        Unwrap a key using Connect RPC.
+        """Unwrap a key using Connect RPC.
 
         Args:
             key_access: Key access information
@@ -631,6 +642,7 @@ class KASClient:
 
         Returns:
             Unwrapped key bytes
+
         """
         # Default to RSA if not specified
         if session_key_type is None:
@@ -655,15 +667,14 @@ class KASClient:
     def _unwrap_with_connect_rpc(
         self, key_access, signed_token, session_key_type=None
     ) -> bytes:
-        """
-        Connect RPC method for unwrapping keys.
+        """Connect RPC method for unwrapping keys.
 
         Args:
             key_access: KeyAccess object
             signed_token: Signed JWT token
             session_key_type: Optional session key type (RSA_KEY_TYPE or EC_KEY_TYPE)
-        """
 
+        """
         # Get access token for authentication if token source is available
         access_token = None
         if self.token_source:
@@ -702,8 +713,8 @@ class KASClient:
 
         except Exception as e:
             logging.error(f"Connect RPC rewrap failed: {e}")
-            raise SDKException(f"Connect RPC rewrap failed: {e}")
+            raise SDKException(f"Connect RPC rewrap failed: {e}") from e
 
     def get_key_cache(self) -> KASKeyCache:
-        """Returns the KAS key cache used for storing and retrieving encryption keys."""
+        """Return the KAS key cache used for storing and retrieving encryption keys."""
         return self.cache

otdf_python/kas_connect_rpc_client.py

@@ -1,13 +1,12 @@
-"""
-KASConnectRPCClient: Handles Connect RPC communication with the Key Access Service (KAS).
+"""KASConnectRPCClient: Handles Connect RPC communication with the Key Access Service (KAS).
 This class encapsulates all interactions with otdf_python_proto.
 """
 
 import logging
 
-import urllib3
+import httpx
 from otdf_python_proto.kas import kas_pb2
-from otdf_python_proto.kas.kas_pb2_connect import AccessServiceClient
+from otdf_python_proto.kas.kas_connect import AccessServiceClientSync
 
 from otdf_python.auth_headers import AuthHeaders
 
@@ -15,45 +14,78 @@ from .sdk_exceptions import SDKException
 
 
 class KASConnectRPCClient:
-    """
-    Handles Connect RPC communication with KAS service using otdf_python_proto.
-    """
+    """Handles Connect RPC communication with KAS service using otdf_python_proto."""
 
     def __init__(self, use_plaintext=False, verify_ssl=True):
-        """
-        Initialize the Connect RPC client.
+        """Initialize the Connect RPC client.
 
         Args:
             use_plaintext: Whether to use plaintext (HTTP) connections
             verify_ssl: Whether to verify SSL certificates
+
         """
         self.use_plaintext = use_plaintext
         self.verify_ssl = verify_ssl
+        self._http_client = None
 
-    def _create_http_client(self):
+    def __enter__(self):
+        """Enter context manager and create HTTP client."""
+        self._http_client = self._create_http_client()
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        """Exit context manager and close HTTP client."""
+        self.close()
+
+    def close(self):
+        """Close HTTP client and release resources.
+
+        This method should be called when the client is no longer needed
+        to properly clean up resources. It's also called automatically
+        when using the client as a context manager.
         """
-        Create HTTP client with SSL verification configuration.
+        if self._http_client is not None:
+            self._http_client.close()
+            self._http_client = None
+
+    def _create_http_client(self):
+        """Create HTTP client with SSL verification configuration.
 
         Returns:
-            urllib3.PoolManager configured for SSL verification settings
+            httpx.Client configured for SSL verification settings
+
         """
         if self.verify_ssl:
             logging.info("Using SSL verification enabled HTTP client")
-            return urllib3.PoolManager()
+            return httpx.Client()
         else:
             logging.info("Using SSL verification disabled HTTP client")
-            urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning)
-            return urllib3.PoolManager(cert_reqs="CERT_NONE")
+            return httpx.Client(verify=False)
+
+    def _get_http_client(self):
+        """Get HTTP client, creating one if needed for backward compatibility.
+
+        Returns:
+            httpx.Client instance
 
-    def _prepare_connect_rpc_url(self, kas_url):
         """
-        Prepare the base URL for Connect RPC client.
+        if self._http_client is None:
+            logging.warning(
+                "KASConnectRPCClient is being used without a context manager. "
+                "Consider using 'with KASConnectRPCClient(...) as client:' to ensure proper resource cleanup."
+            )
+            self._http_client = self._create_http_client()
+        return self._http_client
+
+    def _prepare_connect_rpc_url(self, kas_url):
+        """Prepare the base URL for Connect RPC client.
 
         Args:
             kas_url: The normalized KAS URL
 
         Returns:
             Base URL for Connect RPC client (without /kas suffix)
+
         """
         connect_rpc_base_url = kas_url
         # Remove /kas suffix, if present
@@ -61,14 +93,14 @@ class KASConnectRPCClient:
         return connect_rpc_base_url
 
     def _prepare_auth_headers(self, access_token):
-        """
-        Prepare authentication headers if access token is available.
+        """Prepare authentication headers if access token is available.
 
         Args:
             access_token: Bearer token for authentication
 
         Returns:
             Dictionary with authentication headers or None
+
         """
         if access_token:
             auth_headers = AuthHeaders(
@@ -79,8 +111,7 @@ class KASConnectRPCClient:
         return None
 
     def get_public_key(self, normalized_kas_url, kas_info, access_token=None):
-        """
-        Get KAS public key using Connect RPC.
+        """Get KAS public key using Connect RPC.
 
         Args:
             normalized_kas_url: The normalized KAS URL
@@ -89,6 +120,7 @@ class KASConnectRPCClient:
 
         Returns:
             Updated kas_info with public_key and kid
+
         """
         logging.info(
             f"KAS Connect RPC client settings for public key retrieval: "
@@ -96,8 +128,6 @@ class KASConnectRPCClient:
             f"kas_url={kas_info.url}"
         )
 
-        http_client = self._create_http_client()
-
         try:
             connect_rpc_base_url = self._prepare_connect_rpc_url(normalized_kas_url)
 
@@ -106,9 +136,13 @@ class KASConnectRPCClient:
                 f"for public key retrieval"
             )
 
-            # Create Connect RPC client with configured HTTP client using Connect protocol
-            # Note: gRPC protocol is not supported with urllib3, use default Connect protocol
-            client = AccessServiceClient(connect_rpc_base_url, http_client=http_client)
+            # Get or create HTTP client
+            http_client = self._get_http_client()
+
+            # Create Connect RPC client with configured HTTP client
+            client = AccessServiceClientSync(
+                address=connect_rpc_base_url, session=http_client
+            )
 
             # Create public key request
             algorithm = getattr(kas_info, "algorithm", "") or ""
@@ -119,10 +153,10 @@ class KASConnectRPCClient:
             )
 
             # Prepare headers with authentication if available
-            extra_headers = self._prepare_auth_headers(access_token)
+            headers = self._prepare_auth_headers(access_token)
 
             # Make the public key call with authentication headers
-            response = client.public_key(request, extra_headers=extra_headers)
+            response = client.public_key(request, headers=headers)
 
             # Update kas_info with response
             kas_info.public_key = response.public_key
@@ -138,13 +172,12 @@ class KASConnectRPCClient:
                 f"Connect RPC public key request failed: {type(e).__name__}: {e}"
             )
             logging.error(f"Full traceback: {error_details}")
-            raise SDKException(f"Connect RPC public key request failed: {e}")
+            raise SDKException(f"Connect RPC public key request failed: {e}") from e
 
     def unwrap_key(
         self, normalized_kas_url, key_access, signed_token, access_token=None
     ):
-        """
-        Unwrap a key using Connect RPC.
+        """Unwrap a key using Connect RPC.
 
         Args:
             normalized_kas_url: The normalized KAS URL
@@ -154,6 +187,7 @@ class KASConnectRPCClient:
 
         Returns:
             Unwrapped key bytes from the response
+
         """
         logging.info(
             f"KAS Connect RPC client settings for unwrap: "
@@ -161,8 +195,6 @@ class KASConnectRPCClient:
             f"kas_url={key_access.url}"
         )
 
-        http_client = self._create_http_client()
-
         try:
             kas_service_url = self._prepare_connect_rpc_url(normalized_kas_url)
 
@@ -170,8 +202,13 @@ class KASConnectRPCClient:
                 f"Creating Connect RPC client for base URL: {kas_service_url}, for unwrap"
             )
 
-            # Note: gRPC protocol is not supported with urllib3, use default Connect protocol
-            client = AccessServiceClient(kas_service_url, http_client=http_client)
+            # Get or create HTTP client
+            http_client = self._get_http_client()
+
+            # Create Connect RPC client with configured HTTP client
+            client = AccessServiceClientSync(
+                address=kas_service_url, session=http_client
+            )
 
             # Create rewrap request
             request = kas_pb2.RewrapRequest(
@@ -182,10 +219,10 @@ class KASConnectRPCClient:
             logging.info(f"Connect RPC signed token: {signed_token}")
 
             # Prepare headers with authentication if available
-            extra_headers = self._prepare_auth_headers(access_token)
+            headers = self._prepare_auth_headers(access_token)
 
             # Make the rewrap call with authentication headers
-            response = client.rewrap(request, extra_headers=extra_headers)
+            response = client.rewrap(request, headers=headers)
 
             # Extract the entity wrapped key from v2 response structure
             # The v2 response has responses[] array with results[] for each policy
@@ -210,4 +247,4 @@ class KASConnectRPCClient:
 
         except Exception as e:
             logging.error(f"Connect RPC rewrap failed: {e}")
-            raise SDKException(f"Connect RPC rewrap failed: {e}")
+            raise SDKException(f"Connect RPC rewrap failed: {e}") from e

otdf_python/kas_info.py

@@ -1,10 +1,11 @@
+"""Key Access Service information and configuration."""
+
 from dataclasses import dataclass
 
 
 @dataclass
 class KASInfo:
-    """
-    Configuration for Key Access Server (KAS) information.
+    """Configuration for Key Access Server (KAS) information.
     This class stores details about a Key Access Server including its URL,
     public key, key ID, default status, and cryptographic algorithm.
     """
@@ -16,7 +17,7 @@ class KASInfo:
     algorithm: str | None = None
 
     def clone(self):
-        """Creates a copy of this KASInfo object."""
+        """Create a copy of this KASInfo object."""
         from copy import copy
 
         return copy(self)