Update app.py

app.py

@@ -9,6 +9,8 @@ import asyncio
 import cloudscraper
 from pydantic import BaseModel
 from urllib.parse import urlparse, parse_qs
+from collections import defaultdict
+import threading
 
 app = Flask(__name__)
 ytmusic = YTMusic()
@@ -162,7 +164,7 @@ api_rotator = ApiRotator([
 
 
 
-async def get_track_download_url(track_id: str) -> str:
+async def get_track_download_url(track_id: str, quality: str) -> str:
     apis = api_rotator.get_prioritized_apis()
     session = cloudscraper.create_scraper()  # Requires cloudscraper package
     headers = {
@@ -178,7 +180,7 @@ async def get_track_download_url(track_id: str) -> str:
             response = session.post(
                 api_url,
                 timeout=20,
-                json={"url": y_url, "audioFormat": "mp3", "downloadMode": "audio"},
+                json={"url": y_url, "audioFormat": "mp3", "downloadMode": "audio", "audioBitrate": quality},
                 headers=headers
             )
             logger.info(f"Response status: {response.status_code}")
@@ -206,18 +208,93 @@ async def get_track_download_url(track_id: str) -> str:
 
 
     
+
+
+
+
+# In-memory storage for request counting
+request_counts = defaultdict(list)
+request_lock = threading.Lock()
+DAILY_LIMIT = 20
+
+def cleanup_old_records():
+    """Remove records older than 24 hours"""
+    today = datetime.now()
+    with request_lock:
+        for ip in list(request_counts.keys()):
+            request_counts[ip] = [
+                (timestamp, count) 
+                for timestamp, count in request_counts[ip] 
+                if (today - timestamp).days < 1
+            ]
+            if not request_counts[ip]:
+                del request_counts[ip]
+
+def get_daily_requests(ip):
+    """Get the number of requests made today by this IP"""
+    cleanup_old_records()
+    today = datetime.now()
+    with request_lock:
+        return sum(
+            count for timestamp, count in request_counts[ip]
+            if (today - timestamp).days < 1
+        )
+
+def increment_daily_requests(ip):
+    """Increment the request counter for this IP"""
+    now = datetime.now()
+    with request_lock:
+        request_counts[ip].append((now, 1))
+
 @app.route('/track_dl', methods=['POST'])
 async def track_dl():
-    track_id = request.json.get('track_id')
-    dl_url = await get_track_download_url(track_id)
-    if dl_url and "http" in dl_url:
-        result = {"url": dl_url}
-        return jsonify(result)
-    else:
-        return {
-                "error": "Failed to Fetch the Track."
-            }
+    client_ip = request.remote_addr
     
+    # Check current usage
+    current_requests = get_daily_requests(client_ip)
+    remaining_requests = DAILY_LIMIT - current_requests
+    
+    if remaining_requests <= 0:
+        return jsonify({
+            "error": "Daily limit of 20 downloads exceeded. Get Premium or try again tomorrow.",
+            "remaining_requests": 0
+        }), 429
+    
+    data = request.get_json()
+    track_id = data.get('track_id')
+    quality = data.get('quality', '128')
+    
+    try:
+        quality_num = int(quality)
+        if quality_num > 128 or quality.upper() == 'FLAC':
+            return jsonify({
+                "error": "Quality above 128 or FLAC is for Premium users Only.",
+                "remaining_requests": remaining_requests
+            }), 400
+            
+        dl_url = await get_track_download_url(track_id, quality)
+        
+        if dl_url and "http" in dl_url:
+            # Increment counter only on successful download
+            increment_daily_requests(client_ip)
+            remaining_requests -= 1  # Decrease remaining count after successful increment
+            
+            result = {
+                "url": dl_url,
+                "remaining_requests": remaining_requests
+            }
+            return jsonify(result)
+        else:
+            return jsonify({
+                "error": "Failed to Fetch the Track.",
+                "remaining_requests": remaining_requests
+            }), 400
+            
+    except ValueError:
+        return jsonify({
+            "error": "Invalid quality value provided. It should be a valid integer or FLAC.",
+            "remaining_requests": remaining_requests
+        }), 400