Chrunos commited on
Commit
2dd163d
·
verified ·
1 Parent(s): c23aa3d

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +63 -0
app.py CHANGED
@@ -14,6 +14,9 @@ from urllib.parse import quote
14
  from pathlib import Path
15
  import uuid
16
  from fastapi import BackgroundTasks
 
 
 
17
 
18
 
19
  # Set up logging
@@ -80,10 +83,70 @@ def get_track(track_id: str):
80
  return get_track_info(track_id)
81
 
82
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
83
  # Download a track and return a download URL
84
  @app.post("/download/track")
85
  def download_track(request: DownloadRequest):
86
  try:
 
 
 
 
 
 
 
 
 
 
 
87
  if request.arl is None or request.arl.strip() == "":
88
  ARL = ARL_TOKEN
89
  else:
 
14
  from pathlib import Path
15
  import uuid
16
  from fastapi import BackgroundTasks
17
+ from collections import defaultdict
18
+ import time
19
+ from datetime import timedelta
20
 
21
 
22
  # Set up logging
 
83
  return get_track_info(track_id)
84
 
85
 
86
+ # Rate limiting dictionary
87
+ class RateLimiter:
88
+ def __init__(self, max_requests: int, time_window: timedelta):
89
+ self.max_requests = max_requests
90
+ self.time_window = time_window
91
+ self.requests: Dict[str, list] = defaultdict(list)
92
+
93
+ def _cleanup_old_requests(self, user_ip: str) -> None:
94
+ """Remove requests that are outside the time window."""
95
+ current_time = time.time()
96
+ self.requests[user_ip] = [
97
+ timestamp for timestamp in self.requests[user_ip]
98
+ if current_time - timestamp < self.time_window.total_seconds()
99
+ ]
100
+
101
+ def is_rate_limited(self, user_ip: str) -> bool:
102
+ """Check if the user has exceeded their rate limit."""
103
+ self._cleanup_old_requests(user_ip)
104
+
105
+ # Get current count after cleanup
106
+ current_count = len(self.requests[user_ip])
107
+
108
+ # Add current request timestamp (incrementing the count)
109
+ current_time = time.time()
110
+ self.requests[user_ip].append(current_time)
111
+
112
+ # Check if user has exceeded the maximum requests
113
+ return (current_count + 1) > self.max_requests
114
+
115
+ def get_current_count(self, user_ip: str) -> int:
116
+ """Get the current request count for an IP."""
117
+ self._cleanup_old_requests(user_ip)
118
+ return len(self.requests[user_ip])
119
+
120
+
121
+ # Initialize rate limiter with 25 requests per day
122
+ rate_limiter = RateLimiter(
123
+ max_requests=25,
124
+ time_window=timedelta(days=1)
125
+ )
126
+
127
+ def get_user_ip(request: Request) -> str:
128
+ """Helper function to get user's IP address."""
129
+ forwarded = request.headers.get("X-Forwarded-For")
130
+ if forwarded:
131
+ return forwarded.split(",")[0]
132
+ return request.client.host
133
+
134
+
135
  # Download a track and return a download URL
136
  @app.post("/download/track")
137
  def download_track(request: DownloadRequest):
138
  try:
139
+ user_ip = get_user_ip(request)
140
+
141
+ if rate_limiter.is_rate_limited(user_ip):
142
+ current_count = rate_limiter.get_current_count(user_ip)
143
+ raise HTTPException(
144
+ status_code=429,
145
+ detail={
146
+ "detail": "You have exceeded the maximum number of requests per day. Please try again tomorrow.",
147
+ "help": "https://t.me/chrunoss"
148
+ }
149
+ )
150
  if request.arl is None or request.arl.strip() == "":
151
  ARL = ARL_TOKEN
152
  else: