From 266a360a9742f4abbad07efdc84e896f15b5a093 Mon Sep 17 00:00:00 2001
From: fit2bot <68588906+fit2bot@users.noreply.github.com>
Date: Wed, 30 Aug 2023 17:31:59 +0800
Subject: [PATCH] =?UTF-8?q?feat:=20=E5=8F=AF=E6=9F=A5=E7=9C=8B=E5=85=A8?=
 =?UTF-8?q?=E5=B1=80=E8=A2=AB=E9=99=90=E5=88=B6=E7=9A=84ip=20=E5=B9=B6?=
 =?UTF-8?q?=E4=B8=94=E5=8F=AF=E4=BB=A5=E8=A7=A3=E9=94=81=20(#11394)?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Co-authored-by: feng <1304903146@qq.com>
---
 apps/settings/api/__init__.py         |  1 +
 apps/settings/api/security.py         | 60 +++++++++++++++++++++++++++
 apps/settings/serializers/security.py |  6 +++
 apps/settings/urls/api_urls.py        |  2 +
 4 files changed, 69 insertions(+)
 create mode 100644 apps/settings/api/security.py

diff --git a/apps/settings/api/__init__.py b/apps/settings/api/__init__.py
index 13f4a2b18..510925590 100644
--- a/apps/settings/api/__init__.py
+++ b/apps/settings/api/__init__.py
@@ -3,6 +3,7 @@ from .email import *
 from .feishu import *
 from .ldap import *
 from .public import *
+from .security import *
 from .settings import *
 from .sms import *
 from .vault import *
diff --git a/apps/settings/api/security.py b/apps/settings/api/security.py
new file mode 100644
index 000000000..bed98ed3b
--- /dev/null
+++ b/apps/settings/api/security.py
@@ -0,0 +1,60 @@
+from uuid import uuid4
+
+from django.conf import settings
+from django.core.cache import cache
+from rest_framework.generics import ListAPIView, CreateAPIView
+from rest_framework.views import Response
+
+from users.utils import LoginIpBlockUtil
+from ..serializers import SecurityBlockIPSerializer
+
+
+class BlockIPSecurityAPI(ListAPIView):
+    serializer_class = SecurityBlockIPSerializer
+    rbac_perms = {
+        'GET': 'settings.change_security',
+        'unlock': 'settings.change_security',
+    }
+
+    @staticmethod
+    def get_ips():
+        ips = []
+        prefix = LoginIpBlockUtil.BLOCK_KEY_TMPL.replace('{}', '')
+        keys = cache.keys(f'{prefix}*')
+        for key in keys:
+            ips.append(key.replace(prefix, ''))
+
+        white_list = settings.SECURITY_LOGIN_IP_WHITE_LIST
+        ips = list(set(ips) - set(white_list))
+        ips = settings.SECURITY_LOGIN_IP_BLACK_LIST
+        ips = [ip for ip in ips if ip != '*']
+        return ips
+
+    def get_page_offset_and_limit(self):
+        get_params = self.request.GET
+        offset = get_params.get('offset', 0)
+        limit = get_params.get('limit', 15)
+        return int(offset), int(limit)
+
+    def list(self, request, *args, **kwargs):
+        ips = self.get_ips()
+        offset, limit = self.get_page_offset_and_limit()
+        slice_ips = ips[offset:offset + limit]
+        data = [{'id': str(uuid4()), 'ip': ip} for ip in slice_ips]
+        ser = self.get_serializer(data, many=True)
+        data = {'count': len(ips), 'results': ser.data}
+        return Response(data=data, status=200)
+
+
+class UnlockIPSecurityAPI(CreateAPIView):
+    serializer_class = SecurityBlockIPSerializer
+    rbac_perms = {
+        'POST': 'settings.change_security',
+    }
+
+    def create(self, request, *args, **kwargs):
+        ips = request.data.get('ips')
+        prefix = LoginIpBlockUtil.BLOCK_KEY_TMPL.replace('{}', '')
+        for ip in ips:
+            LoginIpBlockUtil(f'{prefix}{ip}').clean_block_if_need()
+        return Response(status=200)
diff --git a/apps/settings/serializers/security.py b/apps/settings/serializers/security.py
index a5c27f1a1..7c9bf676a 100644
--- a/apps/settings/serializers/security.py
+++ b/apps/settings/serializers/security.py
@@ -7,6 +7,7 @@ __all__ = [
     'SecurityPasswordRuleSerializer', 'SecuritySessionSerializer',
     'SecurityAuthSerializer', 'SecuritySettingSerializer',
     'SecurityLoginLimitSerializer', 'SecurityBasicSerializer',
+    'SecurityBlockIPSerializer'
 ]
 
 
@@ -225,3 +226,8 @@ class SecuritySettingSerializer(
     SecurityLoginLimitSerializer,
 ):
     PREFIX_TITLE = _('Security')
+
+
+class SecurityBlockIPSerializer(serializers.Serializer):
+    id = serializers.UUIDField(required=False)
+    ip = serializers.CharField(max_length=1024, required=False, allow_blank=True)
diff --git a/apps/settings/urls/api_urls.py b/apps/settings/urls/api_urls.py
index ba074d4ee..97055bc83 100644
--- a/apps/settings/urls/api_urls.py
+++ b/apps/settings/urls/api_urls.py
@@ -20,6 +20,8 @@ urlpatterns = [
     path('sms/backend/', api.SMSBackendAPI.as_view(), name='sms-backend'),
     path('vault/testing/', api.VaultTestingAPI.as_view(), name='vault-testing'),
     path('vault/sync/', api.VaultSyncDataAPI.as_view(), name='vault-sync'),
+    path('security/block-ip/', api.BlockIPSecurityAPI.as_view(), name='block-ip'),
+    path('security/unlock-ip/', api.UnlockIPSecurityAPI.as_view(), name='unlock-ip'),
 
     path('setting/', api.SettingsApi.as_view(), name='settings-setting'),
     path('logo/', api.SettingsLogoApi.as_view(), name='settings-logo'),