Update app.py

app.py

@@ -25,6 +25,7 @@ from collections import defaultdict
 from starlette.responses import JSONResponse
 import logging
 import gc
+from typing import Dict, Any
 
 tmp_dir = tempfile.gettempdir()
 BASE_URL = "https://chrunos-zam.hf.space"
@@ -122,29 +123,60 @@ async def get_video_url(youtube_url: str):
 # Define a global temporary download directory
 global_download_dir = tempfile.mkdtemp()
 
-# Rate limiting dictionary
-request_counts = defaultdict(lambda: {"count": 0, "reset_time": datetime.now() + timedelta(days=1)})
-MAX_REQUESTS_PER_DAY = 100  # Set your desired limit
+
+class RateLimiter:
+    def __init__(self, max_requests: int, time_window: timedelta):
+        self.max_requests = max_requests
+        self.time_window = time_window
+        self.requests: Dict[str, list] = defaultdict(list)
+    
+    def _cleanup_old_requests(self, user_ip: str) -> None:
+        """Remove requests that are outside the time window."""
+        current_time = time.time()
+        self.requests[user_ip] = [
+            timestamp for timestamp in self.requests[user_ip]
+            if current_time - timestamp < self.time_window.total_seconds()
+        ]
+    
+    def is_rate_limited(self, user_ip: str) -> bool:
+        """Check if the user has exceeded their rate limit."""
+        self._cleanup_old_requests(user_ip)
+        
+        # Add current request
+        current_time = time.time()
+        self.requests[user_ip].append(current_time)
+        
+        # Check if user has exceeded the maximum requests
+        return len(self.requests[user_ip]) > self.max_requests
+
+
+
+# Initialize rate limiter with 100 requests per day
+rate_limiter = RateLimiter(
+    max_requests=5,
+    time_window=timedelta(days=1)
+)
 
 def get_user_ip(request: Request) -> str:
     """Helper function to get user's IP address."""
+    forwarded = request.headers.get("X-Forwarded-For")
+    if forwarded:
+        return forwarded.split(",")[0]
     return request.client.host
 
 @app.post("/maxs")
 async def download_high_quality_video(request: Request):
     user_ip = get_user_ip(request)
-    user_info = request_counts[user_ip]
     
-    # Check if reset time has passed
-    if datetime.now() > user_info["reset_time"]:
-        user_info["count"] = 0
-        user_info["reset_time"] = datetime.now() + timedelta(days=1)
-
-    # Check if user has exceeded the request limit
-    if user_info["count"] >= MAX_REQUESTS_PER_DAY:
-        error_message = "You have exceeded the maximum number of requests per day. Please try again tomorrow."
-        return {"error": error_message, "url": "https://t.me/chrunoss"}
-
+    if rate_limiter.is_rate_limited(user_ip):
+        raise HTTPException(
+            status_code=429,
+            detail={
+                "error": "You have exceeded the maximum number of requests per day. Please try again tomorrow.",
+                "url": "https://t.me/chrunoss"
+            }
+        )
+        
     data = await request.json()
     video_url = data.get('url')
     quality = data.get('quality', '1080')  # Default to 1080p if not specified
@@ -192,12 +224,10 @@ async def download_high_quality_video(request: Request):
     encoded_filename = urllib.parse.quote(downloaded_file.name)
     download_url = f"{BASE_URL}/file/{encoded_filename}"
 
-    # Increment the user's request count
-    user_info["count"] += 1
 
     gc.collect()
 
-    return {"url": download_url}
+    return {"url": download_url, "requests_remaining": rate_limiter.max_requests - rate_limiter.get_current_count(user_ip)}