Spaces:
Paused
Paused
| import argparse | |
| import struct | |
| import sys | |
| from Crypto.Cipher import AES | |
| from collections import namedtuple | |
| import array | |
| CENCSampleAuxiliaryDataFormat = namedtuple("CENCSampleAuxiliaryDataFormat", ["is_encrypted", "iv", "sub_samples"]) | |
| class MP4Atom: | |
| """ | |
| Represents an MP4 atom, which is a basic unit of data in an MP4 file. | |
| Each atom contains a header (size and type) and data. | |
| """ | |
| __slots__ = ("atom_type", "size", "data") | |
| def __init__(self, atom_type: bytes, size: int, data: memoryview | bytearray): | |
| """ | |
| Initializes an MP4Atom instance. | |
| Args: | |
| atom_type (bytes): The type of the atom. | |
| size (int): The size of the atom. | |
| data (memoryview | bytearray): The data contained in the atom. | |
| """ | |
| self.atom_type = atom_type | |
| self.size = size | |
| self.data = data | |
| def __repr__(self): | |
| return f"<MP4Atom type={self.atom_type}, size={self.size}>" | |
| def pack(self): | |
| """ | |
| Packs the atom into binary data. | |
| Returns: | |
| bytes: Packed binary data with size, type, and data. | |
| """ | |
| return struct.pack(">I", self.size) + self.atom_type + self.data | |
| class MP4Parser: | |
| """ | |
| Parses MP4 data to extract atoms and their structure. | |
| """ | |
| def __init__(self, data: memoryview): | |
| """ | |
| Initializes an MP4Parser instance. | |
| Args: | |
| data (memoryview): The binary data of the MP4 file. | |
| """ | |
| self.data = data | |
| self.position = 0 | |
| def read_atom(self) -> MP4Atom | None: | |
| """ | |
| Reads the next atom from the data. | |
| Returns: | |
| MP4Atom | None: MP4Atom object or None if no more atoms are available. | |
| """ | |
| pos = self.position | |
| if pos + 8 > len(self.data): | |
| return None | |
| size, atom_type = struct.unpack_from(">I4s", self.data, pos) | |
| pos += 8 | |
| if size == 1: | |
| if pos + 8 > len(self.data): | |
| return None | |
| size = struct.unpack_from(">Q", self.data, pos)[0] | |
| pos += 8 | |
| if size < 8 or pos + size - 8 > len(self.data): | |
| return None | |
| atom_data = self.data[pos : pos + size - 8] | |
| self.position = pos + size - 8 | |
| return MP4Atom(atom_type, size, atom_data) | |
| def list_atoms(self) -> list[MP4Atom]: | |
| """ | |
| Lists all atoms in the data. | |
| Returns: | |
| list[MP4Atom]: List of MP4Atom objects. | |
| """ | |
| atoms = [] | |
| original_position = self.position | |
| self.position = 0 | |
| while self.position + 8 <= len(self.data): | |
| atom = self.read_atom() | |
| if not atom: | |
| break | |
| atoms.append(atom) | |
| self.position = original_position | |
| return atoms | |
| def _read_atom_at(self, pos: int, end: int) -> MP4Atom | None: | |
| if pos + 8 > end: | |
| return None | |
| size, atom_type = struct.unpack_from(">I4s", self.data, pos) | |
| pos += 8 | |
| if size == 1: | |
| if pos + 8 > end: | |
| return None | |
| size = struct.unpack_from(">Q", self.data, pos)[0] | |
| pos += 8 | |
| if size < 8 or pos + size - 8 > end: | |
| return None | |
| atom_data = self.data[pos : pos + size - 8] | |
| return MP4Atom(atom_type, size, atom_data) | |
| def print_atoms_structure(self, indent: int = 0): | |
| """ | |
| Prints the structure of all atoms in the data. | |
| Args: | |
| indent (int): The indentation level for printing. | |
| """ | |
| pos = 0 | |
| end = len(self.data) | |
| while pos + 8 <= end: | |
| atom = self._read_atom_at(pos, end) | |
| if not atom: | |
| break | |
| self.print_single_atom_structure(atom, pos, indent) | |
| pos += atom.size | |
| def print_single_atom_structure(self, atom: MP4Atom, parent_position: int, indent: int): | |
| """ | |
| Prints the structure of a single atom. | |
| Args: | |
| atom (MP4Atom): The atom to print. | |
| parent_position (int): The position of the parent atom. | |
| indent (int): The indentation level for printing. | |
| """ | |
| try: | |
| atom_type = atom.atom_type.decode("utf-8") | |
| except UnicodeDecodeError: | |
| atom_type = repr(atom.atom_type) | |
| print(" " * indent + f"Type: {atom_type}, Size: {atom.size}") | |
| child_pos = 0 | |
| child_end = len(atom.data) | |
| while child_pos + 8 <= child_end: | |
| child_atom = self._read_atom_at(parent_position + 8 + child_pos, parent_position + 8 + child_end) | |
| if not child_atom: | |
| break | |
| self.print_single_atom_structure(child_atom, parent_position, indent + 2) | |
| child_pos += child_atom.size | |
| class MP4Decrypter: | |
| """ | |
| Class to handle the decryption of CENC encrypted MP4 segments. | |
| Attributes: | |
| key_map (dict[bytes, bytes]): Mapping of track IDs to decryption keys. | |
| current_key (bytes | None): Current decryption key. | |
| trun_sample_sizes (array.array): Array of sample sizes from the 'trun' box. | |
| current_sample_info (list): List of sample information from the 'senc' box. | |
| encryption_overhead (int): Total size of encryption-related boxes. | |
| """ | |
| def __init__(self, key_map: dict[bytes, bytes]): | |
| """ | |
| Initializes the MP4Decrypter with a key map. | |
| Args: | |
| key_map (dict[bytes, bytes]): Mapping of track IDs to decryption keys. | |
| """ | |
| self.key_map = key_map | |
| self.current_key = None | |
| self.trun_sample_sizes = array.array("I") | |
| self.current_sample_info = [] | |
| self.encryption_overhead = 0 | |
| def decrypt_segment(self, combined_segment: bytes) -> bytes: | |
| """ | |
| Decrypts a combined MP4 segment. | |
| Args: | |
| combined_segment (bytes): Combined initialization and media segment. | |
| Returns: | |
| bytes: Decrypted segment content. | |
| """ | |
| data = memoryview(combined_segment) | |
| parser = MP4Parser(data) | |
| atoms = parser.list_atoms() | |
| atom_process_order = [b"moov", b"moof", b"sidx", b"mdat"] | |
| processed_atoms = {} | |
| for atom_type in atom_process_order: | |
| if atom := next((a for a in atoms if a.atom_type == atom_type), None): | |
| processed_atoms[atom_type] = self._process_atom(atom_type, atom) | |
| result = bytearray() | |
| for atom in atoms: | |
| if atom.atom_type in processed_atoms: | |
| processed_atom = processed_atoms[atom.atom_type] | |
| result.extend(processed_atom.pack()) | |
| else: | |
| result.extend(atom.pack()) | |
| return bytes(result) | |
| def _process_atom(self, atom_type: bytes, atom: MP4Atom) -> MP4Atom: | |
| """ | |
| Processes an MP4 atom based on its type. | |
| Args: | |
| atom_type (bytes): Type of the atom. | |
| atom (MP4Atom): The atom to process. | |
| Returns: | |
| MP4Atom: Processed atom. | |
| """ | |
| match atom_type: | |
| case b"moov": | |
| return self._process_moov(atom) | |
| case b"moof": | |
| return self._process_moof(atom) | |
| case b"sidx": | |
| return self._process_sidx(atom) | |
| case b"mdat": | |
| return self._decrypt_mdat(atom) | |
| case _: | |
| return atom | |
| def _process_moov(self, moov: MP4Atom) -> MP4Atom: | |
| """ | |
| Processes the 'moov' (Movie) atom, which contains metadata about the entire presentation. | |
| This includes information about tracks, media data, and other movie-level metadata. | |
| Args: | |
| moov (MP4Atom): The 'moov' atom to process. | |
| Returns: | |
| MP4Atom: Processed 'moov' atom with updated track information. | |
| """ | |
| parser = MP4Parser(moov.data) | |
| new_moov_data = bytearray() | |
| for atom in iter(parser.read_atom, None): | |
| if atom.atom_type == b"trak": | |
| new_trak = self._process_trak(atom) | |
| new_moov_data.extend(new_trak.pack()) | |
| elif atom.atom_type != b"pssh": | |
| # Skip PSSH boxes as they are not needed in the decrypted output | |
| new_moov_data.extend(atom.pack()) | |
| return MP4Atom(b"moov", len(new_moov_data) + 8, new_moov_data) | |
| def _process_moof(self, moof: MP4Atom) -> MP4Atom: | |
| """ | |
| Processes the 'moov' (Movie) atom, which contains metadata about the entire presentation. | |
| This includes information about tracks, media data, and other movie-level metadata. | |
| Args: | |
| moov (MP4Atom): The 'moov' atom to process. | |
| Returns: | |
| MP4Atom: Processed 'moov' atom with updated track information. | |
| """ | |
| parser = MP4Parser(moof.data) | |
| new_moof_data = bytearray() | |
| for atom in iter(parser.read_atom, None): | |
| if atom.atom_type == b"traf": | |
| new_traf = self._process_traf(atom) | |
| new_moof_data.extend(new_traf.pack()) | |
| else: | |
| new_moof_data.extend(atom.pack()) | |
| return MP4Atom(b"moof", len(new_moof_data) + 8, new_moof_data) | |
| def _process_traf(self, traf: MP4Atom) -> MP4Atom: | |
| """ | |
| Processes the 'traf' (Track Fragment) atom, which contains information about a track fragment. | |
| This includes sample information, sample encryption data, and other track-level metadata. | |
| Args: | |
| traf (MP4Atom): The 'traf' atom to process. | |
| Returns: | |
| MP4Atom: Processed 'traf' atom with updated sample information. | |
| """ | |
| parser = MP4Parser(traf.data) | |
| new_traf_data = bytearray() | |
| tfhd = None | |
| sample_count = 0 | |
| sample_info = [] | |
| atoms = parser.list_atoms() | |
| # calculate encryption_overhead earlier to avoid dependency on trun | |
| self.encryption_overhead = sum(a.size for a in atoms if a.atom_type in {b"senc", b"saiz", b"saio"}) | |
| for atom in atoms: | |
| if atom.atom_type == b"tfhd": | |
| tfhd = atom | |
| new_traf_data.extend(atom.pack()) | |
| elif atom.atom_type == b"trun": | |
| sample_count = self._process_trun(atom) | |
| new_trun = self._modify_trun(atom) | |
| new_traf_data.extend(new_trun.pack()) | |
| elif atom.atom_type == b"senc": | |
| # Parse senc but don't include it in the new decrypted traf data and similarly don't include saiz and saio | |
| sample_info = self._parse_senc(atom, sample_count) | |
| elif atom.atom_type not in {b"saiz", b"saio"}: | |
| new_traf_data.extend(atom.pack()) | |
| if tfhd: | |
| tfhd_track_id = struct.unpack_from(">I", tfhd.data, 4)[0] | |
| self.current_key = self._get_key_for_track(tfhd_track_id) | |
| self.current_sample_info = sample_info | |
| return MP4Atom(b"traf", len(new_traf_data) + 8, new_traf_data) | |
| def _decrypt_mdat(self, mdat: MP4Atom) -> MP4Atom: | |
| """ | |
| Decrypts the 'mdat' (Media Data) atom, which contains the actual media data (audio, video, etc.). | |
| The decryption is performed using the current decryption key and sample information. | |
| Args: | |
| mdat (MP4Atom): The 'mdat' atom to decrypt. | |
| Returns: | |
| MP4Atom: Decrypted 'mdat' atom with decrypted media data. | |
| """ | |
| if not self.current_key or not self.current_sample_info: | |
| return mdat # Return original mdat if we don't have decryption info | |
| decrypted_samples = bytearray() | |
| mdat_data = mdat.data | |
| position = 0 | |
| for i, info in enumerate(self.current_sample_info): | |
| if position >= len(mdat_data): | |
| break # No more data to process | |
| sample_size = self.trun_sample_sizes[i] if i < len(self.trun_sample_sizes) else len(mdat_data) - position | |
| sample = mdat_data[position : position + sample_size] | |
| position += sample_size | |
| decrypted_sample = self._process_sample(sample, info, self.current_key) | |
| decrypted_samples.extend(decrypted_sample) | |
| return MP4Atom(b"mdat", len(decrypted_samples) + 8, decrypted_samples) | |
| def _parse_senc(self, senc: MP4Atom, sample_count: int) -> list[CENCSampleAuxiliaryDataFormat]: | |
| """ | |
| Parses the 'senc' (Sample Encryption) atom, which contains encryption information for samples. | |
| This includes initialization vectors (IVs) and sub-sample encryption data. | |
| Args: | |
| senc (MP4Atom): The 'senc' atom to parse. | |
| sample_count (int): The number of samples. | |
| Returns: | |
| list[CENCSampleAuxiliaryDataFormat]: List of sample auxiliary data formats with encryption information. | |
| """ | |
| data = memoryview(senc.data) | |
| version_flags = struct.unpack_from(">I", data, 0)[0] | |
| version, flags = version_flags >> 24, version_flags & 0xFFFFFF | |
| position = 4 | |
| if version == 0: | |
| sample_count = struct.unpack_from(">I", data, position)[0] | |
| position += 4 | |
| sample_info = [] | |
| for _ in range(sample_count): | |
| if position + 8 > len(data): | |
| break | |
| iv = data[position : position + 8].tobytes() | |
| position += 8 | |
| sub_samples = [] | |
| if flags & 0x000002 and position + 2 <= len(data): # Check if subsample information is present | |
| subsample_count = struct.unpack_from(">H", data, position)[0] | |
| position += 2 | |
| for _ in range(subsample_count): | |
| if position + 6 <= len(data): | |
| clear_bytes, encrypted_bytes = struct.unpack_from(">HI", data, position) | |
| position += 6 | |
| sub_samples.append((clear_bytes, encrypted_bytes)) | |
| else: | |
| break | |
| sample_info.append(CENCSampleAuxiliaryDataFormat(True, iv, sub_samples)) | |
| return sample_info | |
| def _get_key_for_track(self, track_id: int) -> bytes: | |
| """ | |
| Retrieves the decryption key for a given track ID from the key map. | |
| Args: | |
| track_id (int): The track ID. | |
| Returns: | |
| bytes: The decryption key for the specified track ID. | |
| """ | |
| if len(self.key_map) == 1: | |
| return next(iter(self.key_map.values())) | |
| key = self.key_map.get(track_id.pack(4, "big")) | |
| if not key: | |
| raise ValueError(f"No key found for track ID {track_id}") | |
| return key | |
| def _process_sample( | |
| sample: memoryview, sample_info: CENCSampleAuxiliaryDataFormat, key: bytes | |
| ) -> memoryview | bytearray | bytes: | |
| """ | |
| Processes and decrypts a sample using the provided sample information and decryption key. | |
| This includes handling sub-sample encryption if present. | |
| Args: | |
| sample (memoryview): The sample data. | |
| sample_info (CENCSampleAuxiliaryDataFormat): The sample auxiliary data format with encryption information. | |
| key (bytes): The decryption key. | |
| Returns: | |
| memoryview | bytearray | bytes: The decrypted sample. | |
| """ | |
| if not sample_info.is_encrypted: | |
| return sample | |
| # pad IV to 16 bytes | |
| iv = sample_info.iv + b"\x00" * (16 - len(sample_info.iv)) | |
| cipher = AES.new(key, AES.MODE_CTR, initial_value=iv, nonce=b"") | |
| if not sample_info.sub_samples: | |
| # If there are no sub_samples, decrypt the entire sample | |
| return cipher.decrypt(sample) | |
| result = bytearray() | |
| offset = 0 | |
| for clear_bytes, encrypted_bytes in sample_info.sub_samples: | |
| result.extend(sample[offset : offset + clear_bytes]) | |
| offset += clear_bytes | |
| result.extend(cipher.decrypt(sample[offset : offset + encrypted_bytes])) | |
| offset += encrypted_bytes | |
| # If there's any remaining data, treat it as encrypted | |
| if offset < len(sample): | |
| result.extend(cipher.decrypt(sample[offset:])) | |
| return result | |
| def _process_trun(self, trun: MP4Atom) -> int: | |
| """ | |
| Processes the 'trun' (Track Fragment Run) atom, which contains information about the samples in a track fragment. | |
| This includes sample sizes, durations, flags, and composition time offsets. | |
| Args: | |
| trun (MP4Atom): The 'trun' atom to process. | |
| Returns: | |
| int: The number of samples in the 'trun' atom. | |
| """ | |
| trun_flags, sample_count = struct.unpack_from(">II", trun.data, 0) | |
| data_offset = 8 | |
| if trun_flags & 0x000001: | |
| data_offset += 4 | |
| if trun_flags & 0x000004: | |
| data_offset += 4 | |
| self.trun_sample_sizes = array.array("I") | |
| for _ in range(sample_count): | |
| if trun_flags & 0x000100: # sample-duration-present flag | |
| data_offset += 4 | |
| if trun_flags & 0x000200: # sample-size-present flag | |
| sample_size = struct.unpack_from(">I", trun.data, data_offset)[0] | |
| self.trun_sample_sizes.append(sample_size) | |
| data_offset += 4 | |
| else: | |
| self.trun_sample_sizes.append(0) # Using 0 instead of None for uniformity in the array | |
| if trun_flags & 0x000400: # sample-flags-present flag | |
| data_offset += 4 | |
| if trun_flags & 0x000800: # sample-composition-time-offsets-present flag | |
| data_offset += 4 | |
| return sample_count | |
| def _modify_trun(self, trun: MP4Atom) -> MP4Atom: | |
| """ | |
| Modifies the 'trun' (Track Fragment Run) atom to update the data offset. | |
| This is necessary to account for the encryption overhead. | |
| Args: | |
| trun (MP4Atom): The 'trun' atom to modify. | |
| Returns: | |
| MP4Atom: Modified 'trun' atom with updated data offset. | |
| """ | |
| trun_data = bytearray(trun.data) | |
| current_flags = struct.unpack_from(">I", trun_data, 0)[0] & 0xFFFFFF | |
| # If the data-offset-present flag is set, update the data offset to account for encryption overhead | |
| if current_flags & 0x000001: | |
| current_data_offset = struct.unpack_from(">i", trun_data, 8)[0] | |
| struct.pack_into(">i", trun_data, 8, current_data_offset - self.encryption_overhead) | |
| return MP4Atom(b"trun", len(trun_data) + 8, trun_data) | |
| def _process_sidx(self, sidx: MP4Atom) -> MP4Atom: | |
| """ | |
| Processes the 'sidx' (Segment Index) atom, which contains indexing information for media segments. | |
| This includes references to media segments and their durations. | |
| Args: | |
| sidx (MP4Atom): The 'sidx' atom to process. | |
| Returns: | |
| MP4Atom: Processed 'sidx' atom with updated segment references. | |
| """ | |
| sidx_data = bytearray(sidx.data) | |
| current_size = struct.unpack_from(">I", sidx_data, 32)[0] | |
| reference_type = current_size >> 31 | |
| current_referenced_size = current_size & 0x7FFFFFFF | |
| # Remove encryption overhead from referenced size | |
| new_referenced_size = current_referenced_size - self.encryption_overhead | |
| new_size = (reference_type << 31) | new_referenced_size | |
| struct.pack_into(">I", sidx_data, 32, new_size) | |
| return MP4Atom(b"sidx", len(sidx_data) + 8, sidx_data) | |
| def _process_trak(self, trak: MP4Atom) -> MP4Atom: | |
| """ | |
| Processes the 'trak' (Track) atom, which contains information about a single track in the movie. | |
| This includes track header, media information, and other track-level metadata. | |
| Args: | |
| trak (MP4Atom): The 'trak' atom to process. | |
| Returns: | |
| MP4Atom: Processed 'trak' atom with updated track information. | |
| """ | |
| parser = MP4Parser(trak.data) | |
| new_trak_data = bytearray() | |
| for atom in iter(parser.read_atom, None): | |
| if atom.atom_type == b"mdia": | |
| new_mdia = self._process_mdia(atom) | |
| new_trak_data.extend(new_mdia.pack()) | |
| else: | |
| new_trak_data.extend(atom.pack()) | |
| return MP4Atom(b"trak", len(new_trak_data) + 8, new_trak_data) | |
| def _process_mdia(self, mdia: MP4Atom) -> MP4Atom: | |
| """ | |
| Processes the 'mdia' (Media) atom, which contains media information for a track. | |
| This includes media header, handler reference, and media information container. | |
| Args: | |
| mdia (MP4Atom): The 'mdia' atom to process. | |
| Returns: | |
| MP4Atom: Processed 'mdia' atom with updated media information. | |
| """ | |
| parser = MP4Parser(mdia.data) | |
| new_mdia_data = bytearray() | |
| for atom in iter(parser.read_atom, None): | |
| if atom.atom_type == b"minf": | |
| new_minf = self._process_minf(atom) | |
| new_mdia_data.extend(new_minf.pack()) | |
| else: | |
| new_mdia_data.extend(atom.pack()) | |
| return MP4Atom(b"mdia", len(new_mdia_data) + 8, new_mdia_data) | |
| def _process_minf(self, minf: MP4Atom) -> MP4Atom: | |
| """ | |
| Processes the 'minf' (Media Information) atom, which contains information about the media data in a track. | |
| This includes data information, sample table, and other media-level metadata. | |
| Args: | |
| minf (MP4Atom): The 'minf' atom to process. | |
| Returns: | |
| MP4Atom: Processed 'minf' atom with updated media information. | |
| """ | |
| parser = MP4Parser(minf.data) | |
| new_minf_data = bytearray() | |
| for atom in iter(parser.read_atom, None): | |
| if atom.atom_type == b"stbl": | |
| new_stbl = self._process_stbl(atom) | |
| new_minf_data.extend(new_stbl.pack()) | |
| else: | |
| new_minf_data.extend(atom.pack()) | |
| return MP4Atom(b"minf", len(new_minf_data) + 8, new_minf_data) | |
| def _process_stbl(self, stbl: MP4Atom) -> MP4Atom: | |
| """ | |
| Processes the 'stbl' (Sample Table) atom, which contains information about the samples in a track. | |
| This includes sample descriptions, sample sizes, sample times, and other sample-level metadata. | |
| Args: | |
| stbl (MP4Atom): The 'stbl' atom to process. | |
| Returns: | |
| MP4Atom: Processed 'stbl' atom with updated sample information. | |
| """ | |
| parser = MP4Parser(stbl.data) | |
| new_stbl_data = bytearray() | |
| for atom in iter(parser.read_atom, None): | |
| if atom.atom_type == b"stsd": | |
| new_stsd = self._process_stsd(atom) | |
| new_stbl_data.extend(new_stsd.pack()) | |
| else: | |
| new_stbl_data.extend(atom.pack()) | |
| return MP4Atom(b"stbl", len(new_stbl_data) + 8, new_stbl_data) | |
| def _process_stsd(self, stsd: MP4Atom) -> MP4Atom: | |
| """ | |
| Processes the 'stsd' (Sample Description) atom, which contains descriptions of the sample entries in a track. | |
| This includes codec information, sample entry details, and other sample description metadata. | |
| Args: | |
| stsd (MP4Atom): The 'stsd' atom to process. | |
| Returns: | |
| MP4Atom: Processed 'stsd' atom with updated sample descriptions. | |
| """ | |
| parser = MP4Parser(stsd.data) | |
| entry_count = struct.unpack_from(">I", parser.data, 4)[0] | |
| new_stsd_data = bytearray(stsd.data[:8]) | |
| parser.position = 8 # Move past version_flags and entry_count | |
| for _ in range(entry_count): | |
| sample_entry = parser.read_atom() | |
| if not sample_entry: | |
| break | |
| processed_entry = self._process_sample_entry(sample_entry) | |
| new_stsd_data.extend(processed_entry.pack()) | |
| return MP4Atom(b"stsd", len(new_stsd_data) + 8, new_stsd_data) | |
| def _process_sample_entry(self, entry: MP4Atom) -> MP4Atom: | |
| """ | |
| Processes a sample entry atom, which contains information about a specific type of sample. | |
| This includes codec-specific information and other sample entry details. | |
| Args: | |
| entry (MP4Atom): The sample entry atom to process. | |
| Returns: | |
| MP4Atom: Processed sample entry atom with updated information. | |
| """ | |
| # Determine the size of fixed fields based on sample entry type | |
| if entry.atom_type in {b"mp4a", b"enca"}: | |
| fixed_size = 28 # 8 bytes for size, type and reserved, 20 bytes for fixed fields in Audio Sample Entry. | |
| elif entry.atom_type in {b"mp4v", b"encv", b"avc1", b"hev1", b"hvc1"}: | |
| fixed_size = 78 # 8 bytes for size, type and reserved, 70 bytes for fixed fields in Video Sample Entry. | |
| else: | |
| fixed_size = 16 # 8 bytes for size, type and reserved, 8 bytes for fixed fields in other Sample Entries. | |
| new_entry_data = bytearray(entry.data[:fixed_size]) | |
| parser = MP4Parser(entry.data[fixed_size:]) | |
| codec_format = None | |
| for atom in iter(parser.read_atom, None): | |
| if atom.atom_type in {b"sinf", b"schi", b"tenc", b"schm"}: | |
| if atom.atom_type == b"sinf": | |
| codec_format = self._extract_codec_format(atom) | |
| continue # Skip encryption-related atoms | |
| new_entry_data.extend(atom.pack()) | |
| # Replace the atom type with the extracted codec format | |
| new_type = codec_format if codec_format else entry.atom_type | |
| return MP4Atom(new_type, len(new_entry_data) + 8, new_entry_data) | |
| def _extract_codec_format(self, sinf: MP4Atom) -> bytes | None: | |
| """ | |
| Extracts the codec format from the 'sinf' (Protection Scheme Information) atom. | |
| This includes information about the original format of the protected content. | |
| Args: | |
| sinf (MP4Atom): The 'sinf' atom to extract from. | |
| Returns: | |
| bytes | None: The codec format or None if not found. | |
| """ | |
| parser = MP4Parser(sinf.data) | |
| for atom in iter(parser.read_atom, None): | |
| if atom.atom_type == b"frma": | |
| return atom.data | |
| return None | |
| def decrypt_segment(init_segment: bytes, segment_content: bytes, key_id: str, key: str) -> bytes: | |
| """ | |
| Decrypts a CENC encrypted MP4 segment. | |
| Args: | |
| init_segment (bytes): Initialization segment data. | |
| segment_content (bytes): Encrypted segment content. | |
| key_id (str): Key ID in hexadecimal format. | |
| key (str): Key in hexadecimal format. | |
| """ | |
| key_map = {bytes.fromhex(key_id): bytes.fromhex(key)} | |
| decrypter = MP4Decrypter(key_map) | |
| decrypted_content = decrypter.decrypt_segment(init_segment + segment_content) | |
| return decrypted_content | |
| def cli(): | |
| """ | |
| Command line interface for decrypting a CENC encrypted MP4 segment. | |
| """ | |
| init_segment = b"" | |
| if args.init and args.segment: | |
| with open(args.init, "rb") as f: | |
| init_segment = f.read() | |
| with open(args.segment, "rb") as f: | |
| segment_content = f.read() | |
| elif args.combined_segment: | |
| with open(args.combined_segment, "rb") as f: | |
| segment_content = f.read() | |
| else: | |
| print("Usage: python mp4decrypt.py --help") | |
| sys.exit(1) | |
| try: | |
| decrypted_segment = decrypt_segment(init_segment, segment_content, args.key_id, args.key) | |
| print(f"Decrypted content size is {len(decrypted_segment)} bytes") | |
| with open(args.output, "wb") as f: | |
| f.write(decrypted_segment) | |
| print(f"Decrypted segment written to {args.output}") | |
| except Exception as e: | |
| print(f"Error: {e}") | |
| sys.exit(1) | |
| if __name__ == "__main__": | |
| arg_parser = argparse.ArgumentParser(description="Decrypts a MP4 init and media segment using CENC encryption.") | |
| arg_parser.add_argument("--init", help="Path to the init segment file", required=False) | |
| arg_parser.add_argument("--segment", help="Path to the media segment file", required=False) | |
| arg_parser.add_argument( | |
| "--combined_segment", help="Path to the combined init and media segment file", required=False | |
| ) | |
| arg_parser.add_argument("--key_id", help="Key ID in hexadecimal format", required=True) | |
| arg_parser.add_argument("--key", help="Key in hexadecimal format", required=True) | |
| arg_parser.add_argument("--output", help="Path to the output file", required=True) | |
| args = arg_parser.parse_args() | |
| cli() | |