fix: 白名单逻辑修改

apps/application/serializers/application_serializers.py

@@ -126,10 +126,15 @@ class ApplicationSerializer(serializers.Serializer):
                 ApplicationSerializer.Authentication(data={'access_token': self.data.get('token')}).auth()
             except Exception as e:
                 is_auth = 'false'
+            application_access_token = QuerySet(ApplicationAccessToken).filter(
+                access_token=self.data.get('token')).first()
             t = Template(content)
             s = t.render(
                 Context(
-                    {'is_auth': is_auth, 'protocol': 'http', 'host': 'localhost:8000', 'token': '0a8d892c755f1a75'}))
+                    {'is_auth': is_auth, 'protocol': 'http', 'host': 'localhost:8000', 'token': '0a8d892c755f1a75',
+                     'white_list_str': ",".join(
+                         application_access_token.white_list),
+                     'white_active': 'true' if application_access_token.white_active else 'false'}))
             response = HttpResponse(s, status=200, headers={'Content-Type': 'text/javascript'})
             set_embed_identity_cookie(request, response)
             return response

apps/application/template/embed.js

@@ -286,7 +286,10 @@ function initMaxkbStyle(){
 }
 
 function embedChatbot() {
-  if ({{is_auth}}) {
+  white_list_str='{{white_list_str}}'
+  white_list=white_list_str.split(',')
+
+  if ({{is_auth}}&&{{white_active}}?white_list.includes(window.location.origin):true) {
     // 初始化maxkb智能小助手
     initMaxkb()
   } else console.error('invalid parameter')

apps/application/views/application_views.py

@@ -21,7 +21,7 @@ from common.constants.permission_constants import CompareConstants, PermissionCo
 from common.exception.app_exception import AppAuthenticationFailed
 from common.response import result
 from common.swagger_api.common_api import CommonApi
-from common.util.common import query_params_to_single_dict, set_embed_identity_cookie
+from common.util.common import query_params_to_single_dict
 from dataset.serializers.dataset_serializers import DataSetSerializers
 
 
@@ -191,14 +191,12 @@ class Application(APIView):
                              tags=["应用/认证"],
                              security=[])
         def post(self, request: Request):
-            response = result.success(
+            return result.success(
                 ApplicationSerializer.Authentication(data={'access_token': request.data.get("access_token")}).auth(),
                 headers={"Access-Control-Allow-Origin": "*", "Access-Control-Allow-Credentials": "true",
                          "Access-Control-Allow-Methods": "POST",
                          "Access-Control-Allow-Headers": "Origin,Content-Type,Cookie,Accept,Token"}
             )
-            set_embed_identity_cookie(request, response)
-            return response
 
     @action(methods=['POST'], detail=False)
     @swagger_auto_schema(operation_summary="创建应用",

apps/application/views/chat_views.py

@@ -18,7 +18,7 @@ from common.auth import TokenAuth, has_permissions
 from common.constants.permission_constants import Permission, Group, Operate, \
     RoleConstants, ViewPermission, CompareConstants
 from common.response import result
-from common.util.common import query_params_to_single_dict, set_embed_identity_cookie
+from common.util.common import query_params_to_single_dict
 
 
 class ChatView(APIView):
@@ -71,13 +71,11 @@ class ChatView(APIView):
                                                            dynamic_tag=keywords.get('application_id'))])
         )
         def post(self, request: Request, chat_id: str):
