From 64eeb8f9016c83042b078aedb58de9c418abb3bf Mon Sep 17 00:00:00 2001 From: mhdzumair Date: Tue, 12 Nov 2024 07:09:42 +0530 Subject: [PATCH] Add AD speed test & improve speed test UI & refactoring --- mediaflow_proxy/main.py | 34 +- mediaflow_proxy/routes/__init__.py | 3 + mediaflow_proxy/routes/speedtest.py | 43 ++ mediaflow_proxy/speedtest/__init__.py | 0 mediaflow_proxy/speedtest/models.py | 46 ++ .../speedtest/providers/all_debrid.py | 50 ++ mediaflow_proxy/speedtest/providers/base.py | 24 + .../speedtest/providers/real_debrid.py | 32 + mediaflow_proxy/speedtest/service.py | 129 ++++ mediaflow_proxy/static/speedtest.html | 689 ++++++++++++++++++ .../static/speedtest_progress.html | 140 ---- mediaflow_proxy/utils/http_utils.py | 2 +- 12 files changed, 1026 insertions(+), 166 deletions(-) create mode 100644 mediaflow_proxy/routes/speedtest.py create mode 100644 mediaflow_proxy/speedtest/__init__.py create mode 100644 mediaflow_proxy/speedtest/models.py create mode 100644 mediaflow_proxy/speedtest/providers/all_debrid.py create mode 100644 mediaflow_proxy/speedtest/providers/base.py create mode 100644 mediaflow_proxy/speedtest/providers/real_debrid.py create mode 100644 mediaflow_proxy/speedtest/service.py create mode 100644 mediaflow_proxy/static/speedtest.html delete mode 100644 mediaflow_proxy/static/speedtest_progress.html diff --git a/mediaflow_proxy/main.py b/mediaflow_proxy/main.py index c985061..404ed8a 100644 --- a/mediaflow_proxy/main.py +++ b/mediaflow_proxy/main.py @@ -1,19 +1,17 @@ import logging -import uuid from importlib import resources -from fastapi import FastAPI, Depends, Security, HTTPException, BackgroundTasks +from fastapi import FastAPI, Depends, Security, HTTPException from fastapi.security import APIKeyQuery, APIKeyHeader from starlette.middleware.cors import CORSMiddleware -from starlette.responses import RedirectResponse, JSONResponse +from starlette.responses import RedirectResponse from starlette.staticfiles import StaticFiles from mediaflow_proxy.configs import settings -from mediaflow_proxy.routes import proxy_router, extractor_router +from mediaflow_proxy.routes import proxy_router, extractor_router, speedtest_router from mediaflow_proxy.schemas import GenerateUrlRequest from mediaflow_proxy.utils.crypto_utils import EncryptionHandler, EncryptionMiddleware from mediaflow_proxy.utils.http_utils import encode_mediaflow_proxy_url -from mediaflow_proxy.utils.rd_speedtest import run_speedtest, prune_task, results logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(name)s - %(levelname)s - %(message)s") app = FastAPI() @@ -51,31 +49,16 @@ async def health_check(): return {"status": "healthy"} -@app.get("/speedtest") -async def trigger_speedtest(background_tasks: BackgroundTasks, api_password: str = Depends(verify_api_key)): - # Generate a random UUID as task_id - task_id = str(uuid.uuid4()) # Generate unique task ID - background_tasks.add_task(run_speedtest, task_id) - - # Schedule the task to be pruned after 1 hour - background_tasks.add_task(prune_task, task_id) - - return RedirectResponse(url=f"/speedtest_progress.html?task_id={task_id}") - - -@app.get("/speedtest/results/{task_id}", response_class=JSONResponse) -async def get_speedtest_result(task_id: str): - if task_id in results: - return results[task_id] - else: - return {"message": "Speedtest is still running, please wait or the task may have expired."} - - @app.get("/favicon.ico") async def get_favicon(): return RedirectResponse(url="/logo.png") +@app.get("/speedtest") +async def show_speedtest_page(): + return RedirectResponse(url="/speedtest.html") + + @app.post("/generate_encrypted_or_encoded_url") async def generate_encrypted_or_encoded_url(request: GenerateUrlRequest): if "api_password" not in request.query_params: @@ -97,6 +80,7 @@ async def generate_encrypted_or_encoded_url(request: GenerateUrlRequest): app.include_router(proxy_router, prefix="/proxy", tags=["proxy"], dependencies=[Depends(verify_api_key)]) app.include_router(extractor_router, prefix="/extractor", tags=["extractors"], dependencies=[Depends(verify_api_key)]) +app.include_router(speedtest_router, prefix="/speedtest", tags=["speedtest"], dependencies=[Depends(verify_api_key)]) static_path = resources.files("mediaflow_proxy").joinpath("static") app.mount("/", StaticFiles(directory=str(static_path), html=True), name="static") diff --git a/mediaflow_proxy/routes/__init__.py b/mediaflow_proxy/routes/__init__.py index 4a828f6..4a8a05c 100644 --- a/mediaflow_proxy/routes/__init__.py +++ b/mediaflow_proxy/routes/__init__.py @@ -1,2 +1,5 @@ from .proxy import proxy_router from .extractor import extractor_router +from .speedtest import speedtest_router + +__all__ = ["proxy_router", "extractor_router", "speedtest_router"] diff --git a/mediaflow_proxy/routes/speedtest.py b/mediaflow_proxy/routes/speedtest.py new file mode 100644 index 0000000..ec33bfd --- /dev/null +++ b/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() diff --git a/mediaflow_proxy/speedtest/__init__.py b/mediaflow_proxy/speedtest/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/mediaflow_proxy/speedtest/models.py b/mediaflow_proxy/speedtest/models.py new file mode 100644 index 0000000..c789fe1 --- /dev/null +++ b/mediaflow_proxy/speedtest/models.py @@ -0,0 +1,46 @@ +from datetime import datetime +from enum import Enum +from typing import Dict, Optional + +from pydantic import BaseModel, Field + + +class SpeedTestProvider(str, Enum): + REAL_DEBRID = "real_debrid" + ALL_DEBRID = "all_debrid" + + +class ServerInfo(BaseModel): + url: str + name: str + + +class UserInfo(BaseModel): + ip: Optional[str] = None + isp: Optional[str] = None + country: Optional[str] = None + + +class SpeedTestResult(BaseModel): + speed_mbps: float = Field(..., description="Speed in Mbps") + duration: float = Field(..., description="Test duration in seconds") + data_transferred: int = Field(..., description="Data transferred in bytes") + timestamp: datetime = Field(default_factory=datetime.utcnow) + + +class LocationResult(BaseModel): + result: Optional[SpeedTestResult] = None + error: Optional[str] = None + server_name: str + server_url: str + + +class SpeedTestTask(BaseModel): + task_id: str + provider: SpeedTestProvider + results: Dict[str, LocationResult] = {} + started_at: datetime + completed_at: Optional[datetime] = None + status: str = "running" + user_info: Optional[UserInfo] = None + current_location: Optional[str] = None diff --git a/mediaflow_proxy/speedtest/providers/all_debrid.py b/mediaflow_proxy/speedtest/providers/all_debrid.py new file mode 100644 index 0000000..9bf3f7d --- /dev/null +++ b/mediaflow_proxy/speedtest/providers/all_debrid.py @@ -0,0 +1,50 @@ +import random +from typing import Dict, Tuple, Optional + +from mediaflow_proxy.configs import settings +from mediaflow_proxy.speedtest.models import ServerInfo, UserInfo +from mediaflow_proxy.speedtest.providers.base import BaseSpeedTestProvider, SpeedTestProviderConfig +from mediaflow_proxy.utils.http_utils import request_with_retry + + +class SpeedTestError(Exception): + pass + + +class AllDebridSpeedTest(BaseSpeedTestProvider): + """AllDebrid speed test provider implementation.""" + + def __init__(self, api_key: str): + self.api_key = api_key + self.servers: Dict[str, ServerInfo] = {} + + async def get_test_urls(self) -> Tuple[Dict[str, str], Optional[UserInfo]]: + response = await request_with_retry( + "GET", + "https://alldebrid.com/internalapi/v4/speedtest", + headers={"User-Agent": settings.user_agent}, + params={"agent": "service", "version": "1.0-363869a7", "apikey": self.api_key}, + ) + + if response.status_code != 200: + raise SpeedTestError("Failed to fetch AllDebrid servers") + + data = response.json() + if data["status"] != "success": + raise SpeedTestError("AllDebrid API returned error") + + # Create UserInfo + user_info = UserInfo(ip=data["data"]["ip"], isp=data["data"]["isp"], country=data["data"]["country"]) + + # Store server info + self.servers = {server["name"]: ServerInfo(**server) for server in data["data"]["servers"]} + + # Generate URLs with random number + random_number = f"{random.uniform(1, 2):.24f}".replace(".", "") + urls = {name: f"{server.url}/speedtest/{random_number}" for name, server in self.servers.items()} + + return urls, user_info + + async def get_config(self) -> SpeedTestProviderConfig: + urls, _ = await self.get_test_urls() + return SpeedTestProviderConfig(test_duration=10, test_urls=urls) diff --git a/mediaflow_proxy/speedtest/providers/base.py b/mediaflow_proxy/speedtest/providers/base.py new file mode 100644 index 0000000..275c901 --- /dev/null +++ b/mediaflow_proxy/speedtest/providers/base.py @@ -0,0 +1,24 @@ +from abc import ABC, abstractmethod +from typing import Dict, Tuple, Optional +from pydantic import BaseModel + +from mediaflow_proxy.speedtest.models import UserInfo + + +class SpeedTestProviderConfig(BaseModel): + test_duration: int = 10 # seconds + test_urls: Dict[str, str] + + +class BaseSpeedTestProvider(ABC): + """Base class for speed test providers.""" + + @abstractmethod + async def get_test_urls(self) -> Tuple[Dict[str, str], Optional[UserInfo]]: + """Get list of test URLs for the provider and optional user info.""" + pass + + @abstractmethod + async def get_config(self) -> SpeedTestProviderConfig: + """Get provider-specific configuration.""" + pass diff --git a/mediaflow_proxy/speedtest/providers/real_debrid.py b/mediaflow_proxy/speedtest/providers/real_debrid.py new file mode 100644 index 0000000..0d38511 --- /dev/null +++ b/mediaflow_proxy/speedtest/providers/real_debrid.py @@ -0,0 +1,32 @@ +from typing import Dict, Tuple, Optional +import random + +from mediaflow_proxy.speedtest.models import UserInfo +from mediaflow_proxy.speedtest.providers.base import BaseSpeedTestProvider, SpeedTestProviderConfig + + +class RealDebridSpeedTest(BaseSpeedTestProvider): + """RealDebrid speed test provider implementation.""" + + async def get_test_urls(self) -> Tuple[Dict[str, str], Optional[UserInfo]]: + urls = { + "AMS": "https://45.download.real-debrid.com/speedtest/testDefault.rar/", + "RBX": "https://rbx.download.real-debrid.com/speedtest/test.rar/", + "LON1": "https://lon1.download.real-debrid.com/speedtest/test.rar/", + "HKG1": "https://hkg1.download.real-debrid.com/speedtest/test.rar/", + "SGP1": "https://sgp1.download.real-debrid.com/speedtest/test.rar/", + "SGPO1": "https://sgpo1.download.real-debrid.com/speedtest/test.rar/", + "TYO1": "https://tyo1.download.real-debrid.com/speedtest/test.rar/", + "LAX1": "https://lax1.download.real-debrid.com/speedtest/test.rar/", + "TLV1": "https://tlv1.download.real-debrid.com/speedtest/test.rar/", + "MUM1": "https://mum1.download.real-debrid.com/speedtest/test.rar/", + "JKT1": "https://jkt1.download.real-debrid.com/speedtest/test.rar/", + "Cloudflare": "https://45.download.real-debrid.cloud/speedtest/testCloudflare.rar/", + } + # Add random number to prevent caching + urls = {location: f"{base_url}{random.uniform(0, 1):.16f}" for location, base_url in urls.items()} + return urls, None + + async def get_config(self) -> SpeedTestProviderConfig: + urls, _ = await self.get_test_urls() + return SpeedTestProviderConfig(test_duration=10, test_urls=urls) diff --git a/mediaflow_proxy/speedtest/service.py b/mediaflow_proxy/speedtest/service.py new file mode 100644 index 0000000..43f1901 --- /dev/null +++ b/mediaflow_proxy/speedtest/service.py @@ -0,0 +1,129 @@ +import logging +import time +from datetime import datetime +from typing import Dict, Optional, Type + +from cachetools import TTLCache +from httpx import AsyncClient + +from mediaflow_proxy.utils.http_utils import Streamer +from .models import SpeedTestTask, LocationResult, SpeedTestResult, SpeedTestProvider +from .providers.all_debrid import AllDebridSpeedTest +from .providers.base import BaseSpeedTestProvider +from .providers.real_debrid import RealDebridSpeedTest +from ..configs import settings + +logger = logging.getLogger(__name__) + + +class SpeedTestService: + """Service for managing speed tests across different providers.""" + + def __init__(self): + # Cache for speed test results (1 hour TTL) + self._cache: TTLCache[str, SpeedTestTask] = TTLCache(maxsize=100, ttl=3600) + + # Provider mapping + self._providers: Dict[SpeedTestProvider, Type[BaseSpeedTestProvider]] = { + SpeedTestProvider.REAL_DEBRID: RealDebridSpeedTest, + SpeedTestProvider.ALL_DEBRID: AllDebridSpeedTest, + } + + def _get_provider(self, provider: SpeedTestProvider, api_key: Optional[str] = None) -> BaseSpeedTestProvider: + """Get the appropriate provider implementation.""" + provider_class = self._providers.get(provider) + if not provider_class: + raise ValueError(f"Unsupported provider: {provider}") + + if provider == SpeedTestProvider.ALL_DEBRID and not api_key: + raise ValueError("API key required for AllDebrid") + + return provider_class(api_key) if provider == SpeedTestProvider.ALL_DEBRID else provider_class() + + async def create_test( + self, task_id: str, provider: SpeedTestProvider, api_key: Optional[str] = None + ) -> SpeedTestTask: + """Create a new speed test task.""" + provider_impl = self._get_provider(provider, api_key) + + # Get initial URLs and user info + urls, user_info = await provider_impl.get_test_urls() + + task = SpeedTestTask(task_id=task_id, provider=provider, started_at=datetime.utcnow(), user_info=user_info) + + self._cache[task_id] = task + return task + + async def get_test_results(self, task_id: str) -> Optional[SpeedTestTask]: + """Get results for a specific task.""" + return self._cache.get(task_id) + + async def run_speedtest(self, task_id: str, provider: SpeedTestProvider, api_key: Optional[str] = None): + """Run the speed test with real-time updates.""" + try: + task = self._cache.get(task_id) + if not task: + raise ValueError(f"Task {task_id} not found") + + provider_impl = self._get_provider(provider, api_key) + config = await provider_impl.get_config() + + async with AsyncClient(follow_redirects=True, timeout=10, proxy=settings.proxy_url) as client: + streamer = Streamer(client) + + for location, url in config.test_urls.items(): + try: + task.current_location = location + result = await self._test_location(location, url, streamer, config.test_duration, provider_impl) + task.results[location] = result + self._cache[task_id] = task + except Exception as e: + logger.error(f"Error testing {location}: {str(e)}") + task.results[location] = LocationResult( + error=str(e), server_name=location, server_url=config.test_urls[location] + ) + self._cache[task_id] = task + + # Mark task as completed + task.completed_at = datetime.utcnow() + task.status = "completed" + task.current_location = None + self._cache[task_id] = task + + except Exception as e: + logger.error(f"Error in speed test task {task_id}: {str(e)}") + if task := self._cache.get(task_id): + task.status = "failed" + self._cache[task_id] = task + + async def _test_location( + self, location: str, url: str, streamer: Streamer, test_duration: int, provider: BaseSpeedTestProvider + ) -> LocationResult: + """Test speed for a specific location.""" + try: + start_time = time.time() + total_bytes = 0 + + async for chunk in streamer.stream_content(url, headers={}): + if time.time() - start_time >= test_duration: + break + total_bytes += len(chunk) + + duration = time.time() - start_time + speed_mbps = (total_bytes * 8) / (duration * 1_000_000) + + # Get server info if available (for AllDebrid) + server_info = getattr(provider, "servers", {}).get(location) + server_url = server_info.url if server_info else url + + return LocationResult( + result=SpeedTestResult( + speed_mbps=round(speed_mbps, 2), duration=round(duration, 2), data_transferred=total_bytes + ), + server_name=location, + server_url=server_url, + ) + + except Exception as e: + logger.error(f"Error testing {location}: {str(e)}") + raise # Re-raise to be handled by run_speedtest diff --git a/mediaflow_proxy/static/speedtest.html b/mediaflow_proxy/static/speedtest.html new file mode 100644 index 0000000..e88e3a9 --- /dev/null +++ b/mediaflow_proxy/static/speedtest.html @@ -0,0 +1,689 @@ + + + + + + Debrid Speed Test + + + + + + +
+ +
+ +
+ +
+ +
+

+ Enter API Password +

+ +
+
+
+ + +
+
+ + +
+ +
+
+
+ + + + + + + + + + + + + + + + +
+
+ + + + \ No newline at end of file diff --git a/mediaflow_proxy/static/speedtest_progress.html b/mediaflow_proxy/static/speedtest_progress.html deleted file mode 100644 index a180764..0000000 --- a/mediaflow_proxy/static/speedtest_progress.html +++ /dev/null @@ -1,140 +0,0 @@ - - - - - - Speedtest - - - - -
- -
-
-

Speedtest in progress... Please wait up to 3 minutes.

-
-
- - \ No newline at end of file diff --git a/mediaflow_proxy/utils/http_utils.py b/mediaflow_proxy/utils/http_utils.py index 09e928f..5e307e0 100644 --- a/mediaflow_proxy/utils/http_utils.py +++ b/mediaflow_proxy/utils/http_utils.py @@ -245,7 +245,7 @@ async def download_file_with_retry( async def request_with_retry( method: str, url: str, headers: dict, timeout: float = 10.0, use_request_proxy: bool = True, **kwargs -): +) -> httpx.Response: """ Sends an HTTP request with retry logic.