feat: Implement rate limiting with custom throttling classes for different user types

apps/common/drf/throttling.py

@@ -0,0 +1,34 @@
+# -*- coding: utf-8 -*-
+from rest_framework.throttling import SimpleRateThrottle
+
+
+class RateThrottle(SimpleRateThrottle):
+
+    def __init__(self):
+        # Override the usual SimpleRateThrottle, because we can't determine
+        # the rate until called by the view.
+        pass
+
+    def allow_request(self, request, view):
+        if getattr(request, "user", None) and request.user.is_authenticated:
+            if getattr(request.user, "is_service_account", False):
+                self.scope = "service_account"
+            else:
+                self.scope = "user"
+        else:
+            self.scope = "anon"
+
+        self.rate = self.get_rate()
+        self.num_requests, self.duration = self.parse_rate(self.rate)
+        return super().allow_request(request, view)
+
+    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': self.scope,
+            'ident': ident
+        }

apps/jumpserver/api/dashboard.py

@@ -1,6 +1,5 @@
 from collections import defaultdict
 
-from django.core.cache import cache
 from django.db.models import Count, Max, F, CharField, Q
 from django.db.models.functions import Cast
 from django.http.response import JsonResponse
@@ -145,7 +144,6 @@ class DateTimeMixin:
 
 
 class DatesLoginMetricMixin:
-    days: int
     dates_list: list
     date_start_end: tuple
     command_type_queryset_list: list
@@ -157,8 +155,6 @@ class DatesLoginMetricMixin:
     operate_logs_queryset: OperateLog.objects
     password_change_logs_queryset: PasswordChangeLog.objects
 
-    CACHE_TIMEOUT = 60
-
     @lazyproperty
     def get_type_to_assets(self):
         result = Asset.objects.annotate(type=F('platform__type')). \
@@ -218,34 +214,19 @@ class DatesLoginMetricMixin:
         return date_metrics_dict.get('id', [])
 
     def get_dates_login_times_assets(self):
-        cache_key = f"stats:top10_assets:{self.days}"
-        data = cache.get(cache_key)
-        if data is not None:
-            return data
-
         assets = self.sessions_queryset.values("asset") \
             .annotate(total=Count("asset")) \
             .annotate(last=Cast(Max("date_start"), output_field=CharField())) \
             .order_by("-total")
-
+        return list(assets[:10])
-        result = list(assets[:10])
-        cache.set(cache_key, result, self.CACHE_TIMEOUT)
-        return result
 
     def get_dates_login_times_users(self):
-        cache_key = f"stats:top10_users:{self.days}"
-        data = cache.get(cache_key)
-        if data is not None:
-            return data
-
         users = self.sessions_queryset.values("user_id") \
             .annotate(total=Count("user_id")) \
             .annotate(user=Max('user')) \
             .annotate(last=Cast(Max("date_start"), output_field=CharField())) \
             .order_by("-total")
-        result = list(users[:10])
+        return list(users[:10])
-        cache.set(cache_key, result, self.CACHE_TIMEOUT)
-        return result
 
     def get_dates_login_record_sessions(self):
         sessions = self.sessions_queryset.order_by('-date_start')

apps/jumpserver/settings/libs.py

@@ -38,6 +38,14 @@ REST_FRAMEWORK = {
         "oauth2_provider.contrib.rest_framework.OAuth2Authentication",
         'authentication.backends.drf.SessionAuthentication',
     ),
+    'DEFAULT_THROTTLE_CLASSES': (
+        'common.drf.throttling.RateThrottle',
+    ),
+    'DEFAULT_THROTTLE_RATES': {
+        'anon': '60/min',
+        'user': '180/min',
+        'service_account': '300/min',
+    },
     'DEFAULT_FILTER_BACKENDS': (
         'django_filters.rest_framework.DjangoFilterBackend',
         'common.drf.filters.SearchFilter',