Legend017/mediaflow-proxy
0
1import logging2import typing3from dataclasses import dataclass4from functools import partial5from urllib import parse6from urllib.parse import urlencode7 8import anyio9import httpx10import tenacity11from fastapi import Response12from starlette.background import BackgroundTask13from starlette.concurrency import iterate_in_threadpool14from starlette.requests import Request15from starlette.types import Receive, Send, Scope16from tenacity import retry, stop_after_attempt, wait_exponential, retry_if_exception_type17from tqdm.asyncio import tqdm as tqdm_asyncio18 19from mediaflow_proxy.configs import settings20from mediaflow_proxy.const import SUPPORTED_REQUEST_HEADERS21from mediaflow_proxy.utils.crypto_utils import EncryptionHandler22 23logger = logging.getLogger(__name__)24 25 26class DownloadError(Exception):27 def __init__(self, status_code, message):28 self.status_code = status_code29 self.message = message30 super().__init__(message)31 32 33@retry(34 stop=stop_after_attempt(3),35 wait=wait_exponential(multiplier=1, min=4, max=10),36 retry=retry_if_exception_type(DownloadError),37)38async def fetch_with_retry(client, method, url, headers, follow_redirects=True, **kwargs):39 """40 Fetches a URL with retry logic.41 42 Args:43 client (httpx.AsyncClient): The HTTP client to use for the request.44 method (str): The HTTP method to use (e.g., GET, POST).45 url (str): The URL to fetch.46 headers (dict): The headers to include in the request.47 follow_redirects (bool, optional): Whether to follow redirects. Defaults to True.48 **kwargs: Additional arguments to pass to the request.49 50 Returns:51 httpx.Response: The HTTP response.52 53 Raises:54 DownloadError: If the request fails after retries.55 """56 try:57 response = await client.request(method, url, headers=headers, follow_redirects=follow_redirects, **kwargs)58 response.raise_for_status()59 return response60 except httpx.TimeoutException:61 logger.warning(f"Timeout while downloading {url}")62 raise DownloadError(409, f"Timeout while downloading {url}")63 except httpx.HTTPStatusError as e:64 logger.error(f"HTTP error {e.response.status_code} while downloading {url}")65 # if e.response.status_code == 404:66 # logger.error(f"Segment Resource not found: {url}")67 # raise e68 raise DownloadError(e.response.status_code, f"HTTP error {e.response.status_code} while downloading {url}")69 except Exception as e:70 logger.error(f"Error downloading {url}: {e}")71 raise72 73 74class Streamer:75 def __init__(self, client):76 """77 Initializes the Streamer with an HTTP client.78 79 Args:80 client (httpx.AsyncClient): The HTTP client to use for streaming.81 """82 self.client = client83 self.response = None84 self.progress_bar = None85 self.bytes_transferred = 086 self.start_byte = 087 self.end_byte = 088 self.total_size = 089 90 async def stream_content(self, url: str, headers: dict) -> typing.AsyncGenerator[bytes, None]:91 """92 Streams content from a URL.93 94 Args:95 url (str): The URL to stream content from.96 headers (dict): The headers to include in the request.97 98 Yields:99 bytes: Chunks of the streamed content.100 """101 try:102 async with self.client.stream("GET", url, headers=headers, follow_redirects=True) as self.response:103 self.response.raise_for_status()104 self.parse_content_range()105 106 if settings.enable_streaming_progress:107 with tqdm_asyncio(108 total=self.total_size,109 initial=self.start_byte,110 unit="B",111 unit_scale=True,112 unit_divisor=1024,113 desc="Streaming",114 ncols=100,115 mininterval=1,116 ) as self.progress_bar:117 async for chunk in self.response.aiter_bytes():118 yield chunk119 chunk_size = len(chunk)120 self.bytes_transferred += chunk_size121 self.progress_bar.set_postfix_str(122 f"📥 : {self.format_bytes(self.bytes_transferred)}", refresh=False123 )124 self.progress_bar.update(chunk_size)125 else:126 async for chunk in self.response.aiter_bytes():127 yield chunk128 except GeneratorExit:129 logger.info("Streaming session stopped by the user")130 except Exception as e:131 logger.error(f"Error streaming content: {e}")132 finally:133 await self.close()134 135 @staticmethod136 def format_bytes(size) -> str:137 power = 2**10138 n = 0139 units = {0: "B", 1: "KB", 2: "MB", 3: "GB", 4: "TB"}140 while size > power:141 size /= power142 n += 1143 return f"{size:.2f} {units[n]}"144 145 def parse_content_range(self):146 content_range = self.response.headers.get("Content-Range", "")147 if content_range:148 range_info = content_range.split()[-1]149 self.start_byte, self.end_byte, self.total_size = map(int, range_info.replace("/", "-").split("-"))150 else:151 self.start_byte = 0152 self.total_size = int(self.response.headers.get("Content-Length", 0))153 self.end_byte = self.total_size - 1 if self.total_size > 0 else 0154 155 async def head(self, url: str, headers: dict):156 """157 Sends a HEAD request to a URL.158 159 Args:160 url (str): The URL to send the HEAD request to.161 headers (dict): The headers to include in the request.162 163 Returns:164 httpx.Response: The HTTP response.165 """166 try:167 self.response = await fetch_with_retry(self.client, "HEAD", url, headers)168 except tenacity.RetryError as e:169 raise e.last_attempt.result()170 return self.response171 172 async def get_text(self, url: str, headers: dict):173 """174 Sends a GET request to a URL and returns the response text.175 176 Args:177 url (str): The URL to send the GET request to.178 headers (dict): The headers to include in the request.179 180 Returns:181 str: The response text.182 """183 try:184 self.response = await fetch_with_retry(self.client, "GET", url, headers)185 except tenacity.RetryError as e:186 raise e.last_attempt.result()187 return self.response.text188 189 async def close(self):190 """191 Closes the HTTP client and response.192 """193 if self.response:194 await self.response.aclose()195 if self.progress_bar:196 self.progress_bar.close()197 await self.client.aclose()198 199 200async def download_file_with_retry(201 url: str,202 headers: dict,203 timeout: float = 10.0,204 verify_ssl: bool = True,205 use_request_proxy: bool = True,206):207 """208 Downloads a file with retry logic.209 210 Args:211 url (str): The URL of the file to download.212 headers (dict): The headers to include in the request.213 timeout (float, optional): The request timeout. Defaults to 10.0.214 verify_ssl (bool, optional): Whether to verify the SSL certificate of the destination. Defaults to True.215 use_request_proxy (bool, optional): Whether to use the proxy configuration from the user's MediaFlow config. Defaults to True.216 217 Returns:218 bytes: The downloaded file content.219 220 Raises:221 DownloadError: If the download fails after retries.222 """223 async with httpx.AsyncClient(224 follow_redirects=True,225 timeout=timeout,226 proxy=settings.proxy_url if use_request_proxy else None,227 verify=verify_ssl,228 ) as client:229 try:230 response = await fetch_with_retry(client, "GET", url, headers)231 return response.content232 except DownloadError as e:233 logger.error(f"Failed to download file: {e}")234 raise e235 except tenacity.RetryError as e:236 raise DownloadError(502, f"Failed to download file: {e.last_attempt.result()}")237 238 239async def request_with_retry(240 method: str, url: str, headers: dict, timeout: float = 10.0, use_request_proxy: bool = True, **kwargs241):242 """243 Sends an HTTP request with retry logic.244 245 Args:246 method (str): The HTTP method to use (e.g., GET, POST).247 url (str): The URL to send the request to.248 headers (dict): The headers to include in the request.249 timeout (float, optional): The request timeout. Defaults to 10.0.250 use_request_proxy (bool, optional): Whether to use the proxy configuration from the user's MediaFlow config. Defaults to True.251 **kwargs: Additional arguments to pass to the request.252 253 Returns:254 httpx.Response: The HTTP response.255 256 Raises:257 DownloadError: If the request fails after retries.258 """259 async with httpx.AsyncClient(260 follow_redirects=True, timeout=timeout, proxy=settings.proxy_url if use_request_proxy else None261 ) as client:262 try:263 response = await fetch_with_retry(client, method, url, headers, **kwargs)264 return response265 except DownloadError as e:266 logger.error(f"Failed to download file: {e}")267 raise268 269 270def encode_mediaflow_proxy_url(271 mediaflow_proxy_url: str,272 endpoint: str | None = None,273 destination_url: str | None = None,274 query_params: dict | None = None,275 request_headers: dict | None = None,276 response_headers: dict | None = None,277 encryption_handler: EncryptionHandler = None,278 expiration: int = None,279 ip: str = None,280) -> str:281 """282 Encodes & Encrypt (Optional) a MediaFlow proxy URL with query parameters and headers.283 284 Args:285 mediaflow_proxy_url (str): The base MediaFlow proxy URL.286 endpoint (str, optional): The endpoint to append to the base URL. Defaults to None.287 destination_url (str, optional): The destination URL to include in the query parameters. Defaults to None.288 query_params (dict, optional): Additional query parameters to include. Defaults to None.289 request_headers (dict, optional): Headers to include as query parameters. Defaults to None.290 response_headers (dict, optional): Headers to include as query parameters. Defaults to None.291 encryption_handler (EncryptionHandler, optional): The encryption handler to use. Defaults to None.292 expiration (int, optional): The expiration time for the encrypted token. Defaults to None.293 ip (str, optional): The public IP address to include in the query parameters. Defaults to None.294 295 Returns:296 str: The encoded MediaFlow proxy URL.297 """298 query_params = query_params or {}299 if destination_url is not None:300 query_params["d"] = destination_url301 302 # Add headers if provided303 if request_headers:304 query_params.update(305 {key if key.startswith("h_") else f"h_{key}": value for key, value in request_headers.items()}306 )307 if response_headers:308 query_params.update(309 {key if key.startswith("r_") else f"r_{key}": value for key, value in response_headers.items()}310 )311 312 if encryption_handler:313 encrypted_token = encryption_handler.encrypt_data(query_params, expiration, ip)314 encoded_params = urlencode({"token": encrypted_token})315 else:316 encoded_params = urlencode(query_params)317 318 # Construct the full URL319 if endpoint is None:320 return f"{mediaflow_proxy_url}?{encoded_params}"321 322 base_url = parse.urljoin(mediaflow_proxy_url, endpoint)323 return f"{base_url}?{encoded_params}"324 325 326def get_original_scheme(request: Request) -> str:327 """328 Determines the original scheme (http or https) of the request.329 330 Args:331 request (Request): The incoming HTTP request.332 333 Returns:334 str: The original scheme ('http' or 'https')335 """336 # Check the X-Forwarded-Proto header first337 forwarded_proto = request.headers.get("X-Forwarded-Proto")338 if forwarded_proto:339 return forwarded_proto340 341 # Check if the request is secure342 if request.url.scheme == "https" or request.headers.get("X-Forwarded-Ssl") == "on":343 return "https"344 345 # Check for other common headers that might indicate HTTPS346 if (347 request.headers.get("X-Forwarded-Ssl") == "on"348 or request.headers.get("X-Forwarded-Protocol") == "https"349 or request.headers.get("X-Url-Scheme") == "https"350 ):351 return "https"352 353 # Default to http if no indicators of https are found354 return "http"355 356 357@dataclass358class ProxyRequestHeaders:359 request: dict360 response: dict361 362 363def get_proxy_headers(request: Request) -> ProxyRequestHeaders:364 """365 Extracts proxy headers from the request query parameters.366 367 Args:368 request (Request): The incoming HTTP request.369 370 Returns:371 ProxyRequest: A named tuple containing the request headers and response headers.372 """373 request_headers = {k: v for k, v in request.headers.items() if k in SUPPORTED_REQUEST_HEADERS}374 request_headers.update({k[2:].lower(): v for k, v in request.query_params.items() if k.startswith("h_")})375 response_headers = {k[2:].lower(): v for k, v in request.query_params.items() if k.startswith("r_")}376 return ProxyRequestHeaders(request_headers, response_headers)377 378 379class EnhancedStreamingResponse(Response):380 body_iterator: typing.AsyncIterable[typing.Any]381 382 def __init__(383 self,384 content: typing.Union[typing.AsyncIterable[typing.Any], typing.Iterable[typing.Any]],385 status_code: int = 200,386 headers: typing.Optional[typing.Mapping[str, str]] = None,387 media_type: typing.Optional[str] = None,388 background: typing.Optional[BackgroundTask] = None,389 ) -> None:390 if isinstance(content, typing.AsyncIterable):391 self.body_iterator = content392 else:393 self.body_iterator = iterate_in_threadpool(content)394 self.status_code = status_code395 self.media_type = self.media_type if media_type is None else media_type396 self.background = background397 self.init_headers(headers)398 399 @staticmethod400 async def listen_for_disconnect(receive: Receive) -> None:401 try:402 while True:403 message = await receive()404 if message["type"] == "http.disconnect":405 logger.debug("Client disconnected")406 break407 except Exception as e:408 logger.error(f"Error in listen_for_disconnect: {str(e)}")409 410 async def stream_response(self, send: Send) -> None:411 try:412 await send(413 {414 "type": "http.response.start",415 "status": self.status_code,416 "headers": self.raw_headers,417 }418 )419 async for chunk in self.body_iterator:420 if not isinstance(chunk, (bytes, memoryview)):421 chunk = chunk.encode(self.charset)422 try:423 await send({"type": "http.response.body", "body": chunk, "more_body": True})424 except (ConnectionResetError, anyio.BrokenResourceError):425 logger.info("Client disconnected during streaming")426 return427 428 await send({"type": "http.response.body", "body": b"", "more_body": False})429 except Exception as e:430 logger.exception(f"Error in stream_response: {str(e)}")431 432 async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:433 async with anyio.create_task_group() as task_group:434 435 async def wrap(func: typing.Callable[[], typing.Awaitable[None]]) -> None:436 try:437 await func()438 except ExceptionGroup as e:439 if not any(isinstance(exc, anyio.get_cancelled_exc_class()) for exc in e.exceptions):440 logger.exception("Error in streaming task")441 raise442 except Exception as e:443 if not isinstance(e, anyio.get_cancelled_exc_class()):444 logger.exception("Error in streaming task")445 raise446 finally:447 task_group.cancel_scope.cancel()448 449 task_group.start_soon(wrap, partial(self.stream_response, send))450 await wrap(partial(self.listen_for_disconnect, receive))451 452 if self.background is not None:453 await self.background()454