Apply suggestions from code review

Co-authored-by: coderabbitai[bot] <136622811+coderabbitai[bot]@users.noreply.github.com>

mediaflow_proxy/extractors/uqload.py

@@ -3,12 +3,39 @@ import re
 from mediaflow_proxy.configs import settings
 
 
-async def uqload_url(d: str, use_request_proxy: bool):
-    async with httpx.AsyncClient(proxy=settings.proxy_url if use_request_proxy else None) as client:
+from typing import Tuple, Dict, Optional
+
+async def uqload_url(d: str, use_request_proxy: bool) -> Tuple[Optional[str], Dict[str, str]]:
+    """
+    Extract video URL from Uqload.
+    
+    Args:
+        d: The Uqload video URL
+        use_request_proxy: Whether to use proxy for the request
+    
+    Returns:
+        Tuple containing the extracted video URL (or None if not found) and headers dictionary
+    
+    Raises:
+        httpx.HTTPError: If the HTTP request fails
+    """
+    if not d.startswith(('http://', 'https://')):
+        raise ValueError("Invalid URL format")
 
-        response = await client.get(d, follow_redirects=True)
-        video_url_match = re.search(r'sources: \["(.*?)"\]', response.text)
-        if video_url_match:
-            final_url = video_url_match.group(1)
-        uqload_dict = {"Referer": "https://uqload.to/"}
-        return final_url, uqload_dict
+    REFERER = "https://uqload.to/"
+    final_url = None
+
+    async with httpx.AsyncClient(proxy=settings.proxy_url if use_request_proxy else None) as client:
+        try:
+            response = await client.get(d, follow_redirects=True)
+            response.raise_for_status()
+            
+            # Look for video URL in response using a more robust pattern
+            video_url_match = re.search(r'sources:\s*\[(["\'])(.*?)\1\]', response.text)
+            if video_url_match:
+                final_url = video_url_match.group(2)
+            
+            return final_url, {"Referer": REFERER}
+        except httpx.HTTPError as e:
+            # Log the error here if logging is available
+            raise

mediaflow_proxy/extractors_routes.py

@@ -10,7 +10,7 @@ host_map = {"Doodstream": doodstream_url, "Mixdrop": mixdrop_url, "Uqload": uqlo
 
 
 @extractor_router.get("/extractor")
-async def doodstream_extractor(
+async def extract_media_url(
     d: str = Query(..., description="Extract Clean Link from various Hosts"),
     use_request_proxy: bool = Query(False, description="Whether to use the MediaFlow proxy configuration."),
     host: str = Query(
@@ -32,8 +32,15 @@ async def doodstream_extractor(
     """
     try:
         final_url, headers_dict = await host_map[host](d, use_request_proxy)
+    except KeyError:
+        return JSONResponse(
+            status_code=400,
+            content={"error": f"Invalid host type. Available hosts: {', '.join(host_map.keys())}"}
+        )
+    except ValueError as e:
+        return JSONResponse(status_code=400, content={"error": str(e)})
     except Exception as e:
-        return JSONResponse(content={"error": str(e)})
+        return JSONResponse(status_code=500, content={"error": "Internal server error"})
     if redirect_stream == True:
         formatted_headers = format_headers(headers_dict)
         redirected_stream = f"/proxy/stream?api_password={settings.api_password}&d={final_url}&{formatted_headers}"

mediaflow_proxy/static/speedtest_progress.html

@@ -63,34 +63,56 @@
         <script>
             const urlParams = new URLSearchParams(window.location.search);
             const taskId = urlParams.get("task_id");
+            if (!taskId || !/^[a-zA-Z0-9-_]+$/.test(taskId)) {
+                window.location.href = "/speedtest";
+            }
         
+            let statusCheckTimeout;
+            let retryCount = 0;
+            const MAX_RETRIES = 5;
+
             async function checkStatus() {
                 try {
-                    const response = await fetch(`/speedtest/results/${taskId}`);
+                    const controller = new AbortController();
+                    const timeoutId = setTimeout(() => controller.abort(), 5000);
+                    const response = await fetch(`/speedtest/results/${taskId}`, {
+                        signal: controller.signal
+                    });
+                    clearTimeout(timeoutId);
+
                     if (!response.ok) {
                         throw new Error('Network response was not ok');
                     }
                     const data = await response.json();
                     console.log("Fetched data:", data);
+                    retryCount = 0;
 
-                    // Check if the test is still running based on the response data
                     if (data && data.message && data.message.includes("still running")) {
                         console.log("Test still running, polling again...");
-                        // Poll again after 5 seconds if the test is still running
-                        setTimeout(checkStatus, 5000);
+                        statusCheckTimeout = setTimeout(checkStatus, 5000);
                     } else {
                         console.log("Test complete, redirecting after a short delay...");
-                        // Redirect to the results if the test is done after a short delay
                         setTimeout(() => {
                             window.location.href = `/speedtest/results/${taskId}`;
-                        }, 2000); // 2 seconds delay
+                        }, 2000);
                     }
                 } catch (error) {
                     console.error("Error fetching status:", error);
-                    // Retry after 5 seconds in case of error
-                    setTimeout(checkStatus, 5000);
+                    retryCount++;
+                    if (retryCount < MAX_RETRIES) {
+                        statusCheckTimeout = setTimeout(checkStatus, 5000);
+                    } else {
+                        alert("Failed to check status after multiple attempts. Please refresh the page.");
+                    }
                 }
             }
+
+            // Cleanup on page unload
+            window.addEventListener('unload', () => {
+                if (statusCheckTimeout) {
+                    clearTimeout(statusCheckTimeout);
+                }
+            });
         
             // Start the first status check after 120 seconds (120000 milliseconds)
             setTimeout(checkStatus, 120000);
@@ -105,12 +127,14 @@
     </head>
     <body class="light-mode">
         <div class="toggle-switch">
-            <label for="darkModeToggle">Dark Mode</label>
-            <input type="checkbox" id="darkModeToggle" onclick="toggleDarkMode()">
+            <label for="darkModeToggle" class="switch">
+                <input type="checkbox" id="darkModeToggle" onclick="toggleDarkMode()" aria-label="Toggle dark mode">
+                <span class="slider">Dark Mode</span>
+            </label>
         </div>
         <div class="container">
             <h1>Speedtest in progress... Please wait up to 3 minutes.</h1>
-            <div class="progress-bar"></div>
+            <div class="progress-bar" role="progressbar" aria-valuemin="0" aria-valuemax="100" aria-valuenow="0"></div>
         </div>
     </body>
 </html>

mediaflow_proxy/utils/rd_speedtest.py

@@ -37,32 +37,33 @@ async def perform_speed_test():
     async with AsyncClient() as client:
         streamer = Streamer(client)
 
-        for location, base_url in test_urls.items():
-                # Generate a random float with 16 decimal places
-                random_number = f"{random.uniform(0, 1):.16f}"
-                url = f"{base_url}{random_number}"
-
-                logging.info(f"Testing URL: {url}")
-
+        async def test_single_url(location: str, url: str) -> Dict[str, Any]:
+            try:
                 start_time = time.time()
                 total_bytes = 0
                 
-                try:
-                    # Stream the response
-                    async for chunk in streamer.stream_content(url, headers={}):
-                        if time.time() - start_time >= test_duration:
-                            break
-                        total_bytes += len(chunk)
+                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)
-                    speed[location] = {
-                        "speed_mbps": round(speed_mbps, 2),
-                        "duration": round(duration, 2)
-                    }
-                    logging.info(f"Speed for {location}: {speed_mbps} Mbps in {duration} seconds")
-                except Exception as e:
-                    speed[location] = {"error": str(e)}
-                    logging.error(f"Error for {location}: {e}")
+                duration = time.time() - start_time
+                speed_mbps = (total_bytes * 8) / (duration * 1_000_000)
+                return {
+                    "speed_mbps": round(speed_mbps, 2),
+                    "duration": round(duration, 2)
+                }
+            except Exception as e:
+                logging.error(f"Error testing {location}: {e}")
+                return {"error": str(e)}
 
+        for location, base_url in test_urls.items():
+            random_number = f"{random.uniform(0, 1):.16f}"
+            url = f"{base_url}{random_number}"
+            logging.info(f"Testing URL: {url}")
+            
+            speed[location] = await test_single_url(location, url)
+            
+            # Add rate limiting between tests
+            await asyncio.sleep(1)
     return speed