diff --git a/lib/streaming_providers/base/models/drm/pssh_parser.py b/lib/streaming_providers/base/models/drm/pssh_parser.py
index 080858f..a0c2619 100644
--- a/lib/streaming_providers/base/models/drm/pssh_parser.py
+++ b/lib/streaming_providers/base/models/drm/pssh_parser.py
@@ -24,8 +24,11 @@ PSSH Box Structure:
└────────────────────────────────────┘
"""
+import base64
+import binascii
+import re
import struct
-from typing import Optional
+from typing import Optional, List, Tuple
from .constants import (
PSSHOffsets,
@@ -44,7 +47,7 @@ class PSSHParser:
Handles both version 0 and version 1+ PSSH boxes, extracting:
- System ID (DRM system UUID)
- - Key IDs (for version 1+)
+ - Key IDs (from box header for v1+, or from payload for v0 Widevine)
- PSSH version
"""
@@ -98,6 +101,13 @@ class PSSHParser:
f"Not a PSSH box: type is '{box_type.decode('ascii', errors='ignore')}'"
)
+ # Validate box size matches actual data
+ if box_size > len(pssh_bytes):
+ raise InvalidPSSHError(
+ f"PSSH box size mismatch: header claims {box_size} bytes, "
+ f"but only {len(pssh_bytes)} bytes available"
+ )
+
# Parse version and system ID
version = pssh_bytes[PSSHOffsets.VERSION]
system_id_bytes = pssh_bytes[PSSHOffsets.SYSTEM_ID_START:PSSHOffsets.SYSTEM_ID_END]
@@ -106,10 +116,24 @@ class PSSHParser:
# Try to identify DRM system
drm_system = DRMSystem.from_uuid(system_id)
- # Extract Key IDs based on version
+ # Extract Key IDs using multiple strategies:
key_ids = []
+
+ # Strategy 1: From PSSH header (version 1+)
if version > 0:
- key_ids = PSSHParser._extract_kids_from_v1_pssh(pssh_bytes, drm_system)
+ try:
+ key_ids = PSSHParser._extract_kids_from_v1_header(pssh_bytes, drm_system)
+ except InvalidPSSHError:
+ # If v1 header extraction fails, try payload parsing as fallback
+ key_ids = []
+
+ # Strategy 2: From Widevine payload (works for v0 and v1)
+ if not key_ids and drm_system == DRMSystem.WIDEVINE:
+ key_ids = PSSHParser._extract_kids_from_widevine_payload(pssh_bytes)
+
+ # Strategy 3: From PlayReady payload (if needed)
+ if not key_ids and drm_system == DRMSystem.PLAYREADY:
+ key_ids = PSSHParser._extract_kids_from_playready_payload(pssh_bytes)
return {
"system_id": system_id,
@@ -119,12 +143,12 @@ class PSSHParser:
}
@staticmethod
- def _extract_kids_from_v1_pssh(
+ def _extract_kids_from_v1_header(
pssh_bytes: bytes,
drm_system: Optional[DRMSystem]
- ) -> list[str]:
+ ) -> List[str]:
"""
- Extract Key IDs from version 1+ PSSH box.
+ Extract Key IDs from version 1+ PSSH box header.
Args:
pssh_bytes: Raw PSSH box bytes
@@ -169,25 +193,244 @@ class PSSHParser:
end_offset = start_offset + KID_SIZE_BYTES
kid_bytes = pssh_bytes[start_offset:end_offset]
- # Convert to hex based on DRM system
- # Widevine uses raw bytes, others may use UUID byte order
- if drm_system == DRMSystem.WIDEVINE:
- # Widevine: KIDs are raw bytes (big-endian)
- kid_hex = bytes_to_uuid_hex(kid_bytes)
- elif drm_system == DRMSystem.PLAYREADY:
- # PlayReady: KIDs are in UUID byte order (mixed-endian)
- kid_hex = PSSHParser._parse_uuid_bytes(kid_bytes)
- else:
- # Unknown system: try UUID format first, fallback to raw
- try:
- kid_hex = PSSHParser._parse_uuid_bytes(kid_bytes)
- except Exception:
- kid_hex = bytes_to_uuid_hex(kid_bytes)
+ # Handle byte order based on DRM system
+ kid_hex = PSSHParser._format_kid_bytes(kid_bytes, drm_system)
- key_ids.append(kid_hex)
+ # Filter out null/empty KIDs
+ if kid_hex and kid_hex != '0' * 32:
+ key_ids.append(kid_hex)
return key_ids
+ @staticmethod
+ def _extract_kids_from_widevine_payload(pssh_bytes: bytes) -> List[str]:
+ """
+ Extract Key IDs from Widevine PSSH payload (protobuf format).
+
+ This parses the internal Widevine structure to find KIDs (field type 0x12)
+ even in version 0 PSSH boxes.
+
+ Args:
+ pssh_bytes: Raw PSSH box bytes
+
+ Returns:
+ List of normalized Key IDs found in protobuf payload
+ """
+ key_ids = []
+
+ try:
+ # Find where the data payload starts
+ version = pssh_bytes[PSSHOffsets.VERSION]
+ data_start = PSSHParser._calculate_payload_start(pssh_bytes, version)
+
+ if data_start is None or data_start >= len(pssh_bytes):
+ return []
+
+ # Parse Widevine protobuf
+ # Field types in Widevine PSSH:
+ # 0x12 = KID (repeated bytes)
+ # 0x1a = Provider
+ # 0x22 = Content ID
+ # 0x2a = Track Type
+ # 0x32 = Policy
+ # 0x38 = Crypto Period Index
+ # 0x48 = Protection Scheme
+ # 0x50 = Crypto Period Seconds
+
+ data = pssh_bytes[data_start:]
+ pos = 0
+
+ while pos < len(data):
+ if pos + 1 > len(data):
+ break
+
+ field_type = data[pos]
+ pos += 1
+
+ # Parse length (varint in protobuf)
+ if pos >= len(data):
+ break
+
+ length, bytes_read = PSSHParser._read_protobuf_varint(data[pos:])
+ pos += bytes_read
+
+ if pos + length > len(data):
+ break
+
+ # Field type 0x12 = KID (bytes)
+ if field_type == 0x12:
+ kid_bytes = data[pos:pos+length]
+ # KIDs are 16 bytes
+ if len(kid_bytes) == KID_SIZE_BYTES:
+ kid_hex = kid_bytes.hex().lower()
+ # Filter out null KIDs
+ if kid_hex != '0' * 32:
+ key_ids.append(kid_hex)
+
+ pos += length
+
+ except (struct.error, IndexError, ValueError):
+ # Silently handle protobuf parsing errors
+ return []
+
+ return key_ids
+
+ @staticmethod
+ def _extract_kids_from_playready_payload(pssh_bytes: bytes) -> List[str]:
+ """
+ Extract Key IDs from PlayReady PSSH payload.
+
+ PlayReady PSSH contains a WRMHEADER XML with KIDs.
+
+ Args:
+ pssh_bytes: Raw PSSH box bytes
+
+ Returns:
+ List of normalized Key IDs found in XML payload
+ """
+ try:
+ # Find where data starts
+ version = pssh_bytes[PSSHOffsets.VERSION]
+ data_start = PSSHParser._calculate_payload_start(pssh_bytes, version)
+
+ if data_start is None or data_start >= len(pssh_bytes):
+ return []
+
+ # PlayReady data is UTF-16LE XML
+ playready_data = pssh_bytes[data_start:].decode('utf-16le', errors='ignore')
+
+ # Look for KID in WRMHEADER
+ # KID appears as: base64-encoded KID or attribute
+ kid_pattern = r']*>([^<]+)'
+ kid_matches = re.findall(kid_pattern, playready_data, re.IGNORECASE)
+
+ key_ids = []
+ for kid_b64 in kid_matches:
+ try:
+ # KID in PlayReady is often base64-encoded UUID
+ kid_bytes = base64.b64decode(kid_b64)
+ if len(kid_bytes) == KID_SIZE_BYTES:
+ # PlayReady KIDs are in mixed-endian format
+ kid_hex = PSSHParser._parse_uuid_bytes(kid_bytes)
+ # Filter out null KIDs
+ if kid_hex != '0' * 32:
+ key_ids.append(kid_hex)
+ except (ValueError, binascii.Error):
+ # Skip invalid base64
+ continue
+
+ return key_ids
+
+ except (UnicodeDecodeError, struct.error, IndexError):
+ return []
+
+ @staticmethod
+ def _calculate_payload_start(pssh_bytes: bytes, version: int) -> Optional[int]:
+ """
+ Calculate where the payload data starts in a PSSH box.
+
+ Args:
+ pssh_bytes: Raw PSSH box bytes
+ version: PSSH version (0 or 1+)
+
+ Returns:
+ Byte offset where payload starts, or None if invalid
+ """
+ try:
+ if version > 0:
+ # v1+: Data starts after KIDs
+ if len(pssh_bytes) < PSSHOffsets.V1_KID_COUNT_END:
+ return None
+
+ # Get KID count
+ kid_count = struct.unpack(
+ ">I",
+ pssh_bytes[PSSHOffsets.V1_KID_COUNT:PSSHOffsets.V1_KID_COUNT_END]
+ )[0]
+ data_start = PSSHOffsets.V1_KIDS_START + (kid_count * KID_SIZE_BYTES)
+
+ # Skip data size field (4 bytes)
+ if len(pssh_bytes) < data_start + 4:
+ return None
+ data_start += 4
+ else:
+ # v0: Data starts right after system ID
+ data_start = PSSHOffsets.SYSTEM_ID_END
+
+ # Skip data size field (4 bytes)
+ if len(pssh_bytes) < data_start + 4:
+ return None
+ data_size = struct.unpack(">I", pssh_bytes[data_start:data_start+4])[0]
+ data_start += 4
+
+ # Validate data size
+ if data_start + data_size > len(pssh_bytes):
+ return None
+
+ return data_start
+
+ except (struct.error, IndexError):
+ return None
+
+ @staticmethod
+ def _read_protobuf_varint(data: bytes) -> Tuple[int, int]:
+ """
+ Read a protobuf varint from bytes.
+
+ Args:
+ data: Byte array to read from
+
+ Returns:
+ Tuple of (value, bytes_read)
+ """
+ value = 0
+ shift = 0
+ bytes_read = 0
+
+ for byte in data:
+ bytes_read += 1
+ value |= (byte & 0x7F) << shift
+ if not (byte & 0x80):
+ break
+ shift += 7
+
+ # Prevent infinite loops on malformed data
+ if bytes_read > 10: # Varint max is 10 bytes
+ break
+
+ return value, bytes_read
+
+ @staticmethod
+ def _format_kid_bytes(kid_bytes: bytes, drm_system: Optional[DRMSystem]) -> str:
+ """
+ Format KID bytes based on DRM system.
+
+ Args:
+ kid_bytes: 16 bytes of KID data
+ drm_system: DRM system identifier
+
+ Returns:
+ Normalized 32-character hex string (no hyphens)
+ """
+ if len(kid_bytes) != KID_SIZE_BYTES:
+ return ""
+
+ try:
+ if drm_system == DRMSystem.WIDEVINE:
+ # Widevine: raw bytes
+ return kid_bytes.hex().lower()
+ elif drm_system == DRMSystem.PLAYREADY:
+ # PlayReady: mixed-endian UUID format
+ return PSSHParser._parse_uuid_bytes(kid_bytes)
+ else:
+ # Unknown: try UUID format first
+ try:
+ return PSSHParser._parse_uuid_bytes(kid_bytes)
+ except (ValueError, struct.error):
+ return kid_bytes.hex().lower()
+ except Exception:
+ return ""
+
@staticmethod
def _parse_uuid_bytes(data: bytes) -> str:
"""
@@ -205,9 +448,12 @@ class PSSHParser:
Returns:
32-character hex string (normalized, no hyphens)
+
+ Raises:
+ ValueError: If data is not exactly 16 bytes
"""
- if len(data) != 16:
- raise ValueError(f"UUID bytes must be 16 bytes (got {len(data)})")
+ if len(data) != KID_SIZE_BYTES:
+ raise ValueError(f"UUID bytes must be {KID_SIZE_BYTES} bytes (got {len(data)})")
# Parse with mixed-endian format
time_low = struct.unpack("