Add files via upload

mediaflow_proxy/routes/__init__.py

@@ -0,0 +1,5 @@
+from .proxy import proxy_router
+from .extractor import extractor_router
+from .speedtest import speedtest_router
+
+__all__ = ["proxy_router", "extractor_router", "speedtest_router"]

mediaflow_proxy/routes/extractor.py

@@ -0,0 +1,61 @@
+import logging
+from typing import Annotated
+
+from fastapi import APIRouter, Query, HTTPException, Request, Depends
+from fastapi.responses import RedirectResponse
+
+from mediaflow_proxy.extractors.base import ExtractorError
+from mediaflow_proxy.extractors.factory import ExtractorFactory
+from mediaflow_proxy.schemas import ExtractorURLParams
+from mediaflow_proxy.utils.cache_utils import get_cached_extractor_result, set_cache_extractor_result
+from mediaflow_proxy.utils.http_utils import (
+    encode_mediaflow_proxy_url,
+    get_original_scheme,
+    ProxyRequestHeaders,
+    get_proxy_headers,
+)
+
+extractor_router = APIRouter()
+logger = logging.getLogger(__name__)
+
+
+@extractor_router.head("/video")
+@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:
+        cache_key = f"{extractor_params.host}_{extractor_params.model_dump_json()}"
+        response = await get_cached_extractor_result(cache_key)
+        if not response:
+            extractor = ExtractorFactory.get_extractor(extractor_params.host, proxy_headers.request)
+            response = await extractor.extract(extractor_params.destination, **extractor_params.extra_params)
+            await set_cache_extractor_result(cache_key, response)
+        else:
+            response["request_headers"].update(proxy_headers.request)
+
+        response["mediaflow_proxy_url"] = str(
+            request.url_for(response.pop("mediaflow_endpoint")).replace(scheme=get_original_scheme(request))
+        )
+        response["query_params"] = response.get("query_params", {})
+        # Add API password to query params
+        response["query_params"]["api_password"] = request.query_params.get("api_password")
+
+        if extractor_params.redirect_stream:
+            stream_url = encode_mediaflow_proxy_url(
+                **response,
+                response_headers=proxy_headers.response,
+            )
+            return RedirectResponse(url=stream_url, status_code=302)
+
+        return response
+
+    except ExtractorError as e:
+        logger.error(f"Extraction failed: {str(e)}")
+        raise HTTPException(status_code=400, detail=str(e))
+    except Exception as e:
+        logger.exception(f"Extraction failed: {str(e)}")
+        raise HTTPException(status_code=500, detail=f"Extraction failed: {str(e)}")

mediaflow_proxy/routes/proxy.py

