acdc137726/mediaflow-proxy

1import logging2import math3import re4from datetime import datetime, timedelta, timezone5from typing import List, Dict6from urllib.parse import urljoin7 8import xmltodict9 10logger = logging.getLogger(__name__)11 12 13def parse_mpd(mpd_content: str | bytes) -> dict:14    """15    Parses the MPD content into a dictionary.16 17    Args:18        mpd_content (str | bytes): The MPD content to parse.19 20    Returns:21        dict: The parsed MPD content as a dictionary.22    """23    return xmltodict.parse(mpd_content)24 25 26def parse_mpd_dict(27    mpd_dict: dict, mpd_url: str, parse_drm: bool = True, parse_segment_profile_id: str | None = None28) -> dict:29    """30    Parses the MPD dictionary and extracts relevant information.31 32    Args:33        mpd_dict (dict): The MPD content as a dictionary.34        mpd_url (str): The URL of the MPD manifest.35        parse_drm (bool, optional): Whether to parse DRM information. Defaults to True.36        parse_segment_profile_id (str, optional): The profile ID to parse segments for. Defaults to None.37 38    Returns:39        dict: The parsed MPD information including profiles and DRM info.40 41    This function processes the MPD dictionary to extract profiles, DRM information, and other relevant data.42    It handles both live and static MPD manifests.43    """44    profiles = []45    parsed_dict = {}46    source = "/".join(mpd_url.split("/")[:-1])47 48    is_live = mpd_dict["MPD"].get("@type", "static").lower() == "dynamic"49    parsed_dict["isLive"] = is_live50 51    media_presentation_duration = mpd_dict["MPD"].get("@mediaPresentationDuration")52 53    # Parse additional MPD attributes for live streams54    if is_live:55        parsed_dict["minimumUpdatePeriod"] = parse_duration(mpd_dict["MPD"].get("@minimumUpdatePeriod", "PT0S"))56        parsed_dict["timeShiftBufferDepth"] = parse_duration(mpd_dict["MPD"].get("@timeShiftBufferDepth", "PT2M"))57        parsed_dict["availabilityStartTime"] = datetime.fromisoformat(58            mpd_dict["MPD"]["@availabilityStartTime"].replace("Z", "+00:00")59        )60        parsed_dict["publishTime"] = datetime.fromisoformat(61            mpd_dict["MPD"].get("@publishTime", "").replace("Z", "+00:00")62        )63 64    periods = mpd_dict["MPD"]["Period"]65    periods = periods if isinstance(periods, list) else [periods]66 67    for period in periods:68        parsed_dict["PeriodStart"] = parse_duration(period.get("@start", "PT0S"))69        for adaptation in period["AdaptationSet"]:70            representations = adaptation["Representation"]71            representations = representations if isinstance(representations, list) else [representations]72 73            for representation in representations:74                profile = parse_representation(75                    parsed_dict,76                    representation,77                    adaptation,78                    source,79                    media_presentation_duration,80                    parse_segment_profile_id,81                )82                if profile:83                    profiles.append(profile)84    parsed_dict["profiles"] = profiles85 86    if parse_drm:87        drm_info = extract_drm_info(periods, mpd_url)88    else:89        drm_info = {}90    parsed_dict["drmInfo"] = drm_info91 92    return parsed_dict93 94 95def pad_base64(encoded_key_id):96    """97    Pads a base64 encoded key ID to make its length a multiple of 4.98 99    Args:100        encoded_key_id (str): The base64 encoded key ID.101 102    Returns:103        str: The padded base64 encoded key ID.104    """105    return encoded_key_id + "=" * (4 - len(encoded_key_id) % 4)106 107 108def extract_drm_info(periods: List[Dict], mpd_url: str) -> Dict:109    """110    Extracts DRM information from the MPD periods.111 112    Args:113        periods (List[Dict]): The list of periods in the MPD.114        mpd_url (str): The URL of the MPD manifest.115 116    Returns:117        Dict: The extracted DRM information.118 119    This function processes the ContentProtection elements in the MPD to extract DRM system information,120    such as ClearKey, Widevine, and PlayReady.121    """122    drm_info = {"isDrmProtected": False}123 124    for period in periods:125        adaptation_sets: list[dict] | dict = period.get("AdaptationSet", [])126        if not isinstance(adaptation_sets, list):127            adaptation_sets = [adaptation_sets]128 129        for adaptation_set in adaptation_sets:130            # Check ContentProtection in AdaptationSet131            process_content_protection(adaptation_set.get("ContentProtection", []), drm_info)132 133            # Check ContentProtection inside each Representation134            representations: list[dict] | dict = adaptation_set.get("Representation", [])135            if not isinstance(representations, list):136                representations = [representations]137 138            for representation in representations:139                process_content_protection(representation.get("ContentProtection", []), drm_info)140 141    # If we have a license acquisition URL, make sure it's absolute142    if "laUrl" in drm_info and not drm_info["laUrl"].startswith(("http://", "https://")):143        drm_info["laUrl"] = urljoin(mpd_url, drm_info["laUrl"])144 145    return drm_info146 147 148def process_content_protection(content_protection: list[dict] | dict, drm_info: dict):149    """150    Processes the ContentProtection elements to extract DRM information.151 152    Args:153        content_protection (list[dict] | dict): The ContentProtection elements.154        drm_info (dict): The dictionary to store DRM information.155 156    This function updates the drm_info dictionary with DRM system information found in the ContentProtection elements.157    """158    if not isinstance(content_protection, list):159        content_protection = [content_protection]160 161    for protection in content_protection:162        drm_info["isDrmProtected"] = True163        scheme_id_uri = protection.get("@schemeIdUri", "").lower()164 165        if "clearkey" in scheme_id_uri:166            drm_info["drmSystem"] = "clearkey"167            if "clearkey:Laurl" in protection:168                la_url = protection["clearkey:Laurl"].get("#text")169                if la_url and "laUrl" not in drm_info:170                    drm_info["laUrl"] = la_url171 172        elif "widevine" in scheme_id_uri or "edef8ba9-79d6-4ace-a3c8-27dcd51d21ed" in scheme_id_uri:173            drm_info["drmSystem"] = "widevine"174            pssh = protection.get("cenc:pssh", {}).get("#text")175            if pssh:176                drm_info["pssh"] = pssh177 178        elif "playready" in scheme_id_uri or "9a04f079-9840-4286-ab92-e65be0885f95" in scheme_id_uri:179            drm_info["drmSystem"] = "playready"180 181        if "@cenc:default_KID" in protection:182            key_id = protection["@cenc:default_KID"].replace("-", "")183            if "keyId" not in drm_info:184                drm_info["keyId"] = key_id185 186        if "ms:laurl" in protection:187            la_url = protection["ms:laurl"].get("@licenseUrl")188            if la_url and "laUrl" not in drm_info:189                drm_info["laUrl"] = la_url190 191    return drm_info192 193 194def parse_representation(195    parsed_dict: dict,196    representation: dict,197    adaptation: dict,198    source: str,199    media_presentation_duration: str,200    parse_segment_profile_id: str | None,201) -> dict | None:202    """203    Parses a representation and extracts profile information.204 205    Args:206        parsed_dict (dict): The parsed MPD data.207        representation (dict): The representation data.208        adaptation (dict): The adaptation set data.209        source (str): The source URL.210        media_presentation_duration (str): The media presentation duration.211        parse_segment_profile_id (str, optional): The profile ID to parse segments for. Defaults to None.212 213    Returns:214        dict | None: The parsed profile information or None if not applicable.215    """216    mime_type = _get_key(adaptation, representation, "@mimeType") or (217        "video/mp4" if "avc" in representation["@codecs"] else "audio/mp4"218    )219    if "video" not in mime_type and "audio" not in mime_type:220        return None221 222    profile = {223        "id": representation.get("@id") or adaptation.get("@id"),224        "mimeType": mime_type,225        "lang": representation.get("@lang") or adaptation.get("@lang"),226        "codecs": representation.get("@codecs") or adaptation.get("@codecs"),227        "bandwidth": int(representation.get("@bandwidth") or adaptation.get("@bandwidth")),228        "startWithSAP": (_get_key(adaptation, representation, "@startWithSAP") or "1") == "1",229        "mediaPresentationDuration": media_presentation_duration,230    }231 232    if "audio" in profile["mimeType"]:233        profile["audioSamplingRate"] = representation.get("@audioSamplingRate") or adaptation.get("@audioSamplingRate")234        profile["channels"] = representation.get("AudioChannelConfiguration", {}).get("@value", "2")235    else:236        profile["width"] = int(representation["@width"])237        profile["height"] = int(representation["@height"])238        frame_rate = representation.get("@frameRate") or adaptation.get("@maxFrameRate") or "30000/1001"239        frame_rate = frame_rate if "/" in frame_rate else f"{frame_rate}/1"240        profile["frameRate"] = round(int(frame_rate.split("/")[0]) / int(frame_rate.split("/")[1]), 3)241        profile["sar"] = representation.get("@sar", "1:1")242 243    if parse_segment_profile_id is None or profile["id"] != parse_segment_profile_id:244        return profile245 246    item = adaptation.get("SegmentTemplate") or representation.get("SegmentTemplate")247    if item:248        profile["segments"] = parse_segment_template(parsed_dict, item, profile, source)249    else:250        profile["segments"] = parse_segment_base(representation, source)251 252    return profile253 254 255def _get_key(adaptation: dict, representation: dict, key: str) -> str | None:256    """257    Retrieves a key from the representation or adaptation set.258 259    Args:260        adaptation (dict): The adaptation set data.261        representation (dict): The representation data.262        key (str): The key to retrieve.263 264    Returns:265        str | None: The value of the key or None if not found.266    """267    return representation.get(key, adaptation.get(key, None))268 269 270def parse_segment_template(parsed_dict: dict, item: dict, profile: dict, source: str) -> List[Dict]:271    """272    Parses a segment template and extracts segment information.273 274    Args:275        parsed_dict (dict): The parsed MPD data.276        item (dict): The segment template data.277        profile (dict): The profile information.278        source (str): The source URL.279 280    Returns:281        List[Dict]: The list of parsed segments.282    """283    segments = []284    timescale = int(item.get("@timescale", 1))285 286    # Initialization287    if "@initialization" in item:288        media = item["@initialization"]289        media = media.replace("$RepresentationID$", profile["id"])290        media = media.replace("$Bandwidth$", str(profile["bandwidth"]))291        if not media.startswith("http"):292            media = f"{source}/{media}"293        profile["initUrl"] = media294 295    # Segments296    if "SegmentTimeline" in item:297        segments.extend(parse_segment_timeline(parsed_dict, item, profile, source, timescale))298    elif "@duration" in item:299        segments.extend(parse_segment_duration(parsed_dict, item, profile, source, timescale))300 301    return segments302 303 304def parse_segment_timeline(parsed_dict: dict, item: dict, profile: dict, source: str, timescale: int) -> List[Dict]:305    """306    Parses a segment timeline and extracts segment information.307 308    Args:309        parsed_dict (dict): The parsed MPD data.310        item (dict): The segment timeline data.311        profile (dict): The profile information.312        source (str): The source URL.313        timescale (int): The timescale for the segments.314 315    Returns:316        List[Dict]: The list of parsed segments.317    """318    timelines = item["SegmentTimeline"]["S"]319    timelines = timelines if isinstance(timelines, list) else [timelines]320    period_start = parsed_dict["availabilityStartTime"] + timedelta(seconds=parsed_dict.get("PeriodStart", 0))321    presentation_time_offset = int(item.get("@presentationTimeOffset", 0))322    start_number = int(item.get("@startNumber", 1))323 324    segments = [325        create_segment_data(timeline, item, profile, source, timescale)326        for timeline in preprocess_timeline(timelines, start_number, period_start, presentation_time_offset, timescale)327    ]328    return segments329 330 331def preprocess_timeline(332    timelines: List[Dict], start_number: int, period_start: datetime, presentation_time_offset: int, timescale: int333) -> List[Dict]:334    """335    Preprocesses the segment timeline data.336 337    Args:338        timelines (List[Dict]): The list of timeline segments.339        start_number (int): The starting segment number.340        period_start (datetime): The start time of the period.341        presentation_time_offset (int): The presentation time offset.342        timescale (int): The timescale for the segments.343 344    Returns:345        List[Dict]: The list of preprocessed timeline segments.346    """347    processed_data = []348    current_time = 0349    for timeline in timelines:350        repeat = int(timeline.get("@r", 0))351        duration = int(timeline["@d"])352        start_time = int(timeline.get("@t", current_time))353 354        for _ in range(repeat + 1):355            segment_start_time = period_start + timedelta(seconds=(start_time - presentation_time_offset) / timescale)356            segment_end_time = segment_start_time + timedelta(seconds=duration / timescale)357            processed_data.append(358                {359                    "number": start_number,360                    "start_time": segment_start_time,361                    "end_time": segment_end_time,362                    "duration": duration,363                    "time": start_time,364                }365            )366            start_time += duration367            start_number += 1368 369        current_time = start_time370 371    return processed_data372 373 374def parse_segment_duration(parsed_dict: dict, item: dict, profile: dict, source: str, timescale: int) -> List[Dict]:375    """376    Parses segment duration and extracts segment information.377    This is used for static or live MPD manifests.378 379    Args:380        parsed_dict (dict): The parsed MPD data.381        item (dict): The segment duration data.382        profile (dict): The profile information.383        source (str): The source URL.384        timescale (int): The timescale for the segments.385 386    Returns:387        List[Dict]: The list of parsed segments.388    """389    duration = int(item["@duration"])390    start_number = int(item.get("@startNumber", 1))391    segment_duration_sec = duration / timescale392 393    if parsed_dict["isLive"]:394        segments = generate_live_segments(parsed_dict, segment_duration_sec, start_number)395    else:396        segments = generate_vod_segments(profile, duration, timescale, start_number)397 398    return [create_segment_data(seg, item, profile, source, timescale) for seg in segments]399 400 401def generate_live_segments(parsed_dict: dict, segment_duration_sec: float, start_number: int) -> List[Dict]:402    """403    Generates live segments based on the segment duration and start number.404    This is used for live MPD manifests.405 406    Args:407        parsed_dict (dict): The parsed MPD data.408        segment_duration_sec (float): The segment duration in seconds.409        start_number (int): The starting segment number.410 411    Returns:412        List[Dict]: The list of generated live segments.413    """414    time_shift_buffer_depth = timedelta(seconds=parsed_dict.get("timeShiftBufferDepth", 60))415    segment_count = math.ceil(time_shift_buffer_depth.total_seconds() / segment_duration_sec)416    current_time = datetime.now(tz=timezone.utc)417    earliest_segment_number = max(418        start_number419        + math.floor((current_time - parsed_dict["availabilityStartTime"]).total_seconds() / segment_duration_sec)420        - segment_count,421        start_number,422    )423 424    return [425        {426            "number": number,427            "start_time": parsed_dict["availabilityStartTime"]428            + timedelta(seconds=(number - start_number) * segment_duration_sec),429            "duration": segment_duration_sec,430        }431        for number in range(earliest_segment_number, earliest_segment_number + segment_count)432    ]433 434 435def generate_vod_segments(profile: dict, duration: int, timescale: int, start_number: int) -> List[Dict]:436    """437    Generates VOD segments based on the segment duration and start number.438    This is used for static MPD manifests.439 440    Args:441        profile (dict): The profile information.442        duration (int): The segment duration.443        timescale (int): The timescale for the segments.444        start_number (int): The starting segment number.445 446    Returns:447        List[Dict]: The list of generated VOD segments.448    """449    total_duration = profile.get("mediaPresentationDuration") or 0450    if isinstance(total_duration, str):451        total_duration = parse_duration(total_duration)452    segment_count = math.ceil(total_duration * timescale / duration)453 454    return [{"number": start_number + i, "duration": duration / timescale} for i in range(segment_count)]455 456 457def create_segment_data(segment: Dict, item: dict, profile: dict, source: str, timescale: int | None = None) -> Dict:458    """459    Creates segment data based on the segment information. This includes the segment URL and metadata.460 461    Args:462        segment (Dict): The segment information.463        item (dict): The segment template data.464        profile (dict): The profile information.465        source (str): The source URL.466        timescale (int, optional): The timescale for the segments. Defaults to None.467 468    Returns:469        Dict: The created segment data.470    """471    media_template = item["@media"]472    media = media_template.replace("$RepresentationID$", profile["id"])473    media = media.replace("$Number%04d$", f"{segment['number']:04d}")474    media = media.replace("$Number$", str(segment["number"]))475    media = media.replace("$Bandwidth$", str(profile["bandwidth"]))476 477    if "time" in segment and timescale is not None:478        media = media.replace("$Time$", str(int(segment["time"] * timescale)))479 480    if not media.startswith("http"):481        media = f"{source}/{media}"482 483    segment_data = {484        "type": "segment",485        "media": media,486        "number": segment["number"],487    }488 489    if "start_time" in segment and "end_time" in segment:490        segment_data.update(491            {492                "start_time": segment["start_time"],493                "end_time": segment["end_time"],494                "extinf": (segment["end_time"] - segment["start_time"]).total_seconds(),495                "program_date_time": segment["start_time"].isoformat() + "Z",496            }497        )498    elif "start_time" in segment and "duration" in segment:499        duration = segment["duration"]500        segment_data.update(501            {502                "start_time": segment["start_time"],503                "end_time": segment["start_time"] + timedelta(seconds=duration),504                "extinf": duration,505                "program_date_time": segment["start_time"].isoformat() + "Z",506            }507        )508    elif "duration" in segment:509        segment_data["extinf"] = segment["duration"]510 511    return segment_data512 513 514def parse_segment_base(representation: dict, source: str) -> List[Dict]:515    """516    Parses segment base information and extracts segment data. This is used for single-segment representations.517 518    Args:519        representation (dict): The representation data.520        source (str): The source URL.521 522    Returns:523        List[Dict]: The list of parsed segments.524    """525    segment = representation["SegmentBase"]526    start, end = map(int, segment["@indexRange"].split("-"))527    if "Initialization" in segment:528        start, _ = map(int, segment["Initialization"]["@range"].split("-"))529 530    return [531        {532            "type": "segment",533            "range": f"{start}-{end}",534            "media": f"{source}/{representation['BaseURL']}",535        }536    ]537 538 539def parse_duration(duration_str: str) -> float:540    """541    Parses a duration ISO 8601 string into seconds.542 543    Args:544        duration_str (str): The duration string to parse.545 546    Returns:547        float: The parsed duration in seconds.548    """549    pattern = re.compile(r"P(?:(\d+)Y)?(?:(\d+)M)?(?:(\d+)D)?T?(?:(\d+)H)?(?:(\d+)M)?(?:(\d+(?:\.\d+)?)S)?")550    match = pattern.match(duration_str)551    if not match:552        raise ValueError(f"Invalid duration format: {duration_str}")553 554    years, months, days, hours, minutes, seconds = [float(g) if g else 0 for g in match.groups()]555    return years * 365 * 24 * 3600 + months * 30 * 24 * 3600 + days * 24 * 3600 + hours * 3600 + minutes * 60 + seconds556