otdf-python 0.4.1__py3-none-any.whl → 0.4.3__py3-none-any.whl

otdf_python/autoconfigure_utils.py

@@ -39,8 +39,6 @@ class KeySplitStep:
 class AutoConfigureException(Exception):
     """Exception for auto-configuration errors."""
 
-    pass
-
 
 class AttributeNameFQN:
     """Fully qualified attribute name."""

otdf_python/cli.py

@@ -112,33 +112,14 @@ def load_client_credentials(creds_file_path: str) -> tuple[str, str]:
         ) from e
 
 
-def build_sdk(args) -> SDK:
-    """Build SDK instance from CLI arguments."""
-    builder = SDKBuilder()
-
-    if args.platform_url:
-        builder.set_platform_endpoint(args.platform_url)
-
-        # Auto-detect HTTP URLs and enable plaintext mode
-        if args.platform_url.startswith("http://") and (
-            not hasattr(args, "plaintext") or not args.plaintext
-        ):
-            logger.debug(
-                f"Auto-detected HTTP URL {args.platform_url}, enabling plaintext mode"
-            )
-            builder.use_insecure_plaintext_connection(True)
-
-    if args.oidc_endpoint:
-        builder.set_issuer_endpoint(args.oidc_endpoint)
-
+def _configure_auth(builder: SDKBuilder, args) -> None:
+    """Configure authentication on the SDK builder."""
     if args.client_id and args.client_secret:
         builder.client_secret(args.client_id, args.client_secret)
     elif hasattr(args, "with_client_creds_file") and args.with_client_creds_file:
-        # Load credentials from file
         client_id, client_secret = load_client_credentials(args.with_client_creds_file)
         builder.client_secret(client_id, client_secret)
     elif hasattr(args, "auth") and args.auth:
-        # Parse combined auth string (clientId:clientSecret) - legacy support
         auth_parts = args.auth.split(":")
         if len(auth_parts) != 2:
             raise CLIError(
@@ -152,12 +133,49 @@ def build_sdk(args) -> SDK:
             "Authentication required: provide --with-client-creds-file OR --client-id and --client-secret",
         )
 
+
+def _configure_kas_allowlist(builder: SDKBuilder, args) -> None:
+    """Configure KAS allowlist on the SDK builder."""
+    if hasattr(args, "ignore_kas_allowlist") and args.ignore_kas_allowlist:
+        logger.warning(
+            "KAS allowlist validation is disabled. This may leak credentials "
+            "to malicious servers if decrypting untrusted TDF files."
+        )
+        builder.with_ignore_kas_allowlist(True)
+    elif hasattr(args, "kas_allowlist") and args.kas_allowlist:
+        kas_urls = [url.strip() for url in args.kas_allowlist.split(",") if url.strip()]
+        logger.debug(f"Using KAS allowlist: {kas_urls}")
+        builder.with_kas_allowlist(kas_urls)
+
+
+def build_sdk(args) -> SDK:
+    """Build SDK instance from CLI arguments."""
+    builder = SDKBuilder()
+
+    if args.platform_url:
+        builder.set_platform_endpoint(args.platform_url)
+        # Auto-detect HTTP URLs and enable plaintext mode
+        if args.platform_url.startswith("http://") and (
+            not hasattr(args, "plaintext") or not args.plaintext
+        ):
+            logger.debug(
+                f"Auto-detected HTTP URL {args.platform_url}, enabling plaintext mode"
+            )
+            builder.use_insecure_plaintext_connection(True)
+
+    if args.oidc_endpoint:
+        builder.set_issuer_endpoint(args.oidc_endpoint)
+
+    _configure_auth(builder, args)
+
     if hasattr(args, "plaintext") and args.plaintext:
         builder.use_insecure_plaintext_connection(True)
 
     if args.insecure:
         builder.use_insecure_skip_verify(True)
 
+    _configure_kas_allowlist(builder, args)
+
     return builder.build()
 
 
@@ -476,6 +494,17 @@ Where creds.json contains:
     security_group.add_argument(
         "--insecure", action="store_true", help="Skip TLS verification"
     )
