Update app.py

app.py

@@ -7,7 +7,7 @@ import yt_dlp
 import ffmpeg
 import urllib.parse
 import os
-from datetime import datetime
+from datetime import datetime, timedelta
 import schedule
 import requests
 import uvicorn
@@ -26,6 +26,7 @@ from PIL import Image
 from io import BytesIO
 from pathlib import Path
 from fastapi.staticfiles import StaticFiles
+from collections import defaultdict
 
 tmp_dir = tempfile.gettempdir()
 BASE_URL = "https://chrunos-multi.hf.space"
@@ -143,43 +144,159 @@ async def get_video_url(youtube_url: str):
 # Define a global temporary download directory
 global_download_dir = tempfile.mkdtemp()
 
-@app.post("/high")
+# Rate limiting dictionary
+request_counts = defaultdict(lambda: {"count": 0, "reset_time": datetime.now() + timedelta(days=1)})
+MAX_REQUESTS_PER_DAY = 10  # Set your desired limit
+
+def get_user_ip(request: Request) -> str:
+    """Helper function to get user's IP address."""
+    return request.client.host
+
+@app.post("/max")
 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}
+
     data = await request.json()
     video_url = data.get('url')
+    quality = data.get('quality', '1080')  # Default to 1080p if not specified
+
+    # Check if the requested quality is above 1080p
+    if int(quality) > 1080:
+        error_message = "Quality above 1080p is for Premium Members Only. Please check the URL for more information."
+        help_url = "https://chrunos.com/premium-shortcuts/"  # Replace with your actual URL
+        return {"error": error_message, "url": help_url}
+
     cookiefile = "firefox-cookies.txt"
     env_to_cookies_from_env("firefox-cookies.txt")
-    # Generate a unique output filename using timestamp
+
     timestamp = datetime.now().strftime('%Y%m%d%H%M%S')
     output_template = str(Path(global_download_dir) / f'%(title)s_{timestamp}.%(ext)s')
 
+    # Convert quality string to height
+    height_map = {
+        '480': 480,
+        '720': 720,
+        '1080': 1080
+    }
+    max_height = height_map.get(quality, 1080)  # Use the quality variable correctly
+
+    # Determine format string based on quality
+    format_str = f'bestvideo[height<={max_height}][vcodec^=avc]+bestaudio/best'
 
     ydl_opts = {
-        'format': 'bestvideo[height<=720]+bestaudio/best',
+        'format': format_str,
         'outtmpl': output_template,
         'quiet': True,
         'no_warnings': True,
         'noprogress': True,
-        'merge_output_format': 'mp4'
+        'merge_output_format': 'mp4',
+        'cookiefile': cookiefile
     }
-    ydl_opts["cookiefile"] =  "firefox-cookies.txt" #create_temp_cookie_file()
 
     await run_in_threadpool(lambda: yt_dlp.YoutubeDL(ydl_opts).download([video_url]))
-
-    # Find the downloaded file (with the title as filename)
+    
     downloaded_files = list(Path(global_download_dir).glob(f"*_{timestamp}.mp4"))
-
     if not downloaded_files:
         return {"error": "Download failed"}
-    
-    # Assume there is only one matching file
+
     downloaded_file = downloaded_files[0]
     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
+
+    return {"url": download_url}
+
+
+
+@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}
+
+    data = await request.json()
+    video_url = data.get('url')
+    quality = data.get('quality', '1080')  # Default to 1080p if not specified
+
+    # Check if the requested quality is above 1080p
+    if int(quality) > 1080:
+        error_message = "Quality above 1080p is not supported. Please check the URL for more information."
+        help_url = "https://example.com/more-info"  # Replace with your actual URL
+        return {"error": error_message, "url": help_url}
+
+    cookiefile = "firefox-cookies.txt"
+    env_to_cookies_from_env("firefox-cookies.txt")
+
+    timestamp = datetime.now().strftime('%Y%m%d%H%M%S')
+    output_template = str(Path(global_download_dir) / f'%(title)s_{timestamp}.%(ext)s')
+
+    # Convert quality string to height
+    height_map = {
+        '480': 480,
+        '720': 720,
+        '1080': 1080
+    }
+    max_height = height_map.get(quality, 1080)  # Use the quality variable correctly
+
+    # Determine format string based on quality
+    format_str = f'bestvideo[height<={max_height}][vcodec^=avc]+bestaudio/best'
+
+    ydl_opts = {
+        'format': format_str,
+        'outtmpl': output_template,
+        'quiet': True,
+        'no_warnings': True,
+        'noprogress': True,
+        'merge_output_format': 'mp4',
+        'cookiefile': cookiefile
+    }
+
+    await run_in_threadpool(lambda: yt_dlp.YoutubeDL(ydl_opts).download([video_url]))
+    
+    downloaded_files = list(Path(global_download_dir).glob(f"*_{timestamp}.mp4"))
+    if not downloaded_files:
+        return {"error": "Download failed"}
+
+    downloaded_file = downloaded_files[0]
+    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
+
     return {"url": download_url}
 
+# Mount the static files directory
+app.mount("/file", StaticFiles(directory=global_download_dir), name="downloads")
+
+@app.middleware("http")
+async def set_mime_type_middleware(request: Request, call_next):
+    response = await call_next(request)
+    if request.url.path.endswith(".mp4"):
+        response.headers["Content-Type"] = "video/mp4"
+    return response