mirror of
https://github.com/nirvana-7777/script.service.ultimate.git
synced 2026-09-16 06:02:35 +02:00
484 lines
17 KiB
Python
484 lines
17 KiB
Python
import base64
|
|
import struct
|
|
import uuid
|
|
from typing import List, Optional
|
|
|
|
from ..models.drm_models import PSSHData
|
|
from .logger import logger
|
|
|
|
|
|
class MP4PSSHExtractor:
|
|
"""Extract PSSH boxes and key IDs from MP4 segments"""
|
|
|
|
@staticmethod
|
|
def extract_from_url(segment_url: str, timeout: int = 10) -> List[PSSHData]:
|
|
"""Download MP4 segment and extract PSSH data"""
|
|
import requests
|
|
|
|
try:
|
|
response = requests.get(segment_url, timeout=timeout)
|
|
response.raise_for_status()
|
|
|
|
# Only download first ~100KB for efficiency
|
|
chunk_size = 1024 * 100
|
|
data = response.content[:chunk_size]
|
|
|
|
return MP4PSSHExtractor.extract_from_bytes(data)
|
|
|
|
except Exception as e:
|
|
logger.error(f"Failed to extract PSSH from {segment_url}: {e}")
|
|
return []
|
|
|
|
@staticmethod
|
|
def extract_from_bytes(data: bytes) -> List[PSSHData]:
|
|
"""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_all_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]
|
|
if box_size == 0:
|
|
box_size = len(data) - offset # Box extends to end of file
|
|
elif box_size == 1:
|
|
# Extended size (skip for now - rare in practice)
|
|
break
|
|
|
|
if offset + box_size > len(data):
|
|
break
|
|
|
|
# Read box type (4 bytes)
|
|
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]
|
|
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]
|
|
)
|
|
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
|
|
|
|
except Exception as e:
|
|
logger.debug(f"Error parsing MP4 box at offset {offset}: {e}")
|
|
offset += 1 # Try to recover
|
|
|
|
return pssh_data_list
|
|
|
|
@staticmethod
|
|
def _extract_all_tenc_kids(data: bytes) -> List[str]:
|
|
"""Recursively extract all KIDs from 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 or offset + box_size > len(data):
|
|
offset += 1
|
|
continue
|
|
|
|
box_type = data[offset + 4 : offset + 8]
|
|
|
|
if box_type == b"tenc":
|
|
# Parse tenc box
|
|
logger.debug(
|
|
f"Found tenc box at offset {offset} (0x{offset:x}), size {box_size}"
|
|
)
|
|
tenc_kid = MP4PSSHExtractor._parse_tenc_box(
|
|
data[offset : offset + box_size]
|
|
)
|
|
if tenc_kid and tenc_kid not in kids:
|
|
kids.append(tenc_kid)
|
|
elif box_size > 8:
|
|
# Recursively search inside container boxes
|
|
container_boxes = {
|
|
b"moov",
|
|
b"trak",
|
|
b"mdia",
|
|
b"minf",
|
|
b"stbl",
|
|
b"stsd",
|
|
b"encv",
|
|
b"sinf",
|
|
b"schi",
|
|
b"udta",
|
|
}
|
|
if box_type in container_boxes:
|
|
# Recursively search inside the container
|
|
inner_data = data[offset + 8 : offset + box_size]
|
|
inner_kids = MP4PSSHExtractor._extract_all_tenc_kids(inner_data)
|
|
for kid in inner_kids:
|
|
if kid not in kids:
|
|
kids.append(kid)
|
|
|
|
offset += box_size
|
|
except Exception as e:
|
|
logger.debug(f"Error parsing box at offset {offset}: {e}")
|
|
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:
|
|
# tenc box should be at least 32 bytes for version 0 with KID
|
|
if len(tenc_bytes) < 32:
|
|
logger.debug(
|
|
f"tenc box too small: {len(tenc_bytes)} bytes, need at least 32"
|
|
)
|
|
return None
|
|
|
|
# Parse box header
|
|
box_size = struct.unpack(">I", tenc_bytes[0:4])[0]
|
|
box_type = tenc_bytes[4:8]
|
|
|
|
if box_type != b"tenc":
|
|
logger.debug(f"Not a tenc box, type: {box_type.hex()}")
|
|
return None
|
|
|
|
version = tenc_bytes[8]
|
|
logger.debug(f"tenc box version: {version}")
|
|
|
|
if version == 0:
|
|
# For version 0: reserved(24) + is_encrypted(1) is a 32-bit field at bytes 12-15
|
|
# Actually, it's stored as a 32-bit big-endian integer
|
|
# Bytes 12-15: reserved (24 bits) + is_encrypted (1 bit)
|
|
# Byte 16: default_iv_size
|
|
# Bytes 17-32: default_KID (16 bytes)
|
|
|
|
# Read the 32-bit field at bytes 12-15
|
|
reserved_encrypted = struct.unpack(">I", tenc_bytes[12:16])[0]
|
|
|
|
# Extract is_encrypted from bit 24 (0-indexed from MSB)
|
|
# Actually, in big-endian, byte 14 is the 3rd byte of this 32-bit field
|
|
# Let me think: bytes[12], bytes[13], bytes[14], bytes[15]
|
|
# Byte 14 = tenc_bytes[14] = your "01"
|
|
# Byte 15 = tenc_bytes[15] = your "08"
|
|
|
|
# Actually simpler: byte 14 contains is_encrypted in its LSB
|
|
is_encrypted = (tenc_bytes[14] & 0x01) != 0
|
|
default_iv_size = tenc_bytes[15]
|
|
|
|
# KID starts at byte 16 (not 14!)
|
|
kid_offset = 16
|
|
|
|
logger.debug(
|
|
f"tenc v0: is_encrypted={is_encrypted}, default_iv_size={default_iv_size}, kid_offset={kid_offset}"
|
|
)
|
|
|
|
elif version == 1:
|
|
# Version 1 has different structure
|
|
kid_offset = 16 # Adjust as needed for v1
|
|
logger.debug(f"tenc v1: kid_offset={kid_offset}")
|
|
else:
|
|
logger.debug(f"Unknown tenc version: {version}")
|
|
return None
|
|
|
|
if kid_offset + 16 > len(tenc_bytes):
|
|
logger.debug(
|
|
f"tenc box too small for KID: {len(tenc_bytes)} bytes, need {kid_offset + 16}"
|
|
)
|
|
return None
|
|
|
|
# Extract KID bytes
|
|
kid_bytes = tenc_bytes[kid_offset : kid_offset + 16]
|
|
hex_kid = kid_bytes.hex()
|
|
|
|
logger.debug(f"KID bytes at offset {kid_offset}: {hex_kid}")
|
|
|
|
# Format as UUID
|
|
try:
|
|
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"Parsed KID: {clean_kid}")
|
|
return clean_kid
|
|
except Exception as e:
|
|
logger.debug(f"Failed to parse KID bytes as UUID: {e}, hex: {hex_kid}")
|
|
# Return hex string as fallback
|
|
return hex_kid
|
|
|
|
except Exception as e:
|
|
logger.debug(f"Failed to parse tenc box: {e}")
|
|
return None
|
|
|
|
# [Keep all the rest of the methods EXACTLY the same...]
|
|
@staticmethod
|
|
def _extract_from_moov(moov_data: bytes) -> List[PSSHData]:
|
|
"""Extract PSSH boxes from moov container"""
|
|
pssh_list = []
|
|
offset = 8 # Skip moov header
|
|
|
|
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"
|
|
)
|
|
|
|
if box_type == "trak":
|
|
# Parse track for PSSH
|
|
trak_data = moov_data[offset : offset + box_size]
|
|
pssh_in_trak = MP4PSSHExtractor._extract_from_trak(trak_data)
|
|
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]
|
|
)
|
|
if pssh_box:
|
|
pssh_list.append(pssh_box)
|
|
|
|
offset += box_size
|
|
|
|
except:
|
|
break
|
|
|
|
return pssh_list
|
|
|
|
@staticmethod
|
|
def _extract_from_trak(trak_data: bytes) -> List[PSSHData]:
|
|
"""Extract PSSH from trak box"""
|
|
pssh_list = []
|
|
offset = 8
|
|
|
|
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"
|
|
)
|
|
|
|
if box_type == "mdia":
|
|
mdia_data = trak_data[offset : offset + box_size]
|
|
pssh_in_mdia = MP4PSSHExtractor._extract_from_mdia(mdia_data)
|
|
pssh_list.extend(pssh_in_mdia)
|
|
|
|
offset += box_size
|
|
|
|
except:
|
|
break
|
|
|
|
return pssh_list
|
|
|
|
@staticmethod
|
|
def _extract_from_mdia(mdia_data: bytes) -> List[PSSHData]:
|
|
"""Extract PSSH from mdia box"""
|
|
pssh_list = []
|
|
offset = 8
|
|
|
|
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"
|
|
)
|
|
|
|
if box_type == "minf":
|
|
minf_data = mdia_data[offset : offset + box_size]
|
|
pssh_in_minf = MP4PSSHExtractor._extract_from_minf(minf_data)
|
|
pssh_list.extend(pssh_in_minf)
|
|
|
|
offset += box_size
|
|
|
|
except:
|
|
break
|
|
|
|
return pssh_list
|
|
|
|
@staticmethod
|
|
def _extract_from_minf(minf_data: bytes) -> List[PSSHData]:
|
|
"""Extract PSSH from minf box"""
|
|
pssh_list = []
|
|
offset = 8
|
|
|
|
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"
|
|
)
|
|
|
|
if box_type == "stbl":
|
|
stbl_data = minf_data[offset : offset + box_size]
|
|
pssh_in_stbl = MP4PSSHExtractor._extract_from_stbl(stbl_data)
|
|
pssh_list.extend(pssh_in_stbl)
|
|
|
|
offset += box_size
|
|
|
|
except:
|
|
break
|
|
|
|
return pssh_list
|
|
|
|
@staticmethod
|
|
def _extract_from_stbl(stbl_data: bytes) -> List[PSSHData]:
|
|
"""Extract PSSH from stbl box (where protection scheme info usually is)"""
|
|
pssh_list = []
|
|
offset = 8
|
|
|
|
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"
|
|
)
|
|
|
|
if box_type == "sinf":
|
|
sinf_data = stbl_data[offset : offset + box_size]
|
|
pssh_in_sinf = MP4PSSHExtractor._extract_from_sinf(sinf_data)
|
|
pssh_list.extend(pssh_in_sinf)
|
|
|
|
offset += box_size
|
|
|
|
except:
|
|
break
|
|
|
|
return pssh_list
|
|
|
|
@staticmethod
|
|
def _extract_from_sinf(sinf_data: bytes) -> List[PSSHData]:
|
|
"""Extract PSSH from sinf (protection scheme information) box"""
|
|
pssh_list = []
|
|
offset = 8
|
|
|
|
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"
|
|
)
|
|
|
|
if box_type == "schi":
|
|
schi_data = sinf_data[offset : offset + box_size]
|
|
pssh_in_schi = MP4PSSHExtractor._extract_from_schi(schi_data)
|
|
pssh_list.extend(pssh_in_schi)
|
|
|
|
offset += box_size
|
|
|
|
except:
|
|
break
|
|
|
|
return pssh_list
|
|
|
|
@staticmethod
|
|
def _extract_from_schi(schi_data: bytes) -> List[PSSHData]:
|
|
"""Extract PSSH from schi box (where PSSH boxes are typically stored)"""
|
|
pssh_list = []
|
|
offset = 8
|
|
|
|
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"
|
|
)
|
|
|
|
if box_type == "pssh":
|
|
pssh_box = MP4PSSHExtractor._parse_pssh_box(
|
|
schi_data[offset : offset + box_size]
|
|
)
|
|
if pssh_box:
|
|
pssh_list.append(pssh_box)
|
|
|
|
offset += box_size
|
|
|
|
except:
|
|
break
|
|
|
|
return pssh_list
|
|
|
|
@staticmethod
|
|
def _parse_pssh_box(pssh_bytes: bytes) -> Optional[PSSHData]:
|
|
"""Parse a PSSH box and extract system_id, pssh_box, and key_ids"""
|
|
try:
|
|
if len(pssh_bytes) < 32: # Minimum size for PSSH box
|
|
return None
|
|
|
|
# Parse box header
|
|
box_size = struct.unpack(">I", pssh_bytes[0:4])[0]
|
|
box_type = pssh_bytes[4:8].decode("ascii")
|
|
|
|
if box_type != "pssh":
|
|
return None
|
|
|
|
version = pssh_bytes[8]
|
|
flags = struct.unpack(">I", b"\x00" + pssh_bytes[9:12])[0]
|
|
|
|
# Extract system ID (bytes 12-28)
|
|
system_id_bytes = pssh_bytes[12:28]
|
|
system_id = str(uuid.UUID(bytes=system_id_bytes))
|
|
|
|
# Extract key IDs (if version > 0)
|
|
key_ids = []
|
|
current_offset = 28
|
|
|
|
if version > 0:
|
|
# Read KID count
|
|
if current_offset + 4 > len(pssh_bytes):
|
|
return None
|
|
|
|
kid_count = struct.unpack(
|
|
">I", pssh_bytes[current_offset : current_offset + 4]
|
|
)[0]
|
|
current_offset += 4
|
|
|
|
# Read each KID
|
|
for _ in range(kid_count):
|
|
if current_offset + 16 > len(pssh_bytes):
|
|
break
|
|
|
|
kid_bytes = pssh_bytes[current_offset : current_offset + 16]
|
|
# kid_uuid = str(uuid.UUID(bytes=kid_bytes))
|
|
# key_ids.append(kid_uuid.replace("-", "").lower())
|
|
key_ids.append(kid_bytes.hex().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")
|
|
|
|
return PSSHData(
|
|
system_id=system_id,
|
|
pssh_box=pssh_b64,
|
|
key_ids=key_ids,
|
|
source="mp4_segment",
|
|
)
|
|
|
|
except Exception as e:
|
|
logger.debug(f"Failed to parse PSSH box: {e}")
|
|
return None
|