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

apps/accounts/serializers/account/account.py

@@ -341,10 +341,6 @@ class AssetAccountBulkSerializer(
 
     @staticmethod
     def _handle_update_create(vd, lookup):
-        ori = Account.objects.filter(**lookup).first()
-        if ori and ori.secret == vd.get('secret'):
-            return ori, False, 'skipped'
-
         instance, value = Account.objects.update_or_create(defaults=vd, **lookup)
         state = 'created' if value else 'updated'
         return instance, True, state

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/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',

apps/settings/tools/ping.py

@@ -12,6 +12,8 @@ from settings.utils import generate_ips
 
 # From /usr/include/linux/icmp.h; your milage may vary.
 ICMP_ECHO_REQUEST = 8  # Seems to be the same on Solaris.
+ICMPV6_ECHO_REQUEST = 128
+ICMPV6_ECHO_REPLY = 129
 
 
 def checksum(source_string):
@@ -41,7 +43,15 @@ def checksum(source_string):
     return answer
 
 
-def receive_one_ping(my_socket, id, timeout):
+def _get_icmp_header_offset(received_packet, family):
+    if family != socket.AF_INET6:
+        return 20
+    if received_packet and (received_packet[0] >> 4) == 6:
+        return 40
+    return 0
+
+
+def receive_one_ping(my_socket, id, timeout, family):
     """
     Receive the ping from the socket.
     """
@@ -55,11 +65,20 @@ def receive_one_ping(my_socket, id, timeout):
 
         time_received = time.time()
         received_packet, addr = my_socket.recvfrom(1024)
-        icmpHeader = received_packet[20:28]
-        type, code, checksum, packet_id, sequence = struct.unpack("bbHHh", icmpHeader)
+        header_offset = _get_icmp_header_offset(received_packet, family)
+        icmpHeader = received_packet[header_offset:header_offset + 8]
+        if len(icmpHeader) < 8:
+            continue
+        type, code, checksum, packet_id, sequence = struct.unpack("BBHHH", icmpHeader)
+        if family == socket.AF_INET6 and type != ICMPV6_ECHO_REPLY:
+            continue
         if packet_id == id:
             bytes = struct.calcsize("d")
-            time_sent = struct.unpack("d", received_packet[28: 28 + bytes])[0]
+            if len(received_packet) < header_offset + 8 + bytes:
+                continue
+            time_sent = struct.unpack(
+                "d", received_packet[header_offset + 8: header_offset + 8 + bytes]
+            )[0]
             return time_received - time_sent
 
         time_left -= how_long_in_select
@@ -67,11 +86,19 @@ def receive_one_ping(my_socket, id, timeout):
             return
 
 
-def send_one_ping(my_socket, dest_addr, id, psize):
+def send_one_ping(my_socket, dest_addr, id, psize, family):
     """
     Send one ping to the given >dest_addr<.
     """
-    dest_addr = socket.gethostbyname(dest_addr)
+    if family == socket.AF_INET6:
+        dest_addr = dest_addr
+        icmp_type = ICMPV6_ECHO_REQUEST
+    else:
+        if isinstance(dest_addr, tuple):
+            dest_addr = (dest_addr[0], 1)
+        else:
+            dest_addr = (socket.gethostbyname(dest_addr), 1)
+        icmp_type = ICMP_ECHO_REQUEST
 
     # Remove header size from packet size
     # psize = psize - 8
@@ -84,33 +111,45 @@ def send_one_ping(my_socket, dest_addr, id, psize):
     my_checksum = 0
 
     # Make a dummy heder with a 0 checksum.
-    header = struct.pack("bbHHh", ICMP_ECHO_REQUEST, 0, my_checksum, id, 1)
+    header = struct.pack("BBHHH", icmp_type, 0, my_checksum, id, 1)
     bytes = struct.calcsize("d")
     data = (psize - bytes) * b"Q"
     data = struct.pack("d", time.time()) + data
 
-    # Calculate the checksum on the data and the dummy header.
-    my_checksum = checksum(header + data)
+    if family != socket.AF_INET6:
+        # Calculate the checksum on the data and the dummy header.
+        my_checksum = checksum(header + data)
 
     # Now that we have the right checksum, we put that in. It's just easier
     # to make up a new header than to stuff it into the dummy.
     header = struct.pack(
-        "bbHHh", ICMP_ECHO_REQUEST, 0, socket.htons(my_checksum), id, 1
+        "BBHHH", icmp_type, 0, socket.htons(my_checksum), id, 1
     )
     packet = header + data
-    my_socket.sendto(packet, (dest_addr, 1))  # Don't know about the 1
+    my_socket.sendto(packet, dest_addr)
+
+
+def resolve_dest_addr(dest_addr):
+    addrinfos = socket.getaddrinfo(
+        dest_addr, None, socket.AF_UNSPEC, socket.SOCK_DGRAM
+    )
+    family, _, _, _, sockaddr = addrinfos[0]
+    return family, sockaddr
 
 
 def ping(dest_addr, timeout, psize, flag=0):
     """
     Returns either the delay (in seconds) or none on timeout.
     """
-    icmp = socket.getprotobyname("icmp")
+    family, dest_sockaddr = resolve_dest_addr(dest_addr)
+    if family == socket.AF_INET6:
+        icmp = socket.IPPROTO_ICMPV6
+        sock_type = socket.SOCK_DGRAM
+    else:
+        icmp = socket.getprotobyname("icmp")
+        sock_type = socket.SOCK_DGRAM if os.getuid() != 0 else socket.SOCK_RAW
     try:
-        if os.getuid() != 0:
-            my_socket = socket.socket(socket.AF_INET, socket.SOCK_DGRAM, icmp)
-        else:
-            my_socket = socket.socket(socket.AF_INET, socket.SOCK_RAW, icmp)
+        my_socket = socket.socket(family, sock_type, icmp)
     except socket.error as e:
         if e.errno == 1:
             # Operation not permitted
@@ -122,8 +161,8 @@ def ping(dest_addr, timeout, psize, flag=0):
     flag &= 0x00FF
     my_id = process_pre | flag
 
-    send_one_ping(my_socket, dest_addr, my_id, psize)
-    delay = receive_one_ping(my_socket, my_id, timeout)
+    send_one_ping(my_socket, dest_sockaddr, my_id, psize, family)
+    delay = receive_one_ping(my_socket, my_id, timeout, family)
 
     my_socket.close()
     return delay

apps/tickets/handlers/apply_asset.py

@@ -18,6 +18,7 @@ class Handler(BaseHandler):
             self._create_asset_permission()
 
     def _create_asset_permission(self):
+        self.ticket.refresh_from_db()
         org_id = self.ticket.org_id
         with tmp_to_org(org_id):
             asset_permission = AssetPermission.objects.filter(id=self.ticket.id).first()