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
+
+
+
+
+
+
+
+
+ Select Debrid Service for Speed Test
+
+
+
+
+
+
+
+
+
+
+
+
+
+ AllDebrid Setup
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+ Initializing test...
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
\ 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.