Add files via upload

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

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)

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

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)

mediaflow_proxy/speedtest/service.py

@@ -0,0 +1,129 @@
+import logging
+import time
+from datetime import datetime, timezone
+from typing import Dict, Optional, Type
+
+from mediaflow_proxy.utils.cache_utils import get_cached_speedtest, set_cache_speedtest
+from mediaflow_proxy.utils.http_utils import Streamer, create_httpx_client
+from .models import SpeedTestTask, LocationResult, SpeedTestResult, SpeedTestProvider
+from .providers.all_debrid import AllDebridSpeedTest
+from .providers.base import BaseSpeedTestProvider
+from .providers.real_debrid import RealDebridSpeedTest
+
+logger = logging.getLogger(__name__)
+
+
+class SpeedTestService:
+    """Service for managing speed tests across different providers."""
+
+    def __init__(self):
+        # 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.now(tz=timezone.utc), user_info=user_info
+        )
+
+        await set_cache_speedtest(task_id, task)
+        return task
+
+    @staticmethod
+    async def get_test_results(task_id: str) -> Optional[SpeedTestTask]:
+        """Get results for a specific task."""
+        return await get_cached_speedtest(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 = await get_cached_speedtest(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 create_httpx_client() as client:
+                streamer = Streamer(client)
+
+                for location, url in config.test_urls.items():
+                    try:
+                        task.current_location = location
+                        await set_cache_speedtest(task_id, task)
+                        result = await self._test_location(location, url, streamer, config.test_duration, provider_impl)
+                        task.results[location] = result
+                        await set_cache_speedtest(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]
+                        )
+                        await set_cache_speedtest(task_id, task)
+
+            # Mark task as completed
+            task.completed_at = datetime.now(tz=timezone.utc)
+            task.status = "completed"
+            task.current_location = None
+            await set_cache_speedtest(task_id, task)
+
+        except Exception as e:
+            logger.error(f"Error in speed test task {task_id}: {str(e)}")
+            if task := await get_cached_speedtest(task_id):
+                task.status = "failed"
+                await set_cache_speedtest(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
+
+            await streamer.create_streaming_response(url, headers={})
+
+            async for chunk in streamer.stream_content():
+                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