diff --git a/modoboa/lib/throttle.py b/modoboa/lib/throttle.py index c52de265d..d972742be 100644 --- a/modoboa/lib/throttle.py +++ b/modoboa/lib/throttle.py @@ -1,8 +1,12 @@ from rest_framework.throttling import SimpleRateThrottle, UserRateThrottle from django.urls import resolve + class UserDdosPerView(SimpleRateThrottle): - """Custom Throttle class for rest_framework. The throttling is applied on a per view basis for authentificated users.""" + """ + Custom Throttle class for rest_framework. The throttling is + applied on a per view basis for authentificated users. + """ scope = 'ddos' @@ -17,29 +21,22 @@ class UserDdosPerView(SimpleRateThrottle): } -class UserLesserDdosUser(SimpleRateThrottle): - """Custom Throttle class for rest_framework. The throttling is applied on a per view basis for authentificated users.""" +class UserLesserDdosUser(UserDdosPerView): + """ + Custom Throttle class for rest_framework. The throttling is + applied on a per view basis for authentificated users. + """ scope = 'ddos_lesser' - def get_cache_key(self, request, view): - if request.user and request.user.is_authenticated: - ident = request.user.pk - else: - ident = self.get_ident(request) - return self.cache_format % { - 'scope': hash(resolve(request.path).url_name), - 'ident': ident - } - class LoginThrottle(SimpleRateThrottle): - """ Custom throttle to reset the cache counter on success. """ + """Custom throttle to reset the cache counter on success.""" scope = 'login' def get_cache_key(self, request, view): - return self.cache_format % { + return self.cache_format % { 'scope': self.scope, 'ident': self.get_ident(request) } @@ -64,16 +61,22 @@ class PasswordResetApplyThrottle(LoginThrottle): scope = 'password_recovery_apply' -class GetThrottleViewsetMixin(): - """Override default get_throttle behaviour to assign throttle classes to different actions.""" +class GetThrottleViewsetMixin: + """ + Override default get_throttle behaviour to assign throttle + classes to different actions. + """ def get_throttles(self): """Give lesser_ddos to GET type actions and ddos to others.""" throttles = [UserRateThrottle()] - - if self.action in ["list", "retrieve", "validate", "dns_detail", "me", "dns_detail", "applications", "structure"]: + actions = [ + "list", "retrieve", "validate", "dns_detail", "me", + "dns_detail", "applications", "structure" + ] + if self.action in actions: throttles.append(UserLesserDdosUser()) else: throttles.append(UserDdosPerView()) - return throttles \ No newline at end of file + return throttles diff --git a/modoboa/limits/api/v1/viewsets.py b/modoboa/limits/api/v1/viewsets.py index 2e3da9e5d..006cb7454 100644 --- a/modoboa/limits/api/v1/viewsets.py +++ b/modoboa/limits/api/v1/viewsets.py @@ -10,7 +10,8 @@ from modoboa.lib.throttle import GetThrottleViewsetMixin from . import serializers -class ResourcesViewSet(GetThrottleViewsetMixin, +class ResourcesViewSet( + GetThrottleViewsetMixin, mixins.RetrieveModelMixin, mixins.UpdateModelMixin, viewsets.GenericViewSet):