Upload 16 files

mediaflow_proxy/configs.py

@@ -0,0 +1,14 @@
+from pydantic_settings import BaseSettings
+
+
+class Settings(BaseSettings):
+    api_password: str  # The password for accessing the API endpoints.
+    proxy_url: str | None = None  # The URL of the proxy server to route requests through.
+    mpd_live_stream_delay: int = 30  # The delay in seconds for live MPD streams.
+
+    class Config:
+        env_file = ".env"
+        extra = "ignore"
+
+
+settings = Settings()

mediaflow_proxy/const.py

@@ -0,0 +1,24 @@
+SUPPORTED_RESPONSE_HEADERS = [
+    "accept-ranges",
+    "content-type",
+    "content-length",
+    "content-range",
+    "connection",
+    "transfer-encoding",
+    "last-modified",
+    "etag",
+    "cache-control",
+    "expires",
+]
+
+SUPPORTED_REQUEST_HEADERS = [
+    "accept",
+    "accept-encoding",
+    "accept-language",
+    "connection",
+    "range",
+    "if-range",
+    "user-agent",
+    "referer",
+    "origin",
+]

mediaflow_proxy/drm/__init__.py

@@ -0,0 +1,11 @@
+import os
+import tempfile
+
+
+async def create_temp_file(suffix: str, content: bytes = None, prefix: str = None) -> tempfile.NamedTemporaryFile:
+    temp_file = tempfile.NamedTemporaryFile(delete=False, suffix=suffix, prefix=prefix)
+    temp_file.delete_file = lambda: os.unlink(temp_file.name)
+    if content:
+        temp_file.write(content)
+        temp_file.close()
+    return temp_file

mediaflow_proxy/drm/decrypter.py

@@ -0,0 +1,778 @@
+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
+
+    @staticmethod
+    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()

mediaflow_proxy/handlers.py

@@ -0,0 +1,345 @@
+import base64
+import logging
+
+import httpx
+from fastapi import Request, Response, HTTPException
+from pydantic import HttpUrl
+from starlette.background import BackgroundTask
+
+from .configs import settings
+from .const import SUPPORTED_RESPONSE_HEADERS
+from .mpd_processor import process_manifest, process_playlist, process_segment
+from .utils.cache_utils import get_cached_mpd, get_cached_init_segment
+from .utils.http_utils import (
+    Streamer,
+    DownloadError,
+    download_file_with_retry,
+    request_with_retry,
+    EnhancedStreamingResponse,
+)
+from .utils.m3u8_processor import M3U8Processor
+from .utils.mpd_utils import pad_base64
+
+logger = logging.getLogger(__name__)
+
+
+async def handle_hls_stream_proxy(
+    request: Request, destination: str, headers: dict, key_url: HttpUrl = None, verify_ssl: bool = True
+):
+    """
+    Handles the HLS stream proxy request, fetching and processing the m3u8 playlist or streaming the content.
+
+    Args:
+        request (Request): The incoming HTTP request.
+        destination (str): The destination URL to fetch the content from.
+        headers (dict): The headers to include in the request.
+        key_url (str, optional): The HLS Key URL to replace the original key URL. Defaults to None.
+        verify_ssl (bool, optional): Whether to verify the SSL certificate of the destination. Defaults to True.
+
+    Returns:
+        Response: The HTTP response with the processed m3u8 playlist or streamed content.
+    """
+    client = httpx.AsyncClient(
+        follow_redirects=True,
+        timeout=httpx.Timeout(30.0),
+        limits=httpx.Limits(max_keepalive_connections=10, max_connections=20),
+        proxy=settings.proxy_url,
+        verify=verify_ssl,
+    )
+    streamer = Streamer(client)
+    try:
+        if destination.endswith((".m3u", ".m3u8")):
+            return await fetch_and_process_m3u8(streamer, destination, headers, request, key_url)
+
+        response = await streamer.head(destination, headers)
+        if "mpegurl" in response.headers.get("content-type", "").lower():
+            return await fetch_and_process_m3u8(streamer, destination, headers, request, key_url)
+
+        headers.update({"range": headers.get("range", "bytes=0-")})
+        # clean up the headers to only include the necessary headers and remove acl headers
+        response_headers = {k: v for k, v in response.headers.multi_items() if k in SUPPORTED_RESPONSE_HEADERS}
+
+        if transfer_encoding := response_headers.get("transfer-encoding"):
+            if "chunked" not in transfer_encoding:
+                transfer_encoding += ", chunked"
+        else:
+            transfer_encoding = "chunked"
+        response_headers["transfer-encoding"] = transfer_encoding
+
+        return EnhancedStreamingResponse(
+            streamer.stream_content(destination, headers),
+            status_code=response.status_code,
+            headers=response_headers,
+            background=BackgroundTask(streamer.close),
+        )
+    except httpx.HTTPStatusError as e:
+        await client.aclose()
+        logger.error(f"Upstream service error while handling request: {e}")
+        return Response(status_code=e.response.status_code, content=f"Upstream service error: {e}")
+    except DownloadError as e:
+        await client.aclose()
+        logger.error(f"Error downloading {destination}: {e}")
+        return Response(status_code=e.status_code, content=str(e))
+    except Exception as e:
+        await client.aclose()
+        logger.error(f"Internal server error while handling request: {e}")
+        return Response(status_code=502, content=f"Internal server error: {e}")
+
+
+async def proxy_stream(method: str, video_url: str, headers: dict, verify_ssl: bool = True):
+    """
+    Proxies the stream request to the given video URL.
+
+    Args:
+        method (str): The HTTP method (e.g., GET, HEAD).
+        video_url (str): The URL of the video to stream.
+        headers (dict): The headers to include in the request.
+        verify_ssl (bool, optional): Whether to verify the SSL certificate of the destination. Defaults to True.
+
+    Returns:
+        Response: The HTTP response with the streamed content.
+    """
+    return await handle_stream_request(method, video_url, headers, verify_ssl)
+
+
+async def handle_stream_request(method: str, video_url: str, headers: dict, verify_ssl: bool = True):
+    """
+    Handles the stream request, fetching the content from the video URL and streaming it.
+
+    Args:
+        method (str): The HTTP method (e.g., GET, HEAD).
+        video_url (str): The URL of the video to stream.
+        headers (dict): The headers to include in the request.
+        verify_ssl (bool, optional): Whether to verify the SSL certificate of the destination. Defaults to True.
+
+    Returns:
+        Response: The HTTP response with the streamed content.
+    """
+    client = httpx.AsyncClient(
+        follow_redirects=True,
+        timeout=httpx.Timeout(30.0),
+        limits=httpx.Limits(max_keepalive_connections=10, max_connections=20),
+        proxy=settings.proxy_url,
+        verify=verify_ssl,
+    )
+    streamer = Streamer(client)
+    try:
+        response = await streamer.head(video_url, headers)
+        # clean up the headers to only include the necessary headers and remove acl headers
+        response_headers = {k: v for k, v in response.headers.multi_items() if k in SUPPORTED_RESPONSE_HEADERS}
+        if transfer_encoding := response_headers.get("transfer-encoding"):
+            if "chunked" not in transfer_encoding:
+                transfer_encoding += ", chunked"
+        else:
+            transfer_encoding = "chunked"
+        response_headers["transfer-encoding"] = transfer_encoding
+
+        if method == "HEAD":
+            await streamer.close()
+            return Response(headers=response_headers, status_code=response.status_code)
+        else:
+            return EnhancedStreamingResponse(
+                streamer.stream_content(video_url, headers),
+                headers=response_headers,
+                status_code=response.status_code,
+                background=BackgroundTask(streamer.close),
+            )
+    except httpx.HTTPStatusError as e:
+        await client.aclose()
+        logger.error(f"Upstream service error while handling {method} request: {e}")
+        return Response(status_code=e.response.status_code, content=f"Upstream service error: {e}")
+    except DownloadError as e:
+        await client.aclose()
+        logger.error(f"Error downloading {video_url}: {e}")
+        return Response(status_code=e.status_code, content=str(e))
+    except Exception as e:
+        await client.aclose()
+        logger.error(f"Internal server error while handling {method} request: {e}")
+        return Response(status_code=502, content=f"Internal server error: {e}")
+
+
+async def fetch_and_process_m3u8(
+    streamer: Streamer, url: str, headers: dict, request: Request, key_url: HttpUrl = None
+):
+    """
+    Fetches and processes the m3u8 playlist, converting it to an HLS playlist.
+
+    Args:
+        streamer (Streamer): The HTTP client to use for streaming.
+        url (str): The URL of the m3u8 playlist.
+        headers (dict): The headers to include in the request.
+        request (Request): The incoming HTTP request.
+        key_url (HttpUrl, optional): The HLS Key URL to replace the original key URL. Defaults to None.
+
+    Returns:
+        Response: The HTTP response with the processed m3u8 playlist.
+    """
+    try:
+        content = await streamer.get_text(url, headers)
+        processor = M3U8Processor(request, key_url)
+        processed_content = await processor.process_m3u8(content, str(streamer.response.url))
+        return Response(
+            content=processed_content,
+            media_type="application/vnd.apple.mpegurl",
+            headers={
+                "Content-Disposition": "inline",
+                "Accept-Ranges": "none",
+            },
+        )
+    except httpx.HTTPStatusError as e:
+        logger.error(f"HTTP error while fetching m3u8: {e}")
+        return Response(status_code=e.response.status_code, content=str(e))
+    except DownloadError as e:
+        logger.error(f"Error downloading m3u8: {url}")
+        return Response(status_code=502, content=str(e))
+    except Exception as e:
+        logger.exception(f"Unexpected error while processing m3u8: {e}")
+        return Response(status_code=502, content=str(e))
+    finally:
+        await streamer.close()
+
+
+async def handle_drm_key_data(key_id, key, drm_info):
+    """
+    Handles the DRM key data, retrieving the key ID and key from the DRM info if not provided.
+
+    Args:
+        key_id (str): The DRM key ID.
+        key (str): The DRM key.
+        drm_info (dict): The DRM information from the MPD manifest.
+
+    Returns:
+        tuple: The key ID and key.
+    """
+    if drm_info and not drm_info.get("isDrmProtected"):
+        return None, None
+
+    if not key_id or not key:
+        if "keyId" in drm_info and "key" in drm_info:
+            key_id = drm_info["keyId"]
+            key = drm_info["key"]
+        elif "laUrl" in drm_info and "keyId" in drm_info:
+            raise HTTPException(status_code=400, detail="LA URL is not supported yet")
+        else:
+            raise HTTPException(
+                status_code=400, detail="Unable to determine key_id and key, and they were not provided"
+            )
+
+    return key_id, key
+
+
+async def get_manifest(
+    request: Request, mpd_url: str, headers: dict, key_id: str = None, key: str = None, verify_ssl: bool = True
+):
+    """
+    Retrieves and processes the MPD manifest, converting it to an HLS manifest.
+
+    Args:
+        request (Request): The incoming HTTP request.
+        mpd_url (str): The URL of the MPD manifest.
+        headers (dict): The headers to include in the request.
+        key_id (str, optional): The DRM key ID. Defaults to None.
+        key (str, optional): The DRM key. Defaults to None.
+        verify_ssl (bool, optional): Whether to verify the SSL certificate of the destination. Defaults to True.
+
+    Returns:
+        Response: The HTTP response with the HLS manifest.
+    """
+    try:
+        mpd_dict = await get_cached_mpd(
+            mpd_url, headers=headers, parse_drm=not key_id and not key, verify_ssl=verify_ssl
+        )
+    except DownloadError as e:
+        raise HTTPException(status_code=e.status_code, detail=f"Failed to download MPD: {e.message}")
+    drm_info = mpd_dict.get("drmInfo", {})
+
+    if drm_info and not drm_info.get("isDrmProtected"):
+        # For non-DRM protected MPD, we still create an HLS manifest
+        return await process_manifest(request, mpd_dict, None, None)
+
+    key_id, key = await handle_drm_key_data(key_id, key, drm_info)
+
+    # check if the provided key_id and key are valid
+    if key_id and len(key_id) != 32:
+        key_id = base64.urlsafe_b64decode(pad_base64(key_id)).hex()
+    if key and len(key) != 32:
+        key = base64.urlsafe_b64decode(pad_base64(key)).hex()
+
+    return await process_manifest(request, mpd_dict, key_id, key)
+
+
+async def get_playlist(
+    request: Request,
+    mpd_url: str,
+    profile_id: str,
+    headers: dict,
+    key_id: str = None,
+    key: str = None,
+    verify_ssl: bool = True,
+):
+    """
+    Retrieves and processes the MPD manifest, converting it to an HLS playlist for a specific profile.
+
+    Args:
+        request (Request): The incoming HTTP request.
+        mpd_url (str): The URL of the MPD manifest.
+        profile_id (str): The profile ID to generate the playlist for.
+        headers (dict): The headers to include in the request.
+        key_id (str, optional): The DRM key ID. Defaults to None.
+        key (str, optional): The DRM key. Defaults to None.
+        verify_ssl (bool, optional): Whether to verify the SSL certificate of the destination. Defaults to True.
+
+    Returns:
+        Response: The HTTP response with the HLS playlist.
+    """
+    mpd_dict = await get_cached_mpd(
+        mpd_url,
+        headers=headers,
+        parse_drm=not key_id and not key,
+        parse_segment_profile_id=profile_id,
+        verify_ssl=verify_ssl,
+    )
+    return await process_playlist(request, mpd_dict, profile_id)
+
+
+async def get_segment(
+    init_url: str,
+    segment_url: str,
+    mimetype: str,
+    headers: dict,
+    key_id: str = None,
+    key: str = None,
+    verify_ssl: bool = True,
+):
+    """
+    Retrieves and processes a media segment, decrypting it if necessary.
+
+    Args:
+        init_url (str): The URL of the initialization segment.
+        segment_url (str): The URL of the media segment.
+        mimetype (str): The MIME type of the segment.
+        headers (dict): The headers to include in the request.
+        key_id (str, optional): The DRM key ID. Defaults to None.
+        key (str, optional): The DRM key. Defaults to None.
+        verify_ssl (bool, optional): Whether to verify the SSL certificate of the destination. Defaults to True.
+
+    Returns:
+        Response: The HTTP response with the processed segment.
+    """
+    try:
+        init_content = await get_cached_init_segment(init_url, headers, verify_ssl)
+        segment_content = await download_file_with_retry(segment_url, headers, verify_ssl=verify_ssl)
+    except DownloadError as e:
+        raise HTTPException(status_code=e.status_code, detail=f"Failed to download segment: {e.message}")
+    return await process_segment(init_content, segment_content, mimetype, key_id, key)
+
+
+async def get_public_ip():
+    """
+    Retrieves the public IP address of the MediaFlow proxy.
+
+    Returns:
+        Response: The HTTP response with the public IP address.
+    """
+    ip_address_data = await request_with_retry("GET", "https://api.ipify.org?format=json", {})
+    return ip_address_data.json()

mediaflow_proxy/main.py

@@ -0,0 +1,58 @@
+import logging
+from importlib import resources
+
+from fastapi import FastAPI, Depends, Security, HTTPException
+from fastapi.security import APIKeyQuery, APIKeyHeader
+from starlette.responses import RedirectResponse
+from starlette.staticfiles import StaticFiles
+
+from mediaflow_proxy.configs import settings
+from mediaflow_proxy.routes import proxy_router
+
+logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(name)s - %(levelname)s - %(message)s")
+app = FastAPI()
+api_password_query = APIKeyQuery(name="api_password", auto_error=False)
+api_password_header = APIKeyHeader(name="api_password", auto_error=False)
+
+
+async def verify_api_key(api_key: str = Security(api_password_query), api_key_alt: str = Security(api_password_header)):
+    """
+    Verifies the API key for the request.
+
+    Args:
+        api_key (str): The API key to validate.
+        api_key_alt (str): The alternative API key to validate.
+
+    Raises:
+        HTTPException: If the API key is invalid.
+    """
+    if api_key == settings.api_password or api_key_alt == settings.api_password:
+        return
+
+    raise HTTPException(status_code=403, detail="Could not validate credentials")
+
+
+@app.get("/health")
+async def health_check():
+    return {"status": "healthy"}
+
+
+@app.get("/favicon.ico")
+async def get_favicon():
+    return RedirectResponse(url="/logo.png")
+
+
+app.include_router(proxy_router, prefix="/proxy", tags=["proxy"], dependencies=[Depends(verify_api_key)])
+
+static_path = resources.files("mediaflow_proxy").joinpath("static")
+app.mount("/", StaticFiles(directory=str(static_path), html=True), name="static")
+
+
+def run():
+    import uvicorn
+
+    uvicorn.run(app, host="127.0.0.1", port=8888)
+
+
+if __name__ == "__main__":
+    run()

mediaflow_proxy/mpd_processor.py

@@ -0,0 +1,210 @@
+import logging
+import math
+import time
+from datetime import datetime, timezone, timedelta
+
+from fastapi import Request, Response, HTTPException
+
+from mediaflow_proxy.configs import settings
+from mediaflow_proxy.drm.decrypter import decrypt_segment
+from mediaflow_proxy.utils.http_utils import encode_mediaflow_proxy_url, get_original_scheme
+
+logger = logging.getLogger(__name__)
+
+
+async def process_manifest(request: Request, mpd_dict: dict, key_id: str = None, key: str = None) -> Response:
+    """
+    Processes the MPD manifest and converts it to an HLS manifest.
+
+    Args:
+        request (Request): The incoming HTTP request.
+        mpd_dict (dict): The MPD manifest data.
+        key_id (str, optional): The DRM key ID. Defaults to None.
+        key (str, optional): The DRM key. Defaults to None.
+
+    Returns:
+        Response: The HLS manifest as an HTTP response.
+    """
+    hls_content = build_hls(mpd_dict, request, key_id, key)
+    return Response(content=hls_content, media_type="application/vnd.apple.mpegurl")
+
+
+async def process_playlist(request: Request, mpd_dict: dict, profile_id: str) -> Response:
+    """
+    Processes the MPD manifest and converts it to an HLS playlist for a specific profile.
+
+    Args:
+        request (Request): The incoming HTTP request.
+        mpd_dict (dict): The MPD manifest data.
+        profile_id (str): The profile ID to generate the playlist for.
+
+    Returns:
+        Response: The HLS playlist as an HTTP response.
+
+    Raises:
+        HTTPException: If the profile is not found in the MPD manifest.
+    """
+    matching_profiles = [p for p in mpd_dict["profiles"] if p["id"] == profile_id]
+    if not matching_profiles:
+        raise HTTPException(status_code=404, detail="Profile not found")
+
+    hls_content = build_hls_playlist(mpd_dict, matching_profiles, request)
+    return Response(content=hls_content, media_type="application/vnd.apple.mpegurl")
+
+
+async def process_segment(
+    init_content: bytes,
+    segment_content: bytes,
+    mimetype: str,
+    key_id: str = None,
+    key: str = None,
+) -> Response:
+    """
+    Processes and decrypts a media segment.
+
+    Args:
+        init_content (bytes): The initialization segment content.
+        segment_content (bytes): The media segment content.
+        mimetype (str): The MIME type of the segment.
+        key_id (str, optional): The DRM key ID. Defaults to None.
+        key (str, optional): The DRM key. Defaults to None.
+
+    Returns:
+        Response: The decrypted segment as an HTTP response.
+    """
+    if key_id and key:
+        # For DRM protected content
+        now = time.time()
+        decrypted_content = decrypt_segment(init_content, segment_content, key_id, key)
+        logger.info(f"Decryption of {mimetype} segment took {time.time() - now:.4f} seconds")
+    else:
+        # For non-DRM protected content, we just concatenate init and segment content
+        decrypted_content = init_content + segment_content
+
+    return Response(content=decrypted_content, media_type=mimetype)
+
+
+def build_hls(mpd_dict: dict, request: Request, key_id: str = None, key: str = None) -> str:
+    """
+    Builds an HLS manifest from the MPD manifest.
+
+    Args:
+        mpd_dict (dict): The MPD manifest data.
+        request (Request): The incoming HTTP request.
+        key_id (str, optional): The DRM key ID. Defaults to None.
+        key (str, optional): The DRM key. Defaults to None.
+
+    Returns:
+        str: The HLS manifest as a string.
+    """
+    hls = ["#EXTM3U", "#EXT-X-VERSION:6"]
+    query_params = dict(request.query_params)
+
+    video_profiles = {}
+    audio_profiles = {}
+
+    # Get the base URL for the playlist_endpoint endpoint
+    proxy_url = request.url_for("playlist_endpoint")
+    proxy_url = str(proxy_url.replace(scheme=get_original_scheme(request)))
+
+    for profile in mpd_dict["profiles"]:
+        query_params.update({"profile_id": profile["id"], "key_id": key_id or "", "key": key or ""})
+        playlist_url = encode_mediaflow_proxy_url(
+            proxy_url,
+            query_params=query_params,
+        )
+
+        if "video" in profile["mimeType"]:
+            video_profiles[profile["id"]] = (profile, playlist_url)
+        elif "audio" in profile["mimeType"]:
+            audio_profiles[profile["id"]] = (profile, playlist_url)
+
+    # Add audio streams
+    for i, (profile, playlist_url) in enumerate(audio_profiles.values()):
+        is_default = "YES" if i == 0 else "NO"  # Set the first audio track as default
+        hls.append(
+            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}"'
+        )
+
+    # Add video streams
+    for profile, playlist_url in video_profiles.values():
+        hls.append(
+            f'#EXT-X-STREAM-INF:BANDWIDTH={profile["bandwidth"]},RESOLUTION={profile["width"]}x{profile["height"]},CODECS="{profile["codecs"]}",FRAME-RATE={profile["frameRate"]},AUDIO="audio"'
+        )
+        hls.append(playlist_url)
+
+    return "\n".join(hls)
+
+
+def build_hls_playlist(mpd_dict: dict, profiles: list[dict], request: Request) -> str:
+    """
+    Builds an HLS playlist from the MPD manifest for specific profiles.
+
+    Args:
+        mpd_dict (dict): The MPD manifest data.
+        profiles (list[dict]): The profiles to include in the playlist.
+        request (Request): The incoming HTTP request.
+
+    Returns:
+        str: The HLS playlist as a string.
+    """
+    hls = ["#EXTM3U", "#EXT-X-VERSION:6"]
+
+    added_segments = 0
+    current_time = datetime.now(timezone.utc)
+    live_stream_delay = timedelta(seconds=settings.mpd_live_stream_delay)
+    target_end_time = current_time - live_stream_delay
+
+    proxy_url = request.url_for("segment_endpoint")
+    proxy_url = str(proxy_url.replace(scheme=get_original_scheme(request)))
+
+    for index, profile in enumerate(profiles):
+        segments = profile["segments"]
+        if not segments:
+            logger.warning(f"No segments found for profile {profile['id']}")
+            continue
+
+        # Add headers for only the first profile
+        if index == 0:
+            sequence = segments[0]["number"]
+            extinf_values = [f["extinf"] for f in segments if "extinf" in f]
+            target_duration = math.ceil(max(extinf_values)) if extinf_values else 3
+            hls.extend(
+                [
+                    f"#EXT-X-TARGETDURATION:{target_duration}",
+                    f"#EXT-X-MEDIA-SEQUENCE:{sequence}",
+                ]
+            )
+            if mpd_dict["isLive"]:
+                hls.append("#EXT-X-PLAYLIST-TYPE:EVENT")
+            else:
+                hls.append("#EXT-X-PLAYLIST-TYPE:VOD")
+
+        init_url = profile["initUrl"]
+
+        query_params = dict(request.query_params)
+        query_params.pop("profile_id", None)
+        query_params.pop("d", None)
+
+        for segment in segments:
+            if mpd_dict["isLive"]:
+                if segment["end_time"] > target_end_time:
+                    continue
+                hls.append(f"#EXT-X-PROGRAM-DATE-TIME:{segment['program_date_time']}")
+            hls.append(f'#EXTINF:{segment["extinf"]:.3f},')
+            query_params.update(
+                {"init_url": init_url, "segment_url": segment["media"], "mime_type": profile["mimeType"]}
+            )
+            hls.append(
+                encode_mediaflow_proxy_url(
+                    proxy_url,
+                    query_params=query_params,
+                )
+            )
+            added_segments += 1
+
+    if not mpd_dict["isLive"]:
+        hls.append("#EXT-X-ENDLIST")
+
+    logger.info(f"Added {added_segments} segments to HLS playlist")
+    return "\n".join(hls)

mediaflow_proxy/routes.py

@@ -0,0 +1,147 @@
+from fastapi import Request, Depends, APIRouter
+from pydantic import HttpUrl
+
+from .handlers import handle_hls_stream_proxy, proxy_stream, get_manifest, get_playlist, get_segment, get_public_ip
+from .utils.http_utils import get_proxy_headers
+
+proxy_router = APIRouter()
+
+
+@proxy_router.head("/hls")
+@proxy_router.get("/hls")
+async def hls_stream_proxy(
+    request: Request,
+    d: HttpUrl,
+    headers: dict = Depends(get_proxy_headers),
+    key_url: HttpUrl | None = None,
+    verify_ssl: bool = False,
+):
+    """
+    Proxify HLS stream requests, fetching and processing the m3u8 playlist or streaming the content.
+
+    Args:
+        request (Request): The incoming HTTP request.
+        d (HttpUrl): The destination URL to fetch the content from.
+        key_url (HttpUrl, optional): The HLS Key URL to replace the original key URL. Defaults to None. (Useful for bypassing some sneaky protection)
+        headers (dict): The headers to include in the request.
+        verify_ssl (bool, optional): Whether to verify the SSL certificate of the destination. Defaults to False.
+
+    Returns:
+        Response: The HTTP response with the processed m3u8 playlist or streamed content.
+    """
+    destination = str(d)
+    return await handle_hls_stream_proxy(request, destination, headers, key_url, verify_ssl)
+
+
+@proxy_router.head("/stream")
+@proxy_router.get("/stream")
+async def proxy_stream_endpoint(
+    request: Request, d: HttpUrl, headers: dict = Depends(get_proxy_headers), verify_ssl: bool = False
+):
+    """
+    Proxies stream requests to the given video URL.
+
+    Args:
+        request (Request): The incoming HTTP request.
+        d (HttpUrl): The URL of the video to stream.
+        headers (dict): The headers to include in the request.
+        verify_ssl (bool, optional): Whether to verify the SSL certificate of the destination. Defaults to False.
+
+    Returns:
+        Response: The HTTP response with the streamed content.
+    """
+    headers.update({"range": headers.get("range", "bytes=0-")})
+    return await proxy_stream(request.method, str(d), headers, verify_ssl)
+
+
+@proxy_router.get("/mpd/manifest")
+async def manifest_endpoint(
+    request: Request,
+    d: HttpUrl,
+    headers: dict = Depends(get_proxy_headers),
+    key_id: str = None,
+    key: str = None,
+    verify_ssl: bool = False,
+):
+    """
+    Retrieves and processes the MPD manifest, converting it to an HLS manifest.
+
+    Args:
+        request (Request): The incoming HTTP request.
+        d (HttpUrl): The URL of the MPD manifest.
+        headers (dict): The headers to include in the request.
+        key_id (str, optional): The DRM key ID. Defaults to None.
+        key (str, optional): The DRM key. Defaults to None.
+        verify_ssl (bool, optional): Whether to verify the SSL certificate of the destination. Defaults to False.
+
+    Returns:
+        Response: The HTTP response with the HLS manifest.
+    """
+    return await get_manifest(request, str(d), headers, key_id, key, verify_ssl)
+
+
+@proxy_router.get("/mpd/playlist")
+async def playlist_endpoint(
+    request: Request,
+    d: HttpUrl,
+    profile_id: str,
+    headers: dict = Depends(get_proxy_headers),
+    key_id: str = None,
+    key: str = None,
+    verify_ssl: bool = False,
+):
+    """
+    Retrieves and processes the MPD manifest, converting it to an HLS playlist for a specific profile.
+
+    Args:
+        request (Request): The incoming HTTP request.
+        d (HttpUrl): The URL of the MPD manifest.
+        profile_id (str): The profile ID to generate the playlist for.
+        headers (dict): The headers to include in the request.
+        key_id (str, optional): The DRM key ID. Defaults to None.
+        key (str, optional): The DRM key. Defaults to None.
+        verify_ssl (bool, optional): Whether to verify the SSL certificate of the destination. Defaults to False.
+
+    Returns:
+        Response: The HTTP response with the HLS playlist.
+    """
+    return await get_playlist(request, str(d), profile_id, headers, key_id, key, verify_ssl)
+
+
+@proxy_router.get("/mpd/segment")
+async def segment_endpoint(
+    init_url: HttpUrl,
+    segment_url: HttpUrl,
+    mime_type: str,
+    headers: dict = Depends(get_proxy_headers),
+    key_id: str = None,
+    key: str = None,
+    verify_ssl: bool = False,
+):
+    """
+    Retrieves and processes a media segment, decrypting it if necessary.
+
+    Args:
+        init_url (HttpUrl): The URL of the initialization segment.
+        segment_url (HttpUrl): The URL of the media segment.
+        mime_type (str): The MIME type of the segment.
+        headers (dict): The headers to include in the request.
+        key_id (str, optional): The DRM key ID. Defaults to None.
+        key (str, optional): The DRM key. Defaults to None.
+        verify_ssl (bool, optional): Whether to verify the SSL certificate of the destination. Defaults to False.
+
+    Returns:
+        Response: The HTTP response with the processed segment.
+    """
+    return await get_segment(str(init_url), str(segment_url), mime_type, headers, key_id, key, verify_ssl)
+
+
+@proxy_router.get("/ip")
+async def get_mediaflow_proxy_public_ip():
+    """
+    Retrieves the public IP address of the MediaFlow proxy server.
+
+    Returns:
+        Response: The HTTP response with the public IP address in the form of a JSON object. {"ip": "xxx.xxx.xxx.xxx"}
+    """
+    return await get_public_ip()

mediaflow_proxy/static/index.html

@@ -0,0 +1,76 @@
+<!DOCTYPE html>
+<html lang="en">
+<head>
+    <meta charset="UTF-8">
+    <meta name="viewport" content="width=device-width, initial-scale=1.0">
+    <title>MediaFlow Proxy</title>
+    <link rel="icon" href="/logo.png" type="image/x-icon">
+    <style>
+        body {
+            font-family: Arial, sans-serif;
+            line-height: 1.6;
+            color: #333;
+            max-width: 800px;
+            margin: 0 auto;
+            padding: 20px;
+            background-color: #f9f9f9;
+        }
+
+        header {
+            background-color: #90aacc;
+            color: #fff;
+            padding: 10px 0;
+            text-align: center;
+        }
+
+        header img {
+            width: 200px;
+            height: 200px;
+            vertical-align: middle;
+            border-radius: 15px;
+        }
+
+        header h1 {
+            display: inline;
+            margin-left: 20px;
+            font-size: 36px;
+        }
+
+        .feature {
+            background-color: #f4f4f4;
+            border-left: 4px solid #3498db;
+            padding: 10px;
+            margin-bottom: 10px;
+        }
+
+        a {
+            color: #3498db;
+        }
+    </style>
+</head>
+<body>
+<header>
+    <img src="/logo.png" alt="MediaFlow Proxy Logo">
+    <h1>MediaFlow Proxy</h1>
+</header>
+<p>A high-performance proxy server for streaming media, supporting HTTP(S), HLS, and MPEG-DASH with real-time DRM decryption.</p>
+
+<h2>Key Features</h2>
+<div class="feature">Convert MPEG-DASH streams (DRM-protected and non-protected) to HLS</div>
+<div class="feature">Support for Clear Key DRM-protected MPD DASH streams</div>
+<div class="feature">Handle both live and video-on-demand (VOD) DASH streams</div>
+<div class="feature">Proxy HTTP/HTTPS links with custom headers</div>
+<div class="feature">Proxy and modify HLS (M3U8) streams in real-time with custom headers and key URL modifications for bypassing some sneaky restrictions.</div>
+<div class="feature">Protect against unauthorized access and network bandwidth abuses</div>
+
+<h2>Getting Started</h2>
+<p>Visit the <a href="https://github.com/mhdzumair/mediaflow-proxy">GitHub repository</a> for installation instructions and documentation.</p>
+
+<h2>Premium Hosted Service</h2>
+<p>For a hassle-free experience, check out <a href="https://store.elfhosted.com/product/mediaflow-proxy">premium hosted service on ElfHosted</a>.</p>
+
+<h2>API Documentation</h2>
+<p>Explore the <a href="/docs">Swagger UI</a> for comprehensive details about the API endpoints and their usage.</p>
+
+</body>
+</html>

mediaflow_proxy/utils/cache_utils.py

@@ -0,0 +1,60 @@
+import datetime
+import logging
+
+from cachetools import TTLCache
+
+from .http_utils import download_file_with_retry
+from .mpd_utils import parse_mpd, parse_mpd_dict
+
+logger = logging.getLogger(__name__)
+
+# cache dictionary
+mpd_cache = TTLCache(maxsize=100, ttl=300)  # 5 minutes default TTL
+init_segment_cache = TTLCache(maxsize=100, ttl=3600)  # 1 hour default TTL
+
+
+async def get_cached_mpd(
+    mpd_url: str, headers: dict, parse_drm: bool, parse_segment_profile_id: str | None = None, verify_ssl: bool = True
+) -> dict:
+    """
+    Retrieves and caches the MPD manifest, parsing it if not already cached.
+
+    Args:
+        mpd_url (str): The URL of the MPD manifest.
+        headers (dict): The headers to include in the request.
+        parse_drm (bool): Whether to parse DRM information.
+        parse_segment_profile_id (str, optional): The profile ID to parse segments for. Defaults to None.
+        verify_ssl (bool, optional): Whether to verify the SSL certificate of the destination. Defaults to True.
+
+    Returns:
+        dict: The parsed MPD manifest data.
+    """
+    current_time = datetime.datetime.now(datetime.UTC)
+    if mpd_url in mpd_cache and mpd_cache[mpd_url]["expires"] > current_time:
+        logger.info(f"Using cached MPD for {mpd_url}")
+        return parse_mpd_dict(mpd_cache[mpd_url]["mpd"], mpd_url, parse_drm, parse_segment_profile_id)
+
+    mpd_dict = parse_mpd(await download_file_with_retry(mpd_url, headers, verify_ssl=verify_ssl))
+    parsed_mpd_dict = parse_mpd_dict(mpd_dict, mpd_url, parse_drm, parse_segment_profile_id)
+    current_time = datetime.datetime.now(datetime.UTC)
+    expiration_time = current_time + datetime.timedelta(seconds=parsed_mpd_dict.get("minimumUpdatePeriod", 300))
+    mpd_cache[mpd_url] = {"mpd": mpd_dict, "expires": expiration_time}
+    return parsed_mpd_dict
+
+
+async def get_cached_init_segment(init_url: str, headers: dict, verify_ssl: bool = True) -> bytes:
+    """
+    Retrieves and caches the initialization segment.
+
+    Args:
+        init_url (str): The URL of the initialization segment.
+        headers (dict): The headers to include in the request.
+        verify_ssl (bool, optional): Whether to verify the SSL certificate of the destination. Defaults to True.
+
+    Returns:
+        bytes: The initialization segment content.
+    """
+    if init_url not in init_segment_cache:
+        init_content = await download_file_with_retry(init_url, headers, verify_ssl=verify_ssl)
+        init_segment_cache[init_url] = init_content
+    return init_segment_cache[init_url]

mediaflow_proxy/utils/http_utils.py

@@ -0,0 +1,355 @@
+import logging
+import typing
+from functools import partial
+from urllib import parse
+
+import anyio
+import httpx
+import tenacity
+from fastapi import Response
+from starlette.background import BackgroundTask
+from starlette.concurrency import iterate_in_threadpool
+from starlette.requests import Request
+from starlette.types import Receive, Send, Scope
+from tenacity import retry, stop_after_attempt, wait_exponential, retry_if_exception_type
+
+from mediaflow_proxy.configs import settings
+from mediaflow_proxy.const import SUPPORTED_REQUEST_HEADERS
+
+logger = logging.getLogger(__name__)
+
+
+class DownloadError(Exception):
+    def __init__(self, status_code, message):
+        self.status_code = status_code
+        self.message = message
+        super().__init__(message)
+
+
+@retry(
+    stop=stop_after_attempt(3),
+    wait=wait_exponential(multiplier=1, min=4, max=10),
+    retry=retry_if_exception_type(DownloadError),
+)
+async def fetch_with_retry(client, method, url, headers, follow_redirects=True, **kwargs):
+    """
+    Fetches a URL with retry logic.
+
+    Args:
+        client (httpx.AsyncClient): The HTTP client to use for the request.
+        method (str): The HTTP method to use (e.g., GET, POST).
+        url (str): The URL to fetch.
+        headers (dict): The headers to include in the request.
+        follow_redirects (bool, optional): Whether to follow redirects. Defaults to True.
+        **kwargs: Additional arguments to pass to the request.
+
+    Returns:
+        httpx.Response: The HTTP response.
+
+    Raises:
+        DownloadError: If the request fails after retries.
+    """
+    try:
+        response = await client.request(method, url, headers=headers, follow_redirects=follow_redirects, **kwargs)
+        response.raise_for_status()
+        return response
+    except httpx.TimeoutException:
+        logger.warning(f"Timeout while downloading {url}")
+        raise DownloadError(409, f"Timeout while downloading {url}")
+    except httpx.HTTPStatusError as e:
+        logger.error(f"HTTP error {e.response.status_code} while downloading {url}")
+        # if e.response.status_code == 404:
+        #     logger.error(f"Segment Resource not found: {url}")
+        #     raise e
+        raise DownloadError(e.response.status_code, f"HTTP error {e.response.status_code} while downloading {url}")
+    except Exception as e:
+        logger.error(f"Error downloading {url}: {e}")
+        raise
+
+
+class Streamer:
+    def __init__(self, client):
+        """
+        Initializes the Streamer with an HTTP client.
+
+        Args:
+            client (httpx.AsyncClient): The HTTP client to use for streaming.
+        """
+        self.client = client
+        self.response = None
+
+    async def stream_content(self, url: str, headers: dict):
+        """
+        Streams content from a URL.
+
+        Args:
+            url (str): The URL to stream content from.
+            headers (dict): The headers to include in the request.
+
+        Yields:
+            bytes: Chunks of the streamed content.
+        """
+        async with self.client.stream("GET", url, headers=headers, follow_redirects=True) as self.response:
+            self.response.raise_for_status()
+            async for chunk in self.response.aiter_raw():
+                yield chunk
+
+    async def head(self, url: str, headers: dict):
+        """
+        Sends a HEAD request to a URL.
+
+        Args:
+            url (str): The URL to send the HEAD request to.
+            headers (dict): The headers to include in the request.
+
+        Returns:
+            httpx.Response: The HTTP response.
+        """
+        try:
+            self.response = await fetch_with_retry(self.client, "HEAD", url, headers)
+        except tenacity.RetryError as e:
+            raise e.last_attempt.result()
+        return self.response
+
+    async def get_text(self, url: str, headers: dict):
+        """
+        Sends a GET request to a URL and returns the response text.
+
+        Args:
+            url (str): The URL to send the GET request to.
+            headers (dict): The headers to include in the request.
+
+        Returns:
+            str: The response text.
+        """
+        try:
+            self.response = await fetch_with_retry(self.client, "GET", url, headers)
+        except tenacity.RetryError as e:
+            raise e.last_attempt.result()
+        return self.response.text
+
+    async def close(self):
+        """
+        Closes the HTTP client and response.
+        """
+        if self.response:
+            await self.response.aclose()
+        await self.client.aclose()
+
+
+async def download_file_with_retry(url: str, headers: dict, timeout: float = 10.0, verify_ssl: bool = True):
+    """
+    Downloads a file with retry logic.
+
+    Args:
+        url (str): The URL of the file to download.
+        headers (dict): The headers to include in the request.
+        timeout (float, optional): The request timeout. Defaults to 10.0.
+        verify_ssl (bool, optional): Whether to verify the SSL certificate of the destination. Defaults to True.
+
+    Returns:
+        bytes: The downloaded file content.
+
+    Raises:
+        DownloadError: If the download fails after retries.
+    """
+    async with httpx.AsyncClient(
+        follow_redirects=True, timeout=timeout, proxy=settings.proxy_url, verify=verify_ssl
+    ) as client:
+        try:
+            response = await fetch_with_retry(client, "GET", url, headers)
+            return response.content
+        except DownloadError as e:
+            logger.error(f"Failed to download file: {e}")
+            raise e
+        except tenacity.RetryError as e:
+            raise DownloadError(502, f"Failed to download file: {e.last_attempt.result()}")
+
+
+async def request_with_retry(method: str, url: str, headers: dict, timeout: float = 10.0, **kwargs):
+    """
+    Sends an HTTP request with retry logic.
+
+    Args:
+        method (str): The HTTP method to use (e.g., GET, POST).
+        url (str): The URL to send the request to.
+        headers (dict): The headers to include in the request.
+        timeout (float, optional): The request timeout. Defaults to 10.0.
+        **kwargs: Additional arguments to pass to the request.
+
+    Returns:
+        httpx.Response: The HTTP response.
+
+    Raises:
+        DownloadError: If the request fails after retries.
+    """
+    async with httpx.AsyncClient(follow_redirects=True, timeout=timeout, proxy=settings.proxy_url) as client:
+        try:
+            response = await fetch_with_retry(client, method, url, headers, **kwargs)
+            return response
+        except DownloadError as e:
+            logger.error(f"Failed to download file: {e}")
+            raise
+
+
+def encode_mediaflow_proxy_url(
+    mediaflow_proxy_url: str,
+    endpoint: str | None = None,
+    destination_url: str | None = None,
+    query_params: dict | None = None,
+    request_headers: dict | None = None,
+) -> str:
+    """
+    Encodes a MediaFlow proxy URL with query parameters and headers.
+
+    Args:
+        mediaflow_proxy_url (str): The base MediaFlow proxy URL.
+        endpoint (str, optional): The endpoint to append to the base URL. Defaults to None.
+        destination_url (str, optional): The destination URL to include in the query parameters. Defaults to None.
+        query_params (dict, optional): Additional query parameters to include. Defaults to None.
+        request_headers (dict, optional): Headers to include as query parameters. Defaults to None.
+
+    Returns:
+        str: The encoded MediaFlow proxy URL.
+    """
+    query_params = query_params or {}
+    if destination_url is not None:
+        query_params["d"] = destination_url
+
+    # Add headers if provided
+    if request_headers:
+        query_params.update(
+            {key if key.startswith("h_") else f"h_{key}": value for key, value in request_headers.items()}
+        )
+    # Encode the query parameters
+    encoded_params = parse.urlencode(query_params, quote_via=parse.quote)
+
+    # Construct the full URL
+    if endpoint is None:
+        return f"{mediaflow_proxy_url}?{encoded_params}"
+
+    base_url = parse.urljoin(mediaflow_proxy_url, endpoint)
+    return f"{base_url}?{encoded_params}"
+
+
+def get_original_scheme(request: Request) -> str:
+    """
+    Determines the original scheme (http or https) of the request.
+
+    Args:
+        request (Request): The incoming HTTP request.
+
+    Returns:
+        str: The original scheme ('http' or 'https')
+    """
+    # Check the X-Forwarded-Proto header first
+    forwarded_proto = request.headers.get("X-Forwarded-Proto")
+    if forwarded_proto:
+        return forwarded_proto
+
+    # Check if the request is secure
+    if request.url.scheme == "https" or request.headers.get("X-Forwarded-Ssl") == "on":
+        return "https"
+
+    # Check for other common headers that might indicate HTTPS
+    if (
+        request.headers.get("X-Forwarded-Ssl") == "on"
+        or request.headers.get("X-Forwarded-Protocol") == "https"
+        or request.headers.get("X-Url-Scheme") == "https"
+    ):
+        return "https"
+
+    # Default to http if no indicators of https are found
+    return "http"
+
+
+def get_proxy_headers(request: Request) -> dict:
+    """
+    Extracts proxy headers from the request query parameters.
+
+    Args:
+        request (Request): The incoming HTTP request.
+
+    Returns:
+        dict: A dictionary of proxy headers.
+    """
+    request_headers = {k: v for k, v in request.headers.items() if k in SUPPORTED_REQUEST_HEADERS}
+    request_headers.update({k[2:].lower(): v for k, v in request.query_params.items() if k.startswith("h_")})
+    return request_headers
+
+
+class EnhancedStreamingResponse(Response):
+    body_iterator: typing.AsyncIterable[typing.Any]
+
+    def __init__(
+        self,
+        content: typing.Union[typing.AsyncIterable[typing.Any], typing.Iterable[typing.Any]],
+        status_code: int = 200,
+        headers: typing.Optional[typing.Mapping[str, str]] = None,
+        media_type: typing.Optional[str] = None,
+        background: typing.Optional[BackgroundTask] = None,
+    ) -> None:
+        if isinstance(content, typing.AsyncIterable):
+            self.body_iterator = content
+        else:
+            self.body_iterator = iterate_in_threadpool(content)
+        self.status_code = status_code
+        self.media_type = self.media_type if media_type is None else media_type
+        self.background = background
+        self.init_headers(headers)
+
+    @staticmethod
+    async def listen_for_disconnect(receive: Receive) -> None:
+        try:
+            while True:
+                message = await receive()
+                if message["type"] == "http.disconnect":
+                    logger.debug("Client disconnected")
+                    break
+        except Exception as e:
+            logger.error(f"Error in listen_for_disconnect: {str(e)}")
+
+    async def stream_response(self, send: Send) -> None:
+        try:
+            await send(
+                {
+                    "type": "http.response.start",
+                    "status": self.status_code,
+                    "headers": self.raw_headers,
+                }
+            )
+            async for chunk in self.body_iterator:
+                if not isinstance(chunk, (bytes, memoryview)):
+                    chunk = chunk.encode(self.charset)
+                try:
+                    await send({"type": "http.response.body", "body": chunk, "more_body": True})
+                except (ConnectionResetError, anyio.BrokenResourceError):
+                    logger.info("Client disconnected during streaming")
+                    return
+
+            await send({"type": "http.response.body", "body": b"", "more_body": False})
+        except Exception as e:
+            logger.error(f"Error in stream_response: {str(e)}")
+
+    async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
+        async with anyio.create_task_group() as task_group:
+
+            async def wrap(func: typing.Callable[[], typing.Awaitable[None]]) -> None:
+                try:
+                    await func()
+                except ExceptionGroup as e:
+                    if not any(isinstance(exc, anyio.get_cancelled_exc_class()) for exc in e.exceptions):
+                        logger.exception("Error in streaming task")
+                    raise
+                except Exception as e:
+                    if not isinstance(e, anyio.get_cancelled_exc_class()):
+                        logger.exception("Error in streaming task")
+                    raise
+                finally:
+                    task_group.cancel_scope.cancel()
+
+            task_group.start_soon(wrap, partial(self.stream_response, send))
+            await wrap(partial(self.listen_for_disconnect, receive))
+
+        if self.background is not None:
+            await self.background()

mediaflow_proxy/utils/m3u8_processor.py

@@ -0,0 +1,83 @@
+import re
+from urllib import parse
+
+from pydantic import HttpUrl
+
+from mediaflow_proxy.utils.http_utils import encode_mediaflow_proxy_url, get_original_scheme
+
+
+class M3U8Processor:
+    def __init__(self, request, key_url: HttpUrl = None):
+        """
+        Initializes the M3U8Processor with the request and URL prefix.
+
+        Args:
+            request (Request): The incoming HTTP request.
+            key_url (HttpUrl, optional): The URL of the key server. Defaults to None.
+        """
+        self.request = request
+        self.key_url = key_url
+        self.mediaflow_proxy_url = str(request.url_for("hls_stream_proxy").replace(scheme=get_original_scheme(request)))
+
+    async def process_m3u8(self, content: str, base_url: str) -> str:
+        """
+        Processes the m3u8 content, proxying URLs and handling key lines.
+
+        Args:
+            content (str): The m3u8 content to process.
+            base_url (str): The base URL to resolve relative URLs.
+
+        Returns:
+            str: The processed m3u8 content.
+        """
+        lines = content.splitlines()
+        processed_lines = []
+        for line in lines:
+            if "URI=" in line:
+                processed_lines.append(await self.process_key_line(line, base_url))
+            elif not line.startswith("#") and line.strip():
+                processed_lines.append(await self.proxy_url(line, base_url))
+            else:
+                processed_lines.append(line)
+        return "\n".join(processed_lines)
+
+    async def process_key_line(self, line: str, base_url: str) -> str:
+        """
+        Processes a key line in the m3u8 content, proxying the URI.
+
+        Args:
+            line (str): The key line to process.
+            base_url (str): The base URL to resolve relative URLs.
+
+        Returns:
+            str: The processed key line.
+        """
+        uri_match = re.search(r'URI="([^"]+)"', line)
+        if uri_match:
+            original_uri = uri_match.group(1)
+            uri = parse.urlparse(original_uri)
+            if self.key_url:
+                uri = uri._replace(scheme=self.key_url.scheme, netloc=self.key_url.host)
+            new_uri = await self.proxy_url(uri.geturl(), base_url)
+            line = line.replace(f'URI="{original_uri}"', f'URI="{new_uri}"')
+        return line
+
+    async def proxy_url(self, url: str, base_url: str) -> str:
+        """
+        Proxies a URL, encoding it with the MediaFlow proxy URL.
+
+        Args:
+            url (str): The URL to proxy.
+            base_url (str): The base URL to resolve relative URLs.
+
+        Returns:
+            str: The proxied URL.
+        """
+        full_url = parse.urljoin(base_url, url)
+
+        return encode_mediaflow_proxy_url(
+            self.mediaflow_proxy_url,
+            "",
+            full_url,
+            query_params=dict(self.request.query_params),
+        )

mediaflow_proxy/utils/mpd_utils.py

@@ -0,0 +1,555 @@
+import logging
+import math
+import re
+from datetime import datetime, timedelta, timezone
+from typing import List, Dict
+from urllib.parse import urljoin
+
+import xmltodict
+
+logger = logging.getLogger(__name__)
+
+
+def parse_mpd(mpd_content: str | bytes) -> dict:
+    """
+    Parses the MPD content into a dictionary.
+
+    Args:
+        mpd_content (str | bytes): The MPD content to parse.
+
+    Returns:
+        dict: The parsed MPD content as a dictionary.
+    """
+    return xmltodict.parse(mpd_content)
+
+
+def parse_mpd_dict(
+    mpd_dict: dict, mpd_url: str, parse_drm: bool = True, parse_segment_profile_id: str | None = None
+) -> dict:
+    """
+    Parses the MPD dictionary and extracts relevant information.
+
+    Args:
+        mpd_dict (dict): The MPD content as a dictionary.
+        mpd_url (str): The URL of the MPD manifest.
+        parse_drm (bool, optional): Whether to parse DRM information. Defaults to True.
+        parse_segment_profile_id (str, optional): The profile ID to parse segments for. Defaults to None.
+
+    Returns:
+        dict: The parsed MPD information including profiles and DRM info.
+
+    This function processes the MPD dictionary to extract profiles, DRM information, and other relevant data.
+    It handles both live and static MPD manifests.
+    """
+    profiles = []
+    parsed_dict = {}
+    source = "/".join(mpd_url.split("/")[:-1])
+
+    is_live = mpd_dict["MPD"].get("@type", "static").lower() == "dynamic"
+    parsed_dict["isLive"] = is_live
+
+    media_presentation_duration = mpd_dict["MPD"].get("@mediaPresentationDuration")
+
+    # Parse additional MPD attributes for live streams
+    if is_live:
+        parsed_dict["minimumUpdatePeriod"] = parse_duration(mpd_dict["MPD"].get("@minimumUpdatePeriod", "PT0S"))
+        parsed_dict["timeShiftBufferDepth"] = parse_duration(mpd_dict["MPD"].get("@timeShiftBufferDepth", "PT2M"))
+        parsed_dict["availabilityStartTime"] = datetime.fromisoformat(
+            mpd_dict["MPD"]["@availabilityStartTime"].replace("Z", "+00:00")
+        )
+        parsed_dict["publishTime"] = datetime.fromisoformat(
+            mpd_dict["MPD"].get("@publishTime", "").replace("Z", "+00:00")
+        )
+
+    periods = mpd_dict["MPD"]["Period"]
+    periods = periods if isinstance(periods, list) else [periods]
+
+    for period in periods:
+        parsed_dict["PeriodStart"] = parse_duration(period.get("@start", "PT0S"))
+        for adaptation in period["AdaptationSet"]:
+            representations = adaptation["Representation"]
+            representations = representations if isinstance(representations, list) else [representations]
+
+            for representation in representations:
+                profile = parse_representation(
+                    parsed_dict,
+                    representation,
+                    adaptation,
+                    source,
+                    media_presentation_duration,
+                    parse_segment_profile_id,
+                )
+                if profile:
+                    profiles.append(profile)
+    parsed_dict["profiles"] = profiles
+
+    if parse_drm:
+        drm_info = extract_drm_info(periods, mpd_url)
+    else:
+        drm_info = {}
+    parsed_dict["drmInfo"] = drm_info
+
+    return parsed_dict
+
+
+def pad_base64(encoded_key_id):
+    """
+    Pads a base64 encoded key ID to make its length a multiple of 4.
+
+    Args:
+        encoded_key_id (str): The base64 encoded key ID.
+
+    Returns:
+        str: The padded base64 encoded key ID.
+    """
+    return encoded_key_id + "=" * (4 - len(encoded_key_id) % 4)
+
+
+def extract_drm_info(periods: List[Dict], mpd_url: str) -> Dict:
+    """
+    Extracts DRM information from the MPD periods.
+
+    Args:
+        periods (List[Dict]): The list of periods in the MPD.
+        mpd_url (str): The URL of the MPD manifest.
+
+    Returns:
+        Dict: The extracted DRM information.
+
+    This function processes the ContentProtection elements in the MPD to extract DRM system information,
+    such as ClearKey, Widevine, and PlayReady.
+    """
+    drm_info = {"isDrmProtected": False}
+
+    for period in periods:
+        adaptation_sets: list[dict] | dict = period.get("AdaptationSet", [])
+        if not isinstance(adaptation_sets, list):
+            adaptation_sets = [adaptation_sets]
+
+        for adaptation_set in adaptation_sets:
+            # Check ContentProtection in AdaptationSet
+            process_content_protection(adaptation_set.get("ContentProtection", []), drm_info)
+
+            # Check ContentProtection inside each Representation
+            representations: list[dict] | dict = adaptation_set.get("Representation", [])
+            if not isinstance(representations, list):
+                representations = [representations]
+
+            for representation in representations:
+                process_content_protection(representation.get("ContentProtection", []), drm_info)
+
+    # If we have a license acquisition URL, make sure it's absolute
+    if "laUrl" in drm_info and not drm_info["laUrl"].startswith(("http://", "https://")):
+        drm_info["laUrl"] = urljoin(mpd_url, drm_info["laUrl"])
+
+    return drm_info
+
+
+def process_content_protection(content_protection: list[dict] | dict, drm_info: dict):
+    """
+    Processes the ContentProtection elements to extract DRM information.
+
+    Args:
+        content_protection (list[dict] | dict): The ContentProtection elements.
+        drm_info (dict): The dictionary to store DRM information.
+
+    This function updates the drm_info dictionary with DRM system information found in the ContentProtection elements.
+    """
+    if not isinstance(content_protection, list):
+        content_protection = [content_protection]
+
+    for protection in content_protection:
+        drm_info["isDrmProtected"] = True
+        scheme_id_uri = protection.get("@schemeIdUri", "").lower()
+
+        if "clearkey" in scheme_id_uri:
+            drm_info["drmSystem"] = "clearkey"
+            if "clearkey:Laurl" in protection:
+                la_url = protection["clearkey:Laurl"].get("#text")
+                if la_url and "laUrl" not in drm_info:
+                    drm_info["laUrl"] = la_url
+
+        elif "widevine" in scheme_id_uri or "edef8ba9-79d6-4ace-a3c8-27dcd51d21ed" in scheme_id_uri:
+            drm_info["drmSystem"] = "widevine"
+            pssh = protection.get("cenc:pssh", {}).get("#text")
+            if pssh:
+                drm_info["pssh"] = pssh
+
+        elif "playready" in scheme_id_uri or "9a04f079-9840-4286-ab92-e65be0885f95" in scheme_id_uri:
+            drm_info["drmSystem"] = "playready"
+
+        if "@cenc:default_KID" in protection:
+            key_id = protection["@cenc:default_KID"].replace("-", "")
+            if "keyId" not in drm_info:
+                drm_info["keyId"] = key_id
+
+        if "ms:laurl" in protection:
+            la_url = protection["ms:laurl"].get("@licenseUrl")
+            if la_url and "laUrl" not in drm_info:
+                drm_info["laUrl"] = la_url
+
+    return drm_info
+
+
+def parse_representation(
+    parsed_dict: dict,
+    representation: dict,
+    adaptation: dict,
+    source: str,
+    media_presentation_duration: str,
+    parse_segment_profile_id: str | None,
+) -> dict | None:
+    """
+    Parses a representation and extracts profile information.
+
+    Args:
+        parsed_dict (dict): The parsed MPD data.
+        representation (dict): The representation data.
+        adaptation (dict): The adaptation set data.
+        source (str): The source URL.
+        media_presentation_duration (str): The media presentation duration.
+        parse_segment_profile_id (str, optional): The profile ID to parse segments for. Defaults to None.
+
+    Returns:
+        dict | None: The parsed profile information or None if not applicable.
+    """
+    mime_type = _get_key(adaptation, representation, "@mimeType") or (
+        "video/mp4" if "avc" in representation["@codecs"] else "audio/mp4"
+    )
+    if "video" not in mime_type and "audio" not in mime_type:
+        return None
+
+    profile = {
+        "id": representation.get("@id") or adaptation.get("@id"),
+        "mimeType": mime_type,
+        "lang": representation.get("@lang") or adaptation.get("@lang"),
+        "codecs": representation.get("@codecs") or adaptation.get("@codecs"),
+        "bandwidth": int(representation.get("@bandwidth") or adaptation.get("@bandwidth")),
+        "startWithSAP": (_get_key(adaptation, representation, "@startWithSAP") or "1") == "1",
+        "mediaPresentationDuration": media_presentation_duration,
+    }
+
+    if "audio" in profile["mimeType"]:
+        profile["audioSamplingRate"] = representation.get("@audioSamplingRate") or adaptation.get("@audioSamplingRate")
+        profile["channels"] = representation.get("AudioChannelConfiguration", {}).get("@value", "2")
+    else:
+        profile["width"] = int(representation["@width"])
+        profile["height"] = int(representation["@height"])
+        frame_rate = representation.get("@frameRate") or adaptation.get("@maxFrameRate") or "30000/1001"
+        frame_rate = frame_rate if "/" in frame_rate else f"{frame_rate}/1"
+        profile["frameRate"] = round(int(frame_rate.split("/")[0]) / int(frame_rate.split("/")[1]), 3)
+        profile["sar"] = representation.get("@sar", "1:1")
+
+    if parse_segment_profile_id is None or profile["id"] != parse_segment_profile_id:
+        return profile
+
+    item = adaptation.get("SegmentTemplate") or representation.get("SegmentTemplate")
+    if item:
+        profile["segments"] = parse_segment_template(parsed_dict, item, profile, source)
+    else:
+        profile["segments"] = parse_segment_base(representation, source)
+
+    return profile
+
+
+def _get_key(adaptation: dict, representation: dict, key: str) -> str | None:
+    """
+    Retrieves a key from the representation or adaptation set.
+
+    Args:
+        adaptation (dict): The adaptation set data.
+        representation (dict): The representation data.
+        key (str): The key to retrieve.
+
+    Returns:
+        str | None: The value of the key or None if not found.
+    """
+    return representation.get(key, adaptation.get(key, None))
+
+
+def parse_segment_template(parsed_dict: dict, item: dict, profile: dict, source: str) -> List[Dict]:
+    """
+    Parses a segment template and extracts segment information.
+
+    Args:
+        parsed_dict (dict): The parsed MPD data.
+        item (dict): The segment template data.
+        profile (dict): The profile information.
+        source (str): The source URL.
+
+    Returns:
+        List[Dict]: The list of parsed segments.
+    """
+    segments = []
+    timescale = int(item.get("@timescale", 1))
+
+    # Initialization
+    if "@initialization" in item:
+        media = item["@initialization"]
+        media = media.replace("$RepresentationID$", profile["id"])
+        media = media.replace("$Bandwidth$", str(profile["bandwidth"]))
+        if not media.startswith("http"):
+            media = f"{source}/{media}"
+        profile["initUrl"] = media
+
+    # Segments
+    if "SegmentTimeline" in item:
+        segments.extend(parse_segment_timeline(parsed_dict, item, profile, source, timescale))
+    elif "@duration" in item:
+        segments.extend(parse_segment_duration(parsed_dict, item, profile, source, timescale))
+
+    return segments
+
+
+def parse_segment_timeline(parsed_dict: dict, item: dict, profile: dict, source: str, timescale: int) -> List[Dict]:
+    """
+    Parses a segment timeline and extracts segment information.
+
+    Args:
+        parsed_dict (dict): The parsed MPD data.
+        item (dict): The segment timeline data.
+        profile (dict): The profile information.
+        source (str): The source URL.
+        timescale (int): The timescale for the segments.
+
+    Returns:
+        List[Dict]: The list of parsed segments.
+    """
+    timelines = item["SegmentTimeline"]["S"]
+    timelines = timelines if isinstance(timelines, list) else [timelines]
+    period_start = parsed_dict["availabilityStartTime"] + timedelta(seconds=parsed_dict.get("PeriodStart", 0))
+    presentation_time_offset = int(item.get("@presentationTimeOffset", 0))
+    start_number = int(item.get("@startNumber", 1))
+
+    segments = [
+        create_segment_data(timeline, item, profile, source, timescale)
+        for timeline in preprocess_timeline(timelines, start_number, period_start, presentation_time_offset, timescale)
+    ]
+    return segments
+
+
+def preprocess_timeline(
+    timelines: List[Dict], start_number: int, period_start: datetime, presentation_time_offset: int, timescale: int
+) -> List[Dict]:
+    """
+    Preprocesses the segment timeline data.
+
+    Args:
+        timelines (List[Dict]): The list of timeline segments.
+        start_number (int): The starting segment number.
+        period_start (datetime): The start time of the period.
+        presentation_time_offset (int): The presentation time offset.
+        timescale (int): The timescale for the segments.
+
+    Returns:
+        List[Dict]: The list of preprocessed timeline segments.
+    """
+    processed_data = []
+    current_time = 0
+    for timeline in timelines:
+        repeat = int(timeline.get("@r", 0))
+        duration = int(timeline["@d"])
+        start_time = int(timeline.get("@t", current_time))
+
+        for _ in range(repeat + 1):
+            segment_start_time = period_start + timedelta(seconds=(start_time - presentation_time_offset) / timescale)
+            segment_end_time = segment_start_time + timedelta(seconds=duration / timescale)
+            processed_data.append(
+                {
+                    "number": start_number,
+                    "start_time": segment_start_time,
+                    "end_time": segment_end_time,
+                    "duration": duration,
+                    "time": start_time,
+                }
+            )
+            start_time += duration
+            start_number += 1
+
+        current_time = start_time
+
+    return processed_data
+
+
+def parse_segment_duration(parsed_dict: dict, item: dict, profile: dict, source: str, timescale: int) -> List[Dict]:
+    """
+    Parses segment duration and extracts segment information.
+    This is used for static or live MPD manifests.
+
+    Args:
+        parsed_dict (dict): The parsed MPD data.
+        item (dict): The segment duration data.
+        profile (dict): The profile information.
+        source (str): The source URL.
+        timescale (int): The timescale for the segments.
+
+    Returns:
+        List[Dict]: The list of parsed segments.
+    """
+    duration = int(item["@duration"])
+    start_number = int(item.get("@startNumber", 1))
+    segment_duration_sec = duration / timescale
+
+    if parsed_dict["isLive"]:
+        segments = generate_live_segments(parsed_dict, segment_duration_sec, start_number)
+    else:
+        segments = generate_vod_segments(profile, duration, timescale, start_number)
+
+    return [create_segment_data(seg, item, profile, source, timescale) for seg in segments]
+
+
+def generate_live_segments(parsed_dict: dict, segment_duration_sec: float, start_number: int) -> List[Dict]:
+    """
+    Generates live segments based on the segment duration and start number.
+    This is used for live MPD manifests.
+
+    Args:
+        parsed_dict (dict): The parsed MPD data.
+        segment_duration_sec (float): The segment duration in seconds.
+        start_number (int): The starting segment number.
+
+    Returns:
+        List[Dict]: The list of generated live segments.
+    """
+    time_shift_buffer_depth = timedelta(seconds=parsed_dict.get("timeShiftBufferDepth", 60))
+    segment_count = math.ceil(time_shift_buffer_depth.total_seconds() / segment_duration_sec)
+    current_time = datetime.now(tz=timezone.utc)
+    earliest_segment_number = max(
+        start_number
+        + math.floor((current_time - parsed_dict["availabilityStartTime"]).total_seconds() / segment_duration_sec)
+        - segment_count,
+        start_number,
+    )
+
+    return [
+        {
+            "number": number,
+            "start_time": parsed_dict["availabilityStartTime"]
+            + timedelta(seconds=(number - start_number) * segment_duration_sec),
+            "duration": segment_duration_sec,
+        }
+        for number in range(earliest_segment_number, earliest_segment_number + segment_count)
+    ]
+
+
+def generate_vod_segments(profile: dict, duration: int, timescale: int, start_number: int) -> List[Dict]:
+    """
+    Generates VOD segments based on the segment duration and start number.
+    This is used for static MPD manifests.
+
+    Args:
+        profile (dict): The profile information.
+        duration (int): The segment duration.
+        timescale (int): The timescale for the segments.
+        start_number (int): The starting segment number.
+
+    Returns:
+        List[Dict]: The list of generated VOD segments.
+    """
+    total_duration = profile.get("mediaPresentationDuration") or 0
+    if isinstance(total_duration, str):
+        total_duration = parse_duration(total_duration)
+    segment_count = math.ceil(total_duration * timescale / duration)
+
+    return [{"number": start_number + i, "duration": duration / timescale} for i in range(segment_count)]
+
+
+def create_segment_data(segment: Dict, item: dict, profile: dict, source: str, timescale: int | None = None) -> Dict:
+    """
+    Creates segment data based on the segment information. This includes the segment URL and metadata.
+
+    Args:
+        segment (Dict): The segment information.
+        item (dict): The segment template data.
+        profile (dict): The profile information.
+        source (str): The source URL.
+        timescale (int, optional): The timescale for the segments. Defaults to None.
+
+    Returns:
+        Dict: The created segment data.
+    """
+    media_template = item["@media"]
+    media = media_template.replace("$RepresentationID$", profile["id"])
+    media = media.replace("$Number%04d$", f"{segment['number']:04d}")
+    media = media.replace("$Number$", str(segment["number"]))
+    media = media.replace("$Bandwidth$", str(profile["bandwidth"]))
+
+    if "time" in segment and timescale is not None:
+        media = media.replace("$Time$", str(int(segment["time"] * timescale)))
+
+    if not media.startswith("http"):
+        media = f"{source}/{media}"
+
+    segment_data = {
+        "type": "segment",
+        "media": media,
+        "number": segment["number"],
+    }
+
+    if "start_time" in segment and "end_time" in segment:
+        segment_data.update(
+            {
+                "start_time": segment["start_time"],
+                "end_time": segment["end_time"],
+                "extinf": (segment["end_time"] - segment["start_time"]).total_seconds(),
+                "program_date_time": segment["start_time"].isoformat() + "Z",
+            }
+        )
+    elif "start_time" in segment and "duration" in segment:
+        duration = segment["duration"]
+        segment_data.update(
+            {
+                "start_time": segment["start_time"],
+                "end_time": segment["start_time"] + timedelta(seconds=duration),
+                "extinf": duration,
+                "program_date_time": segment["start_time"].isoformat() + "Z",
+            }
+        )
+    elif "duration" in segment:
+        segment_data["extinf"] = segment["duration"]
+
+    return segment_data
+
+
+def parse_segment_base(representation: dict, source: str) -> List[Dict]:
+    """
+    Parses segment base information and extracts segment data. This is used for single-segment representations.
+
+    Args:
+        representation (dict): The representation data.
+        source (str): The source URL.
+
+    Returns:
+        List[Dict]: The list of parsed segments.
+    """
+    segment = representation["SegmentBase"]
+    start, end = map(int, segment["@indexRange"].split("-"))
+    if "Initialization" in segment:
+        start, _ = map(int, segment["Initialization"]["@range"].split("-"))
+
+    return [
+        {
+            "type": "segment",
+            "range": f"{start}-{end}",
+            "media": f"{source}/{representation['BaseURL']}",
+        }
+    ]
+
+
+def parse_duration(duration_str: str) -> float:
+    """
+    Parses a duration ISO 8601 string into seconds.
+
+    Args:
+        duration_str (str): The duration string to parse.
+
+    Returns:
+        float: The parsed duration in seconds.
+    """
+    pattern = re.compile(r"P(?:(\d+)Y)?(?:(\d+)M)?(?:(\d+)D)?T?(?:(\d+)H)?(?:(\d+)M)?(?:(\d+(?:\.\d+)?)S)?")
+    match = pattern.match(duration_str)
+    if not match:
+        raise ValueError(f"Invalid duration format: {duration_str}")
+
+    years, months, days, hours, minutes, seconds = [float(g) if g else 0 for g in match.groups()]
+    return years * 365 * 24 * 3600 + months * 30 * 24 * 3600 + days * 24 * 3600 + hours * 3600 + minutes * 60 + seconds