@@ -0,0 +1,138 @@
+from typing import Annotated
+
+from fastapi import Request, Depends, APIRouter, Query, HTTPException
+
+from mediaflow_proxy.handlers import (
+    handle_hls_stream_proxy,
+    proxy_stream,
+    get_manifest,
+    get_playlist,
+    get_segment,
+    get_public_ip,
+)
+from mediaflow_proxy.schemas import (
+    MPDSegmentParams,
+    MPDPlaylistParams,
+    HLSManifestParams,
+    ProxyStreamParams,
+    MPDManifestParams,
+)
+from mediaflow_proxy.utils.http_utils import get_proxy_headers, ProxyRequestHeaders
+
+proxy_router = APIRouter()
+
+
+@proxy_router.head("/hls/manifest.m3u8")
+@proxy_router.get("/hls/manifest.m3u8")
+async def hls_manifest_proxy(
+    request: Request,
+    hls_params: Annotated[HLSManifestParams, Query()],
+    proxy_headers: Annotated[ProxyRequestHeaders, Depends(get_proxy_headers)],
+):
+    """
+    Proxify HLS stream requests, fetching and processing the m3u8 playlist or streaming the content.
+
+    Args:
+        request (Request): The incoming HTTP request.
+        hls_params (HLSPlaylistParams): The parameters for the HLS stream request.
+        proxy_headers (ProxyRequestHeaders): The headers to include in the request.
+
+    Returns:
+        Response: The HTTP response with the processed m3u8 playlist or streamed content.
+    """
+    return await handle_hls_stream_proxy(request, hls_params, proxy_headers)
+
+
+@proxy_router.head("/stream")
+@proxy_router.get("/stream")
+async def proxy_stream_endpoint(
+    request: Request,
+    stream_params: Annotated[ProxyStreamParams, Query()],
+    proxy_headers: Annotated[ProxyRequestHeaders, Depends(get_proxy_headers)],
+):
+    """
+    Proxies stream requests to the given video URL.
+
+    Args:
+        request (Request): The incoming HTTP request.
+        stream_params (ProxyStreamParams): The parameters for the stream request.
+        proxy_headers (ProxyRequestHeaders): The headers to include in the request.
+
+    Returns:
+        Response: The HTTP response with the streamed content.
+    """
+    content_range = proxy_headers.request.get("range", "bytes=0-")
+    if "nan" in content_range.casefold():
+        # Handle invalid range requests "bytes=NaN-NaN"
+        raise HTTPException(status_code=416, detail="Invalid Range Header")
+    proxy_headers.request.update({"range": content_range})
+    return await proxy_stream(request.method, stream_params, proxy_headers)
+
+
+@proxy_router.get("/mpd/manifest.m3u8")
+async def mpd_manifest_proxy(
+    request: Request,
+    manifest_params: Annotated[MPDManifestParams, Query()],
+    proxy_headers: Annotated[ProxyRequestHeaders, Depends(get_proxy_headers)],
+):
+    """
+    Retrieves and processes the MPD manifest, converting it to an HLS manifest.
+
+    Args:
+        request (Request): The incoming HTTP request.
+        manifest_params (MPDManifestParams): The parameters for the manifest request.
+        proxy_headers (ProxyRequestHeaders): The headers to include in the request.
+
+    Returns:
+        Response: The HTTP response with the HLS manifest.
+    """
+    return await get_manifest(request, manifest_params, proxy_headers)
+
+
+@proxy_router.get("/mpd/playlist.m3u8")
+async def playlist_endpoint(
+    request: Request,
+    playlist_params: Annotated[MPDPlaylistParams, Query()],
+    proxy_headers: Annotated[ProxyRequestHeaders, Depends(get_proxy_headers)],
+):
+    """
+    Retrieves and processes the MPD manifest, converting it to an HLS playlist for a specific profile.
+
+    Args:
+        request (Request): The incoming HTTP request.
+        playlist_params (MPDPlaylistParams): The parameters for the playlist request.
+        proxy_headers (ProxyRequestHeaders): The headers to include in the request.
+
+    Returns:
+        Response: The HTTP response with the HLS playlist.
+    """
+    return await get_playlist(request, playlist_params, proxy_headers)
+
+
+@proxy_router.get("/mpd/segment.mp4")
+async def segment_endpoint(
+    segment_params: Annotated[MPDSegmentParams, Query()],
+    proxy_headers: Annotated[ProxyRequestHeaders, Depends(get_proxy_headers)],
+):
+    """
+    Retrieves and processes a media segment, decrypting it if necessary.
+
+    Args:
+        segment_params (MPDSegmentParams): The parameters for the segment request.
+        proxy_headers (ProxyRequestHeaders): The headers to include in the request.
+
+    Returns:
+        Response: The HTTP response with the processed segment.
+    """
+    return await get_segment(segment_params, proxy_headers)
+
+
+@proxy_router.get("/ip")
+async def get_mediaflow_proxy_public_ip():
+    """
+    Retrieves the public IP address of the MediaFlow proxy server.
+
+    Returns:
+        Response: The HTTP response with the public IP address in the form of a JSON object. {"ip": "xxx.xxx.xxx.xxx"}
+    """
+    return await get_public_ip()

mediaflow_proxy/routes/speedtest.py

@@ -0,0 +1,43 @@
+import uuid
+
+from fastapi import APIRouter, BackgroundTasks, HTTPException, Request
+from fastapi.responses import RedirectResponse
+
+from mediaflow_proxy.speedtest.service import SpeedTestService, SpeedTestProvider
+
+speedtest_router = APIRouter()
+
+# Initialize service
+speedtest_service = SpeedTestService()
+
+
+@speedtest_router.get("/", summary="Show speed test interface")
+async def show_speedtest_page():
+    """Return the speed test HTML interface."""
+    return RedirectResponse(url="/speedtest.html")
+
+
+@speedtest_router.post("/start", summary="Start a new speed test", response_model=dict)
+async def start_speedtest(background_tasks: BackgroundTasks, provider: SpeedTestProvider, request: Request):
+    """Start a new speed test for the specified provider."""
+    task_id = str(uuid.uuid4())
+    api_key = request.headers.get("api_key")
+
+    # Create and initialize the task
+    await speedtest_service.create_test(task_id, provider, api_key)
+
+    # Schedule the speed test
+    background_tasks.add_task(speedtest_service.run_speedtest, task_id, provider, api_key)
+
+    return {"task_id": task_id}
+
+
+@speedtest_router.get("/results/{task_id}", summary="Get speed test results")
+async def get_speedtest_results(task_id: str):
+    """Get the results or current status of a speed test."""
+    task = await speedtest_service.get_test_results(task_id)
+
+    if not task:
+        raise HTTPException(status_code=404, detail="Speed test task not found or expired")
+
+    return task.dict()