otdf-python 0.1.9__py3-none-any.whl → 0.3.0__py3-none-any.whl

otdf_python/kas_connect_rpc_client.py

@@ -0,0 +1,207 @@
+"""
+KASConnectRPCClient: Handles Connect RPC communication with the Key Access Service (KAS).
+This class encapsulates all interactions with otdf_python_proto.
+"""
+
+import logging
+
+import urllib3
+from otdf_python_proto.kas import kas_pb2
+from otdf_python_proto.kas.kas_pb2_connect import AccessServiceClient
+
+from .sdk_exceptions import SDKException
+
+
+class KASConnectRPCClient:
+    """
+    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.
+
+        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.
+
+        Returns:
+            urllib3.PoolManager configured for SSL verification settings
+        """
+        if self.verify_ssl:
+            logging.info("Using SSL verification enabled HTTP client")
+            return urllib3.PoolManager()
+        else:
+            logging.info("Using SSL verification disabled HTTP client")
+            urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning)
+            return urllib3.PoolManager(cert_reqs="CERT_NONE")
+
+    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
+        connect_rpc_base_url = connect_rpc_base_url.removesuffix("/kas")
+        return connect_rpc_base_url
+
+    def _prepare_auth_headers(self, access_token):
+        """
+        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:
+            return {"Authorization": f"Bearer {access_token}"}
+        return None
+
+    def get_public_key(self, normalized_kas_url, kas_info, access_token=None):
+        """
+        Get KAS public key using Connect RPC.
+
+        Args:
+            normalized_kas_url: The normalized KAS URL
+            kas_info: KAS information object with algorithm
+            access_token: Optional access token for authentication
+
+        Returns:
+            Updated kas_info with public_key and kid
+        """
+        logging.info(
+            f"KAS Connect RPC client settings for public key retrieval: "
+            f"verify_ssl={self.verify_ssl}, use_plaintext={self.use_plaintext}, "
+            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)
+
+            logging.info(
+                f"Creating Connect RPC client for base URL: {connect_rpc_base_url}, "
+                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)
+
+            # Create public key request
+            algorithm = getattr(kas_info, "algorithm", "") or ""
+            request = (
+                kas_pb2.PublicKeyRequest(algorithm=algorithm)
+                if algorithm
+                else kas_pb2.PublicKeyRequest()
+            )
+
+            # Prepare headers with authentication if available
+            extra_headers = self._prepare_auth_headers(access_token)
+
+            # Make the public key call with authentication headers
+            response = client.public_key(request, extra_headers=extra_headers)
+
+            # Update kas_info with response
+            kas_info.public_key = response.public_key
+            kas_info.kid = response.kid
+
+            return kas_info
+
+        except Exception as e:
+            import traceback
+
+            error_details = traceback.format_exc()
+            logging.error(
+                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}")
+
+    def unwrap_key(
+        self, normalized_kas_url, key_access, signed_token, access_token=None
+    ):
+        """
+        Unwrap a key using Connect RPC.
+
+        Args:
+            normalized_kas_url: The normalized KAS URL
+            key_access: Key access information
+            signed_token: Signed JWT token for the request
+            access_token: Optional access token for authentication
+
+        Returns:
+            Unwrapped key bytes from the response
+        """
+        logging.info(
+            f"KAS Connect RPC client settings for unwrap: "
+            f"verify_ssl={self.verify_ssl}, use_plaintext={self.use_plaintext}, "
+            f"kas_url={key_access.url}"
+        )
+
+        http_client = self._create_http_client()
+
+        try:
+            kas_service_url = self._prepare_connect_rpc_url(normalized_kas_url)
+
+            logging.info(
+                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)
+
+            # Create rewrap request
+            request = kas_pb2.RewrapRequest(
+                signed_request_token=signed_token,
+            )
+
+            # Debug: Log the signed token details
+            logging.info(f"Connect RPC signed token: {signed_token}")
+
+            # Prepare headers with authentication if available
+            extra_headers = self._prepare_auth_headers(access_token)
+
+            # Make the rewrap call with authentication headers
+            response = client.rewrap(request, extra_headers=extra_headers)
+
+            # Extract the entity wrapped key from v2 response structure
+            # The v2 response has responses[] array with results[] for each policy
+            if response.responses and len(response.responses) > 0:
+                policy_result = response.responses[0]  # First policy
+                if policy_result.results and len(policy_result.results) > 0:
+                    kao_result = policy_result.results[0]  # First KAO result
+                    if kao_result.kas_wrapped_key:
+                        entity_wrapped_key = kao_result.kas_wrapped_key
+                    else:
+                        raise SDKException(f"KAO result error: {kao_result.error}")
+                else:
+                    raise SDKException("No KAO results in policy response")
+            else:
+                # Fallback to legacy entity_wrapped_key field for backward compatibility
+                entity_wrapped_key = response.entity_wrapped_key
+                if not entity_wrapped_key:
+                    raise SDKException("No entity_wrapped_key in Connect RPC response")
+
+            logging.info("Connect RPC rewrap succeeded")
+            return entity_wrapped_key
+
+        except Exception as e:
+            logging.error(f"Connect RPC rewrap failed: {e}")
+            raise SDKException(f"Connect RPC rewrap failed: {e}")

otdf_python/kas_info.py

@@ -0,0 +1,25 @@
+from dataclasses import dataclass
+
+
+@dataclass
+class KASInfo:
+    """
+    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.
+    """
+
+    url: str
+    public_key: str | None = None
+    kid: str | None = None
+    default: bool | None = None
+    algorithm: str | None = None
+
+    def clone(self):
+        """Creates a copy of this KASInfo object."""
+        from copy import copy
+
+        return copy(self)
+
+    def __str__(self):
+        return f"KASInfo(url={self.url}, kid={self.kid}, default={self.default}, algorithm={self.algorithm})"

otdf_python/kas_key_cache.py

@@ -0,0 +1,52 @@
+"""
+KASKeyCache: In-memory cache for KAS (Key Access Service) public keys and info.
+"""
+
+import threading
+from typing import Any
+
+
+class KASKeyCache:
+    def __init__(self):
+        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.
+
+        Args:
+            url: The URL of the KAS
+            algorithm: Optional algorithm identifier
+
+        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.
+
+        Args:
+            kas_info: The KASInfo object to store
+        """
+        cache_key = self._make_key(kas_info.url, getattr(kas_info, "algorithm", None))
+        with self._lock:
+            self._cache[cache_key] = kas_info
+
+    def set(self, key, value):
+        """Store a key-value pair in the cache."""
+        with self._lock:
+            self._cache[key] = value
+
+    def clear(self):
+        """Clears 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"""
+        return f"{url}:{algorithm or ''}"

