diff --git a/lib/streaming_providers/base/models/__init__.py b/lib/streaming_providers/base/models/__init__.py index 087bc73..8ac36cd 100644 --- a/lib/streaming_providers/base/models/__init__.py +++ b/lib/streaming_providers/base/models/__init__.py @@ -1,5 +1,5 @@ # streaming_providers/base/models/__init__.py -from .drm_models import DRMConfig, DRMSystem, LicenseConfig, LicenseUnwrapperParams +from .drm import DRMConfig, DRMSystem, LicenseConfig, LicenseUnwrapperParams from .streaming_channel import StreamingChannel from .subscription import SubscriptionPackage, UserSubscription diff --git a/lib/streaming_providers/base/models/drm/__init__.py b/lib/streaming_providers/base/models/drm/__init__.py new file mode 100644 index 0000000..a06beda --- /dev/null +++ b/lib/streaming_providers/base/models/drm/__init__.py @@ -0,0 +1,137 @@ +""" +DRM Models Package + +Comprehensive DRM (Digital Rights Management) support for streaming providers. + +This package provides: +- DRM system identification and configuration +- PSSH (Protection System Specific Header) parsing +- tenc (Track Encryption) box parsing +- License server configuration +- Multi-DRM support with priority handling + +Usage: + from streaming_providers.base.models.drm import ( + DRMSystem, + DRMConfig, + PSSHData, + LicenseConfig, + ) + + # Create Widevine configuration + drm_config = DRMConfig.create_widevine( + server_url="https://license.example.com/widevine", + priority=1 + ) + + # Parse PSSH from manifest + pssh_data = PSSHData( + system_id="edef8ba9-79d6-4ace-a3c8-27dcd51d21ed", + pssh_box="base64_encoded_pssh_box" + ) +""" + +# Version +__version__ = "2.0.0" + +# Core enums +from .drm_systems import DRMSystem + +# Data models +from .pssh_data import PSSHData +from .license_config import ( + LicenseConfig, + LicenseUnwrapperParams, + WrapperType, + UnwrapperType, +) +from .drm_config import DRMConfig + +# Parsers (for advanced usage) +from .pssh_parser import PSSHParser +from .tenc_parser import TencParser + +# Exceptions +from .exceptions import ( + DRMError, + InvalidPSSHError, + InvalidTencError, + InvalidUUIDError, + InvalidKeyIDError, + PSSHSizeError, + UnsupportedDRMSystemError, + LicenseConfigError, + Base64DecodingError, +) + +# Utilities (for advanced usage) +from .utils import ( + normalize_uuid, + format_uuid, + normalize_key_id, + bytes_to_uuid_hex, + bytes_to_uuid_formatted, + uuid_to_bytes, + safe_base64_decode, + safe_base64_encode, + deduplicate_key_ids, +) + +# Constants (for advanced usage) +from .constants import ( + PSSHOffsets, + TencOffsets, + MAX_PSSH_SIZE, + MAX_TENC_SIZE, + KID_SIZE_BYTES, +) + +# Public API +__all__ = [ + # Version + "__version__", + + # Core classes + "DRMSystem", + "DRMConfig", + "PSSHData", + "LicenseConfig", + "LicenseUnwrapperParams", + + # Enums + "WrapperType", + "UnwrapperType", + + # Parsers + "PSSHParser", + "TencParser", + + # Exceptions + "DRMError", + "InvalidPSSHError", + "InvalidTencError", + "InvalidUUIDError", + "InvalidKeyIDError", + "PSSHSizeError", + "UnsupportedDRMSystemError", + "LicenseConfigError", + "Base64DecodingError", + + # Utilities + "normalize_uuid", + "format_uuid", + "normalize_key_id", + "bytes_to_uuid_hex", + "bytes_to_uuid_formatted", + "uuid_to_bytes", + "safe_base64_decode", + "safe_base64_encode", + "deduplicate_key_ids", + + # Constants + "PSSHOffsets", + "TencOffsets", + "MAX_PSSH_SIZE", + "MAX_TENC_SIZE", + "KID_SIZE_BYTES", +] \ No newline at end of file diff --git a/lib/streaming_providers/base/models/drm/constants.py b/lib/streaming_providers/base/models/drm/constants.py new file mode 100644 index 0000000..8840d0d --- /dev/null +++ b/lib/streaming_providers/base/models/drm/constants.py @@ -0,0 +1,159 @@ +""" +DRM Constants and Configuration + +Contains all magic numbers, offsets, and constants used in DRM processing. +""" + +from typing import FrozenSet + +# Character sets for validation +HEX_CHARS: FrozenSet[str] = frozenset('0123456789abcdef') +BASE64_CHARS: FrozenSet[str] = frozenset( + 'ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/=' +) + +# Size limits for security +MAX_PSSH_SIZE: int = 64 * 1024 # 64KB +MAX_TENC_SIZE: int = 1024 # 1KB +MAX_KEY_ID_COUNT: int = 100 # Maximum KIDs in a single PSSH + +# KID and UUID sizes +KID_SIZE_BYTES: int = 16 # 128 bits +UUID_HEX_LENGTH: int = 32 # 32 hex characters without hyphens +UUID_FORMATTED_LENGTH: int = 36 # 36 characters with hyphens + + +class PSSHOffsets: + """ + PSSH Box Binary Structure Offsets + + PSSH Box Format (ISO/IEC 23001-7): + [0-3] Box size (uint32) + [4-7] Box type ('pssh') + [8] Version (uint8) + [9-11] Flags (uint24) + [12-27] System ID (16 bytes) + + Version 0: + [28-31] Data size (uint32) + [32+] Data + + Version 1+: + [28-31] KID count (uint32) + [32+] KIDs (16 bytes each) + [...] Data size (uint32) + [...] Data + """ + BOX_SIZE = 0 + BOX_TYPE = 4 + BOX_TYPE_END = 8 + VERSION = 8 + FLAGS = 9 + FLAGS_END = 12 + SYSTEM_ID_START = 12 + SYSTEM_ID_END = 28 + + # Version 1+ specific + V1_KID_COUNT = 28 + V1_KID_COUNT_END = 32 + V1_KIDS_START = 32 + + # Version 0 specific + V0_DATA_SIZE = 28 + V0_DATA_SIZE_END = 32 + V0_DATA_START = 32 + + # Minimum sizes + MIN_PSSH_SIZE = 32 # Minimum valid PSSH box + MIN_V1_PSSH_SIZE = 36 # Minimum V1 PSSH with 0 KIDs + + +class TencOffsets: + """ + Track Encryption Box (tenc) Binary Structure Offsets + + tenc Box Format (ISO/IEC 23001-7): + [0-3] Box size (uint32) + [4-7] Box type ('tenc') + [8] Version (uint8) + [9-11] Flags (uint24) + [12-15] Reserved (uint24) + default_crypt_byte_block (uint8) + + Version 0: + [16] Reserved (uint8) + [17] default_is_protected (uint8) + [18] default_per_sample_IV_size (uint8) + [19-34] default_KID (16 bytes) + + Version 1: + [16] default_constant_IV_size (uint8) + [17+] default_constant_IV (if size > 0) + [...] default_KID (16 bytes, after IV) + """ + BOX_SIZE = 0 + BOX_TYPE = 4 + BOX_TYPE_END = 8 + VERSION = 8 + + # Version 0 specific + V0_IS_PROTECTED = 17 + V0_IV_SIZE = 18 + V0_KID_START = 19 + V0_KID_END = 35 + + MIN_TENC_V0_SIZE = 35 # Minimum valid V0 tenc box + + +# DRM System UUID Mappings (normalized, no hyphens) +DRM_SYSTEM_UUIDS: dict[str, str] = { + "edef8ba979d64acea3c827dcd51d21ed": "WIDEVINE", + "9a04f07998404286ab92e65be0885f95": "PLAYREADY", + "e2719d58a985b3c9781ab030af78d30e": "CLEARKEY", + "3d5e6d359b9a41e8b843dd3c6e72c42c": "WISEPLAY", + "94ce86fb07ff4f43adb893d2fa968ca2": "FAIRPLAY", +} + +# Reverse mapping for quick lookups +DRM_SYSTEM_NAMES: dict[str, str] = { + "WIDEVINE": "edef8ba979d64acea3c827dcd51d21ed", + "PLAYREADY": "9a04f07998404286ab92e65be0885f95", + "CLEARKEY": "e2719d58a985b3c9781ab030af78d30e", + "WISEPLAY": "3d5e6d359b9a41e8b843dd3c6e72c42c", + "FAIRPLAY": "94ce86fb07ff4f43adb893d2fa968ca2", + "GENERIC": "", + "NONE": "", +} + +# DRM System Alias Mappings (all normalized - no dots, hyphens, lowercase) +DRM_ALIAS_MAPPING: dict[str, str] = { + # Widevine + "widevine": "WIDEVINE", + "comwidevinealpha": "WIDEVINE", + "edef8ba979d64acea3c827dcd51d21ed": "WIDEVINE", + + # PlayReady + "playready": "PLAYREADY", + "commicrosoftplayready": "PLAYREADY", + "9a04f07998404286ab92e65be0885f95": "PLAYREADY", + + # ClearKey + "clearkey": "CLEARKEY", + "orgw3clearkey": "CLEARKEY", + "e2719d58a985b3c9781ab030af78d30e": "CLEARKEY", + + # FairPlay + "fairplay": "FAIRPLAY", + "comapplefps": "FAIRPLAY", + "skd": "FAIRPLAY", + "94ce86fb07ff4f43adb893d2fa968ca2": "FAIRPLAY", + + # WisePlay + "wiseplay": "WISEPLAY", + "comhuaweiwiseplay": "WISEPLAY", + "3d5e6d359b9a41e8b843dd3c6e72c42c": "WISEPLAY", + + # Generic/None + "generic": "GENERIC", + "none": "NONE", + "unencrypted": "NONE", +} \ No newline at end of file diff --git a/lib/streaming_providers/base/models/drm/drm_config.py b/lib/streaming_providers/base/models/drm/drm_config.py new file mode 100644 index 0000000..43c31f9 --- /dev/null +++ b/lib/streaming_providers/base/models/drm/drm_config.py @@ -0,0 +1,169 @@ +""" +DRM Configuration Model + +Main configuration class that combines DRM system, PSSH data, and license config. +""" + +from dataclasses import dataclass +from typing import Optional + +from .drm_systems import DRMSystem +from .license_config import LicenseConfig +from .exceptions import LicenseConfigError + + +@dataclass +class DRMConfig: + """ + Complete DRM Configuration. + + Combines DRM system identification with license configuration + for a single DRM system. + + Attributes: + system: DRM system type + priority: Priority for multi-DRM scenarios (higher = preferred) + license: License server configuration (optional for unencrypted) + """ + + system: DRMSystem + priority: int = 0 + license: Optional[LicenseConfig] = None + + def validate(self) -> None: + """ + Validate the DRM configuration. + + Raises: + LicenseConfigError: If configuration is invalid + """ + if not isinstance(self.system, DRMSystem): + raise LicenseConfigError( + f"system must be a DRMSystem enum, got {type(self.system)}" + ) + + if self.license: + self.license.validate() + + def to_dict(self) -> dict: + """ + Convert to dictionary format expected by players. + + Returns: + Dictionary with DRM configuration in player-compatible format: + { + "com.widevine.alpha": { + "priority": 1, + "license": {...} + } + } + """ + result = { + self.system.value: { + "priority": self.priority + } + } + + if self.license: + license_dict = self.license.to_dict() + if license_dict: + result[self.system.value]["license"] = license_dict + + return result + + @classmethod + def create_widevine( + cls, + server_url: str, + priority: int = 1, + **license_kwargs + ) -> "DRMConfig": + """ + Helper to create Widevine DRM configuration. + + Args: + server_url: Widevine license server URL + priority: Priority (default: 1) + **license_kwargs: Additional LicenseConfig parameters + + Returns: + DRMConfig instance for Widevine + """ + license_config = LicenseConfig(server_url=server_url, **license_kwargs) + return cls(system=DRMSystem.WIDEVINE, priority=priority, license=license_config) + + @classmethod + def create_playready( + cls, + server_url: str, + priority: int = 1, + **license_kwargs + ) -> "DRMConfig": + """ + Helper to create PlayReady DRM configuration. + + Args: + server_url: PlayReady license server URL + priority: Priority (default: 1) + **license_kwargs: Additional LicenseConfig parameters + + Returns: + DRMConfig instance for PlayReady + """ + license_config = LicenseConfig(server_url=server_url, **license_kwargs) + return cls(system=DRMSystem.PLAYREADY, priority=priority, license=license_config) + + @classmethod + def create_clearkey( + cls, + keyids: dict[str, str], + priority: int = 0, + **license_kwargs + ) -> "DRMConfig": + """ + Helper to create ClearKey DRM configuration. + + Args: + keyids: Mapping of Key IDs to Keys (hex strings) + priority: Priority (default: 0) + **license_kwargs: Additional LicenseConfig parameters + + Returns: + DRMConfig instance for ClearKey + """ + license_config = LicenseConfig(keyids=keyids, **license_kwargs) + return cls(system=DRMSystem.CLEARKEY, priority=priority, license=license_config) + + @classmethod + def create_fairplay( + cls, + server_url: str, + server_certificate: str, + priority: int = 1, + **license_kwargs + ) -> "DRMConfig": + """ + Helper to create FairPlay DRM configuration. + + Args: + server_url: FairPlay license server URL (skd://) + server_certificate: Base64-encoded FairPlay certificate + priority: Priority (default: 1) + **license_kwargs: Additional LicenseConfig parameters + + Returns: + DRMConfig instance for FairPlay + """ + license_config = LicenseConfig( + server_url=server_url, + server_certificate=server_certificate, + **license_kwargs + ) + return cls(system=DRMSystem.FAIRPLAY, priority=priority, license=license_config) + + def __repr__(self) -> str: + """Return detailed representation.""" + has_license = "with license" if self.license else "no license" + return ( + f"" + ) \ No newline at end of file diff --git a/lib/streaming_providers/base/models/drm/drm_systems.py b/lib/streaming_providers/base/models/drm/drm_systems.py new file mode 100644 index 0000000..0b2b46f --- /dev/null +++ b/lib/streaming_providers/base/models/drm/drm_systems.py @@ -0,0 +1,144 @@ +""" +DRM System Definitions + +Enum for supported DRM systems with efficient UUID and alias lookups. +""" + +from enum import Enum +from functools import cached_property +from typing import Optional + +from .constants import DRM_SYSTEM_NAMES, DRM_SYSTEM_UUIDS, DRM_ALIAS_MAPPING +from .utils import normalize_uuid, normalize_alias +from .exceptions import InvalidUUIDError + + +class DRMSystem(str, Enum): + """ + Supported DRM Systems + + Each DRM system has: + - A unique identifier (Android/EME format) + - A standard UUID (ISO/IEC 23001-7) + - Multiple aliases for flexible lookups + """ + + WIDEVINE = "com.widevine.alpha" + PLAYREADY = "com.microsoft.playready" + WISEPLAY = "com.huawei.wiseplay" + CLEARKEY = "org.w3.clearkey" + FAIRPLAY = "com.apple.fps" + GENERIC = "generic" + NONE = "none" + + @cached_property + def system_uuid(self) -> str: + """ + Get the standard UUID for this DRM system (normalized, no hyphens). + + Returns: + 32-character hex string (lowercase, no hyphens), or empty string for GENERIC/NONE + """ + # Get enum name (e.g., 'WIDEVINE') + enum_name = self.name + return DRM_SYSTEM_NAMES.get(enum_name, "") + + @cached_property + def system_uuid_formatted(self) -> str: + """ + Get the formatted UUID with hyphens for this DRM system. + + Returns: + UUID string with hyphens (8-4-4-4-12), or empty string for GENERIC/NONE + """ + uuid = self.system_uuid + if not uuid: + return "" + + return ( + f"{uuid[0:8]}-{uuid[8:12]}-{uuid[12:16]}-" + f"{uuid[16:20]}-{uuid[20:32]}" + ) + + @classmethod + def from_uuid(cls, uuid: str) -> Optional["DRMSystem"]: + """ + Get DRM system from UUID. + + Accepts any format (with or without hyphens). + + Args: + uuid: UUID string (e.g., "edef8ba9-79d6-4ace-a3c8-27dcd51d21ed") + + Returns: + DRMSystem enum value or None if not recognized + + Examples: + >>> DRMSystem.from_uuid("edef8ba9-79d6-4ace-a3c8-27dcd51d21ed") + + + >>> DRMSystem.from_uuid("edef8ba979d64acea3c827dcd51d21ed") + + """ + if not uuid: + return None + + try: + # Normalize UUID (removes hyphens, validates format) + normalized = normalize_uuid(uuid) + except InvalidUUIDError: + return None + + # Lookup in mapping + enum_name = DRM_SYSTEM_UUIDS.get(normalized) + if not enum_name: + return None + + return getattr(cls, enum_name, None) + + @classmethod + def from_alias(cls, alias: str) -> Optional["DRMSystem"]: + """ + Resolve DRM system from human-friendly alias or UUID. + + Handles: + - Short aliases: "clearkey", "widevine", "playready", "fairplay", "wiseplay" + - Full Android identifiers: "com.widevine.alpha", "org.w3.clearkey", etc. + - UUIDs (with or without hyphens): "edef8ba9-79d6-4ace-a3c8-27dcd51d21ed" + + Args: + alias: DRM system alias, identifier, or UUID + + Returns: + DRMSystem enum value or None if unrecognized + + Examples: + >>> DRMSystem.from_alias("widevine") + + + >>> DRMSystem.from_alias("com.widevine.alpha") + + + >>> DRMSystem.from_alias("clearkey") + + """ + if not alias: + return None + + # Normalize alias (remove dots, hyphens, lowercase) + normalized = normalize_alias(alias) + + # Lookup in alias mapping + enum_name = DRM_ALIAS_MAPPING.get(normalized) + if not enum_name: + return None + + return getattr(cls, enum_name, None) + + def __str__(self) -> str: + """Return the DRM system identifier""" + return self.value + + def __repr__(self) -> str: + """Return detailed representation""" + return f"<{self.__class__.__name__}.{self.name}: '{self.value}'>" \ No newline at end of file diff --git a/lib/streaming_providers/base/models/drm/exceptions.py b/lib/streaming_providers/base/models/drm/exceptions.py new file mode 100644 index 0000000..528c573 --- /dev/null +++ b/lib/streaming_providers/base/models/drm/exceptions.py @@ -0,0 +1,50 @@ +""" +DRM-specific Exception Classes + +Custom exceptions for better error handling and debugging in DRM operations. +""" + + +class DRMError(Exception): + """Base exception for all DRM-related errors""" + pass + + +class InvalidPSSHError(DRMError): + """Raised when PSSH box data is invalid or malformed""" + pass + + +class InvalidTencError(DRMError): + """Raised when tenc box data is invalid or malformed""" + pass + + +class InvalidUUIDError(DRMError): + """Raised when UUID format is invalid""" + pass + + +class InvalidKeyIDError(DRMError): + """Raised when Key ID format is invalid""" + pass + + +class PSSHSizeError(DRMError): + """Raised when PSSH box exceeds size limits""" + pass + + +class UnsupportedDRMSystemError(DRMError): + """Raised when DRM system is not supported""" + pass + + +class LicenseConfigError(DRMError): + """Raised when license configuration is invalid""" + pass + + +class Base64DecodingError(DRMError): + """Raised when base64 decoding fails""" + pass \ No newline at end of file diff --git a/lib/streaming_providers/base/models/drm/license_config.py b/lib/streaming_providers/base/models/drm/license_config.py new file mode 100644 index 0000000..6c1a3c5 --- /dev/null +++ b/lib/streaming_providers/base/models/drm/license_config.py @@ -0,0 +1,207 @@ +""" +License Configuration Models + +Data classes for DRM license configuration, including server URLs, +certificates, headers, and unwrapper parameters. +""" + +from dataclasses import dataclass, field +from enum import Enum +from typing import Optional + +from .utils import safe_base64_decode, safe_base64_encode, normalize_key_id +from .exceptions import LicenseConfigError + + +class WrapperType(str, Enum): + """License request wrapper types""" + BASE64 = "base64" + URLENC = "urlenc" + NONE = "none" + + +class UnwrapperType(str, Enum): + """License response unwrapper types""" + AUTO = "auto" + BASE64 = "base64" + JSON = "json" + XML = "xml" + NONE = "none" + + +@dataclass +class LicenseUnwrapperParams: + """ + Parameters for license response unwrapping. + + Used to extract license data and HDCP information from + non-standard license server responses. + + Attributes: + path_data: JSON/XML path to license data (e.g., "license.data") + path_data_traverse: Whether to traverse nested structures for data + path_hdcp_res: JSON/XML path to HDCP resolution restriction + path_hdcp_res_traverse: Whether to traverse for HDCP resolution + path_hdcp_ver: JSON/XML path to HDCP version requirement + path_hdcp_ver_traverse: Whether to traverse for HDCP version + """ + + path_data: Optional[str] = None + path_data_traverse: bool = False + path_hdcp_res: Optional[str] = None + path_hdcp_res_traverse: bool = False + path_hdcp_ver: Optional[str] = None + path_hdcp_ver_traverse: bool = False + + def to_dict(self) -> dict: + """Convert to dictionary, excluding None values.""" + return { + k: v for k, v in { + "path_data": self.path_data, + "path_data_traverse": self.path_data_traverse, + "path_hdcp_res": self.path_hdcp_res, + "path_hdcp_res_traverse": self.path_hdcp_res_traverse, + "path_hdcp_ver": self.path_hdcp_ver, + "path_hdcp_ver_traverse": self.path_hdcp_ver_traverse, + }.items() + if v is not None and (not isinstance(v, bool) or v) + } + + +@dataclass +class LicenseConfig: + """ + DRM License Server Configuration. + + Contains all information needed to acquire and process DRM licenses, + including server URLs, certificates, request customization, and + response processing. + + Attributes: + server_url: License server URL + server_certificate: Base64-encoded server certificate (for FairPlay, etc.) + use_http_get_request: Use GET instead of POST for license requests + req_headers: Custom HTTP headers as JSON string + req_params: URL query parameters as JSON string + req_data: Base64-encoded custom request body data + wrapper: Request body wrapper type + unwrapper: Response unwrapper type + unwrapper_params: Parameters for response unwrapping + keyids: ClearKey key mappings (KID -> Key in hex) + """ + + server_url: Optional[str] = None + server_certificate: Optional[str] = None + use_http_get_request: bool = False + req_headers: Optional[str] = None + req_params: Optional[str] = None + req_data: Optional[str] = None + wrapper: Optional[str] = None + unwrapper: Optional[str] = None + unwrapper_params: Optional[LicenseUnwrapperParams] = None + keyids: dict[str, str] = field(default_factory=dict) + + def __post_init__(self): + """Normalize key IDs in keyids mapping.""" + if self.keyids: + # Normalize all KIDs and keys to lowercase hex + normalized_keyids = {} + for kid, key in self.keyids.items(): + try: + norm_kid = normalize_key_id(kid) + norm_key = key.lower().replace("-", "") + normalized_keyids[norm_kid] = norm_key + except Exception: + # Skip invalid entries + pass + self.keyids = normalized_keyids + + def validate(self) -> None: + """ + Validate the license configuration. + + Raises: + LicenseConfigError: If configuration is invalid + """ + # Validate server_certificate if present + if self.server_certificate: + try: + safe_base64_decode(self.server_certificate) + except Exception as e: + raise LicenseConfigError( + f"server_certificate must be valid base64: {e}" + ) from e + + # Validate req_data if present + if self.req_data: + try: + safe_base64_decode(self.req_data) + except Exception as e: + raise LicenseConfigError( + f"req_data must be valid base64: {e}" + ) from e + + # Validate keyids format (for ClearKey) + if self.keyids: + for kid, key in self.keyids.items(): + if len(kid) != 32: + raise LicenseConfigError( + f"Invalid KID length: {kid} (must be 32 hex chars)" + ) + if len(key) != 32: + raise LicenseConfigError( + f"Invalid KEY length for KID {kid}: {key} (must be 32 hex chars)" + ) + if not all(c in "0123456789abcdef" for c in kid): + raise LicenseConfigError(f"KID contains non-hex characters: {kid}") + if not all(c in "0123456789abcdef" for c in key): + raise LicenseConfigError( + f"KEY contains non-hex characters for KID {kid}: {key}" + ) + + @classmethod + def create_with_req_data(cls, req_data_template: str, **kwargs) -> "LicenseConfig": + """ + Helper to create LicenseConfig with base64-encoded req_data. + + Args: + req_data_template: Plain text request data template + **kwargs: Other LicenseConfig parameters + + Returns: + LicenseConfig instance with encoded req_data + """ + req_data_encoded = safe_base64_encode(req_data_template.encode("utf-8")) + return cls(req_data=req_data_encoded, **kwargs) + + def to_dict(self) -> dict: + """ + Convert to dictionary, excluding None/empty values. + + Returns: + Dictionary representation + """ + result = {} + + if self.server_url: + result["server_url"] = self.server_url + if self.server_certificate: + result["server_certificate"] = self.server_certificate + if self.use_http_get_request: + result["use_http_get_request"] = self.use_http_get_request + if self.req_headers: + result["req_headers"] = self.req_headers + if self.req_params: + result["req_params"] = self.req_params + if self.req_data: + result["req_data"] = self.req_data + if self.wrapper: + result["wrapper"] = self.wrapper + if self.unwrapper: + result["unwrapper"] = self.unwrapper + if self.unwrapper_params: + result["unwrapper_params"] = self.unwrapper_params.to_dict() + if self.keyids: + result["keyids"] = self.keyids + + return result \ No newline at end of file diff --git a/lib/streaming_providers/base/models/drm/pssh_data.py b/lib/streaming_providers/base/models/drm/pssh_data.py new file mode 100644 index 0000000..fc57fbb --- /dev/null +++ b/lib/streaming_providers/base/models/drm/pssh_data.py @@ -0,0 +1,196 @@ +""" +PSSH Data Model + +Data class for storing PSSH (Protection System Specific Header) information. +""" + +from dataclasses import dataclass, field +from functools import cached_property +from typing import Optional + +from .drm_systems import DRMSystem +from .pssh_parser import PSSHParser +from .utils import normalize_uuid, normalize_key_id, format_uuid, deduplicate_key_ids +from .exceptions import InvalidPSSHError, InvalidKeyIDError + + +@dataclass +class PSSHData: + """ + Container for PSSH (Protection System Specific Header) data. + + Attributes: + system_id: DRM system UUID (normalized, 32 hex chars, no hyphens) + pssh_box: Base64-encoded PSSH box (optional if extracting from segments) + key_ids: List of Key IDs (normalized, 32 hex chars, no hyphens) + source: Source of PSSH data (e.g., "manifest", "segment", "init") + """ + + system_id: str + pssh_box: str = "" + key_ids: list[str] = field(default_factory=list) + source: str = "manifest" + + def __post_init__(self): + """Normalize and validate all data at creation time.""" + # Normalize system_id + if self.system_id: + try: + self.system_id = normalize_uuid(self.system_id) + except Exception as e: + raise InvalidPSSHError(f"Invalid system_id: {e}") from e + else: + raise InvalidPSSHError("system_id is required") + + # Normalize key_ids + if self.key_ids: + normalized_kids = [] + for kid in self.key_ids: + try: + normalized_kids.append(normalize_key_id(kid)) + except InvalidKeyIDError: + # Skip invalid KIDs but don't fail + pass + + # Remove duplicates while preserving order + self.key_ids = deduplicate_key_ids(normalized_kids) + + # Auto-extract key_ids from pssh_box if provided but no key_ids + if self.pssh_box and not self.key_ids: + try: + metadata = PSSHParser.parse_pssh_box(self.pssh_box) + self.key_ids = metadata["key_ids"] + except Exception: + # If extraction fails, leave key_ids empty + # They might be available in tenc boxes or segments + pass + + @cached_property + def drm_system(self) -> Optional[DRMSystem]: + """ + Get DRM system enum from system_id. + + Returns: + DRMSystem enum value or None if not recognized + """ + return DRMSystem.from_uuid(self.system_id) + + @cached_property + def system_id_formatted(self) -> str: + """ + Get system_id with hyphens for display. + + Returns: + UUID string with hyphens (8-4-4-4-12) + """ + return format_uuid(self.system_id) + + @cached_property + def pssh_version(self) -> int: + """ + Get PSSH version from the box. + + Returns: + PSSH version number, or -1 if no PSSH box + """ + if not self.pssh_box: + return -1 + + try: + return PSSHParser.get_pssh_version(self.pssh_box) + except Exception: + return -1 + + def needs_tenc_fallback(self) -> bool: + """ + Check if this PSSH needs tenc box fallback for Key IDs. + + Returns: + True if tenc fallback is needed (version 0 PSSH without KIDs) + """ + return PSSHParser.needs_tenc_fallback(self.pssh_box, self.drm_system) + + @property + def needs_extraction(self) -> bool: + """ + Check if PSSH/key_ids need to be extracted from segments. + + Returns: + True if PSSH box or Key IDs are missing + """ + return not self.pssh_box or not self.key_ids + + def add_key_ids(self, new_kids: list[str]) -> None: + """ + Add Key IDs to the existing list (normalized and deduplicated). + + Args: + new_kids: List of Key IDs to add + """ + normalized_kids = [] + for kid in new_kids: + try: + normalized_kids.append(normalize_key_id(kid)) + except InvalidKeyIDError: + # Skip invalid KIDs + pass + + # Merge with existing and deduplicate + all_kids = self.key_ids + normalized_kids + self.key_ids = deduplicate_key_ids(all_kids) + + def validate(self) -> None: + """ + Validate the PSSH data. + + Raises: + InvalidPSSHError: If data is invalid + """ + if not self.system_id: + raise InvalidPSSHError("system_id is required") + + # Validate system_id format + try: + normalize_uuid(self.system_id) + except Exception as e: + raise InvalidPSSHError(f"Invalid system_id format: {e}") from e + + # Validate pssh_box if present + if self.pssh_box: + try: + PSSHParser.parse_pssh_box(self.pssh_box) + except Exception as e: + raise InvalidPSSHError(f"Invalid pssh_box: {e}") from e + + # Validate key_ids format + for kid in self.key_ids: + try: + normalize_key_id(kid) + except Exception as e: + raise InvalidPSSHError(f"Invalid key_id format: {e}") from e + + def to_dict(self) -> dict: + """ + Convert to dictionary representation. + + Returns: + Dictionary with all PSSH data + """ + return { + "system_id": self.system_id, + "system_id_formatted": self.system_id_formatted, + "pssh_box": self.pssh_box, + "key_ids": self.key_ids, + "source": self.source, + "drm_system": self.drm_system.value if self.drm_system else None, + "version": self.pssh_version, + } + + def __repr__(self) -> str: + """Return detailed representation.""" + drm_name = self.drm_system.name if self.drm_system else "UNKNOWN" + kid_count = len(self.key_ids) + return ( + f"" + ) \ No newline at end of file diff --git a/lib/streaming_providers/base/models/drm/pssh_parser.py b/lib/streaming_providers/base/models/drm/pssh_parser.py new file mode 100644 index 0000000..080858f --- /dev/null +++ b/lib/streaming_providers/base/models/drm/pssh_parser.py @@ -0,0 +1,284 @@ +""" +PSSH Box Parser + +Handles parsing of Protection System Specific Header (PSSH) boxes +according to ISO/IEC 23001-7 specification. + +PSSH Box Structure: +┌────────────────────────────────────┐ +│ Box Size (4 bytes) │ +├────────────────────────────────────┤ +│ Box Type 'pssh' (4 bytes) │ +├────────────────────────────────────┤ +│ Version (1 byte) | Flags (3 bytes) │ +├────────────────────────────────────┤ +│ System ID (16 bytes UUID) │ +├────────────────────────────────────┤ +│ [Version 1+] │ +│ KID Count (4 bytes) │ +│ KIDs (16 bytes each) │ +├────────────────────────────────────┤ +│ Data Size (4 bytes) │ +├────────────────────────────────────┤ +│ Data (variable length) │ +└────────────────────────────────────┘ +""" + +import struct +from typing import Optional + +from .constants import ( + PSSHOffsets, + KID_SIZE_BYTES, + MAX_PSSH_SIZE, + MAX_KEY_ID_COUNT, +) +from .drm_systems import DRMSystem +from .exceptions import InvalidPSSHError, PSSHSizeError +from .utils import bytes_to_uuid_hex, safe_base64_decode + + +class PSSHParser: + """ + Parser for PSSH (Protection System Specific Header) boxes. + + Handles both version 0 and version 1+ PSSH boxes, extracting: + - System ID (DRM system UUID) + - Key IDs (for version 1+) + - PSSH version + """ + + @staticmethod + def parse_pssh_box(pssh_base64: str) -> dict: + """ + Parse a PSSH box and extract metadata. + + Args: + pssh_base64: Base64-encoded PSSH box + + Returns: + Dictionary containing: + - system_id: Normalized UUID (32 hex chars, no hyphens) + - version: PSSH version (0 or 1+) + - key_ids: List of Key IDs (normalized, no hyphens) + - drm_system: DRMSystem enum value (if recognized) + + Raises: + InvalidPSSHError: If PSSH box is malformed + PSSHSizeError: If PSSH box exceeds size limits + """ + if not pssh_base64: + raise InvalidPSSHError("PSSH box cannot be empty") + + # Decode base64 + try: + pssh_bytes = safe_base64_decode(pssh_base64) + except Exception as e: + raise InvalidPSSHError(f"Failed to decode PSSH base64: {e}") from e + + # Check size limits + if len(pssh_bytes) > MAX_PSSH_SIZE: + raise PSSHSizeError( + f"PSSH box too large: {len(pssh_bytes)} bytes (max: {MAX_PSSH_SIZE})" + ) + + # Validate minimum size + if len(pssh_bytes) < PSSHOffsets.MIN_PSSH_SIZE: + raise InvalidPSSHError( + f"PSSH box too small: {len(pssh_bytes)} bytes " + f"(minimum: {PSSHOffsets.MIN_PSSH_SIZE})" + ) + + # Parse box header + box_size = struct.unpack(">I", pssh_bytes[PSSHOffsets.BOX_SIZE:4])[0] + box_type = pssh_bytes[PSSHOffsets.BOX_TYPE:PSSHOffsets.BOX_TYPE_END] + + if box_type != b"pssh": + raise InvalidPSSHError( + f"Not a PSSH box: type is '{box_type.decode('ascii', errors='ignore')}'" + ) + + # Parse version and system ID + version = pssh_bytes[PSSHOffsets.VERSION] + system_id_bytes = pssh_bytes[PSSHOffsets.SYSTEM_ID_START:PSSHOffsets.SYSTEM_ID_END] + system_id = bytes_to_uuid_hex(system_id_bytes) + + # Try to identify DRM system + drm_system = DRMSystem.from_uuid(system_id) + + # Extract Key IDs based on version + key_ids = [] + if version > 0: + key_ids = PSSHParser._extract_kids_from_v1_pssh(pssh_bytes, drm_system) + + return { + "system_id": system_id, + "version": version, + "key_ids": key_ids, + "drm_system": drm_system, + } + + @staticmethod + def _extract_kids_from_v1_pssh( + pssh_bytes: bytes, + drm_system: Optional[DRMSystem] + ) -> list[str]: + """ + Extract Key IDs from version 1+ PSSH box. + + Args: + pssh_bytes: Raw PSSH box bytes + drm_system: Identified DRM system (if any) + + Returns: + List of normalized Key IDs (32 hex chars, no hyphens) + + Raises: + InvalidPSSHError: If PSSH structure is invalid + """ + if len(pssh_bytes) < PSSHOffsets.MIN_V1_PSSH_SIZE: + raise InvalidPSSHError( + f"Version 1+ PSSH too small: {len(pssh_bytes)} bytes " + f"(minimum: {PSSHOffsets.MIN_V1_PSSH_SIZE})" + ) + + # Read KID count + kid_count = struct.unpack( + ">I", + pssh_bytes[PSSHOffsets.V1_KID_COUNT:PSSHOffsets.V1_KID_COUNT_END] + )[0] + + # Validate KID count + if kid_count > MAX_KEY_ID_COUNT: + raise InvalidPSSHError( + f"Too many Key IDs in PSSH: {kid_count} (max: {MAX_KEY_ID_COUNT})" + ) + + # Calculate required size + required_size = PSSHOffsets.V1_KIDS_START + (kid_count * KID_SIZE_BYTES) + if len(pssh_bytes) < required_size: + raise InvalidPSSHError( + f"PSSH truncated: claims {kid_count} KIDs but only " + f"{len(pssh_bytes)} bytes available" + ) + + # Extract KIDs + key_ids = [] + for i in range(kid_count): + start_offset = PSSHOffsets.V1_KIDS_START + (i * KID_SIZE_BYTES) + 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) + + key_ids.append(kid_hex) + + return key_ids + + @staticmethod + def _parse_uuid_bytes(data: bytes) -> str: + """ + Parse UUID bytes with mixed-endian format (Microsoft GUID format). + + UUID byte order (RFC 4122): + - time_low (4 bytes): big-endian → little-endian + - time_mid (2 bytes): big-endian → little-endian + - time_hi_version (2 bytes): big-endian → little-endian + - clock_seq (2 bytes): big-endian (no change) + - node (6 bytes): big-endian (no change) + + Args: + data: 16 bytes of UUID data + + Returns: + 32-character hex string (normalized, no hyphens) + """ + if len(data) != 16: + raise ValueError(f"UUID bytes must be 16 bytes (got {len(data)})") + + # Parse with mixed-endian format + time_low = struct.unpack(" int: + """ + Get PSSH version without full parsing. + + Args: + pssh_base64: Base64-encoded PSSH box + + Returns: + PSSH version number (0, 1, etc.) + + Raises: + InvalidPSSHError: If PSSH box is malformed + """ + if not pssh_base64: + raise InvalidPSSHError("PSSH box cannot be empty") + + try: + pssh_bytes = safe_base64_decode(pssh_base64) + except Exception as e: + raise InvalidPSSHError(f"Failed to decode PSSH base64: {e}") from e + + if len(pssh_bytes) < PSSHOffsets.VERSION + 1: + raise InvalidPSSHError(f"PSSH box too small to read version") + + return pssh_bytes[PSSHOffsets.VERSION] + + @staticmethod + def needs_tenc_fallback(pssh_base64: str, drm_system: Optional[DRMSystem]) -> bool: + """ + Check if PSSH needs tenc box fallback for Key IDs. + + Version 0 PSSH boxes don't contain Key IDs, so we need to extract + them from tenc boxes in the MP4 structure. + + Args: + pssh_base64: Base64-encoded PSSH box + drm_system: DRM system (if known) + + Returns: + True if tenc fallback is needed + """ + if not pssh_base64: + return True + + try: + metadata = PSSHParser.parse_pssh_box(pssh_base64) + + # Need fallback if: + # 1. No Key IDs extracted + # 2. Version 0 PSSH (doesn't contain KIDs) + # 3. Widevine system (commonly uses tenc) + return ( + not metadata["key_ids"] and + metadata["version"] == 0 and + metadata["drm_system"] == DRMSystem.WIDEVINE + ) + except Exception: + return True \ No newline at end of file diff --git a/lib/streaming_providers/base/models/drm/tenc_parser.py b/lib/streaming_providers/base/models/drm/tenc_parser.py new file mode 100644 index 0000000..157ef99 --- /dev/null +++ b/lib/streaming_providers/base/models/drm/tenc_parser.py @@ -0,0 +1,217 @@ +""" +tenc Box Parser + +Handles parsing of Track Encryption (tenc) boxes for extracting Key IDs +when PSSH boxes don't contain them (version 0 PSSH). +""" + +import struct + +from .constants import TencOffsets, MAX_TENC_SIZE, KID_SIZE_BYTES +from .exceptions import InvalidTencError +from .utils import bytes_to_uuid_hex + + +class TencParser: + """ + Parser for tenc (Track Encryption) boxes. + + Used as fallback when PSSH version 0 doesn't contain Key IDs. + According to ISO/IEC 23001-7, tenc boxes contain default encryption + information including the default Key ID. + """ + + @staticmethod + def extract_kids_from_tenc(tenc_data: bytes) -> list[str]: + """ + Extract Key IDs from tenc box data. + + This is typically used as a fallback when PSSH version 0 boxes + don't contain Key IDs directly. + + Args: + tenc_data: Raw tenc box bytes + + Returns: + List of normalized Key IDs (32 hex chars, no hyphens) + + Raises: + InvalidTencError: If tenc box is malformed + """ + if not tenc_data: + raise InvalidTencError("tenc box data cannot be empty") + + # Check size limits + if len(tenc_data) > MAX_TENC_SIZE: + raise InvalidTencError( + f"tenc box too large: {len(tenc_data)} bytes (max: {MAX_TENC_SIZE})" + ) + + # Validate minimum size for version 0 + if len(tenc_data) < TencOffsets.MIN_TENC_V0_SIZE: + raise InvalidTencError( + f"tenc box too small: {len(tenc_data)} bytes " + f"(minimum: {TencOffsets.MIN_TENC_V0_SIZE})" + ) + + # Validate box type + box_type = tenc_data[TencOffsets.BOX_TYPE:TencOffsets.BOX_TYPE_END] + if box_type != b"tenc": + raise InvalidTencError( + f"Not a tenc box: type is '{box_type.decode('ascii', errors='ignore')}'" + ) + + # Parse version + version = tenc_data[TencOffsets.VERSION] + + # Extract KID based on version + if version == 0: + return TencParser._extract_kid_from_v0_tenc(tenc_data) + elif version == 1: + return TencParser._extract_kid_from_v1_tenc(tenc_data) + else: + raise InvalidTencError(f"Unsupported tenc version: {version}") + + @staticmethod + def _extract_kid_from_v0_tenc(tenc_data: bytes) -> list[str]: + """ + Extract Key ID from version 0 tenc box. + + Version 0 structure: + [16] Reserved (uint8) + [17] default_is_protected (uint8) + [18] default_per_sample_IV_size (uint8) + [19-34] default_KID (16 bytes) + + Args: + tenc_data: Raw tenc box bytes + + Returns: + List containing single normalized Key ID + + Raises: + InvalidTencError: If tenc structure is invalid + """ + if len(tenc_data) < TencOffsets.V0_KID_END: + raise InvalidTencError( + f"tenc v0 box truncated: {len(tenc_data)} bytes " + f"(need at least {TencOffsets.V0_KID_END})" + ) + + # Check if track is protected + is_protected = tenc_data[TencOffsets.V0_IS_PROTECTED] + if is_protected == 0: + # Unencrypted track, no KID + return [] + + # Extract KID bytes + kid_bytes = tenc_data[TencOffsets.V0_KID_START:TencOffsets.V0_KID_END] + + # tenc KIDs are typically in UUID/GUID format (mixed-endian) + try: + kid_hex = TencParser._parse_guid_bytes(kid_bytes) + except Exception: + # Fallback to raw hex if GUID parsing fails + kid_hex = bytes_to_uuid_hex(kid_bytes) + + return [kid_hex] + + @staticmethod + def _extract_kid_from_v1_tenc(tenc_data: bytes) -> list[str]: + """ + Extract Key ID from version 1 tenc box. + + Version 1 structure: + [16] default_constant_IV_size (uint8) + [17+] default_constant_IV (if IV_size > 0) + [...] default_KID (16 bytes, after IV) + + Args: + tenc_data: Raw tenc box bytes + + Returns: + List containing single normalized Key ID + + Raises: + InvalidTencError: If tenc structure is invalid + """ + # Read IV size + if len(tenc_data) < 17: + raise InvalidTencError("tenc v1 box too small to read IV size") + + iv_size = tenc_data[16] + + # KID starts after IV + kid_start = 17 + iv_size + kid_end = kid_start + KID_SIZE_BYTES + + if len(tenc_data) < kid_end: + raise InvalidTencError( + f"tenc v1 box truncated: {len(tenc_data)} bytes " + f"(need at least {kid_end} for KID)" + ) + + # Extract KID bytes + kid_bytes = tenc_data[kid_start:kid_end] + + # tenc KIDs are typically in UUID/GUID format + try: + kid_hex = TencParser._parse_guid_bytes(kid_bytes) + except Exception: + kid_hex = bytes_to_uuid_hex(kid_bytes) + + return [kid_hex] + + @staticmethod + def _parse_guid_bytes(data: bytes) -> str: + """ + Parse GUID/UUID bytes with mixed-endian format (Microsoft GUID). + + GUID byte order: + - Data1 (4 bytes): little-endian + - Data2 (2 bytes): little-endian + - Data3 (2 bytes): little-endian + - Data4 (8 bytes): big-endian + + Args: + data: 16 bytes of GUID data + + Returns: + 32-character hex string (normalized, no hyphens) + """ + if len(data) != 16: + raise ValueError(f"GUID bytes must be 16 bytes (got {len(data)})") + + # Parse with mixed-endian format (Microsoft GUID) + data1 = struct.unpack(" bool: + """ + Check if track is encrypted based on tenc box. + + Args: + tenc_data: Raw tenc box bytes + + Returns: + True if track is encrypted, False otherwise + """ + if not tenc_data or len(tenc_data) < TencOffsets.V0_IS_PROTECTED + 1: + return False + + version = tenc_data[TencOffsets.VERSION] + + if version == 0: + is_protected = tenc_data[TencOffsets.V0_IS_PROTECTED] + return is_protected != 0 + + # Version 1+ tracks are assumed encrypted if tenc box exists + return True \ No newline at end of file diff --git a/lib/streaming_providers/base/models/drm/utils.py b/lib/streaming_providers/base/models/drm/utils.py new file mode 100644 index 0000000..f12f7ac --- /dev/null +++ b/lib/streaming_providers/base/models/drm/utils.py @@ -0,0 +1,244 @@ +""" +DRM Utility Functions + +Common utility functions for UUID formatting, validation, and normalization. +""" + +import base64 +import binascii +from typing import Optional + +from .constants import HEX_CHARS, UUID_HEX_LENGTH, UUID_FORMATTED_LENGTH +from .exceptions import InvalidUUIDError, InvalidKeyIDError, Base64DecodingError + + +def is_hex_string(s: str, length: Optional[int] = None) -> bool: + """ + Check if string contains only hexadecimal characters. + + Args: + s: String to validate + length: Optional expected length + + Returns: + True if string is valid hex (and matches length if provided) + """ + if not s: + return False + + if length is not None and len(s) != length: + return False + + return all(c in HEX_CHARS for c in s.lower()) + + +def normalize_uuid(uuid: str) -> str: + """ + Normalize UUID to 32 lowercase hex characters without hyphens. + + Args: + uuid: UUID string (with or without hyphens) + + Returns: + Normalized UUID string (32 hex chars, no hyphens, lowercase) + + Raises: + InvalidUUIDError: If UUID format is invalid + """ + if not uuid: + raise InvalidUUIDError("UUID cannot be empty") + + # Remove hyphens and convert to lowercase + normalized = uuid.lower().replace("-", "") + + # Validate format + if len(normalized) != UUID_HEX_LENGTH: + raise InvalidUUIDError( + f"UUID must be {UUID_HEX_LENGTH} hex characters (got {len(normalized)}): {uuid}" + ) + + if not is_hex_string(normalized): + raise InvalidUUIDError(f"UUID contains non-hex characters: {uuid}") + + return normalized + + +def format_uuid(uuid: str) -> str: + """ + Format UUID with standard hyphens (8-4-4-4-12). + + Args: + uuid: UUID string (32 hex chars, no hyphens) + + Returns: + Formatted UUID string with hyphens + + Raises: + InvalidUUIDError: If UUID format is invalid + """ + normalized = normalize_uuid(uuid) + + return ( + f"{normalized[0:8]}-{normalized[8:12]}-{normalized[12:16]}-" + f"{normalized[16:20]}-{normalized[20:32]}" + ) + + +def normalize_key_id(kid: str) -> str: + """ + Normalize Key ID to 32 lowercase hex characters without hyphens. + + Args: + kid: Key ID string (with or without hyphens) + + Returns: + Normalized Key ID string (32 hex chars, no hyphens, lowercase) + + Raises: + InvalidKeyIDError: If Key ID format is invalid + """ + if not kid: + raise InvalidKeyIDError("Key ID cannot be empty") + + # Remove hyphens and convert to lowercase + normalized = kid.lower().replace("-", "") + + # Validate format + if len(normalized) != UUID_HEX_LENGTH: + raise InvalidKeyIDError( + f"Key ID must be {UUID_HEX_LENGTH} hex characters (got {len(normalized)}): {kid}" + ) + + if not is_hex_string(normalized): + raise InvalidKeyIDError(f"Key ID contains non-hex characters: {kid}") + + return normalized + + +def bytes_to_uuid_hex(data: bytes) -> str: + """ + Convert 16 bytes to UUID hex string (no hyphens). + + This is more efficient than using the uuid module. + + Args: + data: 16 bytes of UUID data + + Returns: + 32-character hex string (lowercase, no hyphens) + + Raises: + InvalidUUIDError: If data is not 16 bytes + """ + if len(data) != 16: + raise InvalidUUIDError(f"UUID bytes must be 16 bytes long (got {len(data)})") + + return data.hex().lower() + + +def bytes_to_uuid_formatted(data: bytes) -> str: + """ + Convert 16 bytes to formatted UUID string with hyphens. + + Args: + data: 16 bytes of UUID data + + Returns: + 36-character UUID string with hyphens (8-4-4-4-12) + + Raises: + InvalidUUIDError: If data is not 16 bytes + """ + hex_string = bytes_to_uuid_hex(data) + return format_uuid(hex_string) + + +def uuid_to_bytes(uuid: str) -> bytes: + """ + Convert UUID string to 16 bytes. + + Args: + uuid: UUID string (with or without hyphens) + + Returns: + 16 bytes of UUID data + + Raises: + InvalidUUIDError: If UUID format is invalid + """ + normalized = normalize_uuid(uuid) + return bytes.fromhex(normalized) + + +def safe_base64_decode(data: str) -> bytes: + """ + Safely decode base64 string with better error handling. + + Args: + data: Base64-encoded string + + Returns: + Decoded bytes + + Raises: + Base64DecodingError: If decoding fails + """ + if not data: + raise Base64DecodingError("Cannot decode empty base64 string") + + try: + return base64.b64decode(data) + except (binascii.Error, ValueError) as e: + raise Base64DecodingError(f"Failed to decode base64 data: {e}") from e + + +def safe_base64_encode(data: bytes) -> str: + """ + Safely encode bytes to base64 string. + + Args: + data: Bytes to encode + + Returns: + Base64-encoded string + """ + if not data: + return "" + + return base64.b64encode(data).decode('utf-8') + + +def normalize_alias(alias: str) -> str: + """ + Normalize alias for consistent lookups. + + Removes hyphens, dots, and converts to lowercase. + + Args: + alias: Alias string + + Returns: + Normalized alias + """ + return alias.lower().strip().replace("-", "").replace(".", "") + + +def deduplicate_key_ids(key_ids: list[str]) -> list[str]: + """ + Remove duplicate Key IDs while preserving order. + + Args: + key_ids: List of Key IDs (normalized) + + Returns: + List of unique Key IDs in original order + """ + seen = set() + result = [] + + for kid in key_ids: + if kid not in seen: + seen.add(kid) + result.append(kid) + + return result \ No newline at end of file diff --git a/lib/streaming_providers/base/models/drm_models.py b/lib/streaming_providers/base/models/drm_models.py deleted file mode 100644 index b099c32..0000000 --- a/lib/streaming_providers/base/models/drm_models.py +++ /dev/null @@ -1,242 +0,0 @@ -# streaming_providers/base/models/drm_models.py -import base64 -from dataclasses import dataclass, field -from enum import Enum -from typing import Dict, List, Optional - - -class DRMSystem(str, Enum): - WIDEVINE = "com.widevine.alpha" - PLAYREADY = "com.microsoft.playready" - WISEPLAY = "com.huawei.wiseplay" - CLEARKEY = "org.w3.clearkey" - FAIRPLAY = "com.apple.fps" - GENERIC = "generic" - NONE = "none" - - @property - def system_uuid(self) -> str: - """Get the standard UUID for this DRM system""" - uuid_mapping = { - self.WIDEVINE: "edef8ba9-79d6-4ace-a3c8-27dcd51d21ed", - self.PLAYREADY: "9a04f079-9840-4286-ab92-e65be0885f95", - self.CLEARKEY: "e2719d58-a985-b3c9-781a-b030af78d30e", - self.WISEPLAY: "3d5e6d35-9b9a-41e8-b843-dd3c6e72c42c", - self.FAIRPLAY: "94ce86fb-07ff-4f43-adb8-93d2fa968ca2", - self.GENERIC: "", # No UUID for generic plugins - self.NONE: "", # No UUID for unencrypted - } - return uuid_mapping.get(self, "") - - @classmethod - def from_uuid(cls, uuid: str) -> Optional["DRMSystem"]: - """Get DRM system from UUID""" - uuid_lower = uuid.lower().replace("-", "") - uuid_mapping = { - "edef8ba979d64acea3c827dcd51d21ed": cls.WIDEVINE, - "9a04f07998404286ab92e65be0885f95": cls.PLAYREADY, - "e2719d58a985b3c9781ab030af78d30e": cls.CLEARKEY, - "3d5e6d359b9a41e8b843dd3c6e72c42c": cls.WISEPLAY, - } - return uuid_mapping.get(uuid_lower) - - @classmethod - def from_alias(cls, alias: str) -> Optional["DRMSystem"]: - """ - Resolve DRM system from human-friendly alias or UUID - - Handles: - - Short aliases: "clearkey", "widevine", "playready", "fairplay", "wiseplay" - - Full Android identifiers: "com.widevine.alpha", "org.w3.clearkey", etc. - - UUIDs (with or without hyphens): "edef8ba9-79d6-4ace-a3c8-27dcd51d21ed" - - Returns: - DRMSystem enum value or None if unrecognized - """ - alias_lower = alias.lower().strip().replace("-", "") - - alias_mapping = { - # Widevine - "widevine": cls.WIDEVINE, - "com.widevine.alpha": cls.WIDEVINE, - "edef8ba979d64acea3c827dcd51d21ed": cls.WIDEVINE, - # PlayReady - "playready": cls.PLAYREADY, - "com.microsoft.playready": cls.PLAYREADY, - "9a04f07998404286ab92e65be0885f95": cls.PLAYREADY, - # ClearKey - "clearkey": cls.CLEARKEY, - "org.w3.clearkey": cls.CLEARKEY, - "e2719d58a985b3c9781ab030af78d30e": cls.CLEARKEY, - # FairPlay - "fairplay": cls.FAIRPLAY, - "com.apple.fps": cls.FAIRPLAY, - "skd": cls.FAIRPLAY, - "94ce86fb07ff4f43adb893d2fa968ca2": cls.FAIRPLAY, - # WisePlay - "wiseplay": cls.WISEPLAY, - "com.huawei.wiseplay": cls.WISEPLAY, - "3d5e6d359b9a41e8b843dd3c6e72c42c": cls.WISEPLAY, - # Generic/None - "generic": cls.GENERIC, - "none": cls.NONE, - "unencrypted": cls.NONE, - } - - return alias_mapping.get(alias_lower) - - -class WrapperType(str, Enum): - BASE64 = "base64" - URLENC = "urlenc" - NONE = "none" - - -@dataclass -class PSSHData: - system_id: str - pssh_box: str = "" # Base64 encoded PSSH box - key_ids: List[str] = field(default_factory=list) - source: str = "manifest" # "manifest", "mp4_segment", "unknown" - - @property - def needs_extraction(self) -> bool: - """Check if PSSH/key_ids need to be extracted from segments""" - return not self.pssh_box or not self.key_ids - - @property - def drm_system(self) -> Optional[DRMSystem]: - """Get the corresponding DRM system for this PSSH""" - return DRMSystem.from_uuid(self.system_id) - - def validate(self): - """Validate the PSSH data""" - if not self.system_id: - raise ValueError("system_id is required") - - # FIX: pssh_box can be empty (PSSH in segments) - if self.pssh_box: # Only validate if not empty - try: - base64.b64decode(self.pssh_box) - except Exception: - raise ValueError("pssh_box must be valid base64") - - # Validate key IDs if present - for kid in self.key_ids: - if not all(c in "0123456789abcdefABCDEF-" for c in kid): - raise ValueError(f"Invalid key ID format: {kid}") - - -class UnwrapperType(str, Enum): - AUTO = "auto" - BASE64 = "base64" - JSON = "json" - XML = "xml" - NONE = "none" - - -@dataclass -class LicenseUnwrapperParams: - path_data: Optional[str] = None - path_data_traverse: bool = False - path_hdcp_res: Optional[str] = None - path_hdcp_res_traverse: bool = False - path_hdcp_ver: Optional[str] = None - path_hdcp_ver_traverse: bool = False - - -@dataclass -class LicenseConfig: - server_url: Optional[str] = None - server_certificate: Optional[str] = None - use_http_get_request: bool = False - req_headers: Optional[str] = None - req_params: Optional[str] = None - req_data: Optional[str] = None - wrapper: Optional[str] = None - unwrapper: Optional[str] = None - unwrapper_params: Optional[LicenseUnwrapperParams] = None - keyids: Dict[str, str] = field(default_factory=dict) # For ClearKey - - def validate(self): - """Validate the license configuration""" - if self.server_certificate: - try: - base64.b64decode(self.server_certificate) - except Exception: - raise ValueError("server_certificate must be valid base64") - - if self.req_data: - try: - base64.b64decode(self.req_data) - except Exception: - raise ValueError("req_data must be valid base64") - - if self.keyids: - for kid, key in self.keyids.items(): - if not all(c in "0123456789abcdefABCDEF" for c in kid): - raise ValueError(f"Invalid KID format: {kid}") - if not all(c in "0123456789abcdefABCDEF" for c in key): - raise ValueError(f"Invalid KEY format: {key}") - - @classmethod - def create_with_base64_req_data(cls, req_data_template: str, **kwargs): - """Helper to ensure req_data is base64 encoded""" - import base64 - - req_data_encoded = base64.b64encode(req_data_template.encode("utf-8")).decode( - "utf-8" - ) - return cls(req_data=req_data_encoded, **kwargs) - - -@dataclass -class DRMConfig: - system: DRMSystem - priority: int = 0 - license: Optional[LicenseConfig] = None - - def validate(self): - """Validate the DRM configuration""" - if self.license: - self.license.validate() - - def to_dict(self) -> Dict: - """Convert to dictionary format expected by players""" - result = { - str(self.system.value): { # Use .value to get the actual string - "priority": self.priority - } - } - - if self.license: - license_dict = {} - if self.license.server_url: - license_dict["server_url"] = self.license.server_url - if self.license.server_certificate: - license_dict["server_certificate"] = self.license.server_certificate - if self.license.use_http_get_request: - license_dict["use_http_get_request"] = self.license.use_http_get_request - if self.license.req_headers: - license_dict["req_headers"] = self.license.req_headers - if self.license.req_params: - license_dict["req_params"] = self.license.req_params - if self.license.req_data: - license_dict["req_data"] = self.license.req_data - if self.license.wrapper: - license_dict["wrapper"] = self.license.wrapper - if self.license.unwrapper: - license_dict["unwrapper"] = self.license.unwrapper - if self.license.unwrapper_params: - license_dict["unwrapper_params"] = { - k: v - for k, v in vars(self.license.unwrapper_params).items() - if v is not None - } - if self.license.keyids: - license_dict["keyids"] = self.license.keyids - - if license_dict: - result[str(self.system.value)]["license"] = license_dict - - return result diff --git a/lib/streaming_providers/base/utils/__init__.py b/lib/streaming_providers/base/utils/__init__.py index d893438..bf75f18 100644 --- a/lib/streaming_providers/base/utils/__init__.py +++ b/lib/streaming_providers/base/utils/__init__.py @@ -5,6 +5,7 @@ from .manifest_parser import ManifestParser from .mpd_cache import MPDCacheManager from .mpd_rewriter import MPDRewriter from .timestamp_converter import TimestampConverter +from .mp4_pssh_extractor import MP4PSSHExtractor from .vfs import VFS __all__ = [ @@ -14,5 +15,6 @@ __all__ = [ "VFS", "MPDRewriter", "MPDCacheManager", + "MP4PSSHExtractor", "TimestampConverter", ] diff --git a/lib/streaming_providers/base/utils/manifest_parser.py b/lib/streaming_providers/base/utils/manifest_parser.py index 81f34cb..7665a08 100644 --- a/lib/streaming_providers/base/utils/manifest_parser.py +++ b/lib/streaming_providers/base/utils/manifest_parser.py @@ -3,7 +3,7 @@ import re from typing import List, Optional from urllib.parse import quote, urljoin, urlparse -from ..models.drm_models import DRMSystem, PSSHData +from ..models.drm import PSSHData from .logger import logger @@ -34,7 +34,7 @@ class ManifestParser: @staticmethod def _extract_from_manifest_content(manifest_content: str) -> List[PSSHData]: """Extract PSSH and DRM systems from manifest content""" - # Try regex extraction first (handles PSSH boxes with KIDs) + # Try regex extraction first pssh_list = ManifestParser._extract_with_regex(manifest_content) if pssh_list: return pssh_list @@ -43,7 +43,6 @@ class ManifestParser: drm_systems_found = set() result = [] - # More efficient: compile regex once cp_pattern = re.compile( r']*schemeIdUri="urn:uuid:([^"]+)"[^>]*>', re.IGNORECASE, @@ -52,49 +51,48 @@ class ManifestParser: for match in cp_pattern.finditer(manifest_content): system_id = match.group(1).lower() - # Skip mp4protection scheme - if ( - "mp4protection" - in manifest_content[max(0, match.start() - 100) : match.start()] - ): + # Skip mp4protection + if "mp4protection" in manifest_content[max(0, match.start() - 100):match.start()]: continue - drm_system = DRMSystem.from_uuid(system_id) - if drm_system and system_id not in drm_systems_found: - drm_systems_found.add(system_id) - result.append( - PSSHData( - system_id=system_id, - pssh_box="", # Empty - PSSH in segments - key_ids=[], - source="manifest_scheme_only", - ) - ) - logger.debug(f"Found DRM system from schemeIdUri: {drm_system.value}") + # Let PSSHData handle normalization! + pssh_data = PSSHData( + system_id=system_id, + pssh_box="", + key_ids=[], + source="manifest_scheme_only", + ) + + if pssh_data.drm_system and system_id not in drm_systems_found: + drm_systems_found.add(pssh_data.system_id) # Use normalized ID + result.append(pssh_data) + logger.debug(f"Found DRM system: {pssh_data.drm_system.value}") return result @staticmethod def _extract_from_single_segment( - segment_url: str, expected_system_ids: List[str] = None + segment_url: str, expected_system_ids: List[str] = None ) -> List[PSSHData]: - """Extract PSSH from a single MP4 segment""" - from .mp4_parser import MP4PSSHExtractor + from .mp4_pssh_extractor import MP4PSSHExtractor try: pssh_from_segment = MP4PSSHExtractor.extract_from_url(segment_url) - # Filter for expected DRM systems if provided if expected_system_ids: + # Normalize expected IDs using the model! + normalized_expected = [] + for sys_id in expected_system_ids: + # Create temporary PSSHData to leverage its normalization + temp = PSSHData(system_id=sys_id, source="filter") + normalized_expected.append(temp.system_id) + filtered_pssh = [ - p for p in pssh_from_segment if p.system_id in expected_system_ids + p for p in pssh_from_segment + if p.system_id in normalized_expected ] if filtered_pssh: - logger.debug(f"Found {len(filtered_pssh)} PSSH boxes in segment") return filtered_pssh - elif pssh_from_segment: - logger.debug(f"Found {len(pssh_from_segment)} PSSH boxes in segment") - return pssh_from_segment except Exception as e: logger.warning(f"Failed to extract PSSH from segment: {e}") diff --git a/lib/streaming_providers/base/utils/mp4_parser.py b/lib/streaming_providers/base/utils/mp4_parser.py deleted file mode 100644 index 5fc4b44..0000000 --- a/lib/streaming_providers/base/utils/mp4_parser.py +++ /dev/null @@ -1,483 +0,0 @@ -import base64 -import struct -import uuid -from typing import List, Optional - -from ..models.drm_models import PSSHData -from .logger import logger - - -class MP4PSSHExtractor: - """Extract PSSH boxes and key IDs from MP4 segments""" - - @staticmethod - def extract_from_url(segment_url: str, timeout: int = 10) -> List[PSSHData]: - """Download MP4 segment and extract PSSH data""" - import requests - - try: - response = requests.get(segment_url, timeout=timeout) - response.raise_for_status() - - # Only download first ~100KB for efficiency - chunk_size = 1024 * 100 - data = response.content[:chunk_size] - - return MP4PSSHExtractor.extract_from_bytes(data) - - except Exception as e: - logger.error(f"Failed to extract PSSH from {segment_url}: {e}") - return [] - - @staticmethod - def extract_from_bytes(data: bytes) -> List[PSSHData]: - """Extract PSSH boxes and encryption info from MP4 binary data""" - pssh_data_list = [] - offset = 0 - - # First, extract all tenc boxes to get default KIDs - tenc_kids = MP4PSSHExtractor._extract_all_tenc_kids(data) - - while offset < len(data): - try: - # Read box size (4 bytes, big-endian) - if offset + 4 > len(data): - break - - box_size = struct.unpack(">I", data[offset : offset + 4])[0] - if box_size == 0: - box_size = len(data) - offset # Box extends to end of file - elif box_size == 1: - # Extended size (skip for now - rare in practice) - break - - if offset + box_size > len(data): - break - - # Read box type (4 bytes) - box_type = data[offset + 4 : offset + 8].decode( - "ascii", errors="ignore" - ) - - if box_type == "moov": - # Look for PSSH in moov box - moov_data = data[offset : offset + box_size] - pssh_in_moov = MP4PSSHExtractor._extract_from_moov(moov_data) - - # Enhance PSSH data with tenc KIDs if needed - for pssh in pssh_in_moov: - if not pssh.key_ids and tenc_kids: - pssh.key_ids = tenc_kids.copy() - pssh_data_list.extend(pssh_in_moov) - - elif box_type == "pssh": - # Found standalone PSSH box - pssh_box = MP4PSSHExtractor._parse_pssh_box( - data[offset : offset + box_size] - ) - if pssh_box: - # Add tenc KIDs if PSSH doesn't have its own - if not pssh_box.key_ids and tenc_kids: - pssh_box.key_ids = tenc_kids.copy() - pssh_data_list.append(pssh_box) - - offset += box_size - - except Exception as e: - logger.debug(f"Error parsing MP4 box at offset {offset}: {e}") - offset += 1 # Try to recover - - return pssh_data_list - - @staticmethod - def _extract_all_tenc_kids(data: bytes) -> List[str]: - """Recursively extract all KIDs from tenc boxes in the MP4 data""" - kids = [] - offset = 0 - - while offset < len(data): - try: - if offset + 8 > len(data): - break - - box_size = struct.unpack(">I", data[offset : offset + 4])[0] - if box_size < 8 or offset + box_size > len(data): - offset += 1 - continue - - box_type = data[offset + 4 : offset + 8] - - if box_type == b"tenc": - # Parse tenc box - logger.debug( - f"Found tenc box at offset {offset} (0x{offset:x}), size {box_size}" - ) - tenc_kid = MP4PSSHExtractor._parse_tenc_box( - data[offset : offset + box_size] - ) - if tenc_kid and tenc_kid not in kids: - kids.append(tenc_kid) - elif box_size > 8: - # Recursively search inside container boxes - container_boxes = { - b"moov", - b"trak", - b"mdia", - b"minf", - b"stbl", - b"stsd", - b"encv", - b"sinf", - b"schi", - b"udta", - } - if box_type in container_boxes: - # Recursively search inside the container - inner_data = data[offset + 8 : offset + box_size] - inner_kids = MP4PSSHExtractor._extract_all_tenc_kids(inner_data) - for kid in inner_kids: - if kid not in kids: - kids.append(kid) - - offset += box_size - except Exception as e: - logger.debug(f"Error parsing box at offset {offset}: {e}") - offset += 1 - - return kids - - @staticmethod - def _parse_tenc_box(tenc_bytes: bytes) -> Optional[str]: - """Parse a tenc box and extract the default key ID""" - try: - # tenc box should be at least 32 bytes for version 0 with KID - if len(tenc_bytes) < 32: - logger.debug( - f"tenc box too small: {len(tenc_bytes)} bytes, need at least 32" - ) - return None - - # Parse box header - box_size = struct.unpack(">I", tenc_bytes[0:4])[0] - box_type = tenc_bytes[4:8] - - if box_type != b"tenc": - logger.debug(f"Not a tenc box, type: {box_type.hex()}") - return None - - version = tenc_bytes[8] - logger.debug(f"tenc box version: {version}") - - if version == 0: - # For version 0: reserved(24) + is_encrypted(1) is a 32-bit field at bytes 12-15 - # Actually, it's stored as a 32-bit big-endian integer - # Bytes 12-15: reserved (24 bits) + is_encrypted (1 bit) - # Byte 16: default_iv_size - # Bytes 17-32: default_KID (16 bytes) - - # Read the 32-bit field at bytes 12-15 - reserved_encrypted = struct.unpack(">I", tenc_bytes[12:16])[0] - - # Extract is_encrypted from bit 24 (0-indexed from MSB) - # Actually, in big-endian, byte 14 is the 3rd byte of this 32-bit field - # Let me think: bytes[12], bytes[13], bytes[14], bytes[15] - # Byte 14 = tenc_bytes[14] = your "01" - # Byte 15 = tenc_bytes[15] = your "08" - - # Actually simpler: byte 14 contains is_encrypted in its LSB - is_encrypted = (tenc_bytes[14] & 0x01) != 0 - default_iv_size = tenc_bytes[15] - - # KID starts at byte 16 (not 14!) - kid_offset = 16 - - logger.debug( - f"tenc v0: is_encrypted={is_encrypted}, default_iv_size={default_iv_size}, kid_offset={kid_offset}" - ) - - elif version == 1: - # Version 1 has different structure - kid_offset = 16 # Adjust as needed for v1 - logger.debug(f"tenc v1: kid_offset={kid_offset}") - else: - logger.debug(f"Unknown tenc version: {version}") - return None - - if kid_offset + 16 > len(tenc_bytes): - logger.debug( - f"tenc box too small for KID: {len(tenc_bytes)} bytes, need {kid_offset + 16}" - ) - return None - - # Extract KID bytes - kid_bytes = tenc_bytes[kid_offset : kid_offset + 16] - hex_kid = kid_bytes.hex() - - logger.debug(f"KID bytes at offset {kid_offset}: {hex_kid}") - - # Format as UUID - try: - kid_uuid = str(uuid.UUID(bytes=kid_bytes)) - # Convert to lowercase without dashes (same format as PSSH KIDs) - clean_kid = kid_uuid.replace("-", "").lower() - logger.debug(f"Parsed KID: {clean_kid}") - return clean_kid - except Exception as e: - logger.debug(f"Failed to parse KID bytes as UUID: {e}, hex: {hex_kid}") - # Return hex string as fallback - return hex_kid - - except Exception as e: - logger.debug(f"Failed to parse tenc box: {e}") - return None - - # [Keep all the rest of the methods EXACTLY the same...] - @staticmethod - def _extract_from_moov(moov_data: bytes) -> List[PSSHData]: - """Extract PSSH boxes from moov container""" - pssh_list = [] - offset = 8 # Skip moov header - - while offset < len(moov_data): - try: - box_size = struct.unpack(">I", moov_data[offset : offset + 4])[0] - box_type = moov_data[offset + 4 : offset + 8].decode( - "ascii", errors="ignore" - ) - - if box_type == "trak": - # Parse track for PSSH - trak_data = moov_data[offset : offset + box_size] - pssh_in_trak = MP4PSSHExtractor._extract_from_trak(trak_data) - pssh_list.extend(pssh_in_trak) - - elif box_type == "pssh": - # PSSH directly in moov - pssh_box = MP4PSSHExtractor._parse_pssh_box( - moov_data[offset : offset + box_size] - ) - if pssh_box: - pssh_list.append(pssh_box) - - offset += box_size - - except: - break - - return pssh_list - - @staticmethod - def _extract_from_trak(trak_data: bytes) -> List[PSSHData]: - """Extract PSSH from trak box""" - pssh_list = [] - offset = 8 - - while offset < len(trak_data): - try: - box_size = struct.unpack(">I", trak_data[offset : offset + 4])[0] - box_type = trak_data[offset + 4 : offset + 8].decode( - "ascii", errors="ignore" - ) - - if box_type == "mdia": - mdia_data = trak_data[offset : offset + box_size] - pssh_in_mdia = MP4PSSHExtractor._extract_from_mdia(mdia_data) - pssh_list.extend(pssh_in_mdia) - - offset += box_size - - except: - break - - return pssh_list - - @staticmethod - def _extract_from_mdia(mdia_data: bytes) -> List[PSSHData]: - """Extract PSSH from mdia box""" - pssh_list = [] - offset = 8 - - while offset < len(mdia_data): - try: - box_size = struct.unpack(">I", mdia_data[offset : offset + 4])[0] - box_type = mdia_data[offset + 4 : offset + 8].decode( - "ascii", errors="ignore" - ) - - if box_type == "minf": - minf_data = mdia_data[offset : offset + box_size] - pssh_in_minf = MP4PSSHExtractor._extract_from_minf(minf_data) - pssh_list.extend(pssh_in_minf) - - offset += box_size - - except: - break - - return pssh_list - - @staticmethod - def _extract_from_minf(minf_data: bytes) -> List[PSSHData]: - """Extract PSSH from minf box""" - pssh_list = [] - offset = 8 - - while offset < len(minf_data): - try: - box_size = struct.unpack(">I", minf_data[offset : offset + 4])[0] - box_type = minf_data[offset + 4 : offset + 8].decode( - "ascii", errors="ignore" - ) - - if box_type == "stbl": - stbl_data = minf_data[offset : offset + box_size] - pssh_in_stbl = MP4PSSHExtractor._extract_from_stbl(stbl_data) - pssh_list.extend(pssh_in_stbl) - - offset += box_size - - except: - break - - return pssh_list - - @staticmethod - def _extract_from_stbl(stbl_data: bytes) -> List[PSSHData]: - """Extract PSSH from stbl box (where protection scheme info usually is)""" - pssh_list = [] - offset = 8 - - while offset < len(stbl_data): - try: - box_size = struct.unpack(">I", stbl_data[offset : offset + 4])[0] - box_type = stbl_data[offset + 4 : offset + 8].decode( - "ascii", errors="ignore" - ) - - if box_type == "sinf": - sinf_data = stbl_data[offset : offset + box_size] - pssh_in_sinf = MP4PSSHExtractor._extract_from_sinf(sinf_data) - pssh_list.extend(pssh_in_sinf) - - offset += box_size - - except: - break - - return pssh_list - - @staticmethod - def _extract_from_sinf(sinf_data: bytes) -> List[PSSHData]: - """Extract PSSH from sinf (protection scheme information) box""" - pssh_list = [] - offset = 8 - - while offset < len(sinf_data): - try: - box_size = struct.unpack(">I", sinf_data[offset : offset + 4])[0] - box_type = sinf_data[offset + 4 : offset + 8].decode( - "ascii", errors="ignore" - ) - - if box_type == "schi": - schi_data = sinf_data[offset : offset + box_size] - pssh_in_schi = MP4PSSHExtractor._extract_from_schi(schi_data) - pssh_list.extend(pssh_in_schi) - - offset += box_size - - except: - break - - return pssh_list - - @staticmethod - def _extract_from_schi(schi_data: bytes) -> List[PSSHData]: - """Extract PSSH from schi box (where PSSH boxes are typically stored)""" - pssh_list = [] - offset = 8 - - while offset < len(schi_data): - try: - box_size = struct.unpack(">I", schi_data[offset : offset + 4])[0] - box_type = schi_data[offset + 4 : offset + 8].decode( - "ascii", errors="ignore" - ) - - if box_type == "pssh": - pssh_box = MP4PSSHExtractor._parse_pssh_box( - schi_data[offset : offset + box_size] - ) - if pssh_box: - pssh_list.append(pssh_box) - - offset += box_size - - except: - break - - return pssh_list - - @staticmethod - def _parse_pssh_box(pssh_bytes: bytes) -> Optional[PSSHData]: - """Parse a PSSH box and extract system_id, pssh_box, and key_ids""" - try: - if len(pssh_bytes) < 32: # Minimum size for PSSH box - return None - - # Parse box header - box_size = struct.unpack(">I", pssh_bytes[0:4])[0] - box_type = pssh_bytes[4:8].decode("ascii") - - if box_type != "pssh": - return None - - version = pssh_bytes[8] - flags = struct.unpack(">I", b"\x00" + pssh_bytes[9:12])[0] - - # Extract system ID (bytes 12-28) - system_id_bytes = pssh_bytes[12:28] - system_id = str(uuid.UUID(bytes=system_id_bytes)) - - # Extract key IDs (if version > 0) - key_ids = [] - current_offset = 28 - - if version > 0: - # Read KID count - if current_offset + 4 > len(pssh_bytes): - return None - - kid_count = struct.unpack( - ">I", pssh_bytes[current_offset : current_offset + 4] - )[0] - current_offset += 4 - - # Read each KID - for _ in range(kid_count): - if current_offset + 16 > len(pssh_bytes): - break - - kid_bytes = pssh_bytes[current_offset : current_offset + 16] - # kid_uuid = str(uuid.UUID(bytes=kid_bytes)) - # key_ids.append(kid_uuid.replace("-", "").lower()) - key_ids.append(kid_bytes.hex().lower()) - current_offset += 16 - else: - logger.debug( - f"Version 0 PSSH for system {system_id}, will look for KIDs in tenc box" - ) - - # Encode entire PSSH box as base64 - pssh_b64 = base64.b64encode(pssh_bytes[:box_size]).decode("ascii") - - return PSSHData( - system_id=system_id, - pssh_box=pssh_b64, - key_ids=key_ids, - source="mp4_segment", - ) - - except Exception as e: - logger.debug(f"Failed to parse PSSH box: {e}") - return None diff --git a/lib/streaming_providers/base/utils/mp4_pssh_extractor.py b/lib/streaming_providers/base/utils/mp4_pssh_extractor.py new file mode 100644 index 0000000..e54c71b --- /dev/null +++ b/lib/streaming_providers/base/utils/mp4_pssh_extractor.py @@ -0,0 +1,470 @@ +""" +MP4 PSSH Extractor + +Extracts PSSH boxes and Key IDs from MP4 segments using the refactored DRM models. +""" + +import struct +from typing import Optional + +from ..models.drm import PSSHData, TencParser, PSSHParser +from ..models.drm.exceptions import InvalidPSSHError, InvalidTencError +from .logger import logger + + +class MP4PSSHExtractor: + """ + Extract PSSH boxes and Key IDs from MP4 segments. + + Uses the refactored DRM models for proper parsing and validation. + """ + + @staticmethod + def extract_from_url(segment_url: str, timeout: int = 10) -> list[PSSHData]: + """ + Download MP4 segment and extract PSSH data. + + Args: + segment_url: URL of the MP4 segment + timeout: Request timeout in seconds + + Returns: + List of PSSHData objects with extracted information + """ + import requests + + try: + response = requests.get(segment_url, timeout=timeout) + response.raise_for_status() + + # Only download first ~100KB for efficiency + chunk_size = 1024 * 100 + data = response.content[:chunk_size] + + return MP4PSSHExtractor.extract_from_bytes(data) + + except Exception as e: + logger.error(f"Failed to extract PSSH from {segment_url}: {e}") + return [] + + @staticmethod + def extract_from_bytes(data: bytes) -> list[PSSHData]: + """ + Extract PSSH boxes and encryption info from MP4 binary data. + + Process: + 1. Extract all tenc boxes to get default KIDs + 2. Extract all PSSH boxes + 3. Merge tenc KIDs with PSSH data where needed + + Args: + data: Raw MP4 binary data + + Returns: + List of PSSHData objects with complete information + """ + pssh_data_list = [] + offset = 0 + + # First pass: Extract all tenc boxes to get default KIDs + tenc_kids = MP4PSSHExtractor._extract_all_tenc_kids(data) + if tenc_kids: + logger.debug(f"Extracted {len(tenc_kids)} KIDs from tenc boxes") + + # Second pass: Extract PSSH boxes + while offset < len(data): + try: + # Read box size (4 bytes, big-endian) + if offset + 8 > len(data): + break + + box_size = struct.unpack(">I", data[offset: offset + 4])[0] + + # Handle special box sizes + if box_size == 0: + box_size = len(data) - offset # Box extends to end + elif box_size == 1: + # Extended size (64-bit) - skip for now + break + + if offset + box_size > len(data): + break + + # Read box type (4 bytes) + box_type = data[offset + 4: offset + 8] + + if box_type == b"moov": + # Look for PSSH in moov container + moov_data = data[offset: offset + box_size] + pssh_in_moov = MP4PSSHExtractor._extract_from_moov(moov_data) + pssh_data_list.extend(pssh_in_moov) + + elif box_type == b"pssh": + # Found standalone PSSH box + pssh_box = MP4PSSHExtractor._parse_pssh_box( + data[offset: offset + box_size] + ) + if pssh_box: + pssh_data_list.append(pssh_box) + + offset += box_size + + except Exception as e: + logger.debug(f"Error parsing MP4 box at offset {offset}: {e}") + offset += 1 # Try to recover + + # Third pass: Enhance PSSH data with tenc KIDs if needed + for pssh_data in pssh_data_list: + if pssh_data.needs_tenc_fallback() and tenc_kids: + logger.debug( + f"Adding {len(tenc_kids)} tenc KIDs to {pssh_data.drm_system} PSSH" + ) + pssh_data.add_key_ids(tenc_kids) + + return pssh_data_list + + @staticmethod + def _extract_all_tenc_kids(data: bytes) -> list[str]: + """ + Extract all Key IDs from tenc boxes in the MP4 data. + + Uses TencParser for proper tenc box parsing. + + Args: + data: Raw MP4 binary data + + Returns: + List of normalized Key IDs (32 hex chars, no hyphens) + """ + kids = [] + offset = 0 + + while offset < len(data): + try: + if offset + 8 > len(data): + break + + box_size = struct.unpack(">I", data[offset: offset + 4])[0] + box_type = data[offset + 4: offset + 8] + + if box_size < 8 or offset + box_size > len(data): + offset += 1 + continue + + if box_type == b"tenc": + # Extract tenc box data + tenc_data = data[offset: offset + box_size] + + # Use TencParser for proper parsing + try: + tenc_kids = TencParser.extract_kids_from_tenc(tenc_data) + for kid in tenc_kids: + if kid not in kids: + kids.append(kid) + logger.debug(f"Extracted KID from tenc: {kid[:8]}...") + except InvalidTencError as e: + logger.debug(f"Invalid tenc box at offset {offset}: {e}") + + elif box_size > 8: + # Recursively search container boxes + container_boxes = { + b"moov", b"trak", b"mdia", b"minf", + b"stbl", b"stsd", b"encv", b"enca", + b"sinf", b"schi" + } + if box_type in container_boxes: + # Search inside container + inner_data = data[offset + 8: offset + box_size] + inner_kids = MP4PSSHExtractor._extract_all_tenc_kids(inner_data) + for kid in inner_kids: + if kid not in kids: + kids.append(kid) + + offset += box_size + + except Exception as e: + logger.debug(f"Error at offset {offset}: {e}") + offset += 1 + + return kids + + @staticmethod + def _extract_from_moov(moov_data: bytes) -> list[PSSHData]: + """ + Extract PSSH boxes from moov container. + + Args: + moov_data: Raw moov box data + + Returns: + List of PSSHData objects found in moov + """ + pssh_list = [] + offset = 8 # Skip moov header + + while offset < len(moov_data): + try: + if offset + 8 > len(moov_data): + break + + box_size = struct.unpack(">I", moov_data[offset: offset + 4])[0] + box_type = moov_data[offset + 4: offset + 8] + + if box_size < 8 or offset + box_size > len(moov_data): + break + + if box_type == b"trak": + # Parse track for PSSH + trak_data = moov_data[offset: offset + box_size] + pssh_in_trak = MP4PSSHExtractor._extract_from_trak(trak_data) + pssh_list.extend(pssh_in_trak) + + elif box_type == b"pssh": + # PSSH directly in moov + pssh_box = MP4PSSHExtractor._parse_pssh_box( + moov_data[offset: offset + box_size] + ) + if pssh_box: + pssh_list.append(pssh_box) + + offset += box_size + + except Exception as e: + logger.debug(f"Error in moov at offset {offset}: {e}") + break + + return pssh_list + + @staticmethod + def _extract_from_trak(trak_data: bytes) -> list[PSSHData]: + """Extract PSSH from trak box""" + pssh_list = [] + offset = 8 + + while offset < len(trak_data): + try: + if offset + 8 > len(trak_data): + break + + box_size = struct.unpack(">I", trak_data[offset: offset + 4])[0] + box_type = trak_data[offset + 4: offset + 8] + + if box_size < 8 or offset + box_size > len(trak_data): + break + + if box_type == b"mdia": + mdia_data = trak_data[offset: offset + box_size] + pssh_in_mdia = MP4PSSHExtractor._extract_from_mdia(mdia_data) + pssh_list.extend(pssh_in_mdia) + + offset += box_size + + except Exception: + break + + return pssh_list + + @staticmethod + def _extract_from_mdia(mdia_data: bytes) -> list[PSSHData]: + """Extract PSSH from mdia box""" + pssh_list = [] + offset = 8 + + while offset < len(mdia_data): + try: + if offset + 8 > len(mdia_data): + break + + box_size = struct.unpack(">I", mdia_data[offset: offset + 4])[0] + box_type = mdia_data[offset + 4: offset + 8] + + if box_size < 8 or offset + box_size > len(mdia_data): + break + + if box_type == b"minf": + minf_data = mdia_data[offset: offset + box_size] + pssh_in_minf = MP4PSSHExtractor._extract_from_minf(minf_data) + pssh_list.extend(pssh_in_minf) + + offset += box_size + + except Exception: + break + + return pssh_list + + @staticmethod + def _extract_from_minf(minf_data: bytes) -> list[PSSHData]: + """Extract PSSH from minf box""" + pssh_list = [] + offset = 8 + + while offset < len(minf_data): + try: + if offset + 8 > len(minf_data): + break + + box_size = struct.unpack(">I", minf_data[offset: offset + 4])[0] + box_type = minf_data[offset + 4: offset + 8] + + if box_size < 8 or offset + box_size > len(minf_data): + break + + if box_type == b"stbl": + stbl_data = minf_data[offset: offset + box_size] + pssh_in_stbl = MP4PSSHExtractor._extract_from_stbl(stbl_data) + pssh_list.extend(pssh_in_stbl) + + offset += box_size + + except Exception: + break + + return pssh_list + + @staticmethod + def _extract_from_stbl(stbl_data: bytes) -> list[PSSHData]: + """Extract PSSH from stbl box (where protection scheme info usually is)""" + pssh_list = [] + offset = 8 + + while offset < len(stbl_data): + try: + if offset + 8 > len(stbl_data): + break + + box_size = struct.unpack(">I", stbl_data[offset: offset + 4])[0] + box_type = stbl_data[offset + 4: offset + 8] + + if box_size < 8 or offset + box_size > len(stbl_data): + break + + if box_type == b"sinf": + sinf_data = stbl_data[offset: offset + box_size] + pssh_in_sinf = MP4PSSHExtractor._extract_from_sinf(sinf_data) + pssh_list.extend(pssh_in_sinf) + + offset += box_size + + except Exception: + break + + return pssh_list + + @staticmethod + def _extract_from_sinf(sinf_data: bytes) -> list[PSSHData]: + """Extract PSSH from sinf (protection scheme information) box""" + pssh_list = [] + offset = 8 + + while offset < len(sinf_data): + try: + if offset + 8 > len(sinf_data): + break + + box_size = struct.unpack(">I", sinf_data[offset: offset + 4])[0] + box_type = sinf_data[offset + 4: offset + 8] + + if box_size < 8 or offset + box_size > len(sinf_data): + break + + if box_type == b"schi": + schi_data = sinf_data[offset: offset + box_size] + pssh_in_schi = MP4PSSHExtractor._extract_from_schi(schi_data) + pssh_list.extend(pssh_in_schi) + + offset += box_size + + except Exception: + break + + return pssh_list + + @staticmethod + def _extract_from_schi(schi_data: bytes) -> list[PSSHData]: + """Extract PSSH from schi box (where PSSH boxes are typically stored)""" + pssh_list = [] + offset = 8 + + while offset < len(schi_data): + try: + if offset + 8 > len(schi_data): + break + + box_size = struct.unpack(">I", schi_data[offset: offset + 4])[0] + box_type = schi_data[offset + 4: offset + 8] + + if box_size < 8 or offset + box_size > len(schi_data): + break + + if box_type == b"pssh": + pssh_box = MP4PSSHExtractor._parse_pssh_box( + schi_data[offset: offset + box_size] + ) + if pssh_box: + pssh_list.append(pssh_box) + + offset += box_size + + except Exception: + break + + return pssh_list + + @staticmethod + def _parse_pssh_box(pssh_bytes: bytes) -> Optional[PSSHData]: + """ + Parse PSSH box and create PSSHData object. + + Uses PSSHParser for proper parsing and validation. + + Args: + pssh_bytes: Raw PSSH box bytes + + Returns: + PSSHData object or None if parsing fails + """ + try: + # Basic validation + if len(pssh_bytes) < 32: + logger.debug(f"PSSH box too small: {len(pssh_bytes)} bytes") + return None + + box_type = pssh_bytes[4:8] + if box_type != b"pssh": + logger.debug(f"Not a PSSH box: {box_type}") + return None + + # Encode entire PSSH box as base64 + from ..models.drm.utils import safe_base64_encode + pssh_b64 = safe_base64_encode(pssh_bytes) + + # Parse PSSH to get system_id and metadata + try: + metadata = PSSHParser.parse_pssh_box(pssh_b64) + except InvalidPSSHError as e: + logger.debug(f"Failed to parse PSSH: {e}") + return None + + # Create PSSHData with parsed information + pssh_data = PSSHData( + system_id=metadata["system_id"], + pssh_box=pssh_b64, + key_ids=metadata["key_ids"], + source="mp4_segment", + ) + + # Log what we found + drm_name = pssh_data.drm_system.name if pssh_data.drm_system else "UNKNOWN" + kid_count = len(pssh_data.key_ids) + version = metadata["version"] + + logger.debug( + f"Parsed PSSH: {drm_name} v{version} with {kid_count} KIDs" + ) + + return pssh_data + + except Exception as e: + logger.debug(f"Failed to parse PSSH box: {e}") + return None \ No newline at end of file diff --git a/lib/streaming_providers/base/utils/mpd_rewriter.py b/lib/streaming_providers/base/utils/mpd_rewriter.py index 467b98e..818169a 100644 --- a/lib/streaming_providers/base/utils/mpd_rewriter.py +++ b/lib/streaming_providers/base/utils/mpd_rewriter.py @@ -1,5 +1,6 @@ # streaming_providers/base/utils/mpd_rewriter.py import base64 +import struct import xml.etree.ElementTree as ET import re from typing import Optional, Tuple, Set, Dict, List @@ -417,8 +418,8 @@ class MPDRewriter: return False + @staticmethod def _parse_video_representation( - self, representation: ET.Element, adaptation_set: ET.Element, period: ET.Element, @@ -568,7 +569,8 @@ class MPDRewriter: # Check default_KID attribute kid_attr = cp.get("{urn:mpeg:cenc:2013}default_KID") if kid_attr: - return kid_attr.replace("-", "").lower() + # Let the model normalize it later - just return raw + return kid_attr # Check cenc:pssh for pssh in cp.findall("cenc:pssh", self.CENC_NAMESPACE): @@ -576,8 +578,13 @@ class MPDRewriter: try: pssh_data = base64.b64decode(pssh.text) if len(pssh_data) >= 36: - kid_bytes = pssh_data[32:48] - return kid_bytes.hex() + version = pssh_data[8] if len(pssh_data) > 8 else 0 + if version > 0: + kid_count = struct.unpack(">I", pssh_data[28:32])[0] + if kid_count > 0 and len(pssh_data) >= 48: + kid_bytes = pssh_data[32:48] + # Return raw hex - KeyConfiguration will normalize! + return kid_bytes.hex() except Exception: continue