From 564fec744d75f3ee80697646516410b44ad85ae1 Mon Sep 17 00:00:00 2001 From: Nirvana Date: Fri, 16 Jan 2026 14:26:31 +0100 Subject: [PATCH] Reformat --- .../base/utils/mp4_parser.py | 186 +++++++++++++++--- 1 file changed, 156 insertions(+), 30 deletions(-) diff --git a/lib/streaming_providers/base/utils/mp4_parser.py b/lib/streaming_providers/base/utils/mp4_parser.py index f9671cb..5bf6fc8 100644 --- a/lib/streaming_providers/base/utils/mp4_parser.py +++ b/lib/streaming_providers/base/utils/mp4_parser.py @@ -31,17 +31,20 @@ class MP4PSSHExtractor: @staticmethod def extract_from_bytes(data: bytes) -> List[PSSHData]: - """Extract PSSH boxes from MP4 binary data""" + """Extract PSSH boxes and encryption info from MP4 binary data""" pssh_data_list = [] offset = 0 + # First, extract all tenc boxes to get default KIDs + tenc_kids = MP4PSSHExtractor._extract_tenc_kids(data) + while offset < len(data): try: # Read box size (4 bytes, big-endian) if offset + 4 > len(data): break - box_size = struct.unpack(">I", data[offset : offset + 4])[0] + box_size = struct.unpack(">I", data[offset: offset + 4])[0] if box_size == 0: box_size = len(data) - offset # Box extends to end of file elif box_size == 1: @@ -52,18 +55,26 @@ class MP4PSSHExtractor: break # Read box type (4 bytes) - box_type = data[offset + 4 : offset + 8].decode("ascii", errors="ignore") + box_type = data[offset + 4: offset + 8].decode("ascii", errors="ignore") if box_type == "moov": # Look for PSSH in moov box - moov_data = data[offset : offset + box_size] + moov_data = data[offset: offset + box_size] pssh_in_moov = MP4PSSHExtractor._extract_from_moov(moov_data) + + # Enhance PSSH data with tenc KIDs if needed + for pssh in pssh_in_moov: + if not pssh.key_ids and tenc_kids: + pssh.key_ids = tenc_kids.copy() pssh_data_list.extend(pssh_in_moov) elif box_type == "pssh": # Found standalone PSSH box - pssh_box = MP4PSSHExtractor._parse_pssh_box(data[offset : offset + box_size]) + pssh_box = MP4PSSHExtractor._parse_pssh_box(data[offset: offset + box_size]) if pssh_box: + # Add tenc KIDs if PSSH doesn't have its own + if not pssh_box.key_ids and tenc_kids: + pssh_box.key_ids = tenc_kids.copy() pssh_data_list.append(pssh_box) offset += box_size @@ -74,29 +85,109 @@ class MP4PSSHExtractor: return pssh_data_list + @staticmethod + def _extract_tenc_kids(data: bytes) -> List[str]: + """Extract default KIDs from all tenc boxes in the MP4 data""" + kids = [] + offset = 0 + + while offset < len(data): + try: + if offset + 8 > len(data): + break + + box_size = struct.unpack(">I", data[offset:offset + 4])[0] + if box_size < 8: + offset += 1 + continue + + box_type = data[offset + 4:offset + 8] + + if box_type == b'tenc': + # Parse tenc box + tenc_kid = MP4PSSHExtractor._parse_tenc_box(data[offset:offset + box_size]) + if tenc_kid: + kids.append(tenc_kid) + + offset += box_size + except: + offset += 1 + + return kids + + @staticmethod + def _parse_tenc_box(tenc_bytes: bytes) -> Optional[str]: + """Parse a tenc box and extract the default key ID""" + try: + if len(tenc_bytes) < 32: + return None + + # tenc box structure: + # 0-3: box size (4 bytes) + # 4-7: box type 'tenc' (4 bytes) + # 8: version (1 byte) + # 9-11: flags (3 bytes) + # 12: reserved (24 bits) + is_encrypted (1 bit) + # 13: default_iv_size (1 byte) + # 14-29: default_KID (16 bytes) + + # Check if box is long enough + if len(tenc_bytes) < 30: + return None + + # Skip to default_KID (offset 14 from start of box) + kid_start = 14 + + if kid_start + 16 > len(tenc_bytes): + return None + + kid_bytes = tenc_bytes[kid_start:kid_start + 16] + kid_uuid = str(uuid.UUID(bytes=kid_bytes)) + + # Convert to lowercase without dashes (same format as PSSH KIDs) + clean_kid = kid_uuid.replace("-", "").lower() + + logger.debug(f"Found KID from tenc box: {clean_kid}") + return clean_kid + + except Exception as e: + logger.debug(f"Failed to parse tenc box: {e}") + return None + @staticmethod def _extract_from_moov(moov_data: bytes) -> List[PSSHData]: """Extract PSSH boxes from moov container""" pssh_list = [] offset = 8 # Skip moov header + # First extract tenc KIDs from this moov + tenc_kids = MP4PSSHExtractor._extract_tenc_kids(moov_data) + while offset < len(moov_data): try: - box_size = struct.unpack(">I", moov_data[offset : offset + 4])[0] - box_type = moov_data[offset + 4 : offset + 8].decode("ascii", errors="ignore") + box_size = struct.unpack(">I", moov_data[offset: offset + 4])[0] + box_type = moov_data[offset + 4: offset + 8].decode("ascii", errors="ignore") if box_type == "trak": # Parse track for PSSH - trak_data = moov_data[offset : offset + box_size] + trak_data = moov_data[offset: offset + box_size] pssh_in_trak = MP4PSSHExtractor._extract_from_trak(trak_data) + + # Add tenc KIDs to PSSH boxes if needed + for pssh in pssh_in_trak: + if not pssh.key_ids and tenc_kids: + pssh.key_ids = tenc_kids.copy() pssh_list.extend(pssh_in_trak) elif box_type == "pssh": # PSSH directly in moov pssh_box = MP4PSSHExtractor._parse_pssh_box( - moov_data[offset : offset + box_size] + moov_data[offset: offset + box_size] ) if pssh_box: + # Add tenc KIDs if PSSH doesn't have its own + if not pssh_box.key_ids and tenc_kids: + pssh_box.key_ids = tenc_kids.copy() pssh_list.append(pssh_box) offset += box_size @@ -114,11 +205,11 @@ class MP4PSSHExtractor: while offset < len(trak_data): try: - box_size = struct.unpack(">I", trak_data[offset : offset + 4])[0] - box_type = trak_data[offset + 4 : offset + 8].decode("ascii", errors="ignore") + box_size = struct.unpack(">I", trak_data[offset: offset + 4])[0] + box_type = trak_data[offset + 4: offset + 8].decode("ascii", errors="ignore") if box_type == "mdia": - mdia_data = trak_data[offset : offset + box_size] + mdia_data = trak_data[offset: offset + box_size] pssh_in_mdia = MP4PSSHExtractor._extract_from_mdia(mdia_data) pssh_list.extend(pssh_in_mdia) @@ -137,11 +228,11 @@ class MP4PSSHExtractor: while offset < len(mdia_data): try: - box_size = struct.unpack(">I", mdia_data[offset : offset + 4])[0] - box_type = mdia_data[offset + 4 : offset + 8].decode("ascii", errors="ignore") + box_size = struct.unpack(">I", mdia_data[offset: offset + 4])[0] + box_type = mdia_data[offset + 4: offset + 8].decode("ascii", errors="ignore") if box_type == "minf": - minf_data = mdia_data[offset : offset + box_size] + minf_data = mdia_data[offset: offset + box_size] pssh_in_minf = MP4PSSHExtractor._extract_from_minf(minf_data) pssh_list.extend(pssh_in_minf) @@ -160,11 +251,11 @@ class MP4PSSHExtractor: while offset < len(minf_data): try: - box_size = struct.unpack(">I", minf_data[offset : offset + 4])[0] - box_type = minf_data[offset + 4 : offset + 8].decode("ascii", errors="ignore") + box_size = struct.unpack(">I", minf_data[offset: offset + 4])[0] + box_type = minf_data[offset + 4: offset + 8].decode("ascii", errors="ignore") if box_type == "stbl": - stbl_data = minf_data[offset : offset + box_size] + stbl_data = minf_data[offset: offset + box_size] pssh_in_stbl = MP4PSSHExtractor._extract_from_stbl(stbl_data) pssh_list.extend(pssh_in_stbl) @@ -183,14 +274,25 @@ class MP4PSSHExtractor: while offset < len(stbl_data): try: - box_size = struct.unpack(">I", stbl_data[offset : offset + 4])[0] - box_type = stbl_data[offset + 4 : offset + 8].decode("ascii", errors="ignore") + box_size = struct.unpack(">I", stbl_data[offset: offset + 4])[0] + box_type = stbl_data[offset + 4: offset + 8].decode("ascii", errors="ignore") if box_type == "sinf": - sinf_data = stbl_data[offset : offset + box_size] + sinf_data = stbl_data[offset: offset + box_size] pssh_in_sinf = MP4PSSHExtractor._extract_from_sinf(sinf_data) pssh_list.extend(pssh_in_sinf) + elif box_type == "tenc": + # Direct tenc box in stbl (unusual but possible) + tenc_kid = MP4PSSHExtractor._parse_tenc_box( + stbl_data[offset:offset + box_size] + ) + if tenc_kid: + # If we have existing PSSH boxes without KIDs, add this KID + for pssh in pssh_list: + if not pssh.key_ids: + pssh.key_ids.append(tenc_kid) + offset += box_size except: @@ -206,14 +308,25 @@ class MP4PSSHExtractor: while offset < len(sinf_data): try: - box_size = struct.unpack(">I", sinf_data[offset : offset + 4])[0] - box_type = sinf_data[offset + 4 : offset + 8].decode("ascii", errors="ignore") + box_size = struct.unpack(">I", sinf_data[offset: offset + 4])[0] + box_type = sinf_data[offset + 4: offset + 8].decode("ascii", errors="ignore") if box_type == "schi": - schi_data = sinf_data[offset : offset + box_size] + schi_data = sinf_data[offset: offset + box_size] pssh_in_schi = MP4PSSHExtractor._extract_from_schi(schi_data) pssh_list.extend(pssh_in_schi) + elif box_type == "tenc": + # tenc box directly in sinf + tenc_kid = MP4PSSHExtractor._parse_tenc_box( + sinf_data[offset:offset + box_size] + ) + if tenc_kid: + # If we have existing PSSH boxes without KIDs, add this KID + for pssh in pssh_list: + if not pssh.key_ids: + pssh.key_ids.append(tenc_kid) + offset += box_size except: @@ -229,16 +342,27 @@ class MP4PSSHExtractor: while offset < len(schi_data): try: - box_size = struct.unpack(">I", schi_data[offset : offset + 4])[0] - box_type = schi_data[offset + 4 : offset + 8].decode("ascii", errors="ignore") + box_size = struct.unpack(">I", schi_data[offset: offset + 4])[0] + box_type = schi_data[offset + 4: offset + 8].decode("ascii", errors="ignore") if box_type == "pssh": pssh_box = MP4PSSHExtractor._parse_pssh_box( - schi_data[offset : offset + box_size] + schi_data[offset: offset + box_size] ) if pssh_box: pssh_list.append(pssh_box) + elif box_type == "tenc": + # tenc box in schi (common structure) + tenc_kid = MP4PSSHExtractor._parse_tenc_box( + schi_data[offset:offset + box_size] + ) + if tenc_kid: + # If we have existing PSSH boxes without KIDs, add this KID + for pssh in pssh_list: + if not pssh.key_ids: + pssh.key_ids.append(tenc_kid) + offset += box_size except: @@ -276,7 +400,7 @@ class MP4PSSHExtractor: if current_offset + 4 > len(pssh_bytes): return None - kid_count = struct.unpack(">I", pssh_bytes[current_offset : current_offset + 4])[0] + kid_count = struct.unpack(">I", pssh_bytes[current_offset: current_offset + 4])[0] current_offset += 4 # Read each KID @@ -284,10 +408,12 @@ class MP4PSSHExtractor: if current_offset + 16 > len(pssh_bytes): break - kid_bytes = pssh_bytes[current_offset : current_offset + 16] + kid_bytes = pssh_bytes[current_offset: current_offset + 16] kid_uuid = str(uuid.UUID(bytes=kid_bytes)) key_ids.append(kid_uuid.replace("-", "").lower()) current_offset += 16 + else: + logger.debug(f"Version 0 PSSH for system {system_id}, will look for KIDs in tenc box") # Encode entire PSSH box as base64 pssh_b64 = base64.b64encode(pssh_bytes[:box_size]).decode("ascii") @@ -301,4 +427,4 @@ class MP4PSSHExtractor: except Exception as e: logger.debug(f"Failed to parse PSSH box: {e}") - return None + return None \ No newline at end of file