import asyncio import logging import re from ipaddress import ip_address from typing import Callable, AsyncGenerator, Any, Tuple, Dict from urllib import parse from urllib.parse import urlencode, urlparse import httpx from fastapi.requests import Request from db.config import settings from db.redis_database import REDIS_ASYNC_CLIENT from db.schemas import UserData from utils import crypto class CircuitBreakerOpenException(Exception): """Custom exception to indicate the circuit breaker is open.""" pass class CircuitBreaker: """ A specialized circuit breaker implementation optimized for web scraping scenarios. It implements a more gradual recovery mechanism to handle intermittent failures and ensure stable recovery. States: - CLOSED: Normal operation, all requests allowed - OPEN: Failure threshold exceeded, no requests allowed - HALF-OPEN: Testing recovery with controlled number of requests """ def __init__( self, failure_threshold: int, # Number of failures before opening circuit recovery_timeout: int, # Time in seconds before attempting recovery half_open_attempts: int, # Number of successful attempts needed for recovery ): self.failure_threshold = failure_threshold self.recovery_timeout = recovery_timeout self.half_open_attempts = half_open_attempts self.state = "CLOSED" self.failures = 0 self.successful_attempts = 0 # Track successful attempts in HALF-OPEN self.last_failure_time = None def is_closed(self) -> bool: """ Check if requests should be allowed through. In HALF-OPEN state, allows controlled testing of the service. """ current_time = asyncio.get_event_loop().time() if self.state == "CLOSED": return True elif self.state == "HALF-OPEN": return True elif ( self.state == "OPEN" and self.last_failure_time and (current_time - self.last_failure_time) >= self.recovery_timeout ): self.state = "HALF-OPEN" self.failures = 0 self.successful_attempts = ( 0 # Reset success counter when entering HALF-OPEN ) return True return False def reset(self): """Reset the circuit breaker to its initial state""" self.state = "CLOSED" self.failures = 0 self.successful_attempts = 0 self.last_failure_time = None def record_failure(self): """ Record a failure and update circuit breaker state. In HALF-OPEN, failures don't immediately trigger OPEN state to account for intermittent failures during recovery. """ self.failures += 1 self.last_failure_time = asyncio.get_event_loop().time() self.successful_attempts = 0 # Reset success counter on any failure if self.failures >= self.failure_threshold: self.state = "OPEN" def record_success(self): """ Record a success and update circuit breaker state. In HALF-OPEN, requires multiple successive successes to close, ensuring stable recovery. """ if self.state == "HALF-OPEN": self.successful_attempts += 1 if self.successful_attempts >= self.half_open_attempts: self.reset() # Only reset after proving stability elif self.state == "CLOSED": self.failures = 0 self.successful_attempts = 0 async def call(self, func: Callable, item: Any, *args, **kwargs) -> Tuple[Any, Any]: """ Execute the given function with circuit breaker protection. Returns a tuple of (item, result/exception). """ if not self.is_closed(): return item, CircuitBreakerOpenException( f"Circuit breaker is OPEN. Failures: {self.failures}, " f"Last failure: {self.last_failure_time}" ) try: result = await func(item, *args, **kwargs) self.record_success() return item, result except Exception as e: self.record_failure() return item, e def get_status(self) -> Dict[str, Any]: """Get detailed current status of the circuit breaker""" return { "state": self.state, "failures": self.failures, "successful_attempts": self.successful_attempts, "last_failure_time": self.last_failure_time, "is_accepting_requests": self.is_closed(), } async def batch_process_with_circuit_breaker( process_func: Callable, data: list, batch_size: int, rate_limit_delay: int, cb: CircuitBreaker, max_retries: int = 5, retry_exceptions: list[type[Exception]] = (), *args, **kwargs, ) -> AsyncGenerator[Any, None]: """ Process data in batches using the circuit breaker pattern with a maximum number of retries. Yields results as they become available. """ total_retries = 0 processed_count = 0 # Ensure retry_exceptions is a tuple for the except clause retry_exceptions = tuple(retry_exceptions) + (CircuitBreakerOpenException,) for i in range(0, len(data), batch_size): batch = data[i : i + batch_size] batch_retries = 0 while batch: # Continue until all items in the batch are processed retry_batch = [] # Reset retry list for this iteration async with asyncio.TaskGroup() as tg: # Using TaskGroup to manage tasks task_data = [ tg.create_task(cb.call(process_func, item, *args, **kwargs)) for item in batch ] # Process results as soon as they complete for task in asyncio.as_completed(task_data): item, result = await task if isinstance(result, Exception): if isinstance(result, retry_exceptions): retry_batch.append(item) logging.info( f"Retryable exception occurred for item {item}: {result}" ) else: logging.exception( f"Unexpected error during batch processing {result}", exc_info=result, ) else: processed_count += 1 yield result if retry_batch: if batch_retries >= max_retries: logging.info( f"Reached maximum number of retries ({max_retries}) for this batch." ) break # Move to the next batch else: batch_retries += 1 total_retries += 1 logging.info( f"Retrying {len(retry_batch)} items due to circuit breaker. Retry attempt {batch_retries}" ) batch = retry_batch # Retry only failed items await asyncio.sleep( cb.recovery_timeout ) # Wait for breaker to close else: break # Exit loop if all items in the batch have been processed successfully # Respect the rate limit await asyncio.sleep(rate_limit_delay) logging.info(f"Processed {processed_count} items out of {len(data)} total items.") async def get_redirector_url(url: str, headers: dict) -> str | None: """ Get the final URL after following all redirects. """ try: async with httpx.AsyncClient(proxy=settings.requests_proxy_url) as client: response = await client.head(url, headers=headers, follow_redirects=True) return str(response.url) except httpx.HTTPError as e: return def get_client_ip(request: Request) -> str | None: """ Extract the client's real IP address from the request headers or fallback to the client host. """ x_forwarded_for = request.headers.get("X-Forwarded-For") if x_forwarded_for: # In some cases, this header can contain multiple IPs # separated by commas. # The first one is the original client's IP. return x_forwarded_for.split(",")[0].strip() # Fallback to X-Real-IP if X-Forwarded-For is not available x_real_ip = request.headers.get("X-Real-IP") if x_real_ip: return x_real_ip return request.client.host if request.client else "127.0.0.1" async def get_mediaflow_proxy_public_ip(mediaflow_config) -> str | None: """ Get the public IP address of the MediaFlow proxy server. """ if mediaflow_config.public_ip: return mediaflow_config.public_ip parsed_url = urlparse(mediaflow_config.proxy_url) if is_private_ip(parsed_url.netloc): # MediaFlow proxy URL is a private IP address return None cache_key = crypto.get_text_hash( f"{mediaflow_config.proxy_url}:{mediaflow_config.api_password}", full_hash=True, ) if public_ip := await REDIS_ASYNC_CLIENT.getex(cache_key, ex=300): return public_ip.decode() try: async with httpx.AsyncClient() as client: response = await client.get( parse.urljoin(mediaflow_config.proxy_url, "/proxy/ip"), params={"api_password": mediaflow_config.api_password}, timeout=10, ) response.raise_for_status() public_ip = response.json().get("ip") if public_ip: await REDIS_ASYNC_CLIENT.set(cache_key, public_ip, ex=300) return public_ip except httpx.HTTPStatusError as e: logging.error(f"HTTP error occurred: {e}") except httpx.TimeoutException as e: logging.error(f"Request timed out: {e}") except httpx.RequestError as e: logging.error(f"Request error occurred: {e}") except Exception as e: logging.error(f"An unexpected error occurred: {e}") raise Exception( f"Failed to get MediaFlow proxy public IP address. {mediaflow_config.proxy_url}" ) async def get_user_public_ip( request: Request, user_data: UserData | None = None ) -> str | None: # Check if user has mediaflow config if ( user_data and user_data.mediaflow_config and user_data.mediaflow_config.proxy_debrid_streams ): public_ip = await get_mediaflow_proxy_public_ip(user_data.mediaflow_config) if public_ip: return public_ip # Get the user's public IP address user_ip = get_client_ip(request) # check if the user's IP address is a private IP address if is_private_ip(user_ip): # Use host public IP address. return None return user_ip def get_request_namespace(request: Request) -> str: """ Extract the namespace from the request URL. """ host = request.url.hostname if "elfhosted.com" not in host: return "mediafusion" subdomain = host.split(".")[0] parts = subdomain.rsplit("-mediafusion") if len(parts) == 1: # public namespace return "mediafusion" namespace = f"tenant-{parts[0]}" return namespace def get_user_data(request: Request, secret_str: str | None = None) -> UserData: return request.user def get_secret_str(request: Request) -> str: return request.scope.get("secret_str") def encode_mediaflow_proxy_url( mediaflow_proxy_url: str, endpoint: str, destination_url: str | None = None, query_params: dict | None = None, request_headers: dict | None = None, response_headers: dict | None = None, encryption_api_password: str = None, expiration: int = None, ip: str = None, ) -> str: 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( {f"h_{key}": value for key, value in request_headers.items()} ) if response_headers: query_params.update( {f"r_{key}": value for key, value in response_headers.items()} ) if encryption_api_password: if "api_password" not in query_params: query_params["api_password"] = encryption_api_password encrypted_token = crypto.encrypt_data( encryption_api_password, query_params, expiration, ip ) encoded_params = urlencode({"token": encrypted_token}) else: encoded_params = urlencode(query_params) # Construct the full URL base_url = parse.urljoin(mediaflow_proxy_url, endpoint) return f"{base_url}?{encoded_params}" def is_private_ip(ip_str: str) -> bool: """ Check if an IP address is private, supporting both IPv4 and IPv6 formats. Handles IP addresses with optional port numbers and IPv6 brackets. Examples: - IPv4: '127.0.0.1', '127.0.0.1:8888' - IPv6: '::1', '[::1]', '[::1]:8000', 'fe80::1234' Returns: bool: True if the IP is private, False otherwise """ if not ip_str: return False # Extract IP from the input string ip_part = ip_str # Handle IPv6 with port [::1]:8000 format ipv6_port_match = re.match(r"\[(.*?)\]:(\d+)$", ip_str) if ipv6_port_match: ip_part = ipv6_port_match.group(1) else: # Handle IPv6 with brackets [::1] format if ip_str.startswith("[") and ip_str.endswith("]"): ip_part = ip_str[1:-1] # Handle IPv4 with port format elif ":" in ip_str and ip_str.count(":") == 1: ip_part = ip_str.split(":")[0] try: ip = ip_address(ip_part.strip()) return ip.is_private except ValueError: return False