acdc137726/mediaflow-proxy

1import datetime2import logging3 4from cachetools import TTLCache5 6from .http_utils import download_file_with_retry7from .mpd_utils import parse_mpd, parse_mpd_dict8 9logger = logging.getLogger(__name__)10 11# cache dictionary12mpd_cache = TTLCache(maxsize=100, ttl=300)  # 5 minutes default TTL13init_segment_cache = TTLCache(maxsize=100, ttl=3600)  # 1 hour default TTL14 15 16async def get_cached_mpd(17    mpd_url: str,18    headers: dict,19    parse_drm: bool,20    parse_segment_profile_id: str | None = None,21    verify_ssl: bool = True,22    use_request_proxy: bool = True,23) -> dict:24    """25    Retrieves and caches the MPD manifest, parsing it if not already cached.26 27    Args:28        mpd_url (str): The URL of the MPD manifest.29        headers (dict): The headers to include in the request.30        parse_drm (bool): Whether to parse DRM information.31        parse_segment_profile_id (str, optional): The profile ID to parse segments for. Defaults to None.32        verify_ssl (bool, optional): Whether to verify the SSL certificate of the destination. Defaults to True.33        use_request_proxy (bool, optional): Whether to use the proxy configuration from the user's MediaFlow config. Defaults to True.34 35    Returns:36        dict: The parsed MPD manifest data.37    """38    current_time = datetime.datetime.now(datetime.UTC)39    if mpd_url in mpd_cache and mpd_cache[mpd_url]["expires"] > current_time:40        logger.info(f"Using cached MPD for {mpd_url}")41        return parse_mpd_dict(mpd_cache[mpd_url]["mpd"], mpd_url, parse_drm, parse_segment_profile_id)42 43    mpd_dict = parse_mpd(44        await download_file_with_retry(mpd_url, headers, verify_ssl=verify_ssl, use_request_proxy=use_request_proxy)45    )46    parsed_mpd_dict = parse_mpd_dict(mpd_dict, mpd_url, parse_drm, parse_segment_profile_id)47    current_time = datetime.datetime.now(datetime.UTC)48    expiration_time = current_time + datetime.timedelta(seconds=parsed_mpd_dict.get("minimumUpdatePeriod", 300))49    mpd_cache[mpd_url] = {"mpd": mpd_dict, "expires": expiration_time}50    return parsed_mpd_dict51 52 53async def get_cached_init_segment(54    init_url: str, headers: dict, verify_ssl: bool = True, use_request_proxy: bool = True55) -> bytes:56    """57    Retrieves and caches the initialization segment.58 59    Args:60        init_url (str): The URL of the initialization segment.61        headers (dict): The headers to include in the request.62        verify_ssl (bool, optional): Whether to verify the SSL certificate of the destination. Defaults to True.63        use_request_proxy (bool, optional): Whether to use the proxy configuration from the user's MediaFlow config. Defaults to True.64 65    Returns:66        bytes: The initialization segment content.67    """68    if init_url not in init_segment_cache:69        init_content = await download_file_with_retry(70            init_url, headers, verify_ssl=verify_ssl, use_request_proxy=use_request_proxy71        )72        init_segment_cache[init_url] = init_content73    return init_segment_cache[init_url]74