+    security_group.add_argument(
+        "--kas-allowlist",
+        help="Comma-separated list of trusted KAS URLs. "
+        "By default, only the platform URL's KAS endpoint is trusted.",
+    )
+    security_group.add_argument(
+        "--ignore-kas-allowlist",
+        action="store_true",
+        help="WARNING: Disable KAS allowlist validation. This is insecure and "
+        "should only be used for testing. May leak credentials to malicious servers.",
+    )
 
     # Subcommands
     subparsers = parser.add_subparsers(dest="command", help="Available commands")

otdf_python/collection_store.py

@@ -28,7 +28,6 @@ class NoOpCollectionStore(CollectionStore):
 
     def store(self, header, key: CollectionKey):
         """Discard key operation (no-op)."""
-        pass
 
     def get_key(self, header) -> CollectionKey:
         return self.NO_PRIVATE_KEY

otdf_python/ecdh.py

@@ -35,20 +35,14 @@ NANOTDF_HKDF_SALT = bytes.fromhex(
 class ECDHError(Exception):
     """Base exception for ECDH operations."""
 
-    pass
-
 
 class UnsupportedCurveError(ECDHError):
     """Raised when an unsupported curve is specified."""
 
-    pass
-
 
 class InvalidKeyError(ECDHError):
     """Raised when a key is invalid or malformed."""
 
-    pass
-
 
 def get_curve(curve_name: str) -> ec.EllipticCurve:
     """Get the cryptography curve object for a given curve name.

otdf_python/kas_allowlist.py

@@ -0,0 +1,182 @@
+"""KAS Allowlist: Validates KAS URLs against a list of trusted hosts.
+
+This module provides protection against SSRF attacks where malicious TDF files
+could contain attacker-controlled KAS URLs to steal OIDC credentials.
+"""
+
+import logging
+from urllib.parse import urlparse
+
+
+class KASAllowlist:
+    """Validates KAS URLs against an allowlist of trusted hosts.
+
+    This class prevents credential theft by ensuring the SDK only sends
+    authentication tokens to trusted KAS endpoints.
+
+    Example:
+        allowlist = KASAllowlist(["https://kas.example.com"])
+        allowlist.is_allowed("https://kas.example.com/kas")  # True
+        allowlist.is_allowed("https://evil.com/kas")  # False
+
+    """
+
+    def __init__(self, allowed_urls: list[str] | None = None, allow_all: bool = False):
+        """Initialize the KAS allowlist.
+
+        Args:
+            allowed_urls: List of trusted KAS URLs. Each URL is normalized to
+                its origin (scheme://host:port) for comparison.
+            allow_all: If True, all URLs are allowed. Use only for testing.
+                A warning is logged when this is enabled.
+
+        """
+        self._allowed_origins: set[str] = set()
+        self._allow_all = allow_all
+
+        if allow_all:
+            logging.warning(
+                "KAS allowlist is disabled (allow_all=True). "
+                "This is insecure and should only be used for testing."
+            )
+
+        if allowed_urls:
+            for url in allowed_urls:
+                self.add(url)
+
+    def add(self, url: str) -> None:
+        """Add a URL to the allowlist.
+
+        The URL is normalized to its origin (scheme://host:port) before storage.
+        Paths and query strings are stripped.
+
+        Args:
+            url: The KAS URL to allow. Can include path components which
+                will be stripped for origin comparison.
+
+        """
+        origin = self._get_origin(url)
+        self._allowed_origins.add(origin)
+        logging.debug(f"Added KAS origin to allowlist: {origin}")
+
+    def is_allowed(self, url: str) -> bool:
+        """Check if a URL is allowed by the allowlist.
+
+        Args:
+            url: The KAS URL to check.
+
+        Returns:
+            True if the URL's origin is in the allowlist or allow_all is True.
+            False otherwise.
+
+        """
+        if self._allow_all:
+            logging.debug(f"KAS URL allowed (allow_all=True): {url}")
+            return True
+
+        if not self._allowed_origins:
+            logging.debug(f"KAS URL rejected (empty allowlist): {url}")
+            return False
+
+        origin = self._get_origin(url)
+        allowed = origin in self._allowed_origins
+        if allowed:
+            logging.debug(f"KAS URL allowed: {url} (origin: {origin})")
+        else:
+            logging.debug(
+                f"KAS URL rejected: {url} (origin: {origin}, "
+                f"allowed: {self._allowed_origins})"
+            )
+        return allowed
+
+    def validate(self, url: str) -> None:
+        """Validate a URL against the allowlist, raising an exception if not allowed.
+
+        Args:
+            url: The KAS URL to validate.
+
+        Raises:
+            SDK.KasAllowlistException: If the URL is not in the allowlist.
+
+        """
+        if not self.is_allowed(url):
+            # Import here to avoid circular imports
+            from .sdk import SDK
+
+            raise SDK.KasAllowlistException(url, self._allowed_origins)
+
+    @property
+    def allowed_origins(self) -> set[str]:
+        """Return the set of allowed origins (read-only copy)."""
+        return self._allowed_origins.copy()
+
+    @property
+    def allow_all(self) -> bool:
+        """Return whether all URLs are allowed."""
+        return self._allow_all
+
+    @staticmethod
+    def _get_origin(url: str) -> str:
+        """Extract the origin (scheme://host:port) from a URL.
+
+        This normalizes URLs for comparison by stripping paths and query strings.
+        Default ports (80 for http, 443 for https) are included explicitly.
+
+        Args:
+            url: The URL to extract the origin from.
+
+        Returns:
+            Normalized origin string in format scheme://host:port
+
+        """
+        # Add scheme if missing
+        if "://" not in url:
+            url = "https://" + url
+
+        try:
+            parsed = urlparse(url)
+        except Exception as e:
+            logging.warning(f"Failed to parse URL {url}: {e}")
+            # Return the URL as-is if parsing fails
+            return url.lower()
+
+        scheme = (parsed.scheme or "https").lower()
+        hostname = (parsed.hostname or "").lower()
+
+        if not hostname:
+            # URL might be malformed, return as-is
+            logging.warning(f"Could not extract hostname from URL: {url}")
+            return url.lower()
+
+        # Determine port (use explicit port or default for scheme)
+        if parsed.port:
+            port = parsed.port
+        elif scheme == "http":
+            port = 80
+        else:
+            port = 443
+
+        return f"{scheme}://{hostname}:{port}"
+
+    @classmethod
+    def from_platform_url(cls, platform_url: str) -> "KASAllowlist":
+        """Create an allowlist from a platform URL.
+
+        This is the default behavior: auto-allow the platform's KAS endpoint.
+
+        Args:
+            platform_url: The OpenTDF platform URL. The KAS endpoint is
+                assumed to be at {platform_url}/kas.
+
+        Returns:
+            KASAllowlist configured to allow the platform's KAS endpoint.
+
+        """
+        allowlist = cls()
+        # Add the platform URL itself (KAS might be at root or /kas)
+        allowlist.add(platform_url)
+        # Also construct the /kas endpoint explicitly
+        kas_url = platform_url.rstrip("/") + "/kas"
+        allowlist.add(kas_url)
+        logging.info(f"Created KAS allowlist from platform URL: {platform_url}")
+        return allowlist

otdf_python/kas_client.py

@@ -38,13 +38,26 @@ class KASClient:
         cache=None,
         use_plaintext=False,
         verify_ssl=True,
+        kas_allowlist=None,
     ):
-        """Initialize KAS client."""
+        """Initialize KAS client.
+
+        Args:
+            kas_url: Default KAS URL
+            token_source: Function that returns an authentication token
+            cache: Optional KASKeyCache for caching public keys
+            use_plaintext: Whether to use HTTP instead of HTTPS
+            verify_ssl: Whether to verify SSL certificates
+            kas_allowlist: Optional KASAllowlist for URL validation. If provided,
+                only URLs in the allowlist will be contacted.
+
+        """
         self.kas_url = kas_url
         self.token_source = token_source
         self.cache = cache or KASKeyCache()
         self.use_plaintext = use_plaintext
         self.verify_ssl = verify_ssl
+        self.kas_allowlist = kas_allowlist
         self.decryptor = None
         self.client_public_key = None
 
@@ -65,18 +78,47 @@ class KASClient:
             self._dpop_public_key
         )
 
+    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.
+        """
+        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.
 
+        This method also validates the URL against the KAS allowlist if one
+        is configured. This prevents SSRF attacks where malicious TDF files
+        could contain attacker-controlled KAS URLs to steal OIDC credentials.
+
         Args:
             url: The KAS URL to normalize
 
         Returns:
             Normalized URL with appropriate protocol and port
 
+        Raises:
+            KASAllowlistException: If the URL is not in the allowlist
+
         """
         from urllib.parse import urlparse
 
+        # Validate against allowlist BEFORE making any requests
+        if self.kas_allowlist is not None:
+            self.kas_allowlist.validate(url)
+
         try:
             # Parse the URL
             parsed = urlparse(url)
@@ -136,7 +178,7 @@ class KASClient:
         # Reconstruct URL preserving the path (especially /kas prefix)
         try:
             # Create URL preserving the path component for proper endpoint routing
-            path = parsed.path if parsed.path else ""
+            path = parsed.path or ""
             normalized_url = f"{scheme}://{parsed.hostname}:{port}{path}"
             logging.debug(f"normalized url [{parsed.geturl()}] to [{normalized_url}]")
             return normalized_url

otdf_python/kas_connect_rpc_client.py

@@ -4,9 +4,9 @@ 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
 
@@ -26,21 +26,56 @@ class KASConnectRPCClient:
         """
         self.use_plaintext = use_plaintext
         self.verify_ssl = verify_ssl
+        self._http_client = None
+
+    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.
+        """
+        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
+
+        """
+        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.
@@ -93,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)
 
@@ -103,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 ""
@@ -116,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
@@ -158,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)
 
@@ -167,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(
@@ -179,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

otdf_python/nanotdf.py

@@ -28,26 +28,18 @@ from .asym_crypto import AsymDecryption
 class NanoTDFException(SDKException):
     """Base exception for NanoTDF operations."""
 
-    pass
-
 
 class NanoTDFMaxSizeLimit(NanoTDFException):
     """Exception for NanoTDF size limit exceeded."""
 
-    pass
-
 
 class UnsupportedNanoTDFFeature(NanoTDFException):
     """Exception for unsupported NanoTDF features."""
 
-    pass
-
 
 class InvalidNanoTDFConfig(NanoTDFException):
     """Exception for invalid NanoTDF configuration."""
 
-    pass
-
 
 class NanoTDF:
     """NanoTDF reader and writer for compact TDF format."""
@@ -78,8 +70,8 @@ class NanoTDF:
         if isinstance(obj, PolicyBody):
             # Convert data_attributes to dataAttributes and use null instead of empty array
             result = {
-                "dataAttributes": obj.data_attributes if obj.data_attributes else None,
-                "dissem": obj.dissem if obj.dissem else None,
+                "dataAttributes": obj.data_attributes or None,
+                "dissem": obj.dissem or None,
             }
             return result
         elif isinstance(obj, AttributeObject):
@@ -123,14 +115,12 @@ class NanoTDF:
             tuple: (policy_body, policy_type)
 
         """
-        attributes = config.attributes if config.attributes else []
+        attributes = config.attributes or []
         policy_object = self._create_policy_object(attributes)
         policy_json = json.dumps(
             policy_object, default=self._serialize_policy_object
         ).encode("utf-8")
-        policy_type = (
-            config.policy_type if config.policy_type else "EMBEDDED_POLICY_PLAIN_TEXT"
-        )
+        policy_type = config.policy_type or "EMBEDDED_POLICY_PLAIN_TEXT"
 
         if policy_type == "EMBEDDED_POLICY_PLAIN_TEXT":
             policy_body = policy_json

otdf_python/nanotdf_ecdsa_struct.py

@@ -6,8 +6,6 @@ from dataclasses import dataclass, field
 class IncorrectNanoTDFECDSASignatureSize(Exception):
     """Exception raised when the signature size is incorrect."""
 
-    pass
-
 
 @dataclass
 class NanoTDFECDSAStruct:

otdf_python/nanotdf_type.py

@@ -8,7 +8,7 @@ class ECCurve(Enum):
 
     SECP256R1 = "secp256r1"
     SECP384R1 = "secp384r1"
-    SECP521R1 = "secp384r1"
+    SECP521R1 = "secp521r1"
     SECP256K1 = "secp256k1"
 
     def __str__(self):