forked from github/plane
41 lines
1.3 KiB
Python
41 lines
1.3 KiB
Python
from rest_framework.throttling import SimpleRateThrottle
|
|
|
|
class ApiKeyRateThrottle(SimpleRateThrottle):
|
|
scope = 'api_key'
|
|
rate = '60/minute'
|
|
|
|
def get_cache_key(self, request, view):
|
|
# Retrieve the API key from the request header
|
|
api_key = request.headers.get('X-Api-Key')
|
|
if not api_key:
|
|
return None # Allow the request if there's no API key
|
|
|
|
# Use the API key as part of the cache key
|
|
return f'{self.scope}:{api_key}'
|
|
|
|
def allow_request(self, request, view):
|
|
allowed = super().allow_request(request, view)
|
|
|
|
if allowed:
|
|
now = self.timer()
|
|
# Calculate the remaining limit and reset time
|
|
history = self.cache.get(self.key, [])
|
|
|
|
# Remove old histories
|
|
while history and history[-1] <= now - self.duration:
|
|
history.pop()
|
|
|
|
# Calculate the requests
|
|
num_requests = len(history)
|
|
|
|
# Check available requests
|
|
available = self.num_requests - num_requests
|
|
|
|
# Unix timestamp for when the rate limit will reset
|
|
reset_time = int(now + self.duration)
|
|
|
|
# Add headers
|
|
request.META['X-RateLimit-Remaining'] = max(0, available)
|
|
request.META['X-RateLimit-Reset'] = reset_time
|
|
|
|
return allowed |