Spaces:
Configuration error
Configuration error
| import asyncio | |
| import logging | |
| import math | |
| import time | |
| from fastapi import Request, Response, HTTPException | |
| from mediaflow_proxy.drm.decrypter import decrypt_segment, process_drm_init_segment | |
| from mediaflow_proxy.utils.crypto_utils import encryption_handler | |
| from mediaflow_proxy.utils.http_utils import ( | |
| encode_mediaflow_proxy_url, | |
| get_original_scheme, | |
| ProxyRequestHeaders, | |
| apply_header_manipulation, | |
| ) | |
| from mediaflow_proxy.utils.dash_prebuffer import dash_prebuffer | |
| from mediaflow_proxy.utils.cache_utils import get_cached_processed_init, set_cached_processed_init | |
| from mediaflow_proxy.utils.m3u8_processor import SkipSegmentFilter | |
| from mediaflow_proxy.configs import settings | |
| logger = logging.getLogger(__name__) | |
| async def process_manifest( | |
| request: Request, | |
| mpd_dict: dict, | |
| proxy_headers: ProxyRequestHeaders, | |
| key_id: str = None, | |
| key: str = None, | |
| resolution: str = None, | |
| skip_segments: list = None, | |
| ) -> Response: | |
| """ | |
| Processes the MPD manifest and converts it to an HLS manifest. | |
| Args: | |
| request (Request): The incoming HTTP request. | |
| mpd_dict (dict): The MPD manifest data. | |
| proxy_headers (ProxyRequestHeaders): The headers to include in the request. | |
| key_id (str, optional): The DRM key ID. Defaults to None. | |
| key (str, optional): The DRM key. Defaults to None. | |
| resolution (str, optional): Target resolution (e.g., '1080p', '720p'). Defaults to None. | |
| skip_segments (list, optional): List of time segments to skip. Each item should have 'start' and 'end' keys. | |
| Returns: | |
| Response: The HLS manifest as an HTTP response. | |
| """ | |
| hls_content = build_hls(mpd_dict, request, key_id, key, resolution, skip_segments) | |
| # Start DASH pre-buffering in background if enabled | |
| if settings.enable_dash_prebuffer: | |
| # Extract headers for pre-buffering | |
| headers = {} | |
| for key, value in request.query_params.items(): | |
| if key.startswith("h_"): | |
| headers[key[2:]] = value | |
| # Get the original MPD URL from the request | |
| mpd_url = request.query_params.get("d", "") | |
| if mpd_url: | |
| # Start pre-buffering in background | |
| asyncio.create_task(dash_prebuffer.prebuffer_dash_manifest(mpd_url, headers)) | |
| return Response(content=hls_content, media_type="application/vnd.apple.mpegurl", headers=proxy_headers.response) | |
| async def process_playlist( | |
| request: Request, | |
| mpd_dict: dict, | |
| profile_id: str, | |
| proxy_headers: ProxyRequestHeaders, | |
| skip_segments: list = None, | |
| start_offset: float = None, | |
| ) -> Response: | |
| """ | |
| Processes the MPD manifest and converts it to an HLS playlist for a specific profile. | |
| Args: | |
| request (Request): The incoming HTTP request. | |
| mpd_dict (dict): The MPD manifest data. | |
| profile_id (str): The profile ID to generate the playlist for. | |
| proxy_headers (ProxyRequestHeaders): The headers to include in the request. | |
| skip_segments (list, optional): List of time segments to skip. Each item should have 'start' and 'end' keys. | |
| start_offset (float, optional): Start offset in seconds for live streams. | |
| Returns: | |
| Response: The HLS playlist as an HTTP response. | |
| Raises: | |
| HTTPException: If the profile is not found in the MPD manifest. | |
| """ | |
| matching_profiles = [p for p in mpd_dict["profiles"] if p["id"] == profile_id] | |
| if not matching_profiles: | |
| raise HTTPException(status_code=404, detail="Profile not found") | |
| hls_content = build_hls_playlist(mpd_dict, matching_profiles, request, skip_segments, start_offset) | |
| # Trigger prebuffering of upcoming segments for live streams | |
| if settings.enable_dash_prebuffer and mpd_dict.get("isLive", False): | |
| # Extract headers for pre-buffering | |
| headers = {} | |
| for key, value in request.query_params.items(): | |
| if key.startswith("h_"): | |
| headers[key[2:]] = value | |
| # Use the new prefetch method for live playlists | |
| asyncio.create_task(dash_prebuffer.prefetch_for_live_playlist(matching_profiles, headers)) | |
| # Don't include propagate headers for playlists - they should only apply to segments | |
| response_headers = apply_header_manipulation({}, proxy_headers, include_propagate=False) | |
| return Response(content=hls_content, media_type="application/vnd.apple.mpegurl", headers=response_headers) | |
| async def process_segment( | |
| init_content: bytes, | |
| segment_content: bytes, | |
| mimetype: str, | |
| proxy_headers: ProxyRequestHeaders, | |
| key_id: str = None, | |
| key: str = None, | |
| use_map: bool = False, | |
| ) -> Response: | |
| """ | |
| Processes and decrypts a media segment. | |
| Args: | |
| init_content (bytes): The initialization segment content. | |
| segment_content (bytes): The media segment content. | |
| mimetype (str): The MIME type of the segment. | |
| proxy_headers (ProxyRequestHeaders): The headers to include in the request. | |
| key_id (str, optional): The DRM key ID. Defaults to None. | |
| key (str, optional): The DRM key. Defaults to None. | |
| use_map (bool, optional): If True, init segment is served separately via EXT-X-MAP, | |
| so don't concatenate init with segment. Defaults to False. | |
| Returns: | |
| Response: The decrypted segment as an HTTP response. | |
| """ | |
| if key_id and key: | |
| # For DRM protected content | |
| now = time.time() | |
| decrypted_content = decrypt_segment(init_content, segment_content, key_id, key, include_init=not use_map) | |
| logger.info(f"Decryption of {mimetype} segment took {time.time() - now:.4f} seconds") | |
| else: | |
| # For non-DRM protected content | |
| if use_map: | |
| # Init is served separately via EXT-X-MAP | |
| decrypted_content = segment_content | |
| else: | |
| # Concatenate init and segment content | |
| decrypted_content = init_content + segment_content | |
| response_headers = apply_header_manipulation({}, proxy_headers) | |
| return Response(content=decrypted_content, media_type=mimetype, headers=response_headers) | |
| async def process_init_segment( | |
| init_content: bytes, | |
| mimetype: str, | |
| proxy_headers: ProxyRequestHeaders, | |
| key_id: str = None, | |
| key: str = None, | |
| init_url: str = None, | |
| ) -> Response: | |
| """ | |
| Processes an initialization segment for EXT-X-MAP. | |
| Args: | |
| init_content (bytes): The initialization segment content. | |
| mimetype (str): The MIME type of the segment. | |
| proxy_headers (ProxyRequestHeaders): The headers to include in the request. | |
| key_id (str, optional): The DRM key ID. Defaults to None. | |
| key (str, optional): The DRM key. Defaults to None. | |
| init_url (str, optional): The init URL for caching. Defaults to None. | |
| Returns: | |
| Response: The processed init segment as an HTTP response. | |
| """ | |
| if key_id and key: | |
| # Check if we have a cached processed version | |
| if init_url: | |
| cached_processed = await get_cached_processed_init(init_url, key_id) | |
| if cached_processed: | |
| logger.debug(f"Using cached processed init segment for {init_url}") | |
| response_headers = apply_header_manipulation({}, proxy_headers) | |
| return Response(content=cached_processed, media_type=mimetype, headers=response_headers) | |
| # For DRM protected content, we need to process the init segment | |
| # to remove encryption-related boxes but keep the moov structure | |
| now = time.time() | |
| processed_content = process_drm_init_segment(init_content, key_id, key) | |
| logger.info(f"Processing of {mimetype} init segment took {time.time() - now:.4f} seconds") | |
| # Cache the processed init segment | |
| if init_url: | |
| await set_cached_processed_init(init_url, key_id, processed_content, ttl=3600) | |
| else: | |
| # For non-DRM protected content, just return the init segment as-is | |
| processed_content = init_content | |
| response_headers = apply_header_manipulation({}, proxy_headers) | |
| return Response(content=processed_content, media_type=mimetype, headers=response_headers) | |
| def build_hls( | |
| mpd_dict: dict, | |
| request: Request, | |
| key_id: str = None, | |
| key: str = None, | |
| resolution: str = None, | |
| skip_segments: list = None, | |
| ) -> str: | |
| """ | |
| Builds an HLS manifest from the MPD manifest. | |
| Args: | |
| mpd_dict (dict): The MPD manifest data. | |
| request (Request): The incoming HTTP request. | |
| key_id (str, optional): The DRM key ID. Defaults to None. | |
| key (str, optional): The DRM key. Defaults to None. | |
| resolution (str, optional): Target resolution (e.g., '1080p', '720p'). Defaults to None. | |
| skip_segments (list, optional): List of time segments to skip. Each item should have 'start' and 'end' keys. | |
| Returns: | |
| str: The HLS manifest as a string. | |
| """ | |
| hls = ["#EXTM3U", "#EXT-X-VERSION:6"] | |
| query_params = dict(request.query_params) | |
| # Preserve skip parameter in query params so it propagates to playlists | |
| if skip_segments: | |
| # Convert back to compact format for URL | |
| skip_str = ",".join(f"{s['start']}-{s['end']}" for s in skip_segments) | |
| query_params["skip"] = skip_str | |
| has_encrypted = query_params.pop("has_encrypted", False) | |
| video_profiles = {} | |
| audio_profiles = {} | |
| # Get the base URL for the playlist_endpoint endpoint | |
| proxy_url = request.url_for("playlist_endpoint") | |
| proxy_url = str(proxy_url.replace(scheme=get_original_scheme(request))) | |
| for profile in mpd_dict["profiles"]: | |
| query_params.update({"profile_id": profile["id"], "key_id": key_id or "", "key": key or ""}) | |
| playlist_url = encode_mediaflow_proxy_url( | |
| proxy_url, | |
| query_params=query_params, | |
| encryption_handler=encryption_handler if has_encrypted else None, | |
| ) | |
| if "video" in profile["mimeType"]: | |
| video_profiles[profile["id"]] = (profile, playlist_url) | |
| elif "audio" in profile["mimeType"]: | |
| audio_profiles[profile["id"]] = (profile, playlist_url) | |
| # Filter video profiles by resolution if specified | |
| if resolution and video_profiles: | |
| video_profiles = _filter_video_profiles_by_resolution(video_profiles, resolution) | |
| # Add audio streams | |
| for i, (profile, playlist_url) in enumerate(audio_profiles.values()): | |
| is_default = "YES" if i == 0 else "NO" # Set the first audio track as default | |
| hls.append( | |
| f'#EXT-X-MEDIA:TYPE=AUDIO,GROUP-ID="audio",NAME="{profile["id"]}",DEFAULT={is_default},AUTOSELECT={is_default},LANGUAGE="{profile.get("lang", "und")}",URI="{playlist_url}"' | |
| ) | |
| # Add video streams | |
| for profile, playlist_url in video_profiles.values(): | |
| # Only add AUDIO attribute if there are audio profiles available | |
| audio_attr = ',AUDIO="audio"' if audio_profiles else "" | |
| hls.append( | |
| f'#EXT-X-STREAM-INF:BANDWIDTH={profile["bandwidth"]},RESOLUTION={profile["width"]}x{profile["height"]},CODECS="{profile["codecs"]}",FRAME-RATE={profile["frameRate"]}{audio_attr}' | |
| ) | |
| hls.append(playlist_url) | |
| return "\n".join(hls) | |
| def _filter_video_profiles_by_resolution(video_profiles: dict, target_resolution: str) -> dict: | |
| """ | |
| Filter video profiles to select the one matching the target resolution. | |
| Falls back to closest lower resolution if exact match not found. | |
| Args: | |
| video_profiles: Dictionary of profile_id -> (profile, playlist_url). | |
| target_resolution: Target resolution string (e.g., '1080p', '720p'). | |
| Returns: | |
| Filtered dictionary with only the selected profile. | |
| """ | |
| # Parse target height from "1080p" -> 1080 | |
| target_height = int(target_resolution.rstrip("p")) | |
| # Convert to list and sort by height descending | |
| profiles_list = [ | |
| (profile_id, profile, playlist_url) | |
| for profile_id, (profile, playlist_url) in video_profiles.items() | |
| if profile.get("height", 0) > 0 | |
| ] | |
| if not profiles_list: | |
| logger.warning("No video profiles with valid height found, returning all profiles") | |
| return video_profiles | |
| sorted_profiles = sorted(profiles_list, key=lambda x: x[1]["height"], reverse=True) | |
| # Find exact match or closest lower | |
| selected = None | |
| for profile_id, profile, playlist_url in sorted_profiles: | |
| if profile["height"] <= target_height: | |
| selected = (profile_id, profile, playlist_url) | |
| break | |
| # If all profiles are higher than target, use lowest available | |
| if selected is None: | |
| selected = sorted_profiles[-1] | |
| profile_id, profile, playlist_url = selected | |
| logger.info( | |
| f"Selected MPD video profile with resolution {profile['width']}x{profile['height']} for target {target_resolution}" | |
| ) | |
| return {profile_id: (profile, playlist_url)} | |
| def build_hls_playlist( | |
| mpd_dict: dict, profiles: list[dict], request: Request, skip_segments: list = None, start_offset: float = None | |
| ) -> str: | |
| """ | |
| Builds an HLS playlist from the MPD manifest for specific profiles. | |
| Args: | |
| mpd_dict (dict): The MPD manifest data. | |
| profiles (list[dict]): The profiles to include in the playlist. | |
| request (Request): The incoming HTTP request. | |
| skip_segments (list, optional): List of time segments to skip. Each item should have 'start' and 'end' keys. | |
| start_offset (float, optional): Start offset in seconds for live streams. Defaults to settings.livestream_start_offset for live. | |
| Returns: | |
| str: The HLS playlist as a string. | |
| """ | |
| hls = ["#EXTM3U", "#EXT-X-VERSION:6"] | |
| added_segments = 0 | |
| skipped_segments = 0 | |
| is_live = mpd_dict.get("isLive", False) | |
| # Inject EXT-X-START for live streams (enables prebuffering by starting behind live edge) | |
| # User-provided start_offset always takes precedence; otherwise use default for live streams only | |
| effective_start_offset = ( | |
| start_offset if start_offset is not None else (settings.livestream_start_offset if is_live else None) | |
| ) | |
| if effective_start_offset is not None: | |
| hls.append(f"#EXT-X-START:TIME-OFFSET={effective_start_offset:.1f},PRECISE=YES") | |
| # Initialize skip filter if skip_segments provided | |
| skip_filter = SkipSegmentFilter(skip_segments) if skip_segments else None | |
| # Use EXT-X-MAP for live streams to avoid duplicate moov atoms | |
| use_map = is_live | |
| proxy_url = request.url_for("segment_endpoint") | |
| proxy_url = str(proxy_url.replace(scheme=get_original_scheme(request))) | |
| # Get init endpoint URL for EXT-X-MAP | |
| init_proxy_url = request.url_for("init_endpoint") | |
| init_proxy_url = str(init_proxy_url.replace(scheme=get_original_scheme(request))) | |
| for index, profile in enumerate(profiles): | |
| segments = profile["segments"] | |
| if not segments: | |
| logger.warning(f"No segments found for profile {profile['id']}") | |
| continue | |
| if is_live: | |
| depth = max(settings.mpd_live_playlist_depth, 1) | |
| trimmed_segments = segments[-depth:] | |
| else: | |
| trimmed_segments = segments | |
| # Add headers for only the first profile | |
| if index == 0: | |
| first_segment = trimmed_segments[0] | |
| extinf_values = [f["extinf"] for f in trimmed_segments if "extinf" in f] | |
| target_duration = math.ceil(max(extinf_values)) if extinf_values else 3 | |
| # Align HLS media sequence with MPD-provided numbering when available | |
| mpd_start_number = profile.get("segment_template_start_number") | |
| sequence = first_segment.get("number") | |
| if sequence is None: | |
| # Fallback to MPD template start number | |
| if mpd_start_number is not None: | |
| sequence = mpd_start_number | |
| else: | |
| # As a last resort, derive from timeline information | |
| time_val = first_segment.get("time") | |
| duration_val = first_segment.get("duration_mpd_timescale") | |
| if time_val is not None and duration_val and duration_val > 0: | |
| sequence = math.floor(time_val / duration_val) | |
| else: | |
| sequence = 1 | |
| hls.extend( | |
| [ | |
| f"#EXT-X-TARGETDURATION:{target_duration}", | |
| f"#EXT-X-MEDIA-SEQUENCE:{sequence}", | |
| ] | |
| ) | |
| # For live streams, don't set PLAYLIST-TYPE to allow sliding window | |
| if not is_live: | |
| hls.append("#EXT-X-PLAYLIST-TYPE:VOD") | |
| init_url = profile["initUrl"] | |
| # For SegmentBase profiles, we may have byte range for initialization segment | |
| init_range = profile.get("initRange") | |
| query_params = dict(request.query_params) | |
| query_params.pop("profile_id", None) | |
| query_params.pop("d", None) | |
| has_encrypted = query_params.pop("has_encrypted", False) | |
| # Add EXT-X-MAP for init segment (for live streams or when beneficial) | |
| if use_map: | |
| init_query_params = { | |
| "init_url": init_url, | |
| "mime_type": profile["mimeType"], | |
| "is_live": "true" if is_live else "false", | |
| } | |
| if init_range: | |
| init_query_params["init_range"] = init_range | |
| # Add key parameters | |
| if query_params.get("key_id"): | |
| init_query_params["key_id"] = query_params["key_id"] | |
| if query_params.get("key"): | |
| init_query_params["key"] = query_params["key"] | |
| # Add api_password for authentication | |
| if query_params.get("api_password"): | |
| init_query_params["api_password"] = query_params["api_password"] | |
| init_map_url = encode_mediaflow_proxy_url( | |
| init_proxy_url, | |
| query_params=init_query_params, | |
| encryption_handler=encryption_handler if has_encrypted else None, | |
| ) | |
| hls.append(f'#EXT-X-MAP:URI="{init_map_url}"') | |
| need_discontinuity = False | |
| for segment in trimmed_segments: | |
| duration = segment["extinf"] | |
| # Check if this segment should be skipped | |
| if skip_filter: | |
| if skip_filter.should_skip_segment(duration): | |
| skip_filter.advance_time(duration) | |
| skipped_segments += 1 | |
| need_discontinuity = True | |
| continue | |
| skip_filter.advance_time(duration) | |
| # Add discontinuity marker after skipped segments | |
| if need_discontinuity: | |
| hls.append("#EXT-X-DISCONTINUITY") | |
| need_discontinuity = False | |
| program_date_time = segment.get("program_date_time") | |
| if program_date_time: | |
| hls.append(f"#EXT-X-PROGRAM-DATE-TIME:{program_date_time}") | |
| hls.append(f"#EXTINF:{duration:.3f},") | |
| segment_query_params = { | |
| "init_url": init_url, | |
| "segment_url": segment["media"], | |
| "mime_type": profile["mimeType"], | |
| "is_live": "true" if is_live else "false", | |
| } | |
| # Add use_map flag so segment endpoint knows not to include init | |
| if use_map: | |
| segment_query_params["use_map"] = "true" | |
| # Add byte range parameters for SegmentBase | |
| if init_range: | |
| segment_query_params["init_range"] = init_range | |
| # Segment may also have its own range (for SegmentBase) | |
| if "initRange" in segment: | |
| segment_query_params["init_range"] = segment["initRange"] | |
| query_params.update(segment_query_params) | |
| hls.append( | |
| encode_mediaflow_proxy_url( | |
| proxy_url, | |
| query_params=query_params, | |
| encryption_handler=encryption_handler if has_encrypted else None, | |
| ) | |
| ) | |
| added_segments += 1 | |
| if not mpd_dict["isLive"]: | |
| hls.append("#EXT-X-ENDLIST") | |
| if skip_filter and skipped_segments > 0: | |
| logger.info(f"Added {added_segments} segments to HLS playlist (skipped {skipped_segments} segments)") | |
| else: | |
| logger.info(f"Added {added_segments} segments to HLS playlist") | |
| return "\n".join(hls) | |