47 lines
1.4 KiB
Python
47 lines
1.4 KiB
Python
from __future__ import annotations
|
|
|
|
from django.conf import settings
|
|
from rest_framework.exceptions import Throttled
|
|
from rest_framework.throttling import SimpleRateThrottle
|
|
|
|
|
|
class SettingsRateThrottle(SimpleRateThrottle):
|
|
setting_name = ""
|
|
default_rate = ""
|
|
|
|
def get_rate(self):
|
|
return getattr(settings, self.setting_name, self.default_rate)
|
|
|
|
|
|
class GenerateRateThrottle(SettingsRateThrottle):
|
|
scope = "generate"
|
|
setting_name = "API_GENERATE_THROTTLE_RATE"
|
|
default_rate = "60/min"
|
|
|
|
def get_cache_key(self, request, view):
|
|
if request.auth is not None:
|
|
ident = f"key:{request.auth.pk}"
|
|
elif request.user and request.user.is_authenticated:
|
|
ident = f"user:{request.user.pk}"
|
|
else:
|
|
ident = self.get_ident(request)
|
|
return self.cache_format % {"scope": self.scope, "ident": ident}
|
|
|
|
|
|
class ApiAuthFailureThrottle(SettingsRateThrottle):
|
|
scope = "api_auth_failure"
|
|
setting_name = "API_AUTH_FAILURE_THROTTLE_RATE"
|
|
default_rate = "30/min"
|
|
|
|
def get_cache_key(self, request, view):
|
|
return self.cache_format % {
|
|
"scope": self.scope,
|
|
"ident": self.get_ident(request),
|
|
}
|
|
|
|
|
|
def throttle_api_auth_failure(request) -> None:
|
|
throttle = ApiAuthFailureThrottle()
|
|
if not throttle.allow_request(request, None):
|
|
raise Throttled(wait=throttle.wait())
|