Legend017/mediaflow-proxy
0
1import base642import logging3from urllib.parse import urlparse4 5import httpx6from fastapi import Request, Response, HTTPException7from starlette.background import BackgroundTask8 9from .configs import settings10from .const import SUPPORTED_RESPONSE_HEADERS11from .mpd_processor import process_manifest, process_playlist, process_segment12from .schemas import HLSManifestParams, ProxyStreamParams, MPDManifestParams, MPDPlaylistParams, MPDSegmentParams13from .utils.cache_utils import get_cached_mpd, get_cached_init_segment14from .utils.http_utils import (15 Streamer,16 DownloadError,17 download_file_with_retry,18 request_with_retry,19 EnhancedStreamingResponse,20 ProxyRequestHeaders,21)22from .utils.m3u8_processor import M3U8Processor23from .utils.mpd_utils import pad_base6424 25logger = logging.getLogger(__name__)26 27 28async def setup_client_and_streamer(use_request_proxy: bool, verify_ssl: bool) -> tuple[httpx.AsyncClient, Streamer]:29 """30 Set up an HTTP client and a streamer.31 32 Args:33 use_request_proxy (bool): Whether to use a proxy for the request.34 verify_ssl (bool): Whether to verify SSL certificates.35 36 Returns:37 tuple: An httpx.AsyncClient instance and a Streamer instance.38 """39 client = httpx.AsyncClient(40 follow_redirects=True,41 timeout=httpx.Timeout(30.0),42 limits=httpx.Limits(max_keepalive_connections=10, max_connections=20),43 proxy=settings.proxy_url if use_request_proxy else None,44 verify=verify_ssl,45 )46 return client, Streamer(client)47 48 49def handle_exceptions(exception: Exception) -> Response:50 """51 Handle exceptions and return appropriate HTTP responses.52 53 Args:54 exception (Exception): The exception that was raised.55 56 Returns:57 Response: An HTTP response corresponding to the exception type.58 """59 if isinstance(exception, httpx.HTTPStatusError):60 logger.error(f"Upstream service error while handling request: {exception}")61 return Response(status_code=exception.response.status_code, content=f"Upstream service error: {exception}")62 elif isinstance(exception, DownloadError):63 logger.error(f"Error downloading content: {exception}")64 return Response(status_code=exception.status_code, content=str(exception))65 else:66 logger.exception(f"Internal server error while handling request: {exception}")67 return Response(status_code=502, content=f"Internal server error: {exception}")68 69 70async def handle_hls_stream_proxy(71 request: Request, hls_params: HLSManifestParams, proxy_headers: ProxyRequestHeaders72) -> Response:73 """74 Handle HLS stream proxy requests.75 76 This function processes HLS manifest files and streams content based on the request parameters.77 78 Args:79 request (Request): The incoming FastAPI request object.80 hls_params (HLSManifestParams): Parameters for the HLS manifest.81 proxy_headers (ProxyRequestHeaders): Headers to be used in the proxy request.82 83 Returns:84 Union[Response, EnhancedStreamingResponse]: Either a processed m3u8 playlist or a streaming response.85 """86 client, streamer = await setup_client_and_streamer(hls_params.use_request_proxy, hls_params.verify_ssl)87 88 try:89 if urlparse(hls_params.destination).path.endswith((".m3u", ".m3u8")):90 return await fetch_and_process_m3u8(91 streamer, hls_params.destination, proxy_headers, request, hls_params.key_url92 )93 94 response = await streamer.head(hls_params.destination, proxy_headers.request)95 if "mpegurl" in response.headers.get("content-type", "").lower():96 return await fetch_and_process_m3u8(97 streamer, hls_params.destination, proxy_headers, request, hls_params.key_url98 )99 100 content_range = proxy_headers.request.get("range", "bytes=0-")101 if "NaN" in content_range:102 # Handle invalid range requests "bytes=NaN-NaN"103 raise HTTPException(status_code=416, detail="Invalid Range Header")104 proxy_headers.request.update({"range": content_range})105 response_headers = prepare_response_headers(response.headers, proxy_headers.response)106 107 return EnhancedStreamingResponse(108 streamer.stream_content(hls_params.destination, proxy_headers.request),109 status_code=response.status_code,110 headers=response_headers,111 background=BackgroundTask(streamer.close),112 )113 except Exception as e:114 await client.aclose()115 return handle_exceptions(e)116 117 118async def handle_stream_request(119 method: str,120 video_url: str,121 proxy_headers: ProxyRequestHeaders,122 verify_ssl: bool = True,123 use_request_proxy: bool = True,124) -> Response:125 """126 Handle general stream requests.127 128 This function processes both HEAD and GET requests for video streams.129 130 Args:131 method (str): The HTTP method (e.g., 'GET' or 'HEAD').132 video_url (str): The URL of the video to stream.133 proxy_headers (ProxyRequestHeaders): Headers to be used in the proxy request.134 verify_ssl (bool, optional): Whether to verify SSL certificates. Defaults to True.135 use_request_proxy (bool, optional): Whether to use a proxy for the request. Defaults to True.136 137 Returns:138 Union[Response, EnhancedStreamingResponse]: Either a HEAD response or a streaming response.139 """140 client, streamer = await setup_client_and_streamer(use_request_proxy, verify_ssl)141 142 try:143 response = await streamer.head(video_url, proxy_headers.request)144 response_headers = prepare_response_headers(response.headers, proxy_headers.response)145 146 if method == "HEAD":147 await streamer.close()148 return Response(headers=response_headers, status_code=response.status_code)149 else:150 return EnhancedStreamingResponse(151 streamer.stream_content(video_url, proxy_headers.request),152 headers=response_headers,153 status_code=response.status_code,154 background=BackgroundTask(streamer.close),155 )156 except Exception as e:157 await client.aclose()158 return handle_exceptions(e)159 160 161def prepare_response_headers(original_headers, proxy_response_headers) -> dict:162 """163 Prepare response headers for the proxy response.164 165 This function filters the original headers, ensures proper transfer encoding,166 and merges them with the proxy response headers.167 168 Args:169 original_headers (httpx.Headers): The original headers from the upstream response.170 proxy_response_headers (dict): Additional headers to be included in the proxy response.171 172 Returns:173 dict: The prepared headers for the proxy response.174 """175 response_headers = {k: v for k, v in original_headers.multi_items() if k in SUPPORTED_RESPONSE_HEADERS}176 response_headers.update(proxy_response_headers)177 return response_headers178 179 180async def proxy_stream(method: str, stream_params: ProxyStreamParams, proxy_headers: ProxyRequestHeaders):181 """182 Proxies the stream request to the given video URL.183 184 Args:185 method (str): The HTTP method (e.g., GET, HEAD).186 stream_params (ProxyStreamParams): The parameters for the stream request.187 proxy_headers (ProxyRequestHeaders): The headers to include in the request.188 189 Returns:190 Response: The HTTP response with the streamed content.191 """192 return await handle_stream_request(193 method, stream_params.destination, proxy_headers, stream_params.verify_ssl, stream_params.use_request_proxy194 )195 196 197async def fetch_and_process_m3u8(198 streamer: Streamer, url: str, proxy_headers: ProxyRequestHeaders, request: Request, key_url: str = None199):200 """201 Fetches and processes the m3u8 playlist, converting it to an HLS playlist.202 203 Args:204 streamer (Streamer): The HTTP client to use for streaming.205 url (str): The URL of the m3u8 playlist.206 proxy_headers (ProxyRequestHeaders): The headers to include in the request.207 request (Request): The incoming HTTP request.208 key_url (str, optional): The HLS Key URL to replace the original key URL. Defaults to None.209 210 Returns:211 Response: The HTTP response with the processed m3u8 playlist.212 """213 try:214 content = await streamer.get_text(url, proxy_headers.request)215 processor = M3U8Processor(request, key_url)216 processed_content = await processor.process_m3u8(content, str(streamer.response.url))217 response_headers = {"Content-Disposition": "inline", "Accept-Ranges": "none"}218 response_headers.update(proxy_headers.response)219 return Response(220 content=processed_content,221 media_type="application/vnd.apple.mpegurl",222 headers=response_headers,223 )224 except Exception as e:225 return handle_exceptions(e)226 finally:227 await streamer.close()228 229 230async def handle_drm_key_data(key_id, key, drm_info):231 """232 Handles the DRM key data, retrieving the key ID and key from the DRM info if not provided.233 234 Args:235 key_id (str): The DRM key ID.236 key (str): The DRM key.237 drm_info (dict): The DRM information from the MPD manifest.238 239 Returns:240 tuple: The key ID and key.241 """242 if drm_info and not drm_info.get("isDrmProtected"):243 return None, None244 245 if not key_id or not key:246 if "keyId" in drm_info and "key" in drm_info:247 key_id = drm_info["keyId"]248 key = drm_info["key"]249 elif "laUrl" in drm_info and "keyId" in drm_info:250 raise HTTPException(status_code=400, detail="LA URL is not supported yet")251 else:252 raise HTTPException(253 status_code=400, detail="Unable to determine key_id and key, and they were not provided"254 )255 256 return key_id, key257 258 259async def get_manifest(260 request: Request,261 manifest_params: MPDManifestParams,262 proxy_headers: ProxyRequestHeaders,263):264 """265 Retrieves and processes the MPD manifest, converting it to an HLS manifest.266 267 Args:268 request (Request): The incoming HTTP request.269 manifest_params (MPDManifestParams): The parameters for the manifest request.270 proxy_headers (ProxyRequestHeaders): The headers to include in the request.271 272 Returns:273 Response: The HTTP response with the HLS manifest.274 """275 try:276 mpd_dict = await get_cached_mpd(277 manifest_params.destination,278 headers=proxy_headers.request,279 parse_drm=not manifest_params.key_id and not manifest_params.key,280 verify_ssl=manifest_params.verify_ssl,281 use_request_proxy=manifest_params.use_request_proxy,282 )283 except DownloadError as e:284 raise HTTPException(status_code=e.status_code, detail=f"Failed to download MPD: {e.message}")285 drm_info = mpd_dict.get("drmInfo", {})286 287 if drm_info and not drm_info.get("isDrmProtected"):288 # For non-DRM protected MPD, we still create an HLS manifest289 return await process_manifest(request, mpd_dict, proxy_headers, None, None)290 291 key_id, key = await handle_drm_key_data(manifest_params.key_id, manifest_params.key, drm_info)292 293 # check if the provided key_id and key are valid294 if key_id and len(key_id) != 32:295 key_id = base64.urlsafe_b64decode(pad_base64(key_id)).hex()296 if key and len(key) != 32:297 key = base64.urlsafe_b64decode(pad_base64(key)).hex()298 299 return await process_manifest(request, mpd_dict, proxy_headers, key_id, key)300 301 302async def get_playlist(303 request: Request,304 playlist_params: MPDPlaylistParams,305 proxy_headers: ProxyRequestHeaders,306):307 """308 Retrieves and processes the MPD manifest, converting it to an HLS playlist for a specific profile.309 310 Args:311 request (Request): The incoming HTTP request.312 playlist_params (MPDPlaylistParams): The parameters for the playlist request.313 proxy_headers (ProxyRequestHeaders): The headers to include in the request.314 315 Returns:316 Response: The HTTP response with the HLS playlist.317 """318 mpd_dict = await get_cached_mpd(319 playlist_params.destination,320 headers=proxy_headers.request,321 parse_drm=not playlist_params.key_id and not playlist_params.key,322 parse_segment_profile_id=playlist_params.profile_id,323 verify_ssl=playlist_params.verify_ssl,324 use_request_proxy=playlist_params.use_request_proxy,325 )326 return await process_playlist(request, mpd_dict, playlist_params.profile_id, proxy_headers)327 328 329async def get_segment(330 segment_params: MPDSegmentParams,331 proxy_headers: ProxyRequestHeaders,332):333 """334 Retrieves and processes a media segment, decrypting it if necessary.335 336 Args:337 segment_params (MPDSegmentParams): The parameters for the segment request.338 proxy_headers (ProxyRequestHeaders): The headers to include in the request.339 340 Returns:341 Response: The HTTP response with the processed segment.342 """343 try:344 init_content = await get_cached_init_segment(345 segment_params.init_url, proxy_headers.request, segment_params.verify_ssl, segment_params.use_request_proxy346 )347 segment_content = await download_file_with_retry(348 segment_params.segment_url,349 proxy_headers.request,350 verify_ssl=segment_params.verify_ssl,351 use_request_proxy=segment_params.use_request_proxy,352 )353 except Exception as e:354 return handle_exceptions(e)355 356 return await process_segment(357 init_content,358 segment_content,359 segment_params.mime_type,360 proxy_headers,361 segment_params.key_id,362 segment_params.key,363 )364 365 366async def get_public_ip(use_request_proxy: bool = True):367 """368 Retrieves the public IP address of the MediaFlow proxy.369 370 Args:371 use_request_proxy (bool, optional): Whether to use the proxy configuration from the user's MediaFlow config. Defaults to True.372 373 Returns:374 Response: The HTTP response with the public IP address.375 """376 ip_address_data = await request_with_retry(377 "GET", "https://api.ipify.org?format=json", {}, use_request_proxy=use_request_proxy378 )379 return ip_address_data.json()380