mirror of
https://github.com/jumpserver/jumpserver.git
synced 2026-05-14 19:03:31 +00:00
51 lines
1.5 KiB
Python
51 lines
1.5 KiB
Python
# -*- coding: utf-8 -*-
|
||
from rest_framework.throttling import SimpleRateThrottle
|
||
|
||
__all__ = ['RateThrottle', 'FileTransferThrottle']
|
||
|
||
|
||
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
|
||
}
|
||
|
||
|
||
class FileTransferThrottle(SimpleRateThrottle):
|
||
"""
|
||
文件上传下载限流，防止DOS攻击
|
||
"""
|
||
scope = 'file_transfer'
|
||
|
||
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}
|