otdf-python 0.3.5__py3-none-any.whl → 0.4.1__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()
@@ -63,14 +66,14 @@ class KASClient:
         )
 
     def _normalize_kas_url(self, url: str) -> str:
-        """
-        Normalize KAS URLs based on client security settings.
+        """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 +81,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 +103,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 +119,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 +141,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 +169,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 +200,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 +247,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 +265,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 +276,7 @@ class KASClient:
 
         Returns:
             Dictionary with unsigned rewrap request
+
         """
         import json
 
@@ -316,8 +319,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 +327,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 +357,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 +365,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 +374,7 @@ class KASClient:
 
         Returns:
             DPoP proof JWT string
+
         """
         now = int(time.time())
 
@@ -424,8 +426,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 +449,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 +481,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 +509,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 +526,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 +541,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 +550,7 @@ class KASClient:
 
         Returns:
             Unwrapped key as bytes
+
         """
         if ec_key_pair is None:
             raise SDKException(
@@ -581,9 +579,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 +596,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 +615,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 +624,7 @@ class KASClient:
 
         Returns:
             Unwrapped key bytes
+
         """
         # Default to RSA if not specified
         if session_key_type is None:
@@ -655,15 +649,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 +695,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,5 +1,4 @@
-"""
-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.
 """
 
@@ -15,27 +14,25 @@ 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
 
     def _create_http_client(self):
-        """
-        Create HTTP client with SSL verification configuration.
+        """Create HTTP client with SSL verification configuration.
 
         Returns:
             urllib3.PoolManager configured for SSL verification settings
+
         """
         if self.verify_ssl:
             logging.info("Using SSL verification enabled HTTP client")
@@ -46,14 +43,14 @@ class KASConnectRPCClient:
             return urllib3.PoolManager(cert_reqs="CERT_NONE")
 
     def _prepare_connect_rpc_url(self, kas_url):
-        """
-        Prepare the base URL for Connect RPC client.
+        """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 +58,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 +76,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 +85,7 @@ class KASConnectRPCClient:
 
         Returns:
             Updated kas_info with public_key and kid
+
         """
         logging.info(
             f"KAS Connect RPC client settings for public key retrieval: "
@@ -138,13 +135,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 +150,7 @@ class KASConnectRPCClient:
 
         Returns:
             Unwrapped key bytes from the response
+
         """
         logging.info(
             f"KAS Connect RPC client settings for unwrap: "
@@ -210,4 +207,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)

otdf_python/kas_key_cache.py

@@ -1,19 +1,19 @@
-"""
-KASKeyCache: In-memory cache for KAS (Key Access Service) public keys and info.
-"""
+"""KASKeyCache: In-memory cache for KAS (Key Access Service) public keys and info."""
 
 import threading
 from typing import Any
 
 
 class KASKeyCache:
+    """In-memory cache for KAS public keys and information."""
+
     def __init__(self):
+        """Initialize KAS key cache."""
         self._cache = {}
         self._lock = threading.Lock()
 
     def get(self, url: str, algorithm: str | None = None) -> Any | None:
-        """
-        Gets a KASInfo object from the cache based on URL and algorithm.
+        """Get a KASInfo object from cache based on URL and algorithm.
 
         Args:
             url: The URL of the KAS
@@ -21,17 +21,18 @@ class KASKeyCache:
 
         Returns:
             The cached KASInfo object, or None if not found
+
         """
         cache_key = self._make_key(url, algorithm)
         with self._lock:
             return self._cache.get(cache_key)
 
     def store(self, kas_info) -> None:
-        """
-        Stores a KASInfo object in the cache.
+        """Store a KASInfo object in cache.
 
         Args:
             kas_info: The KASInfo object to store
+
         """
         cache_key = self._make_key(kas_info.url, getattr(kas_info, "algorithm", None))
         with self._lock:
@@ -43,10 +44,10 @@ class KASKeyCache:
             self._cache[key] = value
 
     def clear(self):
-        """Clears the cache"""
+        """Clear the cache."""
         with self._lock:
             self._cache.clear()
 
     def _make_key(self, url: str, algorithm: str | None = None) -> str:
-        """Creates a cache key from URL and algorithm"""
+        """Create a cache key from URL and algorithm."""
         return f"{url}:{algorithm or ''}"

otdf_python/key_type.py

@@ -1,7 +1,11 @@
+"""Key type constants for RSA and EC encryption."""
+
 from enum import Enum
 
 
 class KeyType(Enum):
+    """Key type enumeration for encryption algorithms."""
+
     RSA2048Key = "rsa:2048"
     EC256Key = "ec:secp256r1"
     EC384Key = "ec:secp384r1"

otdf_python/key_type_constants.py

@@ -1,5 +1,4 @@
-"""
-Constants for session key types used in the KAS client.
+"""Constants for session key types used in the KAS client.
 This matches the Java SDK's KeyType enum pattern.
 """
 
@@ -7,9 +6,7 @@ from enum import Enum, auto
 
 
 class KeyType(Enum):
-    """
-    Enum for key types used in the KAS client.
-    """
+    """Enum for key types used in the KAS client."""
 
     RSA2048 = auto()
     EC_P256 = auto()
@@ -18,16 +15,12 @@ class KeyType(Enum):
 
     @property
     def is_ec(self):
-        """
-        Returns True if this key type is an EC key, False otherwise.
-        """
+        """Returns True if this key type is an EC key, False otherwise."""
         return self in [KeyType.EC_P256, KeyType.EC_P384, KeyType.EC_P521]
 
     @property
     def curve_name(self):
-        """
-        Returns the curve name for EC keys.
-        """
+        """Returns the curve name for EC keys."""
         if self == KeyType.EC_P256:
             return "P-256"
         elif self == KeyType.EC_P384: