# ============================================================================ # streaming_providers/base/drm_operations.py """ DRM-related operations with caching and optimized PSSH extraction. """ import time from threading import Lock from typing import Dict, List, Optional, Tuple from .drm import DRMPluginManager from .models import DRMSystem from .utils.logger import logger class PSSHCache: """Thread-safe cache for PSSH data""" def __init__(self, ttl_seconds: int = 3600): self.cache: Dict[str, Tuple[List, float]] = {} self.ttl = ttl_seconds self.lock = Lock() def get(self, key: str) -> Optional[List]: """Get cached PSSH data if not expired""" with self.lock: if key in self.cache: pssh_list, timestamp = self.cache[key] if time.time() - timestamp < self.ttl: logger.debug(f"Cache HIT for {key}") return pssh_list else: logger.debug(f"Cache EXPIRED for {key}") del self.cache[key] return None def set(self, key: str, pssh_list: List): """Cache PSSH data""" with self.lock: self.cache[key] = (pssh_list, time.time()) logger.debug(f"Cache SET for {key}") def clear(self): """Clear all cache entries""" with self.lock: self.cache.clear() logger.debug("Cache CLEARED") class DRMConfigCache: """Thread-safe cache for DRM configurations (ClearKey only, no expiry)""" def __init__(self): self.cache: Dict[str, List] = {} self.lock = Lock() def get(self, key: str) -> Optional[List]: """Get cached DRM configs""" with self.lock: if key in self.cache: logger.debug(f"DRM Config Cache HIT for {key}") return self.cache[key] return None def set(self, key: str, drm_configs: List): """Cache DRM configs (only if contains ClearKey)""" with self.lock: self.cache[key] = drm_configs logger.debug(f"DRM Config Cache SET for {key}") def clear(self): """Clear all cache entries""" with self.lock: self.cache.clear() logger.debug("DRM Config Cache CLEARED") class DRMOperations: """Handles all DRM-related operations.""" def __init__(self, registry, cache_ttl: int = 3600): self.registry = registry self.drm_plugin_manager = DRMPluginManager() self.pssh_cache = PSSHCache(ttl_seconds=cache_ttl) self.drm_config_cache = DRMConfigCache() # New: ClearKey config cache logger.debug("DRMOperations: Initialized with caching") def get_channel_drm_configs(self, provider_name: str, channel_id: str, **kwargs) -> List: """Get DRM configurations for a channel.""" cache_key = f"{provider_name}:{channel_id}" # Check DRM config cache first (ClearKey configs) cached_configs = self.drm_config_cache.get(cache_key) if cached_configs is not None: logger.info(f"Using cached DRM configs for '{channel_id}'") return cached_configs # Cache miss - get from provider provider = self.registry.get_provider(provider_name) if not provider: raise ValueError(f"Provider '{provider_name}' not found or disabled") drm_configs = provider.get_drm(channel_id, **kwargs) # Extract PSSH if needed pssh_data_list = [] if drm_configs and self.drm_plugin_manager.plugins: if self._needs_pssh_extraction(drm_configs): # Try PSSH cache FIRST before fetching manifest pssh_data_list = self.pssh_cache.get(cache_key) if pssh_data_list is None: # Cache miss - now fetch manifest and extract logger.debug(f"PSSH cache miss for {cache_key}, fetching manifest") manifest_url = provider.get_manifest(channel_id, **kwargs) if manifest_url: pssh_data_list = self._extract_pssh_from_manifest(manifest_url) if pssh_data_list: self.pssh_cache.set(cache_key, pssh_data_list) else: logger.debug(f"Using cached PSSH for {cache_key}") # Process through plugins processed = self.drm_plugin_manager.process_drm_configs( drm_configs, pssh_data_list, **kwargs ) # Cache if contains ClearKey config if self._has_clearkey_config(processed): logger.info(f"Caching ClearKey DRM configs for '{channel_id}'") self.drm_config_cache.set(cache_key, processed) logger.info(f"Processed DRM for '{channel_id}': {len(processed)} configs") return processed @staticmethod def _has_clearkey_config(drm_configs: List) -> bool: """Check if any config is a ClearKey config""" return any(config.system == DRMSystem.CLEARKEY for config in drm_configs) def _needs_pssh_extraction(self, drm_configs) -> bool: """Check if PSSH extraction is needed.""" config_systems = {config.system for config in drm_configs} plugin_systems = set(self.drm_plugin_manager.plugins.keys()) return bool(config_systems & plugin_systems or DRMSystem.GENERIC in plugin_systems) @staticmethod def _extract_pssh_from_manifest(manifest_url: str) -> List: """Extract PSSH data from manifest with single segment fallback.""" import requests from .utils.manifest_parser import ManifestParser try: response = requests.get(manifest_url, timeout=10) response.raise_for_status() manifest_content = response.text # Extract from manifest content pssh_list = ManifestParser._extract_from_manifest_content(manifest_content) # Check if we need segment extraction needs_segment_extraction = not pssh_list or any( not p.pssh_box or not p.key_ids for p in pssh_list ) if needs_segment_extraction: logger.debug("PSSH incomplete in manifest, extracting from init segment") # Extract ONE init segment URL init_segment_url = ManifestParser.extract_single_init_segment_url( manifest_content, manifest_url ) if init_segment_url: # Get expected system IDs expected_system_ids = [p.system_id for p in pssh_list] if pssh_list else [] segment_pssh = ManifestParser._extract_from_single_segment( init_segment_url, expected_system_ids ) if segment_pssh: # Merge or replace return ManifestParser._merge_pssh_data(pssh_list, segment_pssh) else: logger.warning("Could not extract init segment URL from manifest") return pssh_list except Exception as e: logger.warning(f"Failed to extract PSSH: {e}") return [] def list_drm_plugins(self) -> Dict: """List registered DRM plugins.""" return self.drm_plugin_manager.list_plugins() def clear_drm_plugins(self): """Clear all DRM plugins.""" self.drm_plugin_manager.clear_plugins() def clear_pssh_cache(self): """Clear PSSH cache.""" self.pssh_cache.clear() def clear_drm_config_cache(self): """Clear DRM config cache.""" self.drm_config_cache.clear()