Update kas_client.py & tdf.py, expand tests

Author: b-long

src/otdf_python/kas_client.py

@@ -222,18 +222,37 @@ def _create_signed_request_jwt(self, policy_json, client_public_key, key_access)
 
         # The server expects a JWT with a requestBody field containing the UnsignedRewrapRequest
         # Create the request body that matches UnsignedRewrapRequest protobuf structure
-        # For legacy v1 SRT format, the policy must be base64-encoded
+        # Use the v2 format with explicit policy ID and requests array for cross-tool compatibility
+
+        # Use "policy" as policy ID for compatibility with otdfctl
+        import json
+
+        policy_uuid = "policy"  # otdfctl uses "policy" as the policy ID
+
+        # For v2 format, the policy body must be base64-encoded
         policy_base64 = base64.b64encode(policy_json.encode("utf-8")).decode("utf-8")
 
         unsigned_rewrap_request = {
             "clientPublicKey": client_public_key,  # Maps to client_public_key
-            "policy": policy_base64,  # Maps to policy (legacy field) - base64-encoded
-            "keyAccess": key_access_dict,  # Maps to key_access (legacy field)
+            "requests": [
+                {  # Maps to requests array (v2 format)
+                    "keyAccessObjects": [
+                        {
+                            "keyAccessObjectId": "kao-0",  # Standard KAO ID
+                            "keyAccessObject": key_access_dict,
+                        }
+                    ],
+                    "policy": {
+                        "id": policy_uuid,  # Use the UUID from policy as the policy ID
+                        "body": policy_base64,  # Base64-encoded policy JSON
+                    },
+                }
+            ],
+            "keyAccess": key_access_dict,  # Keep legacy field for backward compatibility
+            "policy": policy_base64,  # Keep legacy field for backward compatibility
         }
 
         # Convert to JSON string
-        import json
-
         request_body_json = json.dumps(unsigned_rewrap_request)
 
         # JWT payload with requestBody field containing the JSON string

src/otdf_python/tdf.py

@@ -2,6 +2,7 @@
 import io
 import os
 import hashlib
+import hmac
 import base64
 import zipfile
 from otdf_python.manifest import (
@@ -334,7 +335,7 @@ def create_tdf(
         segment_size = (
             getattr(config, "default_segment_size", None) or self.SEGMENT_SIZE
         )
-        hasher = hashlib.sha256()
+        segment_hashes_raw = []
         total = 0
         # Write encrypted payload in segments
         with writer.payload() as f:
@@ -346,9 +347,14 @@ def create_tdf(
                     break
                 encrypted = aesgcm.encrypt(chunk)
                 f.write(encrypted.as_bytes())
-                seg_hash = base64.b64encode(
-                    hashlib.sha256(encrypted.as_bytes()).digest()
-                ).decode()
+                # Calculate segment hash using GMAC (last 16 bytes of encrypted segment)
+                # This matches the platform SDK when segmentHashAlg is "GMAC"
+                encrypted_bytes = encrypted.as_bytes()
+                gmac_length = 16  # kGMACPayloadLength from platform SDK
+                if len(encrypted_bytes) < gmac_length:
+                    raise ValueError("Encrypted segment too short for GMAC")
+                seg_hash_raw = encrypted_bytes[-gmac_length:]  # Take last 16 bytes
+                seg_hash = base64.b64encode(seg_hash_raw).decode()
                 segments.append(
                     ManifestSegment(
                         hash=seg_hash,
@@ -360,14 +366,19 @@ def create_tdf(
                         ),  # Changed from encrypted_segment_size to encryptedSegmentSize
                     )
                 )
-                hasher.update(encrypted.as_bytes())
+                # Collect raw segment hash bytes for root signature calculation
+                segment_hashes_raw.append(seg_hash_raw)
                 total += len(chunk)
         # Use config fields for policy
         policy_json = self._build_policy_json(config)
         # Encode policy as base64 to match Java SDK
         policy_b64 = base64.b64encode(policy_json.encode()).decode()
 
-        root_sig = base64.b64encode(hasher.digest()).decode()
+        # Calculate root signature: HMAC-SHA256 over concatenated segment hash raw bytes
+        # This matches the platform SDK approach
+        aggregate_hash = b"".join(segment_hashes_raw)
+        root_sig_raw = hmac.new(key, aggregate_hash, hashlib.sha256).digest()
+        root_sig = base64.b64encode(root_sig_raw).decode()
         integrity_info = ManifestIntegrityInformation(
             rootSignature=ManifestRootSignature(
                 alg="HS256", sig=root_sig
@@ -454,7 +465,6 @@ def read_payload(
         from otdf_python.aesgcm import AesGcm
         from otdf_python.asym_crypto import AsymDecryption
         import base64
-        import hashlib
 
         with zipfile.ZipFile(io.BytesIO(tdf_bytes), "r") as z:
             manifest_json = z.read("0.manifest.json").decode()
@@ -485,8 +495,15 @@ def read_payload(
             for seg in segments:
                 enc_len = seg.encryptedSegmentSize  # Changed field name
                 enc_bytes = encrypted_payload[offset : offset + enc_len]
-                # Integrity check (SHA256 HMAC)
-                seg_hash = base64.b64encode(hashlib.sha256(enc_bytes).digest()).decode()
+                # Integrity check using GMAC (last 16 bytes of encrypted segment)
+                # This matches how segments are hashed when segmentHashAlg is "GMAC"
+                gmac_length = 16  # kGMACPayloadLength from platform SDK
+                if len(enc_bytes) < gmac_length:
+                    raise ValueError(
+                        "Encrypted segment too short for GMAC verification"
+                    )
+                seg_hash_raw = enc_bytes[-gmac_length:]  # Take last 16 bytes
+                seg_hash = base64.b64encode(seg_hash_raw).decode()
                 if seg.hash != seg_hash:
                     raise ValueError("Segment signature mismatch")
                 iv = enc_bytes[: AesGcm.GCM_NONCE_LENGTH]