-            response = ChatMessageSerializer(data={'chat_id': chat_id}).chat(request.data.get('message'),
+            return ChatMessageSerializer(data={'chat_id': chat_id}).chat(request.data.get('message'),
-                                                                             request.data.get(
+                                                                         request.data.get(
-                                                                                 're_chat') if 're_chat' in request.data else False,
+                                                                             're_chat') if 're_chat' in request.data else False,
-                                                                             request.data.get(
+                                                                         request.data.get(
-                                                                                 'stream') if 'stream' in request.data else True)
+                                                                             'stream') if 'stream' in request.data else True)
-            set_embed_identity_cookie(request, response)
-            return response
 
     @action(methods=['GET'], detail=False)
     @swagger_auto_schema(operation_summary="获取对话列表",

apps/common/auth/authenticate.py

@@ -6,14 +6,11 @@
     @date：2023/9/4 11:16
     @desc:  认证类
 """
-import datetime
 import traceback
-from urllib.parse import urlparse
 
 from django.core import cache
 from django.core import signing
 from django.db.models import QuerySet
-from ipware import get_client_ip
 from rest_framework.authentication import TokenAuthentication
 
 from application.models.api_key_model import ApplicationAccessToken, ApplicationApiKey
@@ -21,13 +18,10 @@ from common.constants.authentication_type import AuthenticationType
 from common.constants.permission_constants import Auth, get_permission_list_by_role, RoleConstants, Permission, Group, \
     Operate
 from common.exception.app_exception import AppAuthenticationFailed, AppEmbedIdentityFailed, AppChatNumOutOfBoundsFailed
-from common.util.common import getRestSeconds
-from common.util.rsa_util import decrypt
 from smartdoc.settings import JWT_AUTH
 from users.models.user import User, get_user_dynamics_permission
 
 token_cache = cache.caches['token_cache']
-chat_cache = cache.caches['chat_cache']
 
 
 class AnonymousAuthentication(TokenAuthentication):
@@ -87,35 +81,6 @@ class TokenAuth(TokenAuthentication):
                     raise AppAuthenticationFailed(1002, "身份验证信息不正确")
                 if not application_access_token.access_token == auth_details.get('access_token'):
                     raise AppAuthenticationFailed(1002, "身份验证信息不正确")
-                if application_access_token.white_active:
-                    referer = request.META.get('HTTP_REFERER')
-                    if referer is not None:
-                        client_ip = urlparse(referer).hostname
-                    else:
-                        client_ip = get_client_ip(request)
-                    if not application_access_token.white_list.__contains__(client_ip):
-                        raise AppAuthenticationFailed(1002, "身份验证信息不正确")
-                if 'embed_identity' in request.COOKIES and request.path.__contains__('/api/application/chat_message/'):
-                    embed_identity = request.COOKIES['embed_identity']
-                    try:
-                        # 如果无法解密 说明embed_identity并非系统颁发
-                        value = decrypt(embed_identity)
-                    except Exception as e:
-                        raise AppEmbedIdentityFailed(1004, '嵌入cookie不正确')
-                    embed_identity_number = chat_cache.get(value)
-                    if embed_identity_number is not None:
-                        if application_access_token.access_num <= embed_identity_number:
-                            raise AppChatNumOutOfBoundsFailed(1003, '访问次数超过今日访问量')
-                    # 对话次数+1
-                    try:
-                        if not chat_cache.incr(value):
-                            # 如果修改失败则设置为1
-                            chat_cache.set(value, 1,
-                                           timeout=getRestSeconds())
-                    except Exception as e:
-                        # 如果修改失败则设置为1 证明 key不存在
-                        chat_cache.add(value, 1,
-                                       timeout=getRestSeconds())
                 return application_access_token.application.user, Auth(
                     role_list=[RoleConstants.APPLICATION_ACCESS_TOKEN],
                     permission_list=[

apps/common/middleware/chat_cookie_middleware.py

@@ -0,0 +1,66 @@
+# coding=utf-8
+"""
+    @project: maxkb
+    @Author：虎
+    @file： chat_cookie_middleware.py
+    @date：2024/3/13 20:13
+    @desc:
+"""
+from django.core import cache
+from django.core import signing
+from django.db.models import QuerySet
+from django.utils.deprecation import MiddlewareMixin
+
+from application.models.api_key_model import ApplicationAccessToken
+from common.exception.app_exception import AppEmbedIdentityFailed
+from common.response import result
+from common.util.common import set_embed_identity_cookie, getRestSeconds
+from common.util.rsa_util import decrypt
+
+chat_cache = cache.caches['chat_cache']
+
+
+class ChatCookieMiddleware(MiddlewareMixin):
+
+    def process_response(self, request, response):
+        if request.path.startswith('/api/application/chat_message') or request.path.startswith(
+                '/api/application/authentication') or request.path.startswith('/api/application/profile'):
+            set_embed_identity_cookie(request, response)
+            if 'embed_identity' in request.COOKIES and request.path.__contains__('/api/application/chat_message/'):
+                embed_identity = request.COOKIES['embed_identity']
+                try:
+                    # 如果无法解密 说明embed_identity并非系统颁发
+                    value = decrypt(embed_identity)
+                except Exception as e:
+                    raise AppEmbedIdentityFailed(1004, '嵌入cookie不正确')
+                # 对话次数+1
+                try:
+                    if not chat_cache.incr(value):
+                        # 如果修改失败则设置为1
+                        chat_cache.set(value, 1,
+                                       timeout=getRestSeconds())
+                except Exception as e:
+                    # 如果修改失败则设置为1 证明 key不存在
+                    chat_cache.set(value, 1,
+                                   timeout=getRestSeconds())
+        return response
+
+    def process_request(self, request):
+        if 'embed_identity' in request.COOKIES and request.path.__contains__('/api/application/chat_message/'):
+            auth = request.META.get('HTTP_AUTHORIZATION', None
+                                    )
+            auth_details = signing.loads(auth)
+            application_access_token = QuerySet(ApplicationAccessToken).filter(
+                application_id=auth_details.get('application_id')).first()
+            embed_identity = request.COOKIES['embed_identity']
+            try:
+                # 如果无法解密 说明embed_identity并非系统颁发
+                value = decrypt(embed_identity)
+            except Exception as e:
+                return result.Result(1003,
+                                     message='访问次数超过今日访问量', response_status=460)
+            embed_identity_number = chat_cache.get(value)
+            if embed_identity_number is not None:
+                if application_access_token.access_num <= embed_identity_number:
+                    return result.Result(1003,
+                                         message='访问次数超过今日访问量', response_status=461)

apps/common/middleware/static_headers_middleware.py

@@ -0,0 +1,23 @@
+# coding=utf-8
+"""
+    @project: maxkb
+    @Author：虎
+    @file： static_headers_middleware.py
+    @date：2024/3/13 18:26
+    @desc:
+"""
+from django.db.models import QuerySet
+from django.utils.deprecation import MiddlewareMixin
+
+from application.models.api_key_model import ApplicationAccessToken
+
+
+class StaticHeadersMiddleware(MiddlewareMixin):
+    def process_response(self, request, response):
+        if request.path.startswith('/ui/chat/'):
+            access_token = request.path.replace('/ui/chat/', '')
+            application_access_token = QuerySet(ApplicationAccessToken).filter(access_token=access_token).first()
+            if application_access_token.white_active:
+                # 添加自定义的响应头
+                response['Content-Security-Policy'] = f'frame-ancestors {" ".join(application_access_token.white_list)}'
+        return response

apps/common/util/common.py

@@ -11,6 +11,7 @@ import importlib
 import uuid
 from functools import reduce
 from typing import Dict, List
+
 from django.core import cache
 
 from .rsa_util import encrypt

apps/smartdoc/settings/base.py

@@ -1,9 +1,9 @@
 import datetime
+import mimetypes
 import os
 from pathlib import Path
 
 from ..const import CONFIG, PROJECT_DIR
-import mimetypes
 
 mimetypes.add_type("text/css", ".css", True)
 mimetypes.add_type("text/javascript", ".js", True)
@@ -46,6 +46,8 @@ MIDDLEWARE = [
     'django.contrib.sessions.middleware.SessionMiddleware',
     'django.middleware.common.CommonMiddleware',
     'django.contrib.messages.middleware.MessageMiddleware',
+    'common.middleware.static_headers_middleware.StaticHeadersMiddleware',
+    'common.middleware.chat_cookie_middleware.ChatCookieMiddleware'
 
 ]