otdf_python/key_type.py

@@ -0,0 +1,31 @@
+from enum import Enum
+
+
+class KeyType(Enum):
+    RSA2048Key = "rsa:2048"
+    EC256Key = "ec:secp256r1"
+    EC384Key = "ec:secp384r1"
+    EC521Key = "ec:secp521r1"
+
+    def __str__(self):
+        return self.value
+
+    def get_curve_name(self):
+        if self == KeyType.EC256Key:
+            return "secp256r1"
+        elif self == KeyType.EC384Key:
+            return "secp384r1"
+        elif self == KeyType.EC521Key:
+            return "secp521r1"
+        else:
+            raise ValueError(f"Unsupported key type: {self}")
+
+    @staticmethod
+    def from_string(key_type):
+        for t in KeyType:
+            if t.value.lower() == key_type.lower():
+                return t
+        raise ValueError(f"No enum constant for key type: {key_type}")
+
+    def is_ec(self):
+        return self != KeyType.RSA2048Key

otdf_python/key_type_constants.py

@@ -0,0 +1,43 @@
+"""
+Constants for session key types used in the KAS client.
+This matches the Java SDK's KeyType enum pattern.
+"""
+
+from enum import Enum, auto
+
+
+class KeyType(Enum):
+    """
+    Enum for key types used in the KAS client.
+    """
+
+    RSA2048 = auto()
+    EC_P256 = auto()
+    EC_P384 = auto()
+    EC_P521 = auto()
+
+    @property
+    def is_ec(self):
+        """
+        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.
+        """
+        if self == KeyType.EC_P256:
+            return "P-256"
+        elif self == KeyType.EC_P384:
+            return "P-384"
+        elif self == KeyType.EC_P521:
+            return "P-521"
+        else:
+            return None
+
+
+# Constants for backward compatibility with string literals
+RSA_KEY_TYPE = KeyType.RSA2048
+EC_KEY_TYPE = KeyType.EC_P256  # Default EC curve

otdf_python/manifest.py

@@ -0,0 +1,215 @@
+import json
+from dataclasses import asdict, dataclass, field
+from typing import Any
+
+
+@dataclass
+class ManifestSegment:
+    hash: str
+    segmentSize: int
+    encryptedSegmentSize: int
+
+
+@dataclass
+class ManifestRootSignature:
+    alg: str
+    sig: str
+
+
+@dataclass
+class ManifestIntegrityInformation:
+    rootSignature: ManifestRootSignature
+    segmentHashAlg: str
+    segmentSizeDefault: int
+    encryptedSegmentSizeDefault: int
+    segments: list[ManifestSegment]
+
+
+@dataclass
+class ManifestPolicyBinding:
+    alg: str
+    hash: str
+
+
+@dataclass
+class ManifestKeyAccess:
+    type: str
+    url: str
+    protocol: str
+    wrappedKey: str
+    policyBinding: Any = None
+    encryptedMetadata: str | None = None
+    kid: str | None = None
+    sid: str | None = None
+    schemaVersion: str | None = None
+    ephemeralPublicKey: str | None = None
+
+
+@dataclass
+class ManifestMethod:
+    algorithm: str
+    iv: str
+    isStreamable: bool | None = None
+
+
+@dataclass
+class ManifestEncryptionInformation:
+    type: str
+    policy: str
+    keyAccess: list[ManifestKeyAccess]
+    method: ManifestMethod
+    integrityInformation: ManifestIntegrityInformation
+
+
+@dataclass
+class ManifestPayload:
+    type: str
+    url: str
+    protocol: str
+    mimeType: str
+    isEncrypted: bool
+
+
+@dataclass
+class ManifestBinding:
+    method: str
+    signature: str
+
+
+@dataclass
+class ManifestAssertion:
+    id: str
+    type: str
+    scope: str
+    appliesTo_state: str
+    statement: Any
+    binding: ManifestBinding | None = None
+
+
+@dataclass
+class Manifest:
+    schemaVersion: str | None = None
+    encryptionInformation: ManifestEncryptionInformation | None = None
+    payload: ManifestPayload | None = None
+    assertions: list[ManifestAssertion] = field(default_factory=list)
+
+    def _remove_none_values_and_empty_lists(self, obj):
+        """Recursively remove None values and empty lists from dictionaries and lists."""
+        if isinstance(obj, dict):
+            cleaned = {}
+            for k, v in obj.items():
+                if v is not None:
+                    # For 'assertions' field, exclude if it's an empty list
+                    if k == "assertions" and isinstance(v, list) and len(v) == 0:
+                        continue
+                    cleaned[k] = self._remove_none_values_and_empty_lists(v)
+            return cleaned
+        elif isinstance(obj, list):
+            return [
+                self._remove_none_values_and_empty_lists(item)
+                for item in obj
+                if item is not None
+            ]
+        else:
+            return obj
+
+    def to_json(self) -> str:
+        # Create manifest dict with fields ordered to match otdfctl expectations
+        # Order: encryptionInformation, payload, schemaVersion, assertions
+        manifest_dict = {}
+
+        # Add fields in the order expected by otdfctl
+        if self.encryptionInformation is not None:
+            manifest_dict["encryptionInformation"] = asdict(self.encryptionInformation)
+
+        if self.payload is not None:
+            manifest_dict["payload"] = asdict(self.payload)
+
+        if self.schemaVersion is not None:
+            manifest_dict["schemaVersion"] = self.schemaVersion
+
+        if self.assertions and len(self.assertions) > 0:
+            manifest_dict["assertions"] = [
+                asdict(assertion) for assertion in self.assertions
+            ]
+
+        cleaned_dict = self._remove_none_values_and_empty_lists(manifest_dict)
+        return json.dumps(cleaned_dict, default=str)
+
+    @staticmethod
+    def from_json(data: str) -> "Manifest":
+        d = json.loads(data)
+
+        # Recursively instantiate nested dataclasses
+        def _payload(p):
+            return ManifestPayload(**p) if p else None
+
+        def _segment(s):
+            return ManifestSegment(**s)
+
+        def _root_sig(rs):
+            return ManifestRootSignature(**rs)
+
+        def _integrity(i):
+            # Handle both snake_case and camelCase fields
+            # TODO: This can probably be simplified to only camelCase
+            return ManifestIntegrityInformation(
+                rootSignature=_root_sig(
+                    i.get("rootSignature", i.get("root_signature"))
+                ),
+                segmentHashAlg=i.get("segmentHashAlg", i.get("segment_hash_alg")),
+                segmentSizeDefault=i.get(
+                    "segmentSizeDefault", i.get("segment_size_default")
+                ),
+                encryptedSegmentSizeDefault=i.get(
+                    "encryptedSegmentSizeDefault",
+                    i.get("encrypted_segment_size_default"),
+                ),
+                segments=[_segment(s) for s in i["segments"]],
+            )
+
+        def _method(m):
+            return ManifestMethod(**m)
+
+        def _key_access(k):
+            return ManifestKeyAccess(**k)
+
+        def _enc_info(e):
+            # Handle both snake_case and camelCase fields
+            # TODO: This can probably be simplified to only camelCase
+            return ManifestEncryptionInformation(
+                type=e.get("type", e.get("key_access_type", "split")),
+                policy=e["policy"],
+                keyAccess=[
+                    _key_access(k)
+                    for k in e.get("keyAccess", e.get("key_access_obj", []))
+                ],
+                method=_method(e["method"]),
+                integrityInformation=_integrity(
+                    e.get("integrityInformation", e.get("integrity_information"))
+                ),
+            )
+
+        def _binding(b):
+            return ManifestBinding(**b) if b else None
+
+        def _assertion(a):
+            return ManifestAssertion(
+                id=a["id"],
+                type=a["type"],
+                scope=a["scope"],
+                appliesTo_state=a.get("appliesTo_state", a.get("applies_to_state")),
+                statement=a["statement"],
+                binding=_binding(a.get("binding")),
+            )
+
+        return Manifest(
+            schemaVersion=d.get("schemaVersion", d.get("tdf_version")),
+            encryptionInformation=_enc_info(
+                d.get("encryptionInformation", d.get("encryption_information"))
+            )
+            if d.get("encryptionInformation") or d.get("encryption_information")
+            else None,
+            payload=_payload(d["payload"]) if d.get("payload") else None,
+            assertions=[_assertion(a) for a in d.get("assertions", [])],
+        )