Spaces:
Running
Running
File size: 9,143 Bytes
1ebbb74 879f4ac 1ebbb74 b78171a 1ebbb74 fb6e6a7 879f4ac fb6e6a7 879f4ac b78171a 879f4ac |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 |
import logging
from urllib import parse
import httpx
import tenacity
from starlette.requests import Request
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_bytes():
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):
"""
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.
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) 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
|