From e97affd2da31d48c7705a0b8d4eee1e45d885ac6 Mon Sep 17 00:00:00 2001 From: Nirvana Date: Thu, 29 Jan 2026 17:20:48 +0100 Subject: [PATCH] Reformat --- .../base/utils/mpd_rewriter.py | 218 +++++++++++++++++- 1 file changed, 210 insertions(+), 8 deletions(-) diff --git a/lib/streaming_providers/base/utils/mpd_rewriter.py b/lib/streaming_providers/base/utils/mpd_rewriter.py index 23bc96c..9449313 100644 --- a/lib/streaming_providers/base/utils/mpd_rewriter.py +++ b/lib/streaming_providers/base/utils/mpd_rewriter.py @@ -2,7 +2,7 @@ import base64 import xml.etree.ElementTree as ET import re -from typing import Optional, Tuple, Set, Dict +from typing import Optional, Tuple, Set, Dict, List from urllib.parse import urljoin, urlparse, quote, urlencode from datetime import datetime, timezone from email.utils import parsedate_to_datetime @@ -51,6 +51,30 @@ class KeyConfiguration: logger.debug(f"Multi-key mode: {len(self.keys)} keys available") +@dataclass +class VideoRepresentation: + """Represents a video representation with its quality metrics.""" + element: ET.Element + adaptation_set: ET.Element + period: ET.Element + bandwidth: int + width: Optional[int] = None + height: Optional[int] = None + frame_rate: Optional[float] = None + representation_id: str = "" + period_id: str = "" + as_id: str = "" + + def get_quality_score(self) -> Tuple[int, int, int]: + """ + Returns a tuple for comparison: (resolution_pixels, bandwidth, frame_rate_score) + Higher values = better quality + """ + resolution = (self.width or 0) * (self.height or 0) + frame_rate_score = int((self.frame_rate or 0) * 100) + return resolution, self.bandwidth, frame_rate_score + + class MPDRewriter: MPD_NAMESPACE = {"mpd": "urn:mpeg:dash:schema:mpd:2011"} CENC_NAMESPACE = {"cenc": "urn:mpeg:cenc:2013"} @@ -60,10 +84,12 @@ class MPDRewriter: media_proxy_url: str, provider_proxy_url: Optional[str] = None, clearkey_keyids: Optional[dict] = None, + highest_quality_video_only: bool = False, # NEW: Enable highest quality filtering ): self.media_proxy_url = media_proxy_url.rstrip("/") self.provider_proxy_url = provider_proxy_url self.key_config = KeyConfiguration(clearkey_keyids or {}) + self.highest_quality_video_only = highest_quality_video_only # Pre-calculate query params that don't change to save cycles during rewrite self._static_params = {} @@ -88,6 +114,7 @@ class MPDRewriter: segment_type: Optional[str] = None, is_encrypted: bool = False, kid: Optional[str] = None, # Specific KID for this AdaptationSet + representation_id: Optional[str] = None, # NEW: For template substitution ) -> str: params = {"url": original_url, **self._static_params} @@ -128,6 +155,10 @@ class MPDRewriter: proxy_url = f"{self.media_proxy_url}/api/{endpoint}/{encoded}" if template_pattern: + # NEW: Substitute $RepresentationID$ if we have it and highest_quality_video_only is enabled + if self.highest_quality_video_only and representation_id and "$RepresentationID$" in template_pattern: + template_pattern = template_pattern.replace("$RepresentationID$", representation_id) + proxy_url += f"/{quote(template_pattern, safe='.-_$')}" return proxy_url @@ -153,19 +184,31 @@ class MPDRewriter: # Single-pass tree preparation with BaseURL extraction encrypted_ids, as_id_to_kid, base_url_map = self._prepare_tree_and_extract_kids(root, mpd_base_url) - # Filter out encrypted AdaptationSets without available keys + # FIRST: Filter out encrypted AdaptationSets without available keys + # This ensures we only consider decryptable content for quality selection if self.key_config.keys: self._remove_adaptationsets_without_keys(root, as_id_to_kid) else: self._remove_all_encrypted_adaptationsets(root) + # THEN: Filter to highest quality video from remaining decryptable content + best_video_info = None + if self.highest_quality_video_only: + best_video_info = self._filter_to_highest_quality_video(root) + if best_video_info: + logger.info( + f"Filtered to highest quality video: {best_video_info.width}x{best_video_info.height} " + f"@ {best_video_info.bandwidth}bps, RepID={best_video_info.representation_id}" + ) + # Verify we have playable content remaining remaining_sets = root.findall(".//mpd:AdaptationSet", self.MPD_NAMESPACE) if not remaining_sets: raise ValueError("No AdaptationSets remain after key filtering - manifest would be empty") # Rewrite URLs with appropriate keys and context-aware base URLs - self._rewrite_node(root, mpd_base_url, encrypted_ids, as_id_to_kid, base_url_map, False, None, "") + self._rewrite_node(root, mpd_base_url, encrypted_ids, as_id_to_kid, base_url_map, + False, None, "", best_video_info) rewritten = ET.tostring(root, encoding="unicode", method="xml") if not rewritten.startswith(" Tuple[Set[str], Dict[str, str], Dict[str, str]]: + def _filter_to_highest_quality_video(self, root: ET.Element) -> Optional[VideoRepresentation]: + """ + Find the absolute highest quality video representation across all periods and adaptation sets. + Remove all other video representations, keep audio/subtitles unchanged. + Returns information about the best video representation found. + """ + all_video_reps: List[VideoRepresentation] = [] + + # Scan all periods and adaptation sets to find video representations + for period in root.findall(".//mpd:Period", self.MPD_NAMESPACE): + period_id = period.get("id", "") + + for adaptation_set in period.findall("mpd:AdaptationSet", self.MPD_NAMESPACE): + # Check if this is a video adaptation set + if not self._is_video_adaptation_set(adaptation_set): + continue + + as_id = adaptation_set.get("id", str(id(adaptation_set))) + + # Collect all representations from this video adaptation set + for representation in adaptation_set.findall("mpd:Representation", self.MPD_NAMESPACE): + video_rep = self._parse_video_representation( + representation, adaptation_set, period, period_id, as_id + ) + if video_rep: + all_video_reps.append(video_rep) + + if not all_video_reps: + logger.warning("No video representations found for filtering") + return None + + # Find the best video representation + best_video = max(all_video_reps, key=lambda v: v.get_quality_score()) + logger.debug( + f"Found {len(all_video_reps)} video representations, " + f"best: {best_video.width}x{best_video.height} @ {best_video.bandwidth}bps" + ) + + # Now remove all video representations EXCEPT the best one + removal_count = 0 + for period in root.findall(".//mpd:Period", self.MPD_NAMESPACE): + period_id = period.get("id", "") + adaptation_sets_to_remove = [] + + for adaptation_set in period.findall("mpd:AdaptationSet", self.MPD_NAMESPACE): + if not self._is_video_adaptation_set(adaptation_set): + continue + + as_id = adaptation_set.get("id", str(id(adaptation_set))) + + # Check if this adaptation set contains the best representation + representations_to_remove = [] + contains_best = False + + for representation in adaptation_set.findall("mpd:Representation", self.MPD_NAMESPACE): + rep_id = representation.get("id", "") + + # Is this the best representation? + if (period_id == best_video.period_id and + as_id == best_video.as_id and + rep_id == best_video.representation_id): + contains_best = True + else: + representations_to_remove.append(representation) + + # Remove non-best representations + for rep in representations_to_remove: + adaptation_set.remove(rep) + removal_count += 1 + + # If this adaptation set no longer has any representations, mark for removal + if not contains_best: + adaptation_sets_to_remove.append(adaptation_set) + + # Remove empty video adaptation sets + for adaptation_set in adaptation_sets_to_remove: + period.remove(adaptation_set) + + logger.info(f"Removed {removal_count} video representation(s), kept highest quality only") + return best_video + + def _is_video_adaptation_set(self, adaptation_set: ET.Element) -> bool: + """Determine if an AdaptationSet is video based on mimeType or contentType.""" + mime_type = adaptation_set.get("mimeType", "") + content_type = adaptation_set.get("contentType", "") + + # Check AdaptationSet level + if mime_type.startswith("video/") or content_type == "video": + return True + + # Check Representation level if not specified at AdaptationSet level + for representation in adaptation_set.findall("mpd:Representation", self.MPD_NAMESPACE): + rep_mime = representation.get("mimeType", "") + if rep_mime.startswith("video/"): + return True + + return False + + def _parse_video_representation( + self, + representation: ET.Element, + adaptation_set: ET.Element, + period: ET.Element, + period_id: str, + as_id: str + ) -> Optional[VideoRepresentation]: + """Parse a Representation element and extract video quality information.""" + try: + # Bandwidth is required + bandwidth = int(representation.get("bandwidth", "0")) + if bandwidth == 0: + return None + + # Width and height (can be on Representation or AdaptationSet) + width = representation.get("width") or adaptation_set.get("width") + height = representation.get("height") or adaptation_set.get("height") + + width = int(width) if width else None + height = int(height) if height else None + + # Frame rate (can be on Representation or AdaptationSet) + frame_rate_str = representation.get("frameRate") or adaptation_set.get("frameRate") + frame_rate = None + if frame_rate_str: + # Handle both "30" and "30000/1001" formats + if "/" in frame_rate_str: + num, denom = frame_rate_str.split("/") + frame_rate = float(num) / float(denom) + else: + frame_rate = float(frame_rate_str) + + representation_id = representation.get("id", "") + + return VideoRepresentation( + element=representation, + adaptation_set=adaptation_set, + period=period, + bandwidth=bandwidth, + width=width, + height=height, + frame_rate=frame_rate, + representation_id=representation_id, + period_id=period_id, + as_id=as_id + ) + except (ValueError, AttributeError) as e: + logger.warning(f"Failed to parse video representation: {e}") + return None + + def _prepare_tree_and_extract_kids(self, root: ET.Element, mpd_base_url: str) -> Tuple[ + Set[str], Dict[str, str], Dict[str, str]]: """ Single-pass optimization: clean tree, identify encrypted sets, extract KIDs and BaseURLs. Returns: (encrypted_adaptation_set_ids, as_id_to_kid_mapping, base_url_mapping) @@ -354,6 +547,7 @@ class MPDRewriter: current_encrypted: bool, current_kid: Optional[str] = None, current_period_id: str = "", + best_video_info: Optional[VideoRepresentation] = None, # NEW ): """Recursive node rewriter with KID-aware key selection and context-aware base URLs.""" # Track period ID as we traverse @@ -362,6 +556,7 @@ class MPDRewriter: # Update state when entering an AdaptationSet current_as_id = None + current_rep_id = None # NEW: Track current representation ID if element.tag.endswith("AdaptationSet"): as_id = element.get("id", str(id(element))) current_as_id = as_id @@ -377,6 +572,10 @@ class MPDRewriter: if current_encrypted and not self.key_config.single_key_mode: current_kid = as_id_to_kid.get(unique_id) + # NEW: Track representation ID for template substitution + if element.tag.endswith("Representation"): + current_rep_id = element.get("id", "") + # Rewrite URL attributes attr_map = { "media": "media", @@ -394,11 +593,13 @@ class MPDRewriter: if "$" in resolved: path, pattern = self.split_template_url(resolved) element.attrib[attr] = self.build_proxy_url( - path, pattern, seg_type, current_encrypted, current_kid + path, pattern, seg_type, current_encrypted, current_kid, + representation_id=current_rep_id # NEW ) else: element.attrib[attr] = self.build_proxy_url( - resolved, None, seg_type, current_encrypted, current_kid + resolved, None, seg_type, current_encrypted, current_kid, + representation_id=current_rep_id # NEW ) # Handle SegmentURL (always 'media' type) @@ -410,14 +611,15 @@ class MPDRewriter: else (resolved, None) ) element.attrib["media"] = self.build_proxy_url( - path, pattern, "media", current_encrypted, current_kid + path, pattern, "media", current_encrypted, current_kid, + representation_id=current_rep_id # NEW ) # Recurse to children for child in element: self._rewrite_node( child, base_url, encrypted_ids, as_id_to_kid, base_url_map, - current_encrypted, current_kid, current_period_id + current_encrypted, current_kid, current_period_id, best_video_info # NEW ) def _extract_mpd_base_url(self, root: ET.Element, manifest_url: str) -> str: