MdeMoLii/mediaflow-proxy

1import logging2import math3import time4 5from fastapi import Request, Response, HTTPException6 7from mediaflow_proxy.drm.decrypter import decrypt_segment8from mediaflow_proxy.utils.crypto_utils import encryption_handler9from mediaflow_proxy.utils.http_utils import encode_mediaflow_proxy_url, get_original_scheme, ProxyRequestHeaders10 11logger = logging.getLogger(__name__)12 13 14async def process_manifest(15    request: Request, mpd_dict: dict, proxy_headers: ProxyRequestHeaders, key_id: str = None, key: str = None16) -> Response:17    """18    Processes the MPD manifest and converts it to an HLS manifest.19 20    Args:21        request (Request): The incoming HTTP request.22        mpd_dict (dict): The MPD manifest data.23        proxy_headers (ProxyRequestHeaders): The headers to include in the request.24        key_id (str, optional): The DRM key ID. Defaults to None.25        key (str, optional): The DRM key. Defaults to None.26 27    Returns:28        Response: The HLS manifest as an HTTP response.29    """30    hls_content = build_hls(mpd_dict, request, key_id, key)31    return Response(content=hls_content, media_type="application/vnd.apple.mpegurl", headers=proxy_headers.response)32 33 34async def process_playlist(35    request: Request, mpd_dict: dict, profile_id: str, proxy_headers: ProxyRequestHeaders36) -> Response:37    """38    Processes the MPD manifest and converts it to an HLS playlist for a specific profile.39 40    Args:41        request (Request): The incoming HTTP request.42        mpd_dict (dict): The MPD manifest data.43        profile_id (str): The profile ID to generate the playlist for.44        proxy_headers (ProxyRequestHeaders): The headers to include in the request.45 46    Returns:47        Response: The HLS playlist as an HTTP response.48 49    Raises:50        HTTPException: If the profile is not found in the MPD manifest.51    """52    matching_profiles = [p for p in mpd_dict["profiles"] if p["id"] == profile_id]53    if not matching_profiles:54        raise HTTPException(status_code=404, detail="Profile not found")55 56    hls_content = build_hls_playlist(mpd_dict, matching_profiles, request)57    return Response(content=hls_content, media_type="application/vnd.apple.mpegurl", headers=proxy_headers.response)58 59 60async def process_segment(61    init_content: bytes,62    segment_content: bytes,63    mimetype: str,64    proxy_headers: ProxyRequestHeaders,65    key_id: str = None,66    key: str = None,67) -> Response:68    """69    Processes and decrypts a media segment.70 71    Args:72        init_content (bytes): The initialization segment content.73        segment_content (bytes): The media segment content.74        mimetype (str): The MIME type of the segment.75        proxy_headers (ProxyRequestHeaders): The headers to include in the request.76        key_id (str, optional): The DRM key ID. Defaults to None.77        key (str, optional): The DRM key. Defaults to None.78 79    Returns:80        Response: The decrypted segment as an HTTP response.81    """82    if key_id and key:83        # For DRM protected content84        now = time.time()85        decrypted_content = decrypt_segment(init_content, segment_content, key_id, key)86        logger.info(f"Decryption of {mimetype} segment took {time.time() - now:.4f} seconds")87    else:88        # For non-DRM protected content, we just concatenate init and segment content89        decrypted_content = init_content + segment_content90 91    return Response(content=decrypted_content, media_type=mimetype, headers=proxy_headers.response)92 93 94def build_hls(mpd_dict: dict, request: Request, key_id: str = None, key: str = None) -> str:95    """96    Builds an HLS manifest from the MPD manifest.97 98    Args:99        mpd_dict (dict): The MPD manifest data.100        request (Request): The incoming HTTP request.101        key_id (str, optional): The DRM key ID. Defaults to None.102        key (str, optional): The DRM key. Defaults to None.103 104    Returns:105        str: The HLS manifest as a string.106    """107    hls = ["#EXTM3U", "#EXT-X-VERSION:6"]108    query_params = dict(request.query_params)109    has_encrypted = query_params.pop("has_encrypted", False)110 111    video_profiles = {}112    audio_profiles = {}113 114    # Get the base URL for the playlist_endpoint endpoint115    proxy_url = request.url_for("playlist_endpoint")116    proxy_url = str(proxy_url.replace(scheme=get_original_scheme(request)))117 118    for profile in mpd_dict["profiles"]:119        query_params.update({"profile_id": profile["id"], "key_id": key_id or "", "key": key or ""})120        playlist_url = encode_mediaflow_proxy_url(121            proxy_url,122            query_params=query_params,123            encryption_handler=encryption_handler if has_encrypted else None,124        )125 126        if "video" in profile["mimeType"]:127            video_profiles[profile["id"]] = (profile, playlist_url)128        elif "audio" in profile["mimeType"]:129            audio_profiles[profile["id"]] = (profile, playlist_url)130 131    # Add audio streams132    for i, (profile, playlist_url) in enumerate(audio_profiles.values()):133        is_default = "YES" if i == 0 else "NO"  # Set the first audio track as default134        hls.append(135            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}"'136        )137 138    # Add video streams139    for profile, playlist_url in video_profiles.values():140        hls.append(141            f'#EXT-X-STREAM-INF:BANDWIDTH={profile["bandwidth"]},RESOLUTION={profile["width"]}x{profile["height"]},CODECS="{profile["codecs"]}",FRAME-RATE={profile["frameRate"]},AUDIO="audio"'142        )143        hls.append(playlist_url)144 145    return "\n".join(hls)146 147 148def build_hls_playlist(mpd_dict: dict, profiles: list[dict], request: Request) -> str:149    """150    Builds an HLS playlist from the MPD manifest for specific profiles.151 152    Args:153        mpd_dict (dict): The MPD manifest data.154        profiles (list[dict]): The profiles to include in the playlist.155        request (Request): The incoming HTTP request.156 157    Returns:158        str: The HLS playlist as a string.159    """160    hls = ["#EXTM3U", "#EXT-X-VERSION:6"]161 162    added_segments = 0163 164    proxy_url = request.url_for("segment_endpoint")165    proxy_url = str(proxy_url.replace(scheme=get_original_scheme(request)))166 167    for index, profile in enumerate(profiles):168        segments = profile["segments"]169        if not segments:170            logger.warning(f"No segments found for profile {profile['id']}")171            continue172 173        # Add headers for only the first profile174        if index == 0:175            sequence = segments[0]["number"]176            extinf_values = [f["extinf"] for f in segments if "extinf" in f]177            target_duration = math.ceil(max(extinf_values)) if extinf_values else 3178            hls.extend(179                [180                    f"#EXT-X-TARGETDURATION:{target_duration}",181                    f"#EXT-X-MEDIA-SEQUENCE:{sequence}",182                ]183            )184            if mpd_dict["isLive"]:185                hls.append("#EXT-X-PLAYLIST-TYPE:EVENT")186            else:187                hls.append("#EXT-X-PLAYLIST-TYPE:VOD")188 189        init_url = profile["initUrl"]190 191        query_params = dict(request.query_params)192        query_params.pop("profile_id", None)193        query_params.pop("d", None)194        has_encrypted = query_params.pop("has_encrypted", False)195 196        for segment in segments:197            hls.append(f'#EXTINF:{segment["extinf"]:.3f},')198            query_params.update(199                {"init_url": init_url, "segment_url": segment["media"], "mime_type": profile["mimeType"]}200            )201            hls.append(202                encode_mediaflow_proxy_url(203                    proxy_url,204                    query_params=query_params,205                    encryption_handler=encryption_handler if has_encrypted else None,206                )207            )208            added_segments += 1209 210    if not mpd_dict["isLive"]:211        hls.append("#EXT-X-ENDLIST")212 213    logger.info(f"Added {added_segments} segments to HLS playlist")214    return "\n".join(hls)215