diff --git a/mediaflow_proxy/extractors/base.py b/mediaflow_proxy/extractors/base.py index f7e6609..2098b4c 100644 --- a/mediaflow_proxy/extractors/base.py +++ b/mediaflow_proxy/extractors/base.py @@ -9,11 +9,12 @@ from mediaflow_proxy.configs import settings class BaseExtractor(ABC): """Base class for all URL extractors.""" - def __init__(self, proxy_enabled: bool = False): + def __init__(self, proxy_enabled: bool, request_headers: dict): self.proxy_url = settings.proxy_url if proxy_enabled else None self.base_headers = { "User-Agent": settings.user_agent, "Accept-Language": "en-US,en;q=0.5", + **request_headers, } async def _make_request( diff --git a/mediaflow_proxy/extractors/doodstream.py b/mediaflow_proxy/extractors/doodstream.py index e0019ce..08fb57f 100644 --- a/mediaflow_proxy/extractors/doodstream.py +++ b/mediaflow_proxy/extractors/doodstream.py @@ -8,8 +8,8 @@ from mediaflow_proxy.extractors.base import BaseExtractor class DoodStreamExtractor(BaseExtractor): """DoodStream URL extractor.""" - def __init__(self, proxy_enabled: bool = False): - super().__init__(proxy_enabled) + def __init__(self, proxy_enabled: bool, request_headers: dict): + super().__init__(proxy_enabled, request_headers) self.base_url = "https://d000d.com" async def extract(self, url: str) -> Tuple[str, Dict[str, str]]: diff --git a/mediaflow_proxy/extractors/factory.py b/mediaflow_proxy/extractors/factory.py index eb7a94a..46ad441 100644 --- a/mediaflow_proxy/extractors/factory.py +++ b/mediaflow_proxy/extractors/factory.py @@ -16,9 +16,9 @@ class ExtractorFactory: } @classmethod - def get_extractor(cls, host: str, proxy_enabled: bool = False) -> BaseExtractor: + def get_extractor(cls, host: str, proxy_enabled: bool, request_headers: dict) -> BaseExtractor: """Get appropriate extractor instance for the given host.""" extractor_class = cls._extractors.get(host) if not extractor_class: raise ValueError(f"Unsupported host: {host}") - return extractor_class(proxy_enabled) + return extractor_class(proxy_enabled, request_headers) diff --git a/mediaflow_proxy/routes/extractor.py b/mediaflow_proxy/routes/extractor.py index 553976d..af8b273 100644 --- a/mediaflow_proxy/routes/extractor.py +++ b/mediaflow_proxy/routes/extractor.py @@ -1,11 +1,17 @@ from typing import Annotated -from fastapi import APIRouter, Query, HTTPException +from fastapi import APIRouter, Query, HTTPException, Request, Depends from fastapi.responses import RedirectResponse from mediaflow_proxy.configs import settings from mediaflow_proxy.extractors.factory import ExtractorFactory from mediaflow_proxy.schemas import ExtractorURLParams +from mediaflow_proxy.utils.http_utils import ( + encode_mediaflow_proxy_url, + get_original_scheme, + ProxyRequestHeaders, + get_proxy_headers, +) extractor_router = APIRouter() @@ -13,15 +19,25 @@ extractor_router = APIRouter() @extractor_router.get("/video") async def extract_url( extractor_params: Annotated[ExtractorURLParams, Query()], + request: Request, + proxy_headers: Annotated[ProxyRequestHeaders, Depends(get_proxy_headers)], ): """Extract clean links from various video hosting services.""" try: - extractor = ExtractorFactory.get_extractor(extractor_params.host, extractor_params.use_request_proxy) + extractor = ExtractorFactory.get_extractor( + extractor_params.host, extractor_params.use_request_proxy, proxy_headers.request + ) final_url, headers = await extractor.extract(extractor_params.destination) if extractor_params.redirect_stream: - formatted_headers = "&".join(f"h_{k}={v}" for k, v in headers.items()) - stream_url = f"/proxy/stream?api_password={settings.api_password}&d={final_url}&{formatted_headers}" + headers.update(proxy_headers.request) + stream_url = encode_mediaflow_proxy_url( + str(request.url_for("proxy_stream_endpoint").replace(scheme=get_original_scheme(request))), + destination_url=final_url, + query_params={"api_password": settings.api_password}, + request_headers=headers, + response_headers=proxy_headers.response, + ) return RedirectResponse(url=stream_url) return {"url": final_url, "headers